[bugfix]: keep sdpa metadata masks 2D; lift to 4D only inside the torch-sdpa impl

SDPAMetadataBuilder.build() lifted 2D [batch, key_len] padding masks to 4D
[batch, 1, 1, key_len], but the metadata is cross-backend: HYWorld and
HunyuanVideo15 build SDPAMetadata while the layer selector picks FLASH_ATTN
when installed. flash_attn.py's equal-length branch padded
attn_mask.shape[1] assuming 2D, so a 4D mask produced a garbage-length pad
and an out-of-bounds gather (CUDA device-side assert, SIGABRT in the CI
transformer tests: test_hyworld crashed, test_matrixgame2/test_wanvideo
never ran).

Store the mask 2D in the builder (the convention every in-tree
producer/consumer assumes) and do the 2D -> [batch, 1, 1, key_len] lift
inside _normalize_attn_mask_for_sdpa, the torch-sdpa-only consumer, so
torch.sdpa still does not misread 2D as its [query_len, key_len]
broadcast. The int -> bool coercion is unchanged. Composes with the
consumer-side normalization in flash_attn.py (98116070): metadata carries
the shared 2D convention, and each impl reshapes for its own kernel.

Adds a CPU-only regression test for the contract and wires it into the CI
unit lane explicitly (fastvideo/tests/attention/ has no lane of its own).

.agents/skills/reseed-performance-baseline/SKILL.md

@@ -69,7 +69,7 @@ approval, then upload reviewed accepted baseline records.
 | `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
 | `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
 | `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
-| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `PERF_MAX_REGRESSION` if set, otherwise `0.05` (5%). |
+| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
 | `intent_rationale` | Yes | One-line explanation for why the baseline shift is legitimate. This is written into provenance and should be reused in the PR. |
 
 Hardcoded defaults:
@@ -148,8 +148,7 @@ For each metric with at least two non-null source values:
 4. Stop if any source record regresses against the batch median by more than
    `max_intra_batch_regression`.
 
-Default `max_intra_batch_regression` to `PERF_MAX_REGRESSION` when set,
-otherwise `0.05`. Print a table with per-source values, batch median, and
+Default `max_intra_batch_regression` to `0.05`. Print a table with per-source values, batch median, and
 worst intra-batch regression.
 
 This check prevents uploading a mixed batch where one JSON is materially
@@ -183,7 +182,7 @@ present, that run is not a valid source for baseline reseeding.
 
 ### 2. Sync and back up existing HF records under /tmp
 
-Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
+Use `fastvideo/performance/hf_store.py` helpers directly. Do **not** use
 `compare_baseline.py` as a sync shortcut; on full main runs it can persist
 records, while this step must only fetch and back up existing history.
 
@@ -192,7 +191,7 @@ The sync command pattern is:
 ```bash
 export PERFORMANCE_TRACKING_ROOT="${PERFORMANCE_TRACKING_ROOT:-/tmp/perf-tracking}"
 export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
-PYTHONPATH=fastvideo/tests/performance python -c 'from hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
+python -c 'from fastvideo.performance.hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
 ```
 
 Then back up only the sanitized model directory under `/tmp`:
@@ -200,8 +199,8 @@ Then back up only the sanitized model directory under `/tmp`:
 ```bash
 SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
 TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)
-MODEL_SAFE=$(PYTHONPATH=fastvideo/tests/performance python - <<'PY'
-from hf_store import sanitize
+MODEL_SAFE=$(python - <<'PY'
+from fastvideo.performance.hf_store import sanitize
 print(sanitize("<model_id>"))
 PY
 )
@@ -235,7 +234,7 @@ first baseline seed. Continue, but report that baseline history was empty.
 Load the last 5 successful records for the target:
 
 ```python
-from hf_store import load_records_for_model
+from fastvideo.performance.hf_store import load_records_for_model
 
 records = load_records_for_model(
     "/tmp/perf-tracking",
@@ -372,7 +371,7 @@ prepared records plus backup on disk.
 Use the shared storage helper so the path and repo type match CI:
 
 ```python
-from hf_store import upload_record
+from fastvideo.performance.hf_store import upload_record
 
 upload_record("<local_record_path>", record, strict=True)
 ```
@@ -460,7 +459,7 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
   intentional baseline replacement.
 - `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
   median comparison, and persistence rules.
-- `fastvideo/tests/performance/hf_store.py` — HF sync, record loading,
+- `fastvideo/performance/hf_store.py` — HF sync, record loading,
   `sanitize()`, and `upload_record()`.
 - `fastvideo/tests/performance/test_inference_performance.py` — source result
   JSON schema.

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

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

.github/workflows/ci-precommit.yml

@@ -1,7 +1,11 @@
 name: pre-commit
 
 on:
-  pull_request:
+  # pull_request_target instead of pull_request: the workflow definition and
+  # the hook config are always taken from the BASE branch, so fork /
+  # first-time-contributor PRs run immediately without a maintainer clicking
+  # "Approve and run". The PR head is checked out as data only.
+  pull_request_target:
     branches: [main]
   workflow_call:
     inputs:
@@ -15,12 +19,25 @@ permissions:
 
 jobs:
   pre-commit:
-    if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
+    if: github.event.pull_request.draft != true
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
       with:
         ref: ${{ inputs.ref || '' }}
+    # For PR events, lint the PR head — but keep the hook definitions from
+    # the base branch so an untrusted PR cannot alter what gets executed.
+    - name: Save trusted hook config
+      if: github.event_name == 'pull_request_target'
+      run: cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
+    - uses: actions/checkout@v4
+      if: github.event_name == 'pull_request_target'
+      with:
+        ref: ${{ github.event.pull_request.head.sha }}
+        persist-credentials: false
+    - name: Restore trusted hook config
+      if: github.event_name == 'pull_request_target'
+      run: cp "$RUNNER_TEMP/trusted-pre-commit-config.yaml" .pre-commit-config.yaml
     - uses: actions/setup-python@v5
       with:
         python-version: "3.12"

apps/performance_dashboard/frontend/src/App.tsx

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

apps/performance_dashboard/frontend/src/api.ts

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

apps/performance_dashboard/frontend/src/styles.css

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

docker/Dockerfile

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

docs/contributing/performance_benchmarks.md

@@ -112,10 +112,11 @@ per-metric policy with direction, percent threshold, absolute threshold, and a
 while it runs so pipeline stage execution times are available in
 `generate_video(...).logging_info`. Stage logs use pipeline-unique keys such as
 `prompt_encoding_stage` so duplicate stage classes do not collide. For
-`PipelineStage` entries, the extractor maps the `stage_class` field:
-`TextEncodingStage` maps to `text_encoder_time_s`, `DenoisingStage` and
-`DmdDenoisingStage` map to `dit_time_s`, and `DecodingStage` maps to
-`vae_decode_time_s`, with a fallback for older logs that used the class name as
+`PipelineStage` entries, shared component stage bases emit a stable
+`component_metric`: text encoding stages map to `text_encoder_time_s`,
+denoising stages and subclasses map to `dit_time_s`, and decoding stages map to
+`vae_decode_time_s`. The extractor falls back to known `stage_class` names for
+older logs that do not include `component_metric` or that used the class name as
 the stage key. Generator-side timings such as `PostDecodeFrameProcessStage`,
 `VideoSaveStage`, and `AudioMuxStage` are intentionally ignored. If a pipeline
 does not report one of the mapped stages, that component metric is stored as
@@ -200,9 +201,9 @@ configs and remain loadable. New or migrated configs should use
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v-1.3b",
-  "variant_id": "canonical",
-  "benchmark_version": 1
+  "workload_id": "wan-t2v",
+  "variant_id": "1.3b-sp2",
+  "benchmark_version": 2
 }
 ```
 
@@ -213,21 +214,25 @@ metadata that make the measured workload explicit:
 
 | Field | Purpose |
 |---|---|
-| `workload_id` | Stable benchmark family, such as `wan-t2v-1.3b`. |
-| `variant_id` | Intentional recipe family, such as `canonical`. |
+| `workload_id` | Stable benchmark family, such as `wan-t2v`. |
+| `variant_id` | Intentional recipe family, including model size and parallelism config, such as `1.3b-sp2`. |
 | `benchmark_version` | Version of the measurement protocol and comparison policy. |
 
 If a config declares `config_schema_version: 2`, loading fails clearly when any
 required v2 identity field is missing. If v2 identity or metadata fields are
 added without `config_schema_version: 2`, loading also fails so partial
 migrations do not silently run as v1 configs. Optional v2 metadata fields
-reserved for follow-up work, such as `recipe`, `metric_threshold_policy`, and
-`quality_metadata`, must be JSON objects when present.
+reserved for follow-up work, such as `metric_threshold_policy` and
+`quality_metadata`, must be JSON objects when present. (`recipe` is emitted
+by the harness and is not config-declarable.)
 
 Recipe fingerprinting, hardware/software profile IDs, exact-identity
-comparison, metric-specific threshold policy behavior, promoted baselines, and
-dashboard regrouping are separate follow-up changes. Until those land, rolling
-baseline comparison remains keyed by `(model_id, gpu_type)`.
+comparison, and dashboard cohort grouping land with this change: v2 records
+compare only within their identity cohort, and a record that opens a NEW
+cohort is marked `baseline_status: "initialized_new_cohort"` (regression
+gating starts once that cohort accumulates history). Legacy v1 configs keep
+the `(model_id, gpu_type)` rolling-baseline comparison. Metric-specific
+threshold policies and promoted baselines remain separate follow-ups.
 
 ### Raw record (`results/perf_*.json`)
 
@@ -237,9 +242,9 @@ Written by `test_inference_performance.py`. One file per benchmark run.
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v-1.3b",
-  "variant_id": "canonical",
-  "benchmark_version": 1,
+  "workload_id": "wan-t2v",
+  "variant_id": "1.3b-sp2",
+  "benchmark_version": 2,
   "model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
   "device": "NVIDIA L40S",
   "num_gpus": 2,
@@ -269,7 +274,43 @@ Written by `test_inference_performance.py`. One file per benchmark run.
   "timestamp": "2026-05-08T22:00:00+00:00",
   "text_encoder_time_s": 2.141,
   "dit_time_s": 8.437,
-  "vae_decode_time_s": 3.208
+  "vae_decode_time_s": 3.208,
+  "recipe": {
+    "recipe_schema_version": 1,
+    "benchmark": {
+      "benchmark_id": "wan-t2v-1.3b-2gpu",
+      "workload_id": "wan-t2v",
+      "variant_id": "1.3b-sp2",
+      "benchmark_version": 2
+    },
+    "model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
+    "init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
+    "generation_kwargs": { "height": 480, "width": 832, "num_frames": 45 },
+    "inputs": { "prompt_count": 1, "prompt_sha256": ["<measured-prompt-sha256>"] },
+    "attention": { "requested_backend": "FLASH_ATTN", "resolved_backend": "FLASH_ATTN" }
+  },
+  "recipe_fingerprint": "<sha256>",
+  "hardware_profile": {
+    "device_type": "cuda",
+    "gpu_count": 2,
+    "gpus": [{ "name": "NVIDIA L40S", "memory_gb": 48, "compute_capability": "8.9" }],
+    "interconnect": "none_or_partial"
+  },
+  "hardware_profile_id": "hw-<sha256-prefix>",
+  "software_profile": {
+    "python": "3.12",
+    "pytorch": "2.12",
+    "cuda": "13.0",
+    "packages": {
+      "fastvideo_kernel": "0.3.2",
+      "flashinfer": "0.2.11",
+      "nvidia_cutlass_dsl": "4.5.0",
+      "triton": "3.4.1"
+    }
+  },
+  "software_profile_id": "sw-<sha256-prefix>",
+  "environment_metadata": { "env": { "IMAGE_VERSION": "py3.12-cuda13.0.0" } },
+  "environment_fingerprint": "env-<sha256-prefix>"
 }
 ```
 
@@ -281,6 +322,9 @@ result, used as the rolling-baseline source of truth.
 ```jsonc
 {
   "model_id": "wan-t2v-1.3b-2gpu",
+  "workload_id": "wan-t2v",
+  "variant_id": "1.3b-sp2",
+  "benchmark_version": 2,
   "timestamp": "2026-05-08T22:00:00+00:00",
   "commit_sha": "<full sha>",
   "gpu_type": "NVIDIA L40S",
@@ -297,6 +341,10 @@ result, used as the rolling-baseline source of truth.
       "gated": true
     }
   },
+  "recipe_fingerprint": "<sha256>",
+  "hardware_profile_id": "hw-<sha256-prefix>",
+  "software_profile_id": "sw-<sha256-prefix>",
+  "environment_fingerprint": "env-<sha256-prefix>",
   "success": true
 }
 ```
@@ -309,6 +357,18 @@ dashboard lists skipped plots for metric series that have no non-null values.
 Records missing both `run_source` and `baseline_eligible` are treated as legacy
 successful main/full-suite uploads and remain eligible for rolling baselines.
 
+New records compare only against the same `model_id`, `gpu_type`,
+`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
+`hardware_profile_id`, and `software_profile_id` cohort.
+`environment_metadata` and `environment_fingerprint` are audit data and are not
+part of the comparison key.
+The recipe prompt digests describe the prompts actually measured by the
+benchmark run; extra configured prompts are ignored unless the benchmark runner
+executes them.
+Software profile package cohorts keep exact versions for relevant
+attention/kernel packages, including FastVideo kernels, FlashAttention,
+FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
+
 ## Environment variable reference
 
 | Variable | Default | Used by | Purpose |
@@ -363,7 +423,7 @@ When the rolling-baseline phase runs, it emits:
      "benchmark_id": "<unique-id>",
      "config_schema_version": 2,
      "workload_id": "<stable-workload-id>",
-     "variant_id": "canonical",
+     "variant_id": "<variant, e.g. 1.3b-sp2>",
      "benchmark_version": 1,
      "model": { "model_path": "...", "model_short_name": "..." },
      "init_kwargs": { "num_gpus": 1, ... },
@@ -390,7 +450,9 @@ When the rolling-baseline phase runs, it emits:
 
    Legacy v1 configs without `config_schema_version` still load, but should not
    gain v2 identity or metadata fields until they are migrated to
-   `config_schema_version: 2`.
+   `config_schema_version: 2`. For v2 configs, `workload_id`, `variant_id`,
+   and `benchmark_version` are part of the comparison key; benchmark runs
+   fail if any of these identity fields are missing.
 
 2. The pytest test auto-discovers all configs — no test code needed. CI
    picks it up on the next `/test performance` run.
@@ -432,5 +494,5 @@ pipelines that did not report a mapped component stage.
 
 **Component timing is `null`** — the generated result did not include a mapped
 stage in `logging_info.stages`. Check that the pipeline emits stage logging
-and that the stage name is listed in `STAGE_METRIC_MAP` in
-`test_inference_performance.py`.
+and that the stage emits `component_metric` or is covered by the legacy
+`STAGE_METRIC_MAP` fallback in `test_inference_performance.py`.

docs/design/inference_schema_parity_inventory.yaml

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

docs/inference/support_matrix.md

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

docs/training/trackers.md

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

docs/training/train_infra.md

@@ -161,6 +161,9 @@ training:
     decay_interval_steps: 0
 ```
 
+See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
+including SwanLab installation and authentication.
+
 ### `callbacks` — Pluggable hooks
 
 Callbacks run at specific points in the training loop (before/after optimizer

examples/inference/basic/basic_dreamx_world.py

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

examples/inference/basic/basic_kandinsky5_i2v.py

@@ -0,0 +1,37 @@
+from fastvideo import VideoGenerator
+
+OUTPUT_PATH = "video_samples_kandinsky5_i2v"
+
+IMAGE_PATH = "assets/girl.png"
+
+
+def main():
+    generator = VideoGenerator.from_pretrained(
+        "kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers", 
+        # "kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"
+        # "kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"
+        num_gpus=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+        # image_encoder_cpu_offload=False,
+    )
+
+    prompt = (
+        "A woman stands up and walks away"
+    )
+    _ = generator.generate_video(
+        prompt,
+        image_path=IMAGE_PATH,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        height=1024,
+        width=1024,
+        num_frames=121,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_kandinsky5_t2v.py

@@ -0,0 +1,37 @@
+from fastvideo import VideoGenerator
+
+OUTPUT_PATH = "video_samples_kandinsky5_t2v"
+
+def main():
+    generator = VideoGenerator.from_pretrained(
+        "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
+        # "kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers"
+        # "kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers"
+        # "kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers"
+        num_gpus=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+        # image_encoder_cpu_offload=False,
+    )
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True,height=512, width=768, num_frames=121)
+
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    _ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=512, width=768, num_frames=121)
+
+
+if __name__ == "__main__":
+    main()

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

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

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

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

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

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

fastvideo/attention/backends/flash_attn.py

@@ -338,7 +338,11 @@ class FlashAttentionImpl(AttentionImpl):
                 )
 
             qkv = torch.stack([query, key, value], dim=2)
-            attn_mask_padded = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True)
+            key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, attn_mask.shape[-1]).to(device=query.device)
+            if key_padding_mask.shape[-1] > qkv.shape[1]:
+                raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
+                                 f"expected at most {qkv.shape[1]}, got {key_padding_mask.shape[-1]}")
+            attn_mask_padded = F.pad(key_padding_mask, (qkv.shape[1] - key_padding_mask.shape[-1], 0), value=True)
             output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)

fastvideo/attention/backends/nabla.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NABLA block-sparse flex-attention backend (Kandinsky5 "nabla" checkpoints).
+
+The block mask is data-dependent: nablaT_v2 mean-pools 64-token blocks of the
+fractal-ordered sequence, thresholds the softmaxed block map, and ORs it with a
+precomputed spatio-temporal-window (STA) mask carried on the attention
+metadata. The mask spans the full sequence, so this backend does not support
+sequence parallelism — use it via LocalAttention only.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+try:
+    from torch.nn.attention.flex_attention import BlockMask, flex_attention
+    flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
+    CAN_USE_FLEX_ATTN = True
+except ImportError:
+    CAN_USE_FLEX_ATTN = False
+
+from fastvideo.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+
+
+def nablaT_v2(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    sta: torch.Tensor,
+    thr: float = 0.9,
+) -> "BlockMask":
+    q = q.transpose(1, 2).contiguous()
+    k = k.transpose(1, 2).contiguous()
+
+    # Map estimation
+    B, h, S, D = q.shape
+    s1 = S // 64
+    qa = q.reshape(B, h, s1, 64, D).mean(-2)
+    ka = k.reshape(B, h, s1, 64, D).mean(-2).transpose(-2, -1)
+    map = qa @ ka
+
+    map = torch.softmax(map / math.sqrt(D), dim=-1)
+    # Map binarization
+    vals, inds = map.sort(-1)
+    cvals = vals.cumsum_(-1)
+    mask = (cvals >= 1 - thr).int()
+    mask = mask.gather(-1, inds.argsort(-1))
+
+    mask = torch.logical_or(mask, sta)
+
+    # BlockMask creation
+    kv_nb = mask.sum(-1).to(torch.int32)
+    kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32)
+    return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None)
+
+
+class NablaAttentionBackend(AttentionBackend):
+
+    @staticmethod
+    def get_name() -> str:
+        return "NABLA_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["NablaAttentionImpl"]:
+        return NablaAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["NablaAttentionMetadata"]:
+        return NablaAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["NablaAttentionMetadataBuilder"]:
+        return NablaAttentionMetadataBuilder
+
+
+@dataclass
+class NablaAttentionMetadata(AttentionMetadata):
+    # Block-level STA window mask [1, 1, S/64, S/64], precomputed once per run.
+    sta_mask: torch.Tensor = None  # type: ignore[assignment]
+    # Cumulative-probability threshold for block-map binarization.
+    P: float = 0.9
+    visual_shape: tuple[int, int, int] = (0, 0, 0)
+
+
+class NablaAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        sta_mask: torch.Tensor,
+        P: float,
+        visual_shape: tuple[int, int, int],
+        **kwargs: Any,
+    ) -> NablaAttentionMetadata:
+        return NablaAttentionMetadata(
+            current_timestep=current_timestep,
+            sta_mask=sta_mask,
+            P=P,
+            visual_shape=visual_shape,
+        )
+
+
+class NablaAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if not CAN_USE_FLEX_ATTN:
+            raise RuntimeError("NABLA attention requires torch.nn.attention.flex_attention, "
+                               "which is unavailable in this PyTorch build.")
+        if causal:
+            raise ValueError("NABLA attention does not support causal masking.")
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: NablaAttentionMetadata,
+    ) -> torch.Tensor:
+        # q/k/v: [B, S, heads, head_dim], fractal-ordered by the model; S % 64 == 0.
+        block_mask = nablaT_v2(query, key, attn_metadata.sta_mask, thr=attn_metadata.P)
+        return flex_attention(
+            query=query.transpose(1, 2),
+            key=key.transpose(1, 2),
+            value=value.transpose(1, 2),
+            block_mask=block_mask,
+        ).transpose(1, 2)

fastvideo/attention/backends/sdpa.py

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

fastvideo/attention/layer.py

@@ -252,6 +252,7 @@ class LocalAttention(nn.Module):
                  causal: bool = False,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
+                 default_backend: AttentionBackendEnum | None = None,
                  **extra_impl_args) -> None:
         super().__init__()
         if softmax_scale is None:
@@ -262,7 +263,10 @@ class LocalAttention(nn.Module):
             num_kv_heads = num_heads
 
         dtype = get_compute_dtype()
-        attn_backend = get_attn_backend(head_size, dtype, supported_attention_backends=supported_attention_backends)
+        attn_backend = get_attn_backend(head_size,
+                                        dtype,
+                                        supported_attention_backends=supported_attention_backends,
+                                        default_backend=default_backend)
         impl_cls = attn_backend.get_impl_cls()
         self.attn_impl = impl_cls(num_heads=num_heads,
                                   head_size=head_size,

fastvideo/attention/selector.py

@@ -84,8 +84,9 @@ def get_attn_backend(
     dtype: torch.dtype,
     supported_attention_backends: tuple[AttentionBackendEnum, ...]
     | None = None,
+    default_backend: AttentionBackendEnum | None = None,
 ) -> type[AttentionBackend]:
-    return _cached_get_attn_backend(head_size, dtype, supported_attention_backends)
+    return _cached_get_attn_backend(head_size, dtype, supported_attention_backends, default_backend)
 
 
 @cache
@@ -94,6 +95,7 @@ def _cached_get_attn_backend(
     dtype: torch.dtype,
     supported_attention_backends: tuple[AttentionBackendEnum, ...]
     | None = None,
+    default_backend: AttentionBackendEnum | None = None,
 ) -> type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.
@@ -112,6 +114,12 @@ def _cached_get_attn_backend(
         if backend_by_env_var is not None:
             selected_backend = backend_name_to_enum(backend_by_env_var)
 
+    # Layer-level default (e.g. a checkpoint that requires a specific sparse
+    # backend). Lower precedence than the global force and the env var, so
+    # users can still override it.
+    if selected_backend is None and default_backend is not None:
+        selected_backend = default_backend
+
     # get device-specific attn_backend
     from fastvideo.platforms import current_platform
 

fastvideo/configs/models/dits/__init__.py

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

fastvideo/configs/models/dits/dreamx_world.py

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

fastvideo/configs/models/dits/kandinsky5.py

@@ -2,14 +2,24 @@
 from dataclasses import dataclass, field
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _is_kandinsky5_transformer_block(n: str, m) -> bool:
+    return ("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
 
 
 @dataclass
 class Kandinsky5ArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [
-        lambda n, m:
-        ("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
-    ])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_kandinsky5_transformer_block])
+
+    # NABLA block-sparse attention for attention_type="nabla" checkpoints, plus
+    # the dense backends every DiT supports.
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.NABLA_ATTN,
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
 
     # Native FastVideo implementation uses the same parameter names as diffusers
     # except FFN internals: Diffusers FFN uses `in_layer/out_layer`, while

fastvideo/configs/models/encoders/base.py

@@ -43,11 +43,16 @@ class TextEncoderArchConfig(EncoderArchConfig):
     require_processor: bool = False
 
     def __post_init__(self) -> None:
-        self.tokenizer_kwargs = {
+        # update_model_arch re-runs __post_init__ after pipeline configs may
+        # have customized tokenizer_kwargs (e.g. kandinsky5/gen3c/longcat set
+        # "padding"); rebuilding the dict here would silently wipe those
+        # customizations, so only fill in defaults for keys not already set.
+        defaults = {
             "truncation": True,
             "max_length": self.text_len,
             "return_tensors": "pt",
         }
+        self.tokenizer_kwargs = defaults | self.tokenizer_kwargs
 
 
 @dataclass

fastvideo/configs/pipelines/__init__.py

@@ -1,10 +1,12 @@
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
+from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
+from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
@@ -16,5 +18,6 @@ __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
-    "HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
+    "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/dreamx_world.py

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

fastvideo/configs/pipelines/kandinsky5.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import Kandinsky5VideoConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, CLIPTextConfig
+from fastvideo.configs.models.encoders.reason1 import Reason1Config
+from fastvideo.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+# Byte-exact copy of the upstream Kandinsky5/diffusers template, including the
+# "promt"/"scren" typos: the checkpoints were trained with this exact system  # codespell:ignore promt,scren
+# prompt, and ENCODE_START_IDX below is the tokenized length of everything
+# before the user prompt. Fixing the typos shifts user content to index 127
+# and mis-conditions every generation.
+KANDINSKY5_PROMPT_TEMPLATE = "\n".join([
+    "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.",  # codespell:ignore promt
+    "Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.",
+    "Describe the location of the video, main characters or objects and their action.",
+    "Describe the dynamism of the video and presented actions.",
+    "Name the visual style of the video: whether it is a professional footage, user generated content, some kind of animation, video game or scren content.",  # codespell:ignore scren
+    "Describe the visual effects, postprocessing and transitions if they are presented in the video.",
+    "Pay attention to the order of key actions shown in the scene.<|im_end|>",
+    "<|im_start|>user\n{}<|im_end|>",
+])
+KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX = 129
+
+
+def kandinsky5_qwen_preprocess_text(prompt: str) -> str:
+    if not prompt.strip():
+        prompt = "."
+    return KANDINSKY5_PROMPT_TEMPLATE.format(prompt)
+
+
+def kandinsky5_qwen_postprocess_text(outputs: BaseEncoderOutput,
+                                     mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if outputs.hidden_states is None:
+        raise RuntimeError("Kandinsky5 Qwen prompt embeddings require hidden_states.")
+    hidden_states = outputs.hidden_states[-1]
+    prompt_embeds = hidden_states[:, KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX:]
+    mask = mask[:, KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX:]
+    if prompt_embeds.shape[1] == 0:
+        prompt_embeds = hidden_states[:, -1:]
+        mask = torch.ones((mask.shape[0], 1), dtype=mask.dtype, device=mask.device)
+    return prompt_embeds, mask
+
+
+def kandinsky5_clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    if outputs.pooler_output is None:
+        raise RuntimeError("Kandinsky5 CLIP pooled output is required.")
+    return outputs.pooler_output
+
+
+@dataclass
+class Kandinsky5T2VConfig(PipelineConfig):
+    """Kandinsky-5.0 Lite text-to-video pipeline configuration."""
+
+    dit_config: DiTConfig = field(default_factory=Kandinsky5VideoConfig)
+    vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig)
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Reason1Config(), CLIPTextConfig()))
+    preprocess_text_funcs: tuple[Callable[[str], Any], ...] = field(
+        default_factory=lambda: (kandinsky5_qwen_preprocess_text, preprocess_text))
+    postprocess_text_funcs: tuple[Callable[..., Any], ...] = field(
+        default_factory=lambda: (kandinsky5_qwen_postprocess_text, kandinsky5_clip_postprocess_text))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", "bf16"))
+    text_encoder_max_lengths: tuple[int, ...] = field(
+        default_factory=lambda: (KANDINSKY5_PROMPT_TEMPLATE_ENCODE_START_IDX + 512, 77))
+
+    flow_shift: float | None = 5.0
+    vae_tiling: bool = True
+
+    def __post_init__(self) -> None:
+        if len(self.text_encoder_configs) != 2:
+            raise ValueError(f"Kandinsky5 pipeline requires exactly 2 text encoders (qwen and clip), "
+                             f"but got {len(self.text_encoder_configs)} encoder(s).")
+        if len(self.text_encoder_precisions) != 2:
+            raise ValueError("Kandinsky5 pipeline requires exactly 2 text encoder precisions, "
+                             f"but got {len(self.text_encoder_precisions)}.")
+        if len(self.text_encoder_max_lengths) != 2:
+            raise ValueError("Kandinsky5 pipeline requires exactly 2 text encoder max lengths, "
+                             f"but got {len(self.text_encoder_max_lengths)}.")
+
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+        qwen_cfg = self.text_encoder_configs[0]
+        qwen_cfg.arch_config.output_hidden_states = True
+        qwen_cfg.arch_config.tokenizer_kwargs.update({
+            "padding": True,
+            "truncation": True,
+            "return_tensors": "pt",
+        })
+
+        clip_cfg = self.text_encoder_configs[1]
+        clip_cfg.arch_config.tokenizer_kwargs.update({
+            "padding": "max_length",
+            "max_length": 77,
+            "truncation": True,
+            "add_special_tokens": True,
+            "return_tensors": "pt",
+        })
+
+
+@dataclass
+class Kandinsky5I2VConfig(Kandinsky5T2VConfig):
+    """Kandinsky-5.0 image-to-video pipeline configuration."""
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # I2V needs the VAE encoder to encode the conditioning image.
+        self.vae_config.load_encoder = True

