[docs] cosmos3: port_feedback.md (pitfalls/lessons) + native-reuse audit

Add the one new native piece for image-conditioned reasoning: a deepstack-capable
reasoner backbone forward. The other parts are already proven/reusable — the
Qwen3-VL vision_encoder (transformers, bit-exact vs framework) and the multimodal
prefill (masked-scatter of image features + get_rope_index, standard transformers
Qwen3-VL boilerplate, reusable like the tokenizer/encoder).

- `Cosmos3VFMTransformer.reason_forward(inputs_embeds, position_ids,
  deepstack_embeds, visual_pos_mask)`: runs the und (causal) layers over a single
  text(+image) sequence, injecting per-layer deepstack visual embeds at the visual
  positions for the first `len(deepstack_embeds)` layers (Qwen3-VL deepstack),
  then final `norm`. Mirrors the framework `reasoner_forward`.
- Parity (`test_cosmos3_reasoning_parity::test_deepstack_reasoner_forward_matches_framework`):
  native `reason_forward` vs framework `reasoner_forward` given identical
  inputs_embeds + positions + per-layer deepstack + visual mask — hidden states
  max abs diff = 0.0, mean = 0.0.

All five omni modalities are now bit-exact vs the framework (video/audio/action
gen, text reasoning, vision encoding, deepstack reasoner). Full suite: 150 passed.

.buildkite/pipeline.yml

@@ -436,7 +436,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 25m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: LoRA Training Tests"
                 env:
                   - TEST_TYPE=training_lora
@@ -481,6 +481,7 @@ steps:
                 - "fastvideo/layers/**"
                 - "fastvideo/worker/**"
                 - "fastvideo/entrypoints/**"
+                - "fastvideo/performance/**"
                 - "fastvideo/tests/performance/**"
                 - ".buildkite/performance-benchmarks/**"
                 - "pyproject.toml"

.buildkite/scripts/pr_test.sh

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

.github/scripts/gate_full_suite.sh

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

@@ -0,0 +1,122 @@
+#!/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,3 +47,6 @@ 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,6 +52,7 @@ jobs:
             core.setOutput('pr_sha', pr.head.sha);
             core.setOutput('pr_branch', pr.head.ref);
             core.setOutput('pr_number', String(prNumber));
+            core.setOutput('pr_title', pr.title);
 
       - name: Trigger Full Suite
         if: steps.perm.outputs.has_write == 'true'
@@ -60,6 +61,7 @@ jobs:
           PR_SHA: ${{ steps.label.outputs.pr_sha }}
           PR_BRANCH: ${{ steps.label.outputs.pr_branch }}
           PR_NUMBER: ${{ steps.label.outputs.pr_number }}
+          PR_TITLE: ${{ steps.label.outputs.pr_title }}
           BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
           BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
         run: |
@@ -71,6 +73,7 @@ jobs:
               --arg commit "$PR_SHA" \
               --arg branch "$PR_BRANCH" \
               --arg message "Full Suite for PR #${PR_NUMBER} (via /merge)" \
+              --arg pr_title "$PR_TITLE" \
               --argjson pr_id "$PR_NUMBER" \
               '{
                 commit: $commit,
@@ -80,11 +83,12 @@ jobs:
                 pull_request_id: $pr_id,
                 pull_request_base_branch: "main",
                 env: {
-                  TEST_SCOPE: "full",
-                  FULL_SUITE: "true",
-                  PR_NUMBER: ($pr_id | tostring)
-                }
-              }')"
+                    TEST_SCOPE: "full",
+                    FULL_SUITE: "true",
+                    PR_NUMBER: ($pr_id | tostring),
+                    PR_TITLE: $pr_title
+                  }
+                }')"
 
   parse-command:
     if: >-
@@ -241,6 +245,7 @@ jobs:
           TEST_SCOPE: ${{ needs.parse-command.outputs.test_scope }}
           FULL_SUITE: ${{ needs.parse-command.outputs.full_suite }}
           TEST_TYPE: ${{ needs.parse-command.outputs.test_type }}
+          PR_TITLE: ${{ github.event.issue.title }}
           BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
           BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
         run: |
@@ -257,6 +262,7 @@ jobs:
               --arg full_suite "$FULL_SUITE" \
               --arg test_type "$TEST_TYPE" \
               --arg pr_number "$PR_NUMBER" \
+              --arg pr_title "$PR_TITLE" \
               '{
                 commit: $commit,
                 branch: $branch,
@@ -266,8 +272,9 @@ jobs:
                 pull_request_base_branch: "main",
                 env: {
                   TEST_SCOPE: $test_scope,
-                  FULL_SUITE: $full_suite,
-                  TEST_TYPE: $test_type,
-                  PR_NUMBER: $pr_number
-                }
-              }')"
+                    FULL_SUITE: $full_suite,
+                    TEST_TYPE: $test_type,
+                    PR_NUMBER: $pr_number,
+                    PR_TITLE: $pr_title
+                  }
+                }')"

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

@@ -7,6 +7,7 @@ on:
 permissions:
   contents: read
   pull-requests: read
+  actions: read
 
 concurrency:
   group: full-suite-${{ github.event.pull_request.number }}
@@ -18,6 +19,8 @@ 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
@@ -49,6 +52,20 @@ 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:
@@ -56,6 +73,7 @@ jobs:
           PR_SHA: ${{ github.event.pull_request.head.sha }}
           PR_BRANCH: ${{ github.event.pull_request.head.ref }}
           PR_NUMBER: ${{ github.event.pull_request.number }}
+          PR_TITLE: ${{ github.event.pull_request.title }}
           BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
           BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
         run: |
@@ -67,6 +85,7 @@ jobs:
               --arg commit "$PR_SHA" \
               --arg branch "$PR_BRANCH" \
               --arg message "Full Suite for PR #${PR_NUMBER}" \
+              --arg pr_title "$PR_TITLE" \
               --argjson pr_id "$PR_NUMBER" \
               '{
                 commit: $commit,
@@ -78,6 +97,7 @@ jobs:
                 env: {
                   TEST_SCOPE: "full",
                   FULL_SUITE: "true",
-                  PR_NUMBER: ($pr_id | tostring)
+                  PR_NUMBER: ($pr_id | tostring),
+                  PR_TITLE: $pr_title
                 }
               }')"

.gitignore

@@ -37,6 +37,11 @@ logs/
 official_weights/
 converted_weights/
 
+# Cosmos3 local parity assets (symlinked from main worktree)
+/official_weights/
+/converted_weights/
+/cosmos-framework
+
 # SSIM test outputs
 fastvideo/tests/ssim/generated_videos/
 **/.cache/**
@@ -72,8 +77,7 @@ docs/distillation/examples/
 # Python pickle files
 *.pkl
 
-# Reference videos
-!fastvideo/tests/ssim/reference_videos/**/*.mp4
+# Reference videos (negations must come after the catch-all on line below)
 
 # Static images
 !docs/assets/images/**/*.png
@@ -127,6 +131,8 @@ apps/dreamverse/web/.env.production.local
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
+!fastvideo/tests/ssim/reference_videos/**/*.mp4
+!fastvideo/tests/ssim/reference_videos/**/*.png
 
 # Editor logs and local Python version pins (accidentally committed)
 *.nvimlog

assets/eval/worldmodel_synthetic_flow_calibration.json

@@ -0,0 +1,22 @@
+{
+  "alpha_yaw": 0.08734091699186919,
+  "alpha_pitch": 0.08169667696275307,
+  "alpha_turn": 5.724587470723463e-17,
+  "beta_fwd": 0.02842768078408099,
+  "beta_strafe": 0.022531015077067108,
+  "focal_length": 457.0,
+  "frame_shape": [
+    352,
+    640
+  ],
+  "calibrated_from": [
+    "1_wasd_only",
+    "camera",
+    "camera4hold_alpha1",
+    "fully_random",
+    "wasdonly_alpha1",
+    "wasd4holdrandview_simple_1key1mouse1"
+  ],
+  "residual_rms": 15.890399609478676,
+  "n_equations": 4125232
+}

docs/contributing/ci_architecture.md

@@ -99,6 +99,13 @@ 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/testing.md

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

docs/design/inference_schema_parity_inventory.yaml

@@ -458,6 +458,8 @@ surfaces:
       guidance_scale: request.sampling.guidance_scale
       guidance_scale_2: request.sampling.guidance_scale_2
       guidance_rescale: request.sampling.guidance_rescale
+      use_embedded_guidance: request.sampling.use_embedded_guidance
+      true_cfg_scale: request.sampling.true_cfg_scale
       boundary_ratio: request.sampling.boundary_ratio
       sigmas: request.sampling.sigmas
       enable_teacache: request.runtime.enable_teacache

docs/distillation/anyflow.md

@@ -0,0 +1,116 @@
+# 🌊 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/optimizations.md

@@ -27,7 +27,6 @@ This page describes the various options for speeding up generation times in Fast
 - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
 - Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`
-- FlashInfer SageAttention3 (SM120 NVFP4): `FASTVIDEO_ATTENTION_BACKEND=FLASHINFER_SAGE_ATTN3`
 - Attn-QAT inference (modified SageAttention3 FP4, sm_120/RTX 5090): `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER`
 - Video MoBA Attention: `FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN`
 - Sparse Linear Attention: `FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN`
@@ -238,60 +237,6 @@ If you are using `uv` and `torch==2.8.0`, make sure that
 
 To use Sage Attention 3 in FastVideo, follow the `README.md` in the linked repository to install the package from source.
 
-### FlashInfer SageAttention3
-
-**`FLASHINFER_SAGE_ATTN3`**
-
-Wraps [FlashInfer](https://github.com/flashinfer-ai/flashinfer)'s dense SM120
-NVFP4 attention (`nvfp4_attention_sm120_quantize_qkv` +
-`nvfp4_attention_sm120_fwd`, added in
-[flashinfer#3640](https://github.com/flashinfer-ai/flashinfer/pull/3640)).
-
-```bash
-FASTVIDEO_ATTENTION_BACKEND=FLASHINFER_SAGE_ATTN3 python example.py
-```
-
-#### Kernel requirements
-
-- Compute capability **12.0 only** (consumer Blackwell, e.g. RTX 5090), CUDA 12.8+
-- `head_dim` 64 or 128, fp16/bf16, equal-shape q/k/v (self-attention)
-- `flashinfer-python` installed (the kernel is JIT-built on first use)
-
-Where any of these do not hold — including when flashinfer is not installed at
-all — the backend logs one warning and transparently falls back to torch SDPA,
-so it can run anywhere.
-
-#### Attention shape logging
-
-The fallback exists on purpose: this backend doubles as a shape-collection
-vehicle for tuning the SageAttention3 kernel on SM 12.1 (DGX Spark). Set
-
-```bash
-export FASTVIDEO_ATTENTION_BACKEND=FLASHINFER_SAGE_ATTN3   # or SAGE_ATTN_THREE
-export FASTVIDEO_ATTN_SHAPE_LOG=/tmp/attn_shapes.jsonl
-```
-
-and run your normal workload (no GPU/kernel required). Each process writes
-`/tmp/attn_shapes.jsonl.rank<N>` (or `.pid<N>` when not distributed) with one
-JSON line per unique attention shape:
-
-```json
-{"backend": "FLASHINFER_SAGE_ATTN3", "batch": 1, "num_q_heads": 12,
- "num_kv_heads": 12, "seq_len_q": 75600, "seq_len_kv": 75600, "head_dim": 128,
- "dtype": "torch.bfloat16", "causal": false, "sm_scale": 0.0884, "count": 1200,
- "torch": "2.8.0", "device": "NVIDIA GB10", "sm": "12.1"}
-```
-
-Counts accumulate per shape; the file is flushed periodically and at exit.
-Hand the JSONL files to the kernel folks. The `SAGE_ATTN_THREE` backend records
-to the same log, so either implementation can be used for collection.
-
-A collected log can directly drive the unified attention kernel benchmark
-(per-shape latency, TFLOPS, and accuracy for the bf16 baseline
-(flash-attn/SDPA), SageAttention FP8, the local SageAttention3 port, and
-FlashInfer NVFP4):
-`python fastvideo-kernel/benchmarks/bench_attn_kernels.py --shapes-jsonl /tmp/attn_shapes.jsonl.pid<N>`.
-
 ### V-MoBA / SLA / SageSLA
 
 These backends are model-specific and require the corresponding kernels and

docs/training/train_infra.md

@@ -161,6 +161,18 @@ training:
     decay_interval_steps: 0
 ```
 
+`training.data.data_path` can also mix multiple preprocessed datasets by using a mapping from dataset path to repeat count:
+
+```yaml
+training:
+  data:
+    data_path:
+      data/zeldam2-clean: 1
+      data/multi3d_games: 2
+```
+
+The repeat count duplicates that dataset's parquet file list before shuffling/sampling, so the example above trains with roughly twice as much `multi3d_games` exposure as `zeldam2-clean`. Paths are just suggested locations; use any local path that contains a FastVideo preprocessed parquet dataset.
+
 See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
 including SwanLab installation and authentication.
 
@@ -326,6 +338,40 @@ Self-Forcing inherits all DMD2 parameters, plus:
 | `enable_gradient_in_rollout` | `true` | Enable backprop through rollout |
 | `start_gradient_frame` | `0` | Frame index where gradients begin |
 