fastvideo/distributed/parallel_state.py

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

fastvideo/models/dits/causal_wanvideo.py

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

fastvideo/models/dits/dreamx_world.py

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

fastvideo/models/dits/dreamx_world_ar.py

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

fastvideo/models/dits/hunyuanvideo15.py

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

fastvideo/models/dits/hyworld/hyworld.py

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

fastvideo/models/dits/kandinsky5.py

@@ -10,13 +10,21 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from fastvideo.attention import LocalAttention
+from fastvideo.attention.backends.nabla import CAN_USE_FLEX_ATTN, flex_attention, nablaT_v2
 from fastvideo.configs.models.dits import Kandinsky5VideoConfig
 from fastvideo.layers.layernorm import LayerNormScaleShift
 from fastvideo.layers.linear import ReplicatedLinear
 from fastvideo.layers.mlp import MLP
+from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum
 
+logger = init_logger(__name__)
+
+if not CAN_USE_FLEX_ATTN:
+    logger.warning("torch.nn.attention.flex_attention is unavailable in this PyTorch build; "
+                   "Kandinsky5 NABLA sparse attention (Pro checkpoints) cannot be used.")
+
 FRACTAL_PIXEL_SIZE = 8
 _ARCH_CONFIG_DEFAULTS = Kandinsky5VideoConfig().arch_config
 
@@ -263,10 +271,9 @@ class Kandinsky5Modulation(nn.Module):
 
 
 def _apply_rotary(x: torch.Tensor, rope: torch.Tensor) -> torch.Tensor:
-    orig_dtype = x.dtype
     x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32)
     x_out = (rope * x_).sum(dim=-1)
-    return x_out.reshape(*x.shape).to(orig_dtype)
+    return x_out.reshape(*x.shape).to(x.dtype)
 
 
 class Kandinsky5Attention(nn.Module):
@@ -277,6 +284,7 @@ class Kandinsky5Attention(nn.Module):
         head_dim: int,
         supported_attention_backends: tuple[AttentionBackendEnum, ...] | None,
         prefix: str = "",
+        use_nabla: bool = False,
     ):
         super().__init__()
         assert num_channels % head_dim == 0
@@ -306,6 +314,17 @@ class Kandinsky5Attention(nn.Module):
             causal=False,
             supported_attention_backends=supported_attention_backends,
         )
+        # NABLA checkpoints get a second attention layer whose backend defaults
+        # to NABLA_ATTN; FASTVIDEO_ATTENTION_BACKEND still overrides it.
+        self.nabla_attention = None
+        if use_nabla:
+            self.nabla_attention = LocalAttention(
+                num_heads=self.num_heads,
+                head_size=head_dim,
+                causal=False,
+                supported_attention_backends=supported_attention_backends,
+                default_backend=AttentionBackendEnum.NABLA_ATTN,
+            )
 
     def forward(
         self,
@@ -339,27 +358,54 @@ class Kandinsky5Attention(nn.Module):
             key = _apply_rotary(key, rotary_emb).type_as(key)
 
         if sparse_params is not None:
-            raise NotImplementedError(
-                "Sparse attention is not yet supported for Kandinsky5 in FastVideo."
-            )
+            if self.nabla_attention is None:
+                raise RuntimeError("sparse_params passed to an attention layer built without use_nabla; "
+                                   "this checkpoint/config combination is inconsistent.")
+            try:
+                # Backend impl reads sta_mask/P from the forward-context
+                # attention metadata built by the denoising stage.
+                hidden_states = self.nabla_attention(query, key, value)
+            except AssertionError as exc:
+                # Standalone parity tests call the model without a pipeline
+                # forward context; run the NABLA kernel directly.
+                if "Forward context is not set" not in str(exc):
+                    raise
+                attn_mask = nablaT_v2(query, key, sparse_params["sta_mask"], thr=sparse_params["P"])
+                hidden_states = flex_attention(
+                    query=query.transpose(1, 2),
+                    key=key.transpose(1, 2),
+                    value=value.transpose(1, 2),
+                    block_mask=attn_mask,
+                ).transpose(1, 2)
+        else:
+            try:
+                hidden_states = self.local_attention(query, key, value)
 
-        try:
-            hidden_states = self.local_attention(query, key, value)
-        except AssertionError as exc:
-            # LocalAttention requires pipeline forward context. Standalone
-            # parity tests call the model directly, so fallback to Torch SDPA.
-            if "Forward context is not set" not in str(exc):
-                raise
-            query = query.transpose(1, 2)
-            key = key.transpose(1, 2)
-            value = value.transpose(1, 2)
-            hidden_states = F.scaled_dot_product_attention(query,
-                                                           key,
-                                                           value,
-                                                           attn_mask=None,
-                                                           is_causal=False)
-            hidden_states = hidden_states.transpose(1, 2)
-        hidden_states = hidden_states.flatten(2)
+            except AssertionError as exc:
+                # LocalAttention requires pipeline forward context. Standalone
+                # parity tests call the model directly, so fallback to Torch SDPA.
+                if "Forward context is not set" not in str(exc):
+                    raise
+
+                query_shape = query.shape[:-2]
+                key_shape = key.shape[:-2]
+                query = query.reshape(query_shape[0], -1, self.num_heads,
+                                      query.shape[-1]).transpose(1, 2)
+                key = key.reshape(key_shape[0], -1, self.num_heads,
+                                  key.shape[-1]).transpose(1, 2)
+                value = value.reshape(key_shape[0], -1, self.num_heads,
+                                      value.shape[-1]).transpose(1, 2)
+                hidden_states = F.scaled_dot_product_attention(
+                    query,
+                    key,
+                    value,
+                    attn_mask=None,
+                    is_causal=False,
+                )
+                hidden_states = hidden_states.transpose(1, 2).reshape(
+                    *query_shape, self.num_heads, -1)
+                
+        hidden_states = hidden_states.flatten(-2, -1)
 
         hidden_states, _ = self.out_layer(hidden_states)
         return hidden_states
@@ -476,7 +522,8 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
                  head_dim: int,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
-                 prefix: str = ""):
+                 prefix: str = "",
+                 use_nabla: bool = False):
         super().__init__()
         self.visual_modulation = Kandinsky5Modulation(time_dim, model_dim, 9)
 
@@ -491,7 +538,8 @@ class Kandinsky5TransformerDecoderBlock(nn.Module):
             model_dim,
             head_dim,
             supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.self_attention")
+            prefix=f"{prefix}.self_attention",
+            use_nabla=use_nabla)
 
         self.cross_attention_norm = LayerNormScaleShift(
             model_dim,
@@ -624,7 +672,8 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                                               arch.ff_dim,
                                               head_dim,
                                               self._supported_attention_backends,
-                                              prefix=f"{config.prefix}.visual_transformer_blocks.{i}")
+                                              prefix=f"{config.prefix}.visual_transformer_blocks.{i}",
+                                              use_nabla=arch.attention_type == "nabla")
             for i in range(arch.num_visual_blocks)
         ])
 
@@ -694,6 +743,7 @@ class Kandinsky5Transformer3DModel(BaseDiT):
                                                   scale_factor)
         to_fractal = sparse_params[
             "to_fractal"] if sparse_params is not None else False
+        
         visual_embed, visual_rope = fractal_flatten(visual_embed, visual_rope,
                                                     visual_shape,
                                                     block_mask=to_fractal)
@@ -724,6 +774,7 @@ class Kandinsky5Transformer3DModel(BaseDiT):
 
         if return_dict:
             return Kandinsky5TransformerOutput(sample=x)
+        
         return x
 
     def materialize_non_persistent_buffers(self, device: torch.device,

fastvideo/models/encoders/qwen2_5_vl_custom.py

@@ -467,6 +467,18 @@ class Qwen2_5_VisionTransformerPretrainedModel(nn.Module):
         return hidden_states
 
 
+def _compute_default_rope_parameters(config, device=None, seq_len=None, **kwargs):
+    # transformers>=5 removes the "default" entry from ROPE_INIT_FUNCTIONS and
+    # moves rope_theta inside rope_parameters; replicate the 4.x default init.
+    rope_params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
+    base = rope_params.get("rope_theta", getattr(config, "rope_theta", 10000.0))
+    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+    partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+    dim = int(head_dim * partial_rotary_factor)
+    inv_freq = 1.0 / (base**(torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
+    return inv_freq, 1.0
+
+
 class Qwen2_5_VLRotaryEmbedding(nn.Module):
     def __init__(self, config: Qwen2_5_VLConfig, device=None):
         super().__init__()
@@ -479,7 +491,14 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        if self.rope_type in ROPE_INIT_FUNCTIONS:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        elif self.rope_type == "default":
+            # transformers>=5 drops the "default" entry from ROPE_INIT_FUNCTIONS.
+            self.rope_init_fn = _compute_default_rope_parameters
+        else:
+            raise KeyError(f"Unsupported rope_type '{self.rope_type}'; available: "
+                           f"{['default', *ROPE_INIT_FUNCTIONS]}")
 
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -948,6 +967,16 @@ QWEN2_5_VL_ATTENTION_CLASSES = {
 # If FlashAttention2 is not available, transparently fall back to SDPA.
 if not is_flash_attn_2_available():
     QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention
+else:
+    # transformers>=5 only resolves the flash-attn functions when the model
+    # preloads them via its attention interface; this module bypasses
+    # PreTrainedModel, so _flash_attention_forward(implementation=None) raises
+    # unless we preload here.
+    try:
+        from transformers.modeling_flash_attention_utils import lazy_import_flash_attention
+        lazy_import_flash_attention("flash_attention_2")
+    except (ImportError, ValueError):
+        QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention
     
 class Qwen2_5_VLDecoderLayer(nn.Module):
     def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int):
@@ -1035,7 +1064,7 @@ class Qwen2_5_VLModel(nn.Module):
     def __init__(self, config: Qwen2_5_VLConfig):
         super().__init__()
         self.config = config
-        self.padding_idx = config.pad_token_id
+        self.padding_idx = getattr(config, "pad_token_id", None)
         self.vocab_size = config.vocab_size
 
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
@@ -1408,6 +1437,18 @@ class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput):
     rope_deltas: Optional[torch.LongTensor] = None
 
 
+def _flatten_text_config(config):
+    # transformers>=5 stops forwarding text-model attributes (hidden_size,
+    # vocab_size, rope_scaling, ...) from the composite Qwen2_5_VLConfig to
+    # config.text_config; this module reads them from the top level.
+    text_config = getattr(config, "text_config", None)
+    if text_config is not None:
+        for key, value in text_config.to_dict().items():
+            if not hasattr(config, key):
+                setattr(config, key, value)
+    return config
+
+
 class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
     _tied_weights_keys = ["lm_head.weight"]
     config_class = Qwen2_5_VLConfig
@@ -1415,6 +1456,7 @@ class Qwen2_5_VLForConditionalGenerationSimple(nn.Module):
 
     def __init__(self, config):
         super().__init__()
+        config = _flatten_text_config(config)
         self.config = config
         self.visual = Qwen2_5_VisionTransformerPretrainedModel(config.vision_config)
 

fastvideo/models/loader/component_loader.py

@@ -320,6 +320,19 @@ class TextEncoderLoader(ComponentLoader):
                 f"text encoder index {idx} out of range for text_encoder_configs (len={len(encoder_configs)}), model_path={model_path}"
             )
         encoder_config = encoder_configs[idx]
+        if (
+            model_config.get("architectures") == ["CLIPModel"]
+            and isinstance(model_config.get("text_config"), dict)
+        ):
+            valid_arch_fields = {
+                f.name for f in dataclasses.fields(encoder_config.arch_config)
+            }
+            model_config = {
+                key: value
+                for key, value in deepcopy(model_config["text_config"]).items()
+                if key in valid_arch_fields
+            }
+            model_config["architectures"] = ["CLIPTextModel"]
         encoder_config.update_model_arch(model_config)
         if idx < 0 or idx >= len(encoder_precisions):
             raise IndexError(
@@ -404,7 +417,9 @@ class TextEncoderLoader(ComponentLoader):
             else:
                 loaded_weights: set[str] = model.load_weights(
                     self._get_all_weights(
-                        model, model_path, to_cpu=use_cpu_offload
+                        model,
+                        model_path,
+                        to_cpu=fastvideo_args.text_encoder_cpu_offload,
                     )
                 )  # type: ignore
 

fastvideo/models/registry.py

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

fastvideo/performance/hf_store.py

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

fastvideo/performance/metric_policy.py

@@ -36,6 +36,7 @@ DEFAULT_METRIC_POLICIES: tuple[MetricPolicy, ...] = (
     MetricPolicy("vae_decode_time_s", "VAE Decode", 3, True, 0.05, 0.25),
 )
 
+
 def _optional_float(value: Any) -> float | None:
     if value is None or isinstance(value, bool):
         return None
@@ -57,9 +58,7 @@ def _optional_bool(value: Any) -> bool | None:
     return None
 
 
-def resolve_metric_policies(
-    threshold_overrides: Mapping[str, Any] | None,
-) -> tuple[MetricPolicy, ...]:
+def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) -> tuple[MetricPolicy, ...]:
     """Return default metric policies with optional per-metric overrides."""
 
     if not isinstance(threshold_overrides, Mapping):
@@ -80,25 +79,15 @@ def resolve_metric_policies(
                 label=base_policy.label,
                 precision=base_policy.precision,
                 lower_is_better=base_policy.lower_is_better,
-                threshold_percent=(
-                    base_policy.threshold_percent
-                    if threshold_percent is None
-                    else threshold_percent
-                ),
-                threshold_absolute=(
-                    base_policy.threshold_absolute
-                    if threshold_absolute is None
-                    else threshold_absolute
-                ),
+                threshold_percent=(base_policy.threshold_percent if threshold_percent is None else threshold_percent),
+                threshold_absolute=(base_policy.threshold_absolute
+                                    if threshold_absolute is None else threshold_absolute),
                 gated=base_policy.gated if gated is None else gated,
-            )
-        )
+            ))
     return tuple(policies)
 
 
-def serialize_metric_thresholds(
-    policies: tuple[MetricPolicy, ...],
-) -> dict[str, dict[str, float | bool]]:
+def serialize_metric_thresholds(policies: tuple[MetricPolicy, ...], ) -> dict[str, dict[str, float | bool]]:
     return {
         policy.key: {
             "threshold_percent": policy.threshold_percent,
@@ -118,10 +107,7 @@ def regression_delta(
         return None
     absolute_delta = current - baseline if policy.lower_is_better else baseline - current
     percent_delta = absolute_delta / baseline
-    threshold_exceeded = (
-        percent_delta > policy.threshold_percent
-        and absolute_delta > policy.threshold_absolute
-    )
+    threshold_exceeded = (percent_delta > policy.threshold_percent and absolute_delta > policy.threshold_absolute)
     return MetricDelta(
         absolute=absolute_delta,
         percent=percent_delta,

fastvideo/performance_dashboard/service.py

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

fastvideo/pipelines/basic/dreamx_world/__init__.py

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

fastvideo/pipelines/basic/dreamx_world/ar_denoising.py

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

fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py

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

fastvideo/pipelines/basic/dreamx_world/config.py

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

fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py

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

fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py

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

fastvideo/pipelines/basic/dreamx_world/presets.py

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

fastvideo/pipelines/basic/dreamx_world/stages.py

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

fastvideo/pipelines/basic/kandinsky5/__init__.py

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

fastvideo/pipelines/basic/kandinsky5/kandinsky5_i2v_pipeline.py

@@ -0,0 +1,89 @@
+# 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.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.kandinsky5 import (
+    Kandinsky5DecodingStage,
+    Kandinsky5DenoisingStage,
+    Kandinsky5ImageEncodingStage,
+    Kandinsky5LatentPreparationStage,
+    Kandinsky5NormalizationStage,
+)
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+
+
+class Kandinsky5I2VPipeline(ComposedPipelineBase):
+    """Kandinsky-5.0 image-to-video pipeline."""
+
+    _required_config_modules = [
+        "scheduler",
+        "text_encoder",
+        "text_encoder_2",
+        "tokenizer",
+        "tokenizer_2",
+        "transformer",
+        "vae",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+
+        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="timestep_preparation_stage",
+            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=Kandinsky5LatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        # Runs AFTER latent preparation so the initial noise is the seeded
+        # generator's first draw (official kandinsky-5 RNG order); this stage
+        # then samples the image latent and places it into the latents.
+        self.add_stage(
+            stage_name="image_encoding_stage",
+            stage=Kandinsky5ImageEncodingStage(vae=self.get_module("vae")),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=Kandinsky5DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="normalization_stage",
+            stage=Kandinsky5NormalizationStage(),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=Kandinsky5DecodingStage(vae=self.get_module("vae"), pipeline=self),
+        )
+
+
+EntryClass = Kandinsky5I2VPipeline

fastvideo/pipelines/basic/kandinsky5/kandinsky5_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.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.kandinsky5 import (
+    Kandinsky5DecodingStage,
+    Kandinsky5DenoisingStage,
+    Kandinsky5LatentPreparationStage,
+)
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+
+
+class Kandinsky5T2VPipeline(ComposedPipelineBase):
+    """Kandinsky-5.0 Lite text-to-video pipeline."""
+
+    _required_config_modules = [
+        "scheduler",
+        "text_encoder",
+        "text_encoder_2",
+        "tokenizer",
+        "tokenizer_2",
+        "transformer",
+        "vae",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+
+        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="timestep_preparation_stage",
+            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=Kandinsky5LatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=Kandinsky5DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=Kandinsky5DecodingStage(vae=self.get_module("vae"), pipeline=self),
+        )
+
+
+EntryClass = Kandinsky5T2VPipeline

fastvideo/pipelines/basic/kandinsky5/presets.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Kandinsky-5 model family pipeline presets."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_NEGATIVE_PROMPT = ("Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, "
+                    "deformed, walking backwards")
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Main Kandinsky-5 denoising pass",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+    }),
+)
+
+KANDINSKY5_T2V_LITE_5S = InferencePreset(
+    name="kandinsky5_t2v_lite_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Lite T2V 5s",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_T2V_LITE_DISTILLED_5S = InferencePreset(
+    name="kandinsky5_t2v_lite_distilled_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Lite T2V Distilled 5s",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 1.0,
+        "num_inference_steps": 16,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_T2V_PRO_5S = InferencePreset(
+    name="kandinsky5_t2v_pro_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Pro T2V 5s",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_T2V_PRO_DISTILLED_5S = InferencePreset(
+    name="kandinsky5_t2v_pro_distilled_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Pro T2V Distilled 5s",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 1.0,
+        "num_inference_steps": 16,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_I2V_LITE_5S = InferencePreset(
+    name="kandinsky5_i2v_lite_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Lite I2V 5s",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_I2V_PRO_5S = InferencePreset(
+    name="kandinsky5_i2v_pro_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Pro I2V 5s",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_I2V_LITE_DISTILLED_5S = InferencePreset(
+    name="kandinsky5_i2v_lite_distilled_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Lite I2V Distilled 5s",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 1.0,
+        "num_inference_steps": 16,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+KANDINSKY5_I2V_PRO_DISTILLED_5S = InferencePreset(
+    name="kandinsky5_i2v_pro_distilled_5s",
+    version=1,
+    model_family="kandinsky5",
+    description="Kandinsky-5.0 Pro I2V Distilled 5s",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "fps": 24,
+        "guidance_scale": 1.0,
+        "num_inference_steps": 16,
+        "negative_prompt": _NEGATIVE_PROMPT,
+    },
+)
+
+ALL_PRESETS = (KANDINSKY5_T2V_LITE_5S, KANDINSKY5_T2V_LITE_DISTILLED_5S, KANDINSKY5_T2V_PRO_5S,
+               KANDINSKY5_T2V_PRO_DISTILLED_5S, KANDINSKY5_I2V_LITE_5S, KANDINSKY5_I2V_LITE_DISTILLED_5S,
+               KANDINSKY5_I2V_PRO_5S, KANDINSKY5_I2V_PRO_DISTILLED_5S)

fastvideo/pipelines/stages/__init__.py

@@ -33,6 +33,8 @@ from fastvideo.pipelines.basic.ltx2.stages import (  # noqa: F401
 from fastvideo.pipelines.stages.matrixgame2_denoising import MatrixGame2CausalDenoisingStage
 from fastvideo.pipelines.stages.matrixgame3_denoising import MatrixGame3DenoisingStage
 from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage
+from fastvideo.pipelines.stages.kandinsky5 import (Kandinsky5DecodingStage, Kandinsky5DenoisingStage,
+                                                   Kandinsky5LatentPreparationStage)
 from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage
 from fastvideo.pipelines.stages.gen3c_stages import (Gen3CCFGPolicyStage, Gen3CConditioningStage, Gen3CDenoisingStage,
                                                      Gen3CLatentPreparationStage)
@@ -65,6 +67,9 @@ __all__ = [
     "MatrixGame2CausalDenoisingStage",
     "MatrixGame3DenoisingStage",
     "HYWorldDenoisingStage",
+    "Kandinsky5DecodingStage",
+    "Kandinsky5DenoisingStage",
+    "Kandinsky5LatentPreparationStage",
     "GameCraftDenoisingStage",
     "Gen3CCFGPolicyStage",
     "Gen3CConditioningStage",

fastvideo/pipelines/stages/base.py

@@ -34,6 +34,7 @@ class PipelineStage(ABC):
     composed with other stages to create a complete pipeline. Each stage is responsible
     for a specific part of the process, such as prompt encoding, latent preparation, etc.
     """
+    performance_component_metric: str | None = None
 
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         """
@@ -155,6 +156,9 @@ class PipelineStage(ABC):
                 logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000)
                 batch.logging_info.add_stage_execution_time(stage_key, execution_time)
                 batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name)
+                component_metric = self.performance_component_metric
+                if component_metric is not None:
+                    batch.logging_info.add_stage_metric(stage_key, "component_metric", component_metric)
             except Exception as e:
                 torch.cuda.synchronize()
                 execution_time = time.perf_counter() - start_time

fastvideo/pipelines/stages/decoding.py

@@ -28,6 +28,7 @@ class DecodingStage(PipelineStage):
     This stage handles the decoding of latent representations into the final
     output format (e.g., pixel values).
     """
+    performance_component_metric = "vae_decode_time_s"
 
     def __init__(self, vae, pipeline=None) -> None:
         self.vae: ParallelTiledVAE = vae

fastvideo/pipelines/stages/denoising.py

@@ -51,6 +51,7 @@ class DenoisingStage(PipelineStage):
     This stage handles the iterative denoising process that transforms
     the initial noise into the final output.
     """
+    performance_component_metric = "dit_time_s"
 
     def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
         super().__init__()
@@ -190,6 +191,19 @@ class DenoisingStage(PipelineStage):
             },
         )
 
+        dreamx_y_camera = batch.extra.get("dreamx_y_camera", batch.extra.get("y_camera"))
+        if isinstance(dreamx_y_camera, dict):
+            dreamx_y_camera = {
+                key: value.to(device=local_device, dtype=target_dtype) if torch.is_tensor(value) else value
+                for key, value in dreamx_y_camera.items()
+            }
+        dreamx_camera_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "y_camera": dreamx_y_camera,
+            },
+        )
+
         for key in ("flux2_txt_ids", "flux2_img_ids"):
             value = batch.extra.get(key)
             if torch.is_tensor(value):
@@ -241,7 +255,11 @@ class DenoisingStage(PipelineStage):
             # the image latent instead of appending along the channel dim
             assert batch.image_latent is None, "TI2V task should not have image latents"
             assert self.vae is not None, "VAE is not provided for TI2V task"
+            vae_device = next(self.vae.parameters()).device
+            self.vae = self.vae.to(local_device)
             z = self.vae.encode(batch.pil_image).mean.float()
+            if getattr(fastvideo_args, "vae_cpu_offload", False):
+                self.vae = self.vae.to(vae_device)
             if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
                 if isinstance(self.vae.shift_factor, torch.Tensor):
                     z -= self.vae.shift_factor.to(z.device, z.dtype)
@@ -494,6 +512,7 @@ class DenoisingStage(PipelineStage):
                             **pos_cond_kwargs,
                             **action_kwargs,
                             **camera_kwargs,
+                            **dreamx_camera_kwargs,
                             **timesteps_r_kwarg,
                             **flux2_id_kwargs,
                         )
@@ -536,6 +555,7 @@ class DenoisingStage(PipelineStage):
                                     **neg_cond_kwargs,
                                     **action_kwargs,
                                     **camera_kwargs,
+                                    **dreamx_camera_kwargs,
                                     **timesteps_r_kwarg,
                                     **flux2_id_kwargs,
                                 )
@@ -1187,6 +1207,7 @@ class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
 
 class Cosmos25AutoDenoisingStage(PipelineStage):
     """Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
+    performance_component_metric = "dit_time_s"
 
     def __init__(self, transformer, scheduler) -> None:
         super().__init__()

fastvideo/pipelines/stages/kandinsky5.py

@@ -0,0 +1,545 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import contextlib
+from typing import Any
+
+import PIL
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+from tqdm.auto import tqdm
+
+from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder
+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.models.loader.component_loader import TransformerLoader, VAELoader
+from fastvideo.models.vaes.common import ParallelTiledVAE
+from fastvideo.models.vision_utils import normalize, numpy_to_pt, pil_to_numpy, resize
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.decoding import DecodingStage
+from fastvideo.pipelines.stages.encoding import EncodingStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class Kandinsky5LatentPreparationStage(PipelineStage):
+
+    def __init__(self, scheduler, transformer) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+        self.transformer = transformer
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.height is None or batch.width is None:
+            raise ValueError("height and width must be provided for Kandinsky5.")
+        height = int(batch.height)
+        width = int(batch.width)
+        num_frames = int(batch.num_frames)
+
+        temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
+        patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+
+        if num_frames % temporal_ratio != 1:
+            num_frames = num_frames // temporal_ratio * temporal_ratio + 1
+            batch.num_frames = num_frames
+
+        required_divisor_h = spatial_ratio * patch_size[1]
+        required_divisor_w = spatial_ratio * patch_size[2]
+        if height % required_divisor_h != 0 or width % required_divisor_w != 0:
+            raise ValueError(f"Kandinsky5 height must be divisible by {required_divisor_h} and width by "
+                             f"{required_divisor_w}; "
+                             f"got height={height}, width={width}.")
+
+        # NABLA sparse attention (Pro checkpoints) reshapes the post-patch grid
+        # into 8x8 blocks; validate here instead of crashing mid-denoise after
+        # all the encoding work is done.
+        arch_cfg = getattr(self.transformer, "config", None) or fastvideo_args.pipeline_config.dit_config.arch_config
+        if getattr(arch_cfg, "attention_type", "regular") == "nabla":
+            nabla_divisor_h = required_divisor_h * 8
+            nabla_divisor_w = required_divisor_w * 8
+            if height % nabla_divisor_h != 0 or width % nabla_divisor_w != 0:
+                raise ValueError(f"Kandinsky5 NABLA checkpoints require height divisible by {nabla_divisor_h} and "
+                                 f"width divisible by {nabla_divisor_w}; "
+                                 f"got height={height}, width={width}.")
+
+        if isinstance(batch.prompt, list):
+            batch_size = len(batch.prompt)
+        elif batch.prompt is not None:
+            batch_size = 1
+        else:
+            batch_size = batch.prompt_embeds[0].shape[0]
+        batch_size *= batch.num_videos_per_prompt
+
+        dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        device = get_local_torch_device()
+        num_latent_frames = (num_frames - 1) // temporal_ratio + 1
+        num_channels = getattr(
+            self.transformer,
+            "in_visual_dim",
+            fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
+        )
+        shape = (
+            batch_size,
+            num_latent_frames,
+            height // spatial_ratio,
+            width // spatial_ratio,
+            num_channels,
+        )
+
+        if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
+            raise ValueError(f"generator list length {len(batch.generator)} does not match batch size {batch_size}.")
+
+        visual_cond = getattr(self.transformer, "visual_cond", False)
+
+        if batch.latents is None:
+            latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+        else:
+            valid_shapes = [shape]
+            if visual_cond:
+                valid_shapes.append((*shape[:-1], num_channels * 2 + 1))
+            if tuple(batch.latents.shape) not in valid_shapes:
+                raise ValueError(f"Provided latents shape {list(batch.latents.shape)} does not match expected "
+                                 f"Kandinsky5 latent shape(s): {[list(s) for s in valid_shapes]}.")
+            latents = batch.latents.to(device=device, dtype=dtype)
+
+        if visual_cond and latents.shape[-1] == num_channels:
+            cond = torch.zeros_like(latents)
+            cond_mask = torch.zeros(
+                (*latents.shape[:-1], 1),
+                device=latents.device,
+                dtype=latents.dtype,
+            )
+            latents = torch.cat([latents, cond, cond_mask], dim=-1)
+
+        # I2V image conditioning is placed by Kandinsky5ImageEncodingStage,
+        # which runs AFTER this stage so the initial noise is the generator's
+        # first draw (matching the official kandinskylab/kandinsky-5 order).
+        batch.latents = latents
+        batch.raw_latent_shape = (
+            batch_size,
+            num_channels,
+            num_latent_frames,
+            height // spatial_ratio,
+            width // spatial_ratio,
+        )
+        return batch
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        result.add_check("latents", batch.latents, V.none_or_tensor)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        return result
+
+
+class Kandinsky5DenoisingStage(PipelineStage):
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    @staticmethod
+    def _scale_factor(height: int, width: int) -> tuple[float, float, float]:
+        if 480 <= height <= 854 and 480 <= width <= 854:
+            return (1.0, 2.0, 2.0)
+        return (1.0, 3.16, 3.16)
+
+    @staticmethod
+    def _text_rope_pos(mask: torch.Tensor, device: torch.device) -> torch.Tensor:
+        seq_len = int(mask.sum(1).max().item())
+        return torch.arange(seq_len, device=device)
+
+    @staticmethod
+    def fast_sta_nabla(
+        T: int,
+        H: int,
+        W: int,
+        wT: int = 3,
+        wH: int = 3,
+        wW: int = 3,
+        device: torch.device | str = "cuda",
+    ) -> torch.Tensor:
+        """
+        Create a sparse temporal attention (STA) mask for efficient video generation.
+
+        This method generates a mask that limits attention to nearby frames and spatial positions, reducing
+        computational complexity for video generation.
+
+        Args:
+            T (int): Number of temporal frames
+            H (int): Height in latent space
+            W (int): Width in latent space
+            wT (int): Temporal attention window size
+            wH (int): Height attention window size
+            wW (int): Width attention window size
+            device (str): Device to create tensor on
+
+        Returns:
+            torch.Tensor: Sparse attention mask of shape (T*H*W, T*H*W)
+        """
+        max_extent = int(torch.tensor([T, H, W], device=device).amax().item())
+        r = torch.arange(0, max_extent, 1, dtype=torch.int16, device=device)
+        mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs()
+        sta_t, sta_h, sta_w = (
+            mat[:T, :T].flatten(),
+            mat[:H, :H].flatten(),
+            mat[:W, :W].flatten(),
+        )
+        sta_t = sta_t <= wT // 2
+        sta_h = sta_h <= wH // 2
+        sta_w = sta_w <= wW // 2
+        sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten()
+        sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2)
+        return sta.reshape(T * H * W, T * H * W)
+
+    def get_sparse_params(self, sample: torch.Tensor, device: torch.device) -> dict[str, Any] | None:
+        """
+        Generate sparse attention parameters for the transformer based on sample dimensions.
+
+        This method computes the sparse attention configuration needed for efficient video processing in the
+        transformer model.
+
+        Args:
+            sample (torch.Tensor): Input sample tensor
+            device (torch.device): Device to place tensors on
+
+        Returns:
+            Dict: Dictionary containing sparse attention parameters
+        """
+        assert self.transformer.config.patch_size[0] == 1
+        _, T, H, W, _ = sample.shape
+        T, H, W = (
+            T // self.transformer.config.patch_size[0],
+            H // self.transformer.config.patch_size[1],
+            W // self.transformer.config.patch_size[2],
+        )
+        if self.transformer.config.attention_type == "nabla":
+            sta_mask = self.fast_sta_nabla(
+                T,
+                H // 8,
+                W // 8,
+                self.transformer.config.attention_wT,
+                self.transformer.config.attention_wH,
+                self.transformer.config.attention_wW,
+                device=device,
+            )
+
+            sparse_params = {
+                "sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0),
+                "attention_type": self.transformer.config.attention_type,
+                "to_fractal": True,
+                "P": self.transformer.config.attention_P,
+                "wT": self.transformer.config.attention_wT,
+                "wW": self.transformer.config.attention_wW,
+                "wH": self.transformer.config.attention_wH,
+                "add_sta": self.transformer.config.attention_add_sta,
+                "visual_shape": (T, H, W),
+                "method": self.transformer.config.attention_method,
+            }
+        else:
+            sparse_params = None
+
+        return sparse_params
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.timesteps is None:
+            raise ValueError("timesteps must be prepared before Kandinsky5 denoising.")
+        if batch.latents is None:
+            raise ValueError("latents must be prepared before Kandinsky5 denoising.")
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        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
+        latents = batch.latents
+        num_channels = getattr(
+            self.transformer,
+            "in_visual_dim",
+            fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
+        )
+
+        prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
+        pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
+        if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
+            raise ValueError("Kandinsky5 requires Qwen prompt attention masks.")
+        text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)
+
+        neg_prompt_embeds = None
+        neg_pooled = None
+        negative_text_rope_pos = None
+        if batch.do_classifier_free_guidance and batch.negative_prompt_embeds:
+            neg_prompt_embeds = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
+            neg_pooled = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
+            if batch.negative_attention_mask is None or not batch.negative_attention_mask:
+                raise ValueError("Kandinsky5 requires Qwen negative attention masks for CFG.")
+            negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask[0].to(device), device)
+
+        height = int(batch.height)
+        width = int(batch.width)
+        temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
+        num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
+        visual_rope_pos = [
+            torch.arange(num_latent_frames, device=device),
+            torch.arange(height // spatial_ratio // 2, device=device),
+            torch.arange(width // spatial_ratio // 2, device=device),
+        ]
+        scale_factor = self._scale_factor(height, width)
+
+        sparse_params = self.get_sparse_params(latents, device)
+
+        # I2V keeps the first (conditioning) frame fixed during denoising.
+        # Key off the actual image conditioning, not transformer.visual_cond:
+        # official T2V checkpoints also ship visual_cond=True, and skipping
+        # frame 0 for them leaves it as undenoised noise.
+        cond_frames = 1 if batch.image_latent is not None else 0
+
+        trajectory_timesteps: list[torch.Tensor] = []
+        trajectory_latents: list[torch.Tensor] = []
+
+        with tqdm(total=batch.num_inference_steps, desc="Kandinsky5 Denoising") as progress_bar:
+            for i, timestep in enumerate(batch.timesteps):
+                if hasattr(self, "interrupt") and self.interrupt:
+                    break
+
+                t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
+                attn_metadata = None
+                if sparse_params is not None:
+                    attn_metadata = NablaAttentionMetadataBuilder().build(
+                        current_timestep=i,
+                        sta_mask=sparse_params["sta_mask"],
+                        P=sparse_params["P"],
+                        visual_shape=sparse_params["visual_shape"],
+                    )
+                autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
+                                if device.type == "cuda" else contextlib.nullcontext())
+                with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
+                                         forward_batch=batch), autocast_ctx:
+                    pred_velocity = self.transformer(
+                        hidden_states=latents.to(dtype=target_dtype),
+                        encoder_hidden_states=prompt_embeds,
+                        pooled_projections=pooled,
+                        timestep=t_expand,
+                        visual_rope_pos=visual_rope_pos,
+                        text_rope_pos=text_rope_pos,
+                        scale_factor=scale_factor,
+                        sparse_params=sparse_params,
+                        return_dict=True,
+                    ).sample
+
+                    if neg_prompt_embeds is not None and neg_pooled is not None:
+                        uncond_pred_velocity = self.transformer(
+                            hidden_states=latents.to(dtype=target_dtype),
+                            encoder_hidden_states=neg_prompt_embeds,
+                            pooled_projections=neg_pooled,
+                            timestep=t_expand,
+                            visual_rope_pos=visual_rope_pos,
+                            text_rope_pos=negative_text_rope_pos,
+                            scale_factor=scale_factor,
+                            sparse_params=sparse_params,
+                            return_dict=True,
+                        ).sample
+                        pred_velocity = uncond_pred_velocity + batch.guidance_scale * (pred_velocity -
+                                                                                       uncond_pred_velocity)
+
+                latents[:, cond_frames:, :, :, :num_channels] = self.scheduler.step(
+                    pred_velocity[:, cond_frames:],
+                    timestep,
+                    latents[:, cond_frames:, :, :, :num_channels],
+                    return_dict=False,
+                )[0]
+
+                if batch.return_trajectory_latents:
+                    trajectory_timesteps.append(timestep)
+                    # latents is mutated in place, so snapshot a channels-first copy.
+                    trajectory_latents.append(latents[..., :num_channels].permute(0, 4, 1, 2, 3).cpu())
+
+                if i == len(batch.timesteps) - 1 or (i + 1) % self.scheduler.order == 0:
+                    progress_bar.update()
+
+        if trajectory_latents:
+            batch.trajectory_latents = torch.stack(trajectory_latents, dim=1)
+            batch.trajectory_timesteps = torch.stack(trajectory_timesteps, dim=0).cpu()
+
+        batch.latents = latents[:, :, :, :, :num_channels]
+        return batch
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.min_list_length(2))
+        return result
+
+
+class Kandinsky5DecodingStage(DecodingStage):
+
+    def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:
+        super().__init__(vae=vae, pipeline=pipeline)
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.latents is None:
+            raise ValueError("latents must be available before Kandinsky5 decoding.")
+        # Kandinsky5 latents are channels-last [B, T, H, W, C]; the base stage
+        # (and the trajectory latents recorded by the denoising stage) work
+        # channels-first.
+        batch.latents = batch.latents.permute(0, 4, 1, 2, 3).contiguous()
+        return super().forward(batch, fastvideo_args)
+
+
+class Kandinsky5ImageEncodingStage(EncodingStage):
+    """Encode the conditioning image into a VAE latent for I2V."""
+
+    def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:
+        super().__init__(vae=vae)
+
+    @staticmethod
+    def _preprocess(image, height: int, width: int) -> torch.Tensor:
+        if isinstance(image, PIL.Image.Image):
+            image = resize(image, height, width)
+            image = numpy_to_pt(pil_to_numpy(image))  # always lands in [0, 1]
+            return normalize(image)  # [0, 1] -> [-1, 1]
+        # Tensor input: no reliable way to tell [0, 1] from an already
+        # normalized [-1, 1] tensor whose values happen to be non-negative,
+        # so mirror diffusers' heuristic and say what we assumed.
+        if image.min() >= 0:
+            logger.warning("Kandinsky5 conditioning image tensor has no negative values; "
+                           "assuming range [0, 1] and normalizing to [-1, 1]. "
+                           "Pass a [-1, 1] tensor with negative values to skip normalization.")
+            image = normalize(image)  # [0, 1] -> [-1, 1]
+        if image.ndim == 3:
+            image = image.unsqueeze(0)
+        if image.shape[-2:] != (height, width):
+            image = torch.nn.functional.interpolate(image.float(),
+                                                    size=(height, width),
+                                                    mode="bilinear",
+                                                    antialias=True)
+        return image
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.pil_image is None:
+            raise ValueError("Kandinsky5 I2V requires an input image.")
+
+        if not fastvideo_args.model_loaded["vae"]:
+            vae = getattr(self, "vae", None)
+            if vae is None:
+                loader = VAELoader()
+                vae = loader.load(fastvideo_args.model_paths["vae"], fastvideo_args)
+                self.vae = vae
+            fastvideo_args.model_loaded["vae"] = True
+
+        device = get_local_torch_device()
+        vae = self.vae.to(device)
+        self.vae = vae
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast
+
+        # [B, C, H, W] -> [B, C, 1, H, W]
+        image = self._preprocess(batch.pil_image, int(batch.height), int(batch.width))
+        image = image.to(device=device, dtype=torch.float32).unsqueeze(2)
+
+        # Encode the single conditioning frame without tiling (matches diffusers).
+        # The untested causal-VAE spatial_tiled_encode path corrupts the latent.
+        prev_use_tiling = vae.use_tiling
+        vae.use_tiling = False
+        try:
+            with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
+                if not vae_autocast_enabled:
+                    image = image.to(vae_dtype)
+                # Sample with the batch generator (diffusers parity); mode()
+                # would make seed-for-seed reproduction of the reference
+                # pipeline impossible.
+                generator = batch.generator
+                if isinstance(generator, list) and len(generator) != image.shape[0]:
+                    generator = generator[0]
+                image_latent = vae.encode(image).sample(generator=generator)
+        finally:
+            vae.use_tiling = prev_use_tiling
+
+        image_latent = image_latent * vae.scaling_factor
+        # [B, C, 1, H, W] -> [B, 1, H, W, C] to match channels-last latents
+        batch.image_latent = image_latent.permute(0, 2, 3, 4, 1).contiguous()
+
+        # Place the conditioning latent into the prepared latents: frame 0 of
+        # the main channels, the visual_cond channel block, and the mask.
+        # NOTE: the official kandinsky-5 repo leaves the visual_cond block
+        # zeros (generation_utils.py generate()), while the diffusers port
+        # copies the image latent into it. A same-seed A/B on
+        # Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers showed the Diffusers
+        # export requires the copy: zeroing the block produces smeared faces
+        # mid-video. Keep the diffusers semantics for Diffusers-format
+        # checkpoints.
+        latents = batch.latents
+        image_latent = batch.image_latent.to(device=latents.device, dtype=latents.dtype)
+        num_channels = image_latent.shape[-1]
+        latents[:, 0:1, :, :, :num_channels] = image_latent
+        if latents.shape[-1] > num_channels:
+            latents[:, 0:1, :, :, num_channels:2 * num_channels] = image_latent
+            latents[:, 0:1, :, :, 2 * num_channels:] = 1.0
+        batch.latents = latents
+
+        if fastvideo_args.vae_cpu_offload:
+            vae.to("cpu")
+        return batch
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("pil_image", batch.pil_image, V.not_none)
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        # This stage runs after latent preparation and writes into its output.
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)])
+        return result
+
+
+class Kandinsky5NormalizationStage(PipelineStage):
+    """Normalize the first latent frames to reduce I2V conditioning artifacts."""
+
+    COND_FRAMES = 4
+    REFERENCE_FRAMES = 5
+
+    @staticmethod
+    def _adaptive_mean_std(source: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
+        source_mean = source.mean(dim=(1, 2, 3, 4), keepdim=True)
+        source_std = source.std(dim=(1, 2, 3, 4), keepdim=True)
+        # Magic constants limit how far the first frames may drift.
+        ref_mean = torch.clamp(reference.mean(dim=(1, 2, 3, 4), keepdim=True), source_mean - 0.05, source_mean + 0.1)
+        ref_std = torch.clamp(reference.std(dim=(1, 2, 3, 4), keepdim=True), source_std - 0.1, source_std + 0.25)
+        normalized = (source - source_mean) / source_std
+        return normalized * ref_std + ref_mean
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        latents = batch.latents
+        n = self.COND_FRAMES
+        if latents is None or latents.shape[1] <= n:
+            return batch
+
+        reference = latents[:, n:n + min(self.REFERENCE_FRAMES, latents.shape[1] - 1)]
+        latents[:, :n] = self._adaptive_mean_std(latents[:, :n].clone(), reference)
+        batch.latents = latents
+        return batch

fastvideo/pipelines/stages/text_encoding.py

@@ -8,6 +8,8 @@ This module contains implementations of prompt encoding stages for diffusion pip
 import torch
 from typing import Any
 
+from torch.distributed.tensor import DTensor
+
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
@@ -24,6 +26,7 @@ class TextEncodingStage(PipelineStage):
     This stage handles the encoding of text prompts into the embedding space
     expected by the diffusion model.
     """
+    performance_component_metric = "text_encoder_time_s"
 
     def __init__(self, text_encoders, tokenizers) -> None:
         """
@@ -202,6 +205,21 @@ class TextEncodingStage(PipelineStage):
             encoder_config = encoder_cfgs[i]
             preprocess_func = preprocess_funcs[i]
             postprocess_func = postprocess_funcs[i]
+            # cpu_offload semantics: params rest on CPU between calls but the
+            # forward computes on GPU. FSDP2-wrapped encoders (CPUOffloadPolicy;
+            # DTensor params) stream themselves per-layer — leave inputs on the
+            # param device and let FSDP's root pre-forward move them. A plain
+            # module parked on CPU by text_encoder_cpu_offload is swapped to the
+            # target device for the forward and back afterwards, mirroring the
+            # image-encoder/VAE offload pattern.
+            first_param = next(text_encoder.parameters(), None)
+            encoder_device = first_param.device if first_param is not None else torch.device(target_device)
+            moved_for_forward = False
+            if (first_param is not None and not isinstance(first_param, DTensor)
+                    and encoder_device.type != torch.device(target_device).type):
+                text_encoder = text_encoder.to(target_device)
+                encoder_device = torch.device(target_device)
+                moved_for_forward = True
 
             tok_kwargs = dict(encoder_config.tokenizer_kwargs)
             if max_length is not None:
@@ -246,7 +264,7 @@ class TextEncodingStage(PipelineStage):
                     # pre-format prompts into message lists upstream and rely on
                     # the inner tokenizer + full tokenizer_kwargs (which include
                     # add_generation_prompt). Preserve that original path exactly.
-                    text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(target_device)
+                    text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(encoder_device)
                 else:
                     # Two-step approach matching Diffusers: format with chat
                     # template first, then tokenize the resulting strings.
@@ -260,9 +278,9 @@ class TextEncodingStage(PipelineStage):
                             enable_thinking=False,
                         )
                         formatted_texts.append(formatted)
-                    text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(target_device)
+                    text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(encoder_device)
             else:
-                text_inputs = tok(processed_texts, **tok_kwargs).to(target_device)
+                text_inputs = tok(processed_texts, **tok_kwargs).to(encoder_device)
 
             input_ids = text_inputs["input_ids"]
             attention_mask = text_inputs["attention_mask"]
@@ -282,15 +300,18 @@ class TextEncodingStage(PipelineStage):
             except Exception:
                 prompt_embeds, attention_mask = postprocess_func(outputs, attention_mask)
             if is_ltx2 and getattr(outputs, "hidden_states", None):
-                audio_embed = outputs.hidden_states[0]
+                audio_embed = outputs.hidden_states[0].to(device=target_device)
                 if dtype is not None:
                     audio_embed = audio_embed.to(dtype=dtype)
                 audio_embeds_list.append(audio_embed)
+            prompt_embeds = prompt_embeds.to(device=target_device)
             if dtype is not None:
                 prompt_embeds = prompt_embeds.to(dtype=dtype)
             embeds_list.append(prompt_embeds)
             if return_attention_mask:
-                attn_masks_list.append(attention_mask)
+                attn_masks_list.append(attention_mask.to(device=target_device))
+            if moved_for_forward and fastvideo_args.text_encoder_cpu_offload:
+                text_encoder.to("cpu")
         self._last_audio_embeds = audio_embeds_list if is_ltx2 else None
         return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices)
 
@@ -350,6 +371,7 @@ class Cosmos25TextEncodingStage(PipelineStage):
     Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's
     `compute_text_embeddings_online()`.
     """
+    performance_component_metric = "text_encoder_time_s"
 
     def __init__(self, text_encoder) -> None:
         super().__init__()

fastvideo/pipelines/stages/validators.py

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

fastvideo/platforms/cuda.py

@@ -157,6 +157,14 @@ class CudaPlatformBase(Platform):
                 "ATTN_QAT_TRAIN selected but fastvideo_kernel.triton_kernels.attn_qat_train is not built. "
                 "Silent fallback would produce a non-QAT training run; refusing to proceed. "
                 "Install the training kernel or pick a different FASTVIDEO_ATTENTION_BACKEND.")
+        elif selected_backend == AttentionBackendEnum.NABLA_ATTN:
+            from fastvideo.attention.backends.nabla import CAN_USE_FLEX_ATTN
+            if CAN_USE_FLEX_ATTN:
+                logger.info("Using NABLA block-sparse flex-attention backend.")
+                return "fastvideo.attention.backends.nabla.NablaAttentionBackend"
+            raise ImportError("NABLA_ATTN selected but torch.nn.attention.flex_attention is unavailable in this "
+                              "PyTorch build. Silent fallback to dense attention would be orders of magnitude "
+                              "slower and diverge from the reference; upgrade PyTorch or pick a different backend.")
         elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
             try:
                 from fastvideo_kernel import video_sparse_attn  # noqa: F401

fastvideo/platforms/interface.py

@@ -22,6 +22,7 @@ class AttentionBackendEnum(enum.Enum):
     VMOBA_ATTN = enum.auto()
     SLA_ATTN = enum.auto()
     SAGE_SLA_ATTN = enum.auto()
+    NABLA_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()
 
 

fastvideo/registry.py

@@ -20,6 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
     Cosmos25Config,
     Cosmos25_14BConfig,
 )