+### Streaming Long Tuning
+
+`StreamingLongTuningMethod` extends Self-Forcing for LongLive-style rollouts. It
+keeps a streaming state, generates overlapping chunks, and trains only the new
+frames while preserving context from earlier chunks.
+
+For the MatrixGame2/Zelda world-model example, self-forcing and long tuning are
+separate runs: first train or load the 1k-step self-forcing checkpoint using
+`examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml`,
+then run
+`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
+from that checkpoint for the 3k-step streaming long-tuning stage.
+
+```yaml
+method:
+  _target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
+  streaming_chunk_size: 9
+  streaming_max_length: 39
+  streaming_fixed_overlap_latents: 3
+  streaming_reencode_overlap_anchor: true
+  streaming_anchor_inject_k: 1
+  streaming_require_full_blocks: true
+  multi_phased_distill_schedule:
+    - stage: streaming_long
+      start_step: 0
+      end_step: 3000
+      num_latent_t: 39
+      streaming_training: true
+```
+
+See
+`examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml`
+for a complete MatrixGame2/Zelda configuration.
+
 ---
 
 ## Callbacks

examples/inference/basic/basic_cosmos3_i2v_new_api.py

@@ -0,0 +1,81 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+# NVIDIA Cosmos3-Nano omni world model — image-to-video (I2V) path through
+# FastVideo's native Cosmos3 pipeline. The input image conditions latent frame 0
+# (kept clean during denoising); the rest of the clip is generated to follow it.
+# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g.
+# ``official_weights/cosmos3``) to skip the Hugging Face download.
+OUTPUT_PATH = "video_samples_cosmos3_i2v"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+    image_path = os.environ.get("COSMOS3_IMAGE_PATH", "assets/images/cyclist.jpg")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                text_encoder=True,
+                pin_cpu_memory=True,
+                dit=False,
+                vae=False,
+            ),
+        ),
+    )
+
+    load_start_time = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start_time
+
+    prompt = (
+        "A mountain biker rides forward along the sunlit forest trail, wheels "
+        "kicking up dust as trees and dappled light sweep past, smooth cinematic "
+        "tracking shot from behind."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        inputs=InputConfig(image_path=image_path),
+        sampling=SamplingConfig(
+            # cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps;
+            # overridable via env for quick smoke runs.
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    start_time = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_cosmos3_new_api.py

@@ -0,0 +1,77 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+# NVIDIA Cosmos3-Nano omni world model — this example exercises the
+# text-to-video (T2V) path through FastVideo's native Cosmos3 pipeline.
+# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g.
+# ``official_weights/cosmos3``) to skip the Hugging Face download.
+OUTPUT_PATH = "video_samples_cosmos3"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                text_encoder=True,
+                pin_cpu_memory=True,
+                dit=False,
+                vae=False,
+            ),
+        ),
+    )
+
+    load_start_time = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start_time
+
+    prompt = (
+        "A golden retriever puppy runs across a sunlit meadow toward the camera, "
+        "ears flopping and wildflowers swaying in the breeze. Shallow depth of "
+        "field, warm afternoon light, smooth cinematic tracking shot."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            # cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps;
+            # overridable via env for quick smoke runs.
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    start_time = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_cosmos3_t2i_new_api.py

@@ -0,0 +1,78 @@
+import os
+import time
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+# NVIDIA Cosmos3-Nano omni world model — text-to-image (T2I) path through
+# FastVideo's native Cosmos3 pipeline. T2I is the single-frame case
+# (num_frames=1); the canonical Cosmos3 T2I resolution is 960x960 (the model's
+# "720" bucket, UniPC flow_shift=10.0). Point COSMOS3_MODEL_PATH at a local
+# diffusers checkpoint (e.g. ``official_weights/cosmos3``) to skip the download.
+OUTPUT_PATH = "video_samples_cosmos3_t2i"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                text_encoder=True,
+                pin_cpu_memory=True,
+                dit=False,
+                vae=False,
+            ),
+        ),
+    )
+
+    load_start_time = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start_time
+
+    prompt = (
+        "A photograph of a red panda sitting on a mossy log in a misty bamboo "
+        "forest, soft golden morning light filtering through the leaves, shallow "
+        "depth of field, crisp fur detail, serene atmosphere."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            # T2I is single-frame; canonical Cosmos3 T2I is 960x960. Overridable
+            # via env for quick smoke runs.
+            num_frames=1,
+            height=int(os.environ.get("COSMOS3_HEIGHT", "960")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "960")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(
+            output_path=OUTPUT_PATH,
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+    start_time = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate image: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_cosmos3_t2vs_new_api.py

@@ -0,0 +1,67 @@
+import os
+import time
+
+# t2vs (text -> video + sound). The Cosmos3 denoise stage generates a joint
+# [vision | sound] latent and AVAE-decodes the sound to a waveform muxed into the
+# mp4. The joint-sound path is gated on COSMOS3_T2VS (set here for the example).
+os.environ.setdefault("COSMOS3_T2VS", "1")
+
+from fastvideo import VideoGenerator  # noqa: E402
+from fastvideo.api import (  # noqa: E402
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    SamplingConfig,
+)
+
+OUTPUT_PATH = "video_samples_cosmos3_t2vs"
+
+
+def main():
+    model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano")
+
+    generator_config = GeneratorConfig(
+        model_path=model_name,
+        engine=EngineConfig(
+            num_gpus=1,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(text_encoder=True, pin_cpu_memory=True, dit=False, vae=False),
+        ),
+    )
+
+    load_start = time.perf_counter()
+    generator = VideoGenerator.from_config(generator_config)
+    load_time = time.perf_counter() - load_start
+
+    prompt = (
+        "Ocean waves crash against a rocky shore at sunset, white foam spraying "
+        "into the air as seagulls wheel overhead. Golden light, cinematic wide "
+        "shot, the rhythmic roar of the surf."
+    )
+    request = GenerationRequest(
+        prompt=prompt,
+        sampling=SamplingConfig(
+            num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")),
+            height=int(os.environ.get("COSMOS3_HEIGHT", "704")),
+            width=int(os.environ.get("COSMOS3_WIDTH", "1280")),
+            num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")),
+            guidance_scale=6.0,
+            fps=24,
+            seed=1024,
+        ),
+        output=OutputConfig(output_path=OUTPUT_PATH, save_video=True, return_frames=False),
+    )
+
+    start = time.perf_counter()
+    result = generator.generate(request)
+    gen_time = time.perf_counter() - start
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video+sound: {gen_time} seconds")
+    print(f"Output written to: {result.video_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_flux_dev.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+import re
+
+DEFAULT_PROMPTS = [
+    "a photo of a cat",
+    (
+        "a cinematic photo of a red panda wearing a tiny backpack, standing on a "
+        "rainy neon-lit street at night, shallow depth of field, sharp focus, "
+        "35mm, bokeh"
+    ),
+]
+
+
+def _safe_filename(text: str, max_len: int = 100) -> str:
+    """Make a stable, filesystem-friendly filename base."""
+    s = text[:max_len].strip()
+    s = s.replace(os.sep, "_")
+    if os.altsep:
+        s = s.replace(os.altsep, "_")
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
+    s = s.strip(" .")
+    return s or "prompt"
+
+
+def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
+    """Delete prior outputs so reruns do not get _1, _2 suffixes."""
+    if not os.path.isdir(out_dir):
+        return
+
+    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
+    for fn in os.listdir(out_dir):
+        if pattern.match(fn):
+            with contextlib.suppress(FileNotFoundError):
+                os.remove(os.path.join(out_dir, fn))
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.",
+    )
+    p.add_argument(
+        "--model-path",
+        default="official_weights/FLUX.1-dev",
+        help="Local Diffusers checkpoint dir or HF repo id.",
+    )
+    p.add_argument(
+        "--out-dir",
+        "--outdir",
+        default="outputs/flux_dev/samples",
+        help="Directory for saved PNG outputs.",
+    )
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=None,
+        help="Prompt. Repeat for multiple images.",
+    )
+    p.add_argument(
+        "--backend",
+        default=None,
+        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA).",
+    )
+    p.add_argument("--seed", type=int, default=42, help="Base seed; each prompt uses seed + index.")
+    p.add_argument("--height", type=int, default=1024, help="Output height.")
+    p.add_argument("--width", type=int, default=1024, help="Output width.")
+    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
+    p.add_argument("--guidance", type=float, default=3.5, help="Guidance scale.")
+    p.add_argument("--num-gpus", type=int, default=1, help="GPU count.")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    prompts: list[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
+
+    if args.backend:
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
+
+    from fastvideo import VideoGenerator
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    init_kwargs = {
+        "num_gpus": args.num_gpus,
+        "workload_type": "t2i",
+        "sp_size": 1,
+        "tp_size": 1,
+        "dit_cpu_offload": False,
+        "dit_layerwise_offload": False,
+        "text_encoder_cpu_offload": False,
+        "vae_cpu_offload": False,
+        "image_encoder_cpu_offload": False,
+        "pin_cpu_memory": False,
+        "use_fsdp_inference": False,
+    }
+
+    generator = VideoGenerator.from_pretrained(
+        model_path=args.model_path,
+        **init_kwargs,
+    )
+    try:
+        for i, prompt in enumerate(prompts):
+            seed = args.seed + i
+            filename_base = (
+                f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
+            )
+            _remove_existing_outputs(args.out_dir, filename_base)
+            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
+            print(f"[flux] prompt_idx={i} seed={seed} output_path={output_path}")
+
+            generation_kwargs = {
+                "output_path": output_path,
+                "height": args.height,
+                "width": args.width,
+                "num_frames": 1,
+                "fps": 1,
+                "num_inference_steps": args.steps,
+                "guidance_scale": args.guidance,
+                "use_embedded_guidance": True,
+                "true_cfg_scale": 1.0,
+                "seed": seed,
+                "save_video": True,
+            }
+
+            generator.generate_video(prompt, **generation_kwargs)
+
+        print(f"[flux] done. outputs written to: {args.out_dir}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_glm_image.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run GLM-Image text-to-image generation through FastVideo.
+
+User story:
+    "I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
+    text-to-image generation command, saved as a PNG."
+"""
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
+    parser.add_argument(
+        "--model-path",
+        default="zai-org/GLM-Image",
+        help="HF id or local diffusers-format GLM-Image weights directory.",
+    )
+    parser.add_argument(
+        "--output",
+        default="image_output/landscape.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default=("A beautiful landscape photography with rolling hills, "
+                 "a winding river, and a vibrant sunset in the background. "
+                 "Warm golden light, photorealistic style."),
+        help="Text prompt.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.5)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
+    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
+
+    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
+    # pipeline class come from the model's registered defaults — don't override.
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        trust_remote_code=True,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=str(output.parent),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+        result = generator.generate(request)
+        if isinstance(result, list):
+            result = result[0]
+
+        frames = result.frames
+        if frames is not None and len(frames):
+            Image.fromarray(frames[0]).save(output)
+            print(f"Saved image to {output}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/edit_glm_image.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run GLM-Image image-to-image (edit) generation through FastVideo.
+
+User story:
+    "I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
+    want a minimal edit command (text + image -> edited image), saved as a PNG."
+
+GLM-Image is a single unified pipeline: passing a condition image switches it
+from text-to-image to the edit path (the condition enters the DiT via a KV-cache
+write pass), so the generator config is identical to `basic_glm_image.py` — the
+`inputs.pil_image` on the request is what selects the edit mode.
+"""
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
+    parser.add_argument(
+        "--model-path",
+        default="zai-org/GLM-Image",
+        help="HF id or local diffusers-format GLM-Image weights directory.",
+    )
+    parser.add_argument(
+        "--image",
+        default="assets/images/couple.jpg",
+        help="Condition image to edit.",
+    )
+    parser.add_argument(
+        "--output",
+        default="image_output/edited.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="Change the background to a snowy mountain landscape at golden hour.",
+        help="Edit instruction.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=1.5)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    condition = Image.open(args.image).convert("RGB")
+    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
+    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
+
+    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
+    # pipeline class come from the model's registered defaults — don't override.
+    # The pipeline is registered as t2i; passing inputs.pil_image below switches
+    # it to the edit path.
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        trust_remote_code=True,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            inputs=InputConfig(pil_image=condition),
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=args.guidance_scale,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=str(output.parent),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+        result = generator.generate(request)
+        if isinstance(result, list):
+            result = result[0]
+
+        frames = result.frames
+        if frames is not None and len(frames):
+            Image.fromarray(frames[0]).save(output)
+            print(f"Saved image to {output}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/optimizations/FastWan_QAD_TAEHV.py

@@ -46,11 +46,8 @@ OUTPUT_PATH = "video_samples"
 # guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass
 # ``--distilled_model ''`` to run the base weights instead.
 DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B"
-# The HF repo moved the transformer weights; try current layout first,
-# then the legacy path (pre-2026-07 uploads).
-DISTILLED_WEIGHTS_FILES = (
-    "transformer/diffusion_pytorch_model.safetensors",
-    "generator_inference_transformer/diffusion_pytorch_model.safetensors",
+DISTILLED_WEIGHTS_FILE = (
+    "generator_inference_transformer/diffusion_pytorch_model.safetensors"
 )
 
 # TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
@@ -104,14 +101,7 @@ def resolve_distilled_weights(hf_id: str) -> str:
     if os.path.exists(hf_id):
         return hf_id
     from huggingface_hub import hf_hub_download
-    from huggingface_hub.errors import EntryNotFoundError
-    for filename in DISTILLED_WEIGHTS_FILES:
-        try:
-            return hf_hub_download(repo_id=hf_id, filename=filename)
-        except EntryNotFoundError:
-            continue
-    raise FileNotFoundError(
-        f"None of {DISTILLED_WEIGHTS_FILES} found in {hf_id!r}")
+    return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)
 
 
 @contextlib.contextmanager

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

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

@@ -0,0 +1,83 @@
+# 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,9 +100,6 @@ callbacks:
     sampling_steps: [4]
     sampling_timesteps: [1000, 750, 500, 250]
     num_frames: 81
-    # Validation/inference uses standard CFG in both clean and Self-Forcing,
-    # so this directly matches Self-Forcing guidance_scale=3.0.
-    guidance_scale: 3.0
 
 pipeline:
   flow_shift: 5

examples/train/scenario/worldmodel/README.md

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

examples/train/scenario/worldmodel/solaris/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

@@ -0,0 +1,94 @@
+# Diffusion-Forcing SFT: Zelda world model I2V Causal
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
+  chunk_size: 3
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path:
+      data/zeldam2-clean: 1
+    dataloader_num_workers: 1
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 9
+    num_height: 480
+    num_width: 832
+    num_frames: 33
+
+  optimizer:
+    learning_rate: 2.0e-5
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 60000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/matrixgame_finetune/checkpoints/zelda_causal_dfsft
+    training_state_checkpointing_steps: 5000
+    checkpoints_total_limit: 3
+
+  tracker:
+    entity: hapo-exp
+    project_name: mg_1.3b_zelda
+    run_name: zelda_causal_dfsft
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 200
+    sampling_steps: [40]
+    sampling_timesteps: [1000, 975, 950, 925, 900, 875, 850, 825, 800, 775,
+                         750, 725, 700, 675, 650, 625, 600, 575, 550, 525,
+                         500, 475, 450, 425, 400, 375, 350, 325, 300, 275,
+                         250, 225, 200, 175, 150, 125, 100, 75, 50, 25]
+    num_frames: 33
+    overlay_actions: true
+    guidance_scale: 6.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1

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

@@ -0,0 +1,88 @@
+# Matrix-Game 2.0 Zelda + multi-game I2V finetune.
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: FastVideo/Matrix-Game-2.0-Base-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path:
+      data/zeldam2-clean: 1
+      # data/multi3d_games: 10
+    dataloader_num_workers: 1
+    train_batch_size: 1
+    training_cfg_rate: 0.0  # unused for MatrixGame2 I2V; no text_embedding CFG dropout
+    seed: 42
+    num_latent_t: 9
+    num_height: 480
+    num_width: 832
+    num_frames: 33
+
+  optimizer:
+    learning_rate: 2.0e-5
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 60000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/matrixgame_finetune/checkpoints/zelda_with_mg_init
+    training_state_checkpointing_steps: 5000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: mg_1.3b_zelda
+    run_name: zelda_with_mg_init
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_i2v_pipeline.MatrixGame2I2VPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 200
+    sampling_steps: [40]
+    num_frames: 33
+    overlay_actions: true
+    guidance_scale: 6.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5

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

@@ -0,0 +1,120 @@
+# Self-forcing distillation: Zelda world model I2V Causal
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
+    init_from: mignonjia/mg_causal_zelda
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+  warp_denoising_step: true
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+
+  chunk_size: 3
+  student_sample_type: sde
+  same_step_across_blocks: true
+  last_step_only: false
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  # Critic optimizer
+  fake_score_learning_rate: 3.0e-7
+  fake_score_betas: [0.9, 0.95]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/zeldam2-clean
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1001
+    num_latent_t: 9
+    num_height: 480
+    num_width: 832
+    num_frames: 33
+
+  optimizer:
+    learning_rate: 3.0e-6
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/zelda_causal_self_forcing
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 1
+
+  tracker:
+    project_name: wangame_sf
+    run_name: mg2_self_forcing_9_latents
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  ema:
+    decay: 0.99
+    start_iter: 200
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 100
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 153
+    overlay_actions: true
+    keyboard_value_scale: 1.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1

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

@@ -0,0 +1,140 @@
+# MatrixGame2 I2V LongLive-style streaming distillation: Zelda world model I2V Causal
+# Student init from self forcing after 1k steps
+
+models:
+  student:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
+    init_from: mignonjia/mg_sf_distilled_zelda_1k_steps 
+    trainable: true
+    # transformer_override_safetensor: outputs/matrixgame_dmd/checkpoints/mg_zelda_sf_m2/checkpoint-1000_weight_only/ema/generator_ema.safetensors
+  teacher:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
+    init_from: mignonjia/mg_bidirectional_zelda
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+  warp_denoising_step: true
+  min_timestep_ratio: 0.02
+  max_timestep_ratio: 0.98
+
+  chunk_size: 3
+  student_sample_type: sde
+  same_step_across_blocks: true
+  last_step_only: false
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  streaming_training: true
+  streaming_chunk_size: 9
+  streaming_max_length: 39
+  streaming_fixed_overlap_latents: 3
+  streaming_reencode_overlap_anchor: true
+  streaming_anchor_inject_k: 1
+  streaming_require_full_blocks: true
+
+  multi_phased_distill_schedule:
+    - stage: streaming_long
+      start_step: 0
+      end_step: 3000
+      num_latent_t: 39
+      streaming_training: true
+      streaming_chunk_size: 9
+      streaming_max_length: 39
+      streaming_fixed_overlap_latents: 3
+
+  # Critic optimizer
+  fake_score_learning_rate: 3.0e-7
+  fake_score_betas: [0.9, 0.95]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/zeldam2-clean
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1001
+    num_latent_t: 39
+    num_height: 480
+    num_width: 832
+    num_frames: 153
+
+  optimizer:
+    learning_rate: 3.0e-6
+    betas: [0.9, 0.95]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 3000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/zelda_causal_long_tuning
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 2
+
+  tracker:
+    project_name: wangame_sf
+    run_name: mg2_39only_streaming_long
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  dit_precision: fp32
+
+callbacks:
+  ema:
+    decay: 0.99
+    start_iter: 200
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: data/zelda_validation_data/validation_zelda.json
+    every_steps: 100
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 153
+    overlay_actions: true
+    keyboard_value_scale: 1.0
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - vbench.temporal_flickering
+        - vbench.motion_smoothness
+        - vbench.subject_consistency
+        - vbench.background_consistency
+        - vbench.dynamic_degree
+        - optical_flow.synthetic_optical_flow
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+pipeline:
+  flow_shift: 5
+  dit_config:
+    local_attn_size: 6
+    sink_size: 1

fastvideo-kernel/benchmarks/bench_attn_kernels.py

@@ -1,542 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-"""Benchmark the four attention kernels: speed, accuracy, TFLOPS.
-
-Lanes (each probes availability + arch gate and skips cleanly, with a
-printed reason, where it cannot run):
-
-  bf16_flash_attn / bf16_sdpa
-                    High-precision baseline: flash-attn (flash_attn_func)
-                    when importable, else torch SDPA — the lane name says
-                    which one is running. (Prefix follows --dtype.)
-  fp8_sageattn      ``sageattention.sageattn`` — the attention op activated by
-                    examples/inference/optimizations/fp8_wan2_1_1_3b.py
-                    (PR #1496, FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN →
-                    fastvideo/attention/backends/sage_attn.py). sageattn
-                    auto-dispatches per-arch; on sm89+ it selects the
-                    qk-int8 / pv-fp8 CUDA kernels.
-  nvfp4_local       fastvideo-kernel/attn_qat_infer (local SageAttention3
-                    port, ``sageattn_blackwell``). Built for sm_120a only.
-  nvfp4_flashinfer  flashinfer.nvfp4_attention_sm120_quantize_qkv +
-                    nvfp4_attention_sm120_fwd. Compute capability 12.0/12.1,
-                    head_dim 64/128.
-
-Accuracy for every lane is measured against an fp32 torch-SDPA ground truth
-(chunked over query length so the score matrix never materializes whole)
-computed on the same fixed-seed inputs. Reported: cosine similarity,
-max abs err, mean rel err.
-
-Inputs are BIASED by default: q and k get a nonzero mean shift (--bias, v
-stays zero-mean so a constant output offset doesn't inflate cosine
-similarity). Zero-mean randn is blind to mean-dependent quantization errors —
-biased inputs caught flashinfer's former per-block-mean qk-correction bug
-(fixed by flashinfer#3838). Use
---include-zeromean to also run the zero-mean rows for contrast.
-
-All lanes run at the [B, H, L, D] interface.
-
-Usage:
-  python bench_attn_kernels.py                          # default shapes
-  python bench_attn_kernels.py --shapes 1x12x32760x128,1x24x4096x64
-  python bench_attn_kernels.py --causal --json out.json
-  python bench_attn_kernels.py --shapes-jsonl /tmp/attn_shapes.jsonl.pid1234
-      # shapes collected via FASTVIDEO_ATTN_SHAPE_LOG (fastvideo/attention/
-      # shape_logger.py); recorded causal/dtype/sm_scale are honored.
-"""
-
-from __future__ import annotations
-
-import argparse
-import dataclasses
-import inspect
-import json
-import statistics
-import sys
-from collections.abc import Callable
-from pathlib import Path
-
-import torch
-
-_REPO_ROOT = Path(__file__).resolve().parents[2]
-
-# ---------------------------------------------------------------------------
-# Pure helpers (CPU-unit-tested in fastvideo/tests/attention/
-# test_attn_kernels_bench.py) — keep them torch-free where possible.
-# ---------------------------------------------------------------------------
-
-
-@dataclasses.dataclass(frozen=True)
-class Shape:
-    batch: int
-    num_q_heads: int
-    num_kv_heads: int
-    seq_len_q: int
-    seq_len_kv: int
-    head_dim: int
-    causal: bool
-    dtype: str = "torch.bfloat16"
-    sm_scale: float | None = None
-
-    def label(self) -> str:
-        base = f"{self.batch}x{self.num_q_heads}x{self.seq_len_q}x{self.head_dim}"
-        if self.seq_len_kv != self.seq_len_q or self.num_kv_heads != self.num_q_heads:
-            base += f" (kv {self.num_kv_heads}x{self.seq_len_kv})"
-        if self.causal:
-            base += " causal"
-        return base
-
-
-def attn_tflops(shape: Shape, ms: float) -> float:
-    """TFLOPS of one dense attention forward.
-
-    FLOPs = 4 * B * H * Lq * Lk * D:
-      QK^T is B*H*Lq*Lk*D multiply-adds = 2*B*H*Lq*Lk*D FLOPs, and P@V is
-      another 2*B*H*Lq*Lk*D. Causal masking computes ~half the score matrix,
-      so FLOPs are halved. TFLOPS = FLOPs / (ms * 1e9).
-    """
-    flops = 4.0 * shape.batch * shape.num_q_heads * shape.seq_len_q * shape.seq_len_kv * shape.head_dim
-    if shape.causal:
-        flops *= 0.5
-    return flops / (ms * 1e9)
-
-
-def parse_shapes_cli(spec: str, causal: bool, dtype: str) -> list[Shape]:
-    """Parse --shapes 'BxHxLxD,BxHxLxD,...' (self-attention shapes)."""
-    shapes = []
-    for part in spec.split(","):
-        try:
-            b, h, length, d = (int(x) for x in part.lower().split("x"))
-        except ValueError as e:
-            raise argparse.ArgumentTypeError(f"bad shape {part!r}, expected BxHxLxD") from e
-        shapes.append(Shape(batch=b, num_q_heads=h, num_kv_heads=h, seq_len_q=length,
-                            seq_len_kv=length, head_dim=d, causal=causal, dtype=dtype))
-    return shapes
-
-
-def parse_shapes_jsonl(paths: list[str], default_dtype: str) -> list[Shape]:
-    """Parse fastvideo/attention/shape_logger.py output (one JSON object per
-    line; keys: batch, num_q_heads, num_kv_heads, seq_len_q, seq_len_kv,
-    head_dim, dtype, causal, sm_scale, count, ...). Dedupes across files,
-    preserving first-seen order."""
-    shapes: list[Shape] = []
-    seen: set[Shape] = set()
-    for path in paths:
-        with open(path) as f:
-            for line in f:
-                line = line.strip()
-                if not line:
-                    continue
-                rec = json.loads(line)
-                shape = Shape(
-                    batch=int(rec["batch"]),
-                    num_q_heads=int(rec["num_q_heads"]),
-                    num_kv_heads=int(rec.get("num_kv_heads", rec["num_q_heads"])),
-                    seq_len_q=int(rec["seq_len_q"]),
-                    seq_len_kv=int(rec.get("seq_len_kv", rec["seq_len_q"])),
-                    head_dim=int(rec["head_dim"]),
-                    causal=bool(rec.get("causal", False)),
-                    dtype=str(rec.get("dtype") or default_dtype),
-                    sm_scale=rec.get("sm_scale"),
-                )
-                if shape not in seen:
-                    seen.add(shape)
-                    shapes.append(shape)
-    return shapes
-
-
-@dataclasses.dataclass(frozen=True)
-class LaneSpec:
-    name: str
-    exact_cap: tuple[int, int] | None = None  # kernel runs on exactly this SM
-    caps: tuple[tuple[int, int], ...] | None = None  # kernel runs only on these SMs
-    min_cap: tuple[int, int] | None = None
-    head_dims: tuple[int, ...] | None = None  # None = any
-    same_shape: bool = False  # requires q/k/v equal shapes (self-attention)
-    same_heads: bool = False  # requires num_q_heads == num_kv_heads
-
-
-def skip_reason(spec: LaneSpec, shape: Shape, cap: tuple[int, int] | None,
-                unavailable: str | None) -> str | None:
-    """Why this lane cannot run this shape (None = it can)."""
-    if unavailable:
-        return unavailable
-    if cap is None:
-        return "no CUDA device"
-    if spec.exact_cap is not None and cap != spec.exact_cap:
-        return (f"requires SM {spec.exact_cap[0]}.{spec.exact_cap[1]}, "
-                f"device is SM {cap[0]}.{cap[1]}")
-    if spec.caps is not None and cap not in spec.caps:
-        allowed = "/".join(f"{a}.{b}" for a, b in spec.caps)
-        return f"requires SM {allowed}, device is SM {cap[0]}.{cap[1]}"
-    if spec.min_cap is not None and cap < spec.min_cap:
-        return (f"requires SM >= {spec.min_cap[0]}.{spec.min_cap[1]}, "
-                f"device is SM {cap[0]}.{cap[1]}")
-    if spec.head_dims is not None and shape.head_dim not in spec.head_dims:
-        return f"head_dim {shape.head_dim} not in {list(spec.head_dims)}"
-    if spec.same_shape and (shape.seq_len_q != shape.seq_len_kv or shape.num_q_heads != shape.num_kv_heads):
-        return "requires equal q/k/v shapes (self-attention)"
-    if spec.same_heads and shape.num_q_heads != shape.num_kv_heads:
-        return "requires num_q_heads == num_kv_heads"
-    return None
-
-
-def shape_workspace_bytes(shape: Shape, dtype_bytes: int, ref_chunk: int) -> int:
-    """Rough peak device memory to bench one shape: q/k/v in the bench dtype,
-    their fp32 copies + fp32 reference output for the ground truth, and the
-    fp32 score-matrix chunk (x2 for the math-backend softmax temporary)."""
-    numel_q = shape.batch * shape.num_q_heads * shape.seq_len_q * shape.head_dim
-    numel_kv = shape.batch * shape.num_kv_heads * shape.seq_len_kv * shape.head_dim
-    lq = shape.seq_len_q if shape.causal else min(ref_chunk, shape.seq_len_q)
-    scores = shape.batch * shape.num_q_heads * lq * shape.seq_len_kv * 4
-    return (numel_q + 2 * numel_kv) * (dtype_bytes + 4) + numel_q * 4 + 2 * scores
-
-
-def format_table(headers: list[str], rows: list[list[str]]) -> str:
-    """Aligned plain-text table (left-justified columns)."""
-    cols = list(zip(headers, *rows)) if rows else [(h,) for h in headers]
-    widths = [max(len(str(c)) for c in col) for col in cols]
-    lines = [
-        "  ".join(str(h).ljust(w) for h, w in zip(headers, widths)),
-        "  ".join("-" * w for w in widths),
-    ]
-    for row in rows:
-        lines.append("  ".join(str(c).ljust(w) for c, w in zip(row, widths)))
-    return "\n".join(lines)
-
-
-# ---------------------------------------------------------------------------
-# Lanes (GPU)
-# ---------------------------------------------------------------------------
-
-AttnFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor, bool, float], torch.Tensor]
-
-_REF_CHUNK = 2048
-
-
-@dataclasses.dataclass
-class Lane:
-    spec: LaneSpec
-    fn: AttnFn | None  # None when unavailable
-    unavailable: str | None = None
-
-
-def _run_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
-    kwargs: dict = {"is_causal": causal, "scale": scale}
-    if q.shape[1] != k.shape[1]:
-        kwargs["enable_gqa"] = True
-    return torch.nn.functional.scaled_dot_product_attention(q, k, v, **kwargs)
-
-
-def _ref_sdpa_fp32(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
-    """fp32 SDPA ground truth, chunked over query length.
-
-    Where flash/mem-efficient SDPA reject fp32 (e.g. SM121), torch falls back
-    to the math backend, which materializes the full [B, H, Lq, Lkv] score
-    matrix -- ~48 GiB for the 32760-token default shape (host OOM-kill on
-    unified-memory devices like GB10). Non-causal attention is exact under
-    query chunking; causal shapes stay unchunked (per-chunk is_causal would
-    mis-align the mask), so they keep the old memory ceiling.
-    """
-    if causal or q.shape[2] <= _REF_CHUNK:
-        return _run_sdpa(q, k, v, causal, scale)
-    return torch.cat([_run_sdpa(q[:, :, i:i + _REF_CHUNK], k, v, False, scale)
-                      for i in range(0, q.shape[2], _REF_CHUNK)], dim=2)
-
-
-def _load_baseline(dtype_tag: str) -> Lane:
-    """High-precision baseline: flash-attn when importable, else torch SDPA.
-
-    The lane name records which implementation is actually timed.
-    """
-    try:
-        from flash_attn import flash_attn_func  # noqa: PLC0415
-    except ImportError:
-        return Lane(LaneSpec(f"{dtype_tag}_sdpa"), _run_sdpa)
-
-    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
-        # flash_attn_func takes [B, L, H, D]; only the head dim must be
-        # contiguous, so plain transposes are accepted without a copy.
-        out = flash_attn_func(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
-                              softmax_scale=scale, causal=causal)
-        return out.transpose(1, 2)
-
-    # flash-attn's fp16/bf16 kernels need sm80+.
-    return Lane(LaneSpec(f"{dtype_tag}_flash_attn", min_cap=(8, 0)), fn)
-
-
-def _load_fp8_sageattn() -> Lane:
-    # sageattn's fp8 (qk-int8 / pv-fp8) kernels need sm89+; the package
-    # dispatches per-arch internally.
-    spec = LaneSpec("fp8_sageattn", min_cap=(8, 9), same_heads=True)
-    try:
-        from sageattention import sageattn  # noqa: PLC0415
-    except ImportError as e:
-        return Lane(spec, None, f"sageattention not installed ({e})")
-
-    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
-        return sageattn(q, k, v, tensor_layout="HND", is_causal=causal, sm_scale=scale)
-
-    return Lane(spec, fn)
-
-
-def _load_nvfp4_local(single_level_p_quant: bool) -> Lane:
-    spec = LaneSpec("nvfp4_local", exact_cap=(12, 0), head_dims=(64, 128), same_shape=True)
-    for p in (_REPO_ROOT, _REPO_ROOT / "fastvideo-kernel", _REPO_ROOT / "fastvideo-kernel" / "python"):
-        if str(p) not in sys.path:
-            sys.path.insert(0, str(p))
-    try:
-        from attn_qat_infer import sageattn_blackwell  # noqa: PLC0415
-    except ImportError as e:
-        return Lane(spec, None, f"attn_qat_infer import failed ({e}); build fastvideo-kernel with sm_120a")
-
-    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
-        # sageattn_blackwell pads L to a multiple of 128 internally and trims.
-        return sageattn_blackwell(q, k, v, is_causal=causal,
-                                  single_level_p_quant=single_level_p_quant, sm_scale=scale)
-
-    return Lane(spec, fn)
-
-
-def _load_nvfp4_flashinfer() -> Lane:
-    # SM121 (GB10 / DGX Spark) is enabled by flashinfer#3897.
-    spec = LaneSpec("nvfp4_flashinfer", caps=((12, 0), (12, 1)), head_dims=(64, 128), same_shape=True)
-    try:
-        from flashinfer import (  # noqa: PLC0415
-            nvfp4_attention_sm120_fwd, nvfp4_attention_sm120_quantize_qkv)
-    except ImportError as e:
-        return Lane(spec, None, f"flashinfer nvfp4_attention_sm120 not installed ({e})")
-
-    quant_kwargs = {}
-    if "per_block_mean" in inspect.signature(nvfp4_attention_sm120_quantize_qkv).parameters:
-        quant_kwargs["per_block_mean"] = True
-
-    def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor:
-        seq_len = q.shape[2]
-        packed = nvfp4_attention_sm120_quantize_qkv(q, k, v, **quant_kwargs)
-        out, _lse = nvfp4_attention_sm120_fwd(*packed, sm_scale=scale, causal=causal, out_dtype=q.dtype)
-        # flashinfer pads seq_len to a multiple of 128; trim it back off.
-        return out[:, :, :seq_len, :]
-
-    return Lane(spec, fn)
-
-
-def build_lanes(args: argparse.Namespace) -> list[Lane]:
-    return [
-        _load_baseline(args.dtype),
-        _load_fp8_sageattn(),
-        _load_nvfp4_local(single_level_p_quant=args.p_quant == "single"),
-        _load_nvfp4_flashinfer(),
-    ]
-
-
-# ---------------------------------------------------------------------------
-# Measurement
-# ---------------------------------------------------------------------------
-
-_DTYPES = {"torch.bfloat16": torch.bfloat16, "torch.float16": torch.float16,
-           "bf16": torch.bfloat16, "fp16": torch.float16}
-
-
-def make_qkv(shape: Shape, seed: int, bias: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Fixed-seed q/k/v; q and k get a +bias mean shift. v stays zero-mean so
-    a constant output offset doesn't inflate the cosine-similarity column."""
-    gen = torch.Generator(device="cuda").manual_seed(seed)  # fixed seed per shape
-    dtype = _DTYPES.get(shape.dtype, torch.bfloat16)
-
-    def rand(heads: int, seq: int, shift: float) -> torch.Tensor:
-        t = torch.randn(shape.batch, heads, seq, shape.head_dim,
-                        generator=gen, device="cuda", dtype=dtype)
-        return t + shift if shift else t
-
-    return (rand(shape.num_q_heads, shape.seq_len_q, bias),
-            rand(shape.num_kv_heads, shape.seq_len_kv, bias),
-            rand(shape.num_kv_heads, shape.seq_len_kv, 0.0))
-
-
-def accuracy_stats(out: torch.Tensor, ref: torch.Tensor) -> dict[str, float]:
-    out = out.float()
-    diff = (out - ref).abs()
-    cos = torch.nn.functional.cosine_similarity(out.flatten(), ref.flatten(), dim=0)
-    return {
-        "cos_sim": cos.item(),
-        "max_abs_err": diff.max().item(),
-        # mean relative error: mean |out-ref| normalized by the reference's
-        # mean magnitude (elementwise |out-ref|/|ref| blows up near ref~0).
-        "mean_rel_err": (diff.mean() / ref.abs().mean().clamp_min(1e-12)).item(),
-    }
-
-
-def time_fn(fn: Callable[[], torch.Tensor], warmup: int, iters: int) -> dict[str, float]:
-    """CUDA-event latency stats (ms) over `iters` timed calls."""
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    times = []
-    for _ in range(iters):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-        start.record()
-        fn()
-        end.record()
-        torch.cuda.synchronize()
-        times.append(start.elapsed_time(end))
-    return {"mean_ms": statistics.fmean(times), "std_ms": statistics.pstdev(times),
-            "median_ms": statistics.median(times)}
-
-
-def bench_shape(shape: Shape, bias: float, lanes: list[Lane], cap: tuple[int, int] | None,
-                args: argparse.Namespace) -> list[dict]:
-    dist = f"biased(+{bias:g})" if bias else "zeromean"
-    results = []
-    runnable = [(lane, skip_reason(lane.spec, shape, cap, lane.unavailable)) for lane in lanes]
-
-    oom_reason = None
-    if any(reason is None for _, reason in runnable):
-        dtype_bytes = 2  # bf16/fp16
-        need = shape_workspace_bytes(shape, dtype_bytes, _REF_CHUNK)
-        free = torch.cuda.mem_get_info()[0]
-        if need > 0.9 * free:
-            oom_reason = (f"would OOM: needs ~{need / 2**30:.1f} GiB workspace, "
-                          f"{free / 2**30:.1f} GiB free")
-            print(f"  SKIP shape: {oom_reason}")
-
-    q = k = v = ref = None
-    if oom_reason is None and any(reason is None for _, reason in runnable):
-        q, k, v = make_qkv(shape, args.seed, bias)
-        scale = shape.sm_scale or shape.head_dim**-0.5
-        # fp32 SDPA ground truth (accuracy baseline only; never timed).
-        ref = _ref_sdpa_fp32(q.float(), k.float(), v.float(), shape.causal, scale)
-
-    for lane, reason in runnable:
-        rec: dict = {"shape": shape.label(), "dist": dist, "lane": lane.spec.name,
-                     **dataclasses.asdict(shape)}
-        reason = reason or oom_reason
-        if reason is not None:
-            rec.update(status="skip", reason=reason)
-            results.append(rec)
-            continue
-        scale = shape.sm_scale or shape.head_dim**-0.5
-        try:
-            out = lane.fn(q, k, v, shape.causal, scale)
-            rec.update(accuracy_stats(out, ref))
-            del out
-            rec.update(status="ok", **time_fn(lambda: lane.fn(q, k, v, shape.causal, scale),
-                                              args.warmup, args.iters))
-            rec["tflops"] = attn_tflops(shape, rec["mean_ms"])
-        except Exception as e:  # lane failed at runtime: report, keep going
-            rec.update(status="fail", reason=f"{type(e).__name__}: {e}")
-        results.append(rec)
-    del q, k, v, ref
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    return results
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-# The real logged workload (FASTVIDEO_ATTN_SHAPE_LOG on 5090/GB10/GB200 runs of
-# Wan2.1-1.3B 480p, 81 frames: latent grid 21x30x52 = 32760 tokens, 12 heads x
-# head_dim 128, bf16, non-causal) plus a seq-len sweep at the same head
-# geometry. 32760 appears once (deduped).
-DEFAULT_SHAPE_DIMS = [
-    (1, 12, 1024, 128),
-    (1, 12, 4096, 128),
-    (1, 12, 16384, 128),
-    (1, 12, 32760, 128),  # Wan2.1-1.3B 480p self-attention (the logged shape)
-]
-
-
-def default_shapes(causal: bool, dtype: str) -> list[Shape]:
-    return [
-        Shape(batch=b, num_q_heads=h, num_kv_heads=h, seq_len_q=length, seq_len_kv=length,
-              head_dim=d, causal=causal, dtype=dtype) for b, h, length, d in DEFAULT_SHAPE_DIMS
-    ]
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    p = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0],
-                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    p.add_argument("--shapes", metavar="BxHxLxD,...",
-                   help="comma-separated self-attention shapes; replaces the default suite")
-    p.add_argument("--shapes-jsonl", nargs="+", metavar="PATH",
-                   help="attention shape log(s) written via FASTVIDEO_ATTN_SHAPE_LOG "
-                   "(fastvideo/attention/shape_logger.py); replaces the default shapes")
-    p.add_argument("--causal", action="store_true",
-                   help="run --shapes/default shapes causally (jsonl shapes keep their recorded flag)")
-    p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16",
-                   help="input dtype for default shapes / jsonl records without one")
-    p.add_argument("--bias", type=float, default=1.0,
-                   help="mean shift added to q and k (biased inputs expose mean-dependent "
-                   "quantization bugs that zero-mean randn hides)")
-    p.add_argument("--include-zeromean", action="store_true",
-                   help="additionally run every shape with zero-mean inputs for contrast")
-    p.add_argument("--p-quant", choices=["single", "two"], default="single",
-                   help="nvfp4_local P-matrix quantization: single-level or two-level "
-                   "(sageattn_blackwell single_level_p_quant knob)")
-    p.add_argument("--warmup", type=int, default=10, help="warmup iterations per lane")
-    p.add_argument("--iters", type=int, default=50, help="timed iterations per lane")
-    p.add_argument("--seed", type=int, default=0, help="input seed (fixed per shape)")
-    p.add_argument("--json", metavar="PATH", help="also dump results as JSON")
-    return p.parse_args(argv)
-
-
-def _fmt(rec: dict, key: str, spec: str) -> str:
-    return format(rec[key], spec) if key in rec else "-"
-
-
-def main(argv: list[str] | None = None) -> None:
-    args = parse_args(argv)
-    dtype_name = f"torch.{'bfloat16' if args.dtype == 'bf16' else 'float16'}"
-    if args.shapes:
-        shapes = parse_shapes_cli(args.shapes, args.causal, dtype_name)
-    elif args.shapes_jsonl:
-        shapes = parse_shapes_jsonl(args.shapes_jsonl, dtype_name)
-    else:
-        shapes = default_shapes(args.causal, dtype_name)
-
-    cap = None
-    meta: dict = {"torch": torch.__version__, "warmup": args.warmup, "iters": args.iters,
-                  "seed": args.seed, "bias": args.bias}
-    if torch.cuda.is_available():
-        cap = torch.cuda.get_device_capability(0)
-        meta.update(device=torch.cuda.get_device_name(0), sm=f"{cap[0]}.{cap[1]}")
-        print(f"Device: {meta['device']} (SM {meta['sm']}), torch {torch.__version__}")
-    else:
-        print("No CUDA device: all lanes will be skipped (shape/skip dry run).")
-
-    lanes = build_lanes(args)
-    biases = [args.bias] + ([0.0] if args.include_zeromean and args.bias else [])
-    results = []
-    for shape in shapes:
-        for bias in biases:
-            print(f"\n== {shape.label()}  dtype={shape.dtype}  bias={bias:g} ==")
-            results.extend(bench_shape(shape, bias, lanes, cap, args))
-
-    headers = ["shape", "dist", "lane", "status", "mean_ms", "std_ms", "TFLOPS",
-               "cos_sim", "max_abs_err", "mean_rel_err"]
-    rows = []
-    for rec in results:
-        status = rec["status"] if rec["status"] == "ok" else f"{rec['status']}: {rec['reason']}"
-        rows.append([
-            rec["shape"], rec["dist"], rec["lane"], status,
-            _fmt(rec, "mean_ms", ".3f"),
-            _fmt(rec, "std_ms", ".3f"),
-            _fmt(rec, "tflops", ".1f"),
-            _fmt(rec, "cos_sim", ".6f"),
-            _fmt(rec, "max_abs_err", ".4e"),
-            _fmt(rec, "mean_rel_err", ".4e"),
-        ])
-    print("\n" + format_table(headers, rows))
-    print("\nAccuracy is vs chunked fp32 torch-SDPA ground truth on identical inputs "
-          "(the fp32 reference itself is never timed).")
-
-    if args.json:
-        with open(args.json, "w") as f:
-            json.dump({"meta": meta, "results": results}, f, indent=2)
-        print(f"Wrote {args.json}")
-
-
-if __name__ == "__main__":
-    main()

fastvideo/api/flux.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from fastvideo.api.sampling_param import SamplingParam
+
+
+@dataclass
+class FluxSamplingParam(SamplingParam):
+
+    prompt: str | None = "a photo of a cat"
+    negative_prompt: str = ""
+
+    num_videos_per_prompt: int = 1
+    seed: int = 0
+
+    num_frames: int = 1
+    height: int = 1024
+    width: int = 1024
+    fps: int = 1
+
+    num_inference_steps: int = 28
+    guidance_scale: float = 3.5
+    use_embedded_guidance: bool = True
+    true_cfg_scale: float = 1.0

fastvideo/api/sampling_param.py

@@ -90,6 +90,10 @@ class SamplingParam:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
+    # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
+    use_embedded_guidance: bool = False
+    # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
+    true_cfg_scale: float = 1.0
     guidance_rescale: float = 0.0
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
@@ -325,6 +329,18 @@ class SamplingParam:
             default=SamplingParam.guidance_rescale,
             help="Guidance rescale factor",
         )
+        parser.add_argument(
+            "--use-embedded-guidance",
+            action="store_true",
+            default=SamplingParam.use_embedded_guidance,
+            help="Use embedded guidance scale (FLUX-style) instead of classic CFG",
+        )
+        parser.add_argument(
+            "--true-cfg-scale",
+            type=float,
+            default=SamplingParam.true_cfg_scale,
+            help="True CFG scale for FLUX when > 1 (requires negative prompt encoding)",
+        )
         parser.add_argument(
             "--boundary-ratio",
             type=float,

fastvideo/api/schema.py

@@ -150,6 +150,7 @@ class SamplingConfig:
     guidance_scale_2: float | None = None
     guidance_rescale: float = 0.0
     true_cfg_scale: float | None = None
+    use_embedded_guidance: bool | None = None
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
 

fastvideo/attention/backends/flash_attn.py

@@ -1,13 +1,15 @@
 # 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 import envs
+from fastvideo.attention.utils.flash_attn_default import (
+    fa_version,
+    flash_attn_func_compilable,
+)
+
 from fastvideo.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
@@ -17,119 +19,6 @@ 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.
@@ -309,9 +198,17 @@ 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,
-                flash_attn_varlen_qk_no_pad,
+                flash_attn_no_pad_compilable as flash_attn_no_pad,
+                flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
             )
 
             attn_mask = attn_metadata.attn_mask

fastvideo/attention/backends/flashinfer_sageattn3.py

@@ -1,159 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""FlashInfer SageAttention3 (SM120 dense NVFP4 attention) backend.
-
-Wraps ``flashinfer.nvfp4_attention_sm120_quantize_qkv`` +
-``flashinfer.nvfp4_attention_sm120_fwd`` (flashinfer-ai/flashinfer#3640).
-The kernel is gated to compute capability 12.0 (consumer Blackwell, e.g.
-RTX 5090), head_dim 64/128, fp16/bf16, and self-attention with equal
-q/k/v shapes. flashinfer pads seq_len up to a multiple of 128 internally;
-this backend trims the padding back off the output. Note flashinfer's own
-accuracy tests only cover seq_len multiples of 128 — padded-key softmax
-mass for other lengths is flashinfer's semantics, unverified here.
-
-Anywhere the kernel cannot run (flashinfer missing, non-SM120 device,
-unsupported head size/dtype/shape) the backend logs one warning and falls
-back to torch SDPA, so it stays usable as a shape-collection vehicle
-(``FASTVIDEO_ATTN_SHAPE_LOG``) on machines without the kernel — e.g.
-DGX Spark (SM121), which the collected shape data is being gathered for.
-
-Works with both ``LocalAttention`` and ``DistributedAttention`` (sequence
-parallel): inputs arrive as [B, L, H, D] and are transposed internally.
-"""
-
-import torch
-
-from fastvideo.attention import shape_logger
-from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
-                                                   AttentionMetadataBuilder)
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-try:
-    from flashinfer import (nvfp4_attention_sm120_fwd, nvfp4_attention_sm120_quantize_qkv)
-except ImportError:
-    nvfp4_attention_sm120_fwd = None  # type: ignore[assignment]
-    nvfp4_attention_sm120_quantize_qkv = None  # type: ignore[assignment]
-
-_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
-_SUPPORTED_HEAD_DIMS = (64, 128)
-
-_fallback_warned = False
-
-
-def _fall_back(reason: str) -> bool:
-    """Log the first fallback reason once per process; returns False."""
-    global _fallback_warned
-    if not _fallback_warned:
-        logger.warning("FLASHINFER_SAGE_ATTN3: %s. Falling back to torch SDPA "
-                       "(subsequent fallbacks are silent).", reason)
-        _fallback_warned = True
-    return False
-
-
-def _can_use_flashinfer(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout: float) -> bool:
-    """Whether the flashinfer SM120 kernel can serve this call ([B, H, L, D] inputs)."""
-    if nvfp4_attention_sm120_fwd is None:
-        return _fall_back("flashinfer with nvfp4_attention_sm120 is not installed")
-    if not q.is_cuda:
-        return _fall_back(f"device {q.device} is not CUDA")
-    cap = torch.cuda.get_device_capability(q.device)
-    # SM121 (GB10 / DGX Spark) is enabled by flashinfer#3897.
-    if cap not in ((12, 0), (12, 1)):
-        return _fall_back(f"compute capability {cap[0]}.{cap[1]} is not 12.0/12.1 (kernel is SM120/SM121-only)")
-    if q.dtype not in _SUPPORTED_DTYPES:
-        return _fall_back(f"dtype {q.dtype} is not fp16/bf16")
-    if q.shape[-1] not in _SUPPORTED_HEAD_DIMS:
-        return _fall_back(f"head_dim {q.shape[-1]} not in {_SUPPORTED_HEAD_DIMS}")
-    if not (q.shape == k.shape == v.shape):
-        return _fall_back(f"q/k/v shapes differ (q={tuple(q.shape)}, k={tuple(k.shape)}, "
-                          f"v={tuple(v.shape)}); kernel requires equal-shape self-attention")
-    if dropout > 0:
-        return _fall_back(f"dropout_p={dropout} is not supported by the kernel")
-    return True
-
-
-class FlashInferSageAttention3Backend(AttentionBackend):
-
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [64, 128]
-
-    @staticmethod
-    def get_name() -> str:
-        return "FLASHINFER_SAGE_ATTN3"
-
-    @staticmethod
-    def get_impl_cls() -> type["FlashInferSageAttention3Impl"]:
-        return FlashInferSageAttention3Impl
-
-    @staticmethod
-    def get_metadata_cls() -> type["AttentionMetadata"]:
-        raise NotImplementedError
-
-    @staticmethod
-    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
-        raise NotImplementedError
-
-
-class FlashInferSageAttention3Impl(AttentionImpl):
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-        self.dropout = extra_impl_args.get("dropout_p", 0.0)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        """Input is [B, L, H, D] (FastVideo impl convention); output matches."""
-        if shape_logger.enabled:
-            shape_logger.record(FlashInferSageAttention3Backend.get_name(),
-                                query,
-                                key,
-                                value,
-                                causal=self.causal,
-                                sm_scale=self.softmax_scale)
-
-        q = query.transpose(1, 2)
-        k = key.transpose(1, 2)
-        v = value.transpose(1, 2)
-
-        if _can_use_flashinfer(q, k, v, self.dropout):
-            seq_len = q.shape[2]
-            # FlashInfer computes qk_correction with an fp32 matmul and its
-            # forward requires that tensor to remain fp32. Ambient pipeline
-            # autocast would otherwise downcast the matmul result to bf16.
-            with torch.autocast(device_type=q.device.type, enabled=False):
-                packed = nvfp4_attention_sm120_quantize_qkv(q.contiguous(), k.contiguous(), v.contiguous())
-            output, _ = nvfp4_attention_sm120_fwd(*packed,
-                                                  sm_scale=self.softmax_scale,
-                                                  causal=self.causal,
-                                                  out_dtype=query.dtype)
-            # flashinfer pads seq_len to a multiple of 128; trim it back off.
-            return output[:, :, :seq_len, :].transpose(1, 2).contiguous()
-
-        attn_kwargs = {
-            "dropout_p": self.dropout,
-            "is_causal": self.causal,
-            "scale": self.softmax_scale,
-        }
-        if q.shape[1] != k.shape[1]:
-            attn_kwargs["enable_gqa"] = True
-        output = torch.nn.functional.scaled_dot_product_attention(q, k, v, **attn_kwargs)
-        return output.transpose(1, 2)

fastvideo/attention/backends/sage_attn3.py

@@ -3,7 +3,6 @@
 import torch
 from sageattn3 import sageattn3_blackwell
 
-from fastvideo.attention import shape_logger
 from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
                                                    AttentionMetadataBuilder)
 from fastvideo.logger import init_logger
@@ -79,14 +78,6 @@ class SageAttention3Impl(AttentionImpl):
     ) -> torch.Tensor:
         """Call sageattn3_blackwell directly. Input is already [B, H, L, D]
         and contiguous from preprocess_qkv."""
-        if shape_logger.enabled:
-            shape_logger.record("SAGE_ATTN_THREE",
-                                query,
-                                key,
-                                value,
-                                causal=self.causal,
-                                sm_scale=self.softmax_scale,
-                                layout="bhsd")
         output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
         return output
 

fastvideo/attention/shape_logger.py

@@ -1,126 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Env-gated attention shape logger.
-
-Set ``FASTVIDEO_ATTN_SHAPE_LOG=/path/to/shapes.jsonl`` to record every unique
-attention shape executed by the backends that call :func:`record`. Each unique
-(backend, batch, heads, seq lens, head_dim, dtype, causal) combination becomes
-one JSONL line with an occurrence count plus first-seen metadata (torch
-version, device name, SM arch).
-
-Purpose: collect the real attention workloads FastVideo runs so kernel authors
-(e.g. FlashInfer SageAttention3 on DGX Spark / SM121) can tune against actual
-shapes instead of guesses.
-
-Workers are separate processes, so each process writes its own file: the
-configured path is suffixed with ``.rank<N>`` when ``WORLD_SIZE>1``, else
-``.pid<N>``. The file is rewritten on flush (counts mutate, so append-only
-JSONL would go stale); flushes happen every ``_FLUSH_EVERY`` new unique shapes
-and at interpreter exit.
-
-Unset env var = zero overhead: call sites guard on the module-level
-``enabled`` bool and never enter this module's code.
-"""
-
-from __future__ import annotations
-
-import atexit
-import json
-import os
-import threading
-
-import torch
-
-_LOG_PATH: str | None = os.environ.get("FASTVIDEO_ATTN_SHAPE_LOG")
-enabled: bool = _LOG_PATH is not None
-
-_FLUSH_EVERY = 1  # rewrite the file every N new unique shapes (crash-safety)
-
-_lock = threading.Lock()
-_entries: dict[tuple, dict] = {}
-_new_since_flush = 0
-_atexit_registered = False
-
-
-def _log_path() -> str:
-    assert _LOG_PATH is not None
-    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
-        return f"{_LOG_PATH}.rank{os.environ.get('RANK', '0')}"
-    return f"{_LOG_PATH}.pid{os.getpid()}"
-
-
-def _device_meta(device: torch.device) -> dict:
-    meta: dict = {"torch": torch.__version__, "device": str(device), "sm": None}
-    if device.type == "cuda" and torch.cuda.is_available():
-        cap = torch.cuda.get_device_capability(device)
-        meta["device"] = torch.cuda.get_device_name(device)
-        meta["sm"] = f"{cap[0]}.{cap[1]}"
-    return meta
-
-
-def _flush_locked() -> None:
-    """Rewrite the whole snapshot; atomic via os.replace."""
-    path = _log_path()
-    tmp = path + ".tmp"
-    with open(tmp, "w") as f:
-        for entry in _entries.values():
-            f.write(json.dumps(entry) + "\n")
-    os.replace(tmp, path)
-
-
-def _flush_at_exit() -> None:
-    with _lock:
-        if _entries:
-            _flush_locked()
-
-
-def record(backend: str,
-           q: torch.Tensor,
-           k: torch.Tensor,
-           v: torch.Tensor,
-           *,
-           causal: bool,
-           sm_scale: float | None,
-           layout: str = "bshd") -> None:
-    """Record one attention call.
-
-    ``layout`` describes what the calling backend sees: ``"bshd"`` for
-    [batch, seq, heads, head_dim] (the FastVideo impl.forward convention) or
-    ``"bhsd"`` for [batch, heads, seq, head_dim] (backends that transpose in
-    ``preprocess_qkv``, e.g. SAGE_ATTN_THREE).
-    """
-    global _new_since_flush, _atexit_registered
-    if not enabled:
-        return
-    if layout == "bhsd":
-        batch, num_q_heads, seq_len_q, head_dim = q.shape
-        num_kv_heads, seq_len_kv = k.shape[1], k.shape[2]
-    else:
-        batch, seq_len_q, num_q_heads, head_dim = q.shape
-        num_kv_heads, seq_len_kv = k.shape[2], k.shape[1]
-    key = (backend, batch, num_q_heads, num_kv_heads, seq_len_q, seq_len_kv, head_dim, str(q.dtype), bool(causal))
-    with _lock:
-        entry = _entries.get(key)
-        if entry is not None:
-            entry["count"] += 1
-            return
-        _entries[key] = {
-            "backend": backend,
-            "batch": batch,
-            "num_q_heads": num_q_heads,
-            "num_kv_heads": num_kv_heads,
-            "seq_len_q": seq_len_q,
-            "seq_len_kv": seq_len_kv,
-            "head_dim": head_dim,
-            "dtype": str(q.dtype),
-            "causal": bool(causal),
-            "sm_scale": float(sm_scale) if sm_scale is not None else None,
-            "count": 1,
-            **_device_meta(q.device),
-        }
-        if not _atexit_registered:
-            atexit.register(_flush_at_exit)
-            _atexit_registered = True
-        _new_since_flush += 1
-        if _new_since_flush >= _FLUSH_EVERY:
-            _flush_locked()
-            _new_since_flush = 0

fastvideo/attention/utils/flash_attn_default.py

@@ -0,0 +1,251 @@
+# 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() -> Any:
+def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
     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,20 +39,28 @@ def _resolve_flash_attn_varlen_func() -> Any:
                                "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
+        return flash_attn_varlen_func_cute, "4"
     try:
         from flash_attn_interface import (
             flash_attn_varlen_func as flash_attn_varlen_func_interface, )
 
-        return flash_attn_varlen_func_interface
+        return flash_attn_varlen_func_interface, "3"
     except ImportError:
         from flash_attn import (
             flash_attn_varlen_func as flash_attn_varlen_func_flash, )
 
-        return flash_attn_varlen_func_flash
+        return flash_attn_varlen_func_flash, "2"
 
 
-flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
+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, )
 
 
 def flash_attn_no_pad(
@@ -191,3 +199,473 @@ 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,7 +1,9 @@
 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 import FluxDiTConfig
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
@@ -15,6 +17,7 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
-    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
-    "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
+    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
+    "StableAudioConfig", "GlmImageDiTConfig"
 ]

fastvideo/configs/models/dits/base.py

@@ -24,7 +24,6 @@ class DiTArchConfig(ArchConfig):
                                                  AttentionBackendEnum.TORCH_SDPA,
                                                  AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                                                  AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
-                                                 AttentionBackendEnum.FLASHINFER_SAGE_ATTN3,
                                                  AttentionBackendEnum.ATTN_QAT_INFER,
                                                  AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN,
                                                  AttentionBackendEnum.SAGE_SLA_ATTN)

fastvideo/configs/models/dits/cosmos3.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 VFM Transformer FastVideo dataclass configs.
+
+Architecture is 1:1 with the published ``nvidia/Cosmos3-Nano`` checkpoint
+(``transformer/config.json``; class ``Cosmos3OmniTransformer`` / framework
+``Cosmos3VFMNetwork``). Field values match that config so the FastVideo native
+DiT builds a parameter tree matching the checkpoint's state-dict surface
+(814 tensors / 44 patterns, validated 2026-06-06).
+
+Reference of record: ``cosmos-framework`` (NVIDIA). The checkpoint is a single
+``layers`` ModuleList of dual-pathway (understanding/text + generation/vision)
+decoder blocks; per layer: ``self_attn`` with und (``to_{q,k,v}``/``to_out``)
+and gen (``add_{q,k,v}_proj``/``to_add_out``) projections + QK-norms, plus
+``mlp`` (und) and ``mlp_moe_gen`` (gen), and four RMSNorms. Top level adds
+``embed_tokens``/``norm``/``norm_moe_gen``/``lm_head``/``proj_in``/``proj_out``/
+``time_embedder`` and dormant ``action_*``/``audio_*`` heads. The checkpoint
+remap lives in ``scripts/checkpoint_conversion/cosmos3_convert.py``.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_cosmos3_transformer_block(name: str, module) -> bool:
+    """FSDP shard boundary: the dual-pathway decoder blocks ``layers.{i}``."""
+    del module
+    parts = name.split(".")
+    return "layers" in parts and parts[-1].isdigit()
+
+
+@dataclass
+class Cosmos3ArchConfig(DiTArchConfig):
+    """Architecture config for the Cosmos3 omni DiT (Cosmos3-Nano).
+
+    1:1 with ``transformer/config.json``. The action/sound heads ship in the
+    checkpoint, so they are constructed for strict-load parity even though the
+    PR1 video path (T2V/I2V/T2I) leaves them dormant.
+    """
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_cosmos3_transformer_block])
+
+    # Conversion is owned by scripts/checkpoint_conversion/cosmos3_convert.py;
+    # the native module tree is the source of truth for parameter names.
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # ---- Backbone (Qwen3-VL-text) ----
+    hidden_size: int = 4096
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8  # GQA (4 query groups)
+    head_dim: int = 128
+    intermediate_size: int = 12288
+    hidden_act: str = "silu"
+    vocab_size: int = 151936
+    rms_norm_eps: float = 1e-6
+    attention_bias: bool = False
+    qk_norm_for_diffusion: bool = True
+    qk_norm_for_text: bool = True
+    use_moe: bool = True  # dual-pathway weights; sparse routing unused
+    joint_attn_implementation: str = "two_way"
+    freeze_und: bool = False
+
+    # ---- Position embedding (unified 3D MRoPE) ----
+    position_embedding_type: str = "unified_3d_mrope"
+    rope_theta: float = 5_000_000.0
+    max_position_embeddings: int = 262144
+    mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20])
+    mrope_interleaved: bool = True
+    unified_3d_mrope_reset_spatial_ids: bool = True
+    temporal_modality_margin: int = 15000  # unified_3d_mrope_temporal_modality_margin
+
+    # ---- VAE / patch geometry ----
+    latent_patch_size: int = 2
+    latent_channel: int = 48
+    patch_latent_dim: int = 192  # latent_patch_size**2 * latent_channel
+
+    # ---- Diffusion conditioning ----
+    timestep_scale: float = 0.001
+
+    # ---- Temporal / FPS modulation ----
+    base_fps: float = 24.0
+    temporal_compression_factor: int = 4
+    enable_fps_modulation: bool = True
+    video_temporal_causal: bool = False
+
+    # ---- Action generation head (dormant in PR1 video path) ----
+    action_gen: bool = True
+    action_dim: int = 64
+    max_action_dim: int = 64
+    num_embodiment_domains: int = 32
+
+    # ---- Sound generation head (dormant in PR1 video path) ----
+    sound_gen: bool = True
+    sound_dim: int = 64
+    sound_latent_fps: float = 25.0
+    temporal_compression_factor_sound: int = 1
+
+    # ---- BaseDiT bookkeeping ----
+    in_channels: int = 48
+    out_channels: int = 48
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # Video DiT contract: latent channels == VAE z_dim.
+        self.num_channels_latents = self.latent_channel
+        if not self.out_channels:
+            self.out_channels = self.in_channels
+        # Derived: patchify packs latent_patch_size**2 spatial patches * channels.
+        self.patch_latent_dim = self.latent_patch_size**2 * self.latent_channel
+
+
+@dataclass
+class Cosmos3VideoConfig(DiTConfig):
+    """Pipeline-level Cosmos3 DiT config (T2V / I2V / T2I share this surface)."""
+
+    arch_config: DiTArchConfig = field(default_factory=Cosmos3ArchConfig)
+    prefix: str = "Cosmos3"

fastvideo/configs/models/dits/flux.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class FluxTransformer2DArchConfig(DiTArchConfig):
+
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int | None = None
+    num_layers: int = 19
+    num_single_layers: int = 38
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096
+    pooled_projection_dim: int = 768
+    guidance_embeds: bool = True
+    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
+
+
+@dataclass
+class FluxDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=FluxTransformer2DArchConfig)
+    prefix: str = "flux"

fastvideo/configs/models/dits/glm_image.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class GlmImageDiTArchConfig(DiTArchConfig):
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    hidden_size: int = 4096
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    num_layers: int = 30
+
+    text_embed_dim: int = 1472
+    time_embed_dim: int = 512
+    condition_dim: int = 256
+
+    prior_vq_quantizer_codebook_size: int = 16384
+
+    patch_size: int = 2
+
+    max_height: int = 2048
+    max_width: int = 2048
+
+    qk_norm: str = "layer_norm"
+    eps: float = 1e-5
+
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["image_projector", "glyph_projector", "prior_token_embedding"])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^glyph_projector\.net\.0\.proj\.(.*)$": r"glyph_projector.fc_in.\1",
+            r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1",
+            r"^prior_projector\.net\.0\.proj\.(.*)$": r"prior_projector.fc_in.\1",
+            r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class GlmImageDiTConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig)
+    prefix: str = "GlmImage"

fastvideo/configs/models/dits/wanvideo.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
+from typing import Literal
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -19,6 +20,11 @@ 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",
@@ -86,6 +92,14 @@ 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/models/vaes/__init__.py

@@ -1,7 +1,9 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
+from fastvideo.configs.models.vaes.cosmos3vae import Cosmos3VAEConfig
 from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
 from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
+from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
@@ -15,10 +17,12 @@ __all__ = [
     "WanVAEConfig",
     "CosmosVAEConfig",
     "Cosmos25VAEConfig",
+    "Cosmos3VAEConfig",
     "Gen3CVAEConfig",
     "Hunyuan15VAEConfig",
     "LTX2VAEConfig",
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
     "Flux2VAEConfig",
+    "GlmImageVAEConfig",
 ]

fastvideo/configs/models/vaes/cosmos3vae.py

@@ -0,0 +1,277 @@
+"""Cosmos3 (Wan2.2-TI2V-5B) VAE config and checkpoint-key mapping.
+
+The Cosmos3 checkpoint VAE is literally ``Wan-AI/Wan2.2-TI2V-5B-Diffusers``
+(diffusers ``AutoencoderKLWan``), so this config locks the Wan2.2 geometry:
+residual down/up blocks, ``patch_size=2``, ``z_dim=48``, ``base_dim=160``,
+``decoder_base_dim=256``, and ``scale_factor_spatial=16``. The 48-dim
+``latents_mean``/``latents_std`` are taken verbatim from the Cosmos3
+checkpoint's ``vae/config.json`` (identical to the canonical Wan2.2-TI2V-5B
+statistics).
+
+Mirrors the :class:`Cosmos25VAEArchConfig` pattern. ``param_names_mapping`` /
+``map_official_key`` translate the *official* Wan2.2 VAE state-dict keys
+(nested-residual naming, e.g. ``encoder.downsamples.{b}.downsamples.{j}`` and
+``decoder.upsamples.{b}.upsamples.{j}``) into FastVideo's ``AutoencoderKLWan``
+key space. The standard diffusers checkpoint already ships native FastVideo
+keys, so these helpers exist for parity tooling and official ``.pth`` loading.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig, WanVAEConfig
+
+
+@dataclass
+class Cosmos3VAEArchConfig(WanVAEArchConfig):
+    # Wan2.2-TI2V-5B geometry (differs from the Wan2.1 WanVAEArchConfig
+    # defaults: residual blocks, patch_size=2, z_dim=48, base_dim=160,
+    # decoder_base_dim=256, scale_factor_spatial=16, 12 patch channels).
+    _name_or_path: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+    base_dim: int = 160
+    decoder_base_dim: int | None = 256
+    z_dim: int = 48
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
+    num_res_blocks: int = 2
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
+    dropout: float = 0.0
+    is_residual: bool = True
+    in_channels: int = 12
+    out_channels: int = 12
+    patch_size: int | None = 2
+    scale_factor_temporal: int = 4
+    scale_factor_spatial: int = 16
+    clip_output: bool = False
+
+    # 48-dim statistics copied verbatim from the Cosmos3 checkpoint
+    # (official_weights/cosmos3/vae/config.json).
+    latents_mean: tuple[float, ...] = (
+        -0.2289,
+        -0.0052,
+        -0.1323,
+        -0.2339,
+        -0.2799,
+        0.0174,
+        0.1838,
+        0.1557,
+        -0.1382,
+        0.0542,
+        0.2813,
+        0.0891,
+        0.157,
+        -0.0098,
+        0.0375,
+        -0.1825,
+        -0.2246,
+        -0.1207,
+        -0.0698,
+        0.5109,
+        0.2665,
+        -0.2108,
+        -0.2158,
+        0.2502,
+        -0.2055,
+        -0.0322,
+        0.1109,
+        0.1567,
+        -0.0729,
+        0.0899,
+        -0.2799,
+        -0.123,
+        -0.0313,
+        -0.1649,
+        0.0117,
+        0.0723,
+        -0.2839,
+        -0.2083,
+        -0.052,
+        0.3748,
+        0.0152,
+        0.1957,
+        0.1433,
+        -0.2944,
+        0.3573,
+        -0.0548,
+        -0.1681,
+        -0.0667,
+    )
+    latents_std: tuple[float, ...] = (
+        0.4765,
+        1.0364,
+        0.4514,
+        1.1677,
+        0.5313,
+        0.499,
+        0.4818,
+        0.5013,
+        0.8158,
+        1.0344,
+        0.5894,
+        1.0901,
+        0.6885,
+        0.6165,
+        0.8454,
+        0.4978,
+        0.5759,
+        0.3523,
+        0.7135,
+        0.6804,
+        0.5833,
+        1.4146,
+        0.8986,
+        0.5659,
+        0.7069,
+        0.5338,
+        0.4889,
+        0.4917,
+        0.4069,
+        0.4999,
+        0.6866,
+        0.4093,
+        0.5709,
+        0.6065,
+        0.6415,
+        0.4944,
+        0.5726,
+        1.2042,
+        0.5458,
+        1.6887,
+        0.3971,
+        1.06,
+        0.3943,
+        0.5537,
+        0.5444,
+        0.4089,
+        0.7468,
+        0.7744,
+    )
+
+    # Simple 1:1 renames. The nested-residual block remapping (encoder
+    # downsamples / decoder upsamples / middle / head) is handled by
+    # ``map_official_key()``.
+    param_names_mapping: dict[str, str] = field(
+        default_factory=lambda: {
+            r"^conv1\.(.*)$": r"quant_conv.\1",
+            r"^conv2\.(.*)$": r"post_quant_conv.\1",
+            r"^encoder\.conv1\.(.*)$": r"encoder.conv_in.\1",
+            r"^decoder\.conv1\.(.*)$": r"decoder.conv_in.\1",
+            r"^encoder\.head\.0\.gamma$": r"encoder.norm_out.gamma",
+            r"^encoder\.head\.2\.(.*)$": r"encoder.conv_out.\1",
+            r"^decoder\.head\.0\.gamma$": r"decoder.norm_out.gamma",
+            r"^decoder\.head\.2\.(.*)$": r"decoder.conv_out.\1",
+        })
+
+    @staticmethod
+    def map_official_key(key: str) -> str | None:
+        """Map a single official Wan2.2 VAE key into FastVideo key space.
+
+        Handles the residual (Wan2.2) module layout where each down/up block
+        is a nested ``Sequential`` (``downsamples.{b}.downsamples.{j}`` /
+        ``upsamples.{b}.upsamples.{j}``) rather than the flat Wan2.1 indexing.
+        Returns ``None`` for keys with no FastVideo counterpart.
+        """
+
+        def map_residual_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^residual\.0\.gamma$", sub):
+                return f"{prefix}.norm1.gamma"
+            m = re.match(r"^residual\.2\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv1.{m.group(1)}"
+            if re.match(r"^residual\.3\.gamma$", sub):
+                return f"{prefix}.norm2.gamma"
+            m = re.match(r"^residual\.6\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv2.{m.group(1)}"
+            m = re.match(r"^shortcut\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv_shortcut.{m.group(1)}"
+            return None
+
+        def map_attn_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^norm\.gamma$", sub):
+                return f"{prefix}.norm.gamma"
+            m = re.match(r"^to_qkv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.to_qkv.{m.group(1)}"
+            m = re.match(r"^proj\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.proj.{m.group(1)}"
+            return None
+
+        def map_resample_subkey(prefix: str, sub: str) -> str | None:
+            m = re.match(r"^resample\.1\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.resample.1.{m.group(1)}"
+            m = re.match(r"^time_conv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.time_conv.{m.group(1)}"
+            return None
+
+        m = re.match(r"^conv1\.(weight|bias)$", key)
+        if m:
+            return f"quant_conv.{m.group(1)}"
+        m = re.match(r"^conv2\.(weight|bias)$", key)
+        if m:
+            return f"post_quant_conv.{m.group(1)}"
+        m = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_in.{m.group(2)}"
+        m = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key)
+        if m:
+            return f"{m.group(1)}.norm_out.gamma"
+        m = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_out.{m.group(2)}"
+
+        m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0", m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
+        if m:
+            return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0", m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1", m.group(2))
+
+        # Encoder: downsamples.{block}.downsamples.{j}.* (nested residual layout)
+        m = re.match(r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.(.*)$", key)
+        if m:
+            block_i, res_i, sub = int(m.group(1)), int(m.group(2)), m.group(3)
+            if sub.startswith("resample.") or sub.startswith("time_conv."):
+                return map_resample_subkey(f"encoder.down_blocks.{block_i}.downsampler", sub)
+            return map_residual_subkey(f"encoder.down_blocks.{block_i}.resnets.{res_i}", sub)
+
+        # Decoder: upsamples.{block}.upsamples.{j}.* (nested residual layout)
+        m = re.match(r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.(.*)$", key)
+        if m:
+            block_i, res_i, sub = int(m.group(1)), int(m.group(2)), m.group(3)
+            if sub.startswith("resample.") or sub.startswith("time_conv."):
+                return map_resample_subkey(f"decoder.up_blocks.{block_i}.upsampler", sub)
+            return map_residual_subkey(f"decoder.up_blocks.{block_i}.resnets.{res_i}", sub)
+
+        return None
+
+    # ``__post_init__`` (scaling_factor / shift_factor / compression ratios) is
+    # inherited unchanged from ``WanVAEArchConfig``.
+
+
+@dataclass
+class Cosmos3VAEConfig(WanVAEConfig):
+    """Cosmos3 VAE config (reuses FastVideo's Wan2.2 ``AutoencoderKLWan``).
+
+    Subclasses :class:`WanVAEConfig` so the model reads the same runtime flags
+    (``use_feature_cache``, ``use_light_vae``, tiling) and only swaps in the
+    Cosmos3 = Wan2.2 ``arch_config``.
+    """
+
+    arch_config: Cosmos3VAEArchConfig = field(default_factory=Cosmos3VAEArchConfig)
+
+    use_feature_cache: bool = True
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    # ``__post_init__`` (blend_num_frames) is inherited from ``WanVAEConfig``.

fastvideo/configs/models/vaes/glm_image.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.autoencoder_kl import (AutoencoderKLArchConfig, AutoencoderKLVAEConfig)
+
+_GLM_IMAGE_LATENTS_MEAN: tuple[float, ...] = (
+    -0.2080078125,
+    1.875,
+    -0.470703125,
+    -1.265625,
+    -1.421875,
+    0.77734375,
+    -0.3671875,
+    -0.9453125,
+    0.318359375,
+    0.7734375,
+    -0.1884765625,
+    -0.022216796875,
+    -0.220703125,
+    -1.59375,
+    -0.81640625,
+    -0.255859375,
+)
+_GLM_IMAGE_LATENTS_STD: tuple[float, ...] = (
+    3.0625,
+    2.203125,
+    2.265625,
+    4.84375,
+    2.5,
+    3.9375,
+    2.203125,
+    3.03125,
+    2.1875,
+    2.046875,
+    2.71875,
+    2.390625,
+    2.390625,
+    2.453125,
+    2.25,
+    2.15625,
+)
+
+
+@dataclass
+class GlmImageVAEArchConfig(AutoencoderKLArchConfig):
+    act_fn: str = "silu"
+    block_out_channels: tuple[int, ...] = (128, 512, 1024, 1024)
+    down_block_types: tuple[str, ...] = (
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+    )
+    up_block_types: tuple[str, ...] = (
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+    )
+    force_upcast: bool = True
+    in_channels: int = 3
+    out_channels: int = 3
+    latent_channels: int = 16
+    latents_mean: tuple[float, ...] = _GLM_IMAGE_LATENTS_MEAN
+    latents_std: tuple[float, ...] = _GLM_IMAGE_LATENTS_STD
+    layers_per_block: int = 3
+    mid_block_add_attention: bool = False
+    norm_num_groups: int = 32
+    sample_size: int = 1024
+    scaling_factor: float = 0.18215
+    shift_factor: float | None = None
+    use_quant_conv: bool = False
+    use_post_quant_conv: bool = False
+
+    temporal_compression_ratio: int = 1
+    spatial_compression_ratio: int = 8
+
+
+@dataclass
+class GlmImageVAEConfig(AutoencoderKLVAEConfig):
+    arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig)
+
+    use_tiling: bool = True
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    tile_sample_min_height: int = 512
+    tile_sample_min_width: int = 512
+    tile_sample_stride_height: int = 384
+    tile_sample_stride_width: int = 384
+
+    load_encoder: bool = True
+    load_decoder: bool = True

fastvideo/configs/pipelines/cosmos3.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline configuration.
+
+Reference of record: the official ``cosmos-framework`` / ``nvidia/Cosmos3-Nano``
+checkpoint (``model_index.json``). Cosmos3 is structurally different from
+Cosmos 2.5:
+
+- Dual-pathway (UND + GEN) DiT lives entirely inside ``Cosmos3VFMTransformer``
+  (``Cosmos3VideoConfig``).
+- No separate text encoder — the Qwen3-VL-text backbone is inside the DiT, so
+  ``text_encoder_configs`` is the empty tuple. The Qwen2 tokenizer is loaded as
+  the ``text_tokenizer`` checkpoint module by the component loader.
+- VAE is Wan2.2 ``AutoencoderKLWan`` (z_dim=48, scale_factor_spatial=16),
+  configured by ``Cosmos3VAEConfig`` (the checkpoint's exact latents_mean/std).
+- Scheduler is FastVideo-native ``UniPCMultistepScheduler`` configured for
+  pure flow matching (flow_prediction, use_flow_sigmas), equivalent to the
+  framework's ``FlowUniPCMultistepScheduler``. The checkpoint's diffusers-style
+  scheduler config (karras/sigma_min/max) is coerced to the flow setup in
+  ``Cosmos3OmniDiffusersPipeline.initialize_pipeline``.
+- T2I default ``flow_shift`` is 3.0 (set per-request by ``_set_flow_shift``);
+  T2V/I2V use the engine-init default of 1.0 baked into this config.
+"""
+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.cosmos3 import (Cosmos3ArchConfig, Cosmos3VideoConfig)
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.vaes import Cosmos3VAEConfig  # Wan2.2 AutoencoderKLWan
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class Cosmos3Config(PipelineConfig):
+    """Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I).
+
+    Wires the framework-parity-verified Cosmos3 components: the native
+    ``Cosmos3VideoConfig`` DiT, the Wan2.2 ``Cosmos3VAEConfig`` VAE, the Qwen2
+    tokenizer (loaded as ``text_tokenizer``), and the UniPC scheduler.
+    """
+
+    dit_config: DiTConfig = field(default_factory=lambda: Cosmos3VideoConfig(arch_config=Cosmos3ArchConfig()))
+
+    vae_config: VAEConfig = field(default_factory=Cosmos3VAEConfig)
+
+    # No separate text encoder: the Qwen3-VL-text backbone lives inside the DiT
+    # and the pipeline tokenizes in Cosmos3DenoisingStage, so all three
+    # text-encoder lists are empty (the generic text-encode stage is not used).
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple)
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=tuple)
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=tuple)
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple)
+
+    embedded_cfg_scale: float = 0.0
+    # T2V/I2V engine-init flow_shift (framework text2video/image2video default);
+    # T2I overrides to 3.0 per request via Cosmos3DenoisingStage._set_flow_shift.
+    flow_shift: float = 10.0
+
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/flux.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import EncoderConfig
+from fastvideo.configs.models.dits.flux import FluxDiTConfig
+from fastvideo.configs.models.encoders import (
+    BaseEncoderOutput,
+    CLIPTextConfig,
+    T5LargeConfig,
+)
+from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+
+def _flux_clip_pooled_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    """CLIP branch for FLUX: Diffusers uses pooled prompt embeddings only."""
+    if outputs.pooler_output is None:
+        raise RuntimeError(
+            "FLUX CLIP conditioning requires pooler_output. Ensure the CLIP text encoder returns pooled features.")
+    return outputs.pooler_output
+
+
+def _flux_t5_sequence_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.last_hidden_state is None:
+        raise RuntimeError("FLUX T5 conditioning requires last_hidden_state.")
+    return outputs.last_hidden_state
+
+
+@dataclass
+class FluxPipelineConfig(PipelineConfig):
+    """Pipeline layout for Diffusers FLUX.1-dev (CLIP + T5 + packed DiT + FlowMatch)."""
+
+    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
+    transformer_arch: str = "FluxTransformer2DModel"
+    vae_arch: str = "AutoencoderKL"
+    text_encoder_archs: tuple[str, ...] = ("CLIPTextModel", "T5EncoderModel")
+    tokenizer_archs: tuple[str, ...] = ("CLIPTokenizer", "T5TokenizerFast")
+
+    dit_config: FluxDiTConfig = field(default_factory=FluxDiTConfig)
+    vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
+
+    embedded_cfg_scale: float = 3.5
+    flow_shift: float | None = None
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (CLIPTextConfig(), T5LargeConfig()))
+    preprocess_text_funcs: tuple[Callable[[str], str],
+                                 ...] = field(default_factory=lambda: (preprocess_text, preprocess_text))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
+        default_factory=lambda: (_flux_clip_pooled_postprocess, _flux_t5_sequence_postprocess))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", "bf16"))
+
+    def __post_init__(self) -> None:
+        te_cfgs = list(self.text_encoder_configs)
+        if len(te_cfgs) >= 1:
+            te_cfgs[0].tokenizer_kwargs.setdefault("padding", "max_length")
+            te_cfgs[0].tokenizer_kwargs.setdefault("max_length", 77)
+            te_cfgs[0].tokenizer_kwargs.setdefault("truncation", True)
+            te_cfgs[0].tokenizer_kwargs.setdefault("return_tensors", "pt")
+        if len(te_cfgs) >= 2:
+            cap = 512
+            te_cfgs[1].tokenizer_kwargs["max_length"] = min(int(te_cfgs[1].tokenizer_kwargs.get("max_length", cap)),
+                                                            cap)
+            te_cfgs[1].tokenizer_kwargs.setdefault("padding", "max_length")
+            te_cfgs[1].tokenizer_kwargs.setdefault("truncation", True)
+            te_cfgs[1].tokenizer_kwargs.setdefault("return_tensors", "pt")

fastvideo/configs/pipelines/glm_image.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+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.glm_image import GlmImageDiTConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
+from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+def glm_image_t5_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+
+    assert torch.isnan(hidden_state).sum() == 0, "T5 hidden states contain NaN"
+
+    max_len = 512
+    prompt_embeds = [u[:min(v, max_len)] for u, v in zip(hidden_state, seq_lens, strict=True)]
+    prompt_embeds_tensor: torch.Tensor = torch.stack(
+        [torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
+
+    return prompt_embeds_tensor
+
+
+@dataclass
+class GlmImageConfig(PipelineConfig):
+
+    dit_config: DiTConfig = field(default_factory=GlmImageDiTConfig)
+    dit_precision: str = "bf16"
+
+    vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig)
+    vae_precision: str = "fp32"
+    vae_tiling: bool = True
+    vae_sp: bool = False
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (glm_image_t5_postprocess, ))
+
+    flow_shift: float | None = 1.0
+    embedded_cfg_scale: float = 7.5

fastvideo/dataset/parquet_dataset_map_style.py

@@ -2,6 +2,7 @@
 import os
 import pickle
 import random
+from collections.abc import Sequence
 from typing import Any
 
 import pyarrow as pa
@@ -94,8 +95,56 @@ class DP_SP_BatchSampler(Sampler[list[int]]):
         return len(self.sp_group_local_indices) // self.batch_size
 
 
-def get_parquet_files_and_length(path: str):
-    dataset_root = os.path.realpath(os.path.expanduser(path))
+def _parse_data_path_specs(path: str | Sequence[str] | dict[str, int]) -> list[tuple[str, int]]:
+    """Parse one or more dataset roots with old-framework repeat counts."""
+    if isinstance(path, dict):
+        return [
+            (str(root), int(repeat))
+            for root, repeat in path.items()
+            if int(repeat) > 0
+        ]
+    if isinstance(path, Sequence) and not isinstance(path, str):
+        return [(str(root), 1) for root in path]
+
+    specs: list[tuple[str, int]] = []
+    for part in str(path).split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if ":" in part:
+            dir_path, count_str = part.rsplit(":", 1)
+            count = int(count_str)
+        else:
+            dir_path, count = part, 1
+        if count > 0:
+            specs.append((dir_path.strip(), count))
+    return specs
+
+
+def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int]):
+    specs = _parse_data_path_specs(path)
+    if len(specs) != 1 or specs[0][1] != 1:
+        all_file_names: list[str] = []
+        all_lengths: list[int] = []
+        for root, repeat in specs:
+            file_names, lengths = get_parquet_files_and_length(str(root))
+            for _ in range(repeat):
+                all_file_names.extend(file_names)
+                all_lengths.extend(lengths)
+        if not all_file_names:
+            raise FileNotFoundError(
+                "No parquet files found under dataset paths: "
+                f"{path}. "
+                "Please verify these paths point to preprocessed parquet data."
+            )
+        file_lengths = sorted(
+            zip(all_file_names, all_lengths, strict=True),
+            key=lambda x: x[0],
+        )
+        file_names_sorted, lengths_sorted = zip(*file_lengths, strict=True)
+        return file_names_sorted, lengths_sorted
+
+    dataset_root = os.path.realpath(os.path.expanduser(specs[0][0]))
     # Check if cached info exists
     cache_dir = os.path.join(dataset_root, "map_style_cache")
     cache_file = os.path.join(cache_dir, "file_info.pkl")
@@ -268,7 +317,7 @@ class LatentsParquetMapStyleDataset(Dataset):
 
     def __init__(
         self,
-        path: str,
+        path: str | Sequence[str] | dict[str, int],
         batch_size: int,
         parquet_schema: pa.Schema,
         cfg_rate: float = 0.0,
@@ -358,7 +407,6 @@ class LatentsParquetMapStyleDataset(Dataset):
     def __len__(self):
         return sum(self.lengths)
 
-
 # ────────────────────────────────────────────────────────────────────────────
 # 3.  Loader helper – everything else stays just like your original trainer
 # ────────────────────────────────────────────────────────────────────────────

fastvideo/envs.py

@@ -208,22 +208,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention
     # - "SAGE_ATTN": use Sage Attention
     # - "SAGE_ATTN_THREE": use Sage Attention 3
-    # - "FLASHINFER_SAGE_ATTN3": use FlashInfer SageAttention3 (SM120 NVFP4);
-    #   falls back to torch SDPA where the kernel cannot run
     # FLASH_ATTN uses FlashAttention-3/2; to run FlashAttention-4 set
     # FASTVIDEO_FA4=1 as well (see below).
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 
-    # If set to a file path, attention backends that support it (currently
-    # FLASHINFER_SAGE_ATTN3 and SAGE_ATTN_THREE) record every unique attention
-    # shape they run to <path>.rank<N>/<path>.pid<N> as JSONL, with occurrence
-    # counts and device metadata. Used to hand real workload shapes to kernel
-    # authors (e.g. SageAttention3 tuning for DGX Spark / SM121).
-    # See fastvideo/attention/shape_logger.py. Read once at import time.
-    "FASTVIDEO_ATTN_SHAPE_LOG":
-    lambda: os.getenv("FASTVIDEO_ATTN_SHAPE_LOG", None),
-
     # If set (=1), the FLASH_ATTN backend uses FlashAttention-4
     # (flash_attn.cute). FA4 is opt-in and never auto-selected just because it
     # is installed. Below sm90, grad-enabled and GQA calls are routed to FA2

fastvideo/eval/README.md

@@ -122,6 +122,44 @@ to run a subset of the Evaluator's registered metrics on this batch
 (useful for scoring different corpora with different metric subsets in
 successive calls without burning model loads).
 
+### Training-time validation metrics
+
+The modular trainer can score validation videos during training through
+`callbacks.validation.metrics`. Each distributed rank that writes local
+validation videos also runs its local metrics on `cuda:<local_rank>`;
+rank 0 only merges scalar summaries and logs artifacts.
+
+```yaml
+callbacks:
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.matrixgame2.matrixgame2_causal_dmd_pipeline.MatrixGame2CausalDMDPipeline
+    dataset_file: examples/distill/MatrixGame2.0/validation.json
+    sampling_steps: [4]
+    metrics:
+      enabled: true
+      names:
+        - vbench.imaging_quality
+        - vbench.aesthetic_quality
+        - optical_flow.synthetic_optical_flow
+        - common.fvd
+      skip_missing_deps: true
+      strict: false
+      unload_after_validation: true
+
+      # `synthetic_optical_flow` also needs actions in the validation
+      # manifest (`action_path`) plus a repo-local calibration file.
+      calibration_path: assets/eval/worldmodel_synthetic_flow_calibration.json
+```
+
+Metric summaries are written under
+`<output_dir>/eval/step_<step>/inference_steps_<n>_rank_<rank>.json` and
+scalar means are logged with the `metrics/validation/...` prefix.
+Install `fastvideo[eval]` for optical-flow dependencies such as
+`ptlflow`; VBench metrics use the pinned submodule under
+`fastvideo/third_party/eval/vbench`. FVD uses `ref_video` entries from
+the validation manifest when present, or the standard
+`FASTVIDEO_FVD_REF_FEATURES` / eval-cache reference feature path.
+
 ### CLI
 
 ```bash