+from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.gen3c import Gen3CConfig
@@ -27,6 +28,7 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu
                                                    Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig,
                                                    Hunyuan15SR1080PConfig)
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
+from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
@@ -117,6 +119,10 @@ class ConfigInfo:
     workload_types: tuple[WorkloadType, ...]
     model_family: str | None = None
     default_preset: str | None = None
+    # When set, overrides the model_index `_class_name` for pipeline resolution.
+    # Lets a model family map to a specific pipeline class by path/detector
+    # (e.g. a T2V and I2V checkpoint that share a `_class_name`).
+    pipeline_cls_name: str | None = None
 
 
 # The central registry mapping a model name to its configuration information
@@ -137,6 +143,7 @@ def register_configs(
     model_detectors: list[Callable[[str], bool]] | None = None,
     model_family: str | None = None,
     default_preset: str | None = None,
+    pipeline_cls_name: str | None = None,
 ) -> None:
     """Register config classes for a model family.
 
@@ -151,6 +158,7 @@ def register_configs(
         workload_types=workload_types,
         model_family=model_family,
         default_preset=default_preset,
+        pipeline_cls_name=pipeline_cls_name,
     )
 
     if hf_model_paths:
@@ -490,18 +498,174 @@ def _register_configs() -> None:
         default_preset="lingbotworld_i2v",
     )
 
+    def _kandinsky5_detector(require: tuple[str, ...] = (), exclude: tuple[str, ...] = ()) -> Callable[[str], bool]:
+
+        def detect(path: str) -> bool:
+            path_lower = path.lower()
+            if "kandinsky5" not in path_lower and "kandinsky-5" not in path_lower:
+                return False
+            return (all(token in path_lower for token in require) and not any(token in path_lower for token in exclude))
+
+        return detect
+
+    # t2v/i2v exclude each other so a checkpoint stored under a directory
+    # containing the other token (e.g. ~/i2v_experiments/kandinsky5-t2v-ft)
+    # falls through to the model_index _class_name fallback detectors below
+    # instead of being misrouted.
+    _is_kandinsky5_t2v = _kandinsky5_detector(require=("t2v", ), exclude=("i2v", ))
+    _is_kandinsky5_i2v = _kandinsky5_detector(require=("i2v", ), exclude=("t2v", ))
+    _is_kandinsky5_t2v_lite = _kandinsky5_detector(require=("t2v", "lite"), exclude=("i2v", "distilled"))
+    _is_kandinsky5_t2v_pro = _kandinsky5_detector(require=("t2v", "pro"), exclude=("i2v", "distilled"))
+    _is_kandinsky5_t2v_lite_distilled = _kandinsky5_detector(require=("t2v", "lite", "distilled"), exclude=("i2v", ))
+    _is_kandinsky5_t2v_pro_distilled = _kandinsky5_detector(require=("t2v", "pro", "distilled"), exclude=("i2v", ))
+    _is_kandinsky5_i2v_lite = _kandinsky5_detector(require=("i2v", "lite"), exclude=("t2v", "distilled"))
+    _is_kandinsky5_i2v_pro = _kandinsky5_detector(require=("i2v", "pro"), exclude=("t2v", "distilled"))
+    _is_kandinsky5_i2v_lite_distilled = _kandinsky5_detector(require=("i2v", "lite", "distilled"), exclude=("t2v", ))
+    _is_kandinsky5_i2v_pro_distilled = _kandinsky5_detector(require=("i2v", "pro", "distilled"), exclude=("t2v", ))
+
     # Kandinsky5 Lite T2V
     register_configs(
         sampling_param_cls=None,
-        pipeline_config_cls=PipelineConfig,
+        pipeline_config_cls=Kandinsky5T2VConfig,
         workload_types=(WorkloadType.T2V, ),
         hf_model_paths=[
             "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
         ],
         model_detectors=[
-            lambda path: any(token in path.lower() for token in ("kandinsky5", "kandinsky-5")),
+            _is_kandinsky5_t2v_lite,
         ],
         model_family="kandinsky5",
+        default_preset="kandinsky5_t2v_lite_5s",
+        pipeline_cls_name="Kandinsky5T2VPipeline",
+    )
+
+    # Kandinsky5 Pro T2V
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5T2VConfig,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "kandinskylab/Kandinsky-5.0-T2V-Pro-sft-5s-Diffusers",
+        ],
+        model_detectors=[
+            _is_kandinsky5_t2v_pro,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_t2v_pro_5s",
+    )
+
+    # Kandinsky5 Lite T2V Distilled
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5T2VConfig,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "kandinskylab/Kandinsky-5.0-T2V-Lite-distilled16steps-5s-Diffusers",
+        ],
+        model_detectors=[
+            _is_kandinsky5_t2v_lite_distilled,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_t2v_lite_distilled_5s",
+    )
+
+    # Kandinsky5 Pro T2V Distilled
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5T2VConfig,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "kandinskylab/Kandinsky-5.0-T2V-Pro-distilled-5s-Diffusers",
+        ],
+        model_detectors=[
+            _is_kandinsky5_t2v_pro_distilled,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_t2v_pro_distilled_5s",
+    )
+
+    # Kandinsky5 Lite I2V
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5I2VConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=["kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers"],
+        model_detectors=[
+            _is_kandinsky5_i2v_lite,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_i2v_lite_5s",
+        pipeline_cls_name="Kandinsky5I2VPipeline",
+    )
+
+    # Kandinsky5 Pro I2V
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5I2VConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=["kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers"],
+        model_detectors=[
+            _is_kandinsky5_i2v_pro,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_i2v_pro_5s",
+    )
+
+    # Kandinsky5 Pro I2V Distilled
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5I2VConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=["kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers"],
+        model_detectors=[
+            _is_kandinsky5_i2v_pro_distilled,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_i2v_pro_distilled_5s",
+    )
+
+    # Kandinsky5 Lite I2V Distilled (no official hub repo yet; local
+    # conversions get distilled sampling defaults instead of the sft ones the
+    # fallback would apply).
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5I2VConfig,
+        workload_types=(),
+        model_detectors=[
+            _is_kandinsky5_i2v_lite_distilled,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_i2v_lite_distilled_5s",
+        pipeline_cls_name="Kandinsky5I2VPipeline",
+    )
+
+    # Kandinsky5 fallbacks — registered AFTER the variant detectors so those
+    # win first-match. Catch checkpoints the variant detectors cannot resolve:
+    # token-less local paths matched via the model_index _class_name
+    # ("kandinsky5t2vpipeline" carries no lite/pro marker), variant combos
+    # without a dedicated entry (e.g. I2V Lite distilled), and t2v+i2v
+    # ambiguous paths resolved by the checkpoint's _class_name.
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5T2VConfig,
+        workload_types=(),
+        model_detectors=[
+            _is_kandinsky5_t2v,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_t2v_lite_5s",
+        pipeline_cls_name="Kandinsky5T2VPipeline",
+    )
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Kandinsky5I2VConfig,
+        workload_types=(),
+        model_detectors=[
+            _is_kandinsky5_i2v,
+        ],
+        model_family="kandinsky5",
+        default_preset="kandinsky5_i2v_lite_5s",
+        pipeline_cls_name="Kandinsky5I2VPipeline",
     )
 
     # LongCat (T2V, I2V, VC use same config; workload varies by path)
@@ -550,7 +714,7 @@ def _register_configs() -> None:
             "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers",
             "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers",
             "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers",
-            # Legacy HF paths (kept for backward compat — pre-rename names):
+            # Legacy HF paths (kept for backward compat - pre-rename names):
             "FastVideo/Matrix-Game-2.0-Base-Diffusers",
             "FastVideo/Matrix-Game-2.0-GTA-Diffusers",
             "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers",
@@ -773,6 +937,40 @@ def _register_configs() -> None:
         model_family="wan",
         default_preset="wan_2_2_ti2v_5b",
     )
+
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=DreamXWorld5BCamPipelineConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=[
+            "FastVideo/DreamX-World-5B-Cam-Diffusers",
+        ],
+        model_detectors=[
+            # Mutually exclusive with the AR detector below: Cam requires an
+            # explicit "cam" marker so hyphenated AR local paths (e.g.
+            # /ckpts/dreamx-world-5b-converted) don't first-match here —
+            # detector resolution is first-match in registration order.
+            lambda path:
+            ("dreamx-world" in path.lower() and "cam" in path.lower()) or "dreamxworldpipeline" in path.lower()
+        ],
+        model_family="dreamx_world",
+        default_preset="dreamx_world_5b_cam",
+    )
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=DreamXWorld5BARPipelineConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=[
+            "FastVideo/DreamX-World-5B-Diffusers",
+        ],
+        model_detectors=[
+            lambda path:
+            ("dreamx-world-5b" in path.lower() and "cam" not in path.lower()) or "dreamxworldarpipeline" in path.lower(
+            )
+        ],
+        model_family="dreamx_world",
+        default_preset="dreamx_world_5b_ar",
+    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
@@ -902,7 +1100,10 @@ def get_model_info(
     assert config_info is not None, "config_info must be resolved"
 
     if override_pipeline_cls_name:
-        pipeline_name = override_pipeline_cls_name
+        # Explicit override: skip config resolution entirely so checkpoints
+        # without a diffusers model_index.json keep working (and no download
+        # is triggered just to log the replaced name).
+        pipeline_name: str | None = override_pipeline_cls_name
         logger.info("Using override pipeline class name %s", pipeline_name)
     else:
         if os.path.exists(model_path):
@@ -911,6 +1112,12 @@ def get_model_info(
             config = maybe_download_model_index(model_path)
 
         pipeline_name = config.get("_class_name")
+        if config_info.pipeline_cls_name is not None:
+            # The resolved (path/detector-based) config pins the pipeline class,
+            # e.g. an I2V checkpoint whose `_class_name` would otherwise resolve
+            # to the T2V pipeline.
+            logger.info("Pinning pipeline class name from %s to %s", pipeline_name, config_info.pipeline_cls_name)
+            pipeline_name = config_info.pipeline_cls_name
 
     if pipeline_name is None:
         raise ValueError("Model config does not contain a _class_name attribute. "
@@ -951,6 +1158,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.dreamx_world.presets import (
+        ALL_PRESETS as DREAMX_WORLD_PRESETS, )
     from fastvideo.pipelines.basic.gamecraft.presets import (
         ALL_PRESETS as GAMECRAFT_PRESETS, )
     from fastvideo.pipelines.basic.gen3c.presets import (
@@ -961,6 +1170,8 @@ def _register_presets() -> None:
         ALL_PRESETS as HUNYUAN15_PRESETS, )
     from fastvideo.pipelines.basic.hyworld.presets import (
         ALL_PRESETS as HYWORLD_PRESETS, )
+    from fastvideo.pipelines.basic.kandinsky5.presets import (
+        ALL_PRESETS as KANDINSKY5_PRESETS, )
     from fastvideo.pipelines.basic.lingbotworld.presets import (
         ALL_PRESETS as LINGBOTWORLD_PRESETS, )
     from fastvideo.pipelines.basic.longcat.presets import (
@@ -984,12 +1195,14 @@ def _register_presets() -> None:
 
     all_preset_groups = (
         COSMOS_PRESETS,
+        DREAMX_WORLD_PRESETS,
         FLUX2_PRESETS,
         GAMECRAFT_PRESETS,
         GEN3C_PRESETS,
         HUNYUAN_PRESETS,
         HUNYUAN15_PRESETS,
         HYWORLD_PRESETS,
+        KANDINSKY5_PRESETS,
         LINGBOTWORLD_PRESETS,
         LONGCAT_PRESETS,
         LTX2_PRESETS,

fastvideo/tests/api/test_compat_translation.py

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

fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py

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

fastvideo/tests/contract/test_ci_test_collection.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Guard: every test directory must be collected by some CI lane or be on
+the explicit allowlist below.
+
+Three separate incidents on 2026-07-05 found test files that no CI lane
+ever collects (fastvideo/tests/stages/, tests/local_tests/ additions in
+PR #1509, and this sweep found seven dark directories in total): the tests
+pass review, merge, and then silently never run. This test makes going
+dark an explicit, reviewed decision instead of an accident: adding a new
+test directory fails CI until it is either wired into a lane or
+allowlisted here with a reason.
+
+Pure text analysis — no fastvideo imports, no GPU, no torch.
+"""
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+TESTS_ROOT = REPO_ROOT / "fastvideo" / "tests"
+
+# Files whose text constitutes "a CI lane references this directory".
+CI_SOURCES = [
+    TESTS_ROOT / "modal" / "pr_test.py",
+    TESTS_ROOT / "modal" / "ssim_test.py",
+    *sorted((REPO_ROOT / ".buildkite").rglob("*.yml")),
+    *sorted((REPO_ROOT / ".buildkite").rglob("*.sh")),
+]
+
+# Directories that intentionally have no CI lane today. Every entry needs a
+# reason; remove the entry when the directory gets wired into a lane.
+# State as found on 2026-07-05 — these SHOULD shrink over time, not grow.
+ALLOWLIST = {
+    "attention": "no lane yet — GPU attention-backend tests, run manually",
+    "audio": "no lane yet — audio encoder tests, run manually",
+    "distributed": "no lane yet — multi-GPU torchrun tests, run manually",
+    "hooks": "no lane yet — run manually",
+    "layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
+    "nightly": "by design: nightly cadence, not per-PR",
+    "modal": "CI infrastructure itself, not a test suite",
+}
+
+
+def _dirs_with_tests() -> list[str]:
+    dirs = []
+    for child in sorted(TESTS_ROOT.iterdir()):
+        if child.is_dir() and any(child.rglob("test_*.py")):
+            dirs.append(child.name)
+    return dirs
+
+
+def _ci_text() -> str:
+    return "\n".join(
+        src.read_text(errors="replace") for src in CI_SOURCES if src.exists())
+
+
+def test_every_test_directory_is_collected_or_allowlisted():
+    ci_text = _ci_text()
+    dark = [
+        name for name in _dirs_with_tests()
+        if f"tests/{name}" not in ci_text and name not in ALLOWLIST
+    ]
+    assert not dark, (
+        f"Test directories not referenced by any CI lane and not "
+        f"allowlisted: {dark}. Wire them into a lane in "
+        f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an "
+        f"allowlist entry with a reason in {__file__}.")
+
+
+def test_local_tests_stays_out_of_ci():
+    # tests/local_tests/ (repo root) is developer-local by design (author
+    # decision, 2026-07-05): parity scaffolds and machine-specific checks
+    # that must never gate CI. Fail if any CI source starts collecting it.
+    assert "tests/local_tests" not in _ci_text(), (
+        "tests/local_tests/ is local-only by design; remove the CI "
+        "reference or move the tests into a fastvideo/tests/ lane.")
+
+
+def test_allowlist_entries_are_still_real_directories():
+    # A stale allowlist hides regressions; entries must track reality.
+    missing = [
+        name for name in ALLOWLIST
+        if name != "modal" and not (TESTS_ROOT / name).is_dir()
+    ]
+    assert not missing, (
+        f"Allowlisted directories no longer exist — remove them: {missing}")

fastvideo/tests/encoders/test_t5_encoder.py

@@ -12,6 +12,7 @@ from fastvideo.configs.pipelines import CosmosConfig, PipelineConfig, WanT2V480P
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TextEncoderLoader
+from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
 from fastvideo.utils import maybe_download_model, PRECISION_TO_TYPE
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.configs.models.encoders import T5Config, T5LargeConfig
@@ -36,9 +37,11 @@ def t5_model_paths_and_config():
 @pytest.fixture
 def t5_large_model_paths_and_config():
     base_model_path = "nvidia/Cosmos-Predict2-2B-Video2World"
-    model_path = maybe_download_model(base_model_path,
-                                      local_dir=os.path.join(
-                                          'data', base_model_path))
+    local_dir = os.path.join('data', base_model_path)
+    skip_if_gated_repo_inaccessible(base_model_path,
+                                    local_path=local_dir,
+                                    test_name="Cosmos T5-large encoder test")
+    model_path = maybe_download_model(base_model_path, local_dir=local_dir)
     text_encoder_path = os.path.join(model_path, "text_encoder")
     tokenizer_path = os.path.join(model_path, "tokenizer")
     return text_encoder_path, tokenizer_path, CosmosConfig()

fastvideo/tests/modal/pr_test.py

@@ -187,6 +187,8 @@ def run_transformer_tests():
 
 
 @app.function(gpu="L40S:4",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -201,6 +203,8 @@ def run_training_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -276,7 +280,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 
@@ -318,6 +322,8 @@ def run_dreamverse_app_tests():
 
 
 @app.function(gpu="L40S:1",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -400,6 +406,8 @@ def run_lora_extraction_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -427,6 +435,8 @@ def run_performance_tests():
         "export PERF_RUN_SOURCE='unknown'; "
         "export PERF_UPLOAD_POLICY='never'; "
         "fi; "
+        "(nvidia-smi --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu "
+        "--format=csv -l 10 > /tmp/gpu_telemetry.csv 2>/dev/null &); "
         "pytest ./fastvideo/tests/performance -vs; "
         "PYTEST_RC=$?; "
         "PERF_RC=0; "
@@ -435,6 +445,8 @@ def run_performance_tests():
         "PERF_RC=$?; "
         "fi; "
         "python ./fastvideo/tests/performance/dashboard.py || true; "
+        "echo '--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---'; "
+        "cat /tmp/gpu_telemetry.csv || true; "
         "FINAL_RC=$PYTEST_RC; "
         "if [ $FINAL_RC -eq 0 ]; then FINAL_RC=$PERF_RC; fi; "
         "exit $FINAL_RC")

fastvideo/tests/ops/quantization/test_absmax_fp8.py

@@ -34,6 +34,9 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
             method._convert_scale(scale)
 
     def test_create_weights_rejects_invalid_dtype(self):
+        # float64 is outside the supported set (bf16/fp16/fp32). The test
+        # originally passed float32, which create_weights has accepted since
+        # PR #981 — it never ran in CI, so the mismatch went unnoticed.
         method = AbsMaxFP8LinearMethod()
         layer = nn.Module()
         with self.assertRaisesRegex(AssertionError, "only supports"):
@@ -43,7 +46,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
                 output_partition_sizes=[3],
                 input_size=2,
                 output_size=3,
-                params_dtype=torch.float32,
+                params_dtype=torch.float64,
             )
 
     def test_absmax_fp8_parameter_weight_loader(self):

fastvideo/tests/performance/compare_baseline.py

@@ -5,7 +5,8 @@ This script:
 1) reads current benchmark results from fastvideo/tests/performance/results,
 2) syncs the canonical baseline from the configured HF dataset repo,
 3) compares each current record against the median of up to 5 prior
-   baseline-eligible successful records (filtered by gpu_type),
+   baseline-eligible successful records in the same workload/variant/version,
+   GPU, recipe, hardware, and software cohort,
 4) writes normalized records back to the HF dataset repo according to
    PERF_UPLOAD_POLICY,
 5) exits non-zero if any gated metric exceeds both its percent and absolute
@@ -64,6 +65,27 @@ PERF_REPORTS_DIR = os.environ.get("PERF_REPORTS_DIR", "/root/data/perf_reports")
 UPLOAD_POLICY = os.environ.get("PERF_UPLOAD_POLICY", "never").strip().lower()
 VALID_UPLOAD_POLICIES = {"never", "pass", "always"}
 VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
+IDENTITY_KEYS = (
+    "workload_id",
+    "variant_id",
+    "benchmark_version",
+    "recipe",
+    "recipe_fingerprint",
+    "hardware_profile",
+    "hardware_profile_id",
+    "software_profile",
+    "software_profile_id",
+    "environment_metadata",
+    "environment_fingerprint",
+)
+COMPARISON_IDENTITY_KEYS = (
+    "workload_id",
+    "variant_id",
+    "benchmark_version",
+    "recipe_fingerprint",
+    "hardware_profile_id",
+    "software_profile_id",
+)
 
 
 def _should_persist_tracking() -> bool:
@@ -136,6 +158,46 @@ def _record_metadata(run_source: str, result: dict[str, Any]) -> dict[str, Any]:
     }
 
 
+def _identity_metadata(result: dict[str, Any]) -> dict[str, Any]:
+    metadata = {
+        key: result[key]
+        for key in IDENTITY_KEYS
+        if key in result and result[key] is not None
+    }
+    recipe = result.get("recipe")
+    benchmark = recipe.get("benchmark") if isinstance(recipe, dict) else None
+    if isinstance(benchmark, dict):
+        for key in ("workload_id", "variant_id", "benchmark_version"):
+            if key not in metadata and benchmark.get(key) is not None:
+                metadata[key] = benchmark[key]
+    return metadata
+
+
+def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
+    missing = [
+        key for key in COMPARISON_IDENTITY_KEYS
+        if key not in record or record[key] is None or (isinstance(record[key], str) and not record[key].strip())
+    ]
+    if len(missing) == len(COMPARISON_IDENTITY_KEYS):
+        # Legacy v1 record: no identity fields at all. Keep the pre-v2
+        # (model_id, gpu_type) rolling-baseline comparison instead of
+        # crashing — v1 configs are documented as loadable.
+        return {}
+    if missing:
+        raise ValueError("Performance record missing required comparison identity fields: " + ", ".join(missing))
+    return {
+        key: str(record[key])
+        for key in COMPARISON_IDENTITY_KEYS
+    }
+
+
+def _format_identity_filters(filters: dict[str, str]) -> str:
+    if not filters:
+        return ""
+    compact = ", ".join(f"{key}={value}" for key, value in filters.items())
+    return f" ({compact})"
+
+
 def _load_current_results() -> list[dict[str, Any]]:
     pattern = os.path.join(RESULTS_DIR, "perf_*.json")
     records: list[dict[str, Any]] = []
@@ -169,7 +231,7 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
     vae_decode_time = safe_float(result.get("vae_decode_time_s"))
     metric_policies = resolve_metric_policies(result.get("regression_thresholds"))
 
-    return {
+    record = {
         "model_id": model_id,
         "timestamp": timestamp,
         "commit_sha": commit_sha,
@@ -184,6 +246,8 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
         "success": True,
         **_record_metadata(_detect_run_source(), result),
     }
+    record.update(_identity_metadata(result))
+    return record
 
 
 def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
@@ -422,17 +486,31 @@ def main() -> int:
             TRACKING_ROOT,
             record["model_id"],
             record["gpu_type"],
+            **_comparison_identity_filters(record),
             last_n=5,
             successful_only=True,
             baseline_eligible_only=True,
         )
 
         if not baseline_records:
-            print(f"No baseline for {record['model_id']} on "
-                  f"{record['gpu_type']}. Initializing...")
+            # A brand-new cohort has NO regression gating until history
+            # accumulates — make that loud and machine-readable instead of
+            # an indistinguishable pass, so a cohort shift (intended or
+            # accidental, e.g. an identity-field change) never silently
+            # blinds the comparison.
+            record["baseline_status"] = "initialized_new_cohort"
+            print("=" * 72)
+            print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
+                  f"{record['model_id']} on {record['gpu_type']}"
+                  f"{_format_identity_filters(_comparison_identity_filters(record))}")
+            print("Regression gating is INACTIVE for this cohort until "
+                  "baseline history accumulates. If this cohort shift is "
+                  "unexpected, check the identity fields above.")
+            print("=" * 72)
             failures: list[str] = []
             record["success"] = True
         else:
+            record["baseline_status"] = "compared"
             failures = _check_regressions(record, baseline_records, metric_policies)
         if static_threshold_failed:
             failures.append(f"{record['model_id']} fixed-threshold phase failed "

fastvideo/tests/performance/dashboard.py

@@ -29,15 +29,76 @@ METRICS = (
     "dit_time_s",
     "vae_decode_time_s",
 )
+COMPARISON_COHORT_KEYS = (
+    "workload_id",
+    "variant_id",
+    "benchmark_version",
+    "recipe_fingerprint",
+    "hardware_profile_id",
+    "software_profile_id",
+)
+GROUP_KEYS = ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)
+
+
+def _cohort_value(value: object) -> str:
+    if value is None:
+        return ""
+    try:
+        if pd.isna(value):
+            return ""
+    except (TypeError, ValueError):
+        pass
+    return str(value)
+
+
+def _display_value(value: object) -> str:
+    return _cohort_value(value) or "legacy"
+
+
+def _short_value(value: object) -> str:
+    text = _display_value(value)
+    if text == "legacy" or len(text) <= 14:
+        return text
+    return text[:12]
+
+
+def _cohort_title(group_key: tuple[object, ...]) -> str:
+    workload = _display_value(group_key[2])
+    variant = _display_value(group_key[3])
+    version = _display_value(group_key[4])
+    version_label = version if version == "legacy" else f"v{version}"
+    return f"{workload} / {variant} / {version_label}"
+
+
+def _cohort_detail(group_key: tuple[object, ...]) -> str:
+    return " | ".join((
+        f"recipe {_short_value(group_key[5])}",
+        _short_value(group_key[6]),
+        _short_value(group_key[7]),
+    ))
+
+
+def _group_record(group_key: tuple[object, ...]) -> dict[str, str]:
+    return {
+        key: _cohort_value(value)
+        for key, value in zip(GROUP_KEYS, group_key)
+    }
+
+
+def _dashboard_frame(df: pd.DataFrame) -> pd.DataFrame:
+    dashboard_df = df.copy()
+    for key in GROUP_KEYS:
+        if key not in dashboard_df.columns:
+            dashboard_df[key] = ""
+        dashboard_df[key] = dashboard_df[key].map(_cohort_value)
+    return dashboard_df
 
 # -----------------------------
 # 1. Grouping
 # -----------------------------
 def group_data(df: pd.DataFrame):