@@ -350,5 +388,3 @@ sweep several baselines into a table, see
 
 - **MIND** metrics. Depend on a separate `vipe` upstream submodule.
 - **VBench-2.0**. Sibling vbench2 package; needs its own port.
-- **Training-time eval callback** (`EvalCallback`) and the
-  `RolloutEvaluator` helper.

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
+    r_z: float = 0.0
     r_y: float = 0.0
     init_pitch: float = 0.0
     notes: str = ""

fastvideo/layers/quantization/nvfp4_config.py

@@ -261,6 +261,22 @@ def _mm_fp4(
     )
 
 
+def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor:
+    """Coerce an activation to a dtype the FP4 linear accepts.
+
+    The pre-attention norm can emit fp32 (e.g. in eager mode, without the
+    torch.compile fusion that keeps it bf16). The FP4 linear emits bf16
+    regardless (see _mm_fp4 out dtype), so cast fp32 -> bf16 rather than
+    failing, matching the sibling fastvideo/layers/fp4linear.py. Non-floating
+    inputs (e.g. int/bool) are a genuine error and are rejected fast.
+    """
+    if not x.is_floating_point():
+        raise TypeError(f"fp4 linear expects floating-point inputs, got {x.dtype}")
+    if x.dtype not in (torch.bfloat16, torch.float16):
+        x = x.to(torch.bfloat16)
+    return x
+
+
 class NVFP4QuantizeMethod(QuantizeMethodBase):
 
     def __init__(self, layer_prefix: str = ""):
@@ -285,8 +301,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
 
     def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         SfLayout, _, _ = _require_flashinfer()
-        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
-            f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
+        x = _coerce_fp4_input_dtype(x)
         x_2d = x.view(-1, x.shape[-1])
         x_fp4, x_scale = _nvfp4_quantize(
             x_2d,
@@ -332,8 +347,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
             if x_scale.dim() > 2:
                 x_scale = x_scale.view(-1, x_scale.shape[-1])
         else:
-            assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
-                f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
+            x = _coerce_fp4_input_dtype(x)
             x = x.view(-1, x.shape[-1])
             x_global_sf = self.x_global_sf
             x_fp4, x_scale = _nvfp4_quantize(

fastvideo/layers/rotary_embedding.py

@@ -295,6 +295,7 @@ def get_1d_rotary_pos_embed(
     interpolation_factor: float = 1.0,
     dtype: torch.dtype = torch.float32,
     use_real: bool = True,
+    freqs_dtype: torch.dtype | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Precompute the frequency tensor for complex exponential (cis) with given dimensions.
@@ -319,12 +320,16 @@ def get_1d_rotary_pos_embed(
     if isinstance(pos, int):
         pos = torch.arange(pos).float()
 
+    # freqs_dtype is an alias for dtype (Diffusers-compatible calling convention).
+    if freqs_dtype is not None:
+        dtype = freqs_dtype
+
     # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
     # has some connection to NTK literature
     if theta_rescale_factor != 1.0:
         theta *= theta_rescale_factor**(dim / (dim - 2))
 
-    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(dtype) / dim))  # [D/2]
+    freqs = 1.0 / (theta**(torch.arange(0, dim, 2, device=pos.device)[:(dim // 2)].to(dtype) / dim))  # [D/2]
     freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
     freqs_cos = freqs.cos()  # [S, D/2]
     freqs_sin = freqs.sin()  # [S, D/2]
@@ -445,6 +450,21 @@ def get_nd_rotary_pos_embed(
     return cos, sin
 
 
+_ROTARY_POS_EMBED_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+# Bound the table cache so long-running servers / causal models (which vary
+# start_frame per frame) cannot grow it without limit; entries are large float64
+# tensors. Least-recently-used eviction keeps the active resolution(s) hot while
+# capping memory.
+_ROTARY_POS_EMBED_CACHE_MAXSIZE = 16
+
+
+def _hashable(value: Any) -> Any:
+    """Return a hashable view of a scalar or sequence for use in a cache key."""
+    if isinstance(value, list | tuple):
+        return tuple(value)
+    return value
+
+
 def get_rotary_pos_embed(
     rope_sizes,
     hidden_size,
@@ -495,6 +515,31 @@ def get_rotary_pos_embed(
         sp_rank = 0
         sp_world_size = 1
 
+    # Memoize on every output-affecting argument; the table is constant across
+    # denoising steps, so this avoids recomputing the float64 cos/sin tables.
+    cache_key = (
+        _hashable(rope_sizes),
+        tuple(rope_dim_list),
+        rope_theta,
+        _hashable(theta_rescale_factor),
+        _hashable(interpolation_factor),
+        shard_dim,
+        sp_rank,
+        sp_world_size,
+        dtype,
+        start_frame,
+        use_real,
+    )
+    cached = _ROTARY_POS_EMBED_CACHE.get(cache_key)
+    if cached is not None:
+        # Move to most-recently-used position so the active table is not evicted
+        # when several resolutions / buckets share the process (LRU recency).
+        # Pop with a default: a concurrent eviction between the get() above and
+        # here would otherwise raise KeyError on the hit path.
+        if _ROTARY_POS_EMBED_CACHE.pop(cache_key, None) is not None:
+            _ROTARY_POS_EMBED_CACHE[cache_key] = cached
+        return cached
+
     freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
         rope_dim_list,
         rope_sizes,
@@ -508,6 +553,14 @@ def get_rotary_pos_embed(
         start_frame=start_frame,
         use_real=use_real,
     )
+    # The returned tensors are shared cache entries: callers must never mutate
+    # them in place. Note .to(device) is an identity alias when the tensor is
+    # already on the target device (e.g. CPU runs), so it does NOT guarantee a
+    # copy — treat the tables as read-only and copy before any in-place op.
+    # Reached only on a miss, so evict the least-recently-used entry at capacity.
+    if len(_ROTARY_POS_EMBED_CACHE) >= _ROTARY_POS_EMBED_CACHE_MAXSIZE:
+        _ROTARY_POS_EMBED_CACHE.pop(next(iter(_ROTARY_POS_EMBED_CACHE)))
+    _ROTARY_POS_EMBED_CACHE[cache_key] = (freqs_cos, freqs_sin)
     return freqs_cos, freqs_sin
 
 

fastvideo/models/audio/cosmos3_avae.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 sound tokenizer (AVAE) — decode path.
+
+The Cosmos3 ``sound_tokenizer`` is an AVAE (audio VAE). Its shipped diffusers
+checkpoint is **decoder-only** (``decoder.*``) in ``AutoencoderOobleck`` naming
+with SnakeBeta activations and ``weight_g``/``weight_v`` weight-norm — exactly
+FastVideo's native :class:`~fastvideo.models.vaes.oobleck.OobleckDecoder`
+(verified bit-exact vs the framework in ``test_cosmos3_avae_parity``). Text-to-
+video+sound (t2vs) only needs DECODE: the DiT generates the sound latent and
+this module decodes it to a waveform, so only the decoder is ported (the
+SpectrogramConvNeXt encoder is not exported in the checkpoint).
+
+Mirrors the framework ``AVAEModel.decode``: run the Oobleck decoder, then clamp
+to [-1, 1]. The VAE bottleneck's decode is the identity (the DiT already emits
+the post-bottleneck latent), so there is no bottleneck step here.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from fastvideo.logger import init_logger
+from fastvideo.models.vaes.oobleck import OobleckDecoder
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class Cosmos3SoundVAEArchConfig:
+    """Cosmos3 AVAE decoder constants (from ``sound_tokenizer/config.json``)."""
+
+    dec_dim: int = 320  # decoder base channels
+    vocoder_input_dim: int = 64  # latent channels in
+    dec_c_mults: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16])
+    dec_strides: list[int] = field(default_factory=lambda: [2, 4, 5, 6, 8])
+    audio_channels: int = 2  # stereo
+    sampling_rate: int = 48000
+
+    @property
+    def hop_size(self) -> int:
+        return int(np.prod(self.dec_strides))  # 1920
+
+
+class Cosmos3SoundVAE(nn.Module):
+    """Decoder-only Cosmos3 AVAE: latent ``[B, z, T]`` -> waveform ``[B, C, N]``."""
+
+    def __init__(self, arch: Cosmos3SoundVAEArchConfig | None = None) -> None:
+        super().__init__()
+        self.arch = arch or Cosmos3SoundVAEArchConfig()
+        self.decoder = OobleckDecoder(
+            channels=self.arch.dec_dim,
+            input_channels=self.arch.vocoder_input_dim,
+            audio_channels=self.arch.audio_channels,
+            # The framework builds decoder blocks from ``reversed(dec_strides)``
+            # (deepest first), so block strides are e.g. [8,6,5,4,2].
+            upsampling_ratios=list(reversed(self.arch.dec_strides)),
+            channel_multiples=list(self.arch.dec_c_mults),
+        )
+
+    @property
+    def sample_rate(self) -> int:
+        return self.arch.sampling_rate
+
+    @property
+    def audio_channels(self) -> int:
+        return self.arch.audio_channels
+
+    @property
+    def hop_size(self) -> int:
+        return self.arch.hop_size
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        """Latent length for a given audio length (``AVAEInterface``: ``N // hop``)."""
+        return int(num_audio_samples) // self.arch.hop_size
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """Decode normalized latent ``[B, z, T]`` to waveform ``[B, C, N]`` in [-1, 1].
+
+        Matches ``AVAEModel.decode``: Oobleck decoder then clamp to [-1, 1] (the
+        VAE bottleneck decode is identity).
+        """
+        audio = self.decoder(latent)  # [B, C, N]
+        return audio.clamp(-1.0, 1.0)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        *,
+        torch_dtype: torch.dtype | None = None,
+    ) -> "Cosmos3SoundVAE":
+        """Build + load the decoder from a ``sound_tokenizer`` directory.
+
+        Reads ``config.json`` (``dec_dim`` / ``vocoder_input_dim`` /
+        ``dec_c_mults`` / ``dec_strides`` / ``sampling_rate`` / ``stereo``) and
+        loads the ``decoder.*`` weights (the checkpoint is decoder-only).
+        """
+        from safetensors.torch import load_file
+
+        cfg_path = os.path.join(model_path, "config.json")
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+        arch = Cosmos3SoundVAEArchConfig(
+            dec_dim=int(cfg["dec_dim"]),
+            vocoder_input_dim=int(cfg["vocoder_input_dim"]),
+            dec_c_mults=list(cfg["dec_c_mults"]),
+            dec_strides=list(cfg["dec_strides"]),
+            audio_channels=2 if cfg.get("stereo", True) else 1,
+            sampling_rate=int(cfg.get("sampling_rate", 48000)),
+        )
+        model = cls(arch)
+
+        weights_path = os.path.join(model_path, "diffusion_pytorch_model.safetensors")
+        state = load_file(weights_path)
+        # Decoder-only checkpoint: strip the ``decoder.`` prefix.
+        dec_state = {k[len("decoder."):]: v for k, v in state.items() if k.startswith("decoder.")}
+        model.decoder.load_state_dict(dec_state, strict=True)
+        logger.info("Loaded Cosmos3 sound AVAE decoder (%d params) from %s",
+                    sum(p.numel() for p in model.parameters()), model_path)
+
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        model.eval()
+        return model
+
+
+EntryClass = Cosmos3SoundVAE

fastvideo/models/dits/causal_wanvideo.py

@@ -665,7 +665,8 @@ class CausalWanTransformer3DModel(BaseDiT):
             else:
                 causal_kwargs = {
                     "kv_cache": kv_cache[block_index],
-                    "crossattn_cache": crossattn_cache[block_index],
+                    "crossattn_cache": (crossattn_cache[block_index]
+                                        if crossattn_cache is not None else None),
                     "current_start": current_start,
                     "cache_start": cache_start,
                     "block_mask": self.block_mask,

fastvideo/models/dits/flux.py

@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.layers.rotary_embedding import apply_rotary_emb, get_1d_rotary_pos_embed
+
+from fastvideo.attention import DistributedAttention
+from fastvideo.configs.models import DiTConfig
+from fastvideo.forward_context import get_forward_context, set_forward_context
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.sd3 import (
+    CombinedTimestepTextProjEmbeddings,
+    SD3AdaLayerNormContinuous,
+    SD3AdaLayerNormZero,
+    SD3FeedForward,
+    SD3TextProjection,
+    SD3TimestepEmbedding,
+)
+from fastvideo.platforms import AttentionBackendEnum
+
+
+@dataclass
+class FluxTransformer2DModelOutput:
+    sample: torch.Tensor
+
+
+class FluxPosEmbed(nn.Module):
+    """1D RoPE axes concatenated per Diffusers `FluxPosEmbed`."""
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        n_axes = ids.shape[-1]
+        cos_out: list[torch.Tensor] = []
+        sin_out: list[torch.Tensor] = []
+        pos = ids.float()
+        is_mps = ids.device.type == "mps"
+        is_npu = ids.device.type == "npu"
+        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                use_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None:
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256,
+            flip_sin_to_cos=True,
+            downscale_freq_shift=0,
+        )
+        self.timestep_embedder = SD3TimestepEmbedding(
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            act_fn="silu",
+        )
+        self.guidance_embedder = SD3TimestepEmbedding(
+            in_channels=256,
+            time_embed_dim=embedding_dim,
+            act_fn="silu",
+        )
+        self.text_embedder = SD3TextProjection(
+            pooled_projection_dim,
+            embedding_dim,
+            act_fn="silu",
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor,
+        pooled_projection: torch.Tensor,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
+        guidance_proj = self.time_proj(guidance)
+        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
+        time_guidance_emb = timesteps_emb + guidance_emb
+        pooled_projections = self.text_embedder(pooled_projection)
+        return time_guidance_emb + pooled_projections
+
+
+class FluxAdaLayerNormZeroSingle(nn.Module):
+    def __init__(self, embedding_dim: int, bias: bool = True) -> None:
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(embedding_dim, 3 * embedding_dim, bias=bias)
+        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        emb, _ = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa
+
+
+class FluxJointAttention(nn.Module):
+    """Joint attention: text tokens precede image tokens (Diffusers order)."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_added_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_added_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_q_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_k_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.add_v_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
+
+        self.to_out = nn.ModuleList(
+            [
+                ReplicatedLinear(self.inner_dim, dim, bias=True),
+                nn.Dropout(0.0),
+            ]
+        )
+        self.to_add_out = ReplicatedLinear(self.inner_dim, dim, bias=True)
+
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = hidden_states.shape[0]
+        text_seq_len = encoder_hidden_states.shape[1]
+        img_seq_len = hidden_states.shape[1]
+
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+        q = q.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        k = k.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        v = v.view(batch_size, img_seq_len, self.heads, self.head_dim)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+
+        enc_q, _ = self.add_q_proj(encoder_hidden_states)
+        enc_k, _ = self.add_k_proj(encoder_hidden_states)
+        enc_v, _ = self.add_v_proj(encoder_hidden_states)
+        enc_q = enc_q.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_k = enc_k.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_v = enc_v.view(batch_size, text_seq_len, self.heads, self.head_dim)
+        enc_q = self.norm_added_q(enc_q)
+        enc_k = self.norm_added_k(enc_k)
+
+        q = torch.cat([enc_q, q], dim=1)
+        k = torch.cat([enc_k, k], dim=1)
+        v = torch.cat([enc_v, v], dim=1)
+
+        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
+        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
+
+        joint_out, _ = self.attn(q, k, v)
+        joint_out = joint_out.reshape(batch_size, text_seq_len + img_seq_len, self.inner_dim)
+
+        enc_out = joint_out[:, :text_seq_len]
+        img_out = joint_out[:, text_seq_len:]
+
+        img_out, _ = self.to_out[0](img_out)
+        img_out = self.to_out[1](img_out)
+        enc_out, _ = self.to_add_out(enc_out)
+        return img_out, enc_out
+
+
+class FluxSingleStreamAttention(nn.Module):
+    """Self-attention on concatenated text+image sequence (single blocks)."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
+        self.attn = DistributedAttention(
+            num_heads=num_attention_heads,
+            head_size=attention_head_dim,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(hidden_states)
+        v, _ = self.to_v(hidden_states)
+        q = q.view(batch_size, seq_len, self.heads, self.head_dim)
+        k = k.view(batch_size, seq_len, self.heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.heads, self.head_dim)
+        q = self.norm_q(q)
+        k = self.norm_k(k)
+        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
+        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
+        out, _ = self.attn(q, k, v)
+        return out.reshape(batch_size, seq_len, self.inner_dim)
+
+
+class FluxTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        self.norm1 = SD3AdaLayerNormZero(dim)
+        self.norm1_context = SD3AdaLayerNormZero(dim)
+        self.attn = FluxJointAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            supported_attention_backends=supported_attention_backends,
+        )
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = SD3FeedForward(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = SD3FeedForward(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del joint_attention_kwargs
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output)
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class FluxSingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_ratio: float = 4.0,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+    ) -> None:
+        super().__init__()
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.norm = FluxAdaLayerNormZeroSingle(dim)
+        self.proj_mlp = ReplicatedLinear(dim, mlp_hidden_dim, bias=True)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = ReplicatedLinear(dim + mlp_hidden_dim, dim, bias=True)
+        self.attn = FluxSingleStreamAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del joint_attention_kwargs
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)[0])
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)[0]
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+        encoder_hidden_states = hidden_states[:, :text_seq_len]
+        hidden_states = hidden_states[:, text_seq_len:]
+        return encoder_hidden_states, hidden_states
+
+
+class FluxTransformer2DModel(BaseDiT):
+    """FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1."""
+
+    _fsdp_shard_conditions = [
+        lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks."))
+        and n.split(".")[-1].isdigit(),
+    ]
+    _compile_conditions = _fsdp_shard_conditions
+    # HF weight names already match this module layout (cf. SGLang regex maps).
+    param_names_mapping: dict[str, Any] = {}
+    reverse_param_names_mapping: dict[str, Any] = {}
+    lora_param_names_mapping: dict[str, Any] = {}
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        del kwargs
+        super().__init__(config=config, hf_config=hf_config)
+        self.fastvideo_config = config
+        self.hf_config = hf_config
+        arch = config.arch_config
+
+        out_ch = arch.out_channels
+        self.out_channels = out_ch if out_ch is not None else arch.in_channels
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.in_channels
+
+        axes_list = list(arch.axes_dims_rope)
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_list)
+        if arch.guidance_embeds:
+            self.time_text_embed = FluxCombinedTimestepGuidanceTextProjEmbeddings(
+                embedding_dim=self.inner_dim,
+                pooled_projection_dim=arch.pooled_projection_dim,
+            )
+        else:
+            self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+                embedding_dim=self.inner_dim,
+                pooled_projection_dim=arch.pooled_projection_dim,
+            )
+        self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, self.inner_dim)
+        self.x_embedder = ReplicatedLinear(arch.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                )
+                for _ in range(arch.num_layers)
+            ]
+        )
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                )
+                for _ in range(arch.num_single_layers)
+            ]
+        )
+
+        self.norm_out = SD3AdaLayerNormContinuous(
+            self.inner_dim,
+            self.inner_dim,
+            elementwise_affine=False,
+            eps=1e-6,
+            bias=True,
+            norm_type="layer_norm",
+        )
+        self.proj_out = ReplicatedLinear(
+            self.inner_dim,
+            arch.patch_size * arch.patch_size * self.out_channels,
+            bias=True,
+        )
+        self.gradient_checkpointing = False
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        pooled_projections: torch.Tensor | None = None,
+        timestep: torch.LongTensor | torch.Tensor | None = None,
+        img_ids: torch.Tensor | None = None,
+        txt_ids: torch.Tensor | None = None,
+        guidance: torch.Tensor | None = None,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+        controlnet_block_samples: Any | None = None,
+        controlnet_single_block_samples: Any | None = None,
+        controlnet_blocks_repeat: bool = False,
+        **kwargs: Any,
+    ) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]:
+        del kwargs
+        if encoder_hidden_states is None:
+            raise ValueError("encoder_hidden_states must be provided")
+        if pooled_projections is None:
+            raise ValueError("pooled_projections must be provided")
+        if timestep is None:
+            raise ValueError("timestep must be provided")
+        if img_ids is None or txt_ids is None:
+            raise ValueError("img_ids and txt_ids must be provided")
+
+        arch = self.fastvideo_config.arch_config
+        if arch.guidance_embeds and guidance is None:
+            raise ValueError("guidance must be provided when guidance_embeds=True")
+
+        if timestep.dim() == 0:
+            timestep = timestep[None]
+        if timestep.dim() > 1:
+            timestep = timestep.reshape(-1)
+        if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
+            timestep = timestep.expand(hidden_states.shape[0])
+
+        try:
+            get_forward_context()
+            forward_context = nullcontext()
+        except AssertionError:
+            if timestep.numel() == 0:
+                ts0 = 0
+            elif torch.is_floating_point(timestep):
+                ts0 = int(round(timestep[0].item() * 1000))
+            else:
+                ts0 = int(timestep[0].item())
+            forward_context = set_forward_context(current_timestep=ts0, attn_metadata=None)
+
+        with forward_context:
+            hidden_states, _ = self.x_embedder(hidden_states)
+
+            ts = timestep.to(hidden_states.dtype) * 1000
+            g = None if guidance is None else guidance.to(hidden_states.dtype) * 1000
+
+            if arch.guidance_embeds:
+                assert g is not None
+                temb = self.time_text_embed(ts, g, pooled_projections)
+            else:
+                temb = self.time_text_embed(timestep=ts, pooled_projection=pooled_projections)
+
+            encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)
+
+            if txt_ids.ndim == 3:
+                txt_ids = txt_ids[0]
+            if img_ids.ndim == 3:
+                img_ids = img_ids[0]
+
+            ids = torch.cat((txt_ids, img_ids), dim=0)
+            image_rotary_emb = self.pos_embed(ids)
+
+            jkwargs = joint_attention_kwargs or {}
+
+            for idx, block in enumerate(self.transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=jkwargs,
+                )
+                if controlnet_block_samples:
+                    interval = len(self.transformer_blocks) / len(controlnet_block_samples)
+                    interval = int(math.ceil(interval))
+                    if controlnet_blocks_repeat:
+                        hidden_states = hidden_states + controlnet_block_samples[idx % len(controlnet_block_samples)]
+                    else:
+                        hidden_states = hidden_states + controlnet_block_samples[idx // interval]
+
+            for idx, block in enumerate(self.single_transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=jkwargs,
+                )
+                if controlnet_single_block_samples:
+                    interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                    interval = int(math.ceil(interval))
+                    hidden_states = hidden_states + controlnet_single_block_samples[idx // interval]
+
+            hidden_states = self.norm_out(hidden_states, temb)
+            output, _ = self.proj_out(hidden_states)
+
+            if not return_dict:
+                return (output,)
+            return FluxTransformer2DModelOutput(sample=output)
+
+
+EntryClass = FluxTransformer2DModel

fastvideo/models/dits/glm_image.py

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from fastvideo.layers.mlp import MLP
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
+from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb
+from fastvideo.layers.visual_embedding import Timesteps
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class GlmImageLayerKVCache:
+
+    def __init__(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode: Optional[str] = None
+
+    def store(self, k: torch.Tensor, v: torch.Tensor):
+        # Append along seq (dim=1).
+        if self.k_cache is None:
+            self.k_cache = k
+            self.v_cache = v
+        else:
+            self.k_cache = torch.cat([self.k_cache, k], dim=1)
+            self.v_cache = torch.cat([self.v_cache, v], dim=1)
+
+    def get(self):
+        return self.k_cache, self.v_cache
+
+    def clear(self):
+        self.k_cache = None
+        self.v_cache = None
+        self.mode = None
+
+
+class GlmImageKVCache:
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
+
+    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
+        return self.caches[layer_idx]
+
+    def set_mode(self, mode: Optional[str]):
+        if mode is not None and mode not in ["write", "read", "skip"]:
+            raise ValueError(
+                f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'"
+            )
+        for cache in self.caches:
+            cache.mode = mode
+
+    def clear(self):
+        for cache in self.caches:
+            cache.clear()
+
+
+# =============================================================================
+# Timestep and Text Projection
+# =============================================================================
+class GlmImageTimestepEmbedding(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+    ):
+        super().__init__()
+        if out_dim is None:
+            out_dim = time_embed_dim
+        self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True)
+        if act_fn == "silu":
+            self.act = nn.SiLU()
+        elif act_fn == "gelu":
+            self.act = nn.GELU(approximate="tanh")
+        else:
+            self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        sample, _ = self.linear_1(sample)
+        sample = self.act(sample)
+        sample, _ = self.linear_2(sample)
+        return sample
+
+
+class GlmImageTextProjection(nn.Module):
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_size: int,
+        out_features: int = None,
+        act_fn: str = "silu",
+    ):
+        super().__init__()
+        if out_features is None:
+            out_features = hidden_size
+        self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True)
+        if act_fn == "silu":
+            self.act_1 = nn.SiLU()
+        elif act_fn == "gelu_tanh":
+            self.act_1 = nn.GELU(approximate="tanh")
+        else:
+            self.act_1 = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True)
+
+    def forward(self, caption: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.linear_1(caption)
+        hidden_states = self.act_1(hidden_states)
+        hidden_states, _ = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        condition_dim: int,
+        pooled_projection_dim: int,
+        timesteps_dim: int = 256,
+    ):
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.condition_proj = Timesteps(
+            num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = GlmImageTimestepEmbedding(
+            in_channels=timesteps_dim, time_embed_dim=embedding_dim
+        )
+        self.condition_embedder = GlmImageTextProjection(
+            pooled_projection_dim, embedding_dim, act_fn="silu"
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        hidden_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        timesteps_proj = self.time_proj(timestep)
+
+        crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(
+            crop_coords.size(0), -1
+        )
+        target_size_proj = self.condition_proj(target_size.flatten()).view(
+            target_size.size(0), -1
+        )
+
+        condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1)
+
+        timesteps_emb = self.timestep_embedder(
+            timesteps_proj.to(dtype=hidden_dtype)
+        )
+        condition_emb = self.condition_embedder(
+            condition_proj.to(dtype=hidden_dtype)
+        )
+
+        conditioning = timesteps_emb + condition_emb
+        return conditioning
+
+
+# =============================================================================
+# Image Projector
+# =============================================================================
+class GlmImageImageProjector(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int = 16,
+        hidden_size: int = 2560,
+        patch_size: int = 2,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, channel, height, width = hidden_states.shape
+        post_patch_height = height // self.patch_size
+        post_patch_width = width // self.patch_size
+
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            channel,
+            post_patch_height,
+            self.patch_size,
+            post_patch_width,
+            self.patch_size,
+        )
+        hidden_states = (
+            hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+        )
+        hidden_states = self.proj(hidden_states)
+
+        return hidden_states
+
+
+# =============================================================================
+# AdaLayerNorm
+# =============================================================================
+class GlmImageAdaLayerNormZero(nn.Module):
+
+    def __init__(self, embedding_dim: int, dim: int) -> None:
+        super().__init__()
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
+        self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+    ) -> Tuple[torch.Tensor, ...]:
+        dtype = hidden_states.dtype
+        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
+        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(
+            dtype=dtype
+        )
+
+        emb, _ = self.linear(temb)
+        (
+            shift_msa,
+            c_shift_msa,
+            scale_msa,
+            c_scale_msa,
+            gate_msa,
+            c_gate_msa,
+            shift_mlp,
+            c_shift_mlp,
+            scale_mlp,
+            c_scale_mlp,
+            gate_mlp,
+            c_gate_mlp,
+        ) = emb.chunk(12, dim=1)
+
+        hidden_states = norm_hidden_states * (
+            1 + scale_msa.unsqueeze(1)
+        ) + shift_msa.unsqueeze(1)
+        encoder_hidden_states = norm_encoder_hidden_states * (
+            1 + c_scale_msa.unsqueeze(1)
+        ) + c_shift_msa.unsqueeze(1)
+
+        return (
+            hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        )
+
+
+# =============================================================================
+# Attention
+# =============================================================================
+class GlmImageAttention(nn.Module):
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int,
+        dim_head: int,
+        out_dim: int,
+        bias: bool = True,
+        qk_norm: str = "layer_norm",
+        elementwise_affine: bool = False,
+        eps: float = 1e-5,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        self.num_kv_heads = self.heads
+        self.dim_head = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
+
+        self.to_out = nn.ModuleList(
+            [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)]
+        )
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "layer_norm":
+            self.norm_q = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+            self.norm_k = nn.LayerNorm(
+                dim_head, eps=eps, elementwise_affine=elementwise_affine
+            )
+        else:
+            raise ValueError(f"unknown qk_norm: {qk_norm}")
+
+        self.attn = LocalAttention(
+            num_heads=self.heads,
+            head_size=dim_head,
+            num_kv_heads=self.heads,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        dtype = encoder_hidden_states.dtype
+
+        batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
+        batch_size, image_seq_length, embed_dim = hidden_states.shape
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # 1. QKV projections
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(hidden_states)
+        value, _ = self.to_v(hidden_states)
+
+        query = query.unflatten(2, (self.heads, -1))
+        key = key.unflatten(2, (self.heads, -1))
+        value = value.unflatten(2, (self.heads, -1))
+
+        # 2. QK normalization
+        if self.norm_q is not None:
+            query = self.norm_q(query).to(dtype=dtype)
+        if self.norm_k is not None:
+            key = self.norm_k(key).to(dtype=dtype)
+
+        # 3. Rotational positional embeddings applied to latent stream
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+
+            query[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+            )
+            key[:, text_seq_length:, :, :] = _apply_rotary_emb(
+                key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
+            )
+
+        # 4. KV Cache handling
+        if kv_cache is not None:
+            if kv_cache.mode == "write":
+                kv_cache.store(key, value)
+            elif kv_cache.mode == "read":
+                # Prepend cached condition k/v along seq (dim=1).
+                k_cache, v_cache = kv_cache.get()
+                key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key
+                value = (
+                    torch.cat([v_cache, value], dim=1) if v_cache is not None else value
+                )
+            elif kv_cache.mode == "skip":
+                pass
+
+        hidden_states = self.attn(query, key, value)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # 6. Output projection
+        hidden_states, _ = self.to_out[0](hidden_states)
+
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+        return hidden_states, encoder_hidden_states
+
+
+# =============================================================================
+# Transformer Block
+# =============================================================================
+class GlmImageTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        dim: int = 2560,
+        num_attention_heads: int = 64,
+        attention_head_dim: int = 40,
+        time_embed_dim: int = 512,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        # 1. Attention
+        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
+
+        self.attn1 = GlmImageAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            out_dim=dim,
+            bias=True,
+            qk_norm="layer_norm",
+            elementwise_affine=False,
+            eps=1e-5,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn1",
+        )
+
+        # 2. Feedforward with fused ScaleResidualLayerNorm
+        self.norm2 = ScaleResidualLayerNormScaleShift(
+            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
+        )
+        self.norm2_context = ScaleResidualLayerNormScaleShift(
+            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
+        )
+        self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_cache: Optional[GlmImageLayerKVCache] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Timestep conditioning
+        (
+            norm_hidden_states,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            norm_encoder_hidden_states,
+            c_gate_msa,
+            c_shift_mlp,
+            c_scale_mlp,
+            c_gate_mlp,
+        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
+
+        # 2. Attention
+        if attention_kwargs is None:
+            attention_kwargs = {}
+
+        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            kv_cache=kv_cache,
+            **attention_kwargs,
+        )
+
+        # 3. Feedforward (fused residual + norm + scale/shift)
+        norm_hidden_states, hidden_states = self.norm2(
+            hidden_states,
+            attn_hidden_states,
+            gate_msa.unsqueeze(1),
+            shift_mlp.unsqueeze(1),
+            scale_mlp.unsqueeze(1),
+        )
+        norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context(
+            encoder_hidden_states,
+            attn_encoder_hidden_states,
+            c_gate_msa.unsqueeze(1),
+            c_shift_mlp.unsqueeze(1),
+            c_scale_mlp.unsqueeze(1),
+        )
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output_context = self.ff(norm_encoder_hidden_states)
+        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = (
+            encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
+        )
+
+        return hidden_states, encoder_hidden_states
+
+
+# =============================================================================
+# Rotary Positional Embedding
+# =============================================================================
+class GlmImageRotaryPosEmbed(nn.Module):
+
+    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.theta = theta
+        self._cache_key: tuple | None = None
+        self._cache_value: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        _, _, raw_h, raw_w = hidden_states.shape
+        height = raw_h // self.patch_size
+        width = raw_w // self.patch_size
+        device = hidden_states.device
+
+        cache_key = (height, width, device.type,
+                     device.index if device.index is not None else -1)
+        if self._cache_key == cache_key and self._cache_value is not None:
+            return self._cache_value
+
+        dim_h, dim_w = self.dim // 2, self.dim // 2
+        h_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[
+                    : (dim_h // 2)
+                ]
+                / dim_h
+            )
+        )
+        w_inv_freq = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[
+                    : (dim_w // 2)
+                ]
+                / dim_w
+            )
+        )
+        h_seq = torch.arange(height, device=device)
+        w_seq = torch.arange(width, device=device)
+        freqs_h = torch.outer(h_seq, h_inv_freq).unsqueeze(1).expand(height, width, -1)
+        freqs_w = torch.outer(w_seq, w_inv_freq).unsqueeze(0).expand(height, width, -1)
+
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1).reshape(height * width, -1)
+        result = (freqs.cos(), freqs.sin())
+        self._cache_key = cache_key
+        self._cache_value = result
+        return result
+
+
+# =============================================================================
+# Final AdaLayerNorm
+# =============================================================================
+class GlmImageAdaLayerNormContinuous(nn.Module):
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine: bool = True,
+        eps: float = 1e-5,
+        bias: bool = True,
+        norm_type: str = "layer_norm",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(
+            conditioning_embedding_dim, embedding_dim * 2, bias=bias
+        )
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor
+    ) -> torch.Tensor:
+        emb = self.linear(conditioning_embedding.to(x.dtype))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# =============================================================================
+# Main Model
+# =============================================================================
+class GlmImageTransformer2DModel(BaseDiT):
+
+    _fsdp_shard_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = GlmImageDiTConfig().arch_config._compile_conditions
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    param_names_mapping = GlmImageDiTConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+    lora_param_names_mapping = {}
+
+    def __init__(
+        self,
+        config: GlmImageDiTConfig,
+        hf_config: dict[str, Any],
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+
+        arch_config = config.arch_config
+
+        self.in_channels = arch_config.in_channels
+        self.out_channels = arch_config.out_channels
+        self.patch_size = arch_config.patch_size
+        self.num_layers = arch_config.num_layers
+        self.attention_head_dim = arch_config.attention_head_dim
+        self.num_attention_heads = arch_config.num_attention_heads
+        self.text_embed_dim = arch_config.text_embed_dim
+        self.time_embed_dim = arch_config.time_embed_dim
+
+        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
+        # Each of these are sincos embeddings of shape 2 * condition_dim
+        pooled_projection_dim = 2 * 2 * arch_config.condition_dim
+        inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
+
+        self.hidden_size = inner_dim
+        self.num_channels_latents = arch_config.out_channels
+
+        # 1. RoPE
+        self.rotary_emb = GlmImageRotaryPosEmbed(
+            arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0
+        )
+
+        # 2. Patch & Text-timestep embedding
+        self.image_projector = GlmImageImageProjector(
+            arch_config.in_channels, inner_dim, arch_config.patch_size
+        )
+        self.glyph_projector = MLP(
+            input_dim=arch_config.text_embed_dim,
+            mlp_hidden_dim=inner_dim,
+            output_dim=inner_dim,
+            act_type="gelu",
+        )
+        self.prior_token_embedding = nn.Embedding(
+            arch_config.prior_vq_quantizer_codebook_size, inner_dim
+        )
+        self.prior_projector = MLP(
+            input_dim=inner_dim,
+            mlp_hidden_dim=inner_dim,
+            output_dim=inner_dim,
+            act_type="silu",
+        )
+
+        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
+            embedding_dim=arch_config.time_embed_dim,
+            condition_dim=arch_config.condition_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            timesteps_dim=arch_config.time_embed_dim,
+        )
+
+        # 3. Transformer blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                GlmImageTransformerBlock(
+                    inner_dim,
+                    arch_config.num_attention_heads,
+                    arch_config.attention_head_dim,
+                    arch_config.time_embed_dim,
+                    supported_attention_backends=self._supported_attention_backends,
+                    prefix=f"transformer_blocks.{i}",
+                )
+                for i in range(arch_config.num_layers)
+            ]
+        )
+
+        # 4. Output projection
+        self.norm_out = GlmImageAdaLayerNormContinuous(
+            inner_dim, arch_config.time_embed_dim, elementwise_affine=False
+        )
+        self.proj_out = nn.Linear(
+            inner_dim,
+            arch_config.patch_size * arch_config.patch_size * arch_config.out_channels,
+            bias=True,
+        )
+
+        self.gradient_checkpointing = False
+
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        prior_token_id: torch.Tensor,
+        prior_token_drop: torch.Tensor,
+        timestep: torch.LongTensor,
+        target_size: torch.Tensor,
+        crop_coords: torch.Tensor,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_caches: Optional[GlmImageKVCache] = None,
+        kv_caches_mode: Optional[str] = None,
+        freqs_cis: Optional[
+            Union[
+                Tuple[torch.Tensor, torch.Tensor],
+                List[Tuple[torch.Tensor, torch.Tensor]],
+            ]
+        ] = None,
+        guidance: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if kv_caches is not None:
+            kv_caches.set_mode(kv_caches_mode)
+
+        batch_size, num_channels, height, width = hidden_states.shape
+
+        if isinstance(encoder_hidden_states, list):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        # 1. RoPE
+        image_rotary_emb = freqs_cis
+        if image_rotary_emb is None:
+            image_rotary_emb = self.rotary_emb(hidden_states)
+
+        # 2. Patch & Timestep embeddings
+        p = self.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.image_projector(hidden_states)
+        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
+        prior_embedding = self.prior_token_embedding(prior_token_id)
+        # Zero dropped priors by multiply: boolean indexing + .any() syncs each step.
+        keep = (~prior_token_drop).to(device=prior_embedding.device, dtype=prior_embedding.dtype)
+        while keep.dim() < prior_embedding.dim():
+            keep = keep.unsqueeze(-1)
+        prior_embedding = prior_embedding * keep
+        prior_hidden_states = self.prior_projector(prior_embedding)
+        hidden_states = hidden_states + prior_hidden_states
+
+        temb = self.time_condition_embed(
+            timestep, target_size, crop_coords, hidden_states.dtype
+        )
+        temb = F.silu(temb)
+
+        # 3. Transformer blocks
+        for idx, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                image_rotary_emb,
+                attention_kwargs,
+                kv_cache=kv_caches[idx] if kv_caches is not None else None,
+            )
+
+        # 4. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # 5. Unpatchify
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_height, post_patch_width, -1, p, p
+        )
+        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+
+        return output.float()
+
+
+EntryClass = GlmImageTransformer2DModel