-    # Group only by model+GPU so each group produces a time-series line.
-    # config_id (commit SHA) is carried as a column for hover/color use.
-    keys = ["model_id", "gpu_type"]
-    return df.groupby(keys, dropna=False)
+    # Group by the same comparison cohort used by baseline gating.
+    return _dashboard_frame(df).groupby(list(GROUP_KEYS), dropna=False)
 
 # -----------------------------
 # 2. Plot builder
@@ -46,15 +107,20 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
     figs = []
     skipped_metrics: list[dict[str, object]] = []
 
-    for (model_id, gpu_type), g in group_data(df):
+    for group_key, g in group_data(df):
+        model_id, gpu_type = group_key[:2]
+        group_record = _group_record(group_key)
+        cohort_title = _cohort_title(group_key)
+        cohort_detail = _cohort_detail(group_key)
         g = g.sort_values("timestamp")
 
         # One chart per metric so the y-axes aren't on wildly different scales
         for metric in METRICS:
             if metric not in g.columns:
                 skipped_metrics.append({
-                    "model_id": model_id,
-                    "gpu_type": gpu_type,
+                    **group_record,
+                    "cohort": cohort_title,
+                    "cohort_detail": cohort_detail,
                     "metric": metric,
                     "reason": "column missing from loaded records",
                     "records": len(g),
@@ -65,8 +131,9 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
             non_null = int(g[metric].notna().sum())
             if non_null == 0:
                 skipped_metrics.append({
-                    "model_id": model_id,
-                    "gpu_type": gpu_type,
+                    **group_record,
+                    "cohort": cohort_title,
+                    "cohort_detail": cohort_detail,
                     "metric": metric,
                     "reason": "no non-null values in loaded records",
                     "records": len(g),
@@ -79,8 +146,8 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
                 x="timestamp",
                 y=metric,
                 markers=True,
-                hover_data=["config_id", "commit_sha"],
-                title=f"{model_id} | {gpu_type} | {metric}",
+                hover_data=["config_id", "commit_sha", *COMPARISON_COHORT_KEYS],
+                title=f"{model_id} | {gpu_type} | {cohort_title} | {cohort_detail} | {metric}",
                 labels={"timestamp": "Time", metric: metric},
             )
             figs.append(fig)
@@ -95,7 +162,7 @@ def render_skipped_metrics(skipped_metrics: list[dict[str, object]]) -> str:
     rows = [
         "<h3>Skipped Metric Plots</h3>",
         "<table>",
-        ("<thead><tr><th>Model</th><th>GPU</th><th>Metric</th>"
+        ("<thead><tr><th>Model</th><th>GPU</th><th>Cohort</th><th>Metric</th>"
          "<th>Records</th><th>Non-null</th><th>Reason</th></tr></thead>"),
         "<tbody>",
     ]
@@ -104,6 +171,7 @@ def render_skipped_metrics(skipped_metrics: list[dict[str, object]]) -> str:
             "<tr>"
             f"<td>{escape(str(item['model_id']))}</td>"
             f"<td>{escape(str(item['gpu_type']))}</td>"
+            f"<td>{escape(str(item['cohort']))}<br><code>{escape(str(item['cohort_detail']))}</code></td>"
             f"<td>{escape(str(item['metric']))}</td>"
             f"<td>{item['records']}</td>"
             f"<td>{item['non_null']}</td>"
@@ -165,6 +233,7 @@ def main() -> None:
         print("Skipped metric plots:")
         for item in skipped_metrics:
             print(f"  - {item['model_id']} | {item['gpu_type']} | "
+                  f"{item['cohort']} | {item['cohort_detail']} | "
                   f"{item['metric']}: {item['reason']} "
                   f"({item['non_null']}/{item['records']} non-null)")
     print(f"Generated {len(figs)} metric plot(s)")

fastvideo/tests/performance/identity.py

@@ -0,0 +1,435 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Comparable identity helpers for performance benchmark records."""
+
+from __future__ import annotations
+
+import dataclasses
+import enum
+import hashlib
+import importlib.metadata
+import json
+import os
+import platform
+import re
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import torch
+
+RECIPE_SCHEMA_VERSION = 1
+PROFILE_ID_LENGTH = 16
+_PATH_EXCLUDED_GENERATION_KEYS = {"output_path", "output_video_name"}
+_PROMPT_KEYS = {"prompt", "negative_prompt", "neg_prompt"}
+REQUIRED_BENCHMARK_IDENTITY_KEYS = (
+    "benchmark_id",
+    "workload_id",
+    "variant_id",
+    "benchmark_version",
+)
+_PROFILE_ENV_VARS = (
+    "CUDA_VISIBLE_DEVICES",
+    "FASTVIDEO_ATTENTION_BACKEND",
+    "IMAGE_VERSION",
+    "UV_TORCH_BACKEND",
+)
+_PACKAGE_DISTRIBUTIONS = {
+    "fastvideo_kernel": ("fastvideo-kernel", "fastvideo_kernel"),
+    "flash_attn": ("flash-attn", "flash_attn"),
+    "flash_attn_4": ("flash-attn-4",),
+    "flash_attention_fp4": ("flash-attention-fp4",),
+    "flashinfer": ("flashinfer-python", "flashinfer"),
+    "nvidia_cutlass_dsl": ("nvidia-cutlass-dsl", "cutlass-dsl"),
+    "sageattention": ("sageattention",),
+    "sageattn3": ("sageattn3",),
+    "triton": ("triton",),
+    "xformers": ("xformers",),
+}
+
+
+def canonical_json(payload: Any) -> str:
+    """Return deterministic, compact JSON for a recipe/profile payload."""
+
+    return json.dumps(
+        _canonicalize(payload),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+    )
+
+
+def sha256_hexdigest(payload: str | bytes) -> str:
+    if isinstance(payload, str):
+        payload = payload.encode("utf-8")
+    return hashlib.sha256(payload).hexdigest()
+
+
+def recipe_fingerprint(recipe: Mapping[str, Any]) -> str:
+    # Runtime-resolved values are audit metadata, not identity: hashing the
+    # resolved HF snapshot revision would churn the cohort (and via the
+    # no-baseline init path, silently reset the rolling baseline) on ANY
+    # upstream repo commit, including model-card edits with unchanged
+    # weights. The fingerprint is pinned to declared inputs only; the
+    # resolved values stay in the stored recipe for auditability.
+    pruned = {
+        key: (dict(value) if isinstance(value, Mapping) else value)
+        for key, value in recipe.items()
+    }
+    model = pruned.get("model")
+    if isinstance(model, dict):
+        model.pop("resolved_revision", None)
+    attention = pruned.get("attention")
+    if isinstance(attention, dict):
+        # Same reasoning as resolved_revision: with requested_backend="auto",
+        # a silent runtime fallback (e.g. a broken flash-attn wheel) must NOT
+        # open a fresh ungated cohort and seed degraded numbers as the new
+        # baseline — it should fail the regression gate in the old cohort.
+        # Intentional backend changes are declared via requested_backend,
+        # which stays in the hash.
+        attention.pop("resolved_backend", None)
+    return sha256_hexdigest(canonical_json(pruned))
+
+
+def profile_id(prefix: str, profile: Mapping[str, Any]) -> str:
+    return f"{prefix}-{sha256_hexdigest(canonical_json(profile))[:PROFILE_ID_LENGTH]}"
+
+
+def hardware_profile_id(profile: Mapping[str, Any]) -> str:
+    return profile_id("hw", profile)
+
+
+def software_profile_id(profile: Mapping[str, Any]) -> str:
+    return profile_id("sw", profile)
+
+
+def environment_fingerprint(metadata: Mapping[str, Any]) -> str:
+    return profile_id("env", metadata)
+
+
+def benchmark_identity_from_config(cfg: Mapping[str, Any]) -> dict[str, Any]:
+    missing = [
+        key for key in REQUIRED_BENCHMARK_IDENTITY_KEYS
+        if _none_if_empty(cfg.get(key)) is None
+    ]
+    if missing:
+        raise ValueError("Benchmark config missing required identity fields: " + ", ".join(missing))
+    return {
+        key: cfg[key]
+        for key in REQUIRED_BENCHMARK_IDENTITY_KEYS
+    }
+
+
+def build_recipe_from_benchmark_config(
+    cfg: Mapping[str, Any],
+    *,
+    attention_backend: str | None = None,
+    resolved_attention_backend: Any | None = None,
+    resolved_model_revision: Any | None = None,
+    measured_prompts: Sequence[Any] | None = None,
+) -> dict[str, Any]:
+    """Build a deterministic recipe document from a benchmark config.
+
+    The recipe captures fields that describe what benchmark workload was run. It
+    intentionally excludes timings, timestamps, commit metadata, and output paths.
+    """
+
+    model = dict(cfg.get("model") or {})
+    init_kwargs = dict(cfg.get("init_kwargs") or {})
+    generation_kwargs = dict(cfg.get("generation_kwargs") or {})
+    benchmark_identity = benchmark_identity_from_config(cfg)
+    prompts = list(measured_prompts) if measured_prompts is not None else list(
+        cfg.get("test_prompts") or ["A cinematic video."])
+    if attention_backend is None:
+        attention_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
+
+    negative_prompt = generation_kwargs.pop("negative_prompt", generation_kwargs.pop("neg_prompt", None))
+    generation_recipe = {
+        key: value
+        for key, value in generation_kwargs.items()
+        if key not in _PATH_EXCLUDED_GENERATION_KEYS and key not in _PROMPT_KEYS
+    }
+
+    return {
+        "recipe_schema_version": RECIPE_SCHEMA_VERSION,
+        "benchmark": benchmark_identity,
+        "model": {
+            "model_path": normalize_model_path(model.get("model_path")),
+            "revision": _none_if_empty(model.get("revision") or cfg.get("revision")),
+            "resolved_revision": _none_if_empty(resolved_model_revision),
+        },
+        "init_kwargs": init_kwargs,
+        "generation_kwargs": generation_recipe,
+        "inputs": {
+            "prompt_count": len(prompts),
+            "prompt_sha256": [_optional_digest(prompt) for prompt in prompts],
+            "negative_prompt_sha256": _optional_digest(negative_prompt),
+        },
+        "attention": {
+            "requested_backend": _none_if_empty(attention_backend) or "auto",
+            "resolved_backend": _none_if_empty(resolved_attention_backend),
+        },
+    }
+
+
+def normalize_model_path(value: Any) -> str | None:
+    if value is None:
+        return None
+    text = os.fspath(value) if isinstance(value, os.PathLike) else str(value)
+    text = text.strip()
+    if not text:
+        return None
+    if _looks_like_local_path(text):
+        return Path(text).expanduser().as_posix()
+    return re.sub(r"/+", "/", text)
+
+
+def resolved_revision_from_model_path(value: Any) -> str | None:
+    normalized = normalize_model_path(value)
+    if normalized is None:
+        return None
+    parts = Path(normalized).parts
+    for index, part in enumerate(parts[:-1]):
+        if part == "snapshots":
+            revision = parts[index + 1]
+            if re.fullmatch(r"[0-9a-f]{40}", revision):
+                return revision
+    return None
+
+
+def hardware_profile(
+    *,
+    num_gpus: int | None = None,
+    gpu_devices: Sequence[Mapping[str, Any]] | None = None,
+    device_type: str | None = None,
+    interconnect: str | None = None,
+) -> dict[str, Any]:
+    """Return normalized hardware cohort metadata.
+
+    ``gpu_devices`` is an injection point for tests and non-CUDA callers.
+    """
+
+    if gpu_devices is not None:
+        devices = [_normalize_gpu_device(device) for device in gpu_devices]
+        return {
+            "device_type": device_type or "cuda",
+            "gpu_count": num_gpus if num_gpus is not None else len(devices),
+            "gpus": devices,
+            "interconnect": interconnect or "unknown",
+        }
+
+    if not torch.cuda.is_available():
+        return {
+            "device_type": device_type or "cpu",
+            "gpu_count": 0,
+            "gpus": [],
+            "interconnect": interconnect or "none",
+        }
+
+    gpu_count = int(num_gpus if num_gpus is not None else torch.cuda.device_count())
+    devices = []
+    cuda_device_count = torch.cuda.device_count()
+    for device_id in range(gpu_count):
+        if device_id >= cuda_device_count:
+            devices.append(_normalize_gpu_device({}))
+            continue
+        props = torch.cuda.get_device_properties(device_id)
+        capability = torch.cuda.get_device_capability(device_id)
+        devices.append(
+            _normalize_gpu_device({
+                "name": props.name,
+                "memory_bytes": props.total_memory,
+                "compute_capability": f"{capability[0]}.{capability[1]}",
+            }))
+
+    return {
+        "device_type": device_type or "cuda",
+        "gpu_count": gpu_count,
+        "gpus": devices,
+        "interconnect": interconnect or _detect_cuda_interconnect(gpu_count),
+    }
+
+
+def software_profile(
+    *,
+    python_version: str | None = None,
+    torch_version: str | None = None,
+    cuda_version: str | None = None,
+    package_versions: Mapping[str, str | None] | None = None,
+) -> dict[str, Any]:
+    """Return normalized software cohort metadata.
+
+    Python, PyTorch, and CUDA use major/minor cohorts. Attention and kernel
+    packages keep exact versions because patch releases can change performance.
+    """
+
+    versions = dict(package_versions) if package_versions is not None else _installed_package_versions()
+    return {
+        "python": _major_minor(python_version or platform.python_version()),
+        "pytorch": _major_minor(torch_version or torch.__version__),
+        "cuda": _major_minor(cuda_version or torch.version.cuda),
+        "packages": {
+            name: str(version)
+            for name, version in sorted(versions.items())
+            if version
+        },
+    }
+
+
+def environment_metadata(
+    *,
+    env: Mapping[str, str] | None = None,
+    package_versions: Mapping[str, str | None] | None = None,
+    hardware: Mapping[str, Any] | None = None,
+    software: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Return audit metadata kept separate from comparable identity keys."""
+
+    source_env = env if env is not None else os.environ
+    full_package_versions = dict(package_versions) if package_versions is not None else _installed_package_versions()
+    return {
+        "python": {
+            "version": platform.python_version(),
+            "implementation": platform.python_implementation(),
+        },
+        "platform": {
+            "system": platform.system(),
+            "machine": platform.machine(),
+            "release": platform.release(),
+        },
+        "torch": {
+            "version": torch.__version__,
+            "cuda": torch.version.cuda,
+            "cudnn": torch.backends.cudnn.version(),
+        },
+        "env": {
+            key: source_env.get(key)
+            for key in _PROFILE_ENV_VARS
+            if source_env.get(key) is not None
+        },
+        "packages": {
+            key: value
+            for key, value in sorted(full_package_versions.items())
+            if value
+        },
+        "hardware_profile": dict(hardware) if hardware is not None else hardware_profile(),
+        "software_profile": dict(software) if software is not None else software_profile(
+            package_versions=full_package_versions),
+    }
+
+
+def _canonicalize(value: Any) -> Any:
+    if dataclasses.is_dataclass(value) and not isinstance(value, type):
+        return _canonicalize(dataclasses.asdict(value))
+    if isinstance(value, Mapping):
+        return {
+            str(key): _canonicalize(value[key])
+            for key in sorted(value, key=lambda item: str(item))
+        }
+    if isinstance(value, (set, frozenset)):
+        return [_canonicalize(item) for item in sorted(value, key=lambda item: str(item))]
+    if isinstance(value, tuple):
+        return [_canonicalize(item) for item in value]
+    if isinstance(value, list):
+        return [_canonicalize(item) for item in value]
+    if isinstance(value, enum.Enum):
+        return value.name
+    if isinstance(value, Path):
+        return value.expanduser().as_posix()
+    if isinstance(value, os.PathLike):
+        return os.fspath(value)
+    if isinstance(value, type):
+        return f"{value.__module__}.{value.__qualname__}"
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    return value
+
+
+def _none_if_empty(value: Any) -> Any | None:
+    if value is None:
+        return None
+    if isinstance(value, str) and not value.strip():
+        return None
+    return value
+
+
+def _optional_digest(value: Any) -> str | None:
+    if value is None:
+        return None
+    return sha256_hexdigest(str(value))
+
+
+def _looks_like_local_path(value: str) -> bool:
+    if value.startswith(("/", "./", "../", "~")):
+        return True
+    looks_like_hf_repo = re.match(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$", value)
+    has_separator = "/" in value or "\\" in value or os.sep in value
+    return has_separator and looks_like_hf_repo is None
+
+
+def _normalize_gpu_device(device: Mapping[str, Any]) -> dict[str, Any]:
+    memory_gb = device.get("memory_gb")
+    if memory_gb is None and device.get("memory_bytes") is not None:
+        memory_gb = round(float(device["memory_bytes"]) / (1024**3))
+    return {
+        "name": _none_if_empty(device.get("name")) or "unknown",
+        "memory_gb": memory_gb,
+        "compute_capability": _none_if_empty(device.get("compute_capability")),
+    }
+
+
+def _detect_cuda_interconnect(gpu_count: int) -> str:
+    if gpu_count <= 1:
+        return "single_gpu"
+    try:
+        from fastvideo.platforms import current_platform
+
+        if current_platform.is_cuda() and current_platform.is_full_nvlink(list(range(gpu_count))):
+            return "full_nvlink"
+        return "none_or_partial"
+    except Exception:
+        return "unknown"
+
+
+def _installed_package_versions() -> dict[str, str | None]:
+    return {
+        name: _distribution_version(*distribution_names)
+        for name, distribution_names in _PACKAGE_DISTRIBUTIONS.items()
+    }
+
+
+def _distribution_version(*names: str) -> str | None:
+    for name in names:
+        try:
+            return importlib.metadata.version(name)
+        except importlib.metadata.PackageNotFoundError:
+            continue
+    return None
+
+
+def _major_minor(version: str | None) -> str | None:
+    if not version:
+        return None
+    match = re.search(r"(\d+)\.(\d+)", version)
+    if match is None:
+        return version
+    return f"{match.group(1)}.{match.group(2)}"
+
+
+__all__ = [
+    "RECIPE_SCHEMA_VERSION",
+    "REQUIRED_BENCHMARK_IDENTITY_KEYS",
+    "benchmark_identity_from_config",
+    "canonical_json",
+    "environment_fingerprint",
+    "environment_metadata",
+    "hardware_profile",
+    "hardware_profile_id",
+    "normalize_model_path",
+    "profile_id",
+    "recipe_fingerprint",
+    "resolved_revision_from_model_path",
+    "build_recipe_from_benchmark_config",
+    "sha256_hexdigest",
+    "software_profile",
+    "software_profile_id",
+]

fastvideo/tests/performance/test_compare_baseline_policy.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import pytest
+
 from fastvideo.tests.performance import compare_baseline
 from fastvideo.performance.metric_policy import resolve_metric_policies
 
@@ -119,6 +121,110 @@ def test_boolean_regression_threshold_values_are_ignored():
     assert latency.threshold_percent == 0.08
     assert latency.threshold_absolute == 0.5
     assert latency.gated is False
+def test_normalized_record_preserves_identity_metadata(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
+    raw = _raw_result()
+    raw.update({
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe": {
+            "recipe_schema_version": 1,
+        },
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile": {
+            "gpu_count": 1,
+        },
+        "hardware_profile_id": "hw-1",
+        "software_profile": {
+            "python": "3.12",
+        },
+        "software_profile_id": "sw-1",
+        "environment_metadata": {
+            "env": {
+                "IMAGE_VERSION": "latest",
+            },
+        },
+        "environment_fingerprint": "env-1",
+    })
+
+    record = compare_baseline.normalize_performance_result(raw)
+
+    assert record["workload_id"] == "wan-t2v"
+    assert record["variant_id"] == "1.3b-sp2"
+    assert record["benchmark_version"] == 2
+    assert record["recipe"] == {"recipe_schema_version": 1}
+    assert record["recipe_fingerprint"] == "recipe-1"
+    assert record["hardware_profile"] == {"gpu_count": 1}
+    assert record["hardware_profile_id"] == "hw-1"
+    assert record["software_profile"] == {"python": "3.12"}
+    assert record["software_profile_id"] == "sw-1"
+    assert record["environment_metadata"] == {"env": {"IMAGE_VERSION": "latest"}}
+    assert record["environment_fingerprint"] == "env-1"
+
+
+def test_normalized_record_reads_identity_labels_from_recipe(monkeypatch):
+    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
+    raw = _raw_result()
+    raw["recipe"] = {
+        "benchmark": {
+            "workload_id": "wan-t2v",
+            "variant_id": "1.3b-sp2",
+            "benchmark_version": 2,
+        },
+    }
+
+    record = compare_baseline.normalize_performance_result(raw)
+
+    assert record["workload_id"] == "wan-t2v"
+    assert record["variant_id"] == "1.3b-sp2"
+    assert record["benchmark_version"] == 2
+
+
+def test_comparison_identity_filters_use_full_issue_key():
+    record = {
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile_id": "hw-1",
+        "software_profile_id": "sw-1",
+    }
+
+    assert compare_baseline._comparison_identity_filters(record) == {
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": "2",
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile_id": "hw-1",
+        "software_profile_id": "sw-1",
+    }
+
+
+def test_comparison_identity_filters_require_full_issue_key():
+    record = {
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 0,
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile_id": "hw-1",
+    }
+
+    with pytest.raises(ValueError, match="software_profile_id"):
+        compare_baseline._comparison_identity_filters(record)
+
+
+def test_comparison_identity_filters_keep_zero_version():
+    record = {
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 0,
+        "recipe_fingerprint": "recipe-1",
+        "hardware_profile_id": "hw-1",
+        "software_profile_id": "sw-1",
+    }
+
+    assert compare_baseline._comparison_identity_filters(record)["benchmark_version"] == "0"
 
 
 def test_baseline_eligibility_only_for_successful_scheduled_main():

fastvideo/tests/performance/test_dashboard_api.py

@@ -56,8 +56,34 @@ def _record(model_id, gpu_type, ts, commit, latency, throughput, success=True, *
 
 def test_summary_endpoint_returns_latest_group_status():
     app = create_app(FakeStore([
-        _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
-        _record("wan", "NVIDIA L40S", "2026-01-02T00:00:00+00:00", "b" * 40, 11.0, 9.0),
+        _record(
+            "wan",
+            "NVIDIA L40S",
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+        ),
+        _record(
+            "wan",
+            "NVIDIA L40S",
+            "2026-01-02T00:00:00+00:00",
+            "b" * 40,
+            11.0,
+            9.0,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+        ),
     ]))
     client = TestClient(app)
 
@@ -71,6 +97,10 @@ def test_summary_endpoint_returns_latest_group_status():
     assert body["rows"][0]["metrics"]["latency"]["threshold_exceeded"] is True
     assert body["rows"][0]["threshold_exceeded_metrics"] == ["latency", "throughput"]
     assert body["rows"][0]["computed_regression_status"] == "fail"
+    assert body["rows"][0]["workload_id"] == "wan-t2v"
+    assert body["rows"][0]["variant_id"] == "1.3b-sp2"
+    assert body["rows"][0]["benchmark_version"] == 2
+    assert body["rows"][0]["recipe_fingerprint"] == "recipe-a"
 
 
 def test_summary_status_is_independent_of_days_window():

fastvideo/tests/performance/test_dashboard_plotly.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pandas as pd
+
+from fastvideo.tests.performance import dashboard
+
+
+def _record(recipe_fingerprint="recipe-a", software_profile_id="sw-a", **overrides):
+    record = {
+        "model_id": "wan-t2v-1.3b-2gpu",
+        "gpu_type": "NVIDIA L40S",
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe_fingerprint": recipe_fingerprint,
+        "hardware_profile_id": "hw-l40s",
+        "software_profile_id": software_profile_id,
+        "timestamp": "2026-01-01T00:00:00+00:00",
+        "commit_sha": "a" * 40,
+        "config_id": "aaaaaaa",
+        "latency": 10.0,
+        "throughput": 4.5,
+        "memory": 10000.0,
+        "text_encoder_time_s": 2.0,
+        "dit_time_s": 8.0,
+        "vae_decode_time_s": 3.0,
+    }
+    record.update(overrides)
+    return record
+
+
+def test_group_data_uses_full_comparison_cohort():
+    df = pd.DataFrame([
+        _record(recipe_fingerprint="recipe-a", software_profile_id="sw-a"),
+        _record(recipe_fingerprint="recipe-b", software_profile_id="sw-b"),
+    ])
+
+    groups = list(dashboard.group_data(df))
+
+    assert len(groups) == 2
+    assert {key[5] for key, _group in groups} == {"recipe-a", "recipe-b"}
+    assert {key[7] for key, _group in groups} == {"sw-a", "sw-b"}
+
+
+def test_group_data_fills_legacy_cohort_columns():
+    df = pd.DataFrame([{
+        "model_id": "legacy-wan",
+        "gpu_type": "NVIDIA L40S",
+        "timestamp": "2026-01-01T00:00:00+00:00",
+        "commit_sha": "a" * 40,
+        "config_id": "aaaaaaa",
+        "latency": 10.0,
+    }])
+
+    groups = list(dashboard.group_data(df))
+
+    assert len(groups) == 1
+    assert groups[0][0][2:] == ("", "", "", "", "", "")
+
+
+def test_build_plots_labels_distinct_cohorts():
+    df = pd.DataFrame([
+        _record(recipe_fingerprint="recipe-a", software_profile_id="sw-a"),
+        _record(recipe_fingerprint="recipe-b", software_profile_id="sw-b"),
+    ])
+
+    figs, skipped_metrics = dashboard.build_plots(df)
+    titles = [fig.layout.title.text for fig in figs]
+
+    assert len(figs) == len(dashboard.METRICS) * 2
+    assert skipped_metrics == []
+    assert any("recipe recipe-a | hw-l40s | sw-a" in title for title in titles)
+    assert any("recipe recipe-b | hw-l40s | sw-b" in title for title in titles)
+
+
+def test_skipped_metric_table_includes_cohort_identity():
+    df = pd.DataFrame([{
+        "model_id": "wan-t2v-1.3b-2gpu",
+        "gpu_type": "NVIDIA L40S",
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "recipe_fingerprint": "recipe-a",
+        "hardware_profile_id": "hw-l40s",
+        "software_profile_id": "sw-a",
+        "timestamp": "2026-01-01T00:00:00+00:00",
+        "commit_sha": "a" * 40,
+        "config_id": "aaaaaaa",
+        "latency": 10.0,
+    }])
+
+    _figs, skipped_metrics = dashboard.build_plots(df)
+    html = dashboard.render_skipped_metrics(skipped_metrics[:1])
+
+    assert skipped_metrics[0]["cohort"] == "wan-t2v / 1.3b-sp2 / v2"
+    assert skipped_metrics[0]["cohort_detail"] == "recipe recipe-a | hw-l40s | sw-a"
+    assert "wan-t2v / 1.3b-sp2 / v2" in html
+    assert "recipe recipe-a | hw-l40s | sw-a" in html

fastvideo/tests/performance/test_dashboard_service.py

@@ -142,6 +142,111 @@ def test_build_latest_summary_separates_informational_threshold_crossing():
     assert rows[0]["threshold_exceeded_metrics"] == ["latency"]
     assert rows[0]["failing_metrics"] == []
     assert rows[0]["computed_regression_status"] == "pass"
+def test_build_latest_summary_keeps_identity_cohorts_separate():
+    records = [
+        _record(
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            recipe_fingerprint="recipe-a",
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+        _record(
+            "2026-01-02T00:00:00+00:00",
+            "b" * 40,
+            20.0,
+            5.0,
+            recipe_fingerprint="recipe-b",
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+        _record(
+            "2026-01-03T00:00:00+00:00",
+            "c" * 40,
+            22.0,
+            4.5,
+            recipe_fingerprint="recipe-b",
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+        ),
+    ]
+
+    rows = build_latest_summary(records)
+    recipe_b_row = next(row for row in rows if row["recipe_fingerprint"] == "recipe-b")
+
+    assert len(rows) == 2
+    assert recipe_b_row["baseline_n"] == 1
+    assert recipe_b_row["metrics"]["latency"]["baseline"] == 20.0
+    assert recipe_b_row["workload_id"] == "wan-t2v"
+    assert recipe_b_row["variant_id"] == "1.3b-sp2"
+    assert recipe_b_row["benchmark_version"] == 2
+
+
+def test_build_latest_summary_keeps_variant_versions_separate():
+    records = [
+        _record(
+            "2026-01-01T00:00:00+00:00",
+            "a" * 40,
+            10.0,
+            10.0,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=1,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+        _record(
+            "2026-01-02T00:00:00+00:00",
+            "b" * 40,
+            20.0,
+            5.0,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+            run_source="scheduled_main",
+            baseline_eligible=True,
+        ),
+        _record(
+            "2026-01-03T00:00:00+00:00",
+            "c" * 40,
+            22.0,
+            4.5,
+            workload_id="wan-t2v",
+            variant_id="1.3b-sp2",
+            benchmark_version=2,
+            recipe_fingerprint="recipe-a",
+            hardware_profile_id="hw-l40s",
+            software_profile_id="sw-cu130",
+        ),
+    ]
+
+    rows = build_latest_summary(records)
+    version_2_row = next(row for row in rows if row["benchmark_version"] == 2)
+
+    assert len(rows) == 2
+    assert version_2_row["baseline_n"] == 1
+    assert version_2_row["metrics"]["latency"]["baseline"] == 20.0
 
 
 def test_filter_records_and_trends_preserve_metric_points():
@@ -218,3 +323,78 @@ def test_load_records_can_filter_baseline_eligible_records(tmp_path):
         "2026-01-02T00:00:00+00:00",
         "2026-01-03T00:00:00+00:00",
     }
+
+
+def test_load_records_for_model_filters_identity_cohort(tmp_path):
+    model_dir = tmp_path / "wan"
+    model_dir.mkdir()
+    (model_dir / "matching.json").write_text(
+        """
+        {
+          "model_id": "wan",
+          "gpu_type": "NVIDIA L40S",
+          "timestamp": "2026-01-01T00:00:00+00:00",
+          "success": true,
+          "baseline_eligible": true,
+          "workload_id": "wan-t2v",
+          "variant_id": "1.3b-sp2",
+          "benchmark_version": 2,
+          "recipe_fingerprint": "recipe-a",
+          "hardware_profile_id": "hw-l40s",
+          "software_profile_id": "sw-cu130"
+        }
+        """,
+        encoding="utf-8",
+    )
+    (model_dir / "other_recipe.json").write_text(
+        """
+        {
+          "model_id": "wan",
+          "gpu_type": "NVIDIA L40S",
+          "timestamp": "2026-01-02T00:00:00+00:00",
+          "success": true,
+          "baseline_eligible": true,
+          "workload_id": "wan-t2v",
+          "variant_id": "1.3b-sp2",
+          "benchmark_version": 2,
+          "recipe_fingerprint": "recipe-b",
+          "hardware_profile_id": "hw-l40s",
+          "software_profile_id": "sw-cu130"
+        }
+        """,
+        encoding="utf-8",
+    )
+    (model_dir / "other_version.json").write_text(
+        """
+        {
+          "model_id": "wan",
+          "gpu_type": "NVIDIA L40S",
+          "timestamp": "2026-01-03T00:00:00+00:00",
+          "success": true,
+          "baseline_eligible": true,
+          "workload_id": "wan-t2v",
+          "variant_id": "1.3b-sp2",
+          "benchmark_version": 3,
+          "recipe_fingerprint": "recipe-a",
+          "hardware_profile_id": "hw-l40s",
+          "software_profile_id": "sw-cu130"
+        }
+        """,
+        encoding="utf-8",
+    )
+
+    records = hf_store.load_records_for_model(
+        str(tmp_path),
+        "wan",
+        "NVIDIA L40S",
+        workload_id="wan-t2v",
+        variant_id="1.3b-sp2",
+        benchmark_version="2",
+        recipe_fingerprint="recipe-a",
+        hardware_profile_id="hw-l40s",
+        software_profile_id="sw-cu130",
+        baseline_eligible_only=True,
+    )
+
+    assert len(records) == 1
+    assert records[0]["recipe_fingerprint"] == "recipe-a"

fastvideo/tests/performance/test_identity.py

@@ -0,0 +1,422 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from copy import deepcopy
+
+import pytest
+import torch
+
+from fastvideo.tests.performance import identity as identity_module
+from fastvideo.tests.performance.identity import (
+    build_recipe_from_benchmark_config,
+    canonical_json,
+    environment_fingerprint,
+    environment_metadata,
+    hardware_profile,
+    hardware_profile_id,
+    recipe_fingerprint,
+    resolved_revision_from_model_path,
+    software_profile,
+    software_profile_id,
+)
+
+
+def _benchmark_config():
+    return {
+        "benchmark_id": "wan-t2v-1.3b-2gpu",
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+        "model": {
+            "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+            "model_short_name": "Wan2.1-T2V-1.3B",
+        },
+        "init_kwargs": {
+            "num_gpus": 2,
+            "sp_size": 2,
+            "tp_size": 1,
+            "vae_sp": True,
+            "vae_tiling": True,
+            "text_encoder_precisions": ["fp32"],
+        },
+        "generation_kwargs": {
+            "height": 480,
+            "width": 832,
+            "num_frames": 45,
+            "num_inference_steps": 4,
+            "guidance_scale": 3,
+            "embedded_cfg_scale": 6,
+            "seed": 1024,
+            "fps": 24,
+            "neg_prompt": "low quality",
+        },
+        "test_prompts": ["A cinematic video."],
+    }
+
+
+def _fingerprint(cfg, attention_backend="FLASH_ATTN"):
+    recipe = build_recipe_from_benchmark_config(cfg, attention_backend=attention_backend)
+    return recipe_fingerprint(recipe)
+
+
+def test_canonical_json_is_stable_for_mapping_order():
+    left = {"b": [2, 1], "a": {"z": "last", "m": "middle"}}
+    right = {"a": {"m": "middle", "z": "last"}, "b": [2, 1]}
+
+    assert canonical_json(left) == canonical_json(right)
+
+
+def test_canonical_json_handles_sets_deterministically():
+    assert canonical_json({"values": {"b", "a"}}) == '{"values":["a","b"]}'
+    assert canonical_json({"values": frozenset(("b", "a"))}) == '{"values":["a","b"]}'
+
+
+def test_same_config_produces_same_recipe_fingerprint():
+    cfg = _benchmark_config()
+
+    assert _fingerprint(cfg) == _fingerprint(deepcopy(cfg))
+
+
+def test_recipe_includes_first_class_benchmark_identity():
+    recipe = build_recipe_from_benchmark_config(_benchmark_config())
+
+    assert recipe["benchmark"] == {
+        "benchmark_id": "wan-t2v-1.3b-2gpu",
+        "workload_id": "wan-t2v",
+        "variant_id": "1.3b-sp2",
+        "benchmark_version": 2,
+    }
+
+
+def test_recipe_requires_first_class_benchmark_identity():
+    cfg = _benchmark_config()
+    del cfg["workload_id"]
+
+    with pytest.raises(ValueError, match="workload_id"):
+        build_recipe_from_benchmark_config(cfg)
+
+
+def test_semantically_equivalent_sequence_forms_match():
+    cfg = _benchmark_config()
+    equivalent = deepcopy(cfg)
+    equivalent["init_kwargs"]["text_encoder_precisions"] = tuple(
+        equivalent["init_kwargs"]["text_encoder_precisions"])
+
+    assert _fingerprint(cfg) == _fingerprint(equivalent)
+
+
+def test_output_path_does_not_change_recipe_fingerprint():
+    cfg = _benchmark_config()
+    with_output_path = deepcopy(cfg)
+    with_output_path["generation_kwargs"]["output_path"] = "/tmp/generated"
+
+    assert _fingerprint(cfg) == _fingerprint(with_output_path)
+
+
+def test_runtime_resolved_values_do_not_change_recipe_fingerprint():
+    # Runtime-resolved values are audit metadata: an upstream repo commit
+    # (even a README-only edit) or a silent attention-backend fallback must
+    # not churn the cohort — that would open a fresh ungated cohort and let
+    # degraded numbers seed a new baseline. Declared inputs (including
+    # requested_backend) stay in the hash.
+    cfg = _benchmark_config()
+    recipe = build_recipe_from_benchmark_config(
+        cfg, attention_backend="FLASH_ATTN",
+        resolved_attention_backend="FLASH_ATTN",
+        resolved_model_revision="aaaa1111")
+    changed = build_recipe_from_benchmark_config(
+        cfg, attention_backend="FLASH_ATTN",
+        resolved_attention_backend="TORCH_SDPA",
+        resolved_model_revision="bbbb2222")
+
+    assert recipe["model"]["resolved_revision"] == "aaaa1111"  # kept for audit
+    assert recipe["attention"]["resolved_backend"] == "FLASH_ATTN"  # kept for audit
+    assert recipe_fingerprint(recipe) == recipe_fingerprint(changed)
+
+
+def test_measured_prompt_override_ignores_unused_configured_prompts():
+    cfg = _benchmark_config()
+    changed = deepcopy(cfg)
+    changed["test_prompts"] = [
+        "Unused prompt that should not affect this measured workload.",
+        cfg["test_prompts"][0],
+    ]
+
+    recipe = build_recipe_from_benchmark_config(cfg, measured_prompts=[cfg["test_prompts"][0]])
+    changed_recipe = build_recipe_from_benchmark_config(changed, measured_prompts=[cfg["test_prompts"][0]])
+
+    assert recipe_fingerprint(recipe) == recipe_fingerprint(changed_recipe)
+
+
+def test_num_inference_steps_changes_recipe_fingerprint():
+    cfg = _benchmark_config()
+    changed = deepcopy(cfg)
+    changed["generation_kwargs"]["num_inference_steps"] = 8
+
+    assert _fingerprint(cfg) != _fingerprint(changed)
+
+
+def test_benchmark_version_changes_recipe_fingerprint():
+    cfg = _benchmark_config()
+    changed = deepcopy(cfg)
+    changed["benchmark_version"] = 3
+
+    assert _fingerprint(cfg) != _fingerprint(changed)
+
+
+def test_attention_backend_changes_recipe_fingerprint():
+    cfg = _benchmark_config()
+
+    assert _fingerprint(cfg, attention_backend="FLASH_ATTN") != _fingerprint(
+        cfg, attention_backend="TORCH_SDPA")
+
+
+def test_precision_changes_recipe_fingerprint():
+    cfg = _benchmark_config()
+    changed = deepcopy(cfg)
+    changed["init_kwargs"]["text_encoder_precisions"] = ["bf16"]
+
+    assert _fingerprint(cfg) != _fingerprint(changed)
+
+
+def test_dimensions_and_seed_change_recipe_fingerprint():
+    cfg = _benchmark_config()
+    different_dimensions = deepcopy(cfg)
+    different_dimensions["generation_kwargs"]["width"] = 1024
+    different_seed = deepcopy(cfg)
+    different_seed["generation_kwargs"]["seed"] = 2048
+
+    assert _fingerprint(cfg) != _fingerprint(different_dimensions)
+    assert _fingerprint(cfg) != _fingerprint(different_seed)
+
+
+def test_distributed_layout_changes_recipe_fingerprint():
+    cfg = _benchmark_config()
+    changed = deepcopy(cfg)
+    changed["init_kwargs"]["sp_size"] = 1
+    changed["init_kwargs"]["tp_size"] = 2
+
+    assert _fingerprint(cfg) != _fingerprint(changed)
+
+
+def test_software_profile_uses_exact_attention_kernel_package_versions_for_id():
+    profile_a = software_profile(
+        python_version="3.12.4",
+        torch_version="2.12.0+cu130",
+        cuda_version="13.0",
+        package_versions={
+            "triton": "3.4.1",
+            "fastvideo_kernel": "0.3.2",
+            "flash_attn_4": "4.0.0.dev0",
+            "flashinfer": "0.2.11",
+            "nvidia_cutlass_dsl": "4.5.0",
+        },
+    )
+    profile_b = software_profile(
+        python_version="3.12.5",
+        torch_version="2.12.1+cu130",
+        cuda_version="13.0",
+        package_versions={
+            "triton": "3.4.9",
+            "fastvideo_kernel": "0.3.7",
+            "flash_attn_4": "4.0.0.dev1",
+            "flashinfer": "0.2.12",
+            "nvidia_cutlass_dsl": "4.5.1",
+        },
+    )
+
+    assert profile_a == {
+        "python": "3.12",
+        "pytorch": "2.12",
+        "cuda": "13.0",
+        "packages": {
+            "fastvideo_kernel": "0.3.2",
+            "flash_attn_4": "4.0.0.dev0",
+            "flashinfer": "0.2.11",
+            "nvidia_cutlass_dsl": "4.5.0",
+            "triton": "3.4.1",
+        },
+    }
+    assert profile_b["python"] == "3.12"
+    assert profile_b["pytorch"] == "2.12"
+    assert profile_b["cuda"] == "13.0"
+    assert software_profile_id(profile_a) != software_profile_id(profile_b)
+
+
+def test_installed_package_versions_checks_attention_kernel_distributions(monkeypatch):
+    versions = {
+        "fastvideo-kernel": "0.3.2",
+        "flash-attn": "2.8.1",
+        "flash-attn-4": "4.0.0.dev0",
+        "flash-attention-fp4": "0.1.0",
+        "flashinfer-python": "0.2.11",
+        "nvidia-cutlass-dsl": "4.5.0",
+    }
+
+    def fake_version(name):
+        if name not in versions:
+            raise identity_module.importlib.metadata.PackageNotFoundError(name)
+        return versions[name]
+
+    monkeypatch.setattr(identity_module.importlib.metadata, "version", fake_version)
+
+    installed = identity_module._installed_package_versions()
+
+    assert installed["fastvideo_kernel"] == "0.3.2"
+    assert installed["flash_attn"] == "2.8.1"
+    assert installed["flash_attn_4"] == "4.0.0.dev0"
+    assert installed["flash_attention_fp4"] == "0.1.0"
+    assert installed["flashinfer"] == "0.2.11"
+    assert installed["nvidia_cutlass_dsl"] == "4.5.0"
+
+
+def test_hardware_profile_id_uses_normalized_gpu_cohort():
+    profile = hardware_profile(
+        num_gpus=2,
+        gpu_devices=[
+            {
+                "name": "NVIDIA L40S",
+                "memory_bytes": 48 * 1024**3,
+                "compute_capability": "8.9",
+            },
+            {
+                "name": "NVIDIA L40S",
+                "memory_gb": 48,
+                "compute_capability": "8.9",
+            },
+        ],
+        interconnect="none_or_partial",
+    )
+
+    assert profile == {
+        "device_type": "cuda",
+        "gpu_count": 2,
+        "gpus": [
+            {
+                "name": "NVIDIA L40S",
+                "memory_gb": 48,
+                "compute_capability": "8.9",
+            },
+            {
+                "name": "NVIDIA L40S",
+                "memory_gb": 48,
+                "compute_capability": "8.9",
+            },
+        ],
+        "interconnect": "none_or_partial",
+    }
+    assert hardware_profile_id(profile).startswith("hw-")
+
+
+def test_hardware_profile_pads_missing_requested_cuda_devices(monkeypatch):
+    class Props:
+        name = "NVIDIA L40S"
+        total_memory = 48 * 1024**3
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device_id: Props())
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device_id: (8, 9))
+
+    profile = hardware_profile(num_gpus=2, interconnect="unknown")
+
+    assert profile["gpu_count"] == 2
+    assert profile["gpus"] == [
+        {
+            "name": "NVIDIA L40S",
+            "memory_gb": 48,
+            "compute_capability": "8.9",
+        },
+        {
+            "name": "unknown",
+            "memory_gb": None,
+            "compute_capability": None,
+        },
+    ]
+
+
+def test_local_path_detection_checks_both_path_separators(monkeypatch):
+    monkeypatch.setattr(identity_module.os, "sep", "\\")
+
+    assert identity_module._looks_like_local_path("models/local/checkpoint") is True
+    assert identity_module._looks_like_local_path(r"C:\models\checkpoint") is True
+    assert identity_module._looks_like_local_path("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") is False
+
+
+def test_environment_metadata_is_separate_audit_fingerprint():
+    cfg = _benchmark_config()
+    recipe_hash = _fingerprint(cfg)
+    audit_a = environment_metadata(
+        env={"IMAGE_VERSION": "py3.12-cuda13.0.0"},
+        package_versions={"triton": "3.4.1"},
+        hardware={"device_type": "cuda", "gpu_count": 2},
+        software={"python": "3.12", "pytorch": "2.12", "cuda": "13.0"},
+    )
+    audit_b = environment_metadata(
+        env={"IMAGE_VERSION": "py3.12-cuda13.0.1"},
+        package_versions={"triton": "3.4.1"},
+        hardware={"device_type": "cuda", "gpu_count": 2},
+        software={"python": "3.12", "pytorch": "2.12", "cuda": "13.0"},
+    )
+
+    assert recipe_hash == _fingerprint(cfg)
+    assert environment_fingerprint(audit_a) != environment_fingerprint(audit_b)
+
+
+def test_resolved_revision_from_hf_snapshot_path():
+    revision = "a" * 40
+
+    assert resolved_revision_from_model_path(
+        f"/root/.cache/huggingface/hub/models--org--repo/snapshots/{revision}") == revision
+    assert resolved_revision_from_model_path("/models/local-repo") is None
+
+
+def test_runtime_identity_from_generator_summarizes_worker_records():
+    from fastvideo.tests.performance.test_inference_performance import _runtime_identity_from_generator
+
+    revision = "b" * 40
+
+    class FakeExecutor:
+        def collective_rpc(self, _method):
+            return [
+                {
+                    "resolved_attention_backends": ["FLASH_ATTN"],
+                    "resolved_model_path": f"/root/.cache/huggingface/hub/models--org--repo/snapshots/{revision}",
+                },
+                {
+                    "resolved_attention_backends": ["FLASH_ATTN"],
+                    "resolved_model_path": f"/root/.cache/huggingface/hub/models--org--repo/snapshots/{revision}",
+                },
+            ]
+
+    class FakeGenerator:
+        executor = FakeExecutor()
+
+    assert _runtime_identity_from_generator(FakeGenerator()) == {
+        "resolved_attention_backend": "FLASH_ATTN",
+        "resolved_model_revision": revision,
+    }
+
+
+def test_collect_worker_identity_handles_torch_module_pipeline():
+    from fastvideo.tests.performance.test_inference_performance import _collect_worker_identity
+
+    class Backend:
+        name = "FLASH_ATTN"
+
+    class Leaf(torch.nn.Module):
+        backend = Backend()
+
+    class Pipeline(torch.nn.Module):
+        model_path = "/models/local"
+
+        def __init__(self):
+            super().__init__()
+            self.leaf = Leaf()
+
+    class Worker:
+        pipeline = Pipeline()
+
+    assert _collect_worker_identity(Worker()) == {
+        "resolved_attention_backends": ["FLASH_ATTN"],
+        "resolved_model_path": "/models/local",
+    }

fastvideo/tests/performance/test_inference_performance.py

@@ -12,12 +12,25 @@ import os
 import time
 from collections.abc import Mapping
 from datetime import datetime, timezone
+from typing import Any
 
 import torch
 import pytest
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
+from fastvideo.tests.performance.identity import (
+    benchmark_identity_from_config,
+    build_recipe_from_benchmark_config,
+    environment_fingerprint,
+    environment_metadata,
+    hardware_profile,
+    hardware_profile_id,
+    recipe_fingerprint,
+    resolved_revision_from_model_path,
+    software_profile,
+    software_profile_id,
+)
 from fastvideo.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)
@@ -34,8 +47,11 @@ V2_REQUIRED_IDENTITY_FIELDS = (
     "variant_id",
     "benchmark_version",
 )
+# "recipe" is no longer config-declarable: the fingerprint follow-up landed,
+# and the generated recipe document owns that key in emitted records. A config
+# declaring it now fails validation loudly instead of being silently
+# overwritten by the generated one.
 V2_OPTIONAL_METADATA_FIELDS = (
-    "recipe",
     "metric_threshold_policy",
     "quality_metadata",
 )
@@ -179,7 +195,11 @@ def _extract_component_times(result: dict) -> dict[str, float | None]:
             logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
             continue
         stage_class = stage_data.get("stage_class", stage_name)
-        metric_key = STAGE_METRIC_MAP.get(stage_class)
+        component_metric = stage_data.get("component_metric")
+        if isinstance(component_metric, str) and component_metric in component_times:
+            metric_key = component_metric
+        else:
+            metric_key = STAGE_METRIC_MAP.get(stage_class)
         if metric_key is None:
             logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
                          stage_name,
@@ -227,6 +247,110 @@ def _write_results(results):
     logger.info("Performance results written to %s", filepath)
 
 
+def _backend_name(value) -> str:
+    if hasattr(value, "name"):
+        return str(value.name)
+    return str(value)
+
+
+def _collect_worker_identity(worker) -> dict[str, Any]:
+    pipeline = getattr(worker, "pipeline", None)
+    model_path = getattr(pipeline, "model_path", None)
+    backends: set[str] = set()
+
+    modules = getattr(pipeline, "modules", None)
+    if isinstance(modules, Mapping):
+        module_iter = modules.values()
+    elif isinstance(pipeline, torch.nn.Module):
+        module_iter = (pipeline,)
+    else:
+        module_iter = ()
+
+    for module in module_iter:
+        if isinstance(module, torch.nn.Module):
+            for submodule in module.modules():
+                backend = getattr(submodule, "backend", None)
+                if backend is not None:
+                    backends.add(_backend_name(backend))
+        else:
+            backend = getattr(module, "backend", None)
+            if backend is not None:
+                backends.add(_backend_name(backend))
+
+    return {
+        "resolved_attention_backends": sorted(backends),
+        "resolved_model_path": model_path,
+    }
+
+
+def _single_or_list(values: set[str]) -> str | list[str] | None:
+    ordered = sorted(values)
+    if not ordered:
+        return None
+    if len(ordered) == 1:
+        return ordered[0]
+    return ordered
+
+
+def _runtime_identity_from_generator(generator) -> dict[str, Any]:
+    worker_records = generator.executor.collective_rpc(_collect_worker_identity)
+    resolved_backends: set[str] = set()
+    resolved_revisions: set[str] = set()
+
+    for record in worker_records:
+        resolved_backends.update(record.get("resolved_attention_backends") or [])
+        revision = resolved_revision_from_model_path(record.get("resolved_model_path"))
+        if revision is not None:
+            resolved_revisions.add(revision)
+
+    return {
+        "resolved_attention_backend": _single_or_list(resolved_backends),
+        "resolved_model_revision": _single_or_list(resolved_revisions),
+    }
+
+
+def _benchmark_identity_fields(cfg):
+    identity = benchmark_identity_from_config(cfg)
+    return {
+        key: identity[key]
+        for key in ("workload_id", "variant_id", "benchmark_version")
+    }
+
+
+def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity):
+    # Legacy v1 configs (no config_schema_version) stay on the legacy
+    # (model_id, gpu_type) cohort: they lack the required identity fields,
+    # and building a recipe for them would raise AFTER the GPU measurement
+    # completed. The validator already enforces that v2 configs carry the
+    # identity fields, so v2 records always get the full identity block.
+    if cfg.get("config_schema_version") is None:
+        return {}
+    recipe = build_recipe_from_benchmark_config(
+        cfg,
+        resolved_attention_backend=runtime_identity.get("resolved_attention_backend"),
+        resolved_model_revision=runtime_identity.get("resolved_model_revision"),
+        measured_prompts=[prompt],
+    )
+    num_gpus = init_kwargs.get("num_gpus", cfg.get("run_config", {}).get("required_gpus", 1))
+    hw_profile = hardware_profile(num_gpus=num_gpus)
+    sw_profile = software_profile()
+    env_metadata = environment_metadata(
+        hardware=hw_profile,
+        software=sw_profile,
+    )
+    return {
+        "recipe": recipe,
+        "recipe_fingerprint": recipe_fingerprint(recipe),
+        "hardware_profile": hw_profile,
+        "hardware_profile_id": hardware_profile_id(hw_profile),
+        "software_profile": sw_profile,
+        "software_profile_id": software_profile_id(sw_profile),
+        "environment_metadata": env_metadata,
+        "environment_fingerprint": environment_fingerprint(env_metadata),
+        **_benchmark_identity_fields(cfg),
+    }
+
+
 # -- Test -------------------------------------------------------------------
 
 def _run_benchmark(cfg):
@@ -264,6 +388,7 @@ def _run_benchmark(cfg):
             model_path=model_info["model_path"],
             **init_kwargs,
         )
+        runtime_identity = _runtime_identity_from_generator(generator)
 
         for i in range(num_warmup):
             logger.info("Warmup run %d/%d", i + 1, num_warmup)
@@ -318,6 +443,7 @@ def _run_benchmark(cfg):
         "dit_time_s": _avg_component(all_component_times, "dit_time_s"),
         "vae_decode_time_s": _avg_component(all_component_times,
                                             "vae_decode_time_s"),
+        **_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity),
     }
 
     logger.info(

fastvideo/tests/performance/test_inference_performance_component_times.py

@@ -1,13 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from fastvideo.pipelines.pipeline_batch_info import PipelineLoggingInfo
+from fastvideo.pipelines.stages.denoising import Cosmos25AutoDenoisingStage, DenoisingStage
+from fastvideo.pipelines.stages.text_encoding import Cosmos25TextEncodingStage
 from fastvideo.tests.performance.test_inference_performance import _extract_component_times
 
 
+class SubclassStyleDenoisingStage(DenoisingStage):
+    pass
+
+
 def test_extract_component_times_handles_pipeline_logging_info_object():
     logging_info = PipelineLoggingInfo()
     logging_info.add_stage_execution_time("prompt_encoding_stage", 1.25)
     logging_info.add_stage_metric("prompt_encoding_stage", "stage_class", "TextEncodingStage")
+    logging_info.add_stage_metric("prompt_encoding_stage", "component_metric", "text_encoder_time_s")
 
     assert _extract_component_times({"logging_info": logging_info}) == {
         "text_encoder_time_s": 1.25,
@@ -45,6 +52,35 @@ def test_extract_component_times_uses_stage_class_for_pipeline_stage_keys():
     }
 
 
+def test_denoising_stage_subclasses_inherit_component_metric():
+    assert SubclassStyleDenoisingStage.performance_component_metric == "dit_time_s"
+
+
+def test_cosmos25_direct_pipeline_stages_define_component_metrics():
+    assert Cosmos25TextEncodingStage.performance_component_metric == "text_encoder_time_s"
+    assert Cosmos25AutoDenoisingStage.performance_component_metric == "dit_time_s"
+
+
+def test_extract_component_times_uses_component_metric_for_stage_subclasses():
+    result = {
+        "logging_info": {
+            "stages": {
+                "denoising_stage": {
+                    "execution_time": 4.2,
+                    "stage_class": "CosmosDenoisingStage",
+                    "component_metric": "dit_time_s",
+                },
+            },
+        },
+    }
+
+    assert _extract_component_times(result) == {
+        "text_encoder_time_s": None,
+        "dit_time_s": 4.2,
+        "vae_decode_time_s": None,
+    }
+
+
 def test_extract_component_times_keeps_legacy_class_name_keys():
     # Backward-compatibility check for logs produced before pipeline-unique
     # stage keys carried a separate stage_class field.

fastvideo/tests/ssim/inference_similarity_utils.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 import os
+from collections.abc import Iterator
 from contextlib import contextmanager
 from logging import Logger
-from typing import Iterator
 
 from fastvideo import VideoGenerator
 from fastvideo.tests.ssim.reference_utils import (
@@ -67,6 +67,12 @@ def _find_reference_video(reference_folder: str, prompt: str) -> str:
     raise FileNotFoundError("Reference video missing")
 
 
+def _remove_stale_generated_video(output_dir: str, output_video_name: str) -> None:
+    stale_path = os.path.join(output_dir, output_video_name)
+    if os.path.exists(stale_path):
+        os.remove(stale_path)
+
+
 def _assert_similarity(
     *,
     logger: Logger,
@@ -214,6 +220,7 @@ def run_text_to_video_similarity_test(
         )
         output_video_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
+        _remove_stale_generated_video(output_dir, output_video_name)
 
         params_map = select_ssim_params(
             default_params_map,
@@ -289,6 +296,7 @@ def run_image_to_video_similarity_test(
         )
         output_video_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
+        _remove_stale_generated_video(output_dir, output_video_name)
 
         params_map = select_ssim_params(
             default_params_map,

fastvideo/tests/ssim/test_dreamx_world_similarity.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+from PIL import Image, ImageDraw
+
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+    run_image_to_video_similarity_test,
+    run_text_to_video_similarity_test,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 1
+
+device_reference_folder = resolve_inference_device_reference_folder(logger)
+
+_MODEL_PATH = os.getenv(
+    "DREAMX_WORLD_SSIM_MODEL_PATH",
+    "FastVideo/DreamX-World-5B-Cam-Diffusers",
+)
+_AR_MODEL_PATH = os.getenv(
+    "DREAMX_WORLD_AR_SSIM_MODEL_PATH",
+    "FastVideo/DreamX-World-5B-Diffusers",
+)
+
+DREAMX_WORLD_PARAMS = {
+    "num_gpus": 1,
+    "model_path": _MODEL_PATH,
+    "height": 64,
+    "width": 64,
+    "num_frames": 9,
+    "num_inference_steps": 1,
+    "guidance_scale": 1.0,
+    "seed": 1024,
+    "fps": 16,
+}
+
+DREAMX_WORLD_FULL_QUALITY_PARAMS = {
+    **DREAMX_WORLD_PARAMS,
+    "height": 480,
+    "width": 832,
+    "num_frames": 161,
+    "num_inference_steps": 30,
+    "guidance_scale": 5.0,
+}
+
+
+DREAMX_WORLD_AR_PARAMS = {
+    "num_gpus": 1,
+    "model_path": _AR_MODEL_PATH,
+    "height": 192,
+    "width": 192,
+    "num_frames": 81,
+    "num_inference_steps": 4,
+    "guidance_scale": 1.0,
+    "seed": 2048,
+    "fps": 16,
+}
+
+DREAMX_WORLD_AR_FULL_QUALITY_PARAMS = {
+    **DREAMX_WORLD_AR_PARAMS,
+    "height": 704,
+    "width": 1280,
+    "num_frames": 1005,
+}
+
+DREAMX_WORLD_MODEL_TO_PARAMS = {
+    "DreamX-World-5B-Cam": DREAMX_WORLD_PARAMS,
+}
+
+DREAMX_WORLD_AR_MODEL_TO_PARAMS = {
+    "DreamX-World-5B": DREAMX_WORLD_AR_PARAMS,
+}
+
+FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS = {
+    "DreamX-World-5B-Cam": DREAMX_WORLD_FULL_QUALITY_PARAMS,
+}
+
+FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS = {
+    "DreamX-World-5B": DREAMX_WORLD_AR_FULL_QUALITY_PARAMS,
+}
+
+DREAMX_WORLD_TEST_CASES = [
+    (
+        "A cinematic first-person drive through a futuristic coastal city at sunrise, "
+        "reflective glass towers, clean streets, soft volumetric light.",
+        ("w", "d", "w"),
+        (4.0, 2.0, 4.0),
+    ),
+]
+
+
+DREAMX_WORLD_AR_TEST_CASES = [
+    (
+        "A long autonomous drive through a futuristic coastal city at sunrise, "
+        "smooth forward camera motion, reflective glass towers, clean streets.",
+        ("w", "d", "w", "a"),
+        (2.0, 1.0, 2.0, 1.0),
+    ),
+]
+
+
+def _write_deterministic_reference_image(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    image = Image.new("RGB", (96, 96), color=(42, 76, 112))
+    draw = ImageDraw.Draw(image)
+    draw.rectangle((0, 56, 96, 96), fill=(36, 44, 52))
+    draw.polygon([(0, 56), (48, 30), (96, 56)], fill=(120, 142, 158))
+    draw.rectangle((34, 42, 62, 72), fill=(178, 198, 212))
+    draw.line((0, 80, 96, 66), fill=(238, 209, 124), width=3)
+    image.save(path)
+
+
+@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_TEST_CASES)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_MODEL_TO_PARAMS.keys()))
+def test_dreamx_world_inference_similarity(
+    prompt: str,
+    action_list: tuple[str, ...],
+    action_speed_list: tuple[float, ...],
+    attention_backend_name: str,
+    model_id: str,
+    tmp_path: Path,
+) -> None:
+    image_path = tmp_path / "dreamx_world_ssim_input.png"
+    _write_deterministic_reference_image(image_path)
+
+    run_image_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        image_path=str(image_path),
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=DREAMX_WORLD_MODEL_TO_PARAMS,
+        full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.98,
+        init_kwargs_override={
+            "use_fsdp_inference": False,
+            "dit_cpu_offload": False,
+            "vae_cpu_offload": True,
+            "text_encoder_cpu_offload": True,
+            "pin_cpu_memory": False,
+            "override_pipeline_cls_name": "DreamXWorldPipeline",
+        },
+        generation_kwargs_override={
+            "action_list": list(action_list),
+            "action_speed_list": list(action_speed_list),
+        },
+    )
+
+@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_AR_TEST_CASES)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_AR_MODEL_TO_PARAMS.keys()))
+def test_dreamx_world_ar_inference_similarity(
+    prompt: str,
+    action_list: tuple[str, ...],
+    action_speed_list: tuple[float, ...],
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=DREAMX_WORLD_AR_MODEL_TO_PARAMS,
+        full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.98,
+        init_kwargs_override={
+            "use_fsdp_inference": False,
+            "dit_cpu_offload": False,
+            "vae_cpu_offload": True,
+            "text_encoder_cpu_offload": True,
+            "pin_cpu_memory": False,
+            "override_pipeline_cls_name": "DreamXWorldARPipeline",
+        },
+        generation_kwargs_override={
+            "action_list": list(action_list),
+            "action_speed_list": list(action_speed_list),
+        },
+    )