fastvideo/models/dits/matrixgame2/action_module.py

@@ -16,6 +16,8 @@ from fastvideo.layers.rotary_embedding import (
 )
 from fastvideo.platforms import AttentionBackendEnum
 
+from .utils import retain_kv_with_sink
+
 
 DISABLE_COMPILE = False
 flex_attention = torch.compile(
@@ -100,6 +102,7 @@ def _update_kv_cache_and_attend(
     start_frame: int,
     num_frame_per_block: int,
     local_attn_size: int,
+    sink_size: int,
     *,
     use_k_for_num_tokens: bool = False,
     store_first_only: bool = False,
@@ -133,7 +136,6 @@ def _update_kv_cache_and_attend(
         k.shape[1] if use_k_for_num_tokens else q.shape[1]
     ) == num_frame_per_block
 
-    sink_size = 0
     max_attention_size = local_attn_size
     sink_tokens = sink_size * 1
     kv_cache_size = kv_cache["k"].shape[1]
@@ -206,16 +208,22 @@ def _update_kv_cache_and_attend(
 
     # Store new k, v in cache
     if store_first_only:
-        kv_cache["k"][:, local_start_index:local_end_index] = k[:1]
-        kv_cache["v"][:, local_start_index:local_end_index] = v[:1]
+        kv_cache["k"][:, local_start_index:local_end_index] = k[
+            : kv_cache["k"].shape[0]
+        ]
+        kv_cache["v"][:, local_start_index:local_end_index] = v[
+            : kv_cache["v"].shape[0]
+        ]
     else:
         kv_cache["k"][:, local_start_index:local_end_index] = k
         kv_cache["v"][:, local_start_index:local_end_index] = v
 
-    # Retrieve from cache and perform attention
-    cache_start = max(0, local_end_index - max_attention_size)
-    cached_k = kv_cache["k"][:, cache_start:local_end_index]
-    cached_v = kv_cache["v"][:, cache_start:local_end_index]
+    cached_k, cached_v = retain_kv_with_sink(
+        kv_cache["k"][:, :local_end_index],
+        kv_cache["v"][:, :local_end_index],
+        min(local_end_index, max_attention_size),
+        sink_tokens,
+    )
 
     if repeat_factor is not None:
         cached_k = cached_k.repeat(repeat_factor, 1, 1, 1)
@@ -261,6 +269,7 @@ class ActionModule(nn.Module):
         enable_mouse=True,
         enable_keyboard=True,
         local_attn_size=6,
+        sink_size=0,
         blocks: list | None = None,
     ):
         super().__init__()
@@ -274,6 +283,7 @@ class ActionModule(nn.Module):
         )
         blocks = blocks if blocks is not None else []
         self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
         self.enable_mouse = enable_mouse
         self.enable_keyboard = enable_keyboard
 
@@ -565,6 +575,7 @@ class ActionModule(nn.Module):
                     start_frame,
                     num_frame_per_block,
                     self.local_attn_size,
+                    self.sink_size,
                     use_k_for_num_tokens=False,
                 )
         else:
@@ -689,6 +700,7 @@ class ActionModule(nn.Module):
                         start_frame,
                         num_frame_per_block,
                         self.local_attn_size,
+                        self.sink_size,
                         use_k_for_num_tokens=True,
                         store_first_only=True,
                         repeat_factor=S,
@@ -719,6 +731,7 @@ class ActionModule(nn.Module):
                         start_frame,
                         num_frame_per_block,
                         self.local_attn_size,
+                        self.sink_size,
                         use_k_for_num_tokens=True,
                     )
             else:

fastvideo/models/dits/matrixgame2/causal_model.py

@@ -46,6 +46,7 @@ from fastvideo.platforms import AttentionBackendEnum, current_platform
 
 from .action_module import ActionModule
 from .model import MatrixGame2CrossAttention
+from .utils import retain_kv_with_sink
 
 logger = init_logger(__name__)
 
@@ -345,9 +346,12 @@ class CausalMatrixGame2SelfAttention(nn.Module):
                 kv_cache["k"][:, local_start_index:local_end_index] = stored_key
                 kv_cache["v"][:, local_start_index:local_end_index] = v
 
-            kv_start = max(0, local_end_index - max_attention_size)
-            k_for_attn = kv_cache["k"][:, kv_start:local_end_index]
-            v_for_attn = kv_cache["v"][:, kv_start:local_end_index]
+            k_for_attn, v_for_attn = retain_kv_with_sink(
+                kv_cache["k"][:, :local_end_index],
+                kv_cache["v"][:, :local_end_index],
+                min(local_end_index, max_attention_size),
+                sink_tokens,
+            )
 
             if relativistic:
                 window_len, query_lo, _ = relativistic_window_offsets(
@@ -418,6 +422,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
         self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
         dim_head = dim // num_heads
         if qk_norm == "rms_norm":
             self.norm_q = RMSNorm(dim_head, eps=eps)
@@ -474,6 +479,7 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
                 ),
                 patch_size=action_config["patch_size"],
                 local_attn_size=local_attn_size,
+                sink_size=sink_size,
                 qk_norm=action_config["qk_norm"],
                 qkv_bias=action_config["qkv_bias"],
                 vae_time_compression_ratio=action_config[
@@ -647,6 +653,8 @@ class CausalMatrixGame2WanModel(BaseDiT):
         self.num_channels_latents = config.num_channels_latents
         self.patch_size = config.patch_size
 
+        # Long tuning controls the causal window from YAML; consume it here
+        # like Wan so train wrappers do not need runtime patching.
         arch_cfg = getattr(config, "arch_config", None)
         self.local_attn_size = (
             getattr(
@@ -662,6 +670,12 @@ class CausalMatrixGame2WanModel(BaseDiT):
             if arch_cfg
             else getattr(config, "sink_size", 0)
         )
+        if self.sink_size < 0:
+            raise ValueError("sink_size must be non-negative")
+        if self.local_attn_size != -1 and self.sink_size >= self.local_attn_size:
+            raise ValueError(
+                "sink_size must be smaller than local_attn_size for "
+                "MatrixGame2 causal attention")
         self.rope_cache_policy = (
             getattr(arch_cfg, "rope_cache_policy",
                     getattr(config, "rope_cache_policy", "absolute"))
@@ -749,6 +763,52 @@ class CausalMatrixGame2WanModel(BaseDiT):
 
         self.__post_init__()
 
+    def set_causal_attention_window(
+        self,
+        *,
+        local_attn_size: int | None = None,
+        sink_size: int | None = None,
+    ) -> None:
+        local_attn_size = (
+            self.local_attn_size if local_attn_size is None else int(local_attn_size)
+        )
+        sink_size = self.sink_size if sink_size is None else int(sink_size)
+        if sink_size < 0:
+            raise ValueError("sink_size must be non-negative")
+        if local_attn_size != -1 and sink_size >= local_attn_size:
+            raise ValueError(
+                "sink_size must be smaller than local_attn_size for "
+                "MatrixGame2 causal attention")
+
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.block_mask = None
+        self.block_mask_keyboard = None
+        self.block_mask_mouse = None
+
+        for block in self.blocks:
+            if hasattr(block, "local_attn_size"):
+                block.local_attn_size = local_attn_size
+            if hasattr(block, "sink_size"):
+                block.sink_size = sink_size
+
+            attn1 = getattr(block, "attn1", None)
+            if attn1 is not None:
+                if hasattr(attn1, "local_attn_size"):
+                    attn1.local_attn_size = local_attn_size
+                if hasattr(attn1, "sink_size"):
+                    attn1.sink_size = sink_size
+
+            action_model = getattr(block, "action_model", None)
+            if action_model is not None:
+                if hasattr(action_model, "local_attn_size"):
+                    action_model.local_attn_size = local_attn_size
+                if hasattr(action_model, "sink_size"):
+                    action_model.sink_size = sink_size
+
+    def set_sink_size(self, sink_size: int) -> None:
+        self.set_causal_attention_window(sink_size=sink_size)
+
     @staticmethod
     def _prepare_blockwise_causal_attn_mask(
         device: torch.device | str,
@@ -756,6 +816,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
         frame_seqlen: int = 880,
         num_frame_per_block: int = 1,
         local_attn_size: int = -1,
+        sink_size: int = 0,
     ) -> BlockMask:
         total_length = num_frames * frame_seqlen
         padded_length = math.ceil(total_length / 128) * 128 - total_length
@@ -782,7 +843,10 @@ class CausalMatrixGame2WanModel(BaseDiT):
             else:
                 return (
                     (kv_idx < ends[q_idx])
-                    & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
+                    & (
+                        (kv_idx < sink_size * frame_seqlen)
+                        | (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))
+                    )
                 ) | (q_idx == kv_idx)
 
         block_mask = create_block_mask(
@@ -796,8 +860,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
         )
 
         if not dist.is_initialized() or dist.get_rank() == 0:
-            print(
-                f" cache a block wise causal mask with block size of {num_frame_per_block} frames"
+            logger.info(
+                "cache a block wise causal mask with block size of %s frames",
+                num_frame_per_block,
             )
 
         return block_mask
@@ -809,6 +874,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
         frame_seqlen: int = 880,
         num_frame_per_block: int = 1,
         local_attn_size: int = -1,
+        sink_size: int = 0,
     ) -> BlockMask:
         total_length2 = num_frames * frame_seqlen
         padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
@@ -834,7 +900,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             else:
                 return (
                     (kv_idx < ends2[q_idx])
-                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
+                    & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
                 ) | (q_idx == kv_idx)
 
         block_mask2 = create_block_mask(
@@ -848,8 +914,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
         )
 
         if not dist.is_initialized() or dist.get_rank() == 0:
-            print(
-                f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames"
+            logger.info(
+                "cache a block wise causal mask for keyboard with block size of %s frames",
+                num_frame_per_block,
             )
 
         return block_mask2
@@ -861,6 +928,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
         frame_seqlen: int = 1,
         num_frame_per_block: int = 1,
         local_attn_size: int = -1,
+        sink_size: int = 0,
     ) -> BlockMask:
         total_length2 = num_frames * frame_seqlen
         padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
@@ -886,7 +954,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             else:
                 return (
                     (kv_idx < ends2[q_idx])
-                    & (kv_idx >= (ends2[q_idx] - local_attn_size))
+                    & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))
                 ) | (q_idx == kv_idx)
 
         block_mask2 = create_block_mask(
@@ -900,8 +968,9 @@ class CausalMatrixGame2WanModel(BaseDiT):
         )
 
         if not dist.is_initialized() or dist.get_rank() == 0:
-            print(
-                f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames"
+            logger.info(
+                "cache a block wise causal mask for action with block size of %s frames",
+                num_frame_per_block,
             )
 
         return block_mask2
@@ -1063,6 +1132,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             frame_seqlen=post_patch_height * post_patch_width,
             num_frame_per_block=self.num_frame_per_block,
             local_attn_size=self.local_attn_size,
+            sink_size=self.sink_size,
         )
         if self.use_rope_keyboard:
             block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action(
@@ -1071,6 +1141,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                 frame_seqlen=1,
                 num_frame_per_block=self.num_frame_per_block,
                 local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
             )
         else:
             block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
@@ -1079,6 +1150,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                 frame_seqlen=post_patch_height * post_patch_width,
                 num_frame_per_block=self.num_frame_per_block,
                 local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
             )
         block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
             device=hidden_states.device,
@@ -1086,6 +1158,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
             frame_seqlen=1,
             num_frame_per_block=self.num_frame_per_block,
             local_attn_size=self.local_attn_size,
+            sink_size=self.sink_size,
         )
         if kv_cache is None:
             kv_cache = [None] * len(self.blocks)
@@ -1239,6 +1312,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                 frame_seqlen=post_patch_height * post_patch_width,
                 num_frame_per_block=self.num_frame_per_block,
                 local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
             )
         if self.block_mask_keyboard is None:
             if self.use_rope_keyboard:
@@ -1248,6 +1322,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                     frame_seqlen=1,
                     num_frame_per_block=self.num_frame_per_block,
                     local_attn_size=self.local_attn_size,
+                    sink_size=self.sink_size,
                 )
             else:
                 self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
@@ -1256,6 +1331,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                     frame_seqlen=post_patch_height * post_patch_width,
                     num_frame_per_block=self.num_frame_per_block,
                     local_attn_size=self.local_attn_size,
+                    sink_size=self.sink_size,
                 )
         if self.block_mask_mouse is None:
             self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
@@ -1264,6 +1340,7 @@ class CausalMatrixGame2WanModel(BaseDiT):
                 frame_seqlen=1,
                 num_frame_per_block=self.num_frame_per_block,
                 local_attn_size=self.local_attn_size,
+                sink_size=self.sink_size,
             )
 
         hidden_states = self.patch_embedding(hidden_states)

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]
-        else:
+        elif not isinstance(encoder_hidden_states_image, torch.Tensor):
             encoder_hidden_states_image = None
 
         batch_size, num_channels, num_frames, height, width = (

fastvideo/models/dits/matrixgame2/utils.py

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

fastvideo/models/dits/wanvideo.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import copy
 import math
 from typing import Any
 
@@ -59,6 +60,11 @@ 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__()
 
@@ -77,14 +83,57 @@ 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:
@@ -595,6 +644,10 @@ 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
@@ -636,6 +689,7 @@ 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):
@@ -684,8 +738,17 @@ 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)
+            timestep,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+            timestep_seq_len=ts_seq_len,
+            r_timestep=r_timestep)
         if ts_seq_len is not None:
             # batch_size, seq_len, 6, inner_dim
             timestep_proj = timestep_proj.unflatten(2, (6, -1))

fastvideo/models/encoders/glm_image_ar_loader.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sole HF-import boundary for GLM-Image's AR encoder (lazy-wrapper exception E001)."""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class GlmImageARLoader(nn.Module):
+
+    def __init__(self, model_path: str, processor_path: str | None = None,
+                 *, torch_dtype: torch.dtype = torch.bfloat16,
+                 trust_remote_code: bool = True) -> None:
+        super().__init__()
+        from transformers import (AutoProcessor,
+                                  GlmImageForConditionalGeneration)
+        logger.info("Loading GLM-Image AR encoder from %s", model_path)
+        self._model = GlmImageForConditionalGeneration.from_pretrained(
+            model_path,
+            torch_dtype=torch_dtype,
+            trust_remote_code=trust_remote_code,
+        )
+        if processor_path is not None:
+            logger.info("Loading GLM-Image processor from %s", processor_path)
+            self.processor = AutoProcessor.from_pretrained(
+                processor_path, trust_remote_code=trust_remote_code)
+        else:
+            self.processor = None
+
+    @torch.no_grad()
+    def generate(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        return self._model.generate(*args, **kwargs)
+
+    @torch.no_grad()
+    def get_image_features(self, pixel_values: torch.Tensor,
+                           image_grid_thw: torch.Tensor) -> Any:
+        return self._model.get_image_features(pixel_values, image_grid_thw)
+
+    @torch.no_grad()
+    def get_image_tokens(self, image_embeds: torch.Tensor,
+                         image_grid_thw: torch.Tensor) -> torch.Tensor:
+        return self._model.get_image_tokens(image_embeds, image_grid_thw)
+
+    @property
+    def config(self):  # type: ignore[no-untyped-def]
+        return self._model.config
+
+    @property
+    def generation_config(self):  # type: ignore[no-untyped-def]
+        return self._model.generation_config
+
+    def to(self, *args, **kwargs):  # type: ignore[override]
+        self._model = self._model.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+    def eval(self):  # type: ignore[override]
+        self._model = self._model.eval()
+        return super().eval()

fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py

@@ -42,7 +42,7 @@ def load_independent(files: list[str], device: str):
     """Before-PR behavior: every rank reads every tensor from disk to GPU."""
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f:
+            for name in f.keys():  # noqa: SIM118
                 param = f.get_tensor(name)
                 yield name, param
 
@@ -54,7 +54,7 @@ def load_broadcast(files: list[str], device: str, node_group,
     handles = []
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f:
+            for name in f.keys():  # noqa: SIM118
                 if local_rank == 0:
                     param = f.get_tensor(name)
                 else:

fastvideo/models/loader/component_loader.py

@@ -92,9 +92,13 @@ class ComponentLoader(ABC):
             "tokenizer": (TokenizerLoader, "transformers"),
             "tokenizer_2": (TokenizerLoader, "transformers"),
             "tokenizer_3": (TokenizerLoader, "transformers"),
+            # Cosmos3's model_index names its Qwen2 tokenizer "text_tokenizer".
+            "text_tokenizer": (TokenizerLoader, "transformers"),
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
+            "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
+            "processor": (ProcessorLoader, "transformers"),
             "upsampler": (UpsamplerLoader, "diffusers"),
             "upsampler_2": (UpsamplerLoader, "diffusers"),
             # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
@@ -523,6 +527,39 @@ class ImageEncoderLoader(TextEncoderLoader):
         )
 
 
+class VisionLanguageEncoderLoader(ComponentLoader):
+    """Loader for vision-language autoregressive encoders."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        from fastvideo.distributed.parallel_state import get_local_torch_device
+        from fastvideo.models.encoders.glm_image_ar_loader import (
+            GlmImageARLoader)
+
+        logger.info("Loading vision-language encoder from %s", model_path)
+        target_device = get_local_torch_device()
+        loader = GlmImageARLoader(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=fastvideo_args.trust_remote_code,
+        ).to(target_device).eval()
+        return loader
+
+
+class ProcessorLoader(ComponentLoader):
+    """Loader for HF processors that pair with vision-language encoders."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        from transformers import AutoProcessor
+
+        logger.info("Loading processor from %s", model_path)
+        processor = AutoProcessor.from_pretrained(
+            model_path,
+            trust_remote_code=fastvideo_args.trust_remote_code,
+        )
+        logger.info("Loaded processor: %s", processor.__class__.__name__)
+        return processor
+
+
 class ImageProcessorLoader(ComponentLoader):
     """Loader for image processor."""
 
@@ -1102,7 +1139,19 @@ class SchedulerLoader(ComponentLoader):
 
         scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
 
-        scheduler = scheduler_cls(**config)
+        # Diffusers checkpoints can carry newer scheduler config keys than the
+        # vendored scheduler accepts (e.g. shift_terminal / sigma_min / sigma_max
+        # from a newer diffusers release). Filter to the class's __init__ params,
+        # mirroring diffusers' ``from_config``, so loading is robust to schema
+        # drift instead of crashing on an unexpected kwarg.
+        import inspect
+        valid_params = set(inspect.signature(scheduler_cls.__init__).parameters)
+        filtered_config = {k: v for k, v in config.items() if k in valid_params}
+        dropped = sorted(set(config) - set(filtered_config))
+        if dropped:
+            logger.warning("Scheduler %s: dropping unsupported config keys %s", class_name, dropped)
+
+        scheduler = scheduler_cls(**filtered_config)
         if fastvideo_args.pipeline_config.flow_shift is not None:
             scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
         return scheduler

fastvideo/models/registry.py

@@ -37,6 +37,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
     "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
+    # Cosmos3-Nano's checkpoint model_index names the DiT "Cosmos3OmniTransformer";
+    # map that HF class name to FastVideo's native Cosmos3VFMTransformer.
+    "Cosmos3OmniTransformer": ("dits", "cosmos3", "Cosmos3VFMTransformer"),
     "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
     "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
     "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
@@ -61,6 +64,11 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
 }
 
+# Text-to-image DiT models (2D image generation)
+_TEXT_TO_IMAGE_DIT_MODELS = {
+    "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"),
+}
+
 _TEXT_ENCODER_MODELS = {
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection":
@@ -134,6 +142,7 @@ _UPSAMPLERS = {
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
+    **_TEXT_TO_IMAGE_DIT_MODELS,
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,

fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py

@@ -0,0 +1,202 @@
+# 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/vaes/oobleck.py

@@ -94,6 +94,9 @@ class OobleckDecoderBlock(nn.Module):
             input_dim, output_dim,
             kernel_size=2 * stride, stride=stride,
             padding=math.ceil(stride / 2),
+            # Clean L*stride upsample for both parities; a no-op (0) for even
+            # strides (Stable Audio), needed for odd strides (Cosmos3: 5).
+            output_padding=stride % 2,
         ))
         self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
         self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)

fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py