fastvideo/tests/ssim/test_kandinsky5_similarity.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SSIM regression test for Kandinsky-5.0 Lite text-to-video.
+
+Runs a reduced-size generation (256x384, 21 frames) with a fixed seed and
+compares against a device-specific reference video via MS-SSIM. The preset
+default is 512x768 x 121 frames; the reduced size keeps CI cheap while the
+patch (1, 2, 2) and VAE temporal-4 / spatial-8 constraints stay satisfied
+(dims divisible by 16, num_frames % 4 == 1).
+"""
+import os
+
+import pytest
+
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.kandinsky5.presets import KANDINSKY5_T2V_LITE_5S
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+    run_text_to_video_similarity_test,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 1
+
+device_reference_folder = resolve_inference_device_reference_folder(logger)
+
+_PRESET_DEFAULTS = KANDINSKY5_T2V_LITE_5S.defaults
+
+KANDINSKY5_T2V_PARAMS = {
+    "num_gpus": 1,
+    "model_path": "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
+    "height": 256,
+    "width": 384,
+    "num_frames": 21,
+    "num_inference_steps": _PRESET_DEFAULTS["num_inference_steps"],
+    "guidance_scale": _PRESET_DEFAULTS["guidance_scale"],
+    "seed": 1024,
+    "sp_size": 1,
+    "tp_size": 1,
+    "fps": _PRESET_DEFAULTS["fps"],
+    "neg_prompt": _PRESET_DEFAULTS["negative_prompt"],
+}
+
+KANDINSKY5_T2V_FULL_QUALITY_PARAMS = {
+    **KANDINSKY5_T2V_PARAMS,
+    "height": _PRESET_DEFAULTS["height"],
+    "width": _PRESET_DEFAULTS["width"],
+    # num_frames stays at the inherited reduced 21 even at full quality
+    # (preset default: 121).
+}
+
+KANDINSKY5_T2V_MODEL_TO_PARAMS = {
+    "Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers": KANDINSKY5_T2V_PARAMS,
+}
+FULL_QUALITY_KANDINSKY5_T2V_MODEL_TO_PARAMS = {
+    "Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers": KANDINSKY5_T2V_FULL_QUALITY_PARAMS,
+}
+
+KANDINSKY5_T2V_TEST_PROMPTS = [
+    "A curious raccoon peers through a vibrant field of yellow sunflowers, its "
+    "eyes wide with interest. The playful yet serene atmosphere is complemented "
+    "by soft natural light filtering through the petals. Mid-shot, warm and "
+    "cheerful tones.",
+]
+
+
+@pytest.mark.parametrize("prompt", KANDINSKY5_T2V_TEST_PROMPTS)
+@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
+@pytest.mark.parametrize("model_id", list(KANDINSKY5_T2V_MODEL_TO_PARAMS.keys()))
+def test_kandinsky5_t2v_inference_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    # The SSIM lane image bakes FASTVIDEO_FA4=1, but on L40S (sm89) the CLIP
+    # encoder's dense MHA inference call routes into the FA4 CuTeDSL JIT,
+    # which fails to compile there (nvvm.fmax signature mismatch vs the
+    # image's nvidia_cutlass_dsl). Pin FA4 off so reference seeding and CI
+    # runs both use the FA2 path with identical numerics. Scoped to this test
+    # and restored afterwards: the other SSIM references are seeded with FA4
+    # on, so a module-level pin would corrupt every test collected in the
+    # same pytest process.
+    saved_fa4 = os.environ.get("FASTVIDEO_FA4")
+    os.environ["FASTVIDEO_FA4"] = "0"
+    try:
+        run_text_to_video_similarity_test(
+            logger=logger,
+            script_dir=os.path.dirname(os.path.abspath(__file__)),
+            device_reference_folder=device_reference_folder,
+            prompt=prompt,
+            attention_backend_name=attention_backend_name,
+            model_id=model_id,
+            default_params_map=KANDINSKY5_T2V_MODEL_TO_PARAMS,
+            full_quality_params_map=FULL_QUALITY_KANDINSKY5_T2V_MODEL_TO_PARAMS,
+            min_acceptable_ssim=0.93,
+            # Match examples/inference/basic/basic_kandinsky5_t2v.py: FSDP
+            # inference is not used for Kandinsky-5, and the Qwen2.5-VL text
+            # encoder stays on CPU between encodes.
+            init_kwargs_override={
+                "use_fsdp_inference": False,
+                "text_encoder_cpu_offload": True,
+                "pin_cpu_memory": True,
+            },
+        )
+    finally:
+        if saved_fa4 is None:
+            os.environ.pop("FASTVIDEO_FA4", None)
+        else:
+            os.environ["FASTVIDEO_FA4"] = saved_fa4

fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# Minimum config to run a single training step of
+# CausalConsistencyDistillationMethod on WanCausalModel for the
+# per-method smoke test. Uses the real Wan 2.1 1.3B checkpoint for both
+# the trainable student and the frozen teacher (AR Euler-step target),
+# with tiny synthetic latents.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+  teacher:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+  ema:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+
+method:
+  _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
+  discrete_cd_N: 12
+  guidance_scale: 3.0
+  ema_decay: 0.95
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 6
+    num_height: 64
+    num_width: 64
+    num_frames: 21
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+pipeline: {}

fastvideo/tests/train/fixtures/wan_causal_t2v_dfsft_framewise_min.yaml

@@ -0,0 +1,46 @@
+# Minimum config to run a single training step of
+# DiffusionForcingSFTMethod on a frame-wise WanCausalModel
+# (num_frames_per_block=1, chunk_size=1) for the per-method smoke
+# test.  Uses the real Wan 2.1 1.3B checkpoint with tiny synthetic
+# latents.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    num_frames_per_block: 1
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
+  chunk_size: 1
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 6
+    num_height: 64
+    num_width: 64
+    num_frames: 21
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+pipeline: {}

fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# Minimum config to run a single training step of
+# TeacherForcingSFTMethod on WanCausalModel for the per-method smoke
+# test. Identical to wan_causal_t2v_dfsft_min.yaml except the method,
+# which feeds clean history to the causal transformer (teacher forcing).
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
+  chunk_size: 3
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 6
+    num_height: 64
+    num_width: 64
+    num_frames: 21
+
+  optimizer:
+    learning_rate: 2.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+pipeline: {}

fastvideo/tests/train/methods/test_wan_causal_cd.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-method GPU smoke test: ``WanCausalModel`` + ``CausalConsistencyDistillationMethod``.
+
+Mirrors ``test_wan_causal_dfsft.py``. Causal consistency distillation
+bootstraps a consistency MSE between the student's ``x0`` at ``t`` and an EMA
+copy of the student at ``t_next``, where ``t_next`` is produced online by a
+single CFG Euler step of a frozen teacher (all under clean-history teacher
+forcing). This test exercises the full step: finite loss, nonzero student
+gradients, frozen teacher, and a post-step EMA update.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from fastvideo.train.methods.consistency_model.causal_cd import (
+    CausalConsistencyDistillationMethod, )
+from fastvideo.train.models.wan import WanCausalModel
+from fastvideo.train.utils.config import load_run_config
+
+
+_FIXTURE = str(
+    Path(__file__).resolve().parent.parent / "fixtures"
+    / "wan_causal_t2v_causal_cd_min.yaml")
+
+
+def _build_synthetic_batch(
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    batch_size = 1
+    return {
+        "text_embedding":
+        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
+        "text_attention_mask":
+        torch.ones(batch_size, 16, device=device),
+        "vae_latent":
+        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
+    }
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_wan_causal_cd_single_train_step(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+
+    cfg = load_run_config(_FIXTURE)
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+
+    monkeypatch.setattr(
+        "fastvideo.train.utils.dataloader."
+        "build_parquet_t2v_train_dataloader",
+        lambda *args, **kwargs: None,
+    )
+
+    student = WanCausalModel(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+    )
+    student.transformer = student.transformer.to(device=device, dtype=dtype)
+
+    teacher = WanCausalModel(
+        init_from=cfg.models["teacher"]["init_from"],
+        training_config=cfg.training,
+        trainable=False,
+    )
+    teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)
+
+    ema = WanCausalModel(
+        init_from=cfg.models["ema"]["init_from"],
+        training_config=cfg.training,
+        trainable=False,
+    )
+    ema.transformer = ema.transformer.to(device=device, dtype=dtype)
+
+    method = CausalConsistencyDistillationMethod(
+        cfg=cfg,
+        role_models={"student": student, "teacher": teacher, "ema": ema},
+    )
+    method.on_train_start()
+
+    batch = _build_synthetic_batch(device, dtype)
+    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
+
+    loss = loss_map["total_loss"]
+    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
+    assert torch.isfinite(loss).item(), (
+        f"total_loss is not finite: {loss.item()}")
+
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+
+    blocks = student.transformer.blocks
+    assert blocks is not None and len(blocks) > 0
+    layer0 = blocks[0]
+    trainable = [p for p in layer0.parameters() if p.requires_grad]
+    assert len(trainable) > 0, "student layer 0 has no trainable parameters"
+    for i, p in enumerate(trainable):
+        assert p.grad is not None, f"student layer 0 param[{i}] has None grad"
+        assert torch.isfinite(p.grad).all().item(), (
+            f"student layer 0 param[{i}] grad contains NaN/Inf")
+    assert any(p.grad.detach().float().norm().item() > 0.0 for p in trainable), (
+        "all student layer-0 grads are exactly zero; consistency loss "
+        "did not reach the first transformer block")
+
+    # Teacher must stay frozen.
+    assert all(not p.requires_grad for p in teacher.transformer.parameters()), (
+        "teacher must be frozen for Causal-CD")
+
+    # The EMA model and student start from the same checkpoint, so the first
+    # parameter must match before any update. FSDP fully_shard params are
+    # DTensors; compare the local shards (torch.equal is unsupported on
+    # DTensor).
+    def _local(p: torch.Tensor) -> torch.Tensor:
+        return p.to_local() if hasattr(p, "to_local") else p
+
+    ema_param = next(ema.transformer.parameters())
+    student_param = next(student.transformer.parameters())
+    assert torch.equal(_local(ema_param), _local(student_param)), (
+        "EMA model should start identical to the student (same checkpoint)")
+
+    # The EMA update must move EMA toward the student. Apply a visibly large
+    # perturbation so the bf16 lerp is well above rounding noise (the real
+    # optimizer step at lr=2e-6 would be sub-ULP in bf16).
+    with torch.no_grad():
+        student_param.add_(1.0)
+    before = _local(ema_param).detach().float().clone()
+    method._update_ema()
+    after = _local(ema_param).detach().float()
+    assert not torch.equal(before, after), (
+        "EMA weights did not move after _update_ema")
+    # EMA = decay*ema + (1-decay)*student moves ~ (1-decay) of the gap.
+    expected = before + (1.0 - method._ema_decay) * (
+        _local(student_param).detach().float() - before)
+    assert torch.allclose(after, expected, atol=1e-2), (
+        "EMA update did not follow the expected lerp")

fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-method GPU smoke test: frame-wise ``WanCausalModel`` + ``DiffusionForcingSFTMethod``.
+
+Mirrors ``test_wan_causal_dfsft.py`` but with a block size of 1 frame
+(``num_frames_per_block=1`` on the model, ``chunk_size=1`` on the method),
+so each frame gets its own independent noise level. The test asserts the
+override took effect and runs one train step: forward, finite loss, and
+nonzero gradients reaching the first transformer block.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29520")
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from fastvideo.train.methods.fine_tuning.dfsft import (
+    DiffusionForcingSFTMethod, )
+from fastvideo.train.models.wan import WanCausalModel
+from fastvideo.train.utils.config import load_run_config
+
+
+_FIXTURE = str(
+    Path(__file__).resolve().parent.parent / "fixtures"
+    / "wan_causal_t2v_dfsft_framewise_min.yaml")
+
+
+def _build_synthetic_batch(
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    batch_size = 1
+    return {
+        "text_embedding":
+        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
+        "text_attention_mask":
+        torch.ones(batch_size, 16, device=device),
+        "vae_latent":
+        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
+    }
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_wan_causal_dfsft_framewise_single_train_step(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+
+    cfg = load_run_config(_FIXTURE)
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+
+    monkeypatch.setattr(
+        "fastvideo.train.utils.dataloader."
+        "build_parquet_t2v_train_dataloader",
+        lambda *args, **kwargs: None,
+    )
+
+    student_cfg = cfg.models["student"]
+    model = WanCausalModel(
+        init_from=student_cfg["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+        num_frames_per_block=student_cfg.get("num_frames_per_block"),
+    )
+    assert model.transformer.num_frame_per_block == 1, (
+        "frame-wise override did not reach the transformer")
+    model.transformer = model.transformer.to(device=device, dtype=dtype)
+
+    method = DiffusionForcingSFTMethod(
+        cfg=cfg,
+        role_models={"student": model},
+    )
+    method.on_train_start()
+
+    batch = _build_synthetic_batch(device, dtype)
+    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
+
+    loss = loss_map["total_loss"]
+    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
+    assert torch.isfinite(loss).item(), (
+        f"total_loss is not finite: {loss.item()}")
+
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+
+    blocks = getattr(model.transformer, "blocks", None)
+    assert blocks is not None and len(blocks) > 0
+    layer0 = blocks[0]
+
+    trainable = [p for p in layer0.parameters() if p.requires_grad]
+    assert len(trainable) > 0, "layer 0 has no trainable parameters"
+
+    for i, p in enumerate(trainable):
+        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
+        assert torch.isfinite(p.grad).all().item(), (
+            f"layer 0 param[{i}] grad contains NaN/Inf")
+
+    any_nonzero = any(
+        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
+    assert any_nonzero, (
+        "all layer-0 grads are exactly zero; backward did not "
+        "reach the first transformer block")

fastvideo/tests/train/methods/test_wan_causal_tfsft.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-method GPU smoke test: ``WanCausalModel`` + ``TeacherForcingSFTMethod``.
+
+Mirrors ``test_wan_causal_dfsft.py``. Teacher forcing concatenates a clean
+context copy of every frame inside the causal transformer (``clean_x``) and
+denoises the current block while attending to *clean* history. This test
+exercises the ``clean_x`` path end-to-end: forward, finite loss, and nonzero
+gradients reaching the first transformer block.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29518")
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from fastvideo.train.methods.fine_tuning.tfsft import (
+    TeacherForcingSFTMethod, )
+from fastvideo.train.models.wan import WanCausalModel
+from fastvideo.train.utils.config import load_run_config
+
+
+_FIXTURE = str(
+    Path(__file__).resolve().parent.parent / "fixtures"
+    / "wan_causal_t2v_tfsft_min.yaml")
+
+
+def _build_synthetic_batch(
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    batch_size = 1
+    return {
+        "text_embedding":
+        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
+        "text_attention_mask":
+        torch.ones(batch_size, 16, device=device),
+        "vae_latent":
+        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
+    }
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_wan_causal_tfsft_single_train_step(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+
+    cfg = load_run_config(_FIXTURE)
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+
+    monkeypatch.setattr(
+        "fastvideo.train.utils.dataloader."
+        "build_parquet_t2v_train_dataloader",
+        lambda *args, **kwargs: None,
+    )
+
+    model = WanCausalModel(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+    )
+    model.transformer = model.transformer.to(device=device, dtype=dtype)
+
+    method = TeacherForcingSFTMethod(
+        cfg=cfg,
+        role_models={"student": model},
+    )
+    method.on_train_start()
+
+    batch = _build_synthetic_batch(device, dtype)
+    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
+
+    loss = loss_map["total_loss"]
+    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
+    assert torch.isfinite(loss).item(), (
+        f"total_loss is not finite: {loss.item()}")
+
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+
+    blocks = getattr(model.transformer, "blocks", None)
+    assert blocks is not None and len(blocks) > 0, (
+        "CausalWanTransformer is expected to expose ``.blocks``")
+    layer0 = blocks[0]
+
+    trainable = [p for p in layer0.parameters() if p.requires_grad]
+    assert len(trainable) > 0, "layer 0 has no trainable parameters"
+
+    for i, p in enumerate(trainable):
+        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
+        assert torch.isfinite(p.grad).all().item(), (
+            f"layer 0 param[{i}] grad contains NaN/Inf")
+
+    any_nonzero = any(
+        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
+    assert any_nonzero, (
+        "all layer-0 grads are exactly zero; backward did not "
+        "reach the first transformer block")
+
+    # Teacher forcing must build its own (concatenated) attention mask and
+    # must not have constructed the diffusion-forcing mask.
+    assert model.transformer.teacher_forcing_block_mask is not None, (
+        "teacher-forcing mask was not constructed")
+    assert model.transformer.block_mask is None, (
+        "diffusion-forcing mask should not be built on the TF path")

fastvideo/tests/training/test_trackers.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import builtins
+from io import BytesIO
+from pathlib import Path
+import sys
+from types import ModuleType
+from typing import Any
+
+import numpy as np
+from PIL import Image
+import pytest
+
+from fastvideo.training import trackers as trackers_module
+from fastvideo.training.trackers import (
+    _coerce_video_fps,
+    _prepare_video_array,
+    BaseTracker,
+    SequentialTracker,
+    SwanlabTracker,
+    WandbTracker,
+)
+
+
+class _FakeSwanlab(ModuleType):
+    def __init__(self) -> None:
+        super().__init__("swanlab")
+        self.init_kwargs: dict[str, Any] | None = None
+        self.logged: list[tuple[dict[str, Any], int]] = []
+        self.video_paths: list[str] = []
+        self.finished = False
+
+    def init(self, **kwargs: Any) -> object:
+        self.init_kwargs = kwargs
+        return object()
+
+    def log(self, metrics: dict[str, Any], step: int) -> None:
+        self.logged.append((metrics, step))
+
+    def finish(self) -> None:
+        self.finished = True
+
+    def Video(self, file_path: str, caption: str | None = None) -> dict[str, Any]:  # noqa: N802
+        self.video_paths.append(file_path)
+        return {
+            "content": Path(file_path).read_bytes(),
+            "caption": caption,
+        }
+
+
+def _install_fake_swanlab(monkeypatch: pytest.MonkeyPatch) -> _FakeSwanlab:
+    swanlab = _FakeSwanlab()
+    monkeypatch.setitem(sys.modules, "swanlab", swanlab)
+    return swanlab
+
+
+def test_prepare_video_array_converts_tchw_to_thwc() -> None:
+    video = np.arange(2 * 3 * 2 * 4, dtype=np.uint8).reshape(2, 3, 2, 4)
+
+    prepared = _prepare_video_array(video)
+
+    np.testing.assert_array_equal(prepared, video.transpose(0, 2, 3, 1))
+
+
+def test_default_video_fps_is_16() -> None:
+    assert _coerce_video_fps(None) == 16
+
+
+def test_wandb_video_uses_shared_default_fps(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    wandb = ModuleType("wandb")
+    video_calls: list[dict[str, Any]] = []
+    wandb.init = lambda **_: object()
+
+    def video(data: Any, **kwargs: Any) -> dict[str, Any]:
+        video_calls.append(kwargs)
+        return {"data": data, **kwargs}
+
+    wandb.Video = video
+    monkeypatch.setitem(sys.modules, "wandb", wandb)
+    tracker = WandbTracker("project", str(tmp_path))
+
+    tracker.video(np.zeros((2, 3, 2, 2), dtype=np.uint8))
+
+    assert video_calls == [{"fps": 16, "format": "mp4"}]
+
+
+def test_prepare_video_array_tiles_and_pads_batches() -> None:
+    video = np.zeros((3, 1, 3, 1, 1), dtype=np.uint8)
+    video[0] = 1
+    video[1] = 2
+    video[2] = 3
+
+    prepared = _prepare_video_array(video)
+
+    assert prepared.shape == (1, 1, 4, 3)
+    np.testing.assert_array_equal(
+        prepared[0, 0],
+        np.array(
+            [
+                [1, 1, 1],
+                [2, 2, 2],
+                [3, 3, 3],
+                [0, 0, 0],
+            ],
+            dtype=np.uint8,
+        ),
+    )
+
+
+def test_swanlab_video_converts_array_to_gif_and_logs_artifact(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    swanlab = _install_fake_swanlab(monkeypatch)
+    tracker = SwanlabTracker("project", str(tmp_path), run_name="run")
+    video = np.zeros((2, 3, 2, 2), dtype=np.uint8)
+    video[0, 0] = 255
+    video[1, 2] = 255
+
+    artifact = tracker.video(video, caption="validation sample", fps=5, format="mp4")
+    tracker.log_artifacts({"validation/videos": [artifact]}, step=7)
+
+    assert artifact["content"].startswith((b"GIF87a", b"GIF89a"))
+    assert artifact["caption"] == "validation sample"
+    assert not Path(swanlab.video_paths[0]).exists()
+    with Image.open(BytesIO(artifact["content"])) as image:
+        assert image.n_frames == 2
+        assert image.info["duration"] == 200
+    assert swanlab.logged == [({"validation/videos": [artifact]}, 7)]
+
+
+@pytest.mark.parametrize(("requested_fps", "expected_fps"), [(None, 12.0), (2, 2.0)])
+def test_swanlab_video_uses_requested_or_source_fps_during_gif_conversion(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    requested_fps: int | None,
+    expected_fps: float,
+) -> None:
+    _install_fake_swanlab(monkeypatch)
+    tracker = SwanlabTracker("project", str(tmp_path))
+    frames = [np.zeros((2, 2, 3), dtype=np.uint8), np.ones((2, 2, 3), dtype=np.uint8)]
+    written_fps: list[float] = []
+    write_gif = trackers_module._write_gif
+
+    monkeypatch.setattr(trackers_module, "_read_video_file", lambda _: (frames, 12.0))
+
+    def capture_fps(file_path: str, video_frames: Any, fps: float) -> None:
+        written_fps.append(fps)
+        write_gif(file_path, video_frames, fps)
+
+    monkeypatch.setattr(trackers_module, "_write_gif", capture_fps)
+
+    artifact = tracker.video(tmp_path / "validation.mp4", fps=requested_fps)
+
+    assert artifact["content"].startswith((b"GIF87a", b"GIF89a"))
+    assert written_fps == [expected_fps]
+
+
+def test_swanlab_video_passes_existing_gif_through(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    swanlab = _install_fake_swanlab(monkeypatch)
+    tracker = SwanlabTracker("project", str(tmp_path))
+    gif_path = tmp_path / "validation.gif"
+    Image.new("RGB", (2, 2)).save(gif_path, format="GIF")
+
+    artifact = tracker.video(gif_path, caption="existing")
+
+    assert swanlab.video_paths == [str(gif_path)]
+    assert gif_path.exists()
+    assert artifact["caption"] == "existing"
+
+
+def test_swanlab_missing_dependency_error_is_actionable(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    monkeypatch.delitem(sys.modules, "swanlab", raising=False)
+    real_import = builtins.__import__
+
+    def missing_swanlab(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name == "swanlab":
+            raise ModuleNotFoundError("No module named 'swanlab'", name="swanlab")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", missing_swanlab)
+
+    with pytest.raises(ModuleNotFoundError, match=r"fastvideo\[swanlab\]"):
+        SwanlabTracker("project", str(tmp_path))
+
+
+class _RecordingTracker(BaseTracker):
+    def __init__(self, name: str) -> None:
+        super().__init__()
+        self.name = name
+        self.artifact_logs: list[dict[str, Any]] = []
+
+    def video(
+        self,
+        data: Any,
+        *,
+        caption: str | None = None,
+        fps: int | None = None,
+        format: str | None = None,
+    ) -> str:
+        return f"{self.name}:{data}"
+
+    def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None:
+        self.artifact_logs.append(artifacts)
+
+
+def test_sequential_tracker_logs_backend_specific_video_artifacts() -> None:
+    first = _RecordingTracker("first")
+    second = _RecordingTracker("second")
+    tracker = SequentialTracker([first, second])
+
+    artifact = tracker.video("sample.mp4")
+    tracker.log_artifacts({"validation/videos": [artifact]}, step=3)
+
+    assert first.artifact_logs == [{"validation/videos": ["first:sample.mp4"]}]
+    assert second.artifact_logs == [{"validation/videos": ["second:sample.mp4"]}]

fastvideo/tests/transformers/test_cosmos.py

@@ -11,6 +11,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
 from fastvideo.utils import maybe_download_model
 from fastvideo.configs.models.dits import CosmosVideoConfig
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@@ -25,18 +26,12 @@ BASE_MODEL_PATH = "nvidia/Cosmos-Predict2-2B-Video2World"
 
 
 def _resolve_model_path() -> str:
-    try:
-        return maybe_download_model(
-            BASE_MODEL_PATH,
-            local_dir=os.path.join("data", BASE_MODEL_PATH),
-        )
-    except ValueError as exc:
-        pytest.skip(
-            "Skipping Cosmos transformer test because the configured "
-            "HuggingFace token cannot access the gated Cosmos weights: "
-            f"{exc}",
-            allow_module_level=True,
-        )
+    local_dir = os.path.join("data", BASE_MODEL_PATH)
+    skip_if_gated_repo_inaccessible(BASE_MODEL_PATH,
+                                    local_path=local_dir,
+                                    test_name="Cosmos transformer test",
+                                    allow_module_level=True)
+    return maybe_download_model(BASE_MODEL_PATH, local_dir=local_dir)
 
 
 MODEL_PATH = _resolve_model_path()

fastvideo/tests/transformers/test_cosmos2_5.py

@@ -19,6 +19,7 @@ if os.path.exists(COSMOS_PREDICT2_5_PATH) and COSMOS_PREDICT2_5_PATH not in sys.
 
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
+from fastvideo.tests.utils import skip_if_gated_repo_inaccessible
 from fastvideo.utils import maybe_download_model
 # Use Cosmos 2.5 specific config
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
@@ -43,7 +44,14 @@ CHECKPOINT_SUBDIR = "base/post-trained"
 CHECKPOINT_FILENAME = "81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt"
 
 
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=None)
+def _resolve_model_path() -> str:
+    skip_if_gated_repo_inaccessible(BASE_MODEL_PATH,
+                                    test_name="Cosmos 2.5 transformer test",
+                                    allow_module_level=True)
+    return maybe_download_model(BASE_MODEL_PATH, local_dir=None)
+
+
+MODEL_PATH = _resolve_model_path()
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, CHECKPOINT_SUBDIR, "transformer")
 if not os.path.exists(TRANSFORMER_PATH):
     # Try without subdirectory
@@ -592,4 +600,3 @@ if __name__ == "__main__":
     # Run tests directly
     test_cosmos25_transformer()
     test_cosmos25_transformer_video()
-

fastvideo/tests/transformers/test_hunyuanvideo.py

@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import glob
-import json
-import os
-import pytest
-import torch
-
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.utils import maybe_download_model
-from fastvideo.forward_context import set_forward_context
-from fastvideo.configs.models.dits import HunyuanVideoConfig
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
-
-LOCAL_RANK = 0
-RANK = 0
-WORLD_SIZE = 1
-
-# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S
-REFERENCE_LATENT = 93.23263549804688
-
-
-@pytest.mark.usefixtures("distributed_setup")
-@pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
-def test_hunyuanvideo_distributed():
-    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
-
-    torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
-
-    # Get tensor parallel info
-    sp_rank = get_sp_parallel_rank()
-    sp_world_size = get_sp_world_size()
-
-    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
-
-    config = json.load(open(CONFIG_PATH))
-    # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
-    # TODO: write normalize config function
-    config.pop("_class_name")
-    config.pop("_diffusers_version")
-
-    precision_str = "bf16"
-    args = FastVideoArgs(
-        model_path=TRANSFORMER_PATH,
-        dit_cpu_offload=True,
-        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
-    )
-    args.device = torch.device(f"cuda:{LOCAL_RANK}")
-
-    loader = TransformerLoader()
-    model = loader.load(TRANSFORMER_PATH, args)
-
-    model.eval()
-
-    # Create random inputs for testing
-    batch_size = 1
-    seq_len = 3
-    device = torch.device(f"cuda:{LOCAL_RANK}")
-
-    # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
-    chunk_per_rank = hidden_states.shape[2] // sp_world_size
-    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
-
-    # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
-
-    # Timestep
-    timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
-    forward_batch = ForwardBatch(
-        data_type="dummy",
-    )
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Run inference on model
-        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
-                output = model(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    timestep=timestep,
-                )
-
-    latent = output.double().sum().item()
-
-    # Check if latents are similar
-    diff_output_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"

fastvideo/tests/transformers/test_hyworld.py

@@ -15,6 +15,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -24,13 +25,25 @@ os.environ["MASTER_PORT"] = "29503"
 
 MODEL_PATH = maybe_download_model("FastVideo/HY-WorldPlay-Bidirectional-Diffusers")
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-REFERENCE_LATENT = -197132.85557549074  # Pre-computed reference value
+REFERENCE_LATENTS = {
+    AttentionBackendEnum.FLASH_ATTN: -197132.85557549074,
+    AttentionBackendEnum.TORCH_SDPA: -211007.95158730447,
+}
 
 
 @pytest.mark.usefixtures("distributed_setup")
 def test_hyworld_transformer():
     transformer_path = TRANSFORMER_PATH
 
+    # An env-forced backend is knowable before the multi-GB download and GPU
+    # load; the authoritative check on the layer's resolved backend still runs
+    # after construction below.
+    forced_backend_name = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
+    if forced_backend_name:
+        forced_backend = getattr(AttentionBackendEnum, forced_backend_name, None)
+        if forced_backend is not None and forced_backend not in REFERENCE_LATENTS:
+            pytest.skip(f"HYWorld transformer test has no reference latent for {forced_backend_name}.")
+
     sp_rank = get_sp_parallel_rank()
     sp_world_size = get_sp_world_size()
 
@@ -51,6 +64,9 @@ def test_hyworld_transformer():
     loader = TransformerLoader()
     model = loader.load(transformer_path, args).to(device, dtype=precision)
     model.eval()
+    attention_backend = model.double_blocks[0].attn.backend
+    if attention_backend not in REFERENCE_LATENTS:
+        pytest.skip(f"HYWorld transformer test has no reference latent for {attention_backend}.")
 
     total_params = sum(p.numel() for p in model.parameters())
     weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters())
@@ -147,9 +163,10 @@ def test_hyworld_transformer():
 
     latent = output.double().sum().item()
 
-    diff = abs(REFERENCE_LATENT - latent)
-    relative_diff = diff / abs(REFERENCE_LATENT)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    reference_latent = REFERENCE_LATENTS[attention_backend]
+    diff = abs(reference_latent - latent)
+    relative_diff = diff / abs(reference_latent)
+    logger.info(f"Reference latent: {reference_latent}, Current latent: {latent}")
     logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%")
 
     # Allow 0.5% relative difference

fastvideo/tests/transformers/test_matrixgame2.py

@@ -1,133 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import pytest
-import torch
-
-from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.distributed.parallel_state import (
-    get_sp_parallel_rank,
-    get_sp_world_size,
-)
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29505"
-os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
-
-MODEL_PATH = maybe_download_model("FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers")
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-REFERENCE_LATENT = -111835.58895772696
-
-
-@pytest.mark.skip(reason="Not reliably reproducible")
-@pytest.mark.usefixtures("distributed_setup")
-def test_matrixgame2_transformer():
-    transformer_path = TRANSFORMER_PATH
-
-    sp_rank = get_sp_parallel_rank()
-    sp_world_size = get_sp_world_size()
-
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-
-    args = FastVideoArgs(
-        model_path=transformer_path,
-        dit_cpu_offload=False,
-        use_fsdp_inference=False,
-        pipeline_config=PipelineConfig(
-            dit_config=MatrixGame2WanVideoConfig(), dit_precision=precision_str
-        ),
-    )
-    args.device = device
-
-    loader = TransformerLoader()
-    model = loader.load(transformer_path, args).to(device, dtype=precision)
-    model.eval()
-
-    total_params = sum(p.numel() for p in model.parameters())
-    weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters())
-    weight_mean = weight_sum / total_params
-    logger.info("Total parameters: %s", total_params)
-    logger.info("Weight sum: %s", weight_sum)
-    logger.info("Weight mean: %s", weight_mean)
-
-    torch.manual_seed(42)
-
-    batch_size = 1
-    latent_frames = 15
-    latent_height = 44
-    latent_width = 80
-    image_dim = 1280
-    image_seq_len = 257
-
-    # Set num_frame_per_block to match actual latent_frames for action_model
-    model.num_frame_per_block = latent_frames
-    # Force full attention mode
-    model.local_attn_size = -1
-    # Clear cached block masks so they regenerate with correct settings
-    model.block_mask = None
-    model.block_mask_mouse = None
-    model.block_mask_keyboard = None
-
-    # Video latents [B, C, T, H, W]
-    x = torch.randn(batch_size, 16, latent_frames, latent_height, latent_width,
-                    device=device, dtype=precision)
-    cond_concat = torch.randn(batch_size, 20, latent_frames, latent_height, latent_width,
-                              device=device, dtype=precision)
-    hidden_states = torch.cat([x, cond_concat], dim=1)
-
-    if sp_world_size > 1:
-        chunk_per_rank = hidden_states.shape[2] // sp_world_size
-        hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank]
-
-    # Image embeddings
-    encoder_hidden_states_image = [
-        torch.randn(batch_size, image_seq_len, image_dim, device=device, dtype=precision)
-    ]
-
-    # Timestep
-    timestep = torch.full((batch_size, latent_frames), 500, device=device, dtype=precision)
-
-    # Action conditions
-    num_pixel_frames = (latent_frames - 1) * 4 + 1
-    mouse_cond = torch.randn(batch_size, num_pixel_frames, 2, device=device, dtype=precision)
-    keyboard_cond = torch.randn(batch_size, num_pixel_frames, 4, device=device, dtype=precision)
-
-    # Text embeddings (empty for i2v)
-    encoder_hidden_states = torch.zeros(batch_size, 0, model.config.hidden_size,
-                                        device=device, dtype=precision)
-
-    forward_batch = ForwardBatch(data_type="dummy")
-    model.num_frame_per_block = latent_frames
-
-    with torch.no_grad():
-        with torch.amp.autocast('cuda', dtype=precision):
-            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
-                output = model(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_hidden_states_image=encoder_hidden_states_image,
-                    timestep=timestep,
-                    mouse_cond=mouse_cond,
-                    keyboard_cond=keyboard_cond,
-                )
-
-    latent = output.double().sum().item()
-
-    diff = abs(REFERENCE_LATENT - latent)
-    relative_diff = diff / abs(REFERENCE_LATENT)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%")
-
-    # The 0.19% diff now
-    assert relative_diff < 0.005, f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)"

fastvideo/tests/utils.py

@@ -6,11 +6,45 @@ from fastvideo.logger import init_logger
 
 import numpy as np
 import torch
-from pytorch_msssim import ms_ssim, ssim
 
 logger = init_logger(__name__)
 
 
+def skip_if_gated_repo_inaccessible(repo_id: str,
+                                    *,
+                                    local_path: str | None = None,
+                                    test_name: str = "test",
+                                    allow_module_level: bool = False) -> None:
+    """Skip a test when ``repo_id`` is gated/private and the token lacks access.
+
+    Local weights win: when ``local_path`` already exists the check is skipped
+    entirely, so cached machines keep running offline. Transient hub failures
+    (offline, DNS, 429/5xx) do NOT skip either — the subsequent download will
+    serve from cache or fail loudly, preserving the CI failure signal. Only a
+    positively-identified authorization problem produces a skip.
+    """
+    import pytest
+
+    if local_path is not None and os.path.exists(local_path):
+        return
+
+    from huggingface_hub import HfApi
+    from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
+
+    try:
+        HfApi().auth_check(repo_id)
+    except (GatedRepoError, RepositoryNotFoundError) as exc:
+        pytest.skip(
+            f"Skipping {test_name}: the configured HuggingFace token cannot "
+            f"access the gated repo {repo_id}: {exc}",
+            allow_module_level=allow_module_level,
+        )
+    except Exception as exc:  # noqa: BLE001 - hub unreachable, proxies, 5xx
+        logger.warning(
+            "Gated-repo access probe for %s failed (%s); proceeding — the "
+            "download will serve from cache or fail loudly.", repo_id, exc)
+
+
 def _read_video_frames(path: str) -> torch.Tensor:
     """Read video frames as a ``(T, C, H, W)`` uint8 tensor.
 