@@ -0,0 +1,708 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Cosmos3 video pipeline (T2V / I2V / T2I).
+
+This replaces the earlier vllm-omni-derived skeleton with a native, stage-based
+:class:`ComposedPipelineBase` pipeline that wires the framework-parity-verified
+Cosmos3 components:
+
+* tokenizer: Qwen2 ``Qwen2TokenizerFast`` + chat template (the only allowed
+  third-party model-adjacent dependency; tokenizers are explicitly permitted),
+* VAE: FastVideo-native ``AutoencoderKLWan`` (Wan2.2) via ``Cosmos3VAEConfig``;
+  encode normalizes ``(mu - mean) * inv_std`` and decode denormalizes + clamps,
+* sequence-packing: :func:`pack_cosmos3_video_sequence` (native, parity-tested),
+* DiT: ``Cosmos3VFMTransformer`` (native, bit-identical to the framework),
+* scheduler: FastVideo-native ``UniPCMultistepScheduler`` configured for pure
+  flow matching (``flow_prediction`` + ``use_flow_sigmas``), numerically
+  equivalent to the framework's ``FlowUniPCMultistepScheduler`` (parity-tested
+  in ``test_cosmos3_scheduler_parity``).
+
+The denoise/CFG glue is a faithful port of the framework's
+``Cosmos3OmniDiffusersPipeline`` math (mirrored in the framework-equivalent
+``diffusers_cosmos3.pipeline``): per UniPC timestep, run a SEQUENTIAL conditional
+then unconditional pass (each repacks the sequence with the prompt / negative
+prompt token ids, forwards the DiT, and zeros the prediction on conditioning
+frames), then combine ``v = uncond + guidance * (cond - uncond)`` and take one
+``scheduler.step(model_output=v, timestep, sample=latent)``. ``timestep_scale``
+is applied to the per-token timesteps *inside* the DiT (its ``forward`` already
+multiplies ``vision_timesteps * timestep_scale`` before the time embedder), so
+the loop passes raw scheduler timesteps to the packer.
+
+The pure denoise math lives in :class:`Cosmos3DenoiseEngine` and the free
+function :func:`cosmos3_get_cfg_velocity` so it can be unit-/parity-tested
+directly against the framework oracle without constructing the full pipeline.
+
+No diffusers/transformers *model* classes are imported at runtime here; only the
+Qwen2 tokenizer (loaded by the component loader) and the UniPC scheduler are
+third-party, both explicitly allowed.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_unipc_multistep import (
+    UniPCMultistepScheduler, )
+from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
+    Cosmos3ActionItem,
+    Cosmos3SampleInputs,
+    Cosmos3SoundItem,
+    Cosmos3VisionItem,
+    pack_cosmos3_video_sequence,
+)
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+logger = init_logger(__name__)
+
+# System prompts, verbatim from the framework (diffusers_cosmos3.pipeline).
+_SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a give prompt."
+_SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a give prompt."
+
+
+# ===========================================================================
+# Special-token resolution (Qwen2 chat tokenizer)
+# ===========================================================================
+def cosmos3_special_tokens(tokenizer: Any) -> dict[str, int]:
+    """Resolve the Cosmos3 generation special tokens from a Qwen2 tokenizer.
+
+    Mirrors the framework's ``llm_special_tokens``:
+    ``start_of_generation=<|vision_start|>``, ``end_of_generation=<|vision_end|>``,
+    ``eos_token_id=tokenizer.eos_token_id``.
+    """
+    return {
+        "start_of_generation": int(tokenizer.convert_tokens_to_ids("<|vision_start|>")),
+        "end_of_generation": int(tokenizer.convert_tokens_to_ids("<|vision_end|>")),
+        "eos_token_id": int(tokenizer.eos_token_id),
+    }
+
+
+def cosmos3_tokenize_caption(
+    tokenizer: Any,
+    caption: str,
+    *,
+    is_video: bool = False,
+    use_system_prompt: bool = False,
+) -> list[int]:
+    """Tokenize a caption with the Qwen2 chat template (framework-faithful).
+
+    Optionally prepends an image/video system prompt; always adds the
+    generation prompt and disables ``add_vision_id`` (matching the framework's
+    ``tokenize_caption``).
+    """
+    conversations: list[dict[str, str]] = []
+    if use_system_prompt:
+        conversations.append({
+            "role": "system",
+            "content": _SYSTEM_PROMPT_VIDEO if is_video else _SYSTEM_PROMPT_IMAGE,
+        })
+    conversations.append({"role": "user", "content": caption})
+    token_ids = tokenizer.apply_chat_template(
+        conversations,
+        tokenize=True,
+        add_generation_prompt=True,
+        add_vision_id=False,
+    )
+    return list(token_ids)
+
+
+# ===========================================================================
+# Reasoning (VLM text generation) — und (causal) pathway + lm_head
+# ===========================================================================
+def cosmos3_generate_reasoner_text(
+    transformer: Any,
+    input_ids: list[int],
+    max_new_tokens: int,
+    *,
+    eos_token_id: int | list[int] | None = None,
+) -> list[int]:
+    """Greedy text reasoning via the und (causal) backbone + ``lm_head``.
+
+    Mirrors the framework ``generate_reasoner_text`` (text-only prefill, greedy):
+    only the und-pathway weights (no ``_moe_gen``) + ``embed_tokens`` / ``norm`` /
+    ``lm_head`` participate; the generation pathway and the VFM multimodal
+    embedders are bypassed (no vision/sound/action tokens). Token-for-token
+    identical to the framework reasoner (``test_cosmos3_reasoning_parity``).
+
+    Re-prefills each step (no KV cache) — correctness-first; a KV-cache fast path
+    is a later optimization. Returns the newly generated token ids.
+    """
+    device = next(transformer.parameters()).device
+    ids = [int(x) for x in input_ids]
+    eos: set[int] = set()
+    if eos_token_id is not None:
+        eos = {int(eos_token_id)} if isinstance(eos_token_id, int) else {int(x) for x in eos_token_id}
+
+    new_tokens: list[int] = []
+    for _ in range(int(max_new_tokens)):
+        n = len(ids)
+        pos = torch.arange(n).unsqueeze(0).expand(3, -1).contiguous().to(device)
+        out = transformer(
+            text_ids=torch.tensor(ids, device=device, dtype=torch.long),
+            text_indexes=torch.arange(n, device=device),
+            position_ids=pos,
+            sequence_length=n,
+            split_lens=[n],
+            attn_modes=["causal"],
+            vision_tokens=[],
+            vision_token_shapes=[],
+            vision_sequence_indexes=torch.empty(0, dtype=torch.long, device=device),
+            vision_timesteps=torch.empty(0, device=device),
+            vision_mse_loss_indexes=torch.empty(0, dtype=torch.long, device=device),
+            vision_noisy_frame_indexes=[],
+        )
+        logits = transformer.lm_head(out["last_hidden_state"][n - 1])  # [vocab]
+        nxt = int(logits.argmax().item())
+        ids.append(nxt)
+        new_tokens.append(nxt)
+        if nxt in eos:
+            break
+    return new_tokens
+
+
+# ===========================================================================
+# VAE encode/decode bridge (normalize / denormalize, matching the framework)
+# ===========================================================================
+@dataclass
+class _VaeNorm:
+    """Cached ``mean`` / ``inv_std`` for VAE (de)normalization."""
+
+    mean: torch.Tensor  # [z_dim]
+    inv_std: torch.Tensor  # [z_dim]
+
+    @classmethod
+    def from_vae(cls, vae: Any, dtype: torch.dtype) -> _VaeNorm:
+        mean = torch.tensor(list(vae.config.latents_mean), dtype=dtype)
+        std = torch.tensor(list(vae.config.latents_std), dtype=dtype)
+        return cls(mean=mean, inv_std=1.0 / std)
+
+
+def cosmos3_vae_encode(vae: Any, video: torch.Tensor, norm: _VaeNorm) -> torch.Tensor:
+    """Encode ``[B, 3, T, H, W]`` pixels in [-1, 1] to NORMALIZED latents.
+
+    Matches the framework ``DiffusersWan22VAE.encode``: take the posterior mode
+    and apply ``(mu - mean) * inv_std``. FastVideo's ``AutoencoderKLWan.encode``
+    returns a ``DiagonalGaussianDistribution``; we read ``.mode()``.
+    """
+    in_dtype = video.dtype
+    device = video.device
+    mean = norm.mean.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    inv_std = norm.inv_std.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    raw_mu = vae.encode(video).mode()
+    return ((raw_mu - mean) * inv_std).to(in_dtype)
+
+
+def cosmos3_vae_decode(vae: Any, latents: torch.Tensor, norm: _VaeNorm) -> torch.Tensor:
+    """Decode NORMALIZED latents ``[B, z, T, H, W]`` to pixels ``[B, 3, T, H, W]``.
+
+    Inverts the normalization (``z / inv_std + mean``) then calls
+    ``vae.decode`` (which already clamps to [-1, 1]).
+    """
+    in_dtype = latents.dtype
+    device = latents.device
+    mean = norm.mean.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    inv_std = norm.inv_std.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1)
+    z_raw = latents / inv_std + mean
+    out = vae.decode(z_raw)
+    if isinstance(out, tuple):
+        out = out[0]
+    if hasattr(out, "sample"):
+        out = out.sample
+    return out.to(in_dtype)
+
+
+# ===========================================================================
+# Per-vision-item packing geometry
+# ===========================================================================
+@dataclass
+class Cosmos3VisionSpec:
+    """Geometry + conditioning for one vision item in a denoise run.
+
+    Args:
+        condition_frame_indexes: Latent-frame indices kept clean.
+        shape: ``(C, T, H, W)`` of the latent for this item.
+    """
+
+    shape: tuple[int, int, int, int]
+    condition_frame_indexes: list[int]
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+# ===========================================================================
+# Pure denoise/CFG math (parity oracle target)
+# ===========================================================================
+def _split_flat_latent(flat: torch.Tensor, specs: list[Any]) -> list[torch.Tensor]:
+    """Split a flat vector into per-item tensors via each spec's ``numel``/``shape``.
+
+    Shared by vision (``[C, T, H, W]``), sound (``[C, T]``), and action
+    (``[T, D]``) specs — every spec exposes ``numel`` and ``shape``.
+    """
+    out: list[torch.Tensor] = []
+    offset = 0
+    for spec in specs:
+        out.append(flat[offset:offset + spec.numel].reshape(spec.shape))
+        offset += spec.numel
+    return out
+
+
+@dataclass
+class Cosmos3SoundSpec:
+    """Geometry + conditioning for one sound item in a denoise run.
+
+    Args:
+        shape: ``(C, T)`` of the sound latent (channels, temporal frames).
+        condition_frame_indexes: Latent-frame indices kept clean (``[]`` for t2vs).
+        fps: Sound latent FPS (``sound_latent_fps``); used iff fps modulation is on.
+    """
+
+    shape: tuple[int, int]
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+@dataclass
+class Cosmos3ActionSpec:
+    """Geometry + conditioning for one action item in a denoise run.
+
+    Args:
+        shape: ``(T, action_dim)`` of the action latent.
+        condition_frame_indexes: Frame indices kept clean (conditioning actions).
+        domain_id: Embodiment domain id for the domain-aware action projection.
+        fps: Action FPS; used iff fps modulation is on.
+    """
+
+    shape: tuple[int, int]
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    domain_id: int = 0
+    fps: float | None = None
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+def cosmos3_get_cfg_velocity(
+    *,
+    transformer: Any,
+    flat_latent: torch.Tensor,
+    timestep: torch.Tensor,
+    guidance: float,
+    specs: list[Cosmos3VisionSpec],
+    cond_token_ids: list[int],
+    uncond_token_ids: list[int],
+    special_tokens: dict[str, int],
+    latent_patch_size: int,
+    temporal_modality_margin: int,
+    reset_spatial_ids: bool,
+    enable_fps_modulation: bool,
+    base_fps: float,
+    temporal_compression_factor: int,
+    include_end_of_generation_token: bool = False,
+    fps_per_item: list[float] | None = None,
+    normalize_cfg: bool = False,
+    sound_specs: list[Cosmos3SoundSpec] | None = None,
+    sound_fps_per_item: list[float] | None = None,
+    action_specs: list[Cosmos3ActionSpec] | None = None,
+    action_fps_per_item: list[float] | None = None,
+) -> torch.Tensor:
+    """Sequential-CFG velocity for one denoise step (framework math).
+
+    Replicates the framework ``get_cfg_velocity``:
+
+    1. split ``flat_latent`` into per-vision-item ``[C, T, H, W]`` latents,
+    2. run a conditional pass (prompt tokens) and an unconditional pass
+       (negative-prompt tokens); each repacks via
+       :func:`pack_cosmos3_video_sequence`, forwards the DiT to obtain
+       ``preds_vision`` (a list of ``[1, C, T, H, W]`` unpatchified noisy-frame
+       predictions), and zeros the prediction on conditioning frames
+       (``pred * (1 - condition_mask)``),
+    3. combine ``v = uncond + guidance * (cond - uncond)`` (optionally
+       norm-rescaled), returned flattened to match ``flat_latent``.
+
+    ``timestep`` is a scalar tensor (raw scheduler timestep); ``timestep_scale``
+    is applied inside the DiT, so it is passed through unscaled here.
+    """
+    assert timestep.numel() == 1, "timestep must be a scalar"
+    timestep_value = float(timestep.reshape(()).item())
+
+    # Combined flat layout: [all vision | all action | all sound], matching the
+    # framework per-sample concat order ([vision_i | action_i | sound_i]); single
+    # sample here.
+    vision_total = sum(spec.numel for spec in specs)
+    action_total = sum(spec.numel for spec in action_specs) if action_specs else 0
+    noise_x_vision = _split_flat_latent(flat_latent[:vision_total], specs)
+    noise_x_action = (_split_flat_latent(flat_latent[vision_total:vision_total +
+                                                     action_total], action_specs) if action_specs else None)
+    noise_x_sound = (_split_flat_latent(flat_latent[vision_total +
+                                                    action_total:], sound_specs) if sound_specs else None)
+    device = next(transformer.parameters()).device
+
+    def _run(token_ids: list[int]) -> torch.Tensor:
+        sound_items: list[Cosmos3SoundItem] = []
+        if sound_specs is not None and noise_x_sound is not None:
+            sound_items = [
+                Cosmos3SoundItem(
+                    latent=noise_x_sound[i],
+                    condition_frame_indexes=list(ss.condition_frame_indexes),
+                    fps=(sound_fps_per_item[i] if sound_fps_per_item is not None else None),
+                ) for i, ss in enumerate(sound_specs)
+            ]
+        action_items: list[Cosmos3ActionItem] = []
+        if action_specs is not None and noise_x_action is not None:
+            action_items = [
+                Cosmos3ActionItem(
+                    latent=noise_x_action[i],
+                    condition_frame_indexes=list(asp.condition_frame_indexes),
+                    domain_id=asp.domain_id,
+                    fps=(action_fps_per_item[i] if action_fps_per_item is not None else None),
+                ) for i, asp in enumerate(action_specs)
+            ]
+        samples = [
+            Cosmos3SampleInputs(
+                text_ids=list(token_ids),
+                vision=Cosmos3VisionItem(
+                    latent=latent,
+                    condition_frame_indexes=list(spec.condition_frame_indexes),
+                    fps=(fps_per_item[i] if fps_per_item is not None else None),
+                ),
+                sound=(sound_items[i] if i < len(sound_items) else None),
+                action=(action_items[i] if i < len(action_items) else None),
+                timestep=timestep_value,
+            ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False))
+        ]
+        packed = pack_cosmos3_video_sequence(
+            samples,
+            special_tokens,
+            latent_patch_size=latent_patch_size,
+            include_end_of_generation_token=include_end_of_generation_token,
+            temporal_modality_margin=temporal_modality_margin,
+            reset_spatial_ids=reset_spatial_ids,
+            enable_fps_modulation=enable_fps_modulation,
+            base_fps=base_fps,
+            temporal_compression_factor=temporal_compression_factor,
+        )
+        out = transformer(**packed.to_dit_kwargs(device=device))
+
+        # Vision velocity: zero on conditioning frames, per item, flattened.
+        vision_vel = torch.zeros(vision_total, device=flat_latent.device, dtype=flat_latent.dtype)
+        preds = out.get("preds_vision")
+        if preds is not None:
+            items: list[torch.Tensor] = []
+            for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False):
+                pred = pred.squeeze(0) if pred.dim() == 5 else pred  # [C, T, H, W]
+                keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device)  # [T,1,1]
+                items.append(pred * keep if keep.sum() > 0 else torch.zeros_like(pred))
+            vision_vel = torch.cat([v.reshape(-1) for v in items]).to(flat_latent.dtype)
+
+        parts = [vision_vel]
+
+        if action_specs:
+            # Action velocity: preds_action are per-item [T, D], already zero on
+            # clean frames; zero on cond frames defensively.
+            action_vel = torch.zeros(action_total, device=flat_latent.device, dtype=flat_latent.dtype)
+            preds_a = out.get("preds_action")
+            if preds_a is not None:
+                a_items: list[torch.Tensor] = []
+                for pred, cond_mask in zip(preds_a, packed.action_condition_mask, strict=False):
+                    pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [T, D]
+                    keep = (1.0 - cond_mask).reshape(-1, 1).to(dtype=pred.dtype, device=pred.device)  # [T, 1]
+                    a_items.append(pred * keep)
+                action_vel = torch.cat([v.reshape(-1) for v in a_items]).to(flat_latent.dtype)
+            parts.append(action_vel)
+
+        if sound_specs:
+            # Sound velocity: preds_sound are per-item [C, T], already zero on clean
+            # frames (unpack fills only noisy frames); zero on cond frames defensively.
+            sound_total = sum(spec.numel for spec in sound_specs)
+            sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype)
+            preds_s = out.get("preds_sound")
+            if preds_s is not None:
+                s_items: list[torch.Tensor] = []
+                for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False):
+                    pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [C, T]
+                    keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device)  # [1, T]
+                    s_items.append(pred * keep)
+                sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype)
+            parts.append(sound_vel)
+
+        return vision_vel if len(parts) == 1 else torch.cat(parts)
+
+    cond_v = _run(cond_token_ids)
+    uncond_v = _run(uncond_token_ids)
+    v_pred = uncond_v + guidance * (cond_v - uncond_v)
+    if normalize_cfg:
+        scale = (torch.norm(cond_v) / (torch.norm(v_pred) + 1e-8)).clamp(min=0.0, max=1.0)
+        v_pred = v_pred * scale
+    return v_pred
+
+
+class Cosmos3DenoiseEngine:
+    """Stateless denoise driver tying CFG velocity to UniPC stepping.
+
+    Holds the transformer + scheduler + packing constants and runs the full
+    UniPC denoise loop. Kept separate from the pipeline so it can be exercised
+    in isolation (smoke + parity tests) with stub or real components.
+    """
+
+    def __init__(
+        self,
+        *,
+        transformer: Any,
+        scheduler: Any,
+        special_tokens: dict[str, int],
+        latent_patch_size: int,
+        temporal_modality_margin: int,
+        reset_spatial_ids: bool,
+        enable_fps_modulation: bool,
+        base_fps: float,
+        temporal_compression_factor: int,
+        include_end_of_generation_token: bool = False,
+    ) -> None:
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.special_tokens = special_tokens
+        self.latent_patch_size = latent_patch_size
+        self.temporal_modality_margin = temporal_modality_margin
+        self.reset_spatial_ids = reset_spatial_ids
+        self.enable_fps_modulation = enable_fps_modulation
+        self.base_fps = base_fps
+        self.temporal_compression_factor = temporal_compression_factor
+        self.include_end_of_generation_token = include_end_of_generation_token
+
+    def velocity(
+        self,
+        *,
+        flat_latent: torch.Tensor,
+        timestep: torch.Tensor,
+        guidance: float,
+        specs: list[Cosmos3VisionSpec],
+        cond_token_ids: list[int],
+        uncond_token_ids: list[int],
+        fps_per_item: list[float] | None = None,
+        sound_specs: list[Cosmos3SoundSpec] | None = None,
+        sound_fps_per_item: list[float] | None = None,
+        action_specs: list[Cosmos3ActionSpec] | None = None,
+        action_fps_per_item: list[float] | None = None,
+    ) -> torch.Tensor:
+        return cosmos3_get_cfg_velocity(
+            transformer=self.transformer,
+            flat_latent=flat_latent,
+            timestep=timestep,
+            guidance=guidance,
+            specs=specs,
+            cond_token_ids=cond_token_ids,
+            uncond_token_ids=uncond_token_ids,
+            special_tokens=self.special_tokens,
+            latent_patch_size=self.latent_patch_size,
+            temporal_modality_margin=self.temporal_modality_margin,
+            reset_spatial_ids=self.reset_spatial_ids,
+            enable_fps_modulation=self.enable_fps_modulation,
+            base_fps=self.base_fps,
+            temporal_compression_factor=self.temporal_compression_factor,
+            include_end_of_generation_token=self.include_end_of_generation_token,
+            fps_per_item=fps_per_item,
+            sound_specs=sound_specs,
+            sound_fps_per_item=sound_fps_per_item,
+            action_specs=action_specs,
+            action_fps_per_item=action_fps_per_item,
+        )
+
+    def denoise(
+        self,
+        *,
+        flat_latent: torch.Tensor,
+        timesteps: torch.Tensor,
+        guidance: float,
+        specs: list[Cosmos3VisionSpec],
+        cond_token_ids: list[int],
+        uncond_token_ids: list[int],
+        fps_per_item: list[float] | None = None,
+        progress_bar: Any | None = None,
+        sound_specs: list[Cosmos3SoundSpec] | None = None,
+        sound_fps_per_item: list[float] | None = None,
+        action_specs: list[Cosmos3ActionSpec] | None = None,
+        action_fps_per_item: list[float] | None = None,
+    ) -> torch.Tensor:
+        """Run the full UniPC denoise loop, returning the final flat latent.
+
+        For each timestep: compute the sequential-CFG velocity, then
+        ``scheduler.step(model_output=v, timestep, sample=latent.unsqueeze(0))``
+        (the framework steps with a leading batch axis), squeezing back to flat.
+        For t2vs the flat latent is ``[vision | sound]`` and the velocity covers
+        both; the scheduler steps the combined vector jointly.
+        """
+        latent = flat_latent
+        iterator = progress_bar(timesteps) if progress_bar is not None else timesteps
+        for t in iterator:
+            v_pred = self.velocity(
+                flat_latent=latent,
+                timestep=t.reshape(1),
+                guidance=guidance,
+                specs=specs,
+                cond_token_ids=cond_token_ids,
+                uncond_token_ids=uncond_token_ids,
+                fps_per_item=fps_per_item,
+                sound_specs=sound_specs,
+                sound_fps_per_item=sound_fps_per_item,
+                action_specs=action_specs,
+                action_fps_per_item=action_fps_per_item,
+            )
+            stepped = self.scheduler.step(
+                model_output=v_pred,
+                timestep=t,
+                sample=latent.unsqueeze(0),
+                return_dict=False,
+            )[0]
+            latent = stepped.squeeze(0)
+        return latent
+
+
+# ===========================================================================
+# Pipeline (ComposedPipelineBase)
+# ===========================================================================
+class Cosmos3OmniDiffusersPipeline(ComposedPipelineBase):
+    """Cosmos3 video generation pipeline (T2V / I2V / T2I).
+
+    Stage-based ``ComposedPipelineBase`` pipeline. The required modules
+    (``transformer`` / ``vae`` / ``scheduler`` / ``text_tokenizer``) are loaded
+    from the ``nvidia/Cosmos3-Nano`` checkpoint by the component loader. The
+    class name matches the checkpoint ``model_index.json`` ``_class_name`` so
+    the registry resolves it directly.
+
+    The denoise/CFG/VAE math is delegated to module-level helpers
+    (:func:`cosmos3_get_cfg_velocity`, :class:`Cosmos3DenoiseEngine`,
+    :func:`cosmos3_vae_encode` / :func:`cosmos3_vae_decode`) which are
+    framework-parity tested in ``tests/local_tests/cosmos3``.
+    """
+
+    is_video_pipeline = True
+    # ``vision_encoder`` / ``sound_tokenizer`` ship in the checkpoint but the
+    # video path does not need them; they are intentionally omitted here.
+    _required_config_modules = ["text_tokenizer", "vae", "transformer", "scheduler"]
+
+    # Engine-init flow_shift (T2V/I2V); T2I overrides to 3.0 per request.
+    _engine_init_flow_shift: float = 1.0
+    # Class-attribute defaults so ``__new__``-based unit tests can read these
+    # before ``initialize_pipeline`` runs.
+    scheduler: Any = None
+    _base_scheduler_config: Any = None
+    _current_flow_shift: float | None = None
+
+    @staticmethod
+    def _flow_scheduler_config(config: Any) -> dict[str, Any]:
+        """Coerce a loaded UniPC config to the framework's flow-matching setup.
+
+        The checkpoint ``scheduler_config.json`` carries diffusers-style fields
+        (``use_karras_sigmas=True``, ``sigma_min``/``sigma_max``, beta schedule)
+        that do not describe the framework sampler. The framework uses
+        ``FlowUniPCMultistepScheduler`` (pure flow matching: ``shift`` +
+        ``num_train_timesteps`` only). FastVideo's vendored UniPC checks
+        ``use_karras_sigmas`` *before* ``use_flow_sigmas``, so leaving karras on
+        builds diffusion-style sigmas and the denoise diverges to NaN. Force the
+        flow config here (parity-verified in ``test_cosmos3_scheduler_parity``).
+        """
+        cfg = dict(config)
+        cfg.update(
+            use_karras_sigmas=False,
+            use_exponential_sigmas=False,
+            use_beta_sigmas=False,
+            use_flow_sigmas=True,
+            prediction_type="flow_prediction",
+            predict_x0=True,
+            final_sigmas_type="zero",
+        )
+        return cfg
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        """Bind the loaded scheduler + snapshot its config so per-request
+        flow_shift rebuilds are cheap and the engine-init shift is applied."""
+        pipeline_config = fastvideo_args.pipeline_config
+        engine_shift = getattr(pipeline_config, "flow_shift", None)
+        if engine_shift is not None:
+            self._engine_init_flow_shift = float(engine_shift)
+        scheduler = self.get_module("scheduler")
+        if scheduler is not None:
+            # Rebuild from a flow-coerced config so the runtime scheduler matches
+            # the framework sampler (the loaded checkpoint config is diffusers-style).
+            flow_config = self._flow_scheduler_config(scheduler.config)
+            self.scheduler = UniPCMultistepScheduler.from_config(flow_config)
+            if isinstance(self.modules, dict):
+                self.modules["scheduler"] = self.scheduler
+            self._base_scheduler_config = self.scheduler.config
+            self._current_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0))
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Wire the Cosmos3 stages.
+
+        The whole text->latent->denoise->decode flow is custom (sequential CFG
+        with per-pass repacking), so a single :class:`Cosmos3DenoisingStage`
+        owns it. ``InputValidationStage`` runs first for the standard checks.
+        """
+        from fastvideo.pipelines.stages import InputValidationStage
+        from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage
+
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=Cosmos3DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+                tokenizer=self.get_module("text_tokenizer"),
+                pipeline=self,
+            ),
+        )
+
+    # -- Scheduler control --------------------------------------------------
+
+    def _set_flow_shift(self, target_shift: float) -> None:
+        """Set UniPC ``flow_shift`` to ``target_shift``.
+
+        Lazily builds a default UniPC scheduler when called before
+        ``initialize_pipeline`` (e.g. the ``__new__``-based scheduler-parity
+        tests); otherwise rebuilds from the snapshotted base config only when
+        the target differs from the current shift.
+        """
+        target = float(target_shift)
+        base_config = self._base_scheduler_config
+        if base_config is None:
+            self.scheduler = UniPCMultistepScheduler(
+                num_train_timesteps=1000,
+                solver_order=2,
+                prediction_type="flow_prediction",
+                use_flow_sigmas=True,
+                flow_shift=target,
+            )
+            self._base_scheduler_config = self.scheduler.config
+            self._current_flow_shift = target
+            return
+        current = self._current_flow_shift
+        if current is not None and target == float(current):
+            return
+        self.scheduler = UniPCMultistepScheduler.from_config(base_config, flow_shift=target)
+        if isinstance(self.modules, dict):
+            self.modules["scheduler"] = self.scheduler
+        self._current_flow_shift = target
+
+    # -- Tokenization -------------------------------------------------------
+
+    def tokenize_caption(self, caption: str, *, is_video: bool = False, use_system_prompt: bool = False) -> list[int]:
+        return cosmos3_tokenize_caption(self.get_module("text_tokenizer"),
+                                        caption,
+                                        is_video=is_video,
+                                        use_system_prompt=use_system_prompt)
+
+
+# Entry point for the pipeline registry. The class name matches the checkpoint
+# ``model_index.json`` ``_class_name`` so ``resolve_pipeline_cls`` finds it.
+EntryClass = Cosmos3OmniDiffusersPipeline

fastvideo/pipelines/basic/cosmos3/presets.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 (Cosmos3-Nano) inference presets.
+
+Defaults track the official ``cosmos-framework`` ``sample_args`` for the video
+paths (``text2video`` / ``image2video``: guidance=6.0, num_steps=35, shift=10.0,
+fps=24, num_frames=189) and ``text2image`` (guidance=4.0, num_steps=50,
+shift=3.0). The default resolution is 16:9 at a VAE-aligned 704x1280 (spatial
+compression 16 -> 44x80 latent grid).
+"""
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Cosmos3 sequential-CFG UniPC denoising pass",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+    }),
+)
+
+# Framework video negative prompt (Cosmos quality prompt).
+COSMOS3_VIDEO_NEGATIVE_PROMPT = (
+    "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+    "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+    "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, "
+    "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, "
+    "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
+    "Overall, the video is of poor quality.")
+
+COSMOS3_NANO = InferencePreset(
+    name="cosmos3_nano",
+    version=1,
+    model_family="cosmos3",
+    description="Cosmos3-Nano text-to-video",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 704,
+        "width": 1280,
+        "num_frames": 189,
+        "fps": 24,
+        "guidance_scale": 6.0,
+        "num_inference_steps": 35,
+        "negative_prompt": COSMOS3_VIDEO_NEGATIVE_PROMPT,
+    },
+)
+
+COSMOS3_NANO_I2V = InferencePreset(
+    name="cosmos3_nano_i2v",
+    version=1,
+    model_family="cosmos3",
+    description="Cosmos3-Nano image-to-video",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 704,
+        "width": 1280,
+        "num_frames": 189,
+        "fps": 24,
+        "guidance_scale": 6.0,
+        "num_inference_steps": 35,
+        "negative_prompt": COSMOS3_VIDEO_NEGATIVE_PROMPT,
+    },
+)
+
+COSMOS3_NANO_T2I = InferencePreset(
+    name="cosmos3_nano_t2i",
+    version=1,
+    model_family="cosmos3",
+    description="Cosmos3-Nano text-to-image",
+    workload_type="t2i",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 24,
+        "guidance_scale": 4.0,
+        "num_inference_steps": 50,
+        "negative_prompt": "",
+    },
+)
+
+ALL_PRESETS = (COSMOS3_NANO, COSMOS3_NANO_I2V, COSMOS3_NANO_T2I)

fastvideo/pipelines/basic/cosmos3/sequence_packing.py

@@ -0,0 +1,549 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastVideo-native Cosmos3 sequence packing (video subset).
+
+Numerical-parity port of the official ``cosmos_framework`` data packer
+(``cosmos_framework.data.vfm.sequence_packing.pack_input_sequence``) restricted
+to the VIDEO generation path that the FastVideo Cosmos3 DiT consumes (T2V / I2V
+/ T2I). It builds, per sample, two splits:
+
+* a ``causal`` text split (prompt token ids, plus the trailing ``eos`` and
+  ``start_of_generation`` markers the framework appends when a generation
+  modality follows), and
+* a ``full`` vision split (VAE latent patch tokens).
+
+The 3D-MRoPE position ids ``[3, seq]`` are produced exactly like the framework:
+text tokens broadcast a single monotone id across the (t, h, w) axes, the
+temporal offset is bumped by ``temporal_modality_margin`` at the text->vision
+boundary, and vision tokens lay out a (T, H, W) grid with spatial ids reset per
+segment. Condition frames (I2V cond frame 0, T2I single conditioned frame, ...)
+are kept in the packed sequence and rope grid but excluded from the MSE-loss /
+timestep bookkeeping, mirroring the framework.
+
+The output ``Cosmos3PackedSequence`` maps 1:1 onto the
+``Cosmos3VFMTransformer.forward`` kwargs via :meth:`to_dit_kwargs`. This module
+is pure torch/python; it imports no diffusers/transformers model classes.
+
+Reference of record: ``cosmos_framework`` (NVIDIA), the parity oracle used by
+``tests/local_tests/cosmos3/test_cosmos3_packing_parity.py``.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from fastvideo.models.dits.cosmos3 import (
+    compute_mrope_position_ids_text,
+    compute_mrope_position_ids_vision,
+)
+
+__all__ = [
+    "Cosmos3VisionItem",
+    "Cosmos3SampleInputs",
+    "Cosmos3PackedSequence",
+    "pack_cosmos3_video_sequence",
+]
+
+
+# ---------------------------------------------------------------------------
+# Inputs
+# ---------------------------------------------------------------------------
+@dataclass
+class Cosmos3VisionItem:
+    """One vision latent for a sample.
+
+    Args:
+        latent: VAE latent ``[C, T, H, W]`` (a leading batch axis of size 1 is
+            accepted and squeezed).
+        condition_frame_indexes: Latent-frame indices that are *conditioned*
+            (clean) rather than noisy. ``[]`` for T2V, ``[0]`` for I2V, and the
+            single conditioned frame for T2I.
+        fps: Frames-per-second for this clip; only used when
+            ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3SoundItem:
+    """One sound latent for a sample (t2vs).
+
+    Args:
+        latent: AVAE sound latent ``[C, T]`` (channels, temporal frames).
+        condition_frame_indexes: Latent-frame indices that are *conditioned*
+            (clean). ``[]`` for t2vs (all frames generated).
+        fps: Sound latent FPS (``sound_latent_fps``, e.g. 25); only used when
+            ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3ActionItem:
+    """One action latent for a sample (action-conditioned world model).
+
+    Args:
+        latent: Action latent ``[T, action_dim]`` (per-frame action vectors).
+        condition_frame_indexes: Frame indices kept clean (conditioning actions).
+        domain_id: Embodiment domain id (scalar / ``[1]``) for the
+            domain-aware action projection.
+        fps: Action FPS; only used when ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    domain_id: int = 0
+    fps: float | None = None
+
+
+@dataclass
+class Cosmos3SampleInputs:
+    """Per-sample packing inputs (text prompt + vision item, +sound, +action)."""
+
+    text_ids: list[int]
+    vision: Cosmos3VisionItem
+    timestep: float
+    sound: Cosmos3SoundItem | None = None
+    action: Cosmos3ActionItem | None = None
+
+
+# ---------------------------------------------------------------------------
+# Output
+# ---------------------------------------------------------------------------
+@dataclass
+class Cosmos3PackedSequence:
+    """Packed-sequence inputs consumed by ``Cosmos3VFMTransformer.forward``.
+
+    Field names mirror the framework ``PackedSequence`` (+ its ``vision``
+    ``ModalityData``) so the parity test can compare field-by-field.
+    """
+
+    # Sequence structure.
+    sample_lens: list[int]
+    split_lens: list[int]
+    attn_modes: list[str]
+    sequence_length: int
+    is_image_batch: bool
+
+    # Text modality.
+    text_ids: torch.Tensor
+    text_indexes: torch.Tensor
+    position_ids: torch.Tensor  # [3, sequence_length]
+
+    # Vision modality.
+    vision_tokens: list[torch.Tensor]
+    vision_token_shapes: list[tuple[int, int, int]]
+    vision_sequence_indexes: torch.Tensor
+    vision_timesteps: torch.Tensor
+    vision_mse_loss_indexes: torch.Tensor
+    vision_noisy_frame_indexes: list[torch.Tensor]
+    vision_condition_mask: list[torch.Tensor]
+    fps_vision: torch.Tensor | None = None
+
+    # Sound modality (t2vs); empty/None when no sound.
+    sound_tokens: list[torch.Tensor] = field(default_factory=list)
+    sound_token_shapes: list[tuple[int, int, int]] = field(default_factory=list)
+    sound_sequence_indexes: torch.Tensor | None = None
+    sound_timesteps: torch.Tensor | None = None
+    sound_mse_loss_indexes: torch.Tensor | None = None
+    sound_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+    sound_condition_mask: list[torch.Tensor] = field(default_factory=list)
+    fps_sound: torch.Tensor | None = None
+
+    # Action modality (action-conditioned world model); empty/None when no action.
+    action_tokens: list[torch.Tensor] = field(default_factory=list)
+    action_token_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    action_sequence_indexes: torch.Tensor | None = None
+    action_timesteps: torch.Tensor | None = None
+    action_mse_loss_indexes: torch.Tensor | None = None
+    action_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+    action_condition_mask: list[torch.Tensor] = field(default_factory=list)
+    action_domain_id: list[torch.Tensor] = field(default_factory=list)
+
+    def to_dit_kwargs(self, device: torch.device | str | None = None) -> dict[str, Any]:
+        """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``.
+
+        Packing is device-agnostic (ids/indexes/position-ids are built on CPU).
+        When ``device`` is given, every tensor input is moved to it so the DiT
+        forward runs on a single device (e.g. the model's GPU at inference).
+        """
+
+        def _mv(x: Any) -> Any:
+            return x.to(device) if (device is not None and torch.is_tensor(x)) else x
+
+        return dict(
+            text_ids=_mv(self.text_ids),
+            text_indexes=_mv(self.text_indexes),
+            position_ids=_mv(self.position_ids),
+            sequence_length=int(self.sequence_length),
+            split_lens=list(self.split_lens),
+            attn_modes=list(self.attn_modes),
+            vision_tokens=[_mv(t) for t in self.vision_tokens],
+            vision_token_shapes=list(self.vision_token_shapes),
+            vision_sequence_indexes=_mv(self.vision_sequence_indexes),
+            vision_timesteps=_mv(self.vision_timesteps),
+            vision_mse_loss_indexes=_mv(self.vision_mse_loss_indexes),
+            vision_noisy_frame_indexes=[_mv(t) for t in self.vision_noisy_frame_indexes],
+            fps_vision=self.fps_vision,
+            sound_tokens=[_mv(t) for t in self.sound_tokens],
+            sound_token_shapes=list(self.sound_token_shapes),
+            sound_sequence_indexes=_mv(self.sound_sequence_indexes),
+            sound_timesteps=_mv(self.sound_timesteps),
+            sound_mse_loss_indexes=_mv(self.sound_mse_loss_indexes),
+            sound_noisy_frame_indexes=[_mv(t) for t in self.sound_noisy_frame_indexes],
+            fps_sound=_mv(self.fps_sound),
+            action_tokens=[_mv(t) for t in self.action_tokens],
+            action_token_shapes=list(self.action_token_shapes),
+            action_sequence_indexes=_mv(self.action_sequence_indexes),
+            action_timesteps=_mv(self.action_timesteps),
+            action_mse_loss_indexes=_mv(self.action_mse_loss_indexes),
+            action_noisy_frame_indexes=[_mv(t) for t in self.action_noisy_frame_indexes],
+            action_domain_id=[_mv(t) for t in self.action_domain_id],
+        )
+
+
+# ---------------------------------------------------------------------------
+# Packing
+# ---------------------------------------------------------------------------
+def pack_cosmos3_video_sequence(
+    samples: list[Cosmos3SampleInputs],
+    special_tokens: dict[str, int],
+    *,
+    latent_patch_size: int = 2,
+    include_end_of_generation_token: bool = False,
+    temporal_modality_margin: int = 15_000,
+    reset_spatial_ids: bool = True,
+    enable_fps_modulation: bool = False,
+    base_fps: float = 24.0,
+    temporal_compression_factor: int = 4,
+    initial_mrope_temporal_offset: int | float = 0,
+) -> Cosmos3PackedSequence:
+    """Pack prompts + vision latents into the Cosmos3 DiT packed-sequence inputs.
+
+    Video subset of ``cosmos_framework`` ``pack_input_sequence`` under
+    ``unified_3d_mrope``: each sample is ``[causal text, full vision]``.
+
+    Args:
+        samples: Per-sample text prompt token ids + vision item + timestep.
+        special_tokens: Must contain ``eos_token_id`` and
+            ``start_of_generation`` (and ``end_of_generation`` if
+            ``include_end_of_generation_token``). ``bos_token_id`` is honored if
+            present (prepended) to match the framework.
+        latent_patch_size: Latent patch size used by the DiT.
+        include_end_of_generation_token: Append the framework's end-of-generation
+            marker after the vision split.
+        temporal_modality_margin: Temporal-offset bump applied at the
+            text->vision boundary (``unified_3d_mrope_temporal_modality_margin``).
+        reset_spatial_ids: Reset vision spatial ids to 0 per segment.
+        enable_fps_modulation: Use float, fps-scaled temporal positions.
+        base_fps: Base FPS used when ``enable_fps_modulation``.
+        temporal_compression_factor: VAE temporal compression factor.
+        initial_mrope_temporal_offset: Per-sample starting temporal offset.
+
+    Returns:
+        A :class:`Cosmos3PackedSequence`.
+    """
+    assert "eos_token_id" in special_tokens, "special_tokens must contain eos_token_id"
+    assert "start_of_generation" in special_tokens, "special_tokens must contain start_of_generation"
+    if latent_patch_size < 1:
+        raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}")
+
+    # Build-time accumulators (concatenated across samples).
+    sample_lens: list[int] = []
+    split_lens: list[int] = []
+    attn_modes: list[str] = []
+
+    text_ids: list[int] = []
+    text_indexes: list[int] = []
+    position_id_blocks: list[torch.Tensor] = []  # each [3, n]
+
+    vision_tokens: list[torch.Tensor] = []
+    vision_token_shapes: list[tuple[int, int, int]] = []
+    vision_sequence_indexes: list[int] = []
+    vision_timesteps: list[float] = []
+    vision_mse_loss_indexes: list[int] = []
+    vision_noisy_frame_indexes: list[torch.Tensor] = []
+    vision_condition_mask: list[torch.Tensor] = []
+    fps_values: list[float] = []
+
+    sound_tokens: list[torch.Tensor] = []
+    sound_token_shapes: list[tuple[int, int, int]] = []
+    sound_sequence_indexes: list[int] = []
+    sound_timesteps: list[float] = []
+    sound_mse_loss_indexes: list[int] = []
+    sound_noisy_frame_indexes: list[torch.Tensor] = []
+    sound_condition_mask: list[torch.Tensor] = []
+    sound_fps_values: list[float] = []
+
+    action_tokens: list[torch.Tensor] = []
+    action_token_shapes: list[tuple[int, ...]] = []
+    action_sequence_indexes: list[int] = []
+    action_timesteps: list[float] = []
+    action_mse_loss_indexes: list[int] = []
+    action_noisy_frame_indexes: list[torch.Tensor] = []
+    action_condition_mask: list[torch.Tensor] = []
+    action_domain_id: list[torch.Tensor] = []
+
+    curr = 0  # running position in the packed sequence
+    is_image_batch = True
+
+    for sample in samples:
+        temporal_offset: int | float = initial_mrope_temporal_offset
+        sample_len = 0
+
+        # ---- 1. Text split (causal) ----
+        if "bos_token_id" in special_tokens:
+            shifted_text_ids = [special_tokens["bos_token_id"], *sample.text_ids]
+        else:
+            shifted_text_ids = list(sample.text_ids)
+        # The video path always has a following generation modality, so the
+        # framework appends eos + start_of_generation.
+        shifted_text_ids = [*shifted_text_ids, special_tokens["eos_token_id"], special_tokens["start_of_generation"]]
+        text_split_len = len(shifted_text_ids)
+
+        text_ids.extend(shifted_text_ids)
+        text_indexes.extend(range(curr, curr + text_split_len))
+
+        text_mrope, temporal_offset = compute_mrope_position_ids_text(
+            num_tokens=text_split_len,
+            temporal_offset=int(temporal_offset),
+        )
+        position_id_blocks.append(text_mrope)
+
+        attn_modes.append("causal")
+        split_lens.append(text_split_len)
+        curr += text_split_len
+        sample_len += text_split_len
+
+        # End of text modality: bump temporal offset before vision.
+        temporal_offset += temporal_modality_margin
+        # Sound shares the vision temporal start (parallel temporal positions).
+        vision_start_temporal_offset = temporal_offset
+
+        # ---- 2. Vision split (full) ----
+        latent = sample.vision.latent
+        latent = latent.squeeze(0) if latent.dim() == 5 else latent  # [C, T, H, W]
+        _c, latent_t, latent_h, latent_w = latent.shape
+        patch_h = math.ceil(latent_h / latent_patch_size)
+        patch_w = math.ceil(latent_w / latent_patch_size)
+        num_vision_tokens = latent_t * patch_h * patch_w
+
+        vision_tokens.append(sample.vision.latent)
+        vision_token_shapes.append((latent_t, patch_h, patch_w))
+        vision_sequence_indexes.extend(range(curr, curr + num_vision_tokens))
+
+        condition_set = {idx for idx in sample.vision.condition_frame_indexes if 0 <= idx < latent_t}
+        cond_mask = torch.zeros((latent_t, 1, 1), device=latent.device, dtype=latent.dtype)
+        for frame_idx in condition_set:
+            cond_mask[frame_idx, 0, 0] = 1.0
+        vision_condition_mask.append(cond_mask)
+
+        noisy_frames = torch.tensor(
+            [idx for idx in range(latent_t) if idx not in condition_set],
+            device=latent.device,
+            dtype=torch.long,
+        )
+        vision_noisy_frame_indexes.append(noisy_frames)
+
+        # MSE-loss indices + per-token timesteps cover only the noisy frames.
+        frame_token_stride = patch_h * patch_w
+        for frame_idx in range(latent_t):
+            if frame_idx in condition_set:
+                continue
+            frame_start = curr + frame_idx * frame_token_stride
+            vision_mse_loss_indexes.extend(range(frame_start, frame_start + frame_token_stride))
+            vision_timesteps.extend([float(sample.timestep)] * frame_token_stride)
+
+        vision_fps = sample.vision.fps if enable_fps_modulation else None
+        if vision_fps is not None:
+            fps_values.append(float(vision_fps))
+        vision_mrope, temporal_offset = compute_mrope_position_ids_vision(
+            grid_t=latent_t,
+            grid_h=patch_h,
+            grid_w=patch_w,
+            temporal_offset=temporal_offset,
+            fps=vision_fps,
+            base_fps=base_fps,
+            temporal_compression_factor=temporal_compression_factor,
+            enable_fps_modulation=enable_fps_modulation,
+        )
+        position_id_blocks.append(vision_mrope)
+
+        curr += num_vision_tokens
+        sample_len += num_vision_tokens
+
+        # ---- 2a2. Action split: shares the vision "full" split ----
+        # Mirrors framework ``_pack_action_tokens``: action latent [T, D] -> T
+        # tokens (token shape (T,)), domain-aware, 3D-MRoPE at the vision temporal
+        # offset with ``start_frame_offset=1`` (parallel to vision; tcf=1; does
+        # not advance the offset).
+        action_split_len = 0
+        if sample.action is not None:
+            action_latent = sample.action.latent  # [T, D]
+            action_t = int(action_latent.shape[0])
+            action_split_len = action_t
+
+            action_tokens.append(action_latent)
+            action_token_shapes.append((action_t, ))
+            action_sequence_indexes.extend(range(curr, curr + action_t))
+            action_domain_id.append(torch.tensor([int(sample.action.domain_id)], dtype=torch.long))
+
+            a_cond_set = {idx for idx in sample.action.condition_frame_indexes if 0 <= idx < action_t}
+            a_cond_mask = torch.zeros((action_t, 1), device=action_latent.device, dtype=action_latent.dtype)
+            for fi in a_cond_set:
+                a_cond_mask[fi, 0] = 1.0
+            action_condition_mask.append(a_cond_mask)
+
+            a_noisy = torch.tensor([idx for idx in range(action_t) if idx not in a_cond_set],
+                                   device=action_latent.device,
+                                   dtype=torch.long)
+            action_noisy_frame_indexes.append(a_noisy)
+
+            for fi in range(action_t):
+                if fi in a_cond_set:
+                    continue
+                action_mse_loss_indexes.append(curr + fi)
+                action_timesteps.append(float(sample.timestep))
+
+            action_fps = sample.action.fps if enable_fps_modulation else None
+            action_mrope, _ = compute_mrope_position_ids_vision(
+                grid_t=action_t,
+                grid_h=1,
+                grid_w=1,
+                temporal_offset=vision_start_temporal_offset,
+                fps=action_fps,
+                base_fps=base_fps,
+                temporal_compression_factor=1,  # action is at frame rate
+                base_temporal_compression_factor=temporal_compression_factor,
+                enable_fps_modulation=enable_fps_modulation,
+                start_frame_offset=1,
+            )
+            position_id_blocks.append(action_mrope)
+            curr += action_t
+            sample_len += action_t
+
+        # ---- 2b. Sound split (t2vs): shares the vision "full" split ----
+        # Mirrors framework ``_pack_sound_tokens``: sound latent [C, T] -> T
+        # tokens (token shape (T,1,1)), packed right after vision, with 3D-MRoPE
+        # temporal positions starting at the vision temporal offset (parallel to
+        # vision, start_frame_offset=0, tcf=1) and NOT advancing it.
+        sound_split_len = 0
+        if sample.sound is not None:
+            sound_latent = sample.sound.latent
+            sound_latent = sound_latent.squeeze(0) if sound_latent.dim() == 3 else sound_latent  # [C, T]
+            _sc, sound_t = sound_latent.shape
+            sound_split_len = sound_t
+
+            sound_tokens.append(sound_latent)
+            sound_token_shapes.append((sound_t, 1, 1))
+            sound_sequence_indexes.extend(range(curr, curr + sound_t))
+
+            s_cond_set = {idx for idx in sample.sound.condition_frame_indexes if 0 <= idx < sound_t}
+            s_cond_mask = torch.zeros((sound_t, 1), device=sound_latent.device, dtype=sound_latent.dtype)
+            for fi in s_cond_set:
+                s_cond_mask[fi, 0] = 1.0
+            sound_condition_mask.append(s_cond_mask)
+
+            s_noisy = torch.tensor([idx for idx in range(sound_t) if idx not in s_cond_set],
+                                   device=sound_latent.device,
+                                   dtype=torch.long)
+            sound_noisy_frame_indexes.append(s_noisy)
+
+            for fi in range(sound_t):
+                if fi in s_cond_set:
+                    continue
+                sound_mse_loss_indexes.append(curr + fi)  # 1 token per sound frame
+                sound_timesteps.append(float(sample.timestep))
+
+            sound_fps = sample.sound.fps if enable_fps_modulation else None
+            if sound_fps is not None:
+                sound_fps_values.append(float(sound_fps))
+            sound_mrope, _ = compute_mrope_position_ids_vision(
+                grid_t=sound_t,
+                grid_h=1,
+                grid_w=1,
+                temporal_offset=vision_start_temporal_offset,
+                fps=sound_fps,
+                base_fps=base_fps,
+                temporal_compression_factor=1,  # sound latent already at sound_latent_fps
+                enable_fps_modulation=enable_fps_modulation,
+                start_frame_offset=0,
+            )
+            position_id_blocks.append(sound_mrope)
+            curr += sound_t
+            sample_len += sound_t
+
+        # ---- 3. Optional end-of-generation marker ----
+        eov_len = 0
+        if include_end_of_generation_token:
+            assert "end_of_generation" in special_tokens, ("special_tokens must contain end_of_generation when "
+                                                           "include_end_of_generation_token=True")
+            text_ids.append(special_tokens["end_of_generation"])
+            text_indexes.append(curr)
+            eov_dtype = torch.float32 if enable_fps_modulation else torch.long
+            eov_ids = torch.full((3, 1), temporal_offset, dtype=eov_dtype)
+            position_id_blocks.append(eov_ids)
+            temporal_offset += 1
+            curr += 1
+            eov_len = 1
+            sample_len += 1
+
+        # Vision + action + sound + any trailing eov marker share one "full" split.
+        attn_modes.append("full")
+        split_lens.append(num_vision_tokens + action_split_len + sound_split_len + eov_len)
+        sample_lens.append(sample_len)
+
+        if latent_t != 1:
+            is_image_batch = False
+
+    sequence_length = sum(sample_lens)
+
+    # position_ids: float iff any block is float (fps modulation path).
+    any_float = any(b.dtype.is_floating_point for b in position_id_blocks)
+    if any_float:
+        position_id_blocks = [b.to(torch.float32) for b in position_id_blocks]
+    position_ids = torch.cat(position_id_blocks, dim=1)  # [3, sequence_length]
+
+    timesteps_dtype = torch.float32
+    return Cosmos3PackedSequence(
+        sample_lens=sample_lens,
+        split_lens=split_lens,
+        attn_modes=attn_modes,
+        sequence_length=sequence_length,
+        is_image_batch=is_image_batch,
+        text_ids=torch.tensor(text_ids, dtype=torch.long),
+        text_indexes=torch.tensor(text_indexes, dtype=torch.long),
+        position_ids=position_ids,
+        vision_tokens=vision_tokens,
+        vision_token_shapes=vision_token_shapes,
+        vision_sequence_indexes=torch.tensor(vision_sequence_indexes, dtype=torch.long),
+        vision_timesteps=torch.tensor(vision_timesteps, dtype=timesteps_dtype),
+        vision_mse_loss_indexes=torch.tensor(vision_mse_loss_indexes, dtype=torch.long),
+        vision_noisy_frame_indexes=vision_noisy_frame_indexes,
+        vision_condition_mask=vision_condition_mask,
+        fps_vision=(torch.tensor(fps_values, dtype=torch.float32) if fps_values else None),
+        sound_tokens=sound_tokens,
+        sound_token_shapes=sound_token_shapes,
+        sound_sequence_indexes=(torch.tensor(sound_sequence_indexes, dtype=torch.long) if sound_tokens else None),
+        sound_timesteps=(torch.tensor(sound_timesteps, dtype=timesteps_dtype) if sound_tokens else None),
+        sound_mse_loss_indexes=(torch.tensor(sound_mse_loss_indexes, dtype=torch.long) if sound_tokens else None),
+        sound_noisy_frame_indexes=sound_noisy_frame_indexes,
+        sound_condition_mask=sound_condition_mask,
+        fps_sound=(torch.tensor(sound_fps_values, dtype=torch.float32) if sound_fps_values else None),
+        action_tokens=action_tokens,
+        action_token_shapes=action_token_shapes,
+        action_sequence_indexes=(torch.tensor(action_sequence_indexes, dtype=torch.long) if action_tokens else None),
+        action_timesteps=(torch.tensor(action_timesteps, dtype=timesteps_dtype) if action_tokens else None),
+        action_mse_loss_indexes=(torch.tensor(action_mse_loss_indexes, dtype=torch.long) if action_tokens else None),
+        action_noisy_frame_indexes=action_noisy_frame_indexes,
+        action_condition_mask=action_condition_mask,
+        action_domain_id=action_domain_id,
+    )

fastvideo/pipelines/basic/flux/__init__.py

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

fastvideo/pipelines/basic/flux/flux_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages.flux_stages import (
+    FluxConditioningStage,
+    FluxDecodingStage,
+    FluxDenoisingStage,
+    FluxInputValidationStage,
+    FluxLatentPreparationStage,
+    FluxTimestepPreparationStage,
+)
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+
+class FluxPipeline(ComposedPipelineBase):
+    """FLUX.1-dev T2I (Diffusers module layout, packed latents, embedded guidance)."""
+
+    _required_config_modules = [
+        "scheduler",
+        "transformer",
+        "vae",
+        "text_encoder",
+        "text_encoder_2",
+        "tokenizer",
+        "tokenizer_2",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=FluxInputValidationStage())
+
+        self.add_stage(
+            stage_name="text_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=[
+                    self.get_module("text_encoder"),
+                    self.get_module("text_encoder_2"),
+                ],
+                tokenizers=[
+                    self.get_module("tokenizer"),
+                    self.get_module("tokenizer_2"),
+                ],
+            ),
+        )
+
+        self.add_stage(stage_name="flux_conditioning_stage", stage=FluxConditioningStage())
+
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=FluxTimestepPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=FluxLatentPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=FluxDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=FluxDecodingStage(vae=self.get_module("vae")),
+        )
+
+
+EntryClass = FluxPipeline

fastvideo/pipelines/basic/glm_image/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-Image pipeline package."""
+
+from fastvideo.pipelines.basic.glm_image.glm_image_pipeline import (
+    GlmImagePipeline, )
+
+__all__ = ["GlmImagePipeline"]

fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+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.glm_image.stages import (
+    GlmImageBeforeDenoisingStage,
+    GlmImageConditionEncodingStage,
+    GlmImageDecodingStage,
+    GlmImageDenoisingStage,
+)
+from fastvideo.pipelines.stages import InputValidationStage
+
+logger = init_logger(__name__)
+
+
+class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
+    pipeline_name = "GlmImagePipeline"
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+        "vision_language_encoder",
+        "processor",
+    ]
+
+    _optional_config_modules: list[str] = []
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=1.0)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage(),
+        )
+
+        self.add_stage(
+            stage_name="glm_image_before_denoising_stage",
+            stage=GlmImageBeforeDenoisingStage(
+                vae=self.get_module("vae"),
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+                processor=self.get_module("processor"),
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vision_language_encoder=self.get_module("vision_language_encoder"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="glm_image_condition_encoding_stage",
+            stage=GlmImageConditionEncodingStage(
+                vae=self.get_module("vae"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=GlmImageDenoisingStage(
+                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=GlmImageDecodingStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+
+EntryClass = GlmImagePipeline

fastvideo/pipelines/basic/glm_image/stages/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-Image pipeline stages."""
+
+from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (GlmImageBeforeDenoisingStage)
+from fastvideo.pipelines.basic.glm_image.stages.condition_encoding import (GlmImageConditionEncodingStage)
+from fastvideo.pipelines.basic.glm_image.stages.decoding import (GlmImageDecodingStage)
+from fastvideo.pipelines.basic.glm_image.stages.denoising import (GlmImageDenoisingStage)
+
+__all__ = [
+    "GlmImageBeforeDenoisingStage",
+    "GlmImageConditionEncodingStage",
+    "GlmImageDecodingStage",
+    "GlmImageDenoisingStage",
+]

fastvideo/pipelines/basic/glm_image/stages/before_denoising.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import re
+from math import sqrt
+
+import numpy as np
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+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 StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+logger = init_logger(__name__)
+
+
+def calculate_shift(
+    image_seq_len: int,
+    base_seq_len: int = 256,
+    base_shift: float = 0.25,
+    max_shift: float = 0.75,
+) -> float:
+    return (image_seq_len / base_seq_len)**0.5 * max_shift + base_shift
+
+
+def get_glyph_texts(prompt: str | list[str]) -> list[str] | list[list[str]]:
+    if isinstance(prompt, str):
+        prompts: list[str] = [prompt]
+        is_batch = False
+    else:
+        prompts = prompt
+        is_batch = True
+    out: list[list[str]] = []
+    for p in prompts:
+        out.append(
+            re.findall(r"'([^']*)'", p) + re.findall(r"“([^“”]*)”", p) + re.findall(r'"([^"]*)"', p) +
+            re.findall(r"「([^「」]*)」", p))
+    return out if is_batch else out[0]
+
+
+def compute_glyph_embeds(
+    prompts: list[str],
+    tokenizer,
+    text_encoder,
+    device: torch.device,
+    dtype: torch.dtype,
+    max_sequence_length: int = 2048,
+) -> torch.Tensor:
+    all_glyph_texts = get_glyph_texts(prompts)
+    all_glyph_embeds = []
+    for glyph_texts in all_glyph_texts:
+        if len(glyph_texts) == 0:
+            glyph_texts = [""]
+        input_ids = tokenizer(
+            glyph_texts,
+            max_length=max_sequence_length,
+            truncation=True,
+        ).input_ids
+        input_ids = [[tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids for ids in input_ids]
+        max_length = max(len(ids) for ids in input_ids)
+        attention_mask = torch.tensor(
+            [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        input_ids_t = torch.tensor(
+            [ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids],
+            device=device,
+        )
+        outputs = text_encoder(input_ids_t, attention_mask=attention_mask)
+        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
+        all_glyph_embeds.append(glyph_embeds)
+
+    max_seq_len = max(emb.size(1) for emb in all_glyph_embeds)
+    padded = []
+    for emb in all_glyph_embeds:
+        if emb.size(1) < max_seq_len:
+            pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype)
+            emb = torch.cat([pad, emb], dim=1)
+        padded.append(emb)
+    return torch.cat(padded, dim=0).to(device=device, dtype=dtype)
+
+
+def _grid_dims(height: int, width: int) -> tuple[int, int, int, int]:
+    th, tw = height // 32, width // 32
+    ratio = th / tw
+    pth = int(sqrt(ratio) * 16)
+    ptw = int(sqrt(1 / ratio) * 16)
+    return th, tw, pth, ptw
+
+
+def _upsample_d32_to_d16(tokens: torch.Tensor, th: int, tw: int) -> torch.Tensor:
+    tokens = tokens.view(1, 1, th, tw).float()
+    tokens = torch.nn.functional.interpolate(tokens, scale_factor=2, mode="nearest").long()
+    return tokens.view(1, -1)
+
+
+class GlmImageBeforeDenoisingStage(PipelineStage):
+
+    def __init__(self,
+                 vae,
+                 text_encoder,
+                 tokenizer,
+                 processor,
+                 transformer,
+                 scheduler,
+                 vision_language_encoder=None) -> None:
+        super().__init__()
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.processor = processor
+        self.transformer = transformer
+        self.scheduler = scheduler
+        if isinstance(vision_language_encoder, tuple):
+            self.vision_language_encoder, self.vl_processor = (vision_language_encoder[0], vision_language_encoder[1]
+                                                               or processor)
+        else:
+            self.vision_language_encoder = vision_language_encoder
+            self.vl_processor = processor
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        th, tw, pth, ptw = _grid_dims(batch.height, batch.width)
+
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+
+        # 1-3. AR token generation. I2I prepends the condition image and uses a
+        # single-scale target grid; T2I is multi-scale.
+        is_t2i = batch.pil_image is None
+        if self.vision_language_encoder is not None:
+            content = [{"type": "text", "text": batch.prompt}]
+            if not is_t2i:
+                content.insert(0, {"type": "image", "image": batch.pil_image})
+            messages = [{"role": "user", "content": content}]
+            inputs = self.vl_processor.apply_chat_template(messages,
+                                                           tokenize=True,
+                                                           target_h=batch.height,
+                                                           target_w=batch.width,
+                                                           return_dict=True,
+                                                           return_tensors="pt").to(device)
+
+            if is_t2i:
+                up_h, up_w = th, tw
+                large_start, large_count = pth * ptw, th * tw
+                max_new = large_count + (pth * ptw) + 1
+            else:
+                # Condition grid(s) first, target grid last.
+                _, t_h, t_w = inputs["image_grid_thw"][-1].tolist()
+                up_h, up_w = int(t_h), int(t_w)
+                large_start, large_count = 0, up_h * up_w
+                max_new = large_count + 1
+
+            outputs = self.vision_language_encoder.generate(**inputs, max_new_tokens=max_new, do_sample=True)
+            gen_tokens = outputs[0][inputs.input_ids.shape[-1]:]
+            if gen_tokens.shape[0] >= large_start + large_count:
+                large_tokens = gen_tokens[large_start:large_start + large_count]
+            else:
+                available = gen_tokens[large_start:]
+                large_tokens = torch.zeros(large_count, dtype=gen_tokens.dtype, device=gen_tokens.device)
+                if available.shape[0] > 0:
+                    large_tokens[:min(available.shape[0], large_count)] = available[:large_count]
+                logger.warning("AR generated %d tokens, expected %d. Padding with zeros.", gen_tokens.shape[0],
+                               large_start + large_count)
+            batch.prior_token_id = _upsample_d32_to_d16(large_tokens, up_h, up_w)
+            batch.prior_token_drop = torch.zeros(batch.prior_token_id.shape, dtype=torch.bool, device=device)
+
+            if not is_t2i:
+                self._compute_source_prior_tokens(batch, inputs)
+        else:
+            num_prior_tokens = 4 * th * tw
+            logger.warning("No vision_language_encoder provided; using random dropped priors.")
+            batch.prior_token_id = torch.randint(0, 16384, (1, num_prior_tokens), device=device)
+            batch.prior_token_drop = torch.ones(batch.prior_token_id.shape, dtype=torch.bool, device=device)
+
+        # 4. Glyph T5 encoding.
+        prompts = [batch.prompt] if isinstance(batch.prompt, str) else list(batch.prompt)
+        prompt_embeds = compute_glyph_embeds(prompts, self.tokenizer, self.text_encoder, device, dtype)
+
+        # 5. CFG-side negative encoding.
+        if batch.do_classifier_free_guidance:
+            neg_prompts = [batch.negative_prompt or ""] * len(prompts)
+            neg_embeds = compute_glyph_embeds(neg_prompts, self.tokenizer, self.text_encoder, device, dtype)
+            L_pos, L_neg = prompt_embeds.shape[1], neg_embeds.shape[1]
+            max_L = max(L_pos, L_neg)
+            if L_pos < max_L:
+                pad = torch.zeros(prompt_embeds.shape[0],
+                                  max_L - L_pos,
+                                  prompt_embeds.shape[2],
+                                  device=device,
+                                  dtype=dtype)
+                prompt_embeds = torch.cat([pad, prompt_embeds], dim=1)
+            if L_neg < max_L:
+                pad = torch.zeros(neg_embeds.shape[0], max_L - L_neg, neg_embeds.shape[2], device=device, dtype=dtype)
+                neg_embeds = torch.cat([pad, neg_embeds], dim=1)
+            # Row 0 conditional (positive), row 1 unconditional (negative).
+            prompt_embeds = torch.cat([prompt_embeds, neg_embeds], dim=0)
+            att_pos = torch.ones((1, max_L), device=device)
+            att_neg = torch.ones((1, max_L), device=device)
+            if L_pos < max_L:
+                att_pos[:, :max_L - L_pos] = 0
+            if L_neg < max_L:
+                att_neg[:, :max_L - L_neg] = 0
+            attention_mask = torch.cat([att_pos, att_neg], dim=0)
+        else:
+            attention_mask = torch.ones((1, prompt_embeds.shape[1]), device=device)
+
+        batch.prompt_embeds = [prompt_embeds]
+        batch.attention_mask = attention_mask
+
+        # 6. Latents + dynamic flow shift.
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+        batch.latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8), device=device, dtype=dtype)
+
+        # Integer-cast linspace timesteps with resolution-dependent shift applied to
+        # sigmas only; the DiT is conditioned on the unshifted integer timesteps.
+        ntt = self.scheduler.config.num_train_timesteps
+        patch_size = self.transformer.patch_size
+        image_seq_len = ((batch.height // 8) * (batch.width // 8)) // (patch_size**2)
+        sched_timesteps = np.linspace(ntt, 1.0, batch.num_inference_steps + 1)[:-1].astype(np.int64).astype(np.float32)
+        sched_sigmas = sched_timesteps / ntt
+        self.scheduler.set_shift(calculate_shift(image_seq_len))
+        self.scheduler.set_timesteps(batch.num_inference_steps,
+                                     device=device,
+                                     sigmas=sched_sigmas.tolist(),
+                                     timesteps=sched_timesteps.tolist())
+        batch.timesteps = self.scheduler.timesteps
+        return batch
+
+    @torch.no_grad()
+    def _compute_source_prior_tokens(self, batch: ForwardBatch, inputs) -> None:
+        image_grid_thw = inputs["image_grid_thw"]
+        num_condition_images = image_grid_thw.shape[0] - 1
+        source_grids = image_grid_thw[:num_condition_images]
+        image_features = self.vision_language_encoder.get_image_features(inputs["pixel_values"], source_grids)
+        image_feature_parts = getattr(image_features, "pooler_output", image_features)
+        embed = torch.cat(image_feature_parts, dim=0)
+        src_ids_d32 = self.vision_language_encoder.get_image_tokens(embed, source_grids)
+        split_sizes = source_grids.prod(dim=-1).tolist()
+        upsampled = [
+            _upsample_d32_to_d16(ids, int(grid[1]), int(grid[2])).squeeze(0)
+            for ids, grid in zip(torch.split(src_ids_d32, split_sizes), source_grids, strict=False)
+        ]
+        src_grids_up = source_grids.clone()
+        src_grids_up[:, 1] *= 2
+        src_grids_up[:, 2] *= 2
+        batch.extra["glm_prior_token_image_ids"] = torch.cat(upsampled, dim=0)
+        batch.extra["glm_source_image_grid_thw"] = src_grids_up
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_not_empty)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prior_token_id", batch.prior_token_id, V.is_tensor)
+        return result

fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.image_processor import ImageProcessor
+from fastvideo.models.dits.glm_image import GlmImageKVCache
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+
+_CONDITION_MULTIPLE_OF = 16  # vae_scale_factor (8) * DiT patch_size (2)
+
+
+class GlmImageConditionEncodingStage(PipelineStage):
+
+    def __init__(self, vae, transformer) -> None:
+        super().__init__()
+        self.vae = vae
+        self.transformer = transformer
+        self.image_processor = ImageProcessor(vae_scale_factor=_CONDITION_MULTIPLE_OF)
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.pil_image is None:
+            return batch
+
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        self.vae.to(device)
+
+        prior_ids = batch.extra["glm_prior_token_image_ids"].to(device)
+        if prior_ids.dim() == 1:
+            prior_ids = prior_ids.unsqueeze(0)
+
+        # Latent patch count must match the source prior tokens; mismatch is fatal.
+        src_grid = batch.extra["glm_source_image_grid_thw"][0]
+        cond_h = int(src_grid[1]) * _CONDITION_MULTIPLE_OF
+        cond_w = int(src_grid[2]) * _CONDITION_MULTIPLE_OF
+        cond_img = self.image_processor.preprocess(batch.pil_image, cond_h, cond_w).to(device=device,
+                                                                                       dtype=torch.float32)
+        latent = self.vae.encode(cond_img).latent_dist.mode()
+        # NOTE: at runtime self.vae.config is a diffusers FrozenDict with flat
+        # latents_mean/latents_std fields. Access the flat fields directly.
+        cfg = self.vae.config
+        mean = torch.tensor(cfg.latents_mean, device=device, dtype=torch.float32).view(1, -1, 1, 1)
+        std = torch.tensor(cfg.latents_std, device=device, dtype=torch.float32).view(1, -1, 1, 1)
+        latent = ((latent - mean) / std).to(dtype)
+
+        kv_caches = GlmImageKVCache(num_layers=self.transformer.num_layers)
+        empty_text = batch.prompt_embeds[0][:1, :0, :].to(device=device, dtype=dtype)
+        with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=batch):
+            self.transformer(
+                hidden_states=latent,
+                encoder_hidden_states=empty_text,
+                prior_token_id=prior_ids,
+                prior_token_drop=torch.zeros((prior_ids.shape[0], ), dtype=torch.bool, device=device),
+                timestep=torch.zeros((1, ), device=device),
+                target_size=torch.tensor([tuple(cond_img.shape[-2:])], device=device, dtype=torch.long),
+                crop_coords=torch.zeros((1, 2), device=device, dtype=torch.long),
+                kv_caches=kv_caches,
+                kv_caches_mode="write",
+            )
+        batch.extra["glm_kv_caches"] = kv_caches
+        return batch

fastvideo/pipelines/basic/glm_image/stages/decoding.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.stages.decoding import DecodingStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+
+class GlmImageDecodingStage(DecodingStage):
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast = (vae_dtype != torch.float32 and not fastvideo_args.disable_autocast)
+
+        latents = self._denormalize_latents(latents)
+        if latents.dim() == 5:
+            latents = latents.squeeze(2)
+
+        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            if not vae_autocast:
+                latents = latents.to(vae_dtype)
+            decoded = self.vae.decode(latents)
+
+        image = decoded.sample if hasattr(decoded, "sample") else decoded
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image.unsqueeze(2)

fastvideo/pipelines/basic/glm_image/stages/denoising.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CFG convention (both denoise paths): row 0 conditional (positive), row 1 unconditional."""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.attention.backends.sdpa import SDPAMetadata
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.platforms import AttentionBackendEnum
+
+
+class GlmImageDenoisingStage(DenoisingStage):
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
+        latents = getattr(batch, "latent", getattr(batch, "latents", None))
+        result.add_check("latents", latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        return result
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+        guidance_scale = batch.guidance_scale
+        do_cfg = guidance_scale > 1.0
+
+        latents = getattr(batch, "latent", getattr(batch, "latents", None))
+        if latents is None:
+            raise ValueError("No latents found in batch.")
+        if latents.dim() == 5:
+            latents = latents.squeeze(2)
+
+        prompt_embeds = batch.prompt_embeds[0]
+        text_attention_mask = getattr(batch, "attention_mask", None)
+        timesteps = batch.timesteps
+
+        patch_size = self.transformer.patch_size
+        _, _, h, w = latents.shape
+        image_seq_length = (h // patch_size) * (w // patch_size)
+        text_seq_length = prompt_embeds.shape[1] if prompt_embeds.dim() >= 2 else 0
+
+        first_block = self.transformer.transformer_blocks[0]
+        backend = getattr(first_block.attn1.attn, "backend", None)
+        sdpa = backend == AttentionBackendEnum.TORCH_SDPA and text_attention_mask is not None
+
+        kv_caches = batch.extra.get("glm_kv_caches")
+        if kv_caches is None:
+            self._denoise_t2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                              text_seq_length, image_seq_length, sdpa, device, dtype)
+        else:
+            self._denoise_i2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                              text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype)
+        return batch
+
+    def _denoise_t2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                     text_seq_length, image_seq_length, sdpa, device, dtype) -> None:
+        num_inference_steps = batch.num_inference_steps
+        bs = 2 if do_cfg else 1
+        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long).repeat(bs, 1)
+        crop_coords = torch.zeros((bs, 2), device=device, dtype=torch.long)
+
+        prior_token_id = batch.prior_token_id
+        if do_cfg and prior_token_id.shape[0] == 1:
+            prior_token_id = prior_token_id.repeat(2, 1)
+        if do_cfg:
+            prior_token_drop = torch.tensor([False, True], device=device)
+        else:
+            prior_token_drop = getattr(batch, "prior_token_drop", torch.tensor([False], device=device))
+
+        attention_mask_kv = None
+        if sdpa:
+            if (text_attention_mask.shape[0] == 1 and bs > 1):
+                text_attention_mask = text_attention_mask.repeat(bs, 1)
+            mix_attn_mask = torch.ones((bs, text_seq_length + image_seq_length), device=device, dtype=torch.float32)
+            mix_attn_mask[:, :text_seq_length] = (text_attention_mask.float().to(device))
+            attention_mask_kv = (mix_attn_mask > 0).unsqueeze(1).unsqueeze(2)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents] * 2) if do_cfg else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t).to(dtype)
+                t_expand = t.expand(latent_model_input.shape[0]) - 1
+
+                attn_metadata = (SDPAMetadata(current_timestep=i, attn_mask=attention_mask_kv)
+                                 if attention_mask_kv is not None else None)
+
+                with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                          attn_metadata=attn_metadata,
+                                                          forward_batch=batch):
+                    noise_pred = self.transformer(
+                        latent_model_input,
+                        prompt_embeds,
+                        prior_token_id,
+                        prior_token_drop,
+                        t_expand,
+                        target_size,
+                        crop_coords,
+                    )
+
+                    if do_cfg:
+                        noise_pred_cond, noise_pred_uncond = noise_pred.float().chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                        guidance_rescale = getattr(batch, "guidance_rescale", 0.0)
+                        if guidance_rescale > 0.0:
+                            dims = list(range(1, noise_pred_cond.ndim))
+                            std_text = noise_pred_cond.std(dim=dims, keepdim=True)
+                            std_cfg = noise_pred.std(dim=dims, keepdim=True)
+                            rescaled = noise_pred * (std_text / std_cfg)
+                            noise_pred = (guidance_rescale * rescaled + (1 - guidance_rescale) * noise_pred)
+
+                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        batch.latents = latents.unsqueeze(2)
+
+    def _denoise_i2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
+                     text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype) -> None:
+        """Two separate transformer calls (cond reads the cache, uncond skips it):
+        the cache mode is one global flag with batch-1 k/v, so a 2-row CFG call
+        cannot express both."""
+        num_inference_steps = batch.num_inference_steps
+        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long)
+        crop_coords = torch.zeros((1, 2), device=device, dtype=torch.long)
+        prior_token_id = batch.prior_token_id[:1]
+        drop_keep = torch.zeros((1, ), dtype=torch.bool, device=device)
+        drop_all = torch.ones((1, ), dtype=torch.bool, device=device)
+        cache_len = kv_caches[0].k_cache.shape[1] if kv_caches[0].k_cache is not None else 0
+
+        def _mask(row: int, with_cache: bool):
+            if not sdpa:
+                return None
+            prefix = cache_len if with_cache else 0
+            m = torch.ones((1, prefix + text_seq_length + image_seq_length), device=device, dtype=torch.float32)
+            m[:, prefix:prefix + text_seq_length] = text_attention_mask[row:row + 1].float().to(device)
+            return (m > 0).unsqueeze(1).unsqueeze(2)
+
+        cond_mask = _mask(0, with_cache=True)
+        uncond_mask = _mask(1, with_cache=False) if do_cfg else None
+        pos_embeds = prompt_embeds[:1]
+        neg_embeds = prompt_embeds[1:2] if do_cfg else None
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = self.scheduler.scale_model_input(latents, t).to(dtype)
+                t_expand = t.expand(1) - 1
+
+                cond_meta = (SDPAMetadata(current_timestep=i, attn_mask=cond_mask) if cond_mask is not None else None)
+                with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                          attn_metadata=cond_meta,
+                                                          forward_batch=batch):
+                    noise_pred = self.transformer(latent_model_input,
+                                                  pos_embeds,
+                                                  prior_token_id,
+                                                  drop_keep,
+                                                  t_expand,
+                                                  target_size,
+                                                  crop_coords,
+                                                  kv_caches=kv_caches,
+                                                  kv_caches_mode="read")
+
+                if do_cfg:
+                    uncond_meta = (SDPAMetadata(current_timestep=i, attn_mask=uncond_mask)
+                                   if uncond_mask is not None else None)
+                    with torch.no_grad(), set_forward_context(current_timestep=i,
+                                                              attn_metadata=uncond_meta,
+                                                              forward_batch=batch):
+                        noise_pred_uncond = self.transformer(latent_model_input,
+                                                             neg_embeds,
+                                                             prior_token_id,
+                                                             drop_all,
+                                                             t_expand,
+                                                             target_size,
+                                                             crop_coords,
+                                                             kv_caches=kv_caches,
+                                                             kv_caches_mode="skip")
+                    noise_pred = noise_pred_uncond.float() + guidance_scale * (noise_pred.float() -
+                                                                               noise_pred_uncond.float())
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        kv_caches.clear()
+        batch.latents = latents.unsqueeze(2)

fastvideo/pipelines/pipeline_batch_info.py

@@ -109,6 +109,10 @@ class ForwardBatch:
     max_sequence_length: int | None = None
     prompt_template: dict[str, Any] | None = None
     do_classifier_free_guidance: bool = False
+    # When True, ``guidance_scale`` is passed into models that use embedded guidance (e.g. FLUX)
+    # and must not imply classic dual-forward CFG. Use ``true_cfg_scale > 1`` for true CFG.
+    use_embedded_guidance: bool = False
+    true_cfg_scale: float = 1.0
 
     # Batch info
     batch_size: int | None = None
@@ -252,9 +256,12 @@ class ForwardBatch:
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 
-        # Enable CFG for standard guidance_scale and LTX-2 text CFG scales.
+        # LTX-2 text CFG scales; FLUX uses ``use_embedded_guidance`` so ``guidance_scale > 1`` alone
+        # does not enable classifier-free guidance.
         ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0)
-        if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
+        if self.use_embedded_guidance:
+            self.do_classifier_free_guidance = (self.true_cfg_scale > 1.0) or ltx2_text_cfg_enabled
+        elif self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
             self.do_classifier_free_guidance = True
         if self.negative_prompt_embeds is None:
             self.negative_prompt_embeds = []