@@ -52,6 +86,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
         video2_path: Path to the second video.
         use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
     """
+    from pytorch_msssim import ms_ssim, ssim
+
     print(f"Computing SSIM between {video1_path} and {video2_path}...")
     if not os.path.exists(video1_path):
         raise FileNotFoundError(f"Video1 not found: {video1_path}")

fastvideo/tests/vaes/test_hunyuan_vae.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import json
-import os
-
-import pytest
-import torch
-from safetensors.torch import load_file
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.logger import init_logger
-
-# from fastvideo.models.vaes.hunyuanvae import (
-#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import HunyuanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
-
-# Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
-# torch 2.7.1
-A40_REFERENCE_LATENT = -106.22467041015625
-
-# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
-L40S_REFERENCE_LATENT = -157.09130859375
-
-
-@pytest.mark.usefixtures("distributed_setup")
-@pytest.mark.skip(reason="Skipping test_hunyuan_vae")
-def test_hunyuan_vae():
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(
-        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
-    )
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model = loader.load(VAE_PATH, args)
-
-    model.enable_tiling(
-        tile_sample_min_height=32,
-        tile_sample_min_width=32,
-        tile_sample_min_num_frames=8,
-        tile_sample_stride_height=16,
-        tile_sample_stride_width=16,
-        tile_sample_stride_num_frames=4,
-    )
-
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        latent = model.encode(input_tensor).mean.double().sum().item()
-
-    # Check if latents are similar
-    device_name = torch.cuda.get_device_name()
-    if "A40" in device_name:
-        REFERENCE_LATENT = A40_REFERENCE_LATENT
-    elif "L40S" in device_name:
-        REFERENCE_LATENT = L40S_REFERENCE_LATENT
-    else:
-        raise ValueError(f"Unknown device: {device_name}")
-
-    diff_encoded_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"

fastvideo/tests/vaes/test_official_wan_vae.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import numpy as np
-import pytest
-import torch
-from diffusers import AutoencoderKLWan
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-try:
-    from wan.modules.vae2_2 import Wan2_2_VAE
-except ImportError:
-    Wan2_2_VAE = None
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-
-
-@pytest.mark.skip(reason="disable test")
-@pytest.mark.usefixtures("distributed_setup")
-def test_official_wan_vae():
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model2 = loader.load(VAE_PATH, args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
-
-    model1 = Wan2_2_VAE(vae_pth="/mnt/weka/home/hao.zhang/wei/Wan2.2_VAE.pth", device=device, dtype=precision)
-
-    # Create identical inputs for both models
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               81,
-                               32,
-                               32,
-                               device=device,
-                               dtype=precision)
-    # latent_tensor = torch.randn(batch_size,
-    #                             16,
-    #                             21,
-    #                             32,
-    #                             32,
-    #                             device=device,
-    #                             dtype=precision)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Test encoding
-        logger.info("Testing encoding...")
-        latent1 = model1.encode([input_tensor.squeeze(0)])[0].unsqueeze(0)
-        print("--------------------------------")
-        latent2 = model2.encode(input_tensor)
-        # Check if latents have the same shape
-        assert latent1.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
-        # Check if latents are similar
-        max_diff_encode = torch.max(torch.abs(latent1 - latent2.mean))
-        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2.mean))
-        logger.info("Maximum difference between encoded latents: %s",
-                    max_diff_encode.item())
-        logger.info("Mean difference between encoded latents: %s",
-                    mean_diff_encode.item())
-        assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}"
-        # Test decoding
-        logger.info("Testing decoding...")
-        # latent1_tensor = latent1.mode()
-        # mean1 = (torch.tensor(model1.config.latents_mean).view(
-        #     1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
-        #                                         input_tensor.dtype))
-        # std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
-        #     1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
-        #                                         input_tensor.dtype)
-        # latent1_tensor = latent1_tensor / std1 + mean1
-        # output1 = model1.decode(latent1_tensor).sample
-
-        output1 = model1.decode([latent1.squeeze(0)])
-        output1 = output1[0].unsqueeze(0)
-
-        mean2 = model2.config.arch_config.shift_factor.to(input_tensor.device, input_tensor.dtype)
-        std2 = model2.config.arch_config.scaling_factor.to(input_tensor.device, input_tensor.dtype)
-        latent2_tensor = latent2.mode()
-        latent2_tensor = latent2_tensor / std2 + mean2
-        output2 = model2.decode(latent2_tensor)
-        # Check if outputs have the same shape
-        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
-
-        # Check if outputs are similar
-        max_diff_decode = torch.max(torch.abs(output1.float() - output2.float()))
-        mean_diff_decode = torch.mean(torch.abs(output1.float() - output2.float()))
-        logger.info("Maximum difference between decoded outputs: %s",
-                    max_diff_decode.item())
-        logger.info("Mean difference between decoded outputs: %s",
-                    mean_diff_decode.item())
-        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"

fastvideo/train/callbacks/validation.py

@@ -404,6 +404,7 @@ class ValidationCallback(Callback):
                         art = self.tracker.video(
                             fname,
                             caption=cap,
+                            fps=sp.fps,
                         )
                         if art is not None:
                             video_logs.append(art)
@@ -456,11 +457,16 @@ class ValidationCallback(Callback):
             None,
         )
 
+        loaded_modules: dict[str, Any] = {"transformer": transformer}
+        # Distillation methods build the flow-match scheduler their few-step DMD
+        # sampler needs; inject it so the pipeline doesn't fall back to UniPC.
+        method_scheduler = getattr(self.method, "_sf_scheduler", None)
+        if method_scheduler is not None:
+            loaded_modules["scheduler"] = method_scheduler
+
         kwargs: dict[str, Any] = {
             "inference_mode": True,
-            "loaded_modules": {
-                "transformer": transformer,
-            },
+            "loaded_modules": loaded_modules,
             "tp_size": tc.distributed.tp_size,
             "sp_size": tc.distributed.sp_size,
             "num_gpus": tc.distributed.num_gpus,
@@ -477,12 +483,6 @@ class ValidationCallback(Callback):
             **kwargs,
         )
 
-        scheduler = self._pipeline.get_module("scheduler")
-        if (scheduler is not None and type(scheduler).__name__ == "SelfForcingFlowMatchScheduler"):
-            scheduler.sigma_min = 0.0
-            scheduler.extra_one_step = True
-            scheduler.set_timesteps(num_inference_steps=1000, training=True)
-
         self._pipeline_key = key
         return self._pipeline