fastvideo/pipelines/stages/cosmos3_stages.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 video denoising stage.
+
+The Cosmos3 video path is monolithic by design: each CFG pass repacks the whole
+text+vision sequence (the conditional pass carries prompt tokens, the
+unconditional pass carries negative-prompt tokens), so the standard
+encode/condition/denoise/decode stage split does not apply. This single stage
+owns the full flow, delegating the framework-parity-tested math to
+``fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline``:
+
+  1. resolve mode (T2I / I2V / T2V) + per-mode defaults, set ``flow_shift``;
+  2. tokenize the prompt + negative prompt with the Qwen2 chat template;
+  3. VAE-encode the conditioning frame(s) for I2V / T2I (kept clean), build the
+     initial noise (clean condition frames + pure noise elsewhere);
+  4. run the UniPC denoise loop with sequential CFG
+     (``Cosmos3DenoiseEngine.denoise``);
+  5. VAE-decode + ``(1 + x) / 2`` clamp to [0, 1].
+
+This mirrors the framework ``Cosmos3OmniDiffusersPipeline.__call__``.
+"""
+from __future__ import annotations
+
+import os
+import weakref
+from typing import Any
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+from tqdm import tqdm
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+    Cosmos3DenoiseEngine,
+    Cosmos3SoundSpec,
+    Cosmos3VisionSpec,
+    _VaeNorm,
+    cosmos3_special_tokens,
+    cosmos3_tokenize_caption,
+    cosmos3_vae_decode,
+    cosmos3_vae_encode,
+)
+from fastvideo.pipelines.basic.cosmos3.presets import (
+    COSMOS3_VIDEO_NEGATIVE_PROMPT, )
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+
+logger = init_logger(__name__)
+
+
+class Cosmos3DenoisingStage(PipelineStage):
+    """Full Cosmos3 video denoise: tokenize + encode + denoise + decode."""
+
+    def __init__(self, *, transformer, scheduler, vae, tokenizer, pipeline=None) -> None:
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.vae = vae
+        self.tokenizer = tokenizer
+        self.pipeline = weakref.ref(pipeline) if pipeline is not None else None
+
+    # ------------------------------------------------------------------
+    # Geometry helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _latent_frames(num_frames: int, temporal_factor: int) -> int:
+        return (int(num_frames) - 1) // int(temporal_factor) + 1
+
+    @staticmethod
+    def _flow_shift_for_resolution(height: int, width: int) -> float:
+        """UniPC ``flow_shift`` for a given pixel resolution.
+
+        Mirrors the framework's ``_RESOLUTION_SHIFT_DEFAULTS`` (8B VLM backbone,
+        which Cosmos3-Nano uses): the shift is keyed by the named resolution
+        bucket the (H, W) belongs to, regardless of task (T2V/I2V/T2I):
+
+            "256" -> 3.0,  "480" -> 5.0,  "704"/"720"/"768" -> 10.0
+
+        We invert the framework's ``{IMAGE,VIDEO}_RES_SIZE_INFO`` tables by the
+        longest side: <=320 is the 256 bucket, 640-832 the 480 bucket, and
+        960-1360 the 704/720/768 buckets.
+        """
+        long_side = max(int(height), int(width))
+        if long_side <= 480:
+            return 3.0
+        if long_side <= 896:
+            return 5.0
+        return 10.0
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        pipeline_config = fastvideo_args.pipeline_config
+        arch = pipeline_config.dit_config.arch_config
+        device = self.transformer.embed_tokens.weight.device
+        dtype = self.transformer.embed_tokens.weight.dtype
+
+        num_frames = int(batch.num_frames) if batch.num_frames is not None else 1
+        height = int(batch.height)
+        width = int(batch.width)
+        fps = float(batch.fps) if batch.fps is not None else float(arch.base_fps)
+        guidance = float(batch.guidance_scale)
+
+        is_t2i = num_frames == 1 and batch.preprocessed_image is None and batch.pil_image is None
+        is_i2v = (batch.preprocessed_image is not None or batch.pil_image is not None) and not is_t2i
+
+        # Resolution-based flow_shift, set on the owning pipeline (rebuilds the
+        # scheduler). The framework picks the UniPC shift purely from the named
+        # resolution bucket (``_RESOLUTION_SHIFT_DEFAULTS``), NOT from the task,
+        # so T2V/I2V/T2I at the same resolution share a shift.
+        pipe = self.pipeline() if self.pipeline is not None else None
+        flow_shift = self._flow_shift_for_resolution(height, width)
+        if pipe is not None and hasattr(pipe, "_set_flow_shift"):
+            pipe._set_flow_shift(flow_shift)
+            scheduler = pipe.scheduler
+        else:
+            scheduler = self.scheduler
+
+        # ---- Tokenize prompt + negative prompt ----
+        prompt = batch.prompt if isinstance(batch.prompt, str) else (batch.prompt[0] if batch.prompt else "")
+        negative_prompt = batch.negative_prompt
+        if negative_prompt is None:
+            negative_prompt = "" if is_t2i else COSMOS3_VIDEO_NEGATIVE_PROMPT
+        if isinstance(negative_prompt, list):
+            negative_prompt = negative_prompt[0] if negative_prompt else ""
+
+        special_tokens = cosmos3_special_tokens(self.tokenizer)
+        is_video = not is_t2i
+        cond_ids = cosmos3_tokenize_caption(self.tokenizer, prompt, is_video=is_video, use_system_prompt=False)
+        uncond_ids = cosmos3_tokenize_caption(self.tokenizer,
+                                              negative_prompt,
+                                              is_video=is_video,
+                                              use_system_prompt=False)
+
+        # ---- VAE normalization constants + geometry ----
+        norm = _VaeNorm.from_vae(self.vae, dtype)
+        temporal_factor = int(arch.temporal_compression_factor)
+        spatial_factor = int(self.vae.config.scale_factor_spatial)
+        latent_t = self._latent_frames(num_frames, temporal_factor)
+        latent_h = height // spatial_factor
+        latent_w = width // spatial_factor
+        latent_channel = int(arch.latent_channel)
+        latent_shape = (latent_channel, latent_t, latent_h, latent_w)
+
+        generator = batch.generator
+        if isinstance(generator, list):
+            generator = generator[0] if generator else None
+
+        # ---- Conditioning latent (I2V / T2I) + condition mask ----
+        condition_frame_indexes: list[int] = []
+        clean_latent: torch.Tensor | None = None
+        if is_i2v or (is_t2i and (batch.preprocessed_image is not None or batch.pil_image is not None)):
+            image = batch.preprocessed_image if batch.preprocessed_image is not None else batch.pil_image
+            cond_pixels = self._image_to_video_tensor(image, num_frames, height, width, device, dtype)
+            clean_latent = cosmos3_vae_encode(self.vae, cond_pixels, norm).squeeze(0).float()  # [C, T, H, W]
+            condition_frame_indexes = [0]
+
+        # ---- Initial noise (clean condition frames + pure noise elsewhere) ----
+        pure_noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype).float()
+        if clean_latent is not None:
+            cond_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=pure_noise.dtype)
+            for idx in condition_frame_indexes:
+                if 0 <= idx < latent_t:
+                    cond_mask[idx, 0, 0] = 1.0
+            clean = clean_latent.to(device=device, dtype=pure_noise.dtype)
+            init_latent = cond_mask * clean + (1.0 - cond_mask) * pure_noise
+        else:
+            init_latent = pure_noise
+
+        spec = Cosmos3VisionSpec(
+            shape=latent_shape,
+            condition_frame_indexes=condition_frame_indexes,
+        )
+
+        # ---- Scheduler timesteps ----
+        scheduler.set_timesteps(int(batch.num_inference_steps), device=device)
+        timesteps = scheduler.timesteps
+
+        engine = Cosmos3DenoiseEngine(
+            transformer=self.transformer,
+            scheduler=scheduler,
+            special_tokens=special_tokens,
+            latent_patch_size=int(arch.latent_patch_size),
+            temporal_modality_margin=int(arch.temporal_modality_margin),
+            reset_spatial_ids=bool(arch.unified_3d_mrope_reset_spatial_ids),
+            enable_fps_modulation=bool(arch.enable_fps_modulation),
+            base_fps=float(arch.base_fps),
+            temporal_compression_factor=temporal_factor,
+            include_end_of_generation_token=False,
+        )
+
+        flat_latent = init_latent.reshape(-1)
+        fps_per_item = [fps] if bool(arch.enable_fps_modulation) else None
+
+        # ---- t2vs: jointly generate sound (combined [vision | sound] latent) ----
+        # Mirrors the framework: a placeholder audio sized to the video duration
+        # sets the sound latent length; sound shares the denoise/CFG with vision.
+        with_audio = is_video and os.environ.get("COSMOS3_T2VS", "") not in ("", "0")
+        sound_specs = None
+        sound_fps_per_item = None
+        sound_vae = None
+        sound_shape: tuple[int, int] | None = None
+        if with_audio:
+            sound_vae = self._get_sound_vae(pipe, device, dtype)
+            sound_dim = int(arch.sound_dim)
+            sound_latent_fps = float(arch.sound_latent_fps)
+            # Framework ``create_placeholder_audio`` + ``get_latent_num_samples``.
+            num_audio_samples = int(num_frames / fps * sound_vae.sample_rate)
+            sound_latent_t = max(1, sound_vae.get_latent_num_samples(num_audio_samples))
+            sound_shape = (sound_dim, sound_latent_t)
+            sound_noise = randn_tensor((sound_dim, sound_latent_t), generator=generator, device=device,
+                                       dtype=dtype).float()
+            flat_latent = torch.cat([flat_latent, sound_noise.reshape(-1)])
+            sound_specs = [Cosmos3SoundSpec(shape=sound_shape, condition_frame_indexes=[], fps=sound_latent_fps)]
+            sound_fps_per_item = [sound_latent_fps] if bool(arch.enable_fps_modulation) else None
+
+        final_flat = engine.denoise(
+            flat_latent=flat_latent,
+            timesteps=timesteps,
+            guidance=guidance,
+            specs=[spec],
+            cond_token_ids=cond_ids,
+            uncond_token_ids=uncond_ids,
+            fps_per_item=fps_per_item,
+            progress_bar=lambda it: tqdm(it, desc="Cosmos3 denoising"),
+            sound_specs=sound_specs,
+            sound_fps_per_item=sound_fps_per_item,
+        )
+
+        # ---- Decode vision: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ----
+        vision_flat = final_flat[:spec.numel]
+        result_latent = vision_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype)
+        decoded = cosmos3_vae_decode(self.vae, result_latent, norm)  # [B, 3, T, H, W] in [-1, 1]
+        video = ((1.0 + decoded) / 2.0).clamp(0.0, 1.0)
+
+        batch.latents = result_latent
+        batch.output = video
+
+        # ---- Decode sound: AVAE latent [C, T] -> waveform [C, N] in [-1, 1] ----
+        if with_audio and sound_vae is not None and sound_shape is not None:
+            sound_latent = final_flat[spec.numel:].reshape(sound_shape).unsqueeze(0).to(device=device, dtype=dtype)
+            waveform = sound_vae.decode(sound_latent)  # [1, C_audio, N]
+            batch.extra["audio"] = waveform[0].detach().float().cpu()  # [C_audio, N]
+            batch.extra["audio_sample_rate"] = int(sound_vae.sample_rate)
+        return batch
+
+    # ------------------------------------------------------------------
+    # Image preprocessing
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _resize_and_center_crop(img: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
+        """Aspect-ratio-preserving resize + center crop, matching the framework
+        (``cosmos_framework.inference.vision._resize_and_center_crop``)."""
+        import math
+
+        import torchvision.transforms.functional as TF
+        orig_h, orig_w = img.shape[-2], img.shape[-1]
+        scaling_ratio = max(target_w / orig_w, target_h / orig_h)
+        resize_h = int(math.ceil(scaling_ratio * orig_h))
+        resize_w = int(math.ceil(scaling_ratio * orig_w))
+        img = TF.resize(img, [resize_h, resize_w])
+        return TF.center_crop(img, [target_h, target_w])
+
+    @staticmethod
+    def _get_sound_vae(pipe: Any, device: torch.device, dtype: torch.dtype) -> Any:
+        """Lazily load + cache the Cosmos3 sound AVAE decoder from the checkpoint.
+
+        The video path does not load ``sound_tokenizer``; t2vs needs only its
+        decoder, so we load it on first use from ``<model_path>/sound_tokenizer``.
+        """
+        cached = getattr(pipe, "_sound_vae", None) if pipe is not None else None
+        if cached is not None:
+            return cached
+        from fastvideo.models.audio.cosmos3_avae import Cosmos3SoundVAE
+        model_path = pipe.model_path
+        sound_dir = os.path.join(model_path, "sound_tokenizer")
+        sound_vae = Cosmos3SoundVAE.from_pretrained(sound_dir, torch_dtype=dtype).to(device)
+        if pipe is not None:
+            pipe._sound_vae = sound_vae
+        return sound_vae
+
+    @classmethod
+    def _image_to_video_tensor(
+        cls,
+        image: Any,
+        num_frames: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Build the I2V conditioning pixel video ``[1, 3, T, H, W]`` in [-1, 1].
+
+        Faithful to the framework (``cosmos_framework.inference.vision``):
+        ``load_conditioning_image`` (aspect-preserving resize + center crop +
+        uint8 quantization, then ``/127.5 - 1``) followed by
+        ``build_conditioned_video_batch``, which fills frame 0 with the image and
+        **repeats the last conditioning frame** for the rest of the clip (a static
+        video), NOT zeros. The whole clip is VAE-encoded by the caller; only the
+        latent condition frame(s) are kept clean by the condition mask, but the
+        VAE is temporal, so the repeated (not zeroed) frames change the condition
+        latent — zero-filling here produces a wrong conditioning latent.
+        """
+        import numpy as np
+
+        if hasattr(image, "convert"):  # PIL.Image: framework-exact preprocessing.
+            arr = np.array(image.convert("RGB"))
+            img = torch.from_numpy(arr).permute(2, 0, 1).float()  # [3, H, W] in [0, 255]
+            # Resize + center crop + uint8 quantization, then -> [-1, 1]
+            # (load_conditioning_image / load_conditioning_image_pixels).
+            img = cls._resize_and_center_crop(img.unsqueeze(0), height, width).squeeze(0)
+            img = img.round().clamp(0, 255) / 127.5 - 1.0  # [3, H, W] in [-1, 1]
+        elif isinstance(image, torch.Tensor):  # already-preprocessed conditioning frame.
+            img = image.float()
+            if img.dim() == 5:  # [B,3,T,H,W]
+                img = img[0]
+            if img.dim() == 4:  # [3,T,H,W] or [B,3,H,W] -> first frame
+                img = img[:, 0]
+            if img.max() > 1.5:  # [0, 255] -> [-1, 1]; otherwise assume already [-1, 1].
+                img = img / 127.5 - 1.0
+            if img.shape[-2:] != (height, width):
+                img = cls._resize_and_center_crop(img.unsqueeze(0), height, width).squeeze(0)
+        else:
+            raise TypeError(f"Unsupported conditioning image type: {type(image)}")
+
+        # Static-repeat video (build_conditioned_video_batch: frame 0 = image,
+        # remaining frames repeat the last conditioning frame). The whole clip is
+        # VAE-encoded by the caller; only the latent condition frame(s) are kept
+        # clean by the condition mask, but the VAE is temporal, so the repeated
+        # (not zeroed) frames change the condition latent — zero-filling here
+        # produces a wrong conditioning latent.
+        img = img.to(device=device, dtype=dtype)
+        video = img.unsqueeze(0).unsqueeze(2).expand(1, 3, num_frames, height, width)
+        return video.contiguous()

fastvideo/pipelines/stages/denoising.py

@@ -64,10 +64,9 @@ class DenoisingStage(PipelineStage):
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
             dtype=torch.float16,  # TODO(will): hack
-            supported_attention_backends=(
-                AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN,
-                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE,
-                AttentionBackendEnum.FLASHINFER_SAGE_ATTN3, AttentionBackendEnum.ATTN_QAT_INFER)  # hack
+            supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN,
+                                          AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN,
+                                          AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
         )
 
     def forward(

fastvideo/pipelines/stages/flux_stages.py

@@ -0,0 +1,424 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+def _pack_latents(
+    latents: torch.Tensor,
+    batch_size: int,
+    num_channels_latents: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Diffusers ``_pack_latents`` for FLUX (2×2 spatial pack in latent space)."""
+    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    return latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+
+def _unpack_latents(
+    latents: torch.Tensor,
+    batch_size: int,
+    num_channels_latents: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Inverse of ``_pack_latents``."""
+    latents = latents.reshape(batch_size, height // 2, width // 2, num_channels_latents, 2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+    return latents.reshape(batch_size, num_channels_latents, height, width)
+
+
+def _prepare_latent_image_ids(
+    patch_height: int,
+    patch_width: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.long,
+) -> torch.Tensor:
+    """Match Diffusers ``FluxPipeline._prepare_latent_image_ids`` (no batch dim)."""
+    latent_image_ids = torch.zeros(patch_height, patch_width, 3, device=device, dtype=torch.float32)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(patch_height, device=device)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(patch_width, device=device)[None, :]
+    h, w, c = latent_image_ids.shape
+    return latent_image_ids.reshape(h * w, c).to(dtype=dtype)
+
+
+class FluxInputValidationStage(InputValidationStage):
+    """Require height/width divisible by 16 (VAE scale × 2 for FLUX packing)."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if (batch.height is not None and batch.width is not None and (batch.height % 16 != 0 or batch.width % 16 != 0)):
+            raise ValueError("FLUX expects height and width divisible by 16 "
+                             f"(VAE latent grid × 2× packing); got {batch.height}×{batch.width}.")
+        return super().forward(batch, fastvideo_args)
+
+
+class FluxConditioningStage(PipelineStage):
+    """Build CLIP pooled + T5 sequence + ``text_ids`` (and optional negative for true CFG)."""
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if len(batch.prompt_embeds) < 2:
+            raise ValueError("FluxConditioningStage expects 2 prompt_embeds (CLIP pooled, T5 sequence), "
+                             f"got {len(batch.prompt_embeds)}")
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+
+        pooled = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
+        enc = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
+        seq_len = enc.shape[1]
+        text_ids = torch.zeros(seq_len, 3, device=device, dtype=torch.long)
+
+        batch.extra["flux_pooled_projections"] = pooled
+        batch.extra["flux_encoder_hidden_states"] = enc
+        batch.extra["flux_text_ids"] = text_ids
+
+        if batch.do_classifier_free_guidance:
+            if not batch.negative_prompt_embeds or len(batch.negative_prompt_embeds) < 2:
+                raise ValueError("True CFG requires two negative_prompt_embeds (CLIP, T5).")
+            neg_pooled = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
+            neg_enc = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
+            batch.extra["flux_negative_pooled_projections"] = neg_pooled
+            batch.extra["flux_negative_encoder_hidden_states"] = neg_enc
+
+        return batch
+
+
+class FluxTimestepPreparationStage(TimestepPreparationStage):
+    """Flow Match with resolution-dependent ``mu`` from packed image sequence length."""
+
+    @staticmethod
+    def _calculate_mu(
+        image_seq_len: int,
+        base_seq_len: int = 256,
+        max_seq_len: int = 4096,
+        base_shift: float = 0.5,
+        max_shift: float = 1.15,
+    ) -> float:
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        return float(image_seq_len) * m + b
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        sig = inspect.signature(self.scheduler.set_timesteps)
+        if "mu" not in sig.parameters:
+            logger.warning(
+                "FLUX timestep prep: scheduler %s.set_timesteps does not accept 'mu'; falling back to the base "
+                "timestep schedule. FLUX expects a FlowMatchEulerDiscreteScheduler with resolution-dependent "
+                "dynamic shifting — output quality may degrade.",
+                type(self.scheduler).__name__)
+            return super().forward(batch, fastvideo_args)
+
+        cfg = getattr(self.scheduler, "config", None)
+        use_dynamic = bool(getattr(cfg, "use_dynamic_shifting", False)) if cfg is not None else False
+        if not use_dynamic:
+            logger.warning(
+                "FLUX timestep prep: scheduler has use_dynamic_shifting=False; falling back to the base timestep "
+                "schedule and skipping the resolution-dependent 'mu' shift. FLUX requires dynamic shifting for "
+                "correct timesteps — output quality may degrade.")
+            return super().forward(batch, fastvideo_args)
+
+        if batch.height is None or batch.width is None:
+            raise ValueError("height/width must be set before FluxTimestepPreparationStage")
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
+        h_lat = batch.height // spatial_ratio
+        w_lat = batch.width // spatial_ratio
+        if h_lat % 2 != 0 or w_lat % 2 != 0:
+            raise ValueError(
+                f"Latent spatial dims must be even for FLUX packing; got {h_lat}×{w_lat} from {batch.height}×{batch.width}."
+            )
+        image_seq_len = (h_lat // 2) * (w_lat // 2)
+
+        base_seq_len = int(getattr(cfg, "base_image_seq_len", 256))
+        max_seq_len = int(getattr(cfg, "max_image_seq_len", 4096))
+        base_shift = float(getattr(cfg, "base_shift", 0.5))
+        max_shift = float(getattr(cfg, "max_shift", 1.15))
+
+        device = get_local_torch_device()
+        mu = self._calculate_mu(
+            image_seq_len=image_seq_len,
+            base_seq_len=base_seq_len,
+            max_seq_len=max_seq_len,
+            base_shift=base_shift,
+            max_shift=max_shift,
+        )
+        self.scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
+        batch.timesteps = self.scheduler.timesteps
+        return batch
+
+
+class FluxLatentPreparationStage(PipelineStage):
+
+    def __init__(self, scheduler) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.height is None or batch.width is None:
+            raise ValueError("height/width required for FluxLatentPreparationStage")
+
+        if isinstance(batch.prompt, list):
+            batch_size = len(batch.prompt)
+        elif batch.prompt is not None:
+            batch_size = 1
+        else:
+            if not batch.prompt_embeds:
+                raise ValueError("prompt or prompt_embeds must be provided")
+            batch_size = batch.prompt_embeds[0].shape[0]
+
+        batch_size *= batch.num_videos_per_prompt
+
+        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}")
+
+        arch = fastvideo_args.pipeline_config.dit_config.arch_config
+        in_channels = int(getattr(arch, "in_channels", 64))
+        num_channels_latents = in_channels // 4
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
+
+        h_lat = batch.height // spatial_ratio
+        w_lat = batch.width // spatial_ratio
+
+        dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        device = get_local_torch_device()
+
+        shape = (batch_size, num_channels_latents, h_lat, w_lat)
+        latents = batch.latents
+        if 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:
+            latents = latents.to(device=device, dtype=dtype)
+            if latents.shape != shape:
+                raise ValueError(f"Expected latents shape {shape}, got {tuple(latents.shape)}")
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+
+        packed = _pack_latents(latents, batch_size, num_channels_latents, h_lat, w_lat)
+
+        patch_h, patch_w = h_lat // 2, w_lat // 2
+        img_ids = _prepare_latent_image_ids(patch_h, patch_w, device, dtype=torch.long)
+
+        batch.latents = packed
+        batch.raw_latent_shape = shape
+        batch.extra["flux_h_lat"] = h_lat
+        batch.extra["flux_w_lat"] = w_lat
+        batch.extra["flux_num_channels_latents"] = num_channels_latents
+        batch.extra["flux_latent_image_ids"] = img_ids
+
+        return batch
+
+
+class FluxDenoisingStage(PipelineStage):
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    @staticmethod
+    def _step_kwargs(scheduler_step, batch: ForwardBatch) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {}
+        sig = inspect.signature(scheduler_step)
+        if "generator" in sig.parameters:
+            gen = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
+            kwargs["generator"] = gen
+        return kwargs
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.timesteps is None:
+            raise ValueError("timesteps must be set before FluxDenoisingStage")
+        if batch.latents is None:
+            raise ValueError("latents must be set before FluxDenoisingStage")
+
+        packed = batch.latents
+        timesteps = batch.timesteps
+
+        pooled = batch.extra["flux_pooled_projections"]
+        enc = batch.extra["flux_encoder_hidden_states"]
+        txt_ids = batch.extra["flux_text_ids"]
+        img_ids = batch.extra["flux_latent_image_ids"]
+
+        neg_pooled = batch.extra.get("flux_negative_pooled_projections")
+        neg_enc = batch.extra.get("flux_negative_encoder_hidden_states")
+
+        true_cfg_scale = float(batch.true_cfg_scale)
+        use_true_cfg = batch.do_classifier_free_guidance and true_cfg_scale > 1.0
+
+        # Prefer the loaded transformer's arch (HF ``guidance_embeds``), not static pipeline defaults.
+        tr_arch = self.transformer.fastvideo_config.arch_config
+        guidance_embeds = bool(getattr(tr_arch, "guidance_embeds", False))
+
+        device = get_local_torch_device()
+        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        bs = packed.shape[0]
+        if guidance_embeds:
+            guidance = torch.full((bs, ), float(batch.guidance_scale), device=device, dtype=torch.float32)
+        else:
+            guidance = None
+
+        step_extras = self._step_kwargs(self.scheduler.step, batch)
+
+        for t in timesteps:
+            t_scalar = t
+            if not isinstance(t_scalar, torch.Tensor):
+                t_scalar = torch.tensor([t_scalar], device=device, dtype=torch.float32)
+            t_scalar = t_scalar.to(device=device, dtype=torch.float32)
+
+            timestep_model = t_scalar.expand(bs).float() / 1000.0
+            timestep_model = timestep_model.to(dtype=target_dtype)
+
+            ts_ctx = int(t_scalar.reshape(-1)[0].item())
+            with (
+                    torch.autocast(
+                        device_type="cuda",
+                        dtype=target_dtype,
+                        enabled=autocast_enabled and device.type == "cuda",
+                    ),
+                    set_forward_context(
+                        current_timestep=ts_ctx,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                    ),
+            ):
+                if use_true_cfg:
+                    assert neg_enc is not None and neg_pooled is not None
+                    n_neg = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=neg_enc,
+                        pooled_projections=neg_pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+                    n_pos = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=enc,
+                        pooled_projections=pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = n_neg + true_cfg_scale * (n_pos - n_neg)
+                else:
+                    noise_pred = self.transformer(
+                        hidden_states=packed,
+                        encoder_hidden_states=enc,
+                        pooled_projections=pooled,
+                        timestep=timestep_model,
+                        guidance=guidance,
+                        txt_ids=txt_ids,
+                        img_ids=img_ids,
+                        return_dict=False,
+                    )[0]
+
+            packed = self.scheduler.step(
+                noise_pred,
+                t_scalar,
+                packed,
+                return_dict=False,
+                **step_extras,
+            )[0]
+
+        batch.latents = packed
+        return batch
+
+
+class FluxDecodingStage(PipelineStage):
+    """Unpack latents, apply VAE scaling/shift, decode to pixels (5D output ``B×3×1×H×W``)."""
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    @staticmethod
+    def _denormalize_latents(latents: torch.Tensor, vae: Any) -> torch.Tensor:
+        cfg = getattr(vae, "config", None)
+        sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None
+        sh = getattr(cfg, "shift_factor", None) if cfg is not None else None
+        if sf is None and hasattr(vae, "scaling_factor"):
+            sf = vae.scaling_factor
+        if sh is None and hasattr(vae, "shift_factor"):
+            sh = vae.shift_factor
+        if sf is not None:
+            latents = latents / (sf.to(latents.device, latents.dtype) if isinstance(sf, torch.Tensor) else sf)
+            if sh is not None:
+                latents = latents + (sh.to(latents.device, latents.dtype) if isinstance(sh, torch.Tensor) else sh)
+        return latents
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        packed = batch.latents
+        if packed is None:
+            raise ValueError("latents must be set before FluxDecodingStage")
+
+        h_lat = int(batch.extra["flux_h_lat"])
+        w_lat = int(batch.extra["flux_w_lat"])
+        num_ch = int(batch.extra["flux_num_channels_latents"])
+        raw_shape = batch.raw_latent_shape
+        if raw_shape is None:
+            raise ValueError("raw_latent_shape missing; FluxLatentPreparationStage must run first.")
+        batch_size = int(raw_shape[0])
+
+        infer_device = get_local_torch_device()
+        packed = packed.to(infer_device)
+
+        latents_4d = _unpack_latents(packed, batch_size, num_ch, h_lat, w_lat)
+        latents_4d = self._denormalize_latents(latents_4d, self.vae)
+
+        vae_device = next(self.vae.parameters()).device
+        latents_4d = latents_4d.to(device=vae_device)
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+        use_cuda_autocast = autocast_enabled and vae_device.type == "cuda"
+
+        with torch.autocast(
+                device_type="cuda",
+                dtype=vae_dtype,
+                enabled=use_cuda_autocast,
+        ):
+            if not autocast_enabled:
+                latents_4d = latents_4d.to(dtype=vae_dtype)
+            dec = self.vae.decode(latents_4d)
+            image = dec.sample if hasattr(dec, "sample") else dec[0]
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        batch.output = image.unsqueeze(2).detach().float().cpu()
+        return batch

fastvideo/pipelines/stages/sr_denoising.py

@@ -61,10 +61,9 @@ class SRDenoisingStage(PipelineStage):
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
             dtype=torch.float16,  # TODO(will): hack
-            supported_attention_backends=(
-                AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
-                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE,
-                AttentionBackendEnum.FLASHINFER_SAGE_ATTN3, AttentionBackendEnum.ATTN_QAT_INFER)  # hack
+            supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
+                                          AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
+                                          AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
         )
 
     def add_noise_to_lq(self, lq_latents: torch.Tensor, strength: float = 0.7) -> torch.Tensor:

fastvideo/platforms/cuda.py

@@ -140,12 +140,6 @@ class CudaPlatformBase(Platform):
             except ImportError as e:
                 logger.info(e)
                 logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
-        elif selected_backend == AttentionBackendEnum.FLASHINFER_SAGE_ATTN3:
-            # No import guard needed: the backend imports without flashinfer
-            # and falls back to torch SDPA internally (it stays usable as a
-            # shape-collection vehicle via FASTVIDEO_ATTN_SHAPE_LOG).
-            logger.info("Using FlashInfer SageAttention3 backend.")
-            return "fastvideo.attention.backends.flashinfer_sageattn3.FlashInferSageAttention3Backend"
         elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
             from fastvideo.attention.backends.attn_qat_infer import (  # noqa: F401
                 AttnQatInferBackend, is_attn_qat_infer_available)

fastvideo/platforms/interface.py

@@ -15,7 +15,6 @@ class AttentionBackendEnum(enum.Enum):
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()
     SAGE_ATTN_THREE = enum.auto()
-    FLASHINFER_SAGE_ATTN3 = enum.auto()
     ATTN_QAT_INFER = enum.auto()
     ATTN_QAT_TRAIN = enum.auto()
     VIDEO_SPARSE_ATTN = enum.auto()

fastvideo/registry.py

@@ -20,6 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
     Cosmos25Config,
     Cosmos25_14BConfig,
 )
+from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config
 from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
@@ -58,11 +59,14 @@ from fastvideo.configs.pipelines.wan import (
     WanT2V480PConfig,
     WanT2V720PConfig,
 )
+from fastvideo.configs.pipelines.glm_image import GlmImageConfig
+from fastvideo.configs.pipelines.flux import FluxPipelineConfig
 from fastvideo.configs.pipelines.sd35 import SD35Config
 from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
 from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
+from fastvideo.api.flux import FluxSamplingParam
 
 from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.logger import init_logger
@@ -718,6 +722,12 @@ def _register_configs() -> None:
             "FastVideo/Matrix-Game-2.0-Base-Diffusers",
             "FastVideo/Matrix-Game-2.0-GTA-Diffusers",
             "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers",
+            # Zelda World Model
+            "mignonjia/mg_longtuning_distilled_zelda",  # distilled using streaming long tuning which init from mg_sf_distilled_zelda_1k_steps
+            "mignonjia/mg_sf_distilled_zelda_1k_steps",  # distilled using self forcing for 1k steps
+            "mignonjia/mg_sf_distilled_zelda",  # distilled using self forcing for 3k steps
+            "mignonjia/mg_causal_zelda",
+            "mignonjia/mg_bidirectional_zelda",
         ],
         model_detectors=[
             lambda path: any(token in path.lower() for token in (
@@ -763,6 +773,22 @@ def _register_configs() -> None:
         default_preset="gen3c_cosmos_7b",
     )
 
+    # Cosmos 3 (must register before Cosmos 2.5 and generic Cosmos detectors
+    # so the cosmos3 path-detection takes precedence)
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Cosmos3Config,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "nvidia/Cosmos3-Nano",
+        ],
+        model_detectors=[
+            lambda path: "cosmos3" in path.lower() or "cosmos-3" in path.lower(),
+        ],
+        model_family="cosmos3",
+        default_preset="cosmos3_nano",
+    )
+
     # Cosmos 2.5 (2B)
     register_configs(
         sampling_param_cls=None,
@@ -1068,6 +1094,33 @@ def _register_configs() -> None:
         default_preset="sd35_medium",
     )
 
+    # GLM-Image
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=GlmImageConfig,
+        hf_model_paths=[
+            "zai-org/GLM-Image",
+        ],
+        model_detectors=[lambda path: "glmimage" in path.lower() or "glm-image" in path.lower()],
+        workload_types=(WorkloadType.T2I, ),
+        model_family="glm_image",
+    )
+
+    # FLUX.1-dev (Diffusers)
+    register_configs(
+        sampling_param_cls=FluxSamplingParam,
+        pipeline_config_cls=FluxPipelineConfig,
+        workload_types=(WorkloadType.T2I, ),
+        hf_model_paths=[
+            "black-forest-labs/FLUX.1-dev",
+        ],
+        model_detectors=[
+            lambda path: "fluxpipeline" in path,
+            lambda path: "flux.1-dev" in path or "flux_1_dev" in path,
+            lambda path: "/flux/" in path or path.endswith("/flux"),
+        ],
+    )
+
 
 # --- Part 3: Main Resolver ---
 
@@ -1158,6 +1211,8 @@ 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.cosmos3.presets import (
+        ALL_PRESETS as COSMOS3_PRESETS, )
     from fastvideo.pipelines.basic.dreamx_world.presets import (
         ALL_PRESETS as DREAMX_WORLD_PRESETS, )
     from fastvideo.pipelines.basic.gamecraft.presets import (
@@ -1195,6 +1250,7 @@ def _register_presets() -> None:
 
     all_preset_groups = (
         COSMOS_PRESETS,
+        COSMOS3_PRESETS,
         DREAMX_WORLD_PRESETS,
         FLUX2_PRESETS,
         GAMECRAFT_PRESETS,

fastvideo/tests/api/test_parser.py

@@ -183,6 +183,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "guidance_scale_2": None,
                 "guidance_rescale": 0.0,
                 "true_cfg_scale": None,
+                "use_embedded_guidance": None,
                 "boundary_ratio": None,
                 "sigmas": None,
             },