[ci][diagnostic]: capture GameCraft FA4 candidate on GB200

Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>

.agents/skills/ci-runner/SKILL.md

@@ -0,0 +1,99 @@
+---
+name: ci-runner
+description: Work on FastVideo's Slurm-only, change-aware GPU CI lanes, static Buildkite graph, trusted ci-runner policy, lane scripts, and GB200 validation.
+---
+
+# Slinky Slurm CI lanes
+
+FastVideo's `ci-runner` Buildkite queue is the control plane for all active
+GPU CI. A private host-owned dispatcher leases GPUs from the Slinky Slurm tray
+and runs the immutable PR SHA inside an isolated Enroot container. Buildkite
+pipeline upload and Slurm submission occur on the login plane; every test
+payload executes on Slurm compute.
+
+The files under `fastvideo/tests/modal/` and `.buildkite/scripts/pr_test.sh`
+are dormant rollback code. Never add an active Buildkite or slash-command
+route to them. `pr_test.sh` must continue to reject Buildkite invocations.
+
+The private operator bundle is deliberately outside this repository because
+it contains site paths and credentials. See
+`docs/contributing/ci_architecture.md`; this skill covers the repository half
+and the coordination contract with that bundle.
+
+## Invariants
+
+- `.buildkite/pipeline.yml` contains exactly one static step for every active
+  GPU lane. Each step pins a unique key and label, a 90-minute timeout, the
+  trusted `/opt/fastvideo-ci-runner/run-ci` command (`run-unit` is the one
+  compatibility wrapper), step-level internal `TEST_TYPE`, and
+  `queue: "ci-runner"`.
+- Active CI contains no `pr_test.sh` command, Modal invocation, default queue,
+  Buildkite plugin, `soft_fail`, or job-controlled artifact glob.
+- The six Fastcheck lanes use `:microscope:` labels. Full-Suite-only lanes use
+  `:test_tube:` or `:bar_chart:` so direct reruns update the right aggregate.
+- SSIM and vanilla training request all four GPUs. Keep both in the
+  `fastvideo/slinky/whole-tray` Buildkite concurrency group with a limit of one
+  so the second job does not consume an agent or command timeout while waiting
+  for the same tray.
+- `/test full` schedules all twenty lanes. `/merge`, `ready`, and new pushes to
+  ready PRs use the trusted base-branch planner in
+  `.github/scripts/plan_merge_ci.py`: automatic Fastcheck remains the universal
+  six-lane baseline, and the merge build adds only path-relevant integration
+  lanes. Unknown source/build paths fail closed to all fourteen additive lanes.
+  The trusted uploader still normalizes and validates the complete static graph
+  before Buildkite evaluates its plan conditions.
+- Focused merge builds may pass allowlisted golden-gate and SSIM test basenames.
+  The private host validates the lane plan and basenames before staging them,
+  and the in-container scripts validate them again. Direct `/test ssim`,
+  explicit `/test full`, and the weekly main-branch schedule run the complete
+  SSIM matrix.
+- The trusted uploader serves exactly three entry pipelines:
+  `pr-fastcheck` for automatic PR builds, `ci` for slash-command/ready-label
+  API builds, and `fastvideo-performance-lane` for the weekly schedule. Keep
+  incoming GitHub webhook processing disabled on `ci` so it cannot duplicate
+  `pr-fastcheck` on every PR update.
+- Test payloads live in `.buildkite/scripts/unit_test.sh` or executable
+  `.buildkite/scripts/lanes/<lane>.sh`. Backend policy (GPU count, extras,
+  secrets, kernel build, artifacts) stays in the agent-owned lane table.
+- Tests must preserve an inherited `MASTER_PORT`. Packed containers share the
+  tray network namespace, so the private runner assigns a distinct port range
+  per GPU lease and the SSIM scheduler assigns task offsets within its range.
+- The ARM64 runner image includes the pinned FA4 CuTe overlay validated on
+  GB200. Keep SSIM at `FASTVIDEO_FA4=1` because its references were seeded with
+  FA4; keep lanes with FA2 baselines at `FASTVIDEO_FA4=0`. A runner image change
+  must revalidate both the FA4 import and an actual GB200 forward kernel.
+- `fastvideo/tests/ssim/ci_runner.py` is the active four-GPU SSIM scheduler.
+  New SSIM files are discovered through `REQUIRED_GPUS` and
+  `*_MODEL_TO_PARAMS`; do not wire them through the dormant Modal scheduler.
+- The host policy fail-closes unknown tuples. A repository-side lane change is
+  inert until the operator updates the private lane table and uploader policy
+  in the same rollout.
+
+## Adding or changing a lane
+
+1. Read the closest `AGENTS.md` and the domain-specific testing guide.
+2. Add or update the executable lane payload under `.buildkite/scripts/`.
+   Keep it deterministic and free of host-specific paths or credential fetches.
+3. Add the static pipeline step and canonical `/test <name>` mapping. Keep the
+   `<name>-ci` alias only when compatibility requires it.
+4. Add its source/test path ownership to `.github/scripts/plan_merge_ci.py`.
+   Prefer the narrowest correctness-preserving lane set; leave unknown paths
+   fail-closed. Extend `fastvideo/tests/contract/test_ci_test_collection.py`,
+   `test_merge_ci_plan.py`, and focused CPU-only scheduler/policy tests.
+5. Coordinate the private lane row: GPU count (1-4), wall time, script, scope
+   pairs, step key, command, HF cache/token, tracking mode, extras, attention
+   backend policy, kernel policy, and artifact relay. Active training lanes
+   keep W&B offline and do not stage a W&B credential.
+6. Update the trusted pipeline-uploader schema. A mismatch must reject the
+   pipeline rather than silently skip a lane.
+7. Run `pre-commit run --files <changed paths>`, the planner's representative
+   diff matrix, contract tests, private driver tests, and a real GB200 canary.
+   Multi-GPU, hardware-reference, training, performance, and SSIM changes need
+   their own target-hardware evidence.
+
+## Rollback
+
+Rollback the Slurm routing/configuration change or pause the `ci-runner` queue.
+Do not silently reactivate Modal. A manual Modal experiment requires the
+explicit local opt-in documented in `ci_architecture.md`; returning it to
+production CI needs a separate reviewed decision.

.agents/skills/reseed-ssim-references/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: reseed-ssim-references
-description: Re-seed HF reference videos for a single existing SSIM test on Modal L40S. Always backs up current refs locally first, regenerates on Modal, pauses for the user to eyeball before-vs-after quality, then overwrites the targeted `<model_id>` subtree on `FastVideo/ssim-reference-videos` with `--force`. Use when an intentional code change (model port fix, attention backend swap, kernel upgrade, hyperparameter change) has invalidated existing refs and they need to be regenerated. Pairs with `seed-ssim-references`, which is for first-time seeding only.
+description: Re-seed HF reference videos for a single existing SSIM test on Modal L40S. Always backs up current refs locally first, regenerates on Modal, pauses for the user to eyeball before-vs-after quality, then overwrites the targeted model subtree on `FastVideo/ssim-reference-videos` with `--force`. Use when an intentional code change (model port fix, attention backend swap, kernel upgrade, hyperparameter change) has invalidated existing refs and they need to be regenerated. Pairs with `seed-ssim-references`, which is for first-time seeding only.
 ---
 
 # Re-seed SSIM Reference Videos
@@ -13,7 +13,7 @@ on HF — the old refs are overwritten — so the skill always:
 
 1. Confirms intent with a one-liner the user has to type.
 2. Downloads the existing refs as a local, timestamped backup.
-3. Regenerates on Modal L40S (same code path that CI uses).
+3. Regenerates through the manual legacy Modal L40S maintenance path.
 4. Pauses for a side-by-side eyeball of backup vs new mp4s.
 5. Uploads with `--force`, scoped to the single `--model-id`.
 6. Reminds the user to keep the backup until the PR lands.
@@ -51,8 +51,9 @@ harder to recover from than failing closed.
 
 Hardcoded:
 
-- Modal GPU: **L40S** (matches CI; re-seeding from another SKU produces refs
-  that L40S CI cannot match).
+- Modal GPU: **L40S**. This is a manual reference-maintenance target, not the
+  active Slurm CI compute path; changing the SKU also changes the historical
+  `L40S_reference_videos` contract.
 - Quality tier: **`default`**. `full_quality` is a separate, deliberate
   operation.
 - HF repo: `FastVideo/ssim-reference-videos` (override via

.agents/skills/seed-ssim-references/SKILL.md

@@ -35,7 +35,8 @@ The skill is run **manually**, once per new test. Before invoking it, the user
 has already sanity-tested the new test locally — it launches `VideoGenerator`
 and writes an artefact without crashing (the missing-reference assertion at
 the end is expected). The skill does not re-test locally; it goes straight
-to Modal L40S (which is what CI uses).
+to the manual legacy Modal L40S reference-maintenance target. Active CI runs
+on the Slinky Slurm cluster and only consumes the resulting references.
 
 ## When to use
 
@@ -61,7 +62,8 @@ Prompt the user for it if they didn't supply it.
 
 Everything else is fixed:
 
-- Modal runner GPU: **L40S** (hardcoded in `fastvideo/tests/modal/ssim_test.py`).
+- Modal maintenance GPU: **L40S** (hardcoded in
+  `fastvideo/tests/modal/ssim_test.py`; this is not the active CI compute path).
 - Device folder: `L40S_reference_videos`.
 - Quality tier: `default` (the tier CI runs). The `full_quality` tier is not
   seeded by this skill.

.buildkite/pipeline.yml

@@ -1,7 +1,8 @@
 env:
   IMAGE_VERSION: "py3.12-latest"
   BUILDKITE_CLEAN_CHECKOUT: true
-  # Buildkite only launches Modal; remote jobs initialize their own submodules.
+  # Slurm workers clone the immutable commit and initialize submodules inside
+  # their isolated container. The Buildkite login-plane checkout is a no-op.
   BUILDKITE_GIT_SUBMODULES: false
 
 notify:
@@ -10,539 +11,458 @@ notify:
     if: build.env("TEST_SCOPE") == "fastcheck" || build.env("TEST_SCOPE") == null
   - github_commit_status:
       context: "full-suite-passed"
-    if: build.env("TEST_SCOPE") == "full"
+    if: build.env("TEST_SCOPE") == "full" || build.env("TEST_SCOPE") == "merge"
   - github_commit_status:
       context: "direct-test-completed"
     if: build.env("TEST_SCOPE") == "direct"
+  - github_commit_status:
+      context: "scheduled-ssim-passed"
+    if: build.env("TEST_SCOPE") == "scheduled"
 
+# This is the complete active GPU CI surface. Every command is a trusted host
+# dispatcher, and every test payload executes inside the Slinky Slurm tray.
+# fastvideo/tests/modal remains available only for an explicit manual rollback;
+# no active pipeline or slash-command route invokes it.
 steps:
-  # ============================================================
-  # Direct test: triggered by /test <name> slash command.
-  # Labels match fastcheck/full-suite counterparts so the GitHub
-  # check status overwrites the original failed check.
-  # Only ONE step executes per build (gated by TEST_TYPE).
-  # ============================================================
+  - label: ":microscope: Encoder Tests"
+    key: "encoder"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,encoder,/) ||
+      build.env("TEST_SCOPE") == "fastcheck" ||
+      build.env("TEST_SCOPE") == null ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "encoder" || build.env("TEST_TYPE") == "encoder_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "encoder_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
 
-  # --- Fastcheck-scope direct tests ---
-    - label: ":microscope: Encoder Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "encoder"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":microscope: VAE Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "vae"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":microscope: Transformer Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "transformer"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":microscope: Kernel Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "kernel_tests"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":vertical_traffic_light: Golden-Gate Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "golden_gate"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":microscope: Unit Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "unit_test"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":microscope: DreamVerse App Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "dreamverse_app"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
+  - label: ":microscope: VAE Tests"
+    key: "vae"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,vae,/) ||
+      build.env("TEST_SCOPE") == "fastcheck" ||
+      build.env("TEST_SCOPE") == null ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "vae" || build.env("TEST_TYPE") == "vae_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "vae_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
 
-  # --- Full-suite-scope direct tests ---
-    - label: ":bar_chart: SSIM Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "ssim"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-          - exit_status: 1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: LoRA Inference Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "inference_lora"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: LoRA Extraction Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "lora_extraction"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Training Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Distillation DMD Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "distillation_dmd"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Self-Forcing Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "self_forcing"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: LoRA Training Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training_lora"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-          - exit_status: 1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Training Tests VSA"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training_vsa"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-          - exit_status: 1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Inference Tests VMoBA"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "inference_vmoba"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Performance Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "performance"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: API Server Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "api_server"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Train Framework Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "train_framework"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
-    - label: ":test_tube: Eval Metrics Tests"
-      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "eval"
-      command: "timeout 90m .buildkite/scripts/pr_test.sh"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      agents:
-        queue: "default"
+  - label: ":microscope: Transformer Tests"
+    key: "transformer"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,transformer,/) ||
+      build.env("TEST_SCOPE") == "fastcheck" ||
+      build.env("TEST_SCOPE") == null ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "transformer" || build.env("TEST_TYPE") == "transformer_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "transformer_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
 
-  # ============================================================
-  # Fastcheck: Runs on every PR (~10-15 min parallel)
-  # Core component validation: encoders, VAEs, transformers,
-  # CUDA kernels, and unit tests.
-  # ============================================================
-    - label: "Trigger Fastcheck"
-      if: build.env("TEST_SCOPE") == "fastcheck" || build.env("TEST_SCOPE") == null
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      plugins:
-        - monorepo-diff#v1.4.0:
-            diff: 'git fetch origin "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" && git diff --name-only "origin/${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}...HEAD"'
-            watch:
-            - path:
-                - "fastvideo/models/encoders/**"
-                - "fastvideo/models/loader/**"
-                - "fastvideo/tests/encoders/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 20m .buildkite/scripts/pr_test.sh"
-                label: ":microscope: Encoder Tests"
-                env:
-                  - TEST_TYPE=encoder
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/models/vaes/**"
-                - "fastvideo/models/loader/**"
-                - "fastvideo/tests/vaes/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 20m .buildkite/scripts/pr_test.sh"
-                label: ":microscope: VAE Tests"
-                env:
-                  - TEST_TYPE=vae
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/models/dits/**"
-                - "fastvideo/models/loader/**"
-                - "fastvideo/tests/transformers/**"
-                - "fastvideo/layers/**"
-                - "fastvideo/attention/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: ":microscope: Transformer Tests"
-                env:
-                  - TEST_TYPE=transformer
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo-kernel/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: ":microscope: Kernel Tests"
-                env:
-                  - TEST_TYPE=kernel_tests
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/**"
-                - ".buildkite/**"
-                - ".github/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: ":microscope: Unit Tests"
-                env:
-                  - TEST_TYPE=unit_test
-                agents:
-                  queue: "default"
-            - path:
-                - "apps/dreamverse/**"
-                - "pyproject.toml"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: ":microscope: DreamVerse App Tests"
-                env:
-                  - TEST_TYPE=dreamverse_app
-                agents:
-                  queue: "default"
+  - label: ":microscope: Kernel Tests"
+    key: "kernel-tests"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,kernel-tests,/) ||
+      build.env("TEST_SCOPE") == "fastcheck" ||
+      build.env("TEST_SCOPE") == null ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "kernel_tests" || build.env("TEST_TYPE") == "kernel_tests_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "kernel_tests_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
 
-  # ============================================================
-  # Full Suite: Runs when TEST_SCOPE=full
-  # Triggered by adding the 'ready' label (via ci-trigger-full-suite.yml)
-  # or on-demand via /test full slash command.
-  # Includes integration tests, SSIM regression, training pipelines,
-  # and performance benchmarks.
-  # ============================================================
-    - label: "Trigger Full Suite"
-      if: build.env("TEST_SCOPE") == "full"
-      retry:
-        automatic:
-          - exit_status: 128
-            limit: 3
-          - exit_status: -1
-            limit: 2
-      plugins:
-        - monorepo-diff#v1.4.0:
-            diff: 'git fetch origin "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" && git diff --name-only "origin/${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}...HEAD"'
-            watch:
-            - path:
-                - "fastvideo/**/*.py"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 90m .buildkite/scripts/pr_test.sh"
-                label: ":bar_chart: SSIM Tests"
-                env:
-                  - TEST_TYPE=ssim
-                retry:
-                  automatic:
-                    - exit_status: 1
-                      limit: 2
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/tests/lora/**"
-                - "fastvideo/models/loader/**"
-                - "fastvideo/tests/transformers/**"
-                - "fastvideo/pipelines/**"
-                - "fastvideo/layers/lora/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 20m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: LoRA Inference Tests"
-                env:
-                  - TEST_TYPE=inference_lora
-                agents:
-                  queue: "default"
-            - path:
-                - "scripts/lora_extraction/**"
-                - "fastvideo/tests/lora_extraction/**"
-                - "fastvideo/models/loader/**"
-                - "fastvideo/training/training_utils.py"
-                - "fastvideo/layers/lora/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 90m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: LoRA Extraction Tests"
-                env:
-                  - TEST_TYPE=lora_extraction
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 25m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Training Tests"
-                env:
-                  - TEST_TYPE=training
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/training/*distillation_pipeline.py"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 25m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Distillation DMD Tests"
-                env:
-                  - TEST_TYPE=distillation_dmd
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/training/*self_forcing_distillation_pipeline.py"
-                - "fastvideo/tests/training/self-forcing/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Self-Forcing Tests"
-                env:
-                  - TEST_TYPE=self_forcing
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 25m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: LoRA Training Tests"
-                env:
-                  - TEST_TYPE=training_lora
-                retry:
-                  automatic:
-                    - exit_status: 1
-                      limit: 2
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/**"
-                - "fastvideo-kernel/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Training Tests VSA"
-                env:
-                  - TEST_TYPE=training_vsa
-                retry:
-                  automatic:
-                    - exit_status: 1
-                      limit: 2
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo-kernel/**"
-                - "fastvideo/attention/backends/vmoba.py"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Inference Tests VMoBA"
-                env:
-                  - TEST_TYPE=inference_vmoba
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/models/dits/**"
-                - "fastvideo/pipelines/**"
-                - "fastvideo/attention/**"
-                - "fastvideo/layers/**"
-                - "fastvideo/worker/**"
-                - "fastvideo/entrypoints/**"
-                - "fastvideo/performance/**"
-                - "fastvideo/tests/performance/**"
-                - ".buildkite/performance-benchmarks/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Performance Tests"
-                env:
-                  - TEST_TYPE=performance
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/entrypoints/openai/**"
-                - "fastvideo/entrypoints/cli/serve.py"
-                - "fastvideo/tests/entrypoints/test_openai_api_integration.py"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: API Server Tests"
-                env:
-                  - TEST_TYPE=api_server
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/train/**"
-                - "fastvideo/tests/train/models/**"
-                - "fastvideo/tests/train/fixtures/**"
-                - "fastvideo/models/dits/**"
-                - "fastvideo/models/loader/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Train Framework Tests"
-                env:
-                  - TEST_TYPE=train_framework
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/eval/**"
-                - "fastvideo/tests/eval/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile"
-              config:
-                command: "timeout 90m .buildkite/scripts/pr_test.sh"
-                label: ":test_tube: Eval Metrics Tests"
-                env:
-                  - TEST_TYPE=eval
-                agents:
-                  queue: "default"
+  - label: ":microscope: Unit Tests"
+    key: "unit"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,unit,/) ||
+      build.env("TEST_SCOPE") == "fastcheck" ||
+      build.env("TEST_SCOPE") == null ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "unit_test" || build.env("TEST_TYPE") == "unit_test_ci"))
+    command: "/opt/fastvideo-ci-runner/run-unit"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "unit_test_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":microscope: DreamVerse App Tests"
+    key: "dreamverse"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,dreamverse,/) ||
+      build.env("TEST_SCOPE") == "fastcheck" ||
+      build.env("TEST_SCOPE") == null ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "dreamverse_app" || build.env("TEST_TYPE") == "dreamverse_app_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "dreamverse_app_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Golden-Gate Tests"
+    key: "golden-gate"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,golden-gate,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "golden_gate" || build.env("TEST_TYPE") == "golden_gate_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "golden_gate_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":bar_chart: SSIM Tests"
+    key: "ssim"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      build.env("TEST_SCOPE") == "scheduled" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,ssim,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "ssim" || build.env("TEST_TYPE") == "ssim_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    concurrency: 1
+    concurrency_group: "fastvideo/slinky/whole-tray"
+    env:
+      TEST_TYPE: "ssim_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+        - exit_status: 1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: LoRA Inference Tests"
+    key: "lora-inference"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,lora-inference,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "inference_lora" || build.env("TEST_TYPE") == "inference_lora_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "inference_lora_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: LoRA Extraction Tests"
+    key: "lora-extraction"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,lora-extraction,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "lora_extraction" || build.env("TEST_TYPE") == "lora_extraction_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "lora_extraction_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Training Tests"
+    key: "training"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,training,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "training" || build.env("TEST_TYPE") == "training_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    concurrency: 1
+    concurrency_group: "fastvideo/slinky/whole-tray"
+    env:
+      TEST_TYPE: "training_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Distillation DMD Tests"
+    key: "distillation"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,distillation,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "distillation_dmd" || build.env("TEST_TYPE") == "distillation_dmd_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "distillation_dmd_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Self-Forcing Tests"
+    key: "self-forcing"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,self-forcing,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "self_forcing" || build.env("TEST_TYPE") == "self_forcing_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "self_forcing_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: LoRA Training Tests"
+    key: "lora-training"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,lora-training,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "training_lora" || build.env("TEST_TYPE") == "training_lora_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "training_lora_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+        - exit_status: 1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Training Tests VSA"
+    key: "training-vsa"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,training-vsa,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "training_vsa" || build.env("TEST_TYPE") == "training_vsa_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "training_vsa_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+        - exit_status: 1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Inference Tests VMoBA"
+    key: "inference-vmoba"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,inference-vmoba,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "inference_vmoba" || build.env("TEST_TYPE") == "inference_vmoba_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "inference_vmoba_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Performance Tests"
+    key: "performance"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,performance,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "performance" || build.env("TEST_TYPE") == "performance_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "performance_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: API Server Tests"
+    key: "api-server"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,api-server,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "api_server" || build.env("TEST_TYPE") == "api_server_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "api_server_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Train Framework Tests"
+    key: "train-framework"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,train-framework,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "train_framework" || build.env("TEST_TYPE") == "train_framework_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "train_framework_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"
+
+  - label: ":test_tube: Eval Metrics Tests"
+    key: "eval"
+    if: |
+      build.env("TEST_SCOPE") == "full" ||
+      (build.env("TEST_SCOPE") == "merge" &&
+       build.env("MERGE_TEST_PLAN") =~ /,eval,/) ||
+      (build.env("TEST_SCOPE") == "direct" &&
+       (build.env("TEST_TYPE") == "eval" || build.env("TEST_TYPE") == "eval_ci"))
+    command: "/opt/fastvideo-ci-runner/run-ci"
+    timeout_in_minutes: 90
+    env:
+      TEST_TYPE: "eval_ci"
+    retry:
+      automatic:
+        - exit_status: 128
+          limit: 3
+        - exit_status: -1
+          limit: 2
+    agents:
+      queue: "ci-runner"

.buildkite/scripts/gamecraft_candidate_log.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+"""Emit and recover a small GameCraft MP4 through a Buildkite job log."""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import binascii
+import hashlib
+import html
+import json
+import re
+import sys
+from pathlib import Path
+
+MARKER = "FV_GAMECRAFT_FA4_MP4"
+CHUNK_BYTES = 3 * 1024
+MAX_MEDIA_BYTES = 1024 * 1024
+
+
+def _sha256(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def emit(media_path: Path) -> None:
+    if media_path.suffix.lower() != ".mp4":
+        raise ValueError(f"Candidate must be an MP4: {media_path}")
+
+    data = media_path.read_bytes()
+    size = len(data)
+    if not 0 < size <= MAX_MEDIA_BYTES:
+        raise ValueError(f"Candidate size must be between 1 and {MAX_MEDIA_BYTES} bytes; got {size}")
+
+    digest = _sha256(data)
+    chunks = [data[offset:offset + CHUNK_BYTES] for offset in range(0, size, CHUNK_BYTES)]
+    print(
+        f"{MARKER}_BEGIN version=1 sha256={digest} size={size} "
+        f"chunks={len(chunks)} chunk_bytes={CHUNK_BYTES}",
+        flush=True,
+    )
+    for index, chunk in enumerate(chunks):
+        encoded = base64.b64encode(chunk).decode("ascii")
+        print(f"{MARKER}_CHUNK index={index:06d} data={encoded}", flush=True)
+    print(
+        f"{MARKER}_END version=1 sha256={digest} size={size} chunks={len(chunks)}",
+        flush=True,
+    )
+
+
+def _buildkite_output(text: str) -> str:
+    """Unwrap Buildkite's public JSON and reverse its HTML entity escaping."""
+    try:
+        payload = json.loads(text)
+    except json.JSONDecodeError:
+        return html.unescape(text)
+    if isinstance(payload, dict) and isinstance(payload.get("output"), str):
+        return html.unescape(payload["output"])
+    return html.unescape(text)
+
+
+def decode(log_text: str) -> tuple[bytes, str]:
+    log_text = _buildkite_output(log_text)
+    begin_pattern = re.compile(
+        rf"{MARKER}_BEGIN version=1 sha256=([0-9a-f]{{64}}) size=([0-9]+) "
+        rf"chunks=([0-9]+) chunk_bytes=([0-9]+)"
+    )
+    end_pattern = re.compile(
+        rf"{MARKER}_END version=1 sha256=([0-9a-f]{{64}}) size=([0-9]+) chunks=([0-9]+)"
+    )
+    chunk_pattern = re.compile(rf"{MARKER}_CHUNK index=([0-9]{{6}}) data=([A-Za-z0-9+/]+={{0,2}})")
+
+    begin_matches = begin_pattern.findall(log_text)
+    end_matches = end_pattern.findall(log_text)
+    if len(begin_matches) != 1 or len(end_matches) != 1:
+        raise ValueError(
+            "Expected exactly one candidate envelope; "
+            f"found {len(begin_matches)} begin and {len(end_matches)} end markers"
+        )
+
+    digest, size_text, count_text, chunk_bytes_text = begin_matches[0]
+    end_digest, end_size_text, end_count_text = end_matches[0]
+    if (digest, size_text, count_text) != (end_digest, end_size_text, end_count_text):
+        raise ValueError("Candidate begin/end metadata does not match")
+
+    size = int(size_text)
+    count = int(count_text)
+    chunk_bytes = int(chunk_bytes_text)
+    if not 0 < size <= MAX_MEDIA_BYTES:
+        raise ValueError(f"Candidate size is outside the accepted range: {size}")
+    if chunk_bytes != CHUNK_BYTES:
+        raise ValueError(f"Unexpected chunk size: {chunk_bytes}")
+
+    matches = chunk_pattern.findall(log_text)
+    if len(matches) != count:
+        raise ValueError(f"Expected {count} candidate chunks; found {len(matches)}")
+
+    encoded_chunks: dict[int, str] = {}
+    for index_text, encoded in matches:
+        index = int(index_text)
+        if index in encoded_chunks:
+            raise ValueError(f"Duplicate candidate chunk: {index}")
+        encoded_chunks[index] = encoded
+    if sorted(encoded_chunks) != list(range(count)):
+        raise ValueError("Candidate chunk sequence is incomplete or out of range")
+
+    try:
+        data = b"".join(base64.b64decode(encoded_chunks[index], validate=True) for index in range(count))
+    except binascii.Error as error:
+        raise ValueError(f"Candidate chunk is not valid base64: {error}") from error
+    if len(data) != size:
+        raise ValueError(f"Decoded candidate size mismatch: expected {size}, got {len(data)}")
+    actual_digest = _sha256(data)
+    if actual_digest != digest:
+        raise ValueError(f"Decoded candidate sha256 mismatch: expected {digest}, got {actual_digest}")
+    return data, digest
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    emit_parser = subparsers.add_parser("emit", help="Encode one MP4 to stdout between strict log markers")
+    emit_parser.add_argument("media_path", type=Path)
+
+    decode_parser = subparsers.add_parser("decode", help="Recover and verify one MP4 from a raw or public JSON log")
+    decode_parser.add_argument("--input", type=Path, help="Log file (default: stdin)")
+    decode_parser.add_argument("--output", type=Path, required=True)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = _parse_args()
+    if args.command == "emit":
+        emit(args.media_path)
+        return 0
+
+    log_text = args.input.read_text() if args.input else sys.stdin.read()
+    data, digest = decode(log_text)
+    if args.output.exists():
+        raise FileExistsError(f"Refusing to overwrite existing output: {args.output}")
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_bytes(data)
+    print(f"Recovered {len(data)} bytes with sha256={digest} to {args.output}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.buildkite/scripts/lanes/api_server.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the OpenAI-compatible API lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/entrypoints/test_openai_api_integration.py -vs

.buildkite/scripts/lanes/distillation_dmd.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the distillation-DMD lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs

.buildkite/scripts/lanes/dreamverse.sh

@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+# DreamVerse needs a GPU for import-time device resolution, but it does not
+# build or exercise fastvideo-kernel. A checksummed Node archive is installed
+# in the disposable Slurm container because the shared CI image is
+# Python/CUDA focused.
+set -euo pipefail
+
+node_version=v22.23.2
+case $(uname -m) in
+  aarch64 | arm64)
+    node_arch=arm64
+    node_archive_sha256=013b59cfd2819703a6f4a14ab891fc46fc2a4e3f5bcd92de3fb4929b43e35b30
+    ;;
+  x86_64 | amd64)
+    node_arch=x64
+    node_archive_sha256=b294a556e639d64338823920e5866c21c02741742d2e1529ee1a225c1ec9252a
+    ;;
+  *)
+    echo "Unsupported architecture for DreamVerse Node runtime: $(uname -m)" >&2
+    exit 2
+    ;;
+esac
+node_archive="node-${node_version}-linux-${node_arch}.tar.gz"
+node_runtime_root=$(mktemp -d -t fastvideo-node.XXXXXX)
+node_archive_path="${node_runtime_root}/${node_archive}"
+node_install_dir="${node_runtime_root}/${node_archive%.tar.gz}"
+curl --proto '=https' --tlsv1.2 --retry 5 --retry-all-errors \
+  --location --fail --silent --show-error \
+  "https://nodejs.org/dist/${node_version}/${node_archive}" \
+  --output "$node_archive_path"
+printf '%s  %s\n' "$node_archive_sha256" "$node_archive_path" | sha256sum --check --status
+tar -xzf "$node_archive_path" -C "$node_runtime_root"
+export PATH="${node_install_dir}/bin:${PATH}"
+node --version
+npm --version
+
+export PYTHONPATH="$(pwd)/apps/dreamverse${PYTHONPATH:+:$PYTHONPATH}"
+pytest apps/dreamverse/dreamverse/tests -q
+
+cd apps/dreamverse/web
+npm ci
+npm run typecheck
+npm test
+machine_arch=$(uname -m)
+if [[ $machine_arch =~ ^(aarch64|arm64)$ ]]; then
+  npx playwright install --with-deps chromium firefox
+else
+  npx playwright install --with-deps chromium webkit firefox
+fi
+
+master_port=${MASTER_PORT:-7959}
+BACKEND_PORT=${BACKEND_PORT:-$((master_port + 50))}
+python -m uvicorn dreamverse.mock_server:app --host 127.0.0.1 --port "$BACKEND_PORT" &
+mock_server_pid=$!
+cleanup() {
+  kill "$mock_server_pid" 2>/dev/null || true
+  wait "$mock_server_pid" 2>/dev/null || true
+}
+trap cleanup EXIT INT TERM
+
+for _ in {1..30}; do
+  curl -fsS "http://127.0.0.1:$BACKEND_PORT/healthz" && break
+  sleep 1
+done
+curl -fsS "http://127.0.0.1:$BACKEND_PORT/healthz"
+
+if [[ $machine_arch =~ ^(aarch64|arm64)$ ]]; then
+  # Playwright WebKit traps before opening a page on Linux ARM64, and its
+  # bundled Chromium lacks the H.264/AAC codecs used by the fMP4 assertions.
+  # Firefox covers every flow, including streaming. Chromium and its mobile
+  # profile still cover all codec-independent UI behavior on GB200.
+  BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
+    npm run e2e -- --project=firefox
+  BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
+    npm run e2e -- \
+      --project=chromium \
+      --project=mobile-chromium \
+      --grep-invert='streams, plays, and surfaces a downloadable clip|starts a new project and switches back to the prior session|saved projects persist across a page reload'
+else
+  BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
+    npm run e2e -- \
+      --project=chromium \
+      --project=webkit \
+      --project=firefox \
+      --project=mobile-safari \
+      --project=mobile-chromium
+fi

.buildkite/scripts/lanes/encoder.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the encoder lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/encoders -vs

.buildkite/scripts/lanes/eval.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the evaluation lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/eval -vs

.buildkite/scripts/lanes/golden_gate.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the golden-gate lane. Environment (HF_HOME
+# and authentication) is the runner's responsibility.
+set -euo pipefail
+
+golden_root=./fastvideo/tests/golden_gate
+selected=${FASTVIDEO_GOLDEN_TEST_FILES-}
+if [ -z "$selected" ]; then
+  if [ "${TEST_SCOPE:-}" = merge ]; then
+    echo "Missing FASTVIDEO_GOLDEN_TEST_FILES for merge scope" >&2
+    exit 2
+  fi
+  selected=all
+fi
+if [ "$selected" = all ]; then
+  exec pytest "$golden_root" -vs
+fi
+
+[[ $selected =~ ^test_[a-z0-9_]+\.py(,test_[a-z0-9_]+\.py)*$ ]] || {
+  echo "Invalid FASTVIDEO_GOLDEN_TEST_FILES selection" >&2
+  exit 2
+}
+
+IFS=, read -r -a golden_files <<< "$selected"
+golden_paths=()
+for golden_file in "${golden_files[@]}"; do
+  golden_path="$golden_root/$golden_file"
+  [ -f "$golden_path" ] || {
+    echo "Selected golden test does not exist: $golden_file" >&2
+    exit 2
+  }
+  golden_paths+=("$golden_path")
+done
+
+exec pytest "${golden_paths[@]}" -vs

.buildkite/scripts/lanes/inference_lora.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the LoRA-inference lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs

.buildkite/scripts/lanes/inference_vmoba.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the VMoBA-inference lane.
+set -euo pipefail
+
+exec python fastvideo/tests/inference/vmoba/test_vmoba_inference.py

.buildkite/scripts/lanes/kernel_tests.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the custom-kernel lane.
+set -euo pipefail
+
+exec pytest fastvideo-kernel/tests/ -vs

.buildkite/scripts/lanes/lora_extraction.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the LoRA-extraction lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py -vs

.buildkite/scripts/lanes/performance.sh

@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+# Canonical Slurm performance lane. Reports are written outside the checkout
+# so the trusted host driver can upload them after untrusted code exits.
+set -uo pipefail
+
+export PERFORMANCE_TRACKING_ROOT=/tmp/perf-tracking
+export PERF_REPORTS_DIR=/workspace/artifacts/performance
+mkdir -p "$PERF_REPORTS_DIR"
+
+if [[ ${BUILDKITE_PULL_REQUEST:-false} =~ ^[1-9][0-9]*$ ]]; then
+  export PERF_RUN_SOURCE=pr
+  export PERF_UPLOAD_POLICY=pass
+elif [ "${BUILDKITE_BRANCH:-}" = main ] \
+  && { [ "${BUILDKITE_SOURCE:-}" = schedule ] || [ "${TEST_SCOPE:-}" = full ]; }; then
+  export PERF_RUN_SOURCE=scheduled_main
+  export PERF_UPLOAD_POLICY=always
+elif [ "${TEST_SCOPE:-}" = direct ]; then
+  export PERF_RUN_SOURCE=unknown
+  export PERF_UPLOAD_POLICY=pass
+else
+  export PERF_RUN_SOURCE=unknown
+  export PERF_UPLOAD_POLICY=never
+fi
+
+nvidia-smi \
+  --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu \
+  --format=csv -l 10 > "$PERF_REPORTS_DIR/gpu_telemetry.csv" 2>/dev/null &
+telemetry_pid=$!
+cleanup() {
+  kill "$telemetry_pid" 2>/dev/null || true
+  wait "$telemetry_pid" 2>/dev/null || true
+}
+trap cleanup EXIT INT TERM
+
+pytest ./fastvideo/tests/performance -vs
+pytest_rc=$?
+compare_rc=0
+if [ "$pytest_rc" -eq 0 ] || [ "$PERF_UPLOAD_POLICY" = always ]; then
+  PERF_PYTEST_RC=$pytest_rc python ./fastvideo/tests/performance/compare_baseline.py
+  compare_rc=$?
+fi
+python ./fastvideo/tests/performance/dashboard.py || true
+cp -f fastvideo/tests/performance/results/*.json "$PERF_REPORTS_DIR/" 2>/dev/null || true
+
+echo "--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---"
+cat "$PERF_REPORTS_DIR/gpu_telemetry.csv" || true
+
+final_rc=$pytest_rc
+if [ "$final_rc" -eq 0 ]; then
+  final_rc=$compare_rc
+fi
+exit "$final_rc"

.buildkite/scripts/lanes/self_forcing.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the self-forcing lane.
+set -euo pipefail
+
+export WANDB_MODE=offline
+exec pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs

.buildkite/scripts/lanes/ssim.sh

@@ -0,0 +1,205 @@
+#!/usr/bin/env bash
+# Canonical four-GPU SSIM lane for the Slinky Slurm worker.
+set -euo pipefail
+
+args=()
+if [ "${FASTVIDEO_SSIM_BOOTSTRAP_MODE:-0}" = 1 ]; then
+  args+=(--bootstrap-mode)
+fi
+selected=${FASTVIDEO_SSIM_TEST_FILES-}
+if [ -z "$selected" ]; then
+  if [ "${TEST_SCOPE:-}" = merge ]; then
+    echo "Missing FASTVIDEO_SSIM_TEST_FILES for merge scope" >&2
+    exit 2
+  fi
+  selected=all
+fi
+if [ "$selected" != all ]; then
+  [[ $selected =~ ^test_[a-z0-9_]+\.py(,test_[a-z0-9_]+\.py)*$ ]] || {
+    echo "Invalid FASTVIDEO_SSIM_TEST_FILES selection" >&2
+    exit 2
+  }
+  IFS=, read -r -a ssim_files <<< "$selected"
+  for ssim_file in "${ssim_files[@]}"; do
+    args+=(--test-file "$ssim_file")
+  done
+fi
+
+# This disposable branch defaults to a focused GameCraft FA4 candidate run.
+# The canonical lane remains available for local inspection by setting this to
+# zero. Always exit 2 in candidate mode so the run cannot be mistaken for a
+# passing quality gate and cannot trigger the production exit-1 retry policy.
+candidate_enabled=${FASTVIDEO_GAMECRAFT_FA4_CANDIDATE_LOG:-1}
+force_candidate_exit() {
+  trap - EXIT
+  exit 2
+}
+if [ "$candidate_enabled" = 1 ]; then
+  trap force_candidate_exit EXIT
+fi
+
+# MoGe's utils3d dependency builds glcontext from source on ARM64. The current
+# runner image predates the baked-in X11 headers below, so keep this guarded
+# bootstrap until every deployed image digest contains libx11-dev.
+if [ ! -f /usr/include/X11/Xlib.h ]; then
+  apt-get -o Acquire::Retries=5 update
+  apt-get -o Acquire::Retries=5 install -y --no-install-recommends libx11-dev
+  rm -rf /var/lib/apt/lists/*
+fi
+
+uv pip install git+https://github.com/microsoft/MoGe.git
+uv pip install k_diffusion einops_exts alias_free_torch torchsde
+
+if [ "$candidate_enabled" = 1 ]; then
+  checkout_sha=$(git rev-parse HEAD)
+  echo "+++ GameCraft T2V current-FA4 candidate on GB200"
+  echo "checkout_sha=$checkout_sha"
+  echo "buildkite_commit=${BUILDKITE_COMMIT:-<unset>}"
+  if [ -n "${BUILDKITE_COMMIT:-}" ] && [ "$checkout_sha" != "$BUILDKITE_COMMIT" ]; then
+    echo "Checkout does not match BUILDKITE_COMMIT" >&2
+    exit 2
+  fi
+
+  visible_gpus=${CUDA_VISIBLE_DEVICES:-0}
+  candidate_gpu=${visible_gpus%%,*}
+  candidate_log=/tmp/fastvideo-gamecraft-fa4-candidate-pytest.log
+  candidate_sentinel=$(mktemp /tmp/fastvideo-gamecraft-fa4-candidate.XXXXXX)
+  candidate_generated_root=fastvideo/tests/ssim/generated_videos/default
+  candidate_env=(
+    "CUDA_VISIBLE_DEVICES=$candidate_gpu"
+    "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False"
+    "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN"
+    "FASTVIDEO_FA4=1"
+    "FASTVIDEO_SSIM_BOOTSTRAP_MODE=0"
+    "FASTVIDEO_SSIM_FULL_QUALITY=0"
+  )
+
+  if env "${candidate_env[@]}" python - <<'PY'
+import os
+
+import torch
+
+print(f"requested_backend={os.environ['FASTVIDEO_ATTENTION_BACKEND']}")
+print(f"requested_FASTVIDEO_FA4={os.environ['FASTVIDEO_FA4']}")
+print(f"torch_version={torch.__version__}")
+print(f"torch_cuda_version={torch.version.cuda}")
+print(f"cuda_available={torch.cuda.is_available()}")
+print(f"cuda_visible_devices={os.environ.get('CUDA_VISIBLE_DEVICES', '<unset>')}")
+print(f"gpu_count={torch.cuda.device_count()}")
+
+if torch.cuda.device_count() != 1:
+    raise RuntimeError(f"Expected exactly one visible candidate GPU, got {torch.cuda.device_count()}")
+gpu_name = torch.cuda.get_device_name(0)
+print(f"gpu[0]={gpu_name} capability={torch.cuda.get_device_capability(0)}")
+if "GB200" not in gpu_name:
+    raise RuntimeError(f"Candidate must run on GB200, got {gpu_name}")
+
+from fastvideo.attention.selector import get_attn_backend
+from fastvideo.attention.utils.flash_attn_default import fa_version
+from fastvideo.platforms import AttentionBackendEnum
+
+requested_backend = AttentionBackendEnum[os.environ["FASTVIDEO_ATTENTION_BACKEND"]]
+resolved_backend = get_attn_backend(
+    128,
+    torch.bfloat16,
+    supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,),
+    requested=requested_backend,
+)
+print(f"resolved_backend={resolved_backend.get_name()}")
+print(f"resolved_flash_attention=FA{fa_version}")
+if fa_version != "4":
+    raise RuntimeError(f"Candidate must use FA4, resolved FA{fa_version}")
+PY
+  then
+    probe_rc=0
+  else
+    probe_rc=$?
+  fi
+  if [ "$probe_rc" -ne 0 ]; then
+    echo "GameCraft candidate environment probe failed with rc=$probe_rc" >&2
+    exit 2
+  fi
+
+  if env "${candidate_env[@]}" python -m pytest \
+      fastvideo/tests/ssim/test_gamecraft_similarity.py::test_gamecraft_t2v_similarity \
+      -vs >"$candidate_log" 2>&1
+  then
+    test_rc=0
+  else
+    test_rc=$?
+  fi
+
+  echo "+++ GameCraft candidate pytest tail"
+  python - "$candidate_log" <<'PY'
+import sys
+from pathlib import Path
+
+log_path = Path(sys.argv[1])
+data = log_path.read_bytes()
+print(f"pytest_log_bytes={len(data)}")
+tail = data[-131072:].decode("utf-8", errors="replace").replace("\r", "\n")
+lines = tail.splitlines()[-160:]
+for line in lines:
+    print(line[:2000])
+PY
+  echo "pytest_rc=$test_rc"
+
+  candidate_videos=()
+  if [ -d "$candidate_generated_root" ]; then
+    mapfile -d '' -t candidate_videos < <(
+      find "$candidate_generated_root" -type f -name '*.mp4' -newer "$candidate_sentinel" -print0
+    )
+  fi
+  if [ "${#candidate_videos[@]}" -ne 1 ]; then
+    echo "Expected exactly one newly generated candidate MP4; found ${#candidate_videos[@]}" >&2
+    printf 'candidate_path=%s\n' "${candidate_videos[@]}" >&2
+    exit 2
+  fi
+  candidate_video=${candidate_videos[0]}
+  case "$candidate_video" in
+    */default/GB200_reference_videos/HunyuanGameCraft-T2V/FLASH_ATTN/*.mp4) ;;
+    *)
+      echo "Candidate path is outside the expected GB200 T2V subtree: $candidate_video" >&2
+      exit 2
+      ;;
+  esac
+
+  candidate_results=()
+  mapfile -d '' -t candidate_results < <(
+    find "$(dirname "$candidate_video")" -type f -name '*_ssim.json' -newer "$candidate_sentinel" -print0
+  )
+  if [ "${#candidate_results[@]}" -ne 1 ]; then
+    echo "Expected exactly one newly generated SSIM JSON; found ${#candidate_results[@]}" >&2
+    exit 2
+  fi
+  candidate_result=${candidate_results[0]}
+
+  python - "$candidate_video" "$candidate_result" <<'PY'
+import json
+import sys
+from pathlib import Path
+
+video_path = Path(sys.argv[1]).resolve()
+result_path = Path(sys.argv[2])
+result = json.loads(result_path.read_text())
+if Path(result["generated_video"]).resolve() != video_path:
+    raise RuntimeError("SSIM JSON does not describe the candidate MP4")
+if result["parameters"]["num_inference_steps"] != 20:
+    raise RuntimeError("Candidate did not use the expected 20 inference steps")
+print(f"candidate_path={video_path}")
+print(f"candidate_ssim_mean={result['mean_ssim']}")
+print(f"candidate_ssim_min={result['min_ssim']}")
+print(f"candidate_ssim_max={result['max_ssim']}")
+print(f"candidate_prompt={result['parameters']['prompt']}")
+PY
+
+  echo "FV_GAMECRAFT_FA4_SSIM_JSON_BEGIN"
+  cat "$candidate_result"
+  echo "FV_GAMECRAFT_FA4_SSIM_JSON_END"
+  echo "+++ GameCraft candidate MP4 log envelope"
+  python .buildkite/scripts/gamecraft_candidate_log.py emit "$candidate_video"
+  echo "--- GameCraft candidate complete: probe_rc=$probe_rc pytest_rc=$test_rc (forced lane rc=2)"
+  exit 2
+fi
+
+exec python fastvideo/tests/ssim/ci_runner.py "${args[@]}"

.buildkite/scripts/lanes/train_framework.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the modular training-framework lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs

.buildkite/scripts/lanes/training.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the legacy vanilla-training lane.
+set -euo pipefail
+
+export WANDB_MODE=offline
+exec pytest ./fastvideo/tests/training/Vanilla -srP

.buildkite/scripts/lanes/training_lora.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the legacy LoRA-training lane.
+set -euo pipefail
+
+export WANDB_MODE=offline
+exec pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP

.buildkite/scripts/lanes/training_vsa.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the legacy VSA-training lane.
+set -euo pipefail
+
+export WANDB_MODE=offline
+exec pytest ./fastvideo/tests/training/VSA -srP

.buildkite/scripts/lanes/transformer.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the transformer lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/transformers -vs

.buildkite/scripts/lanes/vae.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Canonical Slurm CI selection for the VAE lane.
+set -euo pipefail
+
+exec pytest ./fastvideo/tests/vaes -vs

.buildkite/scripts/pr_test.sh

@@ -1,6 +1,19 @@
 #!/bin/bash
 set -uo pipefail
 
+# DORMANT ROLLBACK ONLY. Active CI is Slurm-only and pipeline.yml never calls
+# this launcher. Refuse every Buildkite invocation even if a stale step or
+# operator typo reaches this file; local rollback experiments require an
+# explicit opt-in.
+if [ -n "${BUILDKITE:-}" ]; then
+  echo "Legacy Modal CI is disabled; use the Slinky Slurm runner." >&2
+  exit 2
+fi
+if [ "${FASTVIDEO_ENABLE_LEGACY_MODAL_CI:-0}" != 1 ]; then
+  echo "Legacy Modal CI is dormant. Set FASTVIDEO_ENABLE_LEGACY_MODAL_CI=1 only for a manual rollback test." >&2
+  exit 2
+fi
+
 log() {
   echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
 }

.buildkite/scripts/unit_test.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+exec pytest \
+  ./fastvideo/tests/api/ \
+  ./fastvideo/tests/contract/ \
+  ./fastvideo/tests/dataset/ \
+  ./fastvideo/tests/workflow/ \
+  ./fastvideo/tests/entrypoints/ \
+  ./fastvideo/tests/loader/ \
+  ./fastvideo/tests/pipelines/ \
+  ./fastvideo/tests/platforms/ \
+  ./fastvideo/tests/train/ \
+  ./fastvideo/tests/stages/ \
+  ./fastvideo/tests/ops/ \
+  ./fastvideo/tests/worker/ \
+  ./fastvideo/tests/training/test_trackers.py \
+  ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \
+  ./fastvideo/tests/modal/test_kernel_build_cache.py \
+  ./fastvideo/tests/modal/test_pr_test.py \
+  ./fastvideo/tests/modal/test_ssim_test.py \
+  --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py \
+  --ignore=./fastvideo/tests/train/models \
+  --ignore=./fastvideo/tests/train/methods \
+  -vs

.github/PULL_REQUEST_TEMPLATE.md

@@ -8,10 +8,10 @@ PR TITLE: Must start with a type tag, e.g.:
 MERGE WORKFLOW:
   1. Ensure pre-commit passes and you have at least 1 approval
   2. Comment /merge (or add the "ready" label) to enter the Merge Queue
-  3. Full Test Suite runs automatically on a staging branch → auto-merge on success
+  3. A path-aware merge gate runs only relevant integration tests → auto-merge on success
 
 ON-DEMAND TESTING (write access required):
-  /test full       — Full Test Suite        /test ssim        — SSIM regression
+  /test full       — Explicit all-lane run  /test ssim        — Full SSIM regression
   /test training   — Training pipeline      /test encoder     — Encoder tests
   /test transformer — Transformer tests     /test vae         — VAE tests
   /test kernel     — CUDA kernel tests      /test unit        — Unit tests

.github/scripts/gate_full_suite.sh

@@ -1,14 +1,14 @@
 #!/usr/bin/env bash
-# Gate the expensive Buildkite full suite on the cheap GitHub checks.
+# Gate the path-aware Buildkite merge plan on the cheap GitHub checks.
 #
 # Polls the workflow runs for the PR head commit and only exits 0 once the
 # watched cheap workflows (pre-commit, docs build) have succeeded, so the
-# 'ready' label cannot burn ~20 GPU lanes on a head that a cheap check has
-# already doomed.
+# 'ready' label cannot burn path-selected GPU lanes on a head that a cheap
+# check has already doomed.
 #
 # Semantics:
 #   - watched run completed with a bad conclusion  -> exit 1 (fail CLOSED:
-#     no full suite; the next push re-arms via the 'synchronize' trigger)
+#     no merge gate; the next push re-arms via the 'synchronize' trigger)
 #   - watched run cancelled                         -> still pending: the docs
 #     workflow's repo-global 'pages' concurrency group cancels runs superseded
 #     by unrelated pushes, so 'cancelled' is not a verdict on this PR
@@ -29,7 +29,7 @@ set -euo pipefail
 : "${PR_NUMBER:?PR_NUMBER (pull request number) is required}"
 : "${GITHUB_REPOSITORY:?GITHUB_REPOSITORY is required}"
 
-# Workflow-level `name:` values that must be green before the full suite
+# Workflow-level `name:` values that must be green before the merge gate
 # may start. "Deploy Documentation" is path-filtered on PRs, so its run may
 # legitimately never exist; pre-commit always runs, so it must appear.
 WATCHED_NAMES='["pre-commit", "Deploy Documentation"]'
@@ -56,7 +56,7 @@ recheck_ready_label() {
   if pr_json=$(gh_api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" 2>/dev/null); then
     if ! jq -e '[.labels[]?.name] | index("ready")' <<<"$pr_json" >/dev/null 2>&1; then
       echo "::error::PR #${PR_NUMBER} no longer has the 'ready' label —" \
-        "NOT triggering the Buildkite full suite. Re-add the label to re-arm."
+        "NOT triggering the Buildkite merge gate. Re-add the label to re-arm."
       exit 1
     fi
   else
@@ -84,7 +84,7 @@ while true; do
         | map(.name) | join(", ")' <<<"$state")
     if [ -n "$failed" ]; then
       echo "::error::Cheap check(s) failed on ${PR_SHA}: ${failed}." \
-        "NOT triggering the Buildkite full suite. Push a fix (the 'ready'" \
+        "NOT triggering the Buildkite merge gate. Push a fix (the 'ready'" \
         "label re-arms on every push), or re-run the failed check and then" \
         "re-run this workflow."
       exit 1
@@ -97,7 +97,7 @@ while true; do
     if [ "$pending" -eq 0 ]; then
       if [ -z "$missing" ]; then
         recheck_ready_label
-        echo "All watched cheap checks are green — full suite may proceed."
+        echo "All watched cheap checks are green — merge gate may proceed."
         exit 0
       fi
       case "$missing" in
@@ -119,14 +119,14 @@ while true; do
     echo "::warning::GitHub API error querying workflow runs for ${PR_SHA} (attempt ${api_fails}/3)."
     if [ "$api_fails" -ge 3 ]; then
       recheck_ready_label
-      echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the full suite WITHOUT the cheap-check gate."
+      echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the merge gate WITHOUT the cheap-check gate."
       exit 0
     fi
   fi
 
   if [ "$elapsed" -ge "$MAX_WAIT_SECS" ]; then
     recheck_ready_label
-    echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the full suite anyway."
+    echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the merge gate anyway."
     exit 0
   fi
   sleep "$POLL_SECS"

.github/scripts/plan_merge_ci.py

@@ -0,0 +1,570 @@
+#!/usr/bin/env python3
+"""Select the additive GPU integration lanes needed by a PR diff.
+
+Fastcheck is the universal six-lane baseline and is intentionally not repeated
+here.  This planner selects only the more expensive merge-gate lanes.  Unknown
+source/build paths fail closed to the complete integration set, while explicit
+documentation and repository-metadata paths require no additional GPU work.
+"""
+
+from __future__ import annotations
+
+import argparse
+import fnmatch
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TextIO
+
+MERGE_LANES = (
+    "golden-gate",
+    "ssim",
+    "lora-inference",
+    "lora-extraction",
+    "training",
+    "distillation",
+    "self-forcing",
+    "lora-training",
+    "training-vsa",
+    "inference-vmoba",
+    "performance",
+    "api-server",
+    "train-framework",
+    "eval",
+)
+
+LANE_SCRIPT_TO_KEY = {
+    "api_server.sh": "api-server",
+    "distillation_dmd.sh": "distillation",
+    "eval.sh": "eval",
+    "golden_gate.sh": "golden-gate",
+    "inference_lora.sh": "lora-inference",
+    "inference_vmoba.sh": "inference-vmoba",
+    "lora_extraction.sh": "lora-extraction",
+    "performance.sh": "performance",
+    "self_forcing.sh": "self-forcing",
+    "ssim.sh": "ssim",
+    "train_framework.sh": "train-framework",
+    "training.sh": "training",
+    "training_lora.sh": "lora-training",
+    "training_vsa.sh": "training-vsa",
+}
+
+FASTCHECK_LANE_SCRIPTS = {
+    "dreamverse.sh",
+    "encoder.sh",
+    "kernel_tests.sh",
+    "transformer.sh",
+    "vae.sh",
+}
+
+LEGACY_TRAINING_LANES = (
+    "training",
+    "distillation",
+    "self-forcing",
+    "lora-training",
+    "training-vsa",
+)
+
+ALL_TRAINING_LANES = (*LEGACY_TRAINING_LANES, "train-framework")
+
+SSIM_SMOKE_TESTS = (
+    "test_flux_t2i_similarity.py",
+    "test_wan_t2v_similarity.py",
+)
+
+SAFE_PATTERNS = (
+    "*.md",
+    "*.rst",
+    ".agents/**",
+    ".claude/**",
+    ".codex/**",
+    ".github/ISSUE_TEMPLATE/**",
+    ".github/PULL_REQUEST_TEMPLATE.md",
+    ".github/dependabot.yml",
+    ".github/mergify.yml",
+    ".github/scripts/**",
+    ".github/workflows/**",
+    ".buildkite/scripts/pre_commit.sh",
+    ".git-blame-ignore-revs",
+    ".gitattributes",
+    ".gitignore",
+    ".pre-commit-config.yaml",
+    "AGENTS.md",
+    "CITATION.cff",
+    "CODE_OF_CONDUCT.md",
+    "CONTRIBUTING.md",
+    "LICENSE",
+    "NOTICE",
+    "__init__.py",
+    "collect_env.py",
+    "SECURITY.md",
+    "assets/**",
+    "comfyui/**",
+    "docs/**",
+    "examples/**",
+    "mkdocs.yml",
+    "requirements-mkdocs.in",
+    "requirements-mkdocs.txt",
+    "scripts/**",
+    "tests/__init__.py",
+    "tests/local_tests/**",
+)
+
+ALL_IMPACT_PATTERNS = (
+    ".buildkite/pipeline.yml",
+    "docker/**",
+    "pyproject.toml",
+    "requirements*.txt",
+    "setup.cfg",
+    "setup.py",
+    "uv.lock",
+)
+
+
+@dataclass(frozen=True)
+class FamilyCoverage:
+    pattern: re.Pattern[str]
+    golden_tests: tuple[str, ...]
+    ssim_tests: tuple[str, ...]
+
+
+FAMILY_COVERAGE = (
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])dreamx(_world)?([/_.-]|$)"),
+        ("test_dreamx.py", ),
+        ("test_dreamx_world_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"),
+        ("test_flux2_klein.py", ),
+        ("test_flux2_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])flux(?![_-]?2)([/_.-]|$)"),
+        ("test_flux.py", ),
+        ("test_flux_t2i_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])(hunyuan)?gamecraft([/_.-]|$)"),
+        ("test_gamecraft.py", ),
+        ("test_gamecraft_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])gen3c([/_.-]|$)"),
+        ("test_gen3c.py", ),
+        ("test_gen3c_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])glm[_-]?image([/_.-]|$)"),
+        ("test_glm_image.py", ),
+        ("test_glm_image_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])kandinsky[_-]?5([/_.-]|$)"),
+        ("test_kandinsky5.py", ),
+        ("test_kandinsky5_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])lingbot([a-z0-9_-]*)([/_.-]|$)"),
+        ("test_lingbot.py", ),
+        ("test_lingbot_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])longcat([/_.-]|$)"),
+        ("test_longcat.py", ),
+        ("test_longcat_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])ltx[_-]?2([/_.-]|$)"),
+        ("test_ltx2.py", ),
+        ("test_ltx2_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])matrixgame[_-]?2([/_.-]|$)"),
+        ("test_matrixgame.py", ),
+        ("test_matrixgame2_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])matrixgame[_-]?3([/_.-]|$)"),
+        ("test_matrixgame.py", ),
+        ("test_matrixgame3_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])minimax[_-]?h3([/_.-]|$)"),
+        ("test_minimax_h3_t2v.py", ),
+        ("test_minimax_h3_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])sd[_-]?3([._-]?5)?([/_.-]|$)"),
+        ("test_sd35.py", ),
+        ("test_sd35_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])stable[_-]?audio([/_.-]|$)"),
+        ("test_stable_audio.py", ),
+        ("test_stable_audio_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])turbo(diffusion)?([/_.-]|$)"),
+        (),
+        ("test_turbodiffusion_similarity.py", ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])wan(video)?([/_.-]|$)"),
+        ("test_wan_t2v.py", ),
+        (
+            "test_causal_similarity.py",
+            "test_wan_i2v_similarity.py",
+            "test_wan_t2v_similarity.py",
+        ),
+    ),
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])z[_-]?image([/_.-]|$)"),
+        ("test_zimage.py", ),
+        ("test_zimage_similarity.py", ),
+    ),
+)
+
+
+@dataclass
+class MergePlan:
+    lanes: set[str] = field(default_factory=set)
+    golden_tests: set[str] = field(default_factory=set)
+    ssim_tests: set[str] = field(default_factory=set)
+    golden_all: bool = False
+    ssim_all: bool = False
+    reasons: list[str] = field(default_factory=list)
+
+    def add_lanes(self, *lanes: str, reason: str) -> None:
+        unknown = set(lanes) - set(MERGE_LANES)
+        if unknown:
+            raise ValueError(f"Unknown merge lanes: {sorted(unknown)}")
+        self.lanes.update(lanes)
+        self.reasons.append(reason)
+
+    def add_golden(self, tests: tuple[str, ...], reason: str) -> None:
+        self.add_lanes("golden-gate", reason=reason)
+        self.golden_tests.update(tests)
+
+    def add_ssim(self, tests: tuple[str, ...], reason: str) -> None:
+        self.add_lanes("ssim", reason=reason)
+        self.ssim_tests.update(tests)
+
+    def require_all(self, reason: str) -> None:
+        self.lanes.update(MERGE_LANES)
+        self.golden_all = True
+        self.ssim_all = True
+        self.reasons.append(reason)
+
+    def ordered_lanes(self) -> tuple[str, ...]:
+        return tuple(lane for lane in MERGE_LANES if lane in self.lanes)
+
+    def encoded_lanes(self) -> str:
+        lanes = self.ordered_lanes()
+        return "," + ",".join(lanes or ("none", )) + ","
+
+    def encoded_golden_tests(self) -> str:
+        if "golden-gate" not in self.lanes:
+            return "none"
+        if self.golden_all or not self.golden_tests:
+            return "all"
+        return ",".join(sorted(self.golden_tests))
+
+    def encoded_ssim_tests(self) -> str:
+        if "ssim" not in self.lanes:
+            return "none"
+        if self.ssim_all or not self.ssim_tests:
+            return "all"
+        return ",".join(sorted(self.ssim_tests))
+
+
+def _matches_any(path: str, patterns: tuple[str, ...]) -> bool:
+    return any(fnmatch.fnmatchcase(path, pattern) for pattern in patterns)
+
+
+def _family_coverage(path: str) -> tuple[set[str], set[str]]:
+    normalized = path.lower()
+    golden: set[str] = set()
+    ssim: set[str] = set()
+    for family in FAMILY_COVERAGE:
+        if family.pattern.search(normalized):
+            golden.update(family.golden_tests)
+            ssim.update(family.ssim_tests)
+    return golden, ssim
+
+
+def _select_output_coverage(plan: MergePlan, path: str) -> None:
+    golden, ssim = _family_coverage(path)
+    if golden:
+        plan.add_golden(tuple(sorted(golden)), reason=f"model-family golden coverage: {path}")
+    else:
+        plan.golden_all = True
+        plan.add_lanes("golden-gate", reason=f"shared output golden coverage: {path}")
+    if ssim:
+        plan.add_ssim(tuple(sorted(ssim)), reason=f"model-family SSIM coverage: {path}")
+    else:
+        plan.add_ssim(SSIM_SMOKE_TESTS, reason=f"shared output SSIM smoke coverage: {path}")
+
+
+def classify_paths(paths: list[str]) -> MergePlan:
+    plan = MergePlan()
+    normalized_paths: list[str] = []
+    for raw_path in paths:
+        path = raw_path.strip()
+        while path.startswith("./"):
+            path = path[2:]
+        if path:
+            normalized_paths.append(path)
+    normalized_paths = sorted(set(normalized_paths))
+    if not normalized_paths:
+        plan.require_all("changed-file list was empty; failing closed")
+        return plan
+
+    for path in normalized_paths:
+        if path == "__FASTVIDEO_CI_PLAN_ALL__":
+            plan.require_all("changed-file API failed; failing closed")
+            continue
+
+        if path in {"requirements-mkdocs.in", "requirements-mkdocs.txt"}:
+            plan.reasons.append(f"documentation dependencies need no GPU integration: {path}")
+            continue
+
+        if _matches_any(path, ALL_IMPACT_PATTERNS):
+            plan.require_all(f"cross-cutting build/runtime surface: {path}")
+            continue
+
+        lane_script_prefix = ".buildkite/scripts/lanes/"
+        if path.startswith(lane_script_prefix):
+            script_name = Path(path).name
+            lane = LANE_SCRIPT_TO_KEY.get(script_name)
+            if lane is None:
+                if script_name in FASTCHECK_LANE_SCRIPTS:
+                    plan.reasons.append(f"covered by automatic Fastcheck lane: {path}")
+                else:
+                    plan.require_all(f"unknown lane script: {path}")
+            elif lane == "golden-gate":
+                plan.golden_all = True
+                plan.add_lanes(lane, reason=f"golden lane implementation: {path}")
+            elif lane == "ssim":
+                plan.ssim_all = True
+                plan.add_lanes(lane, reason=f"SSIM lane implementation: {path}")
+            else:
+                plan.add_lanes(lane, reason=f"lane implementation: {path}")
+            continue
+
+        if path.startswith("fastvideo/tests/golden_gate/"):
+            name = Path(path).name
+            if name.startswith("test_") and name.endswith(".py"):
+                plan.add_golden((name, ), reason=f"changed golden test: {path}")
+            elif name in {"AGENTS.md", "README.md"}:
+                plan.reasons.append(f"golden documentation only: {path}")
+            else:
+                plan.golden_all = True
+                plan.add_lanes("golden-gate", reason=f"shared golden harness/reference: {path}")
+            continue
+
+        if path.startswith("fastvideo/tests/ssim/"):
+            name = Path(path).name
+            if name.startswith("test_") and name.endswith(".py"):
+                plan.add_ssim((name, ), reason=f"changed SSIM test: {path}")
+            elif path.endswith((".py", ".json", ".pt", ".png", ".mp4")):
+                plan.ssim_all = True
+                plan.add_lanes("ssim", reason=f"shared SSIM harness/reference: {path}")
+            continue
+
+        if path.startswith("fastvideo/tests/performance/") or path.startswith(".buildkite/performance-benchmarks/"):
+            plan.add_lanes("performance", reason=f"performance coverage: {path}")
+            continue
+        if path.startswith(("fastvideo/performance/", "fastvideo/performance_dashboard/",
+                            "apps/performance_dashboard/")):
+            plan.add_lanes("performance", reason=f"performance implementation: {path}")
+            continue
+        if path.startswith("fastvideo/benchmarks/"):
+            if "/mlx_" in path or Path(path).name.startswith("mlx_"):
+                plan.reasons.append(f"covered by the path-filtered macOS MLX workflow: {path}")
+            else:
+                plan.add_lanes("performance", reason=f"benchmark implementation: {path}")
+            continue
+        if path.startswith("fastvideo/tests/eval/") or path.startswith("fastvideo/eval/"):
+            plan.add_lanes("eval", reason=f"evaluation coverage: {path}")
+            continue
+        if path.startswith("fastvideo/third_party/eval/"):
+            plan.add_lanes("eval", reason=f"vendored evaluation implementation: {path}")
+            continue
+        if path.startswith("fastvideo/tests/lora_extraction/") or path.startswith("scripts/lora_extraction/"):
+            plan.add_lanes("lora-extraction", reason=f"LoRA extraction coverage: {path}")
+            continue
+        if path.startswith("fastvideo/tests/inference/lora/"):
+            plan.add_lanes("lora-inference", reason=f"LoRA inference coverage: {path}")
+            continue
+        if path.startswith("fastvideo/tests/inference/vmoba/"):
+            plan.add_lanes("inference-vmoba", reason=f"VMoBA inference coverage: {path}")
+            continue
+        if path.startswith(("fastvideo/dataset/", "fastvideo/workflow/", "fastvideo/pipelines/preprocess/",
+                            "fastvideo/pipelines/training/")):
+            plan.add_lanes(*ALL_TRAINING_LANES, reason=f"shared data/training input surface: {path}")
+            continue
+        if path.startswith("fastvideo/tests/train/") or path.startswith("fastvideo/train/"):
+            plan.add_lanes("train-framework", reason=f"modular training coverage: {path}")
+            continue
+
+        if path.startswith("fastvideo/tests/training/"):
+            lowered = path.lower()
+            if "/vanilla/" in lowered:
+                plan.add_lanes("training", reason=f"vanilla training coverage: {path}")
+            elif "/distill/" in lowered:
+                plan.add_lanes("distillation", reason=f"distillation coverage: {path}")
+            elif "/self-forcing/" in lowered:
+                plan.add_lanes("self-forcing", reason=f"self-forcing coverage: {path}")
+            elif "/lora/" in lowered:
+                plan.add_lanes("lora-training", reason=f"LoRA training coverage: {path}")
+            elif "/vsa/" in lowered:
+                plan.add_lanes("training-vsa", reason=f"VSA training coverage: {path}")
+            else:
+                plan.add_lanes(*LEGACY_TRAINING_LANES, reason=f"shared legacy training coverage: {path}")
+            continue
+
+        if path.startswith("fastvideo/training/"):
+            lowered = path.lower()
+            if "self_forcing" in lowered:
+                plan.add_lanes("self-forcing", reason=f"self-forcing implementation: {path}")
+            elif "distill" in lowered:
+                plan.add_lanes("distillation", reason=f"distillation implementation: {path}")
+            elif "lora" in lowered:
+                plan.add_lanes("lora-training", reason=f"LoRA training implementation: {path}")
+            else:
+                plan.add_lanes(*LEGACY_TRAINING_LANES, reason=f"shared legacy training implementation: {path}")
+            continue
+
+        lowered = path.lower()
+        if "vmoba" in lowered and path.startswith(("fastvideo/", ".buildkite/")):
+            plan.add_lanes("inference-vmoba", reason=f"VMoBA implementation: {path}")
+            plan.add_golden(("test_wan_t2v.py", ), reason=f"VMoBA end-to-end coverage: {path}")
+            continue
+        if "lora" in lowered and path.startswith("fastvideo/"):
+            plan.add_lanes(
+                "lora-inference",
+                "lora-extraction",
+                "lora-training",
+                reason=f"shared LoRA implementation: {path}",
+            )
+            _select_output_coverage(plan, path)
+            continue
+
+        if path.startswith("fastvideo/entrypoints/") or path.startswith("fastvideo/api/"):
+            plan.add_lanes("api-server", reason=f"API/entrypoint integration: {path}")
+            if "openai" not in lowered and "/cli/" not in lowered:
+                _select_output_coverage(plan, path)
+            continue
+        if path.startswith("fastvideo/worker/"):
+            plan.add_lanes("api-server", reason=f"worker/API integration: {path}")
+            _select_output_coverage(plan, path)
+            continue
+        if path.startswith("fastvideo/distributed/"):
+            plan.add_lanes(
+                "training",
+                "train-framework",
+                reason=f"distributed runtime integration: {path}",
+            )
+            _select_output_coverage(plan, path)
+            continue
+        if path.startswith(("fastvideo/hooks/", "fastvideo/platforms/", "fastvideo/third_party/")):
+            _select_output_coverage(plan, path)
+            continue
+        if path.startswith(("fastvideo/models/", "fastvideo/pipelines/", "fastvideo/configs/",
+                            "fastvideo/layers/", "fastvideo/attention/")):
+            _select_output_coverage(plan, path)
+            continue
+        if path in {
+                "fastvideo/fastvideo_args.py",
+                "fastvideo/forward_context.py",
+                "fastvideo/image_processor.py",
+                "fastvideo/registry.py",
+                "fastvideo/utils.py",
+        }:
+            _select_output_coverage(plan, path)
+            continue
+        if path.startswith("fastvideo/mlx_runtime/"):
+            plan.reasons.append(f"covered by the path-filtered macOS MLX workflow: {path}")
+            continue
+        if path.startswith("fastvideo/logging_utils/") or path in {
+                "fastvideo/__init__.py",
+                "fastvideo/envs.py",
+                "fastvideo/logger.py",
+                "fastvideo/profiler.py",
+                "fastvideo/version.py",
+        }:
+            plan.reasons.append(f"covered by automatic Fastcheck: {path}")
+            continue
+        if path.startswith(("fastvideo-kernel/", "csrc/")):
+            plan.add_golden(("test_wan_t2v.py", ), reason=f"kernel integration smoke: {path}")
+            plan.add_ssim(("test_wan_t2v_similarity.py", ), reason=f"kernel numerical smoke: {path}")
+            continue
+
+        if path.startswith("apps/dreamverse/"):
+            # DreamVerse is already one of the six automatic Fastcheck lanes.
+            plan.reasons.append(f"covered by automatic DreamVerse Fastcheck: {path}")
+            continue
+        if path.startswith("fastvideo/tests/"):
+            # The automatic unit/component Fastcheck lanes own the remaining
+            # package tests. Domain-specific expensive test roots were handled
+            # above.
+            plan.reasons.append(f"covered by automatic Fastcheck: {path}")
+            continue
+        if path in {".buildkite/scripts/unit_test.sh", ".buildkite/scripts/pr_test.sh"}:
+            plan.reasons.append(f"covered by automatic unit Fastcheck: {path}")
+            continue
+        if _matches_any(path, SAFE_PATTERNS):
+            plan.reasons.append(f"no additional GPU integration needed: {path}")
+            continue
+
+        plan.require_all(f"unclassified path; failing closed: {path}")
+
+    return plan
+
+
+def _write_github_output(output: TextIO, plan: MergePlan) -> None:
+    output.write(f"merge_test_plan={plan.encoded_lanes()}\n")
+    output.write(f"merge_golden_tests={plan.encoded_golden_tests()}\n")
+    output.write(f"merge_ssim_tests={plan.encoded_ssim_tests()}\n")
+    output.write(f"merge_plan_label={','.join(plan.ordered_lanes()) or 'none'}\n")
+
+
+def _write_summary(output: TextIO, plan: MergePlan) -> None:
+    output.write("## Change-aware merge test plan\n\n")
+    output.write("| Selection | Value |\n|---|---|\n")
+    output.write(f"| Additional Slurm lanes | `{','.join(plan.ordered_lanes()) or 'none'}` |\n")
+    output.write(f"| Golden tests | `{plan.encoded_golden_tests()}` |\n")
+    output.write(f"| SSIM tests | `{plan.encoded_ssim_tests()}` |\n\n")
+    output.write("Fastcheck remains the universal six-lane baseline.\n")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--paths-file", type=Path, required=True)
+    parser.add_argument("--github-output", type=Path)
+    parser.add_argument("--summary-file", type=Path)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    paths = args.paths_file.read_text(encoding="utf-8").splitlines()
+    plan = classify_paths(paths)
+    print(f"MERGE_TEST_PLAN={plan.encoded_lanes()}")
+    print(f"MERGE_GOLDEN_TESTS={plan.encoded_golden_tests()}")
+    print(f"MERGE_SSIM_TESTS={plan.encoded_ssim_tests()}")
+    for reason in plan.reasons:
+        print(f"- {reason}")
+    if args.github_output:
+        with args.github_output.open("a", encoding="utf-8") as output:
+            _write_github_output(output, plan)
+    if args.summary_file:
+        with args.summary_file.open("a", encoding="utf-8") as output:
+            _write_summary(output, plan)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.github/scripts/test_gate_full_suite.sh

@@ -53,7 +53,7 @@ PC_PENDING='{"name": "pre-commit", "id": 1, "status": "in_progress", "conclusion
 DOCS_OK='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "success"}'
 DOCS_BAD='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "failure"}'
 DOCS_CANCELLED='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "cancelled"}'
-OTHER='{"name": "Trigger Full Suite", "id": 3, "status": "in_progress", "conclusion": null}'
+OTHER='{"name": "Trigger Merge Gate", "id": 3, "status": "in_progress", "conclusion": null}'
 NULL_NAME='{"name": null, "id": 4, "status": "completed", "conclusion": "failure"}'
 PC_OK_RERUN='{"name": "pre-commit", "id": 5, "status": "completed", "conclusion": "success"}'
 

.github/workflows/_template-build-image.yml

@@ -190,6 +190,7 @@ jobs:
         if: ${{ !inputs.push_by_digest }}
         run: |
           echo "✅ Python ${{ inputs.python_version }} image successfully built and pushed to ${{ steps.image.outputs.name }}:${{ inputs.tag_suffix }}-sha-${GITHUB_SHA::7}"
+          echo "Digest: ${{ steps.build-push.outputs.digest }}"
           echo "To run tests with this image, manually trigger the 'Run Tests' workflow."
 
       - name: Digest success message

.github/workflows/ci-aggregate-status.yml

@@ -26,29 +26,48 @@ jobs:
               per_page: 100,
             });
 
-            const bkStatuses = data.statuses.filter(
-              s => s.context.startsWith('buildkite/ci/')
-            );
-
-            const FASTCHECK_PREFIX = 'buildkite/ci/microscope-';
+            // Buildkite derives the GitHub context prefix from the label emoji.
+            // Keep hard Full Suite lanes in test-tube/bar-chart namespaces and
+            // Fastcheck lanes in microscope so targeted reruns cannot clear the
+            // wrong aggregate status. Automatic PR jobs use pr-fastcheck while
+            // slash-command and Full Suite jobs use ci; normalize the suffix
+            // and keep the newest status for each logical lane.
+            const FASTCHECK_PREFIXES = [
+              'buildkite/pr-fastcheck/microscope-',
+              'buildkite/ci/microscope-',
+            ];
             const FULL_SUITE_PREFIXES = [
               'buildkite/ci/test-tube-',
               'buildkite/ci/bar-chart-',
             ];
 
-            const fastcheck = bkStatuses.filter(
-              s => s.context.startsWith(FASTCHECK_PREFIX)
-            );
-            const fullSuite = bkStatuses.filter(
-              s => FULL_SUITE_PREFIXES.some(p => s.context.startsWith(p))
-            );
+            function newestByLane(prefixes) {
+              const statuses = new Map();
+              for (const status of data.statuses) {
+                const prefix = prefixes.find(p => status.context.startsWith(p));
+                if (!prefix) continue;
+                const lane = status.context.slice(prefix.length);
+                const previous = statuses.get(lane);
+                if (!previous || Date.parse(status.updated_at) > Date.parse(previous.updated_at)) {
+                  statuses.set(lane, status);
+                }
+              }
+              return statuses;
+            }
 
-            if (
-              fastcheck.length > 0
-              && fastcheck.every(s => s.state === 'success')
-            ) {
+            const fastcheck = newestByLane(FASTCHECK_PREFIXES);
+            const fullSuiteOnly = newestByLane(FULL_SUITE_PREFIXES);
+            const fastcheckPassed =
+              fastcheck.size === 6
+              && [...fastcheck.values()].every(s => s.state === 'success');
+            const fullSuitePassed =
+              fastcheckPassed
+              && fullSuiteOnly.size === 14
+              && [...fullSuiteOnly.values()].every(s => s.state === 'success');
+
+            if (fastcheckPassed) {
               core.info(
-                `All ${fastcheck.length} fastcheck tests passed — updating fastcheck-passed`
+                `All ${fastcheck.size} fastcheck tests passed — updating fastcheck-passed`
               );
               await github.rest.repos.createCommitStatus({
                 owner: context.repo.owner,
@@ -56,17 +75,13 @@ jobs:
                 sha,
                 state: 'success',
                 context: 'fastcheck-passed',
-                description:
-                  `All ${fastcheck.length} fastcheck tests passed`,
+                description: `All ${fastcheck.size} fastcheck tests passed`,
               });
             }
 
-            if (
-              fullSuite.length > 0
-              && fullSuite.every(s => s.state === 'success')
-            ) {
+            if (fullSuitePassed) {
               core.info(
-                `All ${fullSuite.length} full suite tests passed — updating full-suite-passed`
+                'All 20 full suite tests passed — updating full-suite-passed'
               );
               await github.rest.repos.createCommitStatus({
                 owner: context.repo.owner,
@@ -74,7 +89,6 @@ jobs:
                 sha,
                 state: 'success',
                 context: 'full-suite-passed',
-                description:
-                  `All ${fullSuite.length} full suite tests passed`,
+                description: 'All 20 full suite tests passed',
               });
             }

.github/workflows/ci-macos-mlx.yml

@@ -0,0 +1,151 @@
+name: macOS MLX Smoke
+
+on:
+  pull_request:
+    branches: [main]
+    paths:
+      - ".github/workflows/ci-macos-mlx.yml"
+      - "fastvideo/mlx_runtime/**"
+      - "fastvideo/tests/mlx/**"
+      - "fastvideo/tests/platforms/test_mps_vsa_error.py"
+      - "fastvideo/platforms/mps.py"
+      - "fastvideo/platforms/__init__.py"
+      - "fastvideo/__init__.py"
+      - "examples/inference/basic/mlx_*.py"
+      - "fastvideo/benchmarks/mlx_*.py"
+      - "pyproject.toml"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: macos-mlx-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  mlx-smoke:
+    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true
+    runs-on: macos-15
+    timeout-minutes: 25
+    env:
+      FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA
+      TOKENIZERS_PARALLELISM: "false"
+      MASTER_ADDR: localhost
+      MASTER_PORT: "29513"
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+
+      - uses: astral-sh/setup-uv@v3
+
+      - name: Install lightweight MLX smoke dependencies
+        run: |
+          uv pip install --system \
+            --index-url https://download.pytorch.org/whl/cpu \
+            torch==2.11.0 torchvision torchaudio
+          uv pip install --system \
+            pytest numpy scipy pillow imageio einops cloudpickle filelock \
+            PyYAML diffusers huggingface_hub remote-pdb safetensors loguru mlx \
+            "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0"
+
+      - name: Show Apple runtime
+        run: |
+          python - <<'PY'
+          import platform
+          import mlx.core as mx
+          import torch
+
+          print("machine:", platform.machine())
+          print("processor:", platform.processor())
+          print("mlx default device:", mx.default_device())
+          memory_size = mx.metal.device_info().get("memory_size") if mx.metal.is_available() else "metal unavailable"
+          print("mlx memory_size:", memory_size)
+          print("torch:", torch.__version__)
+          print("torch mps available:", torch.backends.mps.is_available())
+          PY
+
+      - name: Run MLX smoke tests
+        run: |
+          python -m pytest \
+            fastvideo/tests/mlx/test_dmd_sampling.py \
+            fastvideo/tests/mlx/test_memory_limits.py \
+            fastvideo/tests/mlx/test_quant_capability.py \
+            fastvideo/tests/mlx/test_mlx_dit_parity.py \
+            fastvideo/tests/mlx/test_mlx_compile_parity.py \
+            fastvideo/tests/mlx/test_mlx_checkpoint.py \
+            fastvideo/tests/mlx/test_mlx_checkpoint_compat.py \
+            fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
+            fastvideo/tests/mlx/test_taehv_decode.py \
+            fastvideo/tests/mlx/test_frame_upsample.py \
+            fastvideo/tests/mlx/test_mlx_fast_spatial.py \
+            fastvideo/tests/mlx/test_mlx_refine.py \
+            fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
+            fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
+            fastvideo/tests/mlx/test_wan22_sample.py \
+            fastvideo/tests/mlx/test_windowed_attention.py \
+            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \
+            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \
+            fastvideo/tests/platforms/test_mps_vsa_error.py \
+            -q
+
+  # Same tests on MLX's CPU backend. Hosted macOS runners are scarce and
+  # slower to schedule; this Linux job gives fast PR signal on the identical
+  # graph (the parity tests were designed to be backend-agnostic), while the
+  # macOS job above stays the source of truth for Metal behavior.
+  mlx-smoke-linux-cpu:
+    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    env:
+      FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA
+      TOKENIZERS_PARALLELISM: "false"
+      MASTER_ADDR: localhost
+      MASTER_PORT: "29513"
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+
+      - uses: astral-sh/setup-uv@v3
+
+      - name: Install lightweight MLX smoke dependencies (CPU backend)
+        run: |
+          uv pip install --system \
+            --index-url https://download.pytorch.org/whl/cpu \
+            torch==2.11.0 torchvision torchaudio
+          uv pip install --system \
+            pytest numpy scipy pillow imageio einops cloudpickle filelock \
+            PyYAML diffusers huggingface_hub remote-pdb safetensors loguru "mlx[cpu]" \
+            "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0"
+
+      - name: Run MLX smoke tests (CPU backend)
+        run: |
+          python -m pytest \
+            fastvideo/tests/mlx/test_dmd_sampling.py \
+            fastvideo/tests/mlx/test_memory_limits.py \
+            fastvideo/tests/mlx/test_quant_capability.py \
+            fastvideo/tests/mlx/test_mlx_dit_parity.py \
+            fastvideo/tests/mlx/test_mlx_compile_parity.py \
+            fastvideo/tests/mlx/test_mlx_checkpoint.py \
+            fastvideo/tests/mlx/test_mlx_checkpoint_compat.py \
+            fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
+            fastvideo/tests/mlx/test_taehv_decode.py \
+            fastvideo/tests/mlx/test_frame_upsample.py \
+            fastvideo/tests/mlx/test_mlx_fast_spatial.py \
+            fastvideo/tests/mlx/test_mlx_refine.py \
+            fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
+            fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
+            fastvideo/tests/mlx/test_wan22_sample.py \
+            fastvideo/tests/mlx/test_windowed_attention.py \
+            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \
+            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \
+            fastvideo/tests/platforms/test_mps_vsa_error.py \
+            -q

.github/workflows/ci-scheduled-ssim.yml

@@ -0,0 +1,44 @@
+name: Scheduled Full SSIM
+
+on:
+  schedule:
+    - cron: "0 5 * * 0"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  trigger:
+    if: github.repository == 'hao-ai-lab/FastVideo'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Trigger weekly full SSIM on Slinky Slurm
+        env:
+          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
+          SOURCE_SHA: ${{ github.sha }}
+          SOURCE_BRANCH: ${{ github.event.repository.default_branch }}
+          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
+          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
+        run: |
+          set -euo pipefail
+          curl -sS --fail-with-body -X POST \
+            "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
+            -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
+            -H "Content-Type: application/json" \
+            --data-raw "$(jq -n \
+              --arg commit "$SOURCE_SHA" \
+              --arg branch "$SOURCE_BRANCH" \
+              '{
+                commit: $commit,
+                branch: $branch,
+                message: "Weekly full SSIM on Slinky Slurm",
+                ignore_pipeline_branch_filters: true,
+                env: {
+                  TEST_SCOPE: "scheduled",
+                  FULL_SUITE: "false",
+                  TEST_TYPE: "ssim",
+                  PR_NUMBER: "false",
+                  PR_TITLE: "Scheduled full SSIM"
+                }
+              }')"

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

@@ -33,7 +33,6 @@ jobs:
             core.setOutput('has_write', String(hasWrite));
 
       - name: Add ready label and react
-        id: label
         if: steps.perm.outputs.has_write == 'true'
         uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
         with:
@@ -48,47 +47,6 @@ jobs:
               comment_id: context.payload.comment.id,
               content: 'rocket',
             });
-            const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: prNumber });
-            core.setOutput('pr_sha', pr.head.sha);
-            core.setOutput('pr_branch', pr.head.ref);
-            core.setOutput('pr_number', String(prNumber));
-            core.setOutput('pr_title', pr.title);
-
-      - name: Trigger Full Suite
-        if: steps.perm.outputs.has_write == 'true'
-        env:
-          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
-          PR_SHA: ${{ steps.label.outputs.pr_sha }}
-          PR_BRANCH: ${{ steps.label.outputs.pr_branch }}
-          PR_NUMBER: ${{ steps.label.outputs.pr_number }}
-          PR_TITLE: ${{ steps.label.outputs.pr_title }}
-          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
-          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
-        run: |
-          curl -sS --fail-with-body -X POST \
-            "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
-            -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
-            -H "Content-Type: application/json" \
-            --data-raw "$(jq -n \
-              --arg commit "$PR_SHA" \
-              --arg branch "$PR_BRANCH" \
-              --arg message "Full Suite for PR #${PR_NUMBER} (via /merge)" \
-              --arg pr_title "$PR_TITLE" \
-              --argjson pr_id "$PR_NUMBER" \
-              '{
-                commit: $commit,
-                branch: $branch,
-                message: $message,
-                ignore_pipeline_branch_filters: true,
-                pull_request_id: $pr_id,
-                pull_request_base_branch: "main",
-                env: {
-                    TEST_SCOPE: "full",
-                    FULL_SUITE: "true",
-                    PR_NUMBER: ($pr_id | tostring),
-                    PR_TITLE: $pr_title
-                  }
-                }')"
 
   parse-command:
     if: >-
@@ -129,7 +87,7 @@ jobs:
           set -euo pipefail
           TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
 
-          VALID="encoder vae transformer kernel unit dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
+          VALID="encoder vae transformer kernel unit dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval unit-ci kernel-ci dreamverse-ci ssim-ci golden-gate-ci encoder-ci vae-ci transformer-ci lora-inference-ci lora-training-ci lora-extraction-ci training-ci distillation-ci self-forcing-ci vsa-ci vmoba-ci performance-ci api-ci train-framework-ci eval-ci full fastcheck pre-commit"
           if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
             echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
             exit 1
@@ -137,7 +95,17 @@ jobs:
 
           declare -A MAP=(
             [encoder]=encoder [vae]=vae [transformer]=transformer
-            [kernel]=kernel_tests [unit]=unit_test [dreamverse]=dreamverse_app
+            [kernel]=kernel_tests [unit]=unit_test [unit-ci]=unit_test_ci
+            [kernel-ci]=kernel_tests_ci [dreamverse-ci]=dreamverse_app_ci
+            [ssim-ci]=ssim_ci [vmoba-ci]=inference_vmoba_ci
+            [golden-gate-ci]=golden_gate_ci [training-ci]=training_ci
+            [encoder-ci]=encoder_ci [vae-ci]=vae_ci [transformer-ci]=transformer_ci
+            [lora-inference-ci]=inference_lora_ci [lora-training-ci]=training_lora_ci
+            [lora-extraction-ci]=lora_extraction_ci [distillation-ci]=distillation_dmd_ci
+            [self-forcing-ci]=self_forcing_ci [vsa-ci]=training_vsa_ci
+            [performance-ci]=performance_ci [api-ci]=api_server_ci
+            [train-framework-ci]=train_framework_ci [eval-ci]=eval_ci
+            [dreamverse]=dreamverse_app
             [ssim]=ssim [golden-gate]=golden_gate [training]=training
             [lora-inference]=inference_lora [lora-training]=training_lora
             [lora-extraction]=lora_extraction

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

@@ -1,4 +1,4 @@
-name: Trigger Full Suite
+name: Trigger Merge Gate
 
 on:
   pull_request_target:
@@ -10,7 +10,7 @@ permissions:
   actions: read
 
 concurrency:
-  group: full-suite-${{ github.event.pull_request.number }}
+  group: merge-gate-${{ github.event.pull_request.number }}
   cancel-in-progress: false
 
 jobs:
@@ -34,29 +34,72 @@ jobs:
             });
             const hasReady = pr.labels.some(l => l.name === 'ready');
             core.setOutput('has_ready', String(hasReady));
-            if (!hasReady) core.info('No ready label — skipping Full Suite trigger.');
+            core.setOutput('changed_files', String(pr.changed_files));
+            if (!hasReady) core.info('No ready label — skipping merge-gate trigger.');
 
       - name: Cancel previous Buildkite builds
         if: steps.check.outputs.has_ready == 'true'
         env:
           BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
           PR_BRANCH: ${{ github.event.pull_request.head.ref }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
         run: |
-          # Find running builds for this branch with TEST_SCOPE=full and cancel them
-          builds=$(curl -sS -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
-            "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds?branch=${PR_BRANCH}&state=running,scheduled" \
-            | jq -r '.[] | select(try (.env.TEST_SCOPE == "full") catch false) | .number')
+          # Match both branch and PR number: forks can reuse the same branch name.
+          builds=$(curl -sS --get -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
+            --data-urlencode "branch=$PR_BRANCH" \
+            --data-urlencode "state=running,scheduled" \
+            "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds" \
+            | jq -r --arg pr_number "$PR_NUMBER" \
+              '.[] | select((.env.TEST_SCOPE? == "merge") and (.env.PR_NUMBER? == $pr_number)) | .number')
           for build_num in $builds; do
             echo "Cancelling Buildkite build #$build_num"
             curl -sS -X PUT -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
               "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel"
           done
 
-      # Checks out the BASE branch (default for pull_request_target), so PR
-      # authors cannot tamper with the gate script.
-      - name: Checkout gate script
+      # Check out the immutable BASE SHA: pull_request_target must never run a
+      # planner or gate script from the untrusted PR head.
+      - name: Checkout trusted merge planner
         if: steps.check.outputs.has_ready == 'true'
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          persist-credentials: false
+
+      - name: Collect changed paths
+        if: steps.check.outputs.has_ready == 'true'
+        env:
+          GH_TOKEN: ${{ github.token }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          EXPECTED_CHANGED_FILES: ${{ steps.check.outputs.changed_files }}
+        run: |
+          set -euo pipefail
+          changed_json="$RUNNER_TEMP/merge-changed-files.json"
+          changed_paths="$RUNNER_TEMP/merge-changed-paths.txt"
+          if gh api --paginate --slurp \
+              "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}/files?per_page=100" \
+              > "$changed_json"; then
+            observed=$(jq '[.[][] | .filename] | unique | length' "$changed_json")
+            if [ "$observed" = "$EXPECTED_CHANGED_FILES" ]; then
+              jq -r '.[][] | .filename, (.previous_filename // empty)' "$changed_json" \
+                | sort -u > "$changed_paths"
+            else
+              echo "::warning::Changed-file API returned $observed of $EXPECTED_CHANGED_FILES paths; selecting all merge lanes."
+              echo '__FASTVIDEO_CI_PLAN_ALL__' > "$changed_paths"
+            fi
+          else
+            echo "::warning::Changed-file API failed; selecting all merge lanes."
+            echo '__FASTVIDEO_CI_PLAN_ALL__' > "$changed_paths"
+          fi
+
+      - name: Select minimal merge tests
+        id: plan
+        if: steps.check.outputs.has_ready == 'true'
+        run: |
+          python3 .github/scripts/plan_merge_ci.py \
+            --paths-file "$RUNNER_TEMP/merge-changed-paths.txt" \
+            --github-output "$GITHUB_OUTPUT" \
+            --summary-file "$GITHUB_STEP_SUMMARY"
 
       - name: Wait for pre-commit and docs build
         if: steps.check.outputs.has_ready == 'true'
@@ -66,7 +109,7 @@ jobs:
           PR_NUMBER: ${{ github.event.pull_request.number }}
         run: bash .github/scripts/gate_full_suite.sh
 
-      - name: Trigger Buildkite Full Suite
+      - name: Trigger Buildkite merge gate
         if: steps.check.outputs.has_ready == 'true'
         env:
           BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
@@ -76,6 +119,10 @@ jobs:
           PR_TITLE: ${{ github.event.pull_request.title }}
           BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
           BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
+          MERGE_TEST_PLAN: ${{ steps.plan.outputs.merge_test_plan }}
+          MERGE_GOLDEN_TESTS: ${{ steps.plan.outputs.merge_golden_tests }}
+          MERGE_SSIM_TESTS: ${{ steps.plan.outputs.merge_ssim_tests }}
+          MERGE_PLAN_LABEL: ${{ steps.plan.outputs.merge_plan_label }}
         run: |
           curl -sS --fail-with-body -X POST \
             "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
@@ -84,8 +131,11 @@ jobs:
             --data-raw "$(jq -n \
               --arg commit "$PR_SHA" \
               --arg branch "$PR_BRANCH" \
-              --arg message "Full Suite for PR #${PR_NUMBER}" \
+              --arg message "Merge gate [${MERGE_PLAN_LABEL}] for PR #${PR_NUMBER}" \
               --arg pr_title "$PR_TITLE" \
+              --arg merge_test_plan "$MERGE_TEST_PLAN" \
+              --arg merge_golden_tests "$MERGE_GOLDEN_TESTS" \
+              --arg merge_ssim_tests "$MERGE_SSIM_TESTS" \
               --argjson pr_id "$PR_NUMBER" \
               '{
                 commit: $commit,
@@ -95,8 +145,11 @@ jobs:
                 pull_request_id: $pr_id,
                 pull_request_base_branch: "main",
                 env: {
-                  TEST_SCOPE: "full",
+                  TEST_SCOPE: "merge",
                   FULL_SUITE: "true",
+                  MERGE_TEST_PLAN: $merge_test_plan,
+                  MERGE_GOLDEN_TESTS: $merge_golden_tests,
+                  MERGE_SSIM_TESTS: $merge_ssim_tests,
                   PR_NUMBER: ($pr_id | tostring),
                   PR_TITLE: $pr_title
                 }

.github/workflows/community-welcome.yml

@@ -38,17 +38,17 @@ jobs:
 
             **How our CI works:**
 
-            PRs run a two-tier CI system:
+            PRs run a three-tier CI system:
             1. **Pre-commit** — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
-            2. **Fastcheck** — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
-            3. **Full Suite** — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the `ready` label.
+            2. **Fastcheck** — six core GPU lanes run automatically via Buildkite (~10-15 min).
+            3. **Merge gate** — a reviewer adds `ready`; changed paths select only the relevant integration, training, golden, or SSIM coverage.
 
             **Before your PR is reviewed:**
             - [ ] `pre-commit run --all-files` passes locally
             - [ ] You've added or updated tests for your changes
             - [ ] The PR description explains what and why
 
-            If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.
+            If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and merge-gate results appear in the Checks section below.
 
             **Useful links:**
             - [Contributing Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)

.github/workflows/infra-build-image.yml

@@ -13,6 +13,11 @@ on:
         required: false
         default: false
         type: boolean
+      build_ci_runner_image:
+        description: 'Build the ARM64 CUDA 13 CI runner image (sm_100)'
+        required: false
+        default: false
+        type: boolean
   # Auto-rebuild the CUDA images when a repository-controlled image input
   # changes on main. This includes the trusted SM89 kernel artifact's source,
   # metadata/key helper, ABI dependency metadata, and build orchestration.
@@ -198,6 +203,28 @@ jobs:
           docker buildx imagetools create "${TAG_ARGS[@]}" "${IMAGE_REFS[@]}"
           docker buildx imagetools inspect "${TAGS[0]}"
 
+  # The CI runner is ARM64 like DGX Spark, but targets sm_100 rather than sm_121.
+  # Publish a single-architecture variant so the self-hosted CI runner can reuse
+  # the exact prebuilt kernel instead of compiling it in every job.
+  build-ci-runner-image:
+    if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_ci_runner_image == 'true' }}
+    uses: ./.github/workflows/_template-build-image.yml
+    with:
+      python_version: '3.12'
+      dockerfile_path: docker/Dockerfile
+      tag_suffix: py3.12-cuda13.0.0-sm100
+      runner: ubuntu-24.04-arm
+      architecture: arm64
+      build_args: |
+        PYTHON_VERSION=3.12
+        CUDA_VERSION=13.0.0
+        UV_TORCH_BACKEND=cu130
+        TORCH_CUDA_ARCH_LIST=10.0
+        CMAKE_BUILD_PARALLEL_LEVEL=1
+        FLASH_ATTN_WHEEL_TAG=cu130torch2.12
+        FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22
+    secrets: inherit
+
   # Dreamverse matrix: {backend, UI} x {12.6.3, 13.0.0}, Python 3.12. Torch backend
   # matches the base CUDA (cu126 / cu130). Keep these images amd64-only until the
   # required FA4 dependency stack is available and validated on arm64.

.github/workflows/infra-docs.yml

@@ -6,6 +6,7 @@ on:
     paths:
       - 'docs/**'
       - 'examples/**'
+      - 'scripts/inference/**'
       - 'mkdocs.yml'
       - 'requirements-mkdocs.in'
       - 'requirements-mkdocs.txt'
@@ -16,6 +17,7 @@ on:
     paths:
       - 'docs/**'
       - 'examples/**'
+      - 'scripts/inference/**'
       - 'mkdocs.yml'
       - 'requirements-mkdocs.in'
       - 'requirements-mkdocs.txt'

.github/workflows/publish-kernel.yml

@@ -62,8 +62,9 @@ jobs:
               cuda-version: '13.0.0'
               torch-cuda-short: 'cu130'
           platform:
-            # x86_64 builds the full cu126 + cu130 set (cu130 ships the consumer
-            # Blackwell sm_120a FP4 kernels).
+            # x86_64 builds the full cu126 + cu130 set. cu130 ships the
+            # data-center Blackwell sm_100a VSA and consumer sm_120a FP4
+            # kernels.
             - os: ubuntu-22.04
               arch: x86_64
               wheel-plat: manylinux_2_35_x86_64
@@ -124,7 +125,7 @@ jobs:
       - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
         run: |
           sudo apt update
-          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo apt install -y git gcc-11 g++-11 clang-11
           sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
 
           # Allow Git to Access Safe Directory
@@ -168,7 +169,8 @@ jobs:
           #     covers sm_120a; turbodiffusion covers sm_100a+sm_120a. The sm_100 FP4
           #     forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221),
           #     JIT-compiled at runtime — not built into this wheel.
-          #   * x86_64 cu130 = Hopper TK + consumer Blackwell sm_120a FP4.
+          #   * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a VSA
+          #     + consumer Blackwell sm_120a FP4.
           #   * x86_64 cu126 = Hopper TK only (older drivers; CUDA < 12.8 has no FP4).
           # The per-arch split in CMakeLists pins the FP4 targets to sm_120a and builds
           # the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps
@@ -178,7 +180,7 @@ jobs:
             export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON"
             export CMAKE_BUILD_PARALLEL_LEVEL=1
           elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then
-            export TORCH_CUDA_ARCH_LIST="9.0a;12.0a"
+            export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a"
             export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
             # A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize.
             export CMAKE_BUILD_PARALLEL_LEVEL=1
@@ -194,7 +196,11 @@ jobs:
           python -m build --wheel --outdir dist
 
           # Fix the wheel to be manylinux compliant
-          uv pip install --system auditwheel
+          # Ubuntu 22.04 ships patchelf 0.14.3, while current auditwheel
+          # requires at least 0.14.5. Use the stable PyPI binary on both
+          # x86_64 and aarch64 release runners.
+          uv pip install --system auditwheel patchelf==0.17.2.4
+          patchelf --version
           # Point auditwheel at torch libs, but do not vendor them into the wheel.
           TORCH_LIB_DIR=$(python - <<'PY'
           import os
@@ -211,7 +217,8 @@ jobs:
             --exclude libtorch.so \
             --exclude libc10.so \
             --exclude libc10_cuda.so \
-            --exclude libtorch_python.so
+            --exclude libtorch_python.so \
+            --exclude libnccl.so.2
           # Move fixed wheels back to dist for upload consistency
           rm dist/*.whl
           mv fixed_dist/*.whl dist/

.gitignore

@@ -23,6 +23,7 @@ Miniconda3-latest-Linux-x86_64.sh
 *validation/
 data/
 outputs/
+outputs_audio/
 outputs_video
 checkpoints/
 sbatch.sh

README.md

@@ -9,6 +9,8 @@
 **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**
 
 ## NEWS
+- `2026/08/23`: [FastH3 Preview v0.2](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2) is a 4-step DMD2-distilled MiniMax-H3 checkpoint that generates synchronized video and audio. Run the verified [basic FastH3 example](examples/inference/basic/basic_fasth3.py), see the [inference guide](examples/inference/basic/README.md#fasth3-preview), or have a coding agent install FastVideo with the [agent setup prompt](#install-with-an-ai-coding-agent).
+- `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac—follow the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/) and read the [Blog](https://haoailab.com/blogs/fastmetal/).
 - `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. See the [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), [Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/), and [blog](https://haoailab.com/blogs/fastwan-qad/).
 - `2026/03/17`: Release demo: Into the Dreamverse: Vibe Directing in FastVideo, check out the [Blog](https://haoailab.com/blogs/dreamverse/).
 - `2026/03/13`: Release demo: Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU, check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
@@ -62,6 +64,12 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo
 Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the
 [MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
 
+> **On an Apple Silicon Mac?** FastVideo runs FastMetal-QAD through an MLX
+> runtime. Install with `uv pip install -e '.[mlx]'`, download
+> [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD),
+> and follow the
+> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
+
 Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
 
 > **On an NVIDIA DGX Spark (GB10 / ARM64 + CUDA 13)?** There's no prebuilt ARM wheel for the FastVideo CUDA kernel, so it's an editable from-source install (`UV_TORCH_BACKEND=cu130 uv pip install -e .`, which compiles that kernel for you) rather than `UV_TORCH_BACKEND=cu130 uv pip install fastvideo`. A compatible prebuilt ARM64 FlashAttention wheel is available separately. Follow the [DGX Spark install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/spark/).

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d
 
 # 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
@@ -89,11 +89,29 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     zsh \
     vim \
     curl \
+    ffmpeg \
+    libgl1 \
+    libglib2.0-0 \
+    libx11-dev \
     gcc-11 \
     g++-11 \
     clang-11 \
+    cmake \
+    pkg-config \
+    build-essential \
+    libssl-dev \
     && rm -rf /var/lib/apt/lists/*
 
+# Rust toolchain: some dependencies only ship sdists on aarch64 and need cargo
+# to build. The dormant legacy Modal image layers the identical apt set +
+# rustup on top of this image (fastvideo/tests/modal/pr_test.py); baking both
+# here keeps its manual rollback path reproducible without changing the Slurm
+# runner's package surface.
+RUN set -o pipefail && \
+    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable --profile minimal && \
+    /root/.cargo/bin/cargo --version && /root/.cargo/bin/rustc --version
+ENV PATH=/root/.cargo/bin:${PATH}
+
 # Set up C++20 compilers for ThunderKittens
 RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
 
@@ -138,6 +156,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     source /opt/venv/bin/activate && \
     uv pip install --upgrade pip && \
     uv pip install --excludes docker/uv-excludes ".[dev]" && \
+    python -c "import cv2; print('OpenCV', cv2.__version__)" && \
     PYTAG=cp$(echo "${PYTHON_VERSION}" | tr -d .) && \
     case "${TARGETARCH:-amd64}" in \
         amd64) \
@@ -169,26 +188,21 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # flash_attn/__init__.py), so FA2/varlen/bert_padding stay from the install above;
 # rmtree clears the wheel's stale cute files first to avoid an install conflict.
 # Then verify both survive so a broken overlay fails the build instead of shipping
-# an FA2-less image. x86 only: the FA4 stack (quack-kernels etc.) is unvalidated on
-# arm64 / GB10 (sm_121), so there we skip the overlay; FA4 is opt-in
-# (FASTVIDEO_FA4=1) and errors if set without the overlay, so leave it unset on
-# arm64 and the image runs FA3/FA2 as usual.
+# an FA2-less image. The pinned stack is validated on ARM64 GB200 (sm_100) as well
+# as x86; FA4 remains opt-in through FASTVIDEO_FA4=1 so lanes with FA2 baselines
+# keep their existing numerics.
 RUN --mount=type=cache,target=/opt/uv/cache \
     source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    if [ "${TARGETARCH}" = "arm64" ]; then \
-        echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; do not set FASTVIDEO_FA4)"; \
-    else \
-        python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
-        uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
-        python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"; \
-    fi
+    python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
+    uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
+    python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"
 
 COPY . .
 
 # Build immutable FastVideo kernel wheels for the published image. The requested
 # architecture remains installed for normal image users; amd64 images also carry
-# an SM89 artifact so the predominant L40S Modal lanes can reuse it exactly.
+# an SM89 artifact for L40S users and the dormant legacy rollback path.
 ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt
 RUN --mount=type=cache,target=/opt/uv/cache \
     source $HOME/.local/bin/env && \

docs/assets/cookbook-recipes.json

@@ -0,0 +1,52 @@
+{
+  "recipes": [
+    {
+      "id": "fastwan21-t2v",
+      "task": "Text to video",
+      "label": "FastWan2.1 1.3B (distilled + VSA)",
+      "model": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+      "source": "scripts/inference/inference_wan_VSA_DMD_1_3B.yaml",
+      "command": "FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN fastvideo generate --config scripts/inference/inference_wan_VSA_DMD_1_3B.yaml"
+    },
+    {
+      "id": "wan22-t2v",
+      "task": "Text to video",
+      "label": "Wan2.2 A14B",
+      "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+      "source": "examples/inference/basic/basic_wan2_2.py",
+      "command": "python examples/inference/basic/basic_wan2_2.py"
+    },
+    {
+      "id": "wan21-i2v",
+      "task": "Image to video",
+      "label": "Wan2.1 14B 480P",
+      "model": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+      "source": "scripts/inference/inference_wan_i2v.yaml",
+      "command": "fastvideo generate --config scripts/inference/inference_wan_i2v.yaml"
+    },
+    {
+      "id": "turbowan22-i2v",
+      "task": "Image to video",
+      "label": "TurboWan2.2 A14B",
+      "model": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers",
+      "source": "examples/inference/basic/basic_turbodiffusion_i2v.py",
+      "command": "python examples/inference/basic/basic_turbodiffusion_i2v.py"
+    },
+    {
+      "id": "wan22-ti2v",
+      "task": "Text or image to video",
+      "label": "Wan2.2 TI2V 5B",
+      "model": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+      "source": "examples/inference/basic/basic_wan2_2_ti2v.py",
+      "command": "python examples/inference/basic/basic_wan2_2_ti2v.py"
+    },
+    {
+      "id": "matrix-game-2",
+      "task": "Interactive world",
+      "label": "Matrix Game 2.0",
+      "model": "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers",
+      "source": "examples/inference/basic/basic_matrixgame2.py",
+      "command": "python examples/inference/basic/basic_matrixgame2.py"
+    }
+  ]
+}

docs/assets/cookbook.js

@@ -0,0 +1,60 @@
+(() => {
+  let recipesPromise;
+
+  const loadRecipes = (url) => {
+    recipesPromise ||= fetch(url).then((response) => {
+      if (!response.ok) throw new Error(`HTTP ${response.status}`);
+      return response.json();
+    });
+    return recipesPromise;
+  };
+
+  const init = () => {
+    document.querySelectorAll("[data-cookbook]").forEach(async (root) => {
+      if (root.dataset.initialized) return;
+      root.dataset.initialized = "true";
+
+      const select = root.querySelector("[data-cookbook-recipe]");
+      const model = root.querySelector("[data-cookbook-model]");
+      const source = root.querySelector("[data-cookbook-source]");
+      const command = root.querySelector("[data-cookbook-command]");
+      const status = root.querySelector("[data-cookbook-status]");
+
+      try {
+        const { recipes } = await loadRecipes(root.dataset.recipes);
+        const byId = new Map(recipes.map((recipe) => [recipe.id, recipe]));
+        const groups = new Map();
+
+        select.replaceChildren();
+        recipes.forEach((recipe) => {
+          if (!groups.has(recipe.task)) {
+            const group = document.createElement("optgroup");
+            group.label = recipe.task;
+            groups.set(recipe.task, group);
+            select.append(group);
+          }
+          groups.get(recipe.task).append(new Option(recipe.label, recipe.id));
+        });
+
+        const render = () => {
+          const recipe = byId.get(select.value);
+          model.textContent = recipe.model;
+          source.textContent = recipe.source;
+          source.href = `https://github.com/hao-ai-lab/FastVideo/blob/main/${recipe.source}`;
+          command.textContent = recipe.command;
+          status.textContent = `${recipe.label} selected.`;
+        };
+
+        select.addEventListener("change", render);
+        select.disabled = false;
+        render();
+      } catch (error) {
+        status.textContent = "Recipes could not be loaded. Use the examples link below.";
+        console.error("Failed to load FastVideo cookbook recipes", error);
+      }
+    });
+  };
+
+  if (window.document$) window.document$.subscribe(init);
+  else document.addEventListener("DOMContentLoaded", init);
+})();

docs/assets/custom.css

@@ -42,6 +42,46 @@ img {
     margin: 0 auto;
 }
 
+.cookbook-picker {
+    padding: 1rem;
+    border: 0.05rem solid var(--md-default-fg-color--lightest);
+    border-radius: 0.2rem;
+}
+
+.cookbook-picker select {
+    width: 100%;
+    padding: 0.6rem;
+    color: var(--md-default-fg-color);
+    background: var(--md-default-bg-color);
+    border: 0.05rem solid var(--md-default-fg-color--lighter);
+    border-radius: 0.2rem;
+}
+
+.cookbook-picker dl {
+    display: grid;
+    grid-template-columns: max-content 1fr;
+    gap: 0.25rem 1rem;
+}
+
+.cookbook-picker dt {
+    font-weight: 700;
+}
+
+.cookbook-picker dd {
+    margin: 0;
+    min-width: 0;
+    overflow-wrap: anywhere;
+}
+
+.cookbook-picker__status {
+    position: absolute;
+    width: 1px;
+    height: 1px;
+    overflow: hidden;
+    clip: rect(0, 0, 0, 0);
+    white-space: nowrap;
+}
+
 .md-typeset .copy-page-button.md-button {
     float: right;
     margin: 0 0 1rem 1rem;

docs/contributing/ci_architecture.md

@@ -6,7 +6,8 @@ lives in [Testing](testing.md).
 
 ## Overview
 
-FastVideo splits validation across GitHub Actions, Buildkite, and Modal:
+FastVideo splits validation across GitHub Actions, Buildkite, Slinky Slurm,
+and Mergify:
 
 ```text
 PR opened or updated
@@ -16,17 +17,24 @@ PR opened or updated
   |     style, lint, type, spelling, Markdown, workflow syntax, filenames
   |
   |-- Tier 2: Fastcheck
-  |     Buildkite orchestrates Modal GPU jobs
-  |     path-filtered component and unit checks
+  |     Buildkite schedules six lanes on the Slinky Slurm cluster
+  |     encoder, VAE, transformer, kernel, unit, DreamVerse
   |
-  |-- /merge, /test full, or ready label
+  |-- /merge or ready label
         |
-        `-- Tier 3: Full Suite
-              Buildkite orchestrates Modal GPU jobs
-              path-filtered integration, SSIM, training, eval, and performance checks
+        `-- Tier 3: Change-aware merge gate
+              trusted base-branch planner classifies the complete PR diff
+              Buildkite adds only relevant integration, quality, training,
+              API, or performance lanes on Slinky Slurm
               |
               pass -> Mergify squash-merges when all merge conditions pass
               fail -> fix, push, and re-run
+
+  |-- /test full
+        `-- Explicit all-20-lane diagnostic run
+
+  `-- weekly schedule on main
+        `-- Complete four-GPU SSIM matrix
 ```
 
 CI is not one monolithic job:
@@ -34,10 +42,31 @@ CI is not one monolithic job:
 - GitHub Actions owns pre-commit, slash-command handling, aggregate status
   updates, docs deployment, image builds, package publishing, and community
   automations.
-- Buildkite owns the GPU test pipeline and path filtering.
-- Modal owns the actual GPU execution environment for test jobs.
+- Buildkite owns the GPU test graph, statuses, and trusted dispatch control
+  plane. Its agent runs on the Slurm login plane; it does not execute test
+  payloads.
+- Slinky Slurm is the only active CI compute backend. A host-owned dispatcher
+  leases GPUs from a persistent four-GPU allocation and runs each lane in an
+  isolated Enroot container at the immutable PR SHA.
 - Mergify owns merge protection, labeling, and the final squash merge.
 
+The old files under `fastvideo/tests/modal/` are retained as dormant manual
+rollback code. `.buildkite/scripts/pr_test.sh` rejects Buildkite invocations,
+and no pipeline or slash-command route calls Modal.
+
+Three Buildkite entry pipelines share the validated graph:
+
+| Pipeline | Trigger | Scope |
+|---|---|---|
+| `pr-fastcheck` | Automatic pull-request webhook | Six Fastcheck lanes |
+| `ci` | `/merge`, `ready`, schedules, and `/test` API builds | Change-aware merge gates, scheduled SSIM, explicit Full Suite, Fastcheck reruns, or one direct lane |
+| `fastvideo-performance-lane` | Weekly scheduler | Direct performance lane |
+
+Each entry pipeline starts with the same trusted `pipeline-upload` job on the
+`ci-runner` queue. The `ci` pipeline's incoming GitHub webhook is disabled;
+otherwise it would duplicate the automatic `pr-fastcheck` build. API and
+scheduled builds continue to work with webhook processing disabled.
+
 ## CI Tiers
 
 ### Tier 1: Pre-commit
@@ -70,34 +99,37 @@ debugging a hook implementation.
 | Attribute | Value |
 |---|---|
 | Triggered by | Buildkite PR builds with `TEST_SCOPE=fastcheck` or unset |
-| Runner | Buildkite agent that launches Modal GPU jobs |
+| Compute | Slinky Slurm (`ci-runner` queue) |
 | Definition | `.buildkite/pipeline.yml` |
-| Entrypoint | `.buildkite/scripts/pr_test.sh` -> `fastvideo/tests/modal/pr_test.py` |
+| Entrypoint | Trusted host driver -> `.buildkite/scripts/unit_test.sh` or `.buildkite/scripts/lanes/*.sh` |
 
-Fastcheck uses Buildkite's `monorepo-diff` plugin. Jobs whose watched paths did
-not change are skipped and do not block the aggregate `fastcheck-passed`
-status.
+Fastcheck always schedules these six lanes: encoder, VAE, transformer, custom
+kernels, unit tests, and DreamVerse. Static steps replace the former
+host-side path-filter plugin: the login plane never checks out or executes PR
+code.
 
-| Buildkite label | `TEST_TYPE` | Main watched paths |
-|---|---|---|
-| Encoder Tests | `encoder` | `fastvideo/models/encoders/**`, `fastvideo/models/loader/**`, `fastvideo/tests/encoders/**`, `pyproject.toml`, `docker/Dockerfile` |
-| VAE Tests | `vae` | `fastvideo/models/vaes/**`, `fastvideo/models/loader/**`, `fastvideo/tests/vaes/**`, `pyproject.toml`, `docker/Dockerfile` |
-| Transformer Tests | `transformer` | `fastvideo/models/dits/**`, `fastvideo/models/loader/**`, `fastvideo/tests/transformers/**`, `fastvideo/layers/**`, `fastvideo/attention/**`, `pyproject.toml`, `docker/Dockerfile` |
-| Kernel Tests | `kernel_tests` | `fastvideo-kernel/**`, `pyproject.toml`, `docker/Dockerfile` |
-| Unit Tests | `unit_test` | `fastvideo/**`, `.buildkite/**`, `.github/**`, `pyproject.toml`, `docker/Dockerfile` |
-| DreamVerse App Tests | `dreamverse_app` | `apps/dreamverse/**`, `pyproject.toml` |
-
-### Tier 3: Full Suite
+### Tier 3: Change-Aware Merge Gate
 
 | Attribute | Value |
 |---|---|
-| Triggered by | `/merge`, adding `ready`, `/test full`, or a new push to a PR that already has `ready` |
-| Runner | Buildkite agent that launches Modal GPU jobs |
+| Triggered by | `/merge`, adding `ready`, or a new push to a PR that already has `ready` |
+| Compute | Slinky Slurm only (`ci-runner` queue) |
 | Definition | `.buildkite/pipeline.yml` |
-| Entrypoint | `.buildkite/scripts/pr_test.sh` -> `fastvideo/tests/modal/pr_test.py` |
+| Entrypoint | `/opt/fastvideo-ci-runner/run-ci` (`run-unit` is a compatibility wrapper) |
 
-Full Suite is also path-filtered. It validates broader behavior before Mergify
-can merge a PR.
+Fastcheck is the universal six-lane baseline. The merge gate does not repeat
+those jobs: it classifies every changed path and adds only the relevant lanes
+from the fourteen-lane integration set below. Selected jobs are hard gates;
+there are no soft-fail hardware lanes. A documentation-only PR can therefore
+finish its merge build after the trusted uploader, while model-family changes
+typically add focused golden-gate and SSIM files and a training-only change
+adds only its owning training lane.
+
+`.github/scripts/plan_merge_ci.py` is the canonical path policy. It runs from
+the immutable base SHA under `pull_request_target`; PR code is never executed
+on the GitHub runner. The changed-file list includes both sides of renames. An
+API failure, truncated response, empty list, unknown build input, or unknown
+source path fails closed to all fourteen integration lanes.
 
 A `ready`-labeled PR does not hit Buildkite immediately:
 `ci-trigger-full-suite.yml` first runs `.github/scripts/gate_full_suite.sh`,
@@ -106,25 +138,100 @@ head. A red cheap check blocks the suite (fail closed; the next push re-arms
 it), while a GitHub outage or a >25 min wait lets it run anyway (fail open).
 `/test full` bypasses the gate.
 
-| Buildkite label | `TEST_TYPE` | Main watched paths |
-|---|---|---|
-| SSIM Tests | `ssim` | `fastvideo/**/*.py`, `pyproject.toml`, `docker/Dockerfile` |
-| LoRA Inference Tests | `inference_lora` | LoRA tests, loader, transformer tests, pipelines, LoRA layers |
-| LoRA Extraction Tests | `lora_extraction` | LoRA extraction scripts/tests, loader, training utilities, LoRA layers |
-| Training Tests | `training` | `fastvideo/**`, `pyproject.toml`, `docker/Dockerfile` |
-| Distillation DMD Tests | `distillation_dmd` | `fastvideo/training/*distillation_pipeline.py` |
-| Self-Forcing Tests | `self_forcing` | self-forcing distillation pipeline and tests |
-| LoRA Training Tests | `training_lora` | `fastvideo/**`, `pyproject.toml`, `docker/Dockerfile` |
-| Training Tests VSA | `training_vsa` | `fastvideo/**`, `fastvideo-kernel/**`, `pyproject.toml`, `docker/Dockerfile` |
-| Inference Tests VMoBA | `inference_vmoba` | `fastvideo-kernel/**`, `fastvideo/attention/backends/vmoba.py` |
-| Performance Tests | `performance` | DiTs, pipelines, attention, layers, worker, entrypoints, performance tests/configs |
-| API Server Tests | `api_server` | OpenAI entrypoints, serve CLI, OpenAI API integration test |
-| Train Framework Tests | `train_framework` | `fastvideo/train/**`, train model/method tests, model loader, DiTs |
-| Eval Metrics Tests | `eval` | `fastvideo/eval/**`, `fastvideo/tests/eval/**`, `pyproject.toml`, `docker/Dockerfile` |
+The complete static graph remains available through `/test full`; path
+selection never deletes or dynamically invents a Buildkite step.
+
+| Lane | Public `TEST_TYPE` | GPUs | Typical merge trigger |
+|---|---|---:|---|
+| Encoder | `encoder` | 1 | Universal Fastcheck |
+| VAE | `vae` | 1 | Universal Fastcheck |
+| Transformer | `transformer` | 1 | Universal Fastcheck |
+| Kernel | `kernel_tests` | 1 | Universal Fastcheck |
+| Unit | `unit_test` | 1 | Universal Fastcheck |
+| DreamVerse | `dreamverse_app` | 1 | Universal Fastcheck |
+| Golden gate | `golden_gate` | 1 | Model, pipeline, attention, layer, or output changes |
+| SSIM | `ssim` | 4 | Matching model/SSIM paths; focused files when possible |
+| LoRA inference | `inference_lora` | 1 | LoRA inference/shared LoRA paths |
+| LoRA extraction | `lora_extraction` | 1 | LoRA extraction/shared LoRA paths |
+| Vanilla training | `training` | 4 | Legacy vanilla/shared training paths |
+| DMD distillation | `distillation_dmd` | 2 | DMD/shared training paths |
+| Self-forcing | `self_forcing` | 2 | Self-forcing/shared training paths |
+| LoRA training | `training_lora` | 2 | LoRA/shared training paths |
+| VSA training | `training_vsa` | 2 | VSA/shared training paths |
+| VMoBA inference | `inference_vmoba` | 1 | VMoBA backend/config paths |
+| Performance | `performance` | 2 | Performance tests or benchmark policy |
+| API server | `api_server` | 1 | API, worker, or server entrypoint paths |
+| Modular train framework | `train_framework` | 1 | `fastvideo/train/` and its tests |
+| Eval metrics | `eval` | 1 | `fastvideo/eval/` and its tests |
+
+Golden-gate and SSIM selections are basenames, not arbitrary pytest arguments.
+The private host checks the comma-separated allowlist before staging, and the
+container checks it again before invoking pytest. Shared quality-harness
+changes still run the complete owning lane. The full SSIM matrix also runs on
+`main` every Sunday at 05:00 UTC through `ci-scheduled-ssim.yml`; `/test ssim`
+and `/test full` remain available for deliberate complete runs.
+
+The four Buildkite workers may accept multiple jobs concurrently. The
+agent-owned lease broker packs their requested GPU counts onto one persistent
+four-GPU Slurm allocation and waits when capacity is full. A four-GPU lane
+such as SSIM or vanilla training owns the whole tray; two two-GPU lanes or up
+to four one-GPU lanes can overlap without sharing devices. SSIM and vanilla
+training also share the `fastvideo/slinky/whole-tray` Buildkite concurrency
+group. That keeps the second whole-tray lane in Buildkite instead of consuming
+an agent and its command timeout while the first lane waits for four free GPUs.
+Because packed Enroot containers share the node network namespace, each GPU
+lease receives its own 100-port rendezvous range. Tests preserve the
+runner-assigned `MASTER_PORT`, and parallel SSIM tasks use distinct offsets
+inside that range.
 
 See [Performance Benchmarks](performance_benchmarks.md) for the performance
 lane's thresholds, rolling baseline, artifacts, and reseeding process.
 
+## `/merge` Request Flow
+
+The Buildkite agent and Slurm have deliberately separate responsibilities:
+
+```text
+/merge PR comment
+  -> GitHub verifies write permission and refreshes the `ready` label
+  -> base-branch `ci-trigger-full-suite` workflow fetches the PR file list,
+     computes MERGE_TEST_PLAN plus focused golden/SSIM basenames, and gates on
+     cheap checks
+  -> sends the PR SHA to pipeline `ci` with TEST_SCOPE=merge and FULL_SUITE=true
+  -> trusted `pipeline-upload` job on queue `ci-runner`
+       fetch exact-SHA .buildkite/pipeline.yml
+       normalize + validate the complete static 20-lane policy and conditions
+       upload static Buildkite steps
+  -> each accepted step reaches the host policy hook
+       validate org/repo/SHA/ref/step key/command/scope/timeout
+       skip checkout on the login plane
+       stage a mode-0600 request on Lustre
+       lease 1-4 GPUs and attach an `srun` step to the Slinky tray
+  -> Enroot worker container
+       clone and verify the exact PR SHA
+       install the lane's project extras and cached kernel
+       run the repository-owned lane script
+       write numeric exit status and approved artifacts
+  -> trusted host returns that status to Buildkite
+  -> Buildkite publishes `full-suite-passed` only when every selected lane passes
+```
+
+The pipeline uploader and dispatcher run on the Slurm login plane, but those
+are control-plane operations only. Python tests, model loading, CUDA kernels,
+Node/Playwright checks, inference, training, SSIM generation, and performance
+benchmarks all execute in Slurm allocations.
+
+PR-controlled values never become host commands. The host policy accepts only
+the pinned pipeline uploader or a known lane tuple. It rejects plugins,
+artifact globs, shell injection variables, non-immutable commits, and unknown
+commands before checkout. Hugging Face credentials are added only for lanes
+that declare them, passed through a mode-0600 request file, and removed before
+the PR payload starts. Active training lanes keep W&B offline and do not stage
+a W&B credential. The ARM64 image includes the pinned FA4 CuTe overlay validated
+on GB200. SSIM opts into FA4 to preserve its reference-video numerics; lanes
+with FA2 baselines keep `FASTVIDEO_FA4=0`. Performance artifacts are relayed
+afterward by the trusted host from an allowlisted directory and extension set.
+
 ## Slash Commands
 
 Slash commands are handled by `.github/workflows/ci-slash-commands.yml`.
@@ -132,7 +239,7 @@ Repository write permission is required.
 
 | Command | Effect |
 |---|---|
-| `/merge` | Adds `ready` and triggers Full Suite for the PR head branch. |
+| `/merge` | Adds `ready` and triggers the path-aware merge gate for the PR head. |
 | `/test full` | Runs the whole Full Suite with `TEST_SCOPE=full`. |
 | `/test fastcheck` | Runs the whole Fastcheck suite with `TEST_SCOPE=fastcheck`. |
 | `/test pre-commit` | Re-runs the pre-commit workflow on the PR merge ref. |
@@ -149,6 +256,7 @@ Valid direct test names:
 | `/test unit` | `unit_test` |
 | `/test dreamverse` | `dreamverse_app` |
 | `/test ssim` | `ssim` |
+| `/test golden-gate` | `golden_gate` |
 | `/test training` | `training` |
 | `/test lora-inference` | `inference_lora` |
 | `/test lora-training` | `training_lora` |
@@ -162,12 +270,18 @@ Valid direct test names:
 | `/test train-framework` | `train_framework` |
 | `/test eval` | `eval` |
 
+The temporary `<name>-ci` spellings remain accepted as compatibility aliases;
+they select the same Slurm lane and do not identify a second backend.
+
 When a direct test completes successfully, Buildkite posts
 `direct-test-completed`. `.github/workflows/ci-aggregate-status.yml` then reads
 the latest Buildkite statuses for the commit and updates `fastcheck-passed` or
 `full-suite-passed` if all jobs in that group are green.
 
-Skipped path-filtered jobs have no status entry and do not block the aggregate.
+Buildkite label emojis define the status namespace used by that aggregation:
+`:microscope:` is reserved for the six Fastcheck lanes, while Full-Suite-only
+lanes use `:test_tube:` or `:bar_chart:`. Each active lane has exactly one
+label and therefore one status context.
 
 ## Merge Protection
 
@@ -177,7 +291,7 @@ Mergify enforces these conditions before it squash-merges to `main`:
 |---|---|
 | `check-success~=pre-commit` | Tier 1 passed. |
 | `check-success=fastcheck-passed` | All triggered Fastcheck jobs passed. |
-| `check-success=full-suite-passed` | All triggered Full Suite jobs passed. |
+| `check-success=full-suite-passed` | The selected merge gate or explicit Full Suite passed. |
 | `#approved-reviews-by>=1` | At least one approving review. |
 | Valid title regex | PR title starts with an accepted `[type]` tag. |
 | `label=ready` | The PR has entered the merge flow. |
@@ -226,29 +340,40 @@ Process labels:
 
 | Label | Who sets it | Meaning |
 |---|---|---|
-| `ready` | `/merge` or maintainer action | Triggers/keeps Full Suite active and enables auto-merge. |
+| `ready` | `/merge` or maintainer action | Triggers/keeps the change-aware merge gate active and enables auto-merge. |
 | `needs-rebase` | Mergify | PR has merge conflicts. |
 | `do-not-merge` | Maintainer | Blocks merge regardless of CI status. |
 
-## Modal Test Entrypoints
+## Slurm Lane Entrypoints
 
-All Buildkite test jobs go through `.buildkite/scripts/pr_test.sh`, which:
+Every active test selection lives in `.buildkite/scripts/unit_test.sh` or a
+focused `.buildkite/scripts/lanes/<lane>.sh`. The private, agent-owned lane
+table binds each internal `*_ci` type to that script, its GPU count, wall-clock
+limit, dependency extras, kernel-build policy, secrets, and artifacts. The
+internal suffix is an implementation detail; there is only one active backend.
 
-1. Reads Buildkite secrets for Modal, Hugging Face, and W&B when needed.
-2. Selects a Modal function based on `TEST_TYPE`.
-3. Passes Buildkite metadata into the Modal container.
-4. Runs the selected test command from `fastvideo/tests/modal/pr_test.py` or
-   `fastvideo/tests/modal/ssim_test.py`.
-5. Uploads performance artifacts for `TEST_TYPE=performance`.
+SSIM uses `fastvideo/tests/ssim/ci_runner.py` inside a single four-GPU lease.
+It discovers `REQUIRED_GPUS` and `*_MODEL_TO_PARAMS` with AST parsing, then
+packs independent pytest subprocesses across the visible GPUs with fail-fast
+termination. Performance writes reports to a host-mounted artifact directory;
+the host uploads only `.md`, `.html`, `.json`, and `.csv` files after the
+container exits.
+
+The Modal launchers remain in the repository for manual rollback archaeology,
+but they are not CI entrypoints. `pr_test.sh` rejects Buildkite calls and needs
+`FASTVIDEO_ENABLE_LEGACY_MODAL_CI=1` even for a local manual invocation.
 
 If you add a new CI test category:
 
-1. Add the Modal function in `fastvideo/tests/modal/pr_test.py` or a focused
-   companion module.
-2. Add the `TEST_TYPE` case in `.buildkite/scripts/pr_test.sh`.
-3. Add the Buildkite direct-test step and any Fastcheck/Full Suite path filters
-   in `.buildkite/pipeline.yml`.
-4. Add or update the `/test` mapping in `.github/workflows/ci-slash-commands.yml`.
+1. Add an executable `.buildkite/scripts/lanes/<lane>.sh` containing the test
+   payload only.
+2. Add the static Buildkite step in `.buildkite/pipeline.yml`, its source/test
+   ownership in `.github/scripts/plan_merge_ci.py`, and the `/test` mapping in
+   `.github/workflows/ci-slash-commands.yml`.
+3. Extend `fastvideo/tests/contract/test_ci_test_collection.py`,
+   `test_merge_ci_plan.py`, and the trusted private runner's lane table plus
+   uploader policy in the same rollout.
+4. Validate on the target GB200 hardware before making the lane a merge gate.
 5. Document the lane here and link any domain-specific authoring guide.
 
 ## CD And Release Workflows
@@ -282,13 +407,18 @@ the explicit `py3.12-cuda13.0.0-latest` tag. This publication policy does not
 change the unparameterized `docker/Dockerfile` build defaults, which remain CUDA
 13 and `cu130`.
 
-Published amd64 development images keep their configured Hopper kernel wheel
-installed and also carry an immutable SM89 wheel under
-`/opt/fastvideo-kernel-prebuilt`. Modal PR and SSIM jobs select the exact
-source, ABI, and GPU-architecture match from that directory, so L40S jobs reuse
-the trusted image artifact while kernel-changing PRs still build locally. Once
-a kernel or artifact-key change reaches `main`, the image workflow republishes
-the matching trusted artifact before later jobs consume the updated image tag.
+Published development images carry architecture-specific kernel wheels under
+`/opt/fastvideo-kernel-prebuilt`. The Slurm worker selects the exact source,
+ABI, and GPU-architecture match, so normal lanes reuse the trusted artifact
+while kernel-changing PRs still build locally. Once a kernel or artifact-key
+change reaches `main`, the image workflow republishes the matching artifact
+before later jobs consume the updated image pin.
+
+The same workflow publishes a single-architecture ARM64, CUDA 13, SM100 image
+for the self-hosted CI runner under the
+`py3.12-cuda13.0.0-sm100-{latest,sha-*}` tags. It carries the matching prebuilt
+kernel wheel so runner jobs can validate and install the exact source and ABI
+match instead of recompiling it in every lane.
 
 The optional Dreamverse matrix builds backend and UI images for CUDA 12.6 and
 CUDA 13 on `amd64`. Dreamverse remains `amd64`-only because its FA4 dependency
@@ -320,15 +450,17 @@ The reusable implementation lives in
 | `.github/mergify.yml` | Merge protection, PR title validation, PR labels, conflict labels, auto-merge |
 | `.github/workflows/ci-precommit.yml` | Tier 1 pre-commit |
 | `.github/workflows/ci-slash-commands.yml` | `/merge` and `/test` handling |
-| `.github/workflows/ci-trigger-full-suite.yml` | Full Suite trigger for `ready` PRs and new pushes to ready PRs |
-| `.github/workflows/ci-aggregate-status.yml` | Aggregate Fastcheck/Full Suite commit statuses |
-| `.buildkite/pipeline.yml` | Buildkite test graph and path filters |
-| `.buildkite/scripts/pr_test.sh` | Buildkite-to-Modal test dispatcher |
-| `fastvideo/tests/modal/pr_test.py` | Modal functions for most GPU CI lanes |
-| `fastvideo/tests/modal/ssim_test.py` | Modal functions and partitioning for SSIM |
+| `.github/workflows/ci-trigger-full-suite.yml` | Change-aware merge-gate trigger for `ready` PRs and new pushes |
+| `.github/workflows/ci-scheduled-ssim.yml` | Weekly complete SSIM trigger on `main` |
+| `.github/scripts/plan_merge_ci.py` | Trusted changed-path to integration-lane planner |
+| `.github/workflows/ci-aggregate-status.yml` | Aggregate Fastcheck and explicit Full Suite direct-rerun statuses |
+| `.buildkite/pipeline.yml` | Static 20-lane Slurm Buildkite graph |
+| `.buildkite/scripts/unit_test.sh`, `.buildkite/scripts/lanes/*.sh` | Active Slurm lane payloads |
+| `fastvideo/tests/ssim/ci_runner.py` | Four-GPU Slurm SSIM scheduler |
+| `.buildkite/scripts/pr_test.sh`, `fastvideo/tests/modal/*.py` | Dormant manual Modal rollback path (disabled in Buildkite) |
 | `.buildkite/performance-benchmarks/tests/*.json` | Performance benchmark configs and thresholds |
 | `.github/workflows/infra-docs.yml` | Docs build and GitHub Pages deploy |
-| `.github/workflows/infra-build-image.yml` | Automatic CUDA matrix and manual Docker image builds |
+| `.github/workflows/infra-build-image.yml` | CUDA matrix, CI runner image, and manual Docker image builds |
 | `.github/workflows/publish-fastvideo.yml` | FastVideo PyPI publishing |
 | `.github/workflows/publish-kernel.yml` | FastVideo kernel PyPI publishing |
 | `.github/workflows/publish-comfyui.yml` | ComfyUI registry publishing |

docs/contributing/performance_benchmarks.md

@@ -23,8 +23,8 @@ It serves three audiences:
 pytest fastvideo/tests/performance/ -vs
 
 # Optional: compare against the rolling HF baseline.
-# PERF_REPORTS_DIR defaults to /root/data/perf_reports for Modal/CI, so
-# override it when running outside the container.
+# PERF_REPORTS_DIR defaults to /root/data/perf_reports in a CI container, so
+# override it for a local run.
 PERF_REPORTS_DIR=/tmp/fastvideo_perf_reports \
 python fastvideo/tests/performance/compare_baseline.py
 
@@ -459,16 +459,18 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
 | `FASTVIDEO_FA4` | `0` | `test_inference_performance.py` | FlashAttention-4 toggle included in `software_profile_id`. |
 | `FASTVIDEO_PERFORMANCE_PROFILE_VERSION` | unset | `test_inference_performance.py` | Optional explicit software cohort/profile version included in `software_profile_id`. |
 | `IMAGE_VERSION` | unset | `test_inference_performance.py` | CI container image/profile version included in `software_profile_id` when available. |
-| `FASTVIDEO_CONTAINER_IMAGE_REF` | unset | `pr_test.py`, `launch_l40s_job.py`, `test_inference_performance.py` | Resolved CI container image ref or digest recorded in `environment_metadata` for audit without changing `software_profile_id`. |
+| `FASTVIDEO_CONTAINER_IMAGE_REF` | unset | Slurm runner, `test_inference_performance.py` | Pinned CI container image digest recorded in `environment_metadata` and `software_profile_id`. |
 | `FASTVIDEO_STAGE_LOGGING` | set by the pytest test | `test_inference_performance.py` | Enables pipeline stage timing capture for component metrics during benchmark runs. |
 
 ## CI integration
 
-The performance step can run on demand with `/test performance` and as part of
-the Full Suite (see [CI/CD Architecture](ci_architecture.md)). The Modal entry
-point is `fastvideo/tests/modal/pr_test.py:run_performance_tests` and the
-Buildkite artifact upload is in
-`.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.
+The performance step can run on demand with `/test performance`, through a
+merge gate when performance tests or benchmark policy changed, and as part of
+an explicit `/test full` run (see [CI/CD Architecture](ci_architecture.md)).
+The weekly `fastvideo-performance-lane` schedule runs the same Slurm payload.
+The active entry point is `.buildkite/scripts/lanes/performance.sh`; the
+trusted host dispatcher relays its allowlisted reports to Buildkite after the
+isolated container exits.
 
 Each performance build runs pytest first. PR and direct runs only continue to
 `compare_baseline.py` when that fixed-threshold phase passes; if pytest fails,

docs/contributing/pull_requests.md

@@ -42,19 +42,20 @@ The important process labels are:
 
 | Label | Meaning |
 |---|---|
-| `ready` | The PR is ready for Full Suite and auto-merge consideration. |
+| `ready` | The PR is ready for the change-aware merge gate and auto-merge consideration. |
 | `needs-rebase` | The PR has merge conflicts with `main`. |
 | `do-not-merge` | A maintainer has blocked merge. |
 
 ## CI Summary
 
-FastVideo has three validation tiers:
+FastVideo has three routine validation tiers plus an explicit full diagnostic:
 
 | Tier | Runs when | What it does |
 |---|---|---|
 | Pre-commit | Pull requests and `/test pre-commit` | Formatting, linting, typing, spelling, Markdown, workflow syntax, filename checks |
-| Fastcheck | PR Buildkite builds | Path-filtered component and unit checks on Modal GPU runners |
-| Full Suite | `/merge`, `ready`, `/test full`, or new pushes to ready PRs | Path-filtered integration, SSIM, training, eval, API, and performance checks |
+| Fastcheck | PR Buildkite builds | Six component, kernel, unit, and app lanes on Slinky Slurm |
+| Merge gate | `/merge`, `ready`, or new pushes to ready PRs | Only path-relevant integration lanes on Slinky Slurm; Fastcheck remains the baseline |
+| Full Suite | `/test full` | Explicit all-twenty-lane diagnostic run on Slinky Slurm |
 
 See [CI/CD Architecture](ci_architecture.md#ci-tiers) for exact jobs, path
 filters, and workflow files.
@@ -66,10 +67,11 @@ filters, and workflow files.
 3. Fix pre-commit failures locally with `pre-commit run --all-files`.
 4. Wait for at least one approving review.
 5. When the PR is approved and ready, comment `/merge`.
-6. `/merge` adds `ready` and triggers the Full Suite for the PR branch.
+6. `/merge` adds `ready`, waits for cheap checks, and triggers the minimal
+   path-relevant integration lanes for the PR branch.
 7. If all required checks pass, Mergify squash-merges the PR to `main`.
-8. If Full Suite fails, fix the regression, push again, and re-run `/merge` or
-   the failed test.
+8. If the merge gate fails, fix the regression, push again, and re-run
+   `/merge`. Use a targeted `/test` command for diagnosis.
 
 Only contributors with repository write permission can use slash commands. If
 you are an external contributor, ask a maintainer to run `/merge` or add
@@ -128,7 +130,7 @@ git push --force-with-lease
 
 Mergify removes `needs-rebase` after conflicts are resolved.
 
-### Full Suite Fails
+### Merge Gate Or Full Suite Fails
 
 The failing Buildkite step is the source of truth. Common causes are:
 

docs/contributing/testing.md

@@ -138,47 +138,31 @@ pytest fastvideo/tests/ssim/ -vs
 
 Use a machine whose GPU and backend match the reference folder you are testing.
 
-## Modal Runs For SSIM
+## Slurm CI Runs For SSIM
 
-For CI-like SSIM execution, use `fastvideo/tests/modal/ssim_test.py`:
+Comment `/test ssim` on a pull request to run the canonical four-GPU SSIM
+lane on the Slinky Slurm cluster. `fastvideo/tests/ssim/ci_runner.py`
+discovers the suite without importing test modules, packs independent pytest
+processes across the four assigned GPUs, and stops the lane on the first
+failure.
+
+The change-aware `/merge` planner may run only the SSIM test basenames owned
+by the changed model family. Shared SSIM harness changes still select the
+complete lane. Independently, `main` runs the full SSIM matrix every Sunday at
+05:00 UTC so infrequently touched model families retain periodic coverage.
+
+For a focused developer run, invoke pytest directly and optionally select one
+model from a parameterized test through `FASTVIDEO_SSIM_MODEL_ID`:
 
 ```bash
-python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests
+pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py -vs
+
+FASTVIDEO_SSIM_MODEL_ID=Wan2.1-T2V-1.3B-Diffusers \
+pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py -vs
 ```
 
-Target specific files or model ids:
-
-```bash
-python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests \
-  --test-files test_wan_t2v_similarity.py \
-  --model-ids Wan2.1-T2V-1.3B-Diffusers
-```
-
-If `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` is not set, the local
-entrypoint fails fast.
-
-To export raw generated videos from Modal to the shared volume:
-
-```bash
-python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests \
-  --sync-generated-to-volume
-```
-
-The raw export path is quality-tiered:
-
-- default params: `ssim_generated_videos/default/<subdir>/generated_videos`
-- full-quality params: `ssim_generated_videos/full_quality/<subdir>/generated_videos`
-
-The printed `modal volume get` command downloads into
-`./generated_videos_modal/<quality-tier>`. Convert those outputs into local
-references with `copy-local`:
-
-```bash
-python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
-  --quality-tier full_quality \
-  --generated-dir ./generated_videos_modal/full_quality/L40S_reference_videos \
-  --device-folder L40S_reference_videos
-```
+The files under `fastvideo/tests/modal/` are retained only as a disabled
+manual rollback implementation. No active CI trigger invokes them.
 
 ### SSIM Bootstrap Mode
 
@@ -206,28 +190,30 @@ python fastvideo/tests/ssim/reference_videos_cli.py promote-draft \
 
 ## CI Integration
 
-FastVideo CI tests are orchestrated by Buildkite and run on Modal GPU
-instances. The main files are:
+FastVideo GPU CI is orchestrated by Buildkite and runs only on isolated Slinky
+Slurm workers. The main files are:
 
 | File | Purpose |
 |---|---|
-| `.buildkite/pipeline.yml` | Buildkite test graph and path filters. |
-| `.buildkite/scripts/pr_test.sh` | Dispatches `TEST_TYPE` to a Modal function. |
-| `fastvideo/tests/modal/pr_test.py` | Modal functions for most test lanes. |
-| `fastvideo/tests/modal/ssim_test.py` | Modal functions and partitioning for SSIM. |
+| `.buildkite/pipeline.yml` | Static, validated 20-lane Slurm test graph. |
+| `.github/scripts/plan_merge_ci.py` | Trusted path-to-lane and focused quality-test policy for `/merge`. |
+| `.buildkite/scripts/unit_test.sh`, `.buildkite/scripts/lanes/*.sh` | Repository-owned test payloads executed inside Slurm containers. |
+| `fastvideo/tests/ssim/ci_runner.py` | Four-GPU SSIM task discovery and scheduling. |
+| `.buildkite/scripts/pr_test.sh`, `fastvideo/tests/modal/*.py` | Dormant manual rollback path; rejected in Buildkite. |
 
-For exact tier membership, path filters, slash commands, and aggregate statuses,
+For exact tier membership, slash commands, runner isolation, and aggregate statuses,
 see [CI/CD Architecture](ci_architecture.md).
 
 ### Adding A New CI Test Category
 
 If a new test does not fit an existing lane:
 
-1. Add a Modal function in `fastvideo/tests/modal/pr_test.py` or a focused
-   companion module.
-2. Add a matching `TEST_TYPE` case in `.buildkite/scripts/pr_test.sh`.
-3. Add Buildkite direct-test and path-filter entries in `.buildkite/pipeline.yml`.
-4. Add the `/test` mapping in `.github/workflows/ci-slash-commands.yml`.
+1. Put the test payload in an executable `.buildkite/scripts/lanes/<lane>.sh`.
+2. Add its static step to `.buildkite/pipeline.yml`, its changed-path ownership
+   to `.github/scripts/plan_merge_ci.py`, and extend the CI contract tests.
+3. Add the `/test` mapping in `.github/workflows/ci-slash-commands.yml`.
+4. Coordinate the matching GPU, timeout, dependency, secret, and artifact
+   policy in the private Slurm runner allowlist.
 5. Document the new category in [CI/CD Architecture](ci_architecture.md) and add
    authoring notes here if contributors need them.
 

docs/cookbook/index.md

@@ -0,0 +1,44 @@
+# Inference Cookbook
+
+Choose a complete recipe maintained in the FastVideo repository. Each command
+runs its checked-in source directly, so coupled model, GPU, offload, and
+attention settings do not drift into unsupported combinations.
+
+The commands expect a local clone:
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git
+cd FastVideo
+```
+
+<div class="cookbook-picker" data-cookbook data-recipes="../assets/cookbook-recipes.json">
+  <label for="cookbook-recipe"><strong>Recipe</strong></label>
+  <select id="cookbook-recipe" data-cookbook-recipe disabled>
+    <option>Loading recipes…</option>
+  </select>
+  <dl>
+    <dt>Model</dt>
+    <dd data-cookbook-model>Loading…</dd>
+    <dt>Source</dt>
+    <dd><a data-cookbook-source href="../inference/examples/basic/">Browse maintained examples</a></dd>
+  </dl>
+  <pre><code class="language-bash" data-cookbook-command>Loading…</code></pre>
+  <p class="cookbook-picker__status" role="status" aria-live="polite" data-cookbook-status></p>
+  <noscript>
+    JavaScript is needed for the recipe picker. Browse the
+    <a href="../inference/examples/examples_inference_index/">inference examples</a>
+    instead.
+  </noscript>
+</div>
+
+## Customize a recipe
+
+Start from the checked-in source, then change only the settings your model
+supports:
+
+- [Configuration](../inference/configuration.md) covers the Python and CLI
+  config surfaces.
+- [Optimizations](../inference/optimizations.md) covers attention backends,
+  compilation, and memory tradeoffs.
+- [Support matrix](../inference/support_matrix.md) lists supported models and
+  optimizations.

docs/design/apple_silicon_fast_mode.md

@@ -0,0 +1,129 @@
+# Fast mode (RIFE) — Apple Silicon
+
+`--fast` makes local generation ~2.7× faster by **generating fewer frames and
+interpolating the rest** with an Apple-Silicon-native RIFE model, instead of
+denoising every frame. Video-diffusion denoise is dominated by self-attention,
+which is O(tokens²); halving the frames cuts the token count ~2× and the denoise
+compute ~3.7×, so the wall-clock drops far more than 2×. RIFE (which estimates
+its own optical flow — no motion vectors needed) fills the dropped frames back
+in for ~1.4 s, and a light unsharp pass counters its softening.
+
+Measured on the 1.3B INT8 QAD model (480×832×81, M4): generate 41 + RIFE→81
+runs in ~35 s of denoise vs ~90 s full, at reconstruction MS-SSIM **0.97**.
+Reproduce with `python -m fastvideo.benchmarks.eval_metalfx_rife --mode int8`.
+
+> **Note:** Apple's *MetalFX* frame interpolation is **not** usable here — it
+> requires game-engine motion vectors + depth, which diffusion output lacks. We
+> use the video-native **`rife-mlx`** model instead (Metal-backed, torch-free).
+
+## Install
+
+```bash
+uv pip install -e ".[mlx]"   # RIFE ships vendored; this only needs MLX
+```
+
+## Use
+
+```bash
+python examples/inference/basic/mlx_wan_prompt_to_video.py \
+  --model-root ./FastMetal-1.3B-QAD \
+  --mlx-checkpoint ./FastMetal-1.3B-QAD \
+  --prompt "A bird's-eye view of a misty forest valley at dawn." \
+  --num-frames 81 --fast \
+  --output-path video_samples/forest_fast.mp4
+```
+
+`--num-frames` stays the *target* length; fast mode generates the smallest
+VAE-aligned keyframe count that RIFE can interpolate to that target.
+
+| Flag | Default | Meaning |
+|---|---|---|
+| `--fast` / `--no-fast` | off | enable fast mode |
+| `--fast-factor` | 2 | generate 1/factor of the frames (2 = half) |
+| `--fast-sharpen` | 0.6 | light unsharp strength to counter RIFE softness (0 disables) |
+
+Fast mode composes with everything else (`--mlx-quantization int8`,
+`--mlx-compile`, TAEHV vs `--decode-backend wan-vae`). Keep `--fast-factor` at 2
+for quality — larger temporal gaps are where RIFE starts inventing motion.
+
+## Spatial fast mode (`--fast-spatial`)
+
+The spatial twin of `--fast`: instead of dropping frames, drop pixels. Denoise
+*and decode* at `height/width // fast-spatial-scale`, then resample the decoded
+frames up to the requested size. Self-attention is O(tokens²), so halving each
+spatial axis cuts the token count 4× and the denoise time far more than that —
+measured on the 1.3B INT8 QAD model at 480×832×81, M4 Max: **86.1 s → 10.3 s**
+of denoise. It composes with `--fast`; both together run the same clip in
+**4.5 s** of denoise.
+
+```bash
+python examples/inference/basic/mlx_wan_prompt_to_video.py \
+  --prompt "A bird's-eye view of a misty forest valley at dawn." \
+  --height 480 --width 832 --num-frames 81 --fast-spatial \
+  --output-path video_samples/forest_fast_spatial.mp4
+```
+
+| Flag | Default | Meaning |
+|---|---|---|
+| `--fast-spatial` / `--no-fast-spatial` | off | enable spatial fast mode |
+| `--fast-spatial-scale` | 2 | denoise at 1/scale of each spatial axis |
+| `--fast-spatial-upsample-mode` | `lanczos` | pixel interpolation kernel (`lanczos`, `cubic`, `bilinear`, `nearest`) |
+| `--fast-spatial-sharpen` | 0.4 | light unsharp strength to counter resampling softness (0 disables) |
+
+### The upsample must happen in pixel space
+
+This is the one thing to get right. The obvious implementation — bilinearly
+upsample the finished latents and decode at the target size — **does not work**,
+and produces a distinctive failure: correct composition and silhouette under a
+smeared, hazy veil, with ringing along strong edges.
+
+A Wan latent cell is a *learned code* for an 8×8 (Wan2.1) or 16×16 (Wan2.2)
+pixel block, not a low-pass sample of the image. The average of two adjacent
+codes is not the code of the averaged blocks; it is a vector the decoder was
+never trained on. Measured on Wan2.1-1.3B at 480×832, a 2× bilinear latent
+upsample destroys **62%** of the latent's high-frequency energy while leaving
+its overall magnitude intact — exactly the signature of that veil. At Wan2.2-5B
+the same operation degrades to black or noise.
+
+Decoded RGB frames have no such problem: an image *is* a sampled 2-D signal, so
+Lanczos interpolation is the operation it was defined for. The result is soft —
+it carries stage-1's real detail budget and no more — but clean and coherent.
+
+`--refine` gets away with a latent-space upsample only because a second DMD pass
+re-denoises the hand-off; spatial fast mode passes the latent straight to the
+decoder, so it cannot.
+
+## Refine (`--refine`) stage-2 timesteps
+
+`--refine` hands stage 1 to stage 2 as `(1 - sigma) * upsampled + sigma * noise`,
+where `sigma` comes from the *first* stage-2 timestep. FastWan's DMD grid opens
+at `t=1000`, which is `sigma == 1` exactly — so a stage-2 grid that starts there
+weights the stage-1 result at zero and refine silently degrades into a plain
+full-resolution run at twice the cost.
+
+Left unset, `--refine-dmd-denoising-steps` now derives the stage-2 grid from the
+stage-1 one with leading full-noise steps dropped (`1000,757,522` → `757,522`).
+That keeps the pass on timesteps the distilled student was trained on while
+letting stage-1 structure through: hand-off `sigma = 0.757`, stage-1 weight
+`0.243`. Passing a grid that starts at full noise is now an error rather than a
+silently wasted pass.
+
+The run prints the resolved hand-off so it is visible:
+
+```
+[refine] stage-2 hand-off sigma=0.7568 (stage-1 weight 0.2432)
+```
+
+There is a trade-off in choosing that grid. Later start = more of the draft
+survives, but fewer stage-2 steps. On Wan2.1 the default `757,522` gives weight
+0.243 with two steps; `--refine-dmd-denoising-steps 522` gives weight 0.478 with
+one. `--refine-sigma` decouples the noise level from the timestep entirely — it
+logs a warning, because the DiT is then told a timestep that does not match the
+noise it receives.
+
+**Wan2.2-5B has a lower ceiling.** Its warped schedule maps `1000,757,522` to
+sigmas `1.000, 0.940, 0.845`, so the best available stage-1 weight is **0.060**
+(vs 0.243 at 1.3B). Un-warped (`--no-warp`) the same grid gives `1.000, 0.757,
+0.522` and a weight of 0.243 — but warping is what matches the FastVideo
+sampling schedule, so turning it off changes the timesteps the distilled student
+sees. Which is better at 5B is unresolved and needs a run on real 5B weights.

docs/design/inference_schema_parity_inventory.yaml

@@ -76,6 +76,11 @@ surfaces:
       lora_target_modules: "Legacy LoRA configuration surface pending dedicated component API."
       output_type: "Legacy output formatting surface pending GenerationResult cleanup."
       VSA_sparsity: "Model-specific inference optimization not yet represented in the typed public schema."
+      VSA_tile_size: "VSA-H3 tile geometry request; model-specific optimization not yet represented in the typed public schema."
+      inference_torch_compile: "Regional inference compile opt-in currently carried through PipelineSelection.experimental rather than CompileConfig."
+      vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema."
+      vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema."
+      vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema."
       attention_backend: "Process-wide default attention-backend request applied per component at load time; kernel-selection knob not yet represented in the typed public schema."
       moba_config_path: "Model-specific MoBA optimization surface not yet represented in the typed public schema."
       master_port: "Executor/bootstrap compatibility field; not part of the canonical inference schema."
@@ -119,12 +124,14 @@ surfaces:
       vae_precision: "Precision override pending dedicated typed component precision design."
       vae_decode_precision: "Decode-only precision override pending dedicated typed component precision design."
       image_encoder_precision: "Precision override pending dedicated typed component precision design."
+      image_encoder_precisions: "Precision overrides pending dedicated typed component precision design."
       text_encoder_precisions: "Precision override pending dedicated typed component precision design."
     internal_only:
       dit_config: "Legacy internal component config object."
       upsampler_config: "Legacy internal component config object."
       vae_config: "Legacy internal component config object."
       image_encoder_config: "Legacy internal component config object."
+      image_encoder_configs: "Legacy internal component config objects."
       text_encoder_configs: "Legacy internal component config object."
       preprocess_text_funcs: "Internal text preprocessing hooks."
       postprocess_text_funcs: "Internal text postprocessing hooks."
@@ -361,6 +368,30 @@ surfaces:
         sources: [fastvideo.configs.pipelines.matrixgame2.MatrixGame2I2V480PConfig]
       num_frames_per_block:
         sources: [fastvideo.configs.pipelines.matrixgame2.MatrixGame2I2V480PConfig]
+      duration_s:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      spectrogram_frame_rate:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      latent_downsample_rate:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      clip_frame_rate:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      sync_frame_rate:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      sync_segment_size:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      sync_segment_stride:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      sync_downsample_rate:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      clip_image_size:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      sync_image_size:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      clip_batch_size_multiplier:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
+      sync_batch_size_multiplier:
+        sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
       audio_channels:
         sources:
           - fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
@@ -375,6 +406,7 @@ surfaces:
           - fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
       max_audio_duration_s:
         sources:
+          - fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig
           - fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
           - fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
       sample_size:
@@ -383,6 +415,7 @@ surfaces:
           - fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
       sampling_rate:
         sources:
+          - fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig
           - fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
           - fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
       audio_txt_guidance_scale:

docs/generate_examples.py

@@ -2,6 +2,7 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py
 
 import itertools
+import json
 import os
 import re
 from dataclasses import dataclass, field
@@ -19,6 +20,40 @@ GENERATED_DOC_PREFIXES = (
     "training/examples/",
     "distillation/examples/",
 )
+COOKBOOK_DATA = ROOT_DIR / "docs/assets/cookbook-recipes.json"
+COOKBOOK_SOURCE_ROOTS = (
+    ROOT_DIR / "examples/inference",
+    ROOT_DIR / "scripts/inference",
+)
+
+
+def validate_cookbook() -> None:
+    """Keep cookbook entries tied to checked-in runnable sources."""
+    recipes = json.loads(COOKBOOK_DATA.read_text(encoding="utf-8")).get("recipes")
+    if not isinstance(recipes, list) or not recipes:
+        raise ValueError(f"{COOKBOOK_DATA}: recipes must be a non-empty list")
+
+    seen: set[str] = set()
+    for recipe in recipes:
+        required = ("id", "task", "label", "model", "source", "command")
+        missing = {key for key in required if not recipe.get(key)}
+        if missing:
+            raise ValueError(f"Cookbook recipe is missing: {', '.join(sorted(missing))}")
+        if recipe["id"] in seen:
+            raise ValueError(f"Duplicate cookbook recipe id: {recipe['id']}")
+        seen.add(recipe["id"])
+
+        source = (ROOT_DIR / recipe["source"]).resolve()
+        if not any(source.is_relative_to(root.resolve()) for root in COOKBOOK_SOURCE_ROOTS):
+            raise ValueError(f"Cookbook source is outside an approved directory: {recipe['source']}")
+        if not source.is_file():
+            raise ValueError(f"Cookbook source does not exist: {recipe['source']}")
+
+        source_text = source.read_text(encoding="utf-8")
+        if recipe["model"] not in source_text:
+            raise ValueError(f"Cookbook model is not present in {recipe['source']}: {recipe['model']}")
+        if recipe["source"] not in recipe["command"]:
+            raise ValueError(f"Cookbook command does not invoke its source: {recipe['id']}")
 
 
 def fix_case(text: str) -> str:
@@ -536,6 +571,7 @@ def on_pre_build(config, **kwargs):
     MkDocs hook to generate examples before building the documentation.
     This function is called automatically by MkDocs' native hook system.
     """
+    validate_cookbook()
     print("Generating example documentation...")
     generate_examples(generate_main_index=True)
     print("Example documentation generated successfully!")
@@ -549,6 +585,7 @@ def on_page_context(context, page, **kwargs):
 
 
 if __name__ == "__main__":
+    validate_cookbook()
     print("Generating example documentation...")
     generate_examples(generate_main_index=True)
     print("Example documentation generated successfully!")

docs/getting_started/installation.md

@@ -65,6 +65,7 @@ uv pip install flash-attn --no-build-isolation -v
 
 ## Next Steps
 
-- [Quick Start Guide](quick_start.md) - Get started with your first video generation
+- [Quick Start](quick_start.md) - Generate your first video
+- [Inference Cookbook](../cookbook/index.md) - Choose a maintained recipe
 - [Configuration](../inference/configuration.md) - Learn about configuration options
-- [Examples](../inference/examples/examples_inference_index.md) - Explore example scripts and notebooks
+- [Examples](../inference/examples/examples_inference_index.md) - Explore scripts and notebooks

docs/getting_started/installation/mps.md

@@ -1,6 +1,10 @@
 # MPS (Apple Silicon)
 
-Instructions to install FastVideo for Apple Silicon.
+Install FastVideo on Apple Silicon and run FastMetal-QAD.
+
+Apple Silicon uses the MLX runtime and the FastMetal-QAD INT8 checkpoints.
+See the [FastMetal-QAD blog](https://haoailab.com/blogs/fastmetal/) and the
+[FastMetal collection](https://huggingface.co/collections/FastVideo/fastmetal).
 
 ## Requirements
 
@@ -49,10 +53,12 @@ brew install ffmpeg
 
 ### Installation
 
+FastMetal's native Apple Silicon runtime requires the `mlx` extra.
+
 #### With uv (recommended)
 
 ```bash
-uv pip install fastvideo
+uv pip install "fastvideo[mlx]"
 ```
 
 #### With Conda environment (alternative)
@@ -60,7 +66,7 @@ uv pip install fastvideo
 `uv` works inside an active conda env too, so prefer `uv pip` for the actual install:
 
 ```bash
-uv pip install fastvideo
+uv pip install "fastvideo[mlx]"
 ```
 
 ### Installation from Source
@@ -76,15 +82,56 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 Basic installation:
 
 ```bash
-uv pip install -e .
+uv pip install -e ".[mlx]"
 ```
 
 Alternative with Conda environment:
 
 ```bash
-uv pip install -e .
+uv pip install -e ".[mlx]"
 ```
 
+## Run FastMetal-QAD
+
+Each release is self-contained. Download one checkpoint and point both
+`--model-root` and `--mlx-checkpoint` at it (the example also auto-detects
+`mlx_dit.json` under `--model-root`).
+
+| Checkpoint | Script | Mac tier |
+| --- | --- | --- |
+| [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | `mlx_wan_prompt_to_video.py` | 16 GB+ |
+| [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | `mlx_wan22_generate.py` | 16 GB+ |
+| [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | `mlx_wan_prompt_to_video.py` | 36 GB+ |
+
+```bash
+hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD
+
+python examples/inference/basic/mlx_wan_prompt_to_video.py \
+  --model-root ./FastMetal-1.3B-QAD \
+  --mlx-checkpoint ./FastMetal-1.3B-QAD \
+  --height 480 --width 832 --num-frames 81 \
+  --prompt "A bird's-eye view of a misty forest valley at dawn."
+```
+
+14B uses the same script. Point both flags at `./FastMetal-14B-QAD`. That repo also ships an EMA variant: keep `--model-root` at the repo root and set `--mlx-checkpoint ./FastMetal-14B-QAD/ema`.
+
+Wan2.2 5B uses a different latent layout, so it has its own entrypoint:
+
+```bash
+hf download FastVideo/FastMetal-5B-QAD --local-dir ./FastMetal-5B-QAD
+
+python examples/inference/basic/mlx_wan22_generate.py \
+  --mlx-checkpoint ./FastMetal-5B-QAD \
+  --text-encoder-root ./FastMetal-5B-QAD \
+  --vae-root ./FastMetal-5B-QAD/vae \
+  --height 704 --width 1280 --num-frames 81 \
+  --prompt "A cinematic portrait with soft neon lighting and smooth camera motion."
+```
+
+CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`, `FastVideo/FastWan-QAD-FP8-1.3B`) is a separate NVIDIA release. The MLX examples look for FastMetal packed weights (`mlx_dit.json`).
+
+`basic_mps.py` is a generic PyTorch MPS demo. For local video on Mac, use the FastMetal commands above.
+
 ## Development Environment Setup
 
 If you're planning to contribute to FastVideo please see the following page:
@@ -92,9 +139,9 @@ If you're planning to contribute to FastVideo please see the following page:
 
 ## Hardware Requirements
 
-### For Basic Inference
-
-- Mac M1, M2, M3, or M4 (at least 32 GB RAM is preferable for high quality video generation)
+- **1.3B / 5B:** 16 GB unified memory and up (M1 and later)
+- **14B:** 36 GB unified memory and up
+- Fanless 13-inch MacBook Air can run 1.3B and 5B at the same resolutions
 
 ## Troubleshooting
 

docs/getting_started/installation/spark_performance.md

@@ -156,8 +156,11 @@ is power-cycled. To avoid it:
 
 - **Builds** (flash-attn, kernel): `nice -n 19`, `MAX_JOBS=2`, `nohup`. Never a
   bare foreground high-parallelism build.
-- Leave `*_cpu_offload` at the example defaults — "CPU" offload is the *same*
-  unified RAM on the GB10, so the win is tiling + sane resolution, not offloading.
+- FastVideo automatically disables DiT layerwise/CPU offload and encoder/VAE CPU
+  offload after each worker binds its GB10 device. Do not force those modes back
+  on: "CPU" offload uses the same unified RAM. Multi-GPU FSDP sharding remains
+  available because it partitions weights without parking them in a separate
+  host pool.
 
 ## Gotchas specific to the GB10
 

docs/getting_started/quick_start.md

@@ -23,61 +23,21 @@ Also optionally install flash-attn:
 uv pip install flash-attn --no-build-isolation -v
 ```
 
-## Basic Usage
+## Choose a maintained recipe
 
-### Text-to-Video Generation
+The cookbook selects complete, checked-in recipes instead of mixing model,
+parallelism, offload, and attention settings independently.
 
-```python
-from fastvideo import VideoGenerator
+[Open the inference cookbook](../cookbook/index.md){ .md-button .md-button--primary }
 
-def main():
-    # Create a video generator with a pre-trained model
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,  # Adjust based on your hardware
-    )
-
-    # Define a prompt for your video
-    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
-
-    # Generate the video
-    video = generator.generate_video(
-        prompt,
-        output_path="my_videos/",  # Controls where videos are saved
-        save_video=True
-    )
-
-if __name__ == '__main__':
-    main()
-```
-
-### Image-to-Video Generation
-
-```python
-from fastvideo import VideoGenerator, SamplingParam
-
-def main():
-    # Create the generator
-    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
-
-    # Set up parameters with an initial image
-    sampling_param = SamplingParam.from_pretrained(model_name)
-    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
-    sampling_param.num_frames = 107
-
-    # Generate video based on the image
-    prompt = "A photograph coming to life with gentle movement"
-    generator.generate_video(prompt, sampling_param=sampling_param,
-                             output_path="my_videos/",
-                             save_video=True)
-
-if __name__ == '__main__':
-    main()
-```
+!!! tip "Need more control?"
+    Start from a maintained recipe, then use the
+    [configuration](../inference/configuration.md) and
+    [optimization](../inference/optimizations.md) guides for supported changes.
 
 ## Next Steps
 
+- [Inference Cookbook](../cookbook/index.md) - Choose a maintained recipe
 - [Installation Guide](installation.md) - Detailed installation instructions
 - [Configuration](../inference/configuration.md) - Learn about configuration options
 - [Examples](../inference/examples/examples_inference_index.md) - Explore more

docs/inference/offloading.md

@@ -14,6 +14,13 @@ vae_cpu_offload: bool = True
 pin_cpu_memory: bool = True
 ```
 
+On unified-memory accelerators such as NVIDIA GB10 and Apple silicon, FastVideo
+detects the selected device inside each worker and disables all five host-offload
+modes before loading modules. Host and accelerator allocations share one physical
+pool there, so offload adds transfers and duplicate residency instead of freeing
+memory. CUDA FSDP sharding remains enabled when requested; MPS continues to
+disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged.
+
 ## Behavior Explanation
 
 !!! note

docs/inference/optimizations.md

@@ -88,6 +88,7 @@ runtime on some GPU/shape combinations. To use FA4, install the pinned
 `flash-attn-4` build (see the `flash-attn-4` source in `pyproject.toml`) and set:
 
 ```bash
+UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"
 export FASTVIDEO_FA4=1
 ```
 
@@ -97,6 +98,21 @@ sm90+) and GQA attention (FA4's `pack_gqa` fails to JIT-compile below sm90).
 On sm90+ both run on FA4. If FA4 is unusable while `FASTVIDEO_FA4=1` is set,
 FastVideo fails loudly instead of silently falling back.
 
+MiniMax-H3 can additionally use FA4's packed-varlen entry point for its long,
+single-sequence dense DiT self-attention:
+
+```bash
+export FASTVIDEO_FA4=1
+export FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN=1
+```
+
+This route is inference-only and remains disabled by default. Runtime guards
+keep masked attention, batch sizes above one, unequal query/key lengths,
+grad-enabled calls, and NVFP4 on their established paths. It does not apply to
+the Preview checkpoint's sparse VSA blocks. Packed-varlen changes floating-point
+reduction order relative to fixed-length FA4, so treat it as a speed/quality
+evaluation option rather than an exact-parity mode.
+
 ### FP4 Flash Attention 4 (Blackwell only)
 
 **`FLASH_ATTN`** with **`--nvfp4_fa4`**
@@ -337,7 +353,42 @@ Only DiT submodules that declare `_compile_conditions` are compiled
 (most shipped models). The text encoder and VAE are not compiled by this
 flag.
 
-### What to expect
+### Regional fullgraph compile (experimental)
+
+`inference_torch_compile` is a stricter, kwargs-free variant that ports the
+training-side regional compile of
+[#1718](https://github.com/hao-ai-lab/FastVideo/pull/1718) to inference: the
+loader wraps each `_compile_conditions` block in
+`torch.compile(fullgraph=True)` with inductor
+`options={"emulate_precision_casts": True}` right after the transformer
+loads. The ordinary compile path keeps its historical compiler-disabled
+attention boundary by default; regional compile opts in only the compatible
+attention instances owned by this transformer. MiniMax-H3 VSA is supported
+only by the inference-only sm_100a tile-64 route
+(`FASTVIDEO_VSA_SM100A=1` and `VSA_tile_size=64`); the loader probes that
+route before capture and keeps the transformer eager when the kernel or
+device is unsupported. Legacy VSA, MiniMax-H3 tile-256 VSA, and the explicit
+`FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep the transformer
+eager with one warning instead of failing mid-denoise.
+
+```python
+generator = VideoGenerator.from_pretrained(
+    "MiniMaxAI/MiniMax-H3",
+    inference_torch_compile=True,  # or FASTVIDEO_INFERENCE_TORCH_COMPILE=1
+)
+```
+
+Do not combine it with `torch_compile_kwargs['mode']` (the loader injects
+inductor options, and torch.compile forbids mode+options); it is
+independent of `enable_torch_compile`, and when both are set the regional
+compile wins for the DiT.
+
+### What to expect from generic compile
+
+The Wan result below measures the existing generic
+`enable_torch_compile=True` path. It is useful evidence that compile can help,
+but it is **not** a benchmark or numerical gate for the stricter regional
+fullgraph path above.
 
 | Config | Effect |
 |---|---|
@@ -353,6 +404,24 @@ generations with the same input shapes. Always exclude the first
 (warmup) generation when measuring steady-state latency — measuring the
 warmup is the most common way to wrongly conclude "compile is slower".
 
+**Regional MiniMax-H3 accuracy caveat (job 2660).** On one GB200 at
+768×1344×124, the native 50-point schedule ran exactly 49 transformer
+forwards. After one warmup, three fixed-prompt/fixed-seed repeats averaged
+**185.08s → 157.01s end to end** and **174.90s → 147.12s denoising**. Each
+leg was independently pixel-deterministic, but compiled output did **not**
+match eager: mean absolute pixel error **20.247/255**, PSNR **16.67 dB**,
+mean SSIM **0.7108**, and mean MS-SSIM **0.6370** across 124 frames. Treat
+regional MiniMax-H3 compile as an opt-in performance experiment, not an
+eager-parity-safe mode.
+
+The same caveat applies to sparse MiniMax-H3 regional compile. Its mask
+compaction, sm_100a launch, trained compression gates, and inference-only H3
+fusions are fullgraph-compatible, but compilation can still change model
+numerics. The FastH3 `all` profile enables regional compile by default because
+it is the fastest measured route at 124, 243, and 345 frames. Use
+`--no-inference-torch-compile` when comparing against the eager sparse-DiT
+route.
+
 **Numerics.** Inductor's lowering is designed to preserve eager
 semantics within floating-point tolerance, but per-model equivalence is
 not asserted by any standing SSIM regression here — the SSIM tests in
@@ -371,9 +440,14 @@ MS-SSIM gate on *your* config, especially when combining
   hao-ai-lab/FastVideo#1365 — keep that fix to get a clean compiled
   region under the default offload path.
 - **`mode="reduce-overhead"` / CUDA graphs**: not yet supported
-  end-to-end. The attention dispatch is an untraceable custom op and
-  still breaks the graph, which CUDA-graph trees cannot span. Use the
-  default inductor mode (shown above) until that is resolved.
+  by regional compile because that path injects inductor `options`, and
+  PyTorch rejects `mode` together with `options`. The generic compile path
+  can accept `mode`, but CUDA-graph compatibility remains backend- and
+  shape-dependent. FA2/FA3 inference and FA4 expose traceable custom-op
+  boundaries. MiniMax-H3's sm_100a tile-64 inference route is the only VSA
+  path in the regional support envelope; other VSA paths and the FA3
+  grad-enabled path remain outside it. Use the default inductor mode shown
+  above unless your exact configuration has its own gate.
 
 Extra `torch.compile` options are passed through `torch_compile_kwargs`
 (a dict), accepted by `VideoGenerator.from_pretrained(...)` and by the

docs/inference/support_matrix.md

@@ -64,6 +64,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | ltx2 | `FastVideo/LTX2-Distilled-Diffusers`<br>`FastVideo/LTX2.3-Distilled-Diffusers`<br>`FastVideo/LTX-2.3-Distilled-Diffusers` | T2V | [basic_ltx2_distilled.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_distilled.py) |
 | ltx2 | `Lightricks/LTX-2.3`<br>`FastVideo/LTX2.3-base`<br>`FastVideo/LTX2.3-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
 | ltx2 | `Lightricks/LTX-2`<br>`FastVideo/LTX2-base`<br>`FastVideo/LTX2-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
+| mmaudio | `FastVideo/MMAudio-large-44k-v2-Diffusers` | V2A, T2A | [basic_mmaudio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mmaudio.py) |
 | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-Base-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`<br>`mignonjia/mg_longtuning_distilled_zelda`<br>`mignonjia/mg_sf_distilled_zelda_1k_steps`<br>`mignonjia/mg_sf_distilled_zelda`<br>`mignonjia/mg_causal_zelda`<br>`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) |
 | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
 | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)<br>[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)<br>[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
@@ -94,6 +95,10 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 (`StableAudioT2AConfig` / `StableAudioOpenSmallConfig`); they are registered
 under the generic T2V workload option in the registry.
 
+**Note (MMAudio)**: the registered Hugging Face model ID is reserved but not
+yet public. Follow the [MMAudio inference guide](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/pipelines/basic/mmaudio/README.md)
+to convert the official weights locally and set `MMAUDIO_MODEL_PATH`.
+
 **Note (MiniMax H3)**: T2VA, FL2VA, and Ref2VA all generate video with stereo
 audio. Use the Ref2VA example when passing ordered image, video, or audio
 references.
@@ -173,6 +178,19 @@ optimizations: absence means **untested**, not incompatible.
 | Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |
 
+## Apple Silicon native runtime
+
+| Release path | Model | Mode | Validated hardware | Status |
+| --- | --- | --- | --- | --- |
+| MLX FastMetal T2V 1.3B | [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released |
+| MLX FastMetal TI2V 5B | [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released |
+| MLX FastMetal T2V 14B | [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB+ unified memory | Released |
+
+Apple Silicon uses FastMetal-QAD. CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`,
+`FastVideo/FastWan-QAD-FP8-1.3B`) is the NVIDIA release. See the
+[Apple Silicon guide](../getting_started/installation/mps.md) and the
+[FastMetal-QAD blog](https://haoailab.com/blogs/fastmetal/).
+
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 
 ***Lucy Edit Dev uses a non-commercial model license. FastVideo support is
@@ -200,9 +218,10 @@ Per the installation guides:
   [GPU install guide](../getting_started/installation/gpu.md).
 - **NVIDIA DGX Spark (GB10, aarch64)** — CUDA 13, from-source kernel build; see
   the [DGX Spark install guide](../getting_started/installation/spark.md).
-- **Apple silicon (MPS)** — macOS 14 or newer; see the
-  [MPS install guide](../getting_started/installation/mps.md) and
-  [`basic_mps.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mps.py).
+- **Apple silicon** — macOS 14 or newer; FastMetal-QAD via the MLX runtime. See the
+  [Apple Silicon guide](../getting_started/installation/mps.md). The older
+  [`basic_mps.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mps.py)
+  demo is PyTorch MPS only.
 
 Optimization-specific hardware constraints (e.g. STA requiring Hopper) are
 listed under [Special requirements](#special-requirements).

examples/inference/basic/README.md

@@ -18,11 +18,27 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 python examples/inference/basic/basic.py
 ```
 
-For an example on Apple silicon: 
-```
-python examples/inference/basic/basic_mps.py
+### Apple Silicon (FastMetal-QAD)
+
+Use the MLX runtime with FastMetal-QAD. See the
+[Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
+
+```bash
+hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD
+
+python examples/inference/basic/mlx_wan_prompt_to_video.py \
+  --model-root ./FastMetal-1.3B-QAD \
+  --mlx-checkpoint ./FastMetal-1.3B-QAD \
+  --prompt "A bird's-eye view of a misty forest valley at dawn."
 ```
 
+5B uses
+[`mlx_wan22_generate.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/mlx_wan22_generate.py)
+with
+`FastVideo/FastMetal-5B-QAD`.
+
+`examples/inference/basic/basic_mps.py` is the older PyTorch MPS demo.
+
 For an example running DMD+VSA inference:
 ```
 python examples/inference/basic/basic_dmd.py
@@ -33,6 +49,70 @@ For the typed config/request path added during the inference API refactor:
 python examples/inference/basic/basic_dmd_new_api.py
 ```
 
+### FastH3 Preview
+
+The verified [basic FastH3 example](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3.py)
+runs the few-step (4-forward, DMD2-distilled) MiniMax-H3 preview, generating
+synchronized video and audio with its trained block-sparse VSA attention:
+
+```bash
+UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"
+```
+
+This installs the pinned FA4 CuTe package and FastVideo kernel release used by
+the measured GB200 profile. Then run:
+
+```
+python examples/inference/basic/basic_fasth3.py --prompt "your prompt"
+```
+The default checkpoint, [FastH3 Preview v0.2](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2), is public on the Hub under the MiniMax H3 Community License. Review its model card and license before use or redistribution.
+
+The default `all` profile is the fastest measured four-GPU Preview recipe on GB200. It selects VSA sparsity 0.9 with 64-token tiles and the sm_100a sparse kernel, enables FA4 for eligible non-VSA paths, regionally compiles and replicates the sparse DiT, compiles and temporally parallelizes the video VAE with the `gather` strategy, and pins CPU-offloaded component memory. It also pins the benchmark protocol: five sigma-grid points (exactly four DiT forwards), one excluded seed-999 warmup, then three timed seed-1000 requests with distinct output paths.
+
+The equivalent explicit command is:
+
+```bash
+python examples/inference/basic/basic_fasth3.py \
+  --prompt "your prompt" \
+  --profile all \
+  --num-gpus 4 \
+  --steps 5 \
+  --vsa-sparsity 0.9 \
+  --vsa-tile-size 64 \
+  --vsa-kernel sm100a \
+  --compile-vae \
+  --parallel-vae \
+  --replicated-dit \
+  --pin-cpu-memory \
+  --fa4 \
+  --no-torch-compile \
+  --inference-torch-compile \
+  --ulysses-a2a off \
+  --warmup \
+  --repeats 3 \
+  --seed 1000 \
+  --warmup-seed 999
+```
+
+`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median.
+
+One script covers each validated duration; regional compile is the fastest
+measured DiT route for all three:
+
+```bash
+# 5 s
+python examples/inference/basic/basic_fasth3.py \
+  --prompt "your prompt" --output outputs/fasth3_5s
+# 10 s
+python examples/inference/basic/basic_fasth3.py \
+  --prompt "your prompt" --num-frames 243 --output outputs/fasth3_10s
+# 15 s
+python examples/inference/basic/basic_fasth3.py \
+  --prompt "your prompt" --num-frames 345 --output outputs/fasth3_15s
+```
+
+Pass `--no-inference-torch-compile` to recover the eager sparse-DiT route.
+
 ## Basic Walkthrough
 
 All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!

examples/inference/basic/basic_fasth3.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Few-step video+audio generation with the DMD2-distilled MiniMax H3 preview.
+
+The default ``all`` profile reproduces the fastest measured FastH3 Preview
+recipe on four GB200 GPUs. It runs the checkpoint's native five-point sigma
+grid (exactly four DiT forwards), trained VSA policy, Blackwell sparse kernel,
+regional fullgraph DiT compile, compiled/parallel video VAE, and inference-only
+H3 fusions. One compile warmup is excluded before three measured requests.
+
+Both regional compile and the default fusions can change floating-point
+operation order, so ``all`` is a report-only performance profile.
+``--profile strict`` disables the H3 fusions but preserves regional compile;
+combine it with ``--no-inference-torch-compile`` for the eager strict route.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import os
+import statistics
+import time
+from collections.abc import Sequence
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    CompileConfig,
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+DEFAULT_MODEL = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-path", default=DEFAULT_MODEL)
+    # The HF repo may require authentication while the MiniMax H3 Community
+    # License review completes. A local snapshot can be passed here instead.
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--output", default="outputs/fasth3")
+    parser.add_argument("--profile",
+                        choices=("all", "strict"),
+                        default="all",
+                        help="all enables the fastest measured, non-parity H3 fusions; strict disables only them")
+    parser.add_argument("--height", type=int, default=768)
+    parser.add_argument("--width", type=int, default=1344)
+    parser.add_argument("--num-frames", type=int, default=124)
+    # num_inference_steps counts sigma-GRID POINTS. The distilled schedule is
+    # t=1000,750,500,250 -> 0: five points and exactly four DiT forwards.
+    parser.add_argument("--steps",
+                        type=int,
+                        default=5,
+                        help="sigma-grid points; N points run N-1 DiT forwards (the trained default is 5)")
+    parser.add_argument("--seed", type=int, default=1000, help="seed reused for every measured request")
+    parser.add_argument("--warmup-seed", type=int, default=999)
+    parser.add_argument("--repeats", type=int, default=3, help="number of measured requests after warmup")
+    parser.add_argument("--warmup",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="run one excluded request before timing")
+    parser.add_argument("--num-gpus", type=int, default=4)
+    parser.add_argument("--vsa-sparsity",
+                        type=float,
+                        default=0.9,
+                        help="run-level VSA sparsity in [0, 1); 0.9 is the checkpoint's trained policy")
+    parser.add_argument("--vsa-tile-size",
+                        type=int,
+                        choices=(64, 256),
+                        default=64,
+                        help="VSA-H3 tile size; 64 is the checkpoint's trained and measured geometry")
+    parser.add_argument("--vsa-kernel",
+                        choices=("triton", "sm100a"),
+                        default="sm100a",
+                        help="tile-64 sparse kernel; sm100a is the measured GB200 route and requires a compatible "
+                        "fastvideo-kernel build")
+    parser.add_argument("--fa4",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="use FA4 for eligible non-VSA attention paths")
+    parser.add_argument("--h3-fusions",
+                        action=argparse.BooleanOptionalAction,
+                        default=None,
+                        help="override the profile's H3 fusion policy (changes model numerics when enabled)")
+    parser.add_argument("--compile-vae",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="compile the video VAE decoder independently of the DiT")
+    parser.add_argument("--parallel-vae",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="round-robin VAE temporal chunks across sequence-parallel ranks")
+    parser.add_argument("--replicated-dit",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="replicate DiT weights instead of FSDP-sharding them")
+    parser.add_argument("--pin-cpu-memory",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="pin CPU-offloaded text-encoder and VAE weights")
+    parser.add_argument("--torch-compile",
+                        action=argparse.BooleanOptionalAction,
+                        default=False,
+                        help="compile the whole DiT path (off in the fastest FastH3 profile)")
+    parser.add_argument("--inference-torch-compile",
+                        action=argparse.BooleanOptionalAction,
+                        default=True,
+                        help="regionally compile DiT blocks (enabled in the fastest FastH3 profile)")
+    parser.add_argument("--ulysses-a2a",
+                        choices=("off", "auto"),
+                        default="off",
+                        help="sequence-parallel all-to-all route; off reproduces the fastest FastH3 profile, while "
+                        "auto opts into the fused NVLink kernel when the installed kernel package supports it")
+    parser.add_argument("--compile-mode",
+                        default=None,
+                        help='whole-DiT torch.compile mode, e.g. "reduce-overhead"; requires '
+                        "--no-inference-torch-compile")
+    args = parser.parse_args(argv)
+    if args.repeats < 1:
+        parser.error("--repeats must be at least 1")
+    if args.num_gpus < 1:
+        parser.error("--num-gpus must be at least 1")
+    if not 0.0 <= args.vsa_sparsity < 1.0:
+        parser.error("--vsa-sparsity must be in [0, 1)")
+    if args.compile_mode is not None and args.inference_torch_compile:
+        parser.error("--compile-mode cannot be combined with regional compile; pass --no-inference-torch-compile")
+    return args
+
+
+def _h3_fusions_enabled(args: argparse.Namespace) -> bool:
+    if args.h3_fusions is not None:
+        return bool(args.h3_fusions)
+    return args.profile == "all"
+
+
+def profile_environment(args: argparse.Namespace) -> dict[str, str | None]:
+    """Return the complete boot-time environment for this profile.
+
+    ``None`` means the variable must be removed. Values are explicit even for
+    disabled features so a shell's inherited experiment settings cannot
+    silently change the advertised profile.
+    """
+    return {
+        "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3",
+        "FASTVIDEO_VSA_SM100A": "1" if args.vsa_kernel == "sm100a" else "0",
+        "FASTVIDEO_VSA_CUTEDSL": "0",
+        # A non-empty output path enables the diagnostic probe.
+        "FASTVIDEO_H3_VSA_PROBE": None,
+        "FASTVIDEO_DISABLE_ATTENTION_COMPILE": "0",
+        "FASTVIDEO_FA4": "1" if args.fa4 else "0",
+        "FASTVIDEO_NVFP4_FA4": "0",
+        "FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN": "0",
+        "FASTVIDEO_MINIMAX_H3_FUSIONS": "all" if _h3_fusions_enabled(args) else "0",
+        "FASTVIDEO_INFERENCE_TORCH_COMPILE": "1" if args.inference_torch_compile else "0",
+        "FASTVIDEO_VAE_PARALLEL_DECODE": "1" if args.parallel_vae else "0",
+        "FASTVIDEO_VAE_PARALLEL_ENCODE": "0",
+        "FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY": "gather",
+        "FASTVIDEO_ULYSSES_A2A": args.ulysses_a2a,
+        "FASTVIDEO_STAGE_LOGGING": "1",
+    }
+
+
+def configure_environment(args: argparse.Namespace) -> dict[str, str | None]:
+    environment = profile_environment(args)
+    for name, value in environment.items():
+        if value is None:
+            os.environ.pop(name, None)
+        else:
+            os.environ[name] = value
+    return environment
+
+
+def _fa4_is_installed() -> bool:
+    try:
+        return importlib.util.find_spec("flash_attn.cute") is not None
+    except (ImportError, ModuleNotFoundError):
+        return False
+
+
+def _sm100a_kernel_is_installed() -> bool:
+    try:
+        from fastvideo_kernel import block_sparse_attn_sm100a
+    except ImportError:
+        return False
+    return bool(getattr(block_sparse_attn_sm100a, "_HAS_VSA_SM100A", False))
+
+
+def validate_profile_dependencies(args: argparse.Namespace) -> None:
+    """Fail before model loading when the selected measured route is absent."""
+    if args.fa4 and not _fa4_is_installed():
+        raise RuntimeError(
+            "FastH3's FA4 profile requires the pinned flash-attn-4 package. Install it with "
+            "`UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"`, or pass --no-fa4.")
+    if args.vsa_kernel == "sm100a" and not _sm100a_kernel_is_installed():
+        raise RuntimeError(
+            "FastH3's sm100a profile requires fastvideo-kernel 0.3.4 built with the Blackwell VSA extension. "
+            "Install this checkout with `UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"` (or run "
+            "`cd fastvideo-kernel && ./build.sh`), or pass --vsa-kernel triton.")
+
+
+def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
+    experimental: dict[str, object] = {
+        "attention_backend": "VIDEO_SPARSE_ATTN_H3",
+        "VSA_sparsity": args.vsa_sparsity,
+        "VSA_tile_size": args.vsa_tile_size,
+        "inference_torch_compile": args.inference_torch_compile,
+        "vae_parallel_decode": args.parallel_vae,
+        "vae_parallel_decode_strategy": "gather",
+    }
+    return GeneratorConfig(
+        model_path=args.model_path,
+        pipeline=PipelineSelection(experimental=experimental),
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            use_fsdp_inference=args.num_gpus > 1 and not args.replicated_dit,
+            parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
+            offload=OffloadConfig(
+                dit=False,
+                dit_layerwise=False,
+                text_encoder=True,
+                vae=True,
+                pin_cpu_memory=args.pin_cpu_memory,
+            ),
+            compile=CompileConfig(
+                enabled=args.torch_compile,
+                mode=args.compile_mode,
+                vae_enabled=args.compile_vae,
+            ),
+        ),
+    )
+
+
+def build_request(args: argparse.Namespace, output_path: Path, seed: int) -> GenerationRequest:
+    return GenerationRequest(
+        prompt=args.prompt,
+        negative_prompt="",
+        sampling=SamplingConfig(
+            height=args.height,
+            width=args.width,
+            num_frames=args.num_frames,
+            fps=24,
+            num_inference_steps=args.steps,
+            # MiniMax-H3 is guidance-distilled; FastH3 inherits that contract.
+            guidance_scale=1.0,
+            batch_cfg=False,
+            seed=seed,
+        ),
+        output=OutputConfig(
+            output_path=str(output_path),
+            save_video=True,
+            return_frames=False,
+        ),
+    )
+
+
+def _actual_output_path(result: object, requested: Path) -> Path:
+    video_path = getattr(result, "video_path", None)
+    return Path(video_path) if video_path else requested
+
+
+def _denoise_seconds(result: object) -> float | None:
+    stages = getattr(getattr(result, "logging_info", None), "stages", None)
+    if not stages:
+        return None
+    for stage_name, metrics in stages.items():
+        if "denois" not in stage_name.lower():
+            continue
+        execution_time = metrics.get("execution_time")
+        return float(execution_time) if execution_time is not None else None
+    return None
+
+
+def run(args: argparse.Namespace) -> list[float]:
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    environment = configure_environment(args)
+    validate_profile_dependencies(args)
+
+    print(f"Profile: {args.profile} ({'non-parity fusions' if _h3_fusions_enabled(args) else 'fusions off'})")
+    print(f"Output directory: {output_dir.resolve()}")
+    print("Denoising contract: 5 sigma points = 4 DiT forwards" if args.steps == 5 else
+          f"Denoising contract override: {args.steps} sigma points = {args.steps - 1} DiT forwards")
+    print("Profile environment: " + " ".join(f"{key}={value if value is not None else '<unset>'}"
+                                                  for key, value in environment.items()))
+
+    generator = VideoGenerator.from_config(build_generator_config(args))
+    measured_wall_times: list[float] = []
+    measured_denoise_times: list[float] = []
+    try:
+        if args.warmup:
+            warmup_path = output_dir / "_fasth3_warmup.mp4"
+            print(f"[warmup] generating (excluded from timing summary): {warmup_path}")
+            started = time.perf_counter()
+            warmup_result = generator.generate(build_request(args, warmup_path, args.warmup_seed))
+            warmup_wall = time.perf_counter() - started
+            actual_warmup_path = _actual_output_path(warmup_result, warmup_path)
+            print(f"[warmup] wall={warmup_wall:.3f}s (excluded)")
+            print(f"Warmup output written to: {actual_warmup_path}")
+
+        for index in range(1, args.repeats + 1):
+            requested_path = output_dir / f"fasth3_{args.profile}_run_{index:02d}.mp4"
+            print(f"[measured {index}/{args.repeats}] generating: {requested_path}")
+            started = time.perf_counter()
+            result = generator.generate(build_request(args, requested_path, args.seed))
+            wall = time.perf_counter() - started
+            measured_wall_times.append(wall)
+            actual_path = _actual_output_path(result, requested_path)
+            print(f"Output written to: {actual_path}")
+            print(f"E2E wall time: {wall:.3f}s")
+            generation_time = getattr(result, "generation_time", None)
+            if generation_time is not None:
+                print(f"Generation time: {float(generation_time):.3f}s")
+            denoise_time = _denoise_seconds(result)
+            if denoise_time is not None:
+                measured_denoise_times.append(denoise_time)
+                print(f"Denoising time: {denoise_time:.3f}s")
+
+        median = statistics.median(measured_wall_times)
+        print(f"Measured E2E wall times (n={len(measured_wall_times)}, warmup excluded): "
+              f"{[round(value, 3) for value in measured_wall_times]}")
+        print(f"Median E2E wall time: {median:.3f}s")
+        if measured_denoise_times:
+            print(f"Median denoising time: {statistics.median(measured_denoise_times):.3f}s")
+        return measured_wall_times
+    finally:
+        generator.shutdown()
+
+
+def main() -> None:
+    run(parse_args())
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_minimax_h3_fl2va.py

@@ -27,7 +27,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
     # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
     #   --model-path noctuashap/MiniMax-H3-pruned-r16
-    # (or a local dir converted with tools/minimax_h3/fit_adaln_basis.py).
+    # (or a local dir produced by
+    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
     # adaln_rank is read from the checkpoint config; no other flags needed.
     # Rank-reduced checkpoints are inference-only: training needs the
     # full-rank release.

examples/inference/basic/basic_minimax_h3_ref2va.py

@@ -27,7 +27,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
     # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
     #   --model-path noctuashap/MiniMax-H3-pruned-r16
-    # (or a local dir converted with tools/minimax_h3/fit_adaln_basis.py).
+    # (or a local dir produced by
+    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
     # adaln_rank is read from the checkpoint config; no other flags needed.
     # Rank-reduced checkpoints are inference-only: training needs the
     # full-rank release.

examples/inference/basic/basic_minimax_h3_t2v.py

@@ -15,6 +15,7 @@ from fastvideo.api import (
     OffloadConfig,
     OutputConfig,
     ParallelismConfig,
+    PipelineSelection,
     SamplingConfig,
 )
 
@@ -24,7 +25,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
     # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
     #   --model-path noctuashap/MiniMax-H3-pruned-r16
-    # (or a local dir converted with tools/minimax_h3/fit_adaln_basis.py).
+    # (or a local dir produced by
+    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
     # adaln_rank is read from the checkpoint config; no other flags needed.
     # Rank-reduced checkpoints are inference-only: training needs the
     # full-rank release.
@@ -40,6 +42,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--compile-mode",
                         default=None,
                         help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
+    parser.add_argument("--inference-torch-compile",
+                        action="store_true",
+                        help="regional fullgraph torch.compile of each DiT block after load (the #1718 "
+                        "training-port semantics: no kwargs; fullgraph + emulate_precision_casts injected). "
+                        "First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time "
+                        "the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent")
     parser.add_argument("--repeats",
                         type=int,
                         default=1,
@@ -53,9 +61,16 @@ def main() -> None:
     output_dir = Path(args.output)
     output_dir.mkdir(parents=True, exist_ok=True)
 
+    # Boot-time run configuration folded into FastVideoArgs (the same
+    # experimental-dict route basic_fasth3.py uses for the VSA knobs).
+    experimental: dict[str, object] = {}
+    if args.inference_torch_compile:
+        experimental["inference_torch_compile"] = True
+
     generator = VideoGenerator.from_config(
         GeneratorConfig(
             model_path=args.model_path,
+            pipeline=PipelineSelection(experimental=experimental),
             engine=EngineConfig(
                 num_gpus=args.num_gpus,
                 use_fsdp_inference=args.num_gpus > 1,

examples/inference/basic/basic_mmaudio.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio large-44k-v2 video-to-audio example."""
+
+import argparse
+import os
+
+from fastvideo import VideoGenerator
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--video-path", required=True)
+    parser.add_argument("--output-path", default="outputs_audio/mmaudio.wav")
+    parser.add_argument("--duration-seconds", type=float, default=8.0)
+    parser.add_argument("--prompt", default="")
+    parser.add_argument("--negative-prompt", default="music")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    generator = VideoGenerator.from_pretrained(
+        os.environ.get(
+            "MMAUDIO_MODEL_PATH",
+            "converted_weights/mmaudio/large_44k_v2",
+        ),
+        workload_type="v2a",
+        num_gpus=1,
+    )
+    result = generator.generate_video(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        video_path=args.video_path,
+        audio_end_in_s=args.duration_seconds,
+        output_path=args.output_path,
+        save_video=True,
+        return_frames=False,
+    )
+    print(result["video_path"])
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/mlx_rife_smoke_test.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tiny MLX RIFE frame-interpolation smoke test."""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import numpy as np
+
+from fastvideo.mlx_runtime.rife_interp import interpolate, load_model
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="MLX RIFE 4.25 frame interpolation smoke test."
+    )
+    parser.add_argument(
+        "--self-test",
+        action="store_true",
+        help="Run a tiny two-frame interpolation test.",
+    )
+    args = parser.parse_args()
+    if not args.self_test:
+        raise SystemExit("Nothing to do; pass --self-test")
+
+    frame0 = np.zeros((64, 96, 3), dtype=np.uint8)
+    frame1 = np.zeros((64, 96, 3), dtype=np.uint8)
+    frame1[:, :, 0] = 255
+    start = time.perf_counter()
+    model = load_model()
+    load_s = time.perf_counter() - start
+    start = time.perf_counter()
+    frames = interpolate([frame0, frame1], factor=2, model=model)
+    interp_s = time.perf_counter() - start
+    assert len(frames) == 3
+    assert frames[1].shape == frame0.shape
+    assert frames[1].dtype == np.uint8
+    print(
+        "MLX RIFE self-test passed: "
+        f"load_s={load_s:.3f} interp_s={interp_s:.3f} shape={frames[1].shape}"
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/mlx_wan22_generate.py

@@ -0,0 +1,533 @@
+# SPDX-License-Identifier: Apache-2.0
+"""End-to-end FastMetal-5B-QAD generation on Apple Silicon (MLX DiT + MLX TAEHV).
+
+This is the Wan2.2 TI2V entrypoint. Use FastVideo/FastMetal-5B-QAD:
+
+    hf download FastVideo/FastMetal-5B-QAD --local-dir ./FastMetal-5B-QAD
+    python examples/inference/basic/mlx_wan22_generate.py \\
+      --mlx-checkpoint ./FastMetal-5B-QAD \\
+      --text-encoder-root ./FastMetal-5B-QAD \\
+      --vae-root ./FastMetal-5B-QAD/vae
+
+Pipeline: torch/MPS UMT5 encode (shared with 1.3B) → MLXWan22DiT 3-step DMD
+(warped schedule, flow_shift=5) → MLX TAEHV decode (taew2_2.pth). Fully MLX
+on the heavy DiT + decode path.
+
+Decoder backends: ``taehv`` (default, MLX, ~seconds), ``taehv-torch`` (parity),
+``wan-vae`` (full AutoencoderKLWan on MPS, slow).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+
+from fastvideo.mlx_runtime.fast_spatial import DEFAULT_FAST_SPATIAL_SHARPEN
+from fastvideo.mlx_runtime.frame_upsample import DEFAULT_PIXEL_UPSAMPLE_MODE, PIXEL_UPSAMPLE_MODES
+from fastvideo.mlx_runtime.memory import cleanup_mlx
+from fastvideo.mlx_runtime.prompt_cache import (
+    fingerprint_digest,
+    load_prompt_cache,
+    save_prompt_cache,
+    text_encoder_fingerprint,
+)
+from fastvideo.mlx_runtime.checkpoint_compat import (
+    UnsupportedMLXCheckpointError,
+    raise_if_unsupported_mlx_checkpoint,
+    resolve_mlx_checkpoint,
+)
+from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
+
+FASTWAN21_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
+FASTWAN22_MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+DEFAULT_HEIGHT = 448
+DEFAULT_WIDTH = 832
+DEFAULT_NUM_FRAMES = 121
+
+def _resolve_model_paths(
+    *,
+    text_encoder_root: Path | None,
+    dit_checkpoint: Path | None,
+    dit_config: Path | None,
+    vae_root: Path | None,
+    mlx_checkpoint: Path | None,
+    decode_backend: str,
+) -> tuple[Path, Path | None, Path | None, Path | None]:
+    """Download only the missing assets required by the selected Wan2.2 path."""
+    from huggingface_hub import snapshot_download
+
+    if text_encoder_root is None:
+        text_encoder_root = Path(snapshot_download(
+            FASTWAN21_MODEL_ID,
+            allow_patterns=["tokenizer/*", "text_encoder/*"],
+        ))
+    if mlx_checkpoint is None and (dit_checkpoint is None or dit_config is None):
+        patterns = []
+        if dit_checkpoint is None:
+            patterns.append("transformer/diffusion_pytorch_model.safetensors")
+        if dit_config is None:
+            patterns.append("transformer/config.json")
+        model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=patterns))
+        dit_checkpoint = dit_checkpoint or model_root / "transformer/diffusion_pytorch_model.safetensors"
+        dit_config = dit_config or model_root / "transformer/config.json"
+    if decode_backend == "wan-vae" and vae_root is None:
+        model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=["vae/*"]))
+        vae_root = model_root / "vae"
+    return text_encoder_root, dit_checkpoint, dit_config, vae_root
+
+
+def _prompt_cache_fingerprint(
+    *,
+    prompt: str,
+    prompt_used: str,
+    enhance_prompt: bool,
+    enhance_prompt_backend: str,
+    text_encoder_root: Path,
+    max_sequence_length: int,
+    dtype: str,
+) -> dict[str, object]:
+    return {
+        "prompt": prompt,
+        "prompt_used": prompt_used,
+        "enhance_prompt": enhance_prompt,
+        "enhance_prompt_backend": enhance_prompt_backend,
+        "text_encoder": text_encoder_fingerprint(text_encoder_root),
+        "max_sequence_length": max_sequence_length,
+        "dtype": dtype,
+    }
+
+
+def _default_prompt_cache_path(fingerprint: dict[str, object]) -> Path:
+    """Content-addressed default cache file for a prompt fingerprint.
+
+    The Wan2.1 entrypoint caches prompt embeddings by default; this one only
+    did so when handed an explicit ``--prompt-embeds-cache`` path, so every 5B
+    run paid a full UMT5 encode (~45s on an M4 Max) even for a repeat prompt.
+    The fingerprint already covers everything that changes the embedding, so
+    hash it for the filename.
+    """
+    digest = fingerprint_digest(fingerprint)[:32]
+    return Path.home() / ".cache" / "fastvideo" / "prompt_embeds" / f"wan22_{digest}.npy"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="MLX Wan2.2-5B T2V (encode → DiT DMD → TAEHV/VAE decode)"
+    )
+    parser.add_argument(
+        "--prompt",
+        default="A red fox trotting through a snowy pine forest at golden hour, cinematic",
+    )
+    parser.add_argument(
+        "--output-path",
+        type=Path,
+        default=Path("video_samples/demo_5b/fox_5b_mlx.mp4"),
+    )
+    parser.add_argument(
+        "--text-encoder-root",
+        type=Path,
+        default=None,
+        help="Root with text_encoder/ + tokenizer/",
+    )
+    parser.add_argument(
+        "--prompt-embeds-cache",
+        type=Path,
+        default=None,
+        help="Explicit .npy UMT5 embedding cache file. Overrides the automatic "
+        "content-addressed cache (--prompt-cache).",
+    )
+    parser.add_argument(
+        "--prompt-cache",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Cache prompt embeddings under ~/.cache/fastvideo/prompt_embeds so "
+        "repeat runs skip the text encoder entirely. Default: on.",
+    )
+    parser.add_argument(
+        "--text-encoder-device",
+        choices=("auto", "cpu", "mps"),
+        default="cpu",
+        help="Device for UMT5 encoding. CPU is safest beside the 5B MLX DiT.",
+    )
+    parser.add_argument(
+        "--enhance-prompt",
+        action="store_true",
+        help="Apply deterministic local cinematic prompt enrichment before UMT5.",
+    )
+    parser.add_argument(
+        "--enhance-prompt-backend",
+        choices=("template",),
+        default="template",
+        help="Prompt enrichment backend.",
+    )
+    parser.add_argument(
+        "--dit-checkpoint",
+        type=Path,
+        default=None,
+    )
+    parser.add_argument("--dit-config", type=Path, default=None)
+    parser.add_argument(
+        "--mlx-checkpoint",
+        type=Path,
+        default=None,
+        help="Packed FastMetal-5B-QAD MLX DiT directory (mlx_dit.json + mlx_dit.safetensors). "
+        "If omitted, a FastMetal directory passed as --text-encoder-root is used when it "
+        "already contains those files.",
+    )
+    parser.add_argument("--vae-root", type=Path, default=None)
+    parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
+    parser.add_argument("--width", type=int, default=DEFAULT_WIDTH)
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=DEFAULT_NUM_FRAMES,
+        help="Pixel frames (121 at 24fps = 5.04 seconds)",
+    )
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--renoise-seed", type=int, default=0)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--flow-shift", type=float, default=5.0)
+    parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
+    parser.add_argument(
+        "--no-warp",
+        action="store_true",
+        help="Disable schedule warping (debug only).",
+    )
+    parser.add_argument(
+        "--fast",
+        action="store_true",
+        help="Generate fewer frames then RIFE-interpolate to --num-frames.",
+    )
+    parser.add_argument("--fast-factor", type=int, default=2)
+    parser.add_argument("--fast-sharpen", type=float, default=0.6)
+    parser.add_argument(
+        "--fast-spatial",
+        action="store_true",
+        help="Denoise and decode at reduced spatial resolution, then resample "
+        "the decoded frames up to the target size.",
+    )
+    parser.add_argument("--fast-spatial-scale", type=int, default=2)
+    parser.add_argument(
+        "--fast-spatial-upsample-mode",
+        choices=PIXEL_UPSAMPLE_MODES,
+        default=DEFAULT_PIXEL_UPSAMPLE_MODE,
+    )
+    parser.add_argument("--fast-spatial-sharpen", type=float, default=DEFAULT_FAST_SPATIAL_SHARPEN)
+    parser.add_argument(
+        "--refine",
+        action="store_true",
+        help="Two-pass DMD: coarse denoise, upsample/re-noise, full-res denoise.",
+    )
+    parser.add_argument("--refine-scale", type=int, default=2)
+    parser.add_argument(
+        "--refine-upsample-mode",
+        choices=("bilinear", "nearest"),
+        default="bilinear",
+    )
+    parser.add_argument("--no-refine-add-noise", action="store_true")
+    parser.add_argument(
+        "--decode-backend",
+        choices=("taehv", "taehv-torch", "wan-vae"),
+        default="taehv",
+    )
+    parser.add_argument("--save-latents", type=Path, default=None)
+    parser.add_argument("--metrics-json", type=Path, default=None,
+                        help="Write measured run metadata as JSON for reports or galleries.")
+    parser.add_argument(
+        "--compile",
+        action="store_true",
+        help="Compile the DiT forward with mx.compile; fallback to eager on failure.",
+    )
+    args = parser.parse_args()
+
+    if args.fast_factor < 2:
+        parser.error("--fast-factor must be at least 2")
+    # --fast-spatial used to be rejected here because it upsampled the completed
+    # 48-channel latent, which is out of distribution for the decoder and gave
+    # black or noisy video. The upsample now runs on decoded frames, so the
+    # latent never leaves the grid it was denoised on and the mode is usable.
+    if args.refine and args.fast_spatial:
+        print("[wan22] --refine takes precedence over --fast-spatial")
+
+    args.mlx_checkpoint = resolve_mlx_checkpoint(args.mlx_checkpoint, args.text_encoder_root)
+    if args.mlx_checkpoint is not None:
+        if args.text_encoder_root is None and (args.mlx_checkpoint / "text_encoder").is_dir():
+            args.text_encoder_root = args.mlx_checkpoint
+        if args.vae_root is None and (args.mlx_checkpoint / "vae").is_dir():
+            args.vae_root = args.mlx_checkpoint / "vae"
+    try:
+        raise_if_unsupported_mlx_checkpoint(args.mlx_checkpoint, args.dit_checkpoint)
+    except UnsupportedMLXCheckpointError as exc:
+        raise SystemExit(str(exc)) from exc
+
+    args.text_encoder_root, args.dit_checkpoint, args.dit_config, args.vae_root = _resolve_model_paths(
+        text_encoder_root=args.text_encoder_root,
+        dit_checkpoint=args.dit_checkpoint,
+        dit_config=args.dit_config,
+        vae_root=args.vae_root,
+        mlx_checkpoint=args.mlx_checkpoint,
+        decode_backend=args.decode_backend,
+    )
+    target_frames = args.num_frames
+    if args.fast:
+        args.num_frames = aligned_keyframe_count(target_frames, args.fast_factor)
+        print(
+            f"[wan22 fast] generating {args.num_frames} frames, "
+            f"RIFE {args.fast_factor}x -> {target_frames}"
+        )
+
+    import mlx.core as mx
+    import torch
+
+    from examples.inference.basic.mlx_wan_prompt_to_video import (
+        _postprocess_video,
+        encode_prompt,
+        make_rotary_embeddings,
+    )
+    from fastvideo.mlx_runtime.fast_spatial import plan_fast_spatial
+    from fastvideo.mlx_runtime.refine import (
+        default_refine_timesteps,
+        plan_refine_resolutions,
+        prepare_refine_latents,
+    )
+    from fastvideo.mlx_runtime.wan22 import (
+        mlx_wan22_dit_from_diffusers_safetensors,
+        mlx_wan22_dit_from_mlx_checkpoint,
+    )
+    from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule, sample_wan22_dmd
+    from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video
+
+    if args.mlx_checkpoint is not None:
+        config = json.loads((args.mlx_checkpoint / "mlx_dit.json").read_text())["config"]
+    else:
+        config = json.loads(args.dit_config.read_text())
+    patch_size = tuple(config.get("patch_size", (1, 2, 2)))
+    if args.refine:
+        active_plan = plan_refine_resolutions(
+            height=args.height, width=args.width, num_frames=args.num_frames,
+            spatial_scale=args.refine_scale, vae_spatial_compression=16,
+            vae_temporal_compression=4, patch_size=patch_size, enabled=True,
+        )
+        spatial_mode = "refine"
+    elif args.fast_spatial:
+        fast_spatial_plan = plan_fast_spatial(
+            height=args.height, width=args.width, num_frames=args.num_frames,
+            spatial_scale=args.fast_spatial_scale, vae_spatial_compression=16,
+            vae_temporal_compression=4, patch_size=patch_size,
+            upsample_mode=args.fast_spatial_upsample_mode,
+            sharpen=args.fast_spatial_sharpen, enabled=True,
+        )
+        active_plan = fast_spatial_plan.plan
+        spatial_mode = "fast_spatial"
+    else:
+        active_plan = plan_refine_resolutions(
+            height=args.height, width=args.width, num_frames=args.num_frames,
+            spatial_scale=1, vae_spatial_compression=16, vae_temporal_compression=4,
+            patch_size=patch_size, enabled=False,
+        )
+        spatial_mode = "off"
+    lat_h, lat_w = active_plan.stage1_latent_height, active_plan.stage1_latent_width
+    lat_t = active_plan.latent_frames
+    in_ch = int(config["in_channels"])
+    print(f"[5B] latent {in_ch}x{lat_t}x{lat_h}x{lat_w}", flush=True)
+
+    total_start = time.perf_counter()
+    prompt_for_encode = args.prompt
+    enhance_backend = None
+    enhance_elapsed_s = 0.0
+    if args.enhance_prompt:
+        from fastvideo.mlx_runtime.prompt_enhance import enhance_prompt
+
+        enhancement = enhance_prompt(args.prompt, backend=args.enhance_prompt_backend)
+        prompt_for_encode = enhancement.enhanced
+        enhance_backend = enhancement.backend
+        enhance_elapsed_s = enhancement.elapsed_s
+        print(f"[enhance] backend={enhance_backend} in {enhance_elapsed_s:.2f}s", flush=True)
+        print(f"[enhance] prompt: {prompt_for_encode}", flush=True)
+
+    t0 = time.perf_counter()
+    prompt_cache_fingerprint = _prompt_cache_fingerprint(
+        prompt=args.prompt,
+        prompt_used=prompt_for_encode,
+        enhance_prompt=args.enhance_prompt,
+        enhance_prompt_backend=args.enhance_prompt_backend,
+        text_encoder_root=args.text_encoder_root,
+        max_sequence_length=512,
+        dtype="fp16",
+    )
+    prompt_cache_path = args.prompt_embeds_cache
+    if prompt_cache_path is None and args.prompt_cache:
+        prompt_cache_path = _default_prompt_cache_path(prompt_cache_fingerprint)
+    cached_embeds = load_prompt_cache(
+        prompt_cache_path,
+        prompt_cache_fingerprint,
+    )
+    if cached_embeds is not None:
+        embeds = torch.from_numpy(cached_embeds).contiguous()
+    else:
+        embeds = encode_prompt(
+            model_root=args.text_encoder_root,
+            prompt=prompt_for_encode,
+            max_sequence_length=512,
+            device_arg=args.text_encoder_device,
+            dtype_arg="fp16",
+        )
+        save_prompt_cache(
+            prompt_cache_path,
+            embeds.cpu().numpy(),
+            prompt_cache_fingerprint,
+        )
+    ehs = mx.array(embeds.numpy()).astype(mx.float16)
+    prompt_encode_s = time.perf_counter() - t0
+    print(f"[5B] prompt encoded {tuple(ehs.shape)} in {prompt_encode_s:.1f}s", flush=True)
+
+    t1 = time.perf_counter()
+    if args.mlx_checkpoint is not None:
+        dit = mlx_wan22_dit_from_mlx_checkpoint(
+            args.mlx_checkpoint,
+            compile=args.compile,
+        )
+    else:
+        dit = mlx_wan22_dit_from_diffusers_safetensors(
+            args.dit_checkpoint,
+            args.dit_config,
+            dtype="fp16",
+            compile=args.compile,
+        )
+    dit_load_s = time.perf_counter() - t1
+    print(f"[5B] DiT loaded in {dit_load_s:.1f}s", flush=True)
+
+    freqs = make_rotary_embeddings(config, latent_frames=lat_t, latent_height=lat_h, latent_width=lat_w)
+    gen = torch.Generator().manual_seed(args.seed)
+    noise = mx.array(
+        torch.randn(1, in_ch, lat_t, lat_h, lat_w, generator=gen, dtype=torch.float32).numpy()).astype(mx.float16)
+
+    steps = [int(s) for s in args.dmd_denoising_steps.split(",") if s.strip()]
+    t2 = time.perf_counter()
+    mx.reset_peak_memory()
+    latents = sample_wan22_dmd(
+        dit,
+        ehs,
+        noise,
+        freqs,
+        dmd_denoising_steps=steps,
+        flow_shift=args.flow_shift,
+        warp_denoising_step=not args.no_warp,
+        seed=args.renoise_seed,
+    )
+    if spatial_mode == "refine":
+        schedule, warped_steps = build_wan22_dmd_schedule(
+            steps, flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp,
+        )
+        # The grid opens at sigma == 1, where the hand-off
+        # `(1 - sigma) * upsampled + sigma * noise` weights stage 1 at zero and
+        # refine silently becomes a plain full-res run. Drop the leading
+        # full-noise steps so stage 1 actually reaches stage 2.
+        stage2_warped = default_refine_timesteps(schedule, warped_steps)
+        stage2_steps = steps[len(warped_steps) - len(stage2_warped):]
+        sigma = schedule.sigma_for(stage2_warped[0])
+        print(f"[5B refine] stage-2 steps={stage2_steps} sigma={sigma:.4f} "
+              f"(stage-1 weight {1.0 - sigma:.4f})", flush=True)
+        latents = prepare_refine_latents(
+            latents, scale=args.refine_scale, sigma=sigma,
+            add_noise_flag=not args.no_refine_add_noise,
+            upsample_mode=args.refine_upsample_mode, seed=args.renoise_seed + 1,
+        )
+        freqs_stage2 = make_rotary_embeddings(
+            config, latent_frames=lat_t,
+            latent_height=active_plan.stage2_latent_height,
+            latent_width=active_plan.stage2_latent_width,
+        )
+        latents = sample_wan22_dmd(
+            dit, ehs, latents, freqs_stage2, dmd_denoising_steps=stage2_steps,
+            flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp,
+            seed=args.renoise_seed + 2,
+        )
+    # spatial_mode == "fast_spatial" leaves the latents on the stage-1 grid;
+    # the resample happens after decode, in _postprocess_video.
+    denoise_s = time.perf_counter() - t2
+    peak = mx.get_peak_memory() / (1024**3)
+    print(f"[5B] denoise {len(steps)} steps in {denoise_s:.1f}s, peak {peak:.2f} GiB", flush=True)
+
+    latents_np = np.array(latents.astype(mx.float32))
+    if args.save_latents is not None:
+        args.save_latents.parent.mkdir(parents=True, exist_ok=True)
+        np.savez(args.save_latents, latents=latents_np, prompt=args.prompt, seed=args.seed)
+        print(f"[5B] wrote latents {args.save_latents}", flush=True)
+
+    if spatial_mode == "refine":
+        del freqs_stage2
+    del dit, latents, ehs, noise, freqs
+    cleanup_mlx()
+
+    metrics = decode_latents_to_video(
+        latents_np,
+        args.output_path,
+        fps=args.fps,
+        backend=args.decode_backend,
+        vae_dir=args.vae_root if args.decode_backend == "wan-vae" else None,
+        z_dim=in_ch,
+    )
+    # One h264 round-trip for both post-decode passes (see _postprocess_video).
+    rife_s = 0.0
+    rife_request = ({
+        "factor": args.fast_factor,
+        "target_frames": target_frames,
+        "sharpen": args.fast_sharpen,
+    } if args.fast else None)
+    spatial_request = fast_spatial_plan if spatial_mode == "fast_spatial" else None
+    if rife_request is not None or spatial_request is not None:
+        rife_start = time.perf_counter()
+        _postprocess_video(
+            video_path=args.output_path, fps=args.fps,
+            rife=rife_request, spatial=spatial_request,
+        )
+        rife_s = time.perf_counter() - rife_start
+    print(f"[5B] decoded via {metrics['backend']} in {metrics['decode_s']:.1f}s → {args.output_path}", flush=True)
+    summary = {
+        "output_path": str(args.output_path.resolve()),
+        "prompt": args.prompt,
+        "prompt_used": prompt_for_encode,
+        "enhance_prompt": args.enhance_prompt,
+        "enhance_backend": enhance_backend,
+        "enhance_elapsed_s": round(enhance_elapsed_s, 3),
+        "height": args.height,
+        "width": args.width,
+        "fps": args.fps,
+        "target_frames": target_frames,
+        "generated_frames": args.num_frames,
+        "seed": args.seed,
+        "renoise_seed": args.renoise_seed,
+        "dmd_denoising_steps": steps,
+        "flow_shift": args.flow_shift,
+        "warp": not args.no_warp,
+        "spatial_mode": spatial_mode,
+        "fast": args.fast,
+        "fast_factor": args.fast_factor if args.fast else None,
+        "fast_spatial_scale": args.fast_spatial_scale if args.fast_spatial else None,
+        "refine_scale": args.refine_scale if args.refine else None,
+        "decode_backend": args.decode_backend,
+        "prompt_encode_s": round(prompt_encode_s, 3),
+        "dit_load_s": round(dit_load_s, 3),
+        "denoise_s": round(denoise_s, 3),
+        "decode_s": round(metrics["decode_s"], 3),
+        "rife_s": round(rife_s, 3),
+        "wall_total_s": round(time.perf_counter() - total_start, 3),
+        "peak_gib": round(peak, 3),
+        "latent_shape": [in_ch, lat_t, lat_h, lat_w],
+        "stage2_latent_shape": [in_ch, lat_t, active_plan.stage2_latent_height, active_plan.stage2_latent_width],
+        "mlx_checkpoint": str(args.mlx_checkpoint.resolve()) if args.mlx_checkpoint else None,
+    }
+    if args.metrics_json is not None:
+        args.metrics_json.parent.mkdir(parents=True, exist_ok=True)
+        args.metrics_json.write_text(json.dumps(summary, indent=2) + "\n")
+        print(f"[5B] wrote metrics {args.metrics_json}", flush=True)
+    print(json.dumps(summary, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/mlx_wan_decode_benchmark.py

@@ -0,0 +1,130 @@
+"""Compare Wan VAE and TAEHV decode on saved FastWan latents."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+
+from examples.inference.basic.mlx_wan_prompt_to_video import DEFAULT_MODEL_ROOT, decode_latents_to_video
+
+
+def _torch_mps_memory() -> dict[str, int | None]:
+    """
+    Report current and recommended memory usage for the MPS backend.
+
+    Returns:
+        dict[str, int | None]: Memory metrics in bytes, or `None` values when
+            PyTorch or MPS is unavailable.
+    """
+    try:
+        import torch
+    except ImportError:
+        return {
+            "current_allocated_bytes": None,
+            "driver_allocated_bytes": None,
+            "recommended_max_bytes": None,
+        }
+    if not torch.backends.mps.is_available():
+        return {
+            "current_allocated_bytes": None,
+            "driver_allocated_bytes": None,
+            "recommended_max_bytes": None,
+        }
+    return {
+        "current_allocated_bytes": int(torch.mps.current_allocated_memory()),
+        "driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
+        "recommended_max_bytes": int(torch.mps.recommended_max_memory()),
+    }
+
+
+def _parse_backends(raw: str) -> list[str]:
+    """
+    Parse and validate a comma-separated list of decoding backends.
+
+    Parameters:
+        raw (str): Comma-separated backend names.
+
+    Returns:
+        list[str]: Trimmed, supported backend names in input order.
+
+    Raises:
+        ValueError: If the input contains an unsupported backend.
+    """
+    backends = [backend.strip() for backend in raw.split(",") if backend.strip()]
+    allowed = {"wan-vae", "taehv"}
+    unknown = sorted(set(backends) - allowed)
+    if unknown:
+        raise ValueError(f"Unsupported decode backends: {unknown}")
+    return backends
+
+
+def main() -> None:
+    """
+    Benchmark selected Wan latent decoding backends and record their performance metrics.
+
+    Loads the specified latent array, decodes it with each selected backend, exports the
+    results as MP4 files, and writes per-backend timing and Torch MPS memory metrics to
+    `metrics.json`.
+    """
+    parser = argparse.ArgumentParser(description="Benchmark decode backends on saved Wan/FastWan latents.")
+    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
+    parser.add_argument("--latents-path", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_decode_benchmark"))
+    parser.add_argument("--backends", default="wan-vae,taehv")
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--torch-device", default="auto")
+    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
+    parser.add_argument("--taehv-source-path", type=Path, default=None)
+    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
+    parser.add_argument("--taehv-parallel", action="store_true")
+    args = parser.parse_args()
+
+    latents = np.load(args.latents_path)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    rows = []
+
+    for backend in _parse_backends(args.backends):
+        print(f"=== Decode backend: {backend} ===")
+        before = _torch_mps_memory()
+        start = time.perf_counter()
+        output_path = args.output_dir / f"{args.latents_path.stem}_{backend}.mp4"
+        decode_latents_to_video(
+            model_root=args.model_root,
+            latents_np=latents,
+            output_path=output_path,
+            fps=args.fps,
+            device_arg=args.torch_device,
+            dtype_arg=args.torch_dtype,
+            backend=backend,
+            taehv_source_path=args.taehv_source_path,
+            taehv_checkpoint_path=args.taehv_checkpoint_path,
+            taehv_parallel=args.taehv_parallel,
+        )
+        elapsed = time.perf_counter() - start
+        after = _torch_mps_memory()
+        metrics = {
+            "backend": backend,
+            "latents_path": str(args.latents_path),
+            "latents_shape": list(latents.shape),
+            "decode_export_s": elapsed,
+            "torch_mps_current_before_bytes": before["current_allocated_bytes"],
+            "torch_mps_current_after_bytes": after["current_allocated_bytes"],
+            "torch_mps_driver_before_bytes": before["driver_allocated_bytes"],
+            "torch_mps_driver_after_bytes": after["driver_allocated_bytes"],
+            "torch_mps_recommended_max_bytes": after["recommended_max_bytes"],
+            "output_path": str(output_path),
+        }
+        rows.append(metrics)
+        print(json.dumps(metrics, indent=2))
+
+    metrics_path = args.output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(rows, indent=2))
+    print(f"Wrote decode metrics to: {metrics_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/mlx_wan_quant_benchmark.py

@@ -0,0 +1,365 @@
+"""Benchmark MLX FastWan quantization modes with one shared prompt encode."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+from typing import cast
+
+import numpy as np
+
+from examples.inference.basic.mlx_wan_prompt_to_video import (
+    DEFAULT_MODEL_ROOT,
+    decode_latents_to_video,
+    encode_prompt,
+    make_rotary_embeddings,
+)
+from fastvideo.mlx_runtime.memory import cleanup_mlx
+
+
+def _parse_modes(raw: str) -> list[str]:
+    """
+    Parse and validate a comma-separated list of quantization modes.
+
+    Parameters:
+        raw (str): Comma-separated mode names.
+
+    Returns:
+        list[str]: Normalized, whitespace-trimmed mode names.
+
+    Raises:
+        ValueError: If any mode is unsupported.
+    """
+    modes = [mode.strip() for mode in raw.split(",") if mode.strip()]
+    allowed = {"none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"}
+    unknown = sorted(set(modes) - allowed)
+    if unknown:
+        raise ValueError(f"Unsupported modes: {unknown}")
+    return modes
+
+
+def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]:
+    """
+    Compare candidate and baseline latent arrays using error and signal-quality metrics.
+
+    Parameters:
+        candidate (np.ndarray): Latent array to evaluate.
+        baseline (np.ndarray): Reference latent array for comparison.
+
+    Returns:
+        dict[str, float]: Mean squared error, mean absolute error, maximum absolute
+        error, and signal-to-noise ratio in decibels between the arrays.
+    """
+    diff = candidate.astype(np.float32) - baseline.astype(np.float32)
+    mse = float(np.mean(np.square(diff)))
+    mae = float(np.mean(np.abs(diff)))
+    max_abs = float(np.max(np.abs(diff)))
+    signal = float(np.mean(np.square(baseline.astype(np.float32))))
+    return {
+        "latent_mse_vs_fp16": mse,
+        "latent_mae_vs_fp16": mae,
+        "latent_max_abs_vs_fp16": max_abs,
+        "latent_snr_db_vs_fp16": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"),
+    }
+
+
+def _torch_mps_memory() -> dict[str, int | None]:
+    """
+    Report PyTorch MPS memory statistics when PyTorch MPS is available.
+
+    Returns:
+        dict[str, int | None]: A mapping of MPS memory metric names to byte counts, or `None` values when PyTorch or MPS is unavailable.
+    """
+    try:
+        import torch
+    except ImportError:
+        return {
+            "torch_mps_current_allocated_bytes": None,
+            "torch_mps_driver_allocated_bytes": None,
+            "torch_mps_recommended_max_bytes": None,
+        }
+    if not torch.backends.mps.is_available():
+        return {
+            "torch_mps_current_allocated_bytes": None,
+            "torch_mps_driver_allocated_bytes": None,
+            "torch_mps_recommended_max_bytes": None,
+        }
+    return {
+        "torch_mps_current_allocated_bytes": int(torch.mps.current_allocated_memory()),
+        "torch_mps_driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
+        "torch_mps_recommended_max_bytes": int(torch.mps.recommended_max_memory()),
+    }
+
+
+def _decode_with_metrics(*, args, latents: np.ndarray, output_path: Path) -> dict[str, float | int | None | str]:
+    """
+    Decode latents to a video and collect export timing and PyTorch MPS memory metrics.
+
+    Parameters:
+        args: Configuration values for decoding and video export.
+        latents (np.ndarray): Latent representation to decode.
+        output_path (Path): Destination path for the exported video.
+
+    Returns:
+        dict[str, float | int | None | str]: Video export duration and PyTorch MPS memory measurements.
+    """
+    before = _torch_mps_memory()
+    decode_start = time.perf_counter()
+    decode_latents_to_video(
+        model_root=args.model_root,
+        latents_np=latents,
+        output_path=output_path,
+        fps=args.fps,
+        device_arg=args.torch_device,
+        dtype_arg=args.torch_dtype,
+        backend=args.decode_backend,
+        taehv_source_path=args.taehv_source_path,
+        taehv_checkpoint_path=args.taehv_checkpoint_path,
+        taehv_parallel=args.taehv_parallel,
+    )
+    decode_time = time.perf_counter() - decode_start
+    after = _torch_mps_memory()
+    return {
+        "decode_export_s": decode_time,
+        "decode_torch_mps_current_before_bytes": before["torch_mps_current_allocated_bytes"],
+        "decode_torch_mps_current_after_bytes": after["torch_mps_current_allocated_bytes"],
+        "decode_torch_mps_driver_before_bytes": before["torch_mps_driver_allocated_bytes"],
+        "decode_torch_mps_driver_after_bytes": after["torch_mps_driver_allocated_bytes"],
+        "decode_torch_mps_recommended_max_bytes": after["torch_mps_recommended_max_bytes"],
+    }
+
+
+def _run_one_mode(
+    *,
+    mode: str,
+    args,
+    config: dict,
+    checkpoint_path: Path,
+    config_path: Path,
+    prompt_embeds,
+    freqs_cis,
+):
+    """
+    Run denoising for one quantization mode and collect performance and memory metrics.
+
+    Parameters:
+        mode (str): Quantization mode to benchmark.
+        args: Benchmark configuration, including dtype, dimensions, seed, scheduler, and denoising settings.
+        config (dict): Model configuration containing the input channel count.
+        checkpoint_path (Path): Path to the transformer checkpoint.
+        config_path (Path): Path to the transformer configuration.
+        prompt_embeds: Encoded prompt embeddings shared across benchmark modes.
+        freqs_cis: Rotary positional embeddings used during denoising.
+
+    Returns:
+        dict: The mode name, generated latent array, and metrics for model loading,
+        denoising, step timing, and MLX memory usage.
+    """
+    import mlx.core as mx
+    import torch
+
+    from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+    from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
+    from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step
+
+    mx_dtype = mx.float16 if args.mlx_dtype == "fp16" else mx.float32
+    quantization = None if mode == "none" else mode
+    latent_frames = (args.num_frames - 1) // 4 + 1
+    latent_height = args.height // 8
+    latent_width = args.width // 8
+
+    load_start = time.perf_counter()
+    mx.clear_cache()
+    mx.reset_peak_memory()
+    dit = mlx_dit_from_diffusers_safetensors(
+        checkpoint_path,
+        config_path,
+        dtype=args.mlx_dtype,
+        quantization=quantization,
+    )
+    load_time = time.perf_counter() - load_start
+    load_peak_memory = mx.get_peak_memory()
+
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
+    schedule = MLXDMDSchedule.from_torch_scheduler(scheduler)
+    timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()]
+    # Same torch generator sequence as the original host-round-trip loop
+    # (initial latents first, then one re-noise draw per intermediate step),
+    # so every mode still shares identical stochasticity.
+    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    latents_seed = torch.randn(
+        (1, int(config["in_channels"]), latent_frames, latent_height, latent_width),
+        generator=generator,
+        dtype=torch.float32,
+    ).numpy()
+    renoise_by_step = [
+        torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy()
+        for _ in range(max(0, len(timesteps) - 1))
+    ]
+    latents = mx.array(latents_seed).astype(mx_dtype)
+    encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx_dtype)
+
+    denoise_start = time.perf_counter()
+    mx.reset_peak_memory()
+    latents_np, step_times = denoise_dmd_on_device(
+        mx=mx,
+        dit=dit,
+        latents=latents,
+        encoder_hidden_states=encoder_hidden_states,
+        freqs_cis=freqs_cis,
+        timesteps=timesteps,
+        renoise_by_step=renoise_by_step,
+        schedule=schedule,
+        dmd_step=dmd_step,
+        mx_dtype=mx_dtype,
+    )
+    denoise_time = time.perf_counter() - denoise_start
+    denoise_peak_memory = mx.get_peak_memory()
+    active_memory = mx.get_active_memory()
+    return {
+        "mode": mode,
+        "latents": latents_np,
+        "metrics": {
+            "mlx_dit_load_s": load_time,
+            "mlx_denoise_s": denoise_time,
+            "mlx_denoise_first_step_s": step_times[0] if step_times else None,
+            "mlx_load_peak_bytes": int(load_peak_memory),
+            "mlx_denoise_peak_bytes": int(denoise_peak_memory),
+            "mlx_active_after_denoise_bytes": int(active_memory),
+        },
+    }
+
+
+def main() -> None:
+    """
+    Run the MLX FastWan quantization benchmark for the selected modes and write latency, memory, output, and latent-difference metrics to the output directory.
+    """
+    parser = argparse.ArgumentParser(description="Benchmark MLX FastWan quantization modes.")
+    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
+    parser.add_argument("--prompt", default="A snow leopard walks across a windy mountain ridge.")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=17)
+    parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
+    parser.add_argument("--flow-shift", type=float, default=8.0)
+    parser.add_argument("--max-sequence-length", type=int, default=256)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--torch-device", default="auto")
+    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
+    parser.add_argument("--mlx-dtype", choices=("fp16", "fp32"), default="fp16")
+    parser.add_argument("--modes", default="none,int8,int4,mxfp8,mxfp4,nvfp4")
+    parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_quant_benchmark"))
+    parser.add_argument("--decode-backend", choices=("none", "wan-vae", "taehv"), default="taehv")
+    parser.add_argument("--taehv-source-path", type=Path, default=None)
+    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
+    parser.add_argument("--taehv-parallel", action="store_true")
+    args = parser.parse_args()
+
+    import mlx.core as mx
+    import torch
+
+    mx.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    config_path = args.model_root / "transformer/config.json"
+    checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors"
+    config = json.loads(config_path.read_text())
+    latent_frames = (args.num_frames - 1) // 4 + 1
+    latent_height = args.height // 8
+    latent_width = args.width // 8
+
+    prompt_start = time.perf_counter()
+    prompt_embeds = encode_prompt(
+        model_root=args.model_root,
+        prompt=args.prompt,
+        max_sequence_length=args.max_sequence_length,
+        device_arg=args.torch_device,
+        dtype_arg=args.torch_dtype,
+    )
+    prompt_time = time.perf_counter() - prompt_start
+    freqs_cis = make_rotary_embeddings(
+        config,
+        latent_frames=latent_frames,
+        latent_height=latent_height,
+        latent_width=latent_width,
+    )
+
+    from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError
+
+    baseline_latents = None
+    rows = []
+    for mode in _parse_modes(args.modes):
+        print(f"=== MLX quant mode: {mode} ===")
+        mode_start = time.perf_counter()
+        try:
+            result = _run_one_mode(
+                mode=mode,
+                args=args,
+                config=config,
+                checkpoint_path=checkpoint_path,
+                config_path=config_path,
+                prompt_embeds=prompt_embeds,
+                freqs_cis=freqs_cis,
+            )
+        except UnsupportedMLXQuantizationError as exc:
+            print(f"skipping mode (unsupported by this MLX build): {exc}")
+            rows.append({"mode": mode, "status": "unsupported_by_mlx", "error": str(exc)})
+            continue
+        cleanup_mlx(mx)
+        latents = result["latents"]
+        if baseline_latents is None:
+            baseline_latents = latents
+        latent_path = args.output_dir / f"latents_{mode}.npy"
+        np.save(latent_path, latents)
+
+        decode_time = 0.0
+        decode_metrics = {}
+        output_path = None
+        if args.decode_backend != "none":
+            output_path = args.output_dir / f"video_{mode}_{args.decode_backend}_{args.height}x{args.width}x{args.num_frames}.mp4"
+            decode_metrics = _decode_with_metrics(args=args, latents=latents, output_path=output_path)
+            decode_time = cast(float, decode_metrics["decode_export_s"])
+
+        mode_total = time.perf_counter() - mode_start
+        mlx_denoise_peak_bytes = int(result["metrics"]["mlx_denoise_peak_bytes"])
+        mlx_active_bytes = int(result["metrics"]["mlx_active_after_denoise_bytes"])
+        metrics = {
+            "mode": mode,
+            "status": "ok",
+            "prompt_encode_shared_s": prompt_time,
+            "height": args.height,
+            "width": args.width,
+            "num_frames": args.num_frames,
+            "decode_backend": args.decode_backend,
+            "decode_export_s": decode_time,
+            "mode_total_excluding_shared_prompt_s": mode_total,
+            "mode_total_including_shared_prompt_s": mode_total + prompt_time,
+            "latents_path": str(latent_path),
+            "output_path": str(output_path) if output_path else None,
+            "mlx_denoise_peak_gib": mlx_denoise_peak_bytes / (1024**3),
+            "mlx_active_after_denoise_gib": mlx_active_bytes / (1024**3),
+            "mlx_dit_peak_under_16gb": mlx_denoise_peak_bytes < 16 * 1024**3,
+            "mlx_dit_active_under_16gb": mlx_active_bytes < 16 * 1024**3,
+            "mac_16gb_status": (
+                "dit_memory_fits_16gb_measured_decode_separately"
+                if mlx_denoise_peak_bytes < 16 * 1024**3 else "dit_memory_exceeds_16gb"
+            ),
+            **result["metrics"],
+            **decode_metrics,
+            **_latent_delta_metrics(latents, baseline_latents),
+        }
+        rows.append(metrics)
+        print(json.dumps(metrics, indent=2))
+
+    metrics_path = args.output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(rows, indent=2))
+    print(f"Wrote benchmark metrics to: {metrics_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/mlx_wan_video_quality.py

@@ -0,0 +1,104 @@
+"""Compare generated MP4s against a reference MP4 with simple pixel metrics."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+def _read_video(path: Path) -> np.ndarray:
+    """
+    Read all frames from a video file as an RGB NumPy array.
+
+    Parameters:
+        path (Path): Path to the video file.
+
+    Returns:
+        np.ndarray: Video frames stacked along the first axis.
+
+    Raises:
+        ValueError: If the video contains no readable frames.
+    """
+    import cv2
+
+    cap = cv2.VideoCapture(str(path))
+    frames = []
+    try:
+        while True:
+            ok, frame_bgr = cap.read()
+            if not ok:
+                break
+            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+            frames.append(frame_rgb)
+    finally:
+        cap.release()
+    if not frames:
+        raise ValueError(f"No frames read from {path}")
+    return np.stack(frames, axis=0)
+
+
+def _metrics(candidate: np.ndarray, reference: np.ndarray) -> dict[str, float | int | list[int]]:
+    """
+    Compute pixel-level comparison metrics between candidate and reference video frames.
+
+    Parameters:
+        candidate (np.ndarray): Candidate video frames in frame, height, width, and channel order.
+        reference (np.ndarray): Reference video frames with the same shape as the candidate.
+
+    Returns:
+        dict[str, float | int | list[int]]: Frame dimensions and pixel comparison metrics, including MSE, MAE, maximum absolute difference, and PSNR in decibels.
+
+    Raises:
+        ValueError: If the candidate and reference arrays have different shapes.
+    """
+    if candidate.shape != reference.shape:
+        raise ValueError(f"Shape mismatch: candidate={candidate.shape}, reference={reference.shape}")
+    candidate_f = candidate.astype(np.float32)
+    reference_f = reference.astype(np.float32)
+    diff = candidate_f - reference_f
+    mse = float(np.mean(np.square(diff)))
+    mae = float(np.mean(np.abs(diff)))
+    max_abs = float(np.max(np.abs(diff)))
+    psnr = float(20.0 * np.log10(255.0 / np.sqrt(mse))) if mse > 0 else float("inf")
+    return {
+        "frames": int(candidate.shape[0]),
+        "height": int(candidate.shape[1]),
+        "width": int(candidate.shape[2]),
+        "channels": int(candidate.shape[3]),
+        "mse_vs_reference": mse,
+        "mae_vs_reference": mae,
+        "max_abs_vs_reference": max_abs,
+        "psnr_db_vs_reference": psnr,
+    }
+
+
+def main() -> None:
+    """Compare candidate MP4 videos with a reference and write pixel-level metrics to a JSON file."""
+    parser = argparse.ArgumentParser(description="Compare MP4s against a reference MP4.")
+    parser.add_argument("--reference", type=Path, required=True)
+    parser.add_argument("--candidates", type=Path, nargs="+", required=True)
+    parser.add_argument("--metrics-json", type=Path, required=True)
+    args = parser.parse_args()
+
+    reference = _read_video(args.reference)
+    rows = []
+    for candidate_path in args.candidates:
+        candidate = _read_video(candidate_path)
+        row = {
+            "reference_path": str(args.reference),
+            "candidate_path": str(candidate_path),
+            **_metrics(candidate, reference),
+        }
+        rows.append(row)
+        print(json.dumps(row, indent=2))
+
+    args.metrics_json.parent.mkdir(parents=True, exist_ok=True)
+    args.metrics_json.write_text(json.dumps(rows, indent=2))
+    print(f"Wrote video quality metrics to: {args.metrics_json}")
+
+
+if __name__ == "__main__":
+    main()

examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh

@@ -19,16 +19,13 @@ export FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1
 export FASTVIDEO_TORCH_PROFILER_WITH_STACK=1  
 export FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1
 export FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY=1
-export FASTVIDEO_TORCH_PROFILER_WAIT_STEPS=2
-export FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS=1
-export FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS=1
 export FASTVIDEO_TORCH_PROFILER_DIR="../profiler_traces/wan_t2v_finetune/"
 
 # Training arguments
 training_args=(
   --tracker_project_name "wan_t2v_finetune"
   --output_dir "checkpoints/wan_t2v_finetune"
-  --max_train_steps 5000
+  --max_train_steps 1  # Profiler captures every selected region until shutdown.
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 8

fastvideo-kernel/CMakeLists.txt

@@ -114,7 +114,63 @@ list(APPEND TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_PATHS})
 # Find Torch package (still useful for libraries)
 find_package(Torch REQUIRED)
 
-# Include directories
+# The Ulysses kernel needs NCCL's 2.29 device API. Keep it optional so ROCm,
+# older NCCL installs, and minimal CUDA builders still produce a usable wheel.
+set(FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A "AUTO" CACHE STRING
+    "Build the NCCL-device-API Ulysses all-to-all: AUTO/ON/OFF")
+set_property(CACHE FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A PROPERTY STRINGS AUTO ON OFF)
+set(ENABLE_ULYSSES_A2A OFF)
+
+if(NOT GPU_BACKEND STREQUAL "ROCM" AND NOT FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A STREQUAL "OFF")
+    # nvidia.nccl is a namespace package, so find_spec rather than __file__.
+    execute_process(COMMAND "${Python_EXECUTABLE}" -c
+        "import importlib.util as u;s=u.find_spec('nvidia.nccl');print(list(s.submodule_search_locations)[0] if s else '')"
+        OUTPUT_VARIABLE NCCL_PIP_ROOT OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET)
+    find_path(NCCL_INCLUDE_DIR nccl_device/coop.h
+        HINTS $ENV{NCCL_HOME}/include ${NCCL_PIP_ROOT}/include)
+    find_library(NCCL_LIBRARY NAMES nccl
+        HINTS $ENV{NCCL_HOME}/lib $ENV{NCCL_HOME}/lib64 ${NCCL_PIP_ROOT}/lib)
+    # PyPI's nvidia-nccl-cu* wheels ship the SONAME but not the development
+    # symlink (libnccl.so.2, no libnccl.so). Accept that exact versioned name;
+    # target_link_libraries can link an absolute SONAME path directly.
+    if(NOT NCCL_LIBRARY)
+        find_file(NCCL_VERSIONED_LIBRARY NAMES libnccl.so.2
+            HINTS $ENV{NCCL_HOME}/lib $ENV{NCCL_HOME}/lib64 ${NCCL_PIP_ROOT}/lib)
+        if(NCCL_VERSIONED_LIBRARY)
+            set(NCCL_LIBRARY "${NCCL_VERSIONED_LIBRARY}")
+        endif()
+    endif()
+
+    if(NCCL_INCLUDE_DIR AND NCCL_LIBRARY)
+        include(CheckCXXSourceCompiles)
+        set(_FASTVIDEO_REQUIRED_INCLUDES "${CMAKE_REQUIRED_INCLUDES}")
+        # nccl.h includes cuda_runtime.h, so a host-compiler feature probe
+        # needs the toolkit includes explicitly even though CUDA is enabled.
+        set(CMAKE_REQUIRED_INCLUDES "${NCCL_INCLUDE_DIR};${CUDAToolkit_INCLUDE_DIRS}")
+        unset(NCCL_HAS_REQUIRED_DEVICE_API CACHE)
+        check_cxx_source_compiles("\
+#define NCCL_HOSTLIB_ONLY
+#include <cstddef>
+#include <nccl_device.h>
+int main() {
+  ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
+  ncclCommProperties props = NCCL_COMM_PROPERTIES_INITIALIZER;
+  return (reqs.size == 0 || props.size == 0);
+}"
+            NCCL_HAS_REQUIRED_DEVICE_API)
+        set(CMAKE_REQUIRED_INCLUDES "${_FASTVIDEO_REQUIRED_INCLUDES}")
+    endif()
+
+    if(NCCL_HAS_REQUIRED_DEVICE_API)
+        set(ENABLE_ULYSSES_A2A ON)
+    elseif(FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A STREQUAL "ON")
+        message(FATAL_ERROR
+            "FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON requires NCCL device headers/library "
+            "with NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER (NCCL 2.29+). "
+            "Resolved include='${NCCL_INCLUDE_DIR}', library='${NCCL_LIBRARY}'.")
+    endif()
+endif()
+
 include_directories(
     ${CMAKE_SOURCE_DIR}/include
     ${CMAKE_SOURCE_DIR}/include/cutlass/include
@@ -124,6 +180,9 @@ include_directories(
     ${CMAKE_SOURCE_DIR}/csrc/turbodiffusion
     ${TORCH_INCLUDE_DIRS}
 )
+if(ENABLE_ULYSSES_A2A)
+    include_directories(${NCCL_INCLUDE_DIR})
+endif()
 
 # ---------------------------
 # ThunderKittens (TK) toggles
@@ -157,6 +216,7 @@ endif()
 message(STATUS "TORCH_CUDA_ARCH_LIST (cmake/env): ${TORCH_CUDA_ARCH_LIST}")
 message(STATUS "FASTVIDEO_KERNEL_BUILD_TK: ${FASTVIDEO_KERNEL_BUILD_TK}")
 message(STATUS "FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER: ${FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER}")
+message(STATUS "FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A: ${FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A}")
 
 set(ENABLE_TK_KERNELS OFF)
 if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON")
@@ -307,6 +367,9 @@ if(BUILD_CXX_KERNELS)
         csrc/turbodiffusion/norm/layernorm.cu
         csrc/turbodiffusion/quant/quant.cu
     )
+    if(ENABLE_ULYSSES_A2A)
+        list(APPEND EXTENSION_SOURCES csrc/comm/ulysses_all_to_all.cu)
+    endif()
     
     # Conditionally add TK kernels
     if(ENABLE_TK_KERNELS)
@@ -318,6 +381,38 @@ if(BUILD_CXX_KERNELS)
 
     # Combined FastVideo Extension
     # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace
+    # ---------------------------------------------------------------------------
+    # VSA block-sparse attention forward, Blackwell (sm_100a) only.
+    #
+    # NOTE the "a" suffix: -arch=sm_100a is NOT enough -- it emits a plain sm_100 target and
+    # ptxas rejects every tcgen05 / setmaxnreg instruction. The explicit gencode spelling
+    # below is required, and matches the 10.0a entry in TORCH_CUDA_ARCH_LIST.
+    #
+    # Built for 64-token sparse blocks -- FastVideo's default (4,4,4) tiling, so no
+    # tile-size change and no top-k granularity change is needed. VSA_BHSD
+    # selects [B, H, S, D]. Both are compile-time; the Python is_supported() checks incoming
+    # tensors against them so callers fall back to Triton rather than getting a wrong answer.
+    # Read the ENVIRONMENT as well as the cache variable. When TORCH_CUDA_ARCH_LIST is
+    # exported (build.sh, and `pip install` with it set) the branch above only prints it --
+    # the cmake variable stays empty, so testing that alone silently skips the kernel and
+    # leaves a build that succeeds with the op missing.
+    set(ENABLE_VSA_SM100A OFF)
+    set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
+    if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST})
+        set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
+    endif()
+    if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)")
+        set(ENABLE_VSA_SM100A ON)
+    endif()
+    if(ENABLE_VSA_SM100A)
+        message(STATUS "fastvideo-kernel: building block_sparse_sm100a (Blackwell, 64- and 128-token blocks)")
+        list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu
+                                      csrc/attention/block_sparse_blk128_sm100a.cu)
+        set_source_files_properties(csrc/attention/block_sparse_sm100a.cu
+                                    csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES
+            COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true")
+    endif()
+
     Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
         ${EXTENSION_SOURCES}
     )
@@ -333,9 +428,16 @@ if(BUILD_CXX_KERNELS)
 
     # Build compile definitions list
     set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
+    if(ENABLE_VSA_SM100A)
+        list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A)
+    endif()
     if(ENABLE_TK_KERNELS)
         list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
     endif()
+    if(ENABLE_ULYSSES_A2A)
+        list(APPEND COMPILE_DEFS FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A)
+    endif()
+
 
     target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS})
 
@@ -346,6 +448,9 @@ if(BUILD_CXX_KERNELS)
     # Link against Torch libraries to avoid undefined symbols at import time
     # (e.g., torch::autograd vtables) when loading the extension module.
     target_link_libraries(fastvideo_kernel_ops PRIVATE ${TORCH_LIBRARIES})
+    if(ENABLE_ULYSSES_A2A)
+        target_link_libraries(fastvideo_kernel_ops PRIVATE "${NCCL_LIBRARY}")
+    endif()
 
     # Also link against libtorch_python to satisfy Python-binding symbols
     # (e.g., torch::PyWarningHandler) required by torch/extension.h.
@@ -449,6 +554,7 @@ message(STATUS "host / backend:       ${CMAKE_SYSTEM_PROCESSOR} / ${GPU_BACKEND}
 message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
 message(STATUS "fastvideo_kernel_ops: ON  (turbodiffusion int8-gemm/quant/rmsnorm/layernorm, all listed archs)")
 message(STATUS "  + TK sta/block_sparse (sm_90a only):          ${ENABLE_TK_KERNELS}")
+message(STATUS "  + Ulysses NCCL-device all-to-all:             ${ENABLE_ULYSSES_A2A}")
 message(STATUS "fp4attn/fp4quant (sm_120a only, CUDA >= 12.8):  ${ENABLE_ATTN_QAT_INFER}")
 message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.")
 message(STATUS "============================================================")

fastvideo-kernel/README.md

@@ -10,6 +10,8 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve
 |---|---|---|---|---|
 | `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built |
 | same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) |
+| same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_sm100a.cu` | Blackwell `sm_100a` only | ON iff `10.0a` is in `TORCH_CUDA_ARCH_LIST` |
+| same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) |
 | `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer Blackwell `sm_120a` only, CUDA ≥ 12.8 | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` is in the arch list) |
 
 Runtime-JIT kernels (no build step, ship in every wheel/image):
@@ -17,25 +19,28 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
 | Kernels | Where | Used when |
 |---|---|---|
 | Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
-| FA4 CuTe-DSL block-sparse (VSA-256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
+| FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
 | VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |
 
 ## What gets built where, and when
 
-| Surface | Trigger | Leg | `TORCH_CUDA_ARCH_LIST` | TK | FP4 |
-|---|---|---|---|---|---|
-| PyPI wheels (`.github/workflows/publish-kernel.yml`) | version bump in `fastvideo-kernel/pyproject.toml` on main, or manual dispatch | x86_64 cu126 | `9.0a` | ON | — (CUDA < 12.8) |
-| | | x86_64 cu130 | `9.0a;12.0a` | ON | ON |
-| | | aarch64 cu130 | `10.0a;12.0a` | — | ON |
-| Docker images `ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev` (`.github/workflows/infra-build-image.yml`) | `docker/Dockerfile` changes on main, or manual dispatch | amd64 cuda12.6.3 + cuda13.0.0 | `9.0a` | ON | — |
-| | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | — |
-| | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | — |
-| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | ON iff sm_120 |
+| Surface | Trigger | Leg | `TORCH_CUDA_ARCH_LIST` | TK | Ulysses | FP4 |
+|---|---|---|---|---|---|---|
+| PyPI wheels (`.github/workflows/publish-kernel.yml`) | version bump in `fastvideo-kernel/pyproject.toml` on main, or manual dispatch | x86_64 cu126 | `9.0a` | ON | AUTO | — (CUDA < 12.8) |
+| | | x86_64 cu130 | `9.0a;10.0a;12.0a` | ON | AUTO | ON |
+| | | aarch64 cu130 | `10.0a;12.0a` | — | AUTO | ON |
+| Docker images `ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev` (`.github/workflows/infra-build-image.yml`) | `docker/Dockerfile` changes on main, or manual dispatch | amd64 cuda12.6.3 + cuda13.0.0 | `9.0a` | ON | AUTO | — |
+| | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | AUTO | — |
+| | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | AUTO | — |
+| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | AUTO | ON iff sm_120 |
 
 Notes:
 
 - No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do.
+- Both cu130 PyPI wheels ship the MiniMax-H3 sm_100a VSA forward.
 - On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64.
+- Ulysses AUTO builds only when CMake finds a NCCL library and device-API headers with the 2.29 initializers. Use
+  `-DFASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON` to require it or `OFF` to test the portable build.
 - Kernel tests run on Buildkite GPU CI for PRs touching `fastvideo-kernel/**` (see `.buildkite/pipeline.yml`).
 
 ## Installation
@@ -64,28 +69,49 @@ cd fastvideo-kernel
 ./build.sh --rocm
 ```
 
-### Optional: FA4 CuTe block-sparse backend (VSA-256 fastpath)
+### Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath)
 
-The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes to the
+The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the
 FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is
 an **optional** dependency: it is imported lazily, and `video_sparse_attn`
 transparently falls back to the Triton backend when it is absent (so the package is
 fully usable without it).
 
-The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`,
-`flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by
+The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`
+and the public/private forward-backward bridges in `flash_attn.cute.interface`) are provided upstream by
 [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
-commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
+commit `14c377950125c70b7a9dabf9c561fca53715ac7d`, the revision FastVideo pins as
 the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
-have an incompatible `_flash_attn_fwd` signature.
+have incompatible block-sparse forward/backward interfaces.
+
+Install it under its distribution name so its own runtime stack resolves with it.
+Do **not** pre-install `nvidia-cutlass-dsl` by hand: this revision pins
+`nvidia-cutlass-dsl==4.6.0.dev0` exactly, and a hand-installed 4.5.x floor either
+gets silently upgraded or, if something else holds it back, leaves the CuTe
+kernels broken.
 
 ```bash
-pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
-pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
+pip install torchvision
+pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@14c377950125c70b7a9dabf9c561fca53715ac7d#subdirectory=flash_attn/cute"
 ```
 
-The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against
-`tests/test_vsa256_forward*.py`.
+That resolves `nvidia-cutlass-dsl` to 4.6.0.dev0 and `quack-kernels` to 0.5.3, a
+combination this revision works with. A mismatched CuTe DSL only surfaces when the
+kernel JIT-compiles, so the error points at CuTe internals rather than at the
+install:
+
+| Error on first VSA-128/256 CuTe call | Cause |
+|---|---|
+| `TypeError: fmax() missing 1 required positional argument: 'b'` | `nvidia-cutlass-dsl` 4.5.x |
+| `AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'` | `quack-kernels` older than 0.5.1 |
+| `ImportError: cannot import name 'alloc_reserved_mbarrier'` | `quack-kernels` 0.6.2 or newer |
+
+An environment whose `flash_attn.cute` came from a prebuilt flash-attn wheel rather
+than from this pin hits the first row; that is what the overlay step in
+`docker/Dockerfile` works around.
+
+The CuTe kernels JIT-compile on first use. Forward and backward are verified on
+Blackwell (sm_100) against `tests/test_vsa128_*.py` and `tests/test_vsa256_*.py`.
 
 ## Usage
 
@@ -142,6 +168,14 @@ After building/installing `fastvideo-kernel`, run:
 ```bash
 cd fastvideo-kernel
 python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
+
+# VSA-256 FA4 CuTe forward/backward on Blackwell
+python benchmarks/bench_vsa.py --block_size 256 --use_cute \
+  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 20
+
+# VSA-128 FA4 CuTe forward/backward on Blackwell
+python benchmarks/bench_vsa.py --block_size 128 --use_cute \
+  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 40
 ```
 
 ### TurboDiffusion Kernels

fastvideo-kernel/benchmarks/bench_vsa.py

@@ -2,8 +2,9 @@
 """
 Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs.
 
-This script benchmarks the autograd-enabled wrapper:
-  - fastvideo_kernel.block_sparse_attn.block_sparse_attn
+This script benchmarks the autograd-enabled wrappers:
+  - 64-token TK/Triton: fastvideo_kernel.block_sparse_attn.block_sparse_attn
+  - 128/256-token Triton/CuTe: fastvideo_kernel.block_sparse_attn_256
 
 So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time.
 """
@@ -23,9 +24,6 @@ try:
 except Exception as e:  # pragma: no cover
     raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
 
-BLOCK_M = 64
-BLOCK_N = 64
-
 
 def set_seed(seed: int = 42) -> None:
     random.seed(seed)
@@ -41,7 +39,11 @@ def parse_arguments() -> argparse.Namespace:
     p.add_argument("--num_heads", type=int, default=12)
     p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
     p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)")
-    p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)")
+    p.add_argument("--q_seq_lens",
+                   type=int,
+                   nargs="+",
+                   default=[49152],
+                   help="Q sequence lengths (must be divisible by --block_size)")
     p.add_argument("--kv_seq_lens",
                    type=int,
                    nargs="+",
@@ -51,9 +53,13 @@ def parse_arguments() -> argparse.Namespace:
     p.add_argument("--rep", type=int, default=20)
     p.add_argument("--seed", type=int, default=42)
     p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
+    p.add_argument("--block_size", type=int, default=64, choices=[64, 128, 256])
     p.add_argument("--force_triton",
                    action="store_true",
                    help="Force wrapper to use Triton path (if supported by shapes).")
+    p.add_argument("--use_cute",
+                   action="store_true",
+                   help="Use the optional FA4 CuTe forward/backward path (requires --block_size 128 or 256).")
     return p.parse_args()
 
 
@@ -84,18 +90,38 @@ def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float:
     return do_bench(fn, warmup=warmup, rep=rep, quantiles=None)
 
 
+def _configure_backend(args: argparse.Namespace) -> None:
+    if args.use_cute and args.block_size not in (128, 256):
+        raise ValueError("--use_cute requires --block_size 128 or 256")
+    if args.use_cute and args.force_triton:
+        raise ValueError("--use_cute and --force_triton are mutually exclusive")
+
+    if args.force_triton:
+        os.environ.pop("FASTVIDEO_VSA_CUTEDSL", None)
+        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
+    elif args.use_cute:
+        os.environ.pop("FASTVIDEO_VSA_TRITON", None)
+        os.environ.pop("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", None)
+        os.environ["FASTVIDEO_VSA_CUTEDSL"] = "1"
+
+
 def main() -> None:
     args = parse_arguments()
     set_seed(args.seed)
+    _configure_backend(args)
 
     dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
 
-    if args.force_triton:
-        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
-
     from fastvideo_kernel.block_sparse_attn import block_sparse_attn
+    from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128, block_sparse_attn_256
 
     bs, h, d = args.batch_size, args.num_heads, args.head_dim
+    block_size = args.block_size
+    attention = {
+        64: block_sparse_attn,
+        128: block_sparse_attn_128,
+        256: block_sparse_attn_256,
+    }[block_size]
     kv_seq_lens = args.kv_seq_lens
     if kv_seq_lens is None:
         kv_seq_lens = args.q_seq_lens
@@ -105,20 +131,22 @@ def main() -> None:
     print("VSA Block-Sparse Attention Benchmark (WRAPPER)")
     print(f"device: {torch.cuda.get_device_name(0)}")
     print(f"batch={bs}, heads={h}, head_dim={d}, dtype={args.dtype}")
-    print(f"BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}")
+    print(f"block_size={block_size}")
     print("NOTE: timings include wrapper overhead (map->index + dispatch).")
-    if args.force_triton:
+    if args.use_cute:
+        print("dispatch: FA4 CuTe")
+    elif args.force_triton:
         print("dispatch: forced Triton (FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1)")
     else:
         print("dispatch: SM90 if available, else Triton")
 
     for q_len, kv_len in zip(args.q_seq_lens, kv_seq_lens):
-        if q_len % BLOCK_M != 0 or kv_len % BLOCK_N != 0:
-            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by 64")
+        if q_len % block_size != 0 or kv_len % block_size != 0:
+            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by {block_size}")
             continue
 
-        num_q_blocks = q_len // BLOCK_M
-        num_kv_blocks = kv_len // BLOCK_N
+        num_q_blocks = q_len // block_size
+        num_kv_blocks = kv_len // block_size
         topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10)
         topk = min(topk, num_kv_blocks)
 
@@ -129,11 +157,11 @@ def main() -> None:
         q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype)
         block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk)
 
-        # Variable block sizes: default full blocks (64 tokens per KV block)
-        variable_block_sizes = torch.full((num_kv_blocks, ), BLOCK_N, dtype=torch.int32, device="cuda")
+        # Variable block sizes: default full logical blocks.
+        variable_block_sizes = torch.full((num_kv_blocks, ), block_size, dtype=torch.int32, device="cuda")
 
         def _fwd():
-            return block_sparse_attn(q, k, v, block_map, variable_block_sizes)
+            return attention(q, k, v, block_map, variable_block_sizes)
 
         fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep)
 
@@ -142,7 +170,7 @@ def main() -> None:
         q_ = q.detach().requires_grad_(True)
         k_ = k.detach().requires_grad_(True)
         v_ = v.detach().requires_grad_(True)
-        o_, _aux_ = block_sparse_attn(q_, k_, v_, block_map, variable_block_sizes)
+        o_, _aux_ = attention(q_, k_, v_, block_map, variable_block_sizes)
         og = torch.randn_like(o_)
         loss = (o_ * og).sum()
 
@@ -156,7 +184,7 @@ def main() -> None:
             rep=max(5, args.rep // 2),
         )
 
-        flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N)
+        flops = flops_sparse_attention(bs, h, d, q_len, topk, block_size)
         fwd_tflops = flops / fwd_ms * 1e-12 * 1e3
         # Rough backward multiplier (attention backward typically ~2-3x forward)
         bwd_tflops = (2.5 * flops) / bwd_ms * 1e-12 * 1e3

fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu

@@ -0,0 +1,7 @@
+// block_sparse_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding.
+//
+// Same source as block_sparse_sm100a.cu with VSA_BLK128 set: the kernel and launch land in
+// namespace vsa_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the
+// exported entry point becomes block_sparse_sm100a_blk128_fwd.
+#define VSA_BLK128 true
+#include "block_sparse_sm100a.cu"

fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh

@@ -0,0 +1,201 @@
+#ifndef BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH
+#define BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH
+
+// Launch surface for the sm_100a VSA block-sparse FMHA forward.
+//
+// Everything a caller needs: a POD argument struct, a predicate saying whether this build can
+// run those arguments, and one launch entry point. The benchmark in
+// block_sparse_bench_sm100a.cu and the torch binding both go through here, so there is
+// one tensormap construction and one launch configuration rather than two that can drift.
+//
+// Two compile-time knobs select the four builds:
+//   VSA_BLK128  false -> 64-token sparse blocks, true -> 128-token
+//   VSA_BHSD    false -> [token][head][dim] (BSHD), true -> [batch][head][token][dim] (BHSD)
+
+#include "block_sparse_kernel_sm100a.cuh"
+
+namespace VSA_NAMESPACE {
+
+struct BlockSparseVsaArgs {
+  const __nv_bfloat16* q;
+  const __nv_bfloat16* k;
+  const __nv_bfloat16* v;      // natural layout; only blk128 reads it (blk64 still needs v_t)
+  const __nv_bfloat16* v_t;    // unused: kept so the bench's V_T buffer still binds
+  __nv_bfloat16* o;
+  float* lse;                  // [batch, num_heads, seqlen] fp32, or nullptr
+
+  const int* q2k_idx;              // [batch*num_heads*num_blocks, max_kv] int32
+  const int* q2k_num;              // [batch*num_heads*num_blocks] int32
+  const int* variable_block_sizes; // [num_blocks] int32, valid tokens per block
+
+  int batch;
+  int num_heads;
+  int seqlen;
+  int head_dim;
+  int num_blocks;
+  int max_kv;
+  float sm_scale;
+};
+
+// cudaSuccess iff this build can run `a`. Deliberately conservative: the caller is expected
+// to fall back to its own implementation rather than get a wrong answer.
+__host__ inline cudaError_t block_sparse_supported(const BlockSparseVsaArgs& a) {
+  if (a.head_dim != HEAD_DIM) return cudaErrorInvalidValue;      // compile-time in the kernel
+  if (a.num_blocks % 2 != 0) return cudaErrorInvalidValue;       // a CTA owns an adjacent pair
+  if (a.seqlen != a.num_blocks * BLOCK) return cudaErrorInvalidValue;
+  if (a.max_kv < 1 || a.num_blocks < 1) return cudaErrorInvalidValue;
+  if (a.q == nullptr || a.k == nullptr || a.o == nullptr) return cudaErrorInvalidValue;
+  if (a.q2k_idx == nullptr || a.q2k_num == nullptr) return cudaErrorInvalidValue;
+  // FastVideo always supplies this; without it padded keys would be attended as real zeros.
+  if (a.variable_block_sizes == nullptr) return cudaErrorInvalidValue;
+  // V is read MN-major at BOTH block sizes now, so no pre-transposed V_T is ever needed.
+  if (a.v == nullptr) return cudaErrorInvalidValue;
+  return cudaSuccess;
+}
+
+__host__ inline cudaError_t launch_block_sparse_sm100a(const BlockSparseVsaArgs& a,
+                                                           cudaStream_t stream) {
+  const cudaError_t sup = block_sparse_supported(a);
+  if (sup != cudaSuccess) return sup;
+
+  const int B = a.batch, H = a.num_heads, S = a.seqlen, hd = a.head_dim;
+  const int num_blocks = a.num_blocks, max_kv = a.max_kv;
+  const long tq = (long)B * S;
+  const int packed_mtiles_per_seq = num_blocks / 2;
+  const int total_work = B * H * packed_mtiles_per_seq;
+  constexpr bool BHSD = VSA_BHSD;
+
+  CUtensorMap tq_, tk_, tvt_, tv_, to_;
+  {
+    uint64_t gd[4] = { (uint64_t)SUB_COLS_BF16, BHSD ? (uint64_t)((long)B * H) : (uint64_t)H,
+                       BHSD ? (uint64_t)S : (uint64_t)tq, (uint64_t)Q_SUBTILES };
+    uint64_t gs[3] = { BHSD ? (uint64_t)((long)S * hd) * 2u : (uint64_t)hd * 2u,
+                       BHSD ? (uint64_t)hd * 2u : (uint64_t)((long)H * hd) * 2u,
+                       (uint64_t)SUB_COLS_BF16 * 2u };
+    uint32_t bd[4] = { (uint32_t)SUB_COLS_BF16, 1u, (uint32_t)M_TILE, (uint32_t)Q_SUBTILES };
+    uint32_t es[4] = { 1u, 1u, 1u, 1u };
+    if (cuTensorMapEncodeTiled(&tq_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4,
+                               const_cast<__nv_bfloat16*>(a.q), gd, gs, bd, es,
+                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
+                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
+      return cudaErrorInvalidValue;
+    if (cuTensorMapEncodeTiled(&to_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, a.o, gd, gs, bd, es,
+                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
+                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
+      return cudaErrorInvalidValue;
+  }
+  {
+    uint64_t gd[4] = { (uint64_t)SUB_COLS_BF16,
+                       BHSD ? (uint64_t)S : (uint64_t)tq,
+                       BHSD ? (uint64_t)(hd / SUB_COLS_BF16)
+                            : (uint64_t)((long)H * hd / SUB_COLS_BF16),
+                       (uint64_t)((long)B * H) };
+    uint64_t gs[3] = { BHSD ? (uint64_t)hd * 2u : (uint64_t)((long)H * hd) * 2u,
+                       (uint64_t)SUB_COLS_BF16 * 2u,
+                       (uint64_t)((long)S * hd) * 2u };
+    uint32_t bd[4] = { (uint32_t)SUB_COLS_BF16, (uint32_t)BLOCK,
+                       BLK128 ? (uint32_t)K_SUBTILES : 1u, 1u };
+    uint32_t es[4] = { 1u, 1u, 1u, 1u };
+    if (cuTensorMapEncodeTiled(&tk_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, BHSD ? 4 : 3,
+                               const_cast<__nv_bfloat16*>(a.k), gd, gs, bd, es,
+                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
+                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
+      return cudaErrorInvalidValue;
+    // V map is byte-for-byte the K map over a.v: MN-major V needs no transpose (blk128).
+    const __nv_bfloat16* vbase = a.v ? a.v : a.k;
+    if (cuTensorMapEncodeTiled(&tv_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, BHSD ? 4 : 3,
+                               const_cast<__nv_bfloat16*>(vbase), gd, gs, bd, es,
+                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
+                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
+      return cudaErrorInvalidValue;
+  }
+  // V_T map: blk64 only. Unused at blk128 but must still be a valid tensormap to pass by value.
+  {
+    const __nv_bfloat16* vt = a.v_t ? a.v_t : a.k;
+    if constexpr (BLK128) {
+      uint64_t gd[3] = { (uint64_t)SUB_COLS_BF16, (uint64_t)((long)H * hd),
+                         (uint64_t)((long)tq / SUB_COLS_BF16) };
+      uint64_t gs[2] = { (uint64_t)tq * 2u, (uint64_t)SUB_COLS_BF16 * 2u };
+      uint32_t bd[3] = { (uint32_t)SUB_COLS_BF16, (uint32_t)hd, (uint32_t)V_SUBTILES };
+      uint32_t es[3] = { 1u, 1u, 1u };
+      if (cuTensorMapEncodeTiled(&tvt_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3,
+                                 const_cast<__nv_bfloat16*>(vt), gd, gs, bd, es,
+                                 CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                                 CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
+                                 CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
+        return cudaErrorInvalidValue;
+    } else {
+      if (make_tma_2d_tiled(&tvt_, const_cast<__nv_bfloat16*>(vt), (long)H * hd, (int)tq, hd,
+                            SUB_COLS_BF16, 2, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
+                            CU_TENSOR_MAP_SWIZZLE_128B) != cudaSuccess)
+        return cudaErrorInvalidValue;
+    }
+  }
+
+  const size_t smem =
+        (size_t)2 * Q_TILE_BYTES + NUM_KV_STAGES * KV_RING_SLOT_BYTES
+      + (size_t)2 * M_TILE * HEAD_DIM * sizeof(__nv_bfloat16)
+      + (2 * NUM_KV_STAGES + 22) * 8
+      + (size_t)CLC_STAGES * (2 * 8 + 16) + 16
+      + 8
+      + (size_t)2 * STAT_REGIONS * STATS * sizeof(float)
+      + 256;
+
+#ifndef VSA_NAMED_BAR
+#define VSA_NAMED_BAR false
+#endif
+#ifndef VSA_THROTTLE
+#define VSA_THROTTLE false
+#endif
+#ifndef VSA_USE_CLC
+#define VSA_USE_CLC true
+#endif
+  constexpr bool FULL_NAMED_BAR = VSA_NAMED_BAR, EX2_EMU = true, SPLIT_P = true,
+                 SOFTMAX_THROTTLE = VSA_THROTTLE, USE_CLC = VSA_USE_CLC,
+                 Q_RASTER = true, MHA = true;
+  auto kfn = &fmha_context_bf16_gen_kernel<32, FULL_NAMED_BAR, EX2_EMU, SPLIT_P,
+                                           SOFTMAX_THROTTLE, USE_CLC, Q_RASTER, MHA,
+                                           /*RESCALE_THRESHOLD=*/8, /*BHSD=*/VSA_BHSD>;
+  cudaError_t e = cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
+  if (e != cudaSuccess) return e;
+
+  const unsigned long long magic0 = make_magic((unsigned)(H * packed_mtiles_per_seq));
+  const unsigned long long magic1 = make_magic((unsigned)H);
+  const unsigned long long magic2 = make_magic((unsigned)packed_mtiles_per_seq);
+  const float scale_log2 = a.sm_scale * (float)M_LOG2E;
+
+  int numSM = 0;
+  e = cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, 0);
+  if (e != cudaSuccess) return e;
+  const int num_ctas = USE_CLC ? total_work : (total_work < numSM ? total_work : numSM);
+  dim3 grid(num_ctas, 1, 1), block(N_WARPS * 32, 1, 1);
+
+  if (USE_CLC) {
+    cudaLaunchConfig_t cfg = {};
+    cfg.gridDim = grid; cfg.blockDim = block; cfg.dynamicSmemBytes = smem; cfg.stream = stream;
+    cudaLaunchAttribute cfgAttr[1];
+    cfgAttr[0].id = cudaLaunchAttributeClusterDimension;
+    cfgAttr[0].val.clusterDim.x = 1; cfgAttr[0].val.clusterDim.y = 1;
+    cfgAttr[0].val.clusterDim.z = 1;
+    cfg.attrs = cfgAttr; cfg.numAttrs = 1;
+    return cudaLaunchKernelEx(&cfg, kfn, tq_, tk_, tvt_, tv_, to_, S, H, scale_log2, B,
+                              num_blocks, packed_mtiles_per_seq, max_kv, magic0, magic1, magic2,
+                              a.q2k_idx, a.q2k_num, a.variable_block_sizes, a.lse);
+  }
+  kfn<<<grid, block, smem, stream>>>(tq_, tk_, tvt_, tv_, to_, S, H, scale_log2, B, num_blocks,
+                                     packed_mtiles_per_seq, max_kv, magic0, magic1, magic2,
+                                     a.q2k_idx, a.q2k_num, a.variable_block_sizes, a.lse);
+  return cudaGetLastError();
+}
+
+}  // namespace VSA_NAMESPACE
+
+// Callers (the bench, the torch binding) keep using unqualified names; each translation unit
+// only ever sees the one configuration its VSA_BLK128 selected.
+using namespace VSA_NAMESPACE;
+
+#endif  // BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH

fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu

@@ -0,0 +1,114 @@
+// block_sparse_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA forward.
+//
+// Forward only: returns (out, lse) so FastVideo's existing Triton backward keeps working
+// unchanged. lse is exactly the M tensor triton_block_sparse_attn_forward writes --
+// max(qk * qk_scale) + log2(l), [B, H, S] fp32 -- which is what lets
+// block_sparse_attn_backward_triton run against our forward untouched.
+//
+// The build is fixed at compile time by two flags, so one extension carries one configuration:
+//   VSA_BLK128  false -> 64-token sparse blocks, true -> 128-token
+//   VSA_BHSD    false -> [B, S, H, D], true -> [B, H, S, D]
+#include <torch/extension.h>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include "block_sparse_launch_sm100a.cuh"
+
+namespace {
+
+void check_qkv(const torch::Tensor& t, const char* name, int64_t B, int64_t H, int64_t S,
+               int64_t D) {
+  TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(t.scalar_type() == at::kBFloat16, name, " must be bfloat16, got ", t.scalar_type());
+  TORCH_CHECK(t.dim() == 4, name, " must be 4-D, got ", t.dim(), " dims");
+  TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+  if (VSA_BHSD) {
+    TORCH_CHECK(t.size(0) == B && t.size(1) == H && t.size(2) == S && t.size(3) == D, name,
+                " has shape ", t.sizes(), ", expected [", B, ",", H, ",", S, ",", D, "]");
+  } else {
+    TORCH_CHECK(t.size(0) == B && t.size(1) == S && t.size(2) == H && t.size(3) == D, name,
+                " has shape ", t.sizes(), ", expected [", B, ",", S, ",", H, ",", D, "]");
+  }
+}
+
+void check_index(const torch::Tensor& t, const char* name) {
+  TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(t.scalar_type() == at::kInt, name, " must be int32, got ", t.scalar_type());
+  TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+}
+
+}  // namespace
+
+// The exported symbol carries the block size: block_sparse_sm100a_fwd is the 64-token build,
+// block_sparse_sm100a_blk128_fwd the 128-token one (block_sparse_blk128_sm100a.cu re-includes
+// this file with VSA_BLK128 set). The python backend picks by the metadata's block size.
+#if VSA_BLK128
+#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_blk128_fwd
+#else
+#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_fwd
+#endif
+
+// Returns {out} or {out, lse}. Layout of out matches the inputs.
+std::vector<torch::Tensor> BLOCK_SPARSE_SM100A_FWD(torch::Tensor q, torch::Tensor k,
+                                                       torch::Tensor v,
+                                                       c10::optional<torch::Tensor> v_t,
+                                                       torch::Tensor q2k_idx,
+                                                       torch::Tensor q2k_num,
+                                                       torch::Tensor variable_block_sizes,
+                                                       double sm_scale, bool need_lse) {
+  const at::cuda::OptionalCUDAGuard guard(device_of(q));
+
+  const int64_t B = q.size(0);
+  const int64_t H = VSA_BHSD ? q.size(1) : q.size(2);
+  const int64_t S = VSA_BHSD ? q.size(2) : q.size(1);
+  const int64_t D = q.size(3);
+
+  check_qkv(q, "q", B, H, S, D);
+  check_qkv(k, "k", B, H, S, D);
+  check_qkv(v, "v", B, H, S, D);
+  check_index(q2k_idx, "q2k_idx");
+  check_index(q2k_num, "q2k_num");
+  check_index(variable_block_sizes, "variable_block_sizes");
+
+  const int64_t num_blocks = variable_block_sizes.numel();
+  const int64_t max_kv = q2k_idx.size(-1);
+  TORCH_CHECK(S == num_blocks * BLOCK, "seqlen ", S, " must equal num_blocks (", num_blocks,
+              ") * ", BLOCK, "; FastVideo pads the sequence up to whole blocks");
+
+  auto out = torch::empty_like(q);
+  torch::Tensor lse;
+  if (need_lse) lse = torch::empty({B, H, S}, q.options().dtype(torch::kFloat32));
+
+  BlockSparseVsaArgs a{};
+  a.q = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
+  a.k = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
+  a.v = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
+  a.v_t = v_t.has_value() ? reinterpret_cast<const __nv_bfloat16*>(v_t->data_ptr()) : nullptr;
+  a.o = reinterpret_cast<__nv_bfloat16*>(out.data_ptr());
+  a.lse = need_lse ? lse.data_ptr<float>() : nullptr;
+  a.q2k_idx = q2k_idx.data_ptr<int>();
+  a.q2k_num = q2k_num.data_ptr<int>();
+  a.variable_block_sizes = variable_block_sizes.data_ptr<int>();
+  a.batch = (int)B;
+  a.num_heads = (int)H;
+  a.seqlen = (int)S;
+  a.head_dim = (int)D;
+  a.num_blocks = (int)num_blocks;
+  a.max_kv = (int)max_kv;
+  a.sm_scale = (float)sm_scale;
+
+  // Report an unsupported regime loudly rather than returning plausible-looking wrong values.
+  TORCH_CHECK(block_sparse_supported(a) == cudaSuccess,
+              "block_sparse_sm100a: unsupported configuration -- requires head_dim==",
+              HEAD_DIM, ", an even num_blocks, seqlen == num_blocks*", BLOCK,
+              ", and a variable_block_sizes tensor. Got head_dim=", D, " num_blocks=",
+              num_blocks, " seqlen=", S);
+
+  const cudaError_t err = launch_block_sparse_sm100a(a, at::cuda::getCurrentCUDAStream());
+  TORCH_CHECK(err == cudaSuccess,
+              "block_sparse_sm100a launch failed: ", cudaGetErrorString(err));
+
+  if (need_lse) return {out, lse};
+  return {out};
+}

fastvideo-kernel/csrc/attention/primitives.cuh

@@ -0,0 +1,877 @@
+// primitives.cuh -- device primitives for the sm_100a VSA block-sparse attention
+// forward: tcgen05 (alloc / mma / ld / st / commit / wait / fence), TMA load / store /
+// tensormap, mbarrier, cluster launch control, setmaxnreg, fast math, and the FMHA helpers.
+//
+// Generated and pruned to what the kernel reaches -- do not edit by hand.
+#pragma once
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cmath>
+#include <cassert>
+#include <cstring>
+#include <vector_types.h>
+
+#ifndef CUDA_CHECK
+#define CUDA_CHECK(stmt) do {                                                 \
+    cudaError_t _e = (stmt);                                                  \
+    if (_e != cudaSuccess) {                                                  \
+      fprintf(stderr, "CUDA error %s:%d: %s -> %s\n",                         \
+              __FILE__, __LINE__, #stmt, cudaGetErrorString(_e));             \
+      std::exit(1);                                                           \
+    }                                                                         \
+  } while (0)
+#endif
+
+__device__ __forceinline__
+uint64_t mbarrier_arrive(uint32_t mbar_smem) {
+  uint64_t state;
+  asm volatile("mbarrier.arrive.shared::cta.b64 %0, [%1];\n"
+               : "=l"(state) : "r"(mbar_smem) : "memory");
+  return state;
+}
+
+__device__ __forceinline__
+void mbarrier_arrive_nostate(uint32_t mbar_smem) {
+  asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n"
+               :: "r"(mbar_smem) : "memory");
+}
+
+__device__ __forceinline__
+void mbarrier_arrive_cluster_default(uint32_t cluster_smem_addr) {
+  asm volatile("mbarrier.arrive.shared::cluster.b64 _, [%0];\n"
+               :: "r"(cluster_smem_addr) : "memory");
+}
+
+__device__ __forceinline__
+void mbarrier_arrive_expect_tx(uint32_t mbar_smem, uint32_t expected_bytes) {
+  asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;\n"
+               :: "r"(mbar_smem), "r"(expected_bytes) : "memory");
+}
+
+__device__ __forceinline__
+void mbarrier_wait_parity_suspend(uint32_t mbar_smem, uint32_t phase_parity) {
+  asm volatile(
+    "{\n"
+    ".reg .pred P1;\n"
+    "LAB_WAIT:\n"
+    "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, 10000000;\n"
+    "@!P1 bra.uni LAB_WAIT;\n"
+    "}\n"
+    :: "r"(mbar_smem), "r"(phase_parity) : "memory");
+}
+
+__device__ __forceinline__
+void mbarrier_wait_parity(uint32_t mbar_smem, uint32_t phase_parity) {
+  asm volatile(
+    "{\n"
+    ".reg .pred P1;\n"
+    "LAB_WAIT_HOT:\n"
+    "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
+    "@!P1 bra.uni LAB_WAIT_HOT;\n"
+    "}\n"
+    :: "r"(mbar_smem), "r"(phase_parity) : "memory");
+}
+
+__device__ __forceinline__
+void fence_proxy_async_shared_cta() {
+  asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
+}
+
+__device__ __forceinline__
+void fence_proxy_async_shared() {
+  asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void clc_try_cancel_async(
+    uint32_t smem_dst, uint32_t mbar_smem) {
+  asm volatile(
+    "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.b128"
+    " [%0], [%1];\n"
+    :: "r"(smem_dst), "r"(mbar_smem) : "memory");
+}
+
+__device__ __forceinline__ void clc_load_response(
+    uint32_t smem_slot, uint32_t& r0, uint32_t& r1,
+    uint32_t& r2, uint32_t& r3) {
+  asm volatile("ld.shared::cta.v4.b32 {%0, %1, %2, %3}, [%4];\n"
+               : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
+               : "r"(smem_slot));
+}
+
+template <int NUM_STAGES>
+struct MbarrierPhaseTracker {
+
+  uint32_t phase[NUM_STAGES];
+  int idx;
+
+  __device__ __forceinline__
+  void init() {
+    for (int i = 0; i < NUM_STAGES; ++i) phase[i] = 0;
+    idx = 0;
+  }
+
+  __device__ __forceinline__
+  uint32_t current_phase() const { return phase[idx]; }
+
+  __device__ __forceinline__
+  void advance() {
+    phase[idx] ^= 1u;
+    idx = (idx + 1) % NUM_STAGES;
+  }
+
+  __device__ __forceinline__
+  int stage() const { return idx; }
+};
+
+template <int NUM_STAGES>
+struct PhaseTracker {
+  int stage;
+  uint32_t phase;
+
+  __device__ __forceinline__
+  PhaseTracker() : stage(0), phase(0) {}
+
+  __device__ __forceinline__
+  void advance() {
+    stage++;
+    if (stage == NUM_STAGES) {
+      stage = 0;
+      phase ^= 1;
+    }
+  }
+
+  __device__ __forceinline__
+  int get_stage() const { return stage; }
+
+  __device__ __forceinline__
+  uint32_t get_phase() const { return phase; }
+};
+
+template <int NUM_STAGES>
+struct EmptyPhaseTracker {
+  int stage;
+  uint32_t phase;
+
+  __device__ __forceinline__
+  EmptyPhaseTracker() : stage(0), phase(1) {}
+
+  __device__ __forceinline__
+  void advance() {
+    stage++;
+    if (stage == NUM_STAGES) {
+      stage = 0;
+      phase ^= 1;
+    }
+  }
+
+  __device__ __forceinline__
+  int get_stage() const { return stage; }
+
+  __device__ __forceinline__
+  uint32_t get_phase() const { return phase; }
+};
+
+template <int STAGES>
+__device__ __forceinline__
+void advance_stage_phase(int& stage, uint32_t& phase) {
+  ++stage;
+  if (stage == STAGES) {
+    stage = 0;
+    phase ^= 1u;
+  }
+}
+
+static constexpr uint32_t SM100_CLC_PEER_MASK = 0xFEFFFFFF;
+
+struct ClcTileInfo {
+  int  m_tile;
+  int  n_tile;
+  bool valid;
+};
+
+enum class ClcRasterOrder { AlongN, AlongM };
+
+__device__ __forceinline__
+void clc_arrive_expect_tx_cta(uint32_t clc_full_local_addr, uint32_t tx_bytes) {
+  if ((threadIdx.x & 31) == 0) {
+    mbarrier_arrive_expect_tx(clc_full_local_addr, tx_bytes);
+  }
+}
+
+__device__ __forceinline__
+void clc_consumer_release(uint32_t clc_empty_local_addr) {
+  uint32_t peer0_addr = clc_empty_local_addr & SM100_CLC_PEER_MASK;
+  mbarrier_arrive_cluster_default(peer0_addr);
+}
+
+__device__ __forceinline__
+void clc_consumer_release_cta(uint32_t clc_empty_local_addr) {
+  mbarrier_arrive_nostate(clc_empty_local_addr);
+}
+
+template <int CLUSTER_SHAPE_M, int CLUSTER_SHAPE_N, ClcRasterOrder ORDER>
+__device__ __forceinline__
+ClcTileInfo clc_parse_response(uint32_t resp_smem_addr) {
+  uint32_t d0, d1, d2, d3;
+  fence_proxy_async_shared_cta();
+  clc_load_response(resp_smem_addr, d0, d1, d2, d3);
+  const int  ctaid_x = static_cast<int>(d0);
+  const int  ctaid_y = static_cast<int>(d1 & 0xFFFFu);
+  const bool valid   = (d2 & 1u) != 0u;
+  (void)d3;
+
+  ClcTileInfo info;
+  info.valid = valid;
+  if constexpr (ORDER == ClcRasterOrder::AlongN) {
+    info.m_tile = ctaid_y / CLUSTER_SHAPE_M;
+    info.n_tile = ctaid_x / CLUSTER_SHAPE_N;
+  } else {
+    info.m_tile = ctaid_x / CLUSTER_SHAPE_M;
+    info.n_tile = ctaid_y / CLUSTER_SHAPE_N;
+  }
+  return info;
+}
+
+template <int CLUSTER_SHAPE_M, int CLUSTER_SHAPE_N, ClcRasterOrder ORDER,
+          int CTA_GROUP = 2, bool SUSPEND = false>
+__device__ __forceinline__
+ClcTileInfo clc_fetch_next_tile(
+    uint64_t* clc_full_bar, uint64_t* clc_empty_bar, uint32_t* clc_response,
+    int clc_cons_stage, uint32_t clc_cons_phase, bool do_release) {
+  uint32_t full_addr = static_cast<uint32_t>(
+      __cvta_generic_to_shared(&clc_full_bar[clc_cons_stage]));
+  if constexpr (SUSPEND) mbarrier_wait_parity_suspend(full_addr, clc_cons_phase);
+  else                   mbarrier_wait_parity(full_addr, clc_cons_phase);
+  uint32_t resp_addr = static_cast<uint32_t>(
+      __cvta_generic_to_shared(&clc_response[clc_cons_stage * 4]));
+  ClcTileInfo t = clc_parse_response<
+      CLUSTER_SHAPE_M, CLUSTER_SHAPE_N, ORDER>(resp_addr);
+  if (do_release) {
+    uint32_t empty_local = static_cast<uint32_t>(
+        __cvta_generic_to_shared(&clc_empty_bar[clc_cons_stage]));
+    if constexpr (CTA_GROUP == 1) {
+      clc_consumer_release_cta(empty_local);
+    } else {
+      clc_consumer_release(empty_local);
+    }
+  }
+  return t;
+}
+
+template <int STAGES = 2>
+__device__ __forceinline__
+void clc_fetch_next_tile_advance(int& clc_cons_stage,
+                                 uint32_t& clc_cons_phase) {
+  advance_stage_phase<STAGES>(clc_cons_stage, clc_cons_phase);
+}
+
+__device__ __forceinline__ unsigned fdiv(unsigned n, unsigned long long pk) {
+  unsigned M = (unsigned)pk;
+  if (M == 0u) return n;
+  return __umulhi(n, M) >> (unsigned)(pk >> 32);
+}
+__host__ inline unsigned long long make_magic(unsigned d) {
+  if (d <= 1u) return 0ULL;
+  unsigned l = 0; while ((1u << (l + 1)) <= d) ++l;
+  unsigned p = 31u + l;
+  unsigned long long m = ((1ull << p) + (unsigned long long)d - 1ull) / d;
+  return (m & 0xffffffffULL) | ((unsigned long long)(p - 32u) << 32);
+}
+
+template <int BARRIER_ID>
+__device__ __forceinline__
+void bar_sync(uint32_t thread_count) {
+  static_assert(BARRIER_ID >= 0 && BARRIER_ID <= 15,
+                "bar.sync: BARRIER_ID must be in [0, 15]");
+  asm volatile("bar.sync %0, %1;\n"
+               :: "n"(BARRIER_ID), "r"(thread_count) : "memory");
+}
+
+template <int BARRIER_ID>
+__device__ __forceinline__
+void bar_arrive(uint32_t thread_count) {
+  static_assert(BARRIER_ID >= 0 && BARRIER_ID <= 15,
+                "bar.arrive: BARRIER_ID must be in [0, 15]");
+  asm volatile("bar.arrive %0, %1;\n"
+               :: "n"(BARRIER_ID), "r"(thread_count) : "memory");
+}
+
+__device__ __forceinline__ void full_bar_arrive(int m_tile, int band) {
+  switch (1 + m_tile * 4 + band) {
+    case 1: bar_arrive<1>(64); break;
+    case 2: bar_arrive<2>(64); break;
+    case 3: bar_arrive<3>(64); break;
+    case 4: bar_arrive<4>(64); break;
+    case 5: bar_arrive<5>(64); break;
+    case 6: bar_arrive<6>(64); break;
+    case 7: bar_arrive<7>(64); break;
+    case 8: bar_arrive<8>(64); break;
+  }
+}
+__device__ __forceinline__ void full_bar_wait(int m_tile, int band) {
+  switch (1 + m_tile * 4 + band) {
+    case 1: bar_sync<1>(64); break;
+    case 2: bar_sync<2>(64); break;
+    case 3: bar_sync<3>(64); break;
+    case 4: bar_sync<4>(64); break;
+    case 5: bar_sync<5>(64); break;
+    case 6: bar_sync<6>(64); break;
+    case 7: bar_sync<7>(64); break;
+    case 8: bar_sync<8>(64); break;
+  }
+}
+
+template <bool IS_CAUSAL, int K_TILE>
+__device__ __forceinline__ void mask_s_row_r2p(float* scores, int k_offset, int q_pos, int seqlen_k) {
+  int n_keep = seqlen_k - k_offset;
+  if constexpr (IS_CAUSAL) {
+    const int causal = q_pos - k_offset + 1;
+    n_keep = n_keep < causal ? n_keep : causal;
+  }
+  #pragma unroll
+  for (int s = 0; s < K_TILE / 32; ++s) {
+    int m = (s + 1) * 32 - n_keep;
+    m = m < 0 ? 0 : (m > 32 ? 32 : m);
+    const uint32_t keep = (m >= 32) ? 0u : (0xFFFFFFFFu >> m);
+    #pragma unroll
+    for (int i = 0; i < 32; ++i)
+      if (!(keep & (1u << i))) scores[s * 32 + i] = -INFINITY;
+  }
+}
+
+template <int CTA_GROUP>
+__device__ __forceinline__ void tcgen05_alloc(uint32_t smem_dst_ptr,
+                                              uint32_t n_cols) {
+  static_assert(CTA_GROUP == 1 || CTA_GROUP == 2,
+                "tcgen05_alloc: CTA_GROUP must be 1 or 2");
+  if constexpr (CTA_GROUP == 1) {
+    asm volatile(
+      "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n"
+      :: "r"(smem_dst_ptr), "r"(n_cols));
+  } else {
+    asm volatile(
+      "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;\n"
+      :: "r"(smem_dst_ptr), "r"(n_cols));
+  }
+}
+
+__device__ __forceinline__ void tcgen05_st_32x32b_x16(
+    uint32_t tmem_addr, const uint32_t (&r)[16]) {
+  asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 "
+               "[%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,"
+               "%11,%12,%13,%14,%15,%16};\n"
+               :: "r"(tmem_addr),
+                  "r"(r[0]),"r"(r[1]),"r"(r[2]),"r"(r[3]),
+                  "r"(r[4]),"r"(r[5]),"r"(r[6]),"r"(r[7]),
+                  "r"(r[8]),"r"(r[9]),"r"(r[10]),"r"(r[11]),
+                  "r"(r[12]),"r"(r[13]),"r"(r[14]),"r"(r[15]));
+}
+
+__device__ __forceinline__ void tcgen05_st_32x32b_x32(
+    uint32_t tmem_addr, const uint32_t (&r)[32]) {
+  asm volatile("tcgen05.st.sync.aligned.32x32b.x32.b32 "
+               "[%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,"
+               "%11,%12,%13,%14,%15,%16,%17,%18,%19,%20,"
+               "%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,"
+               "%31,%32};\n"
+               :: "r"(tmem_addr),
+                  "r"(r[0]),"r"(r[1]),"r"(r[2]),"r"(r[3]),
+                  "r"(r[4]),"r"(r[5]),"r"(r[6]),"r"(r[7]),
+                  "r"(r[8]),"r"(r[9]),"r"(r[10]),"r"(r[11]),
+                  "r"(r[12]),"r"(r[13]),"r"(r[14]),"r"(r[15]),
+                  "r"(r[16]),"r"(r[17]),"r"(r[18]),"r"(r[19]),
+                  "r"(r[20]),"r"(r[21]),"r"(r[22]),"r"(r[23]),
+                  "r"(r[24]),"r"(r[25]),"r"(r[26]),"r"(r[27]),
+                  "r"(r[28]),"r"(r[29]),"r"(r[30]),"r"(r[31]));
+}
+
+__device__ __forceinline__ void tcgen05_commit1_lead(uint32_t lead, uint32_t mbar_smem_addr) {
+  asm volatile(
+    "{\n\t"
+    ".reg .pred q;\n\t"
+    "setp.ne.b32 q, %0, 0;\n\t"
+    "@q tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%1];\n\t"
+    "}\n"
+    :: "r"(lead), "r"(mbar_smem_addr));
+}
+
+__device__ __forceinline__ void tcgen05_wait_st() {
+  asm volatile("tcgen05.wait::st.sync.aligned;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void tcgen05_fence_before_thread_sync() {
+  asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory");
+}
+
+__device__ __forceinline__
+void tma_load_2d(uint32_t smem_dst, const void* tensormap_ptr,
+                 uint32_t mbar_smem, int coord_x, int coord_y) {
+  asm volatile(
+    "cp.async.bulk.tensor.2d.shared::cluster.global.tile"
+    ".mbarrier::complete_tx::bytes"
+    " [%0], [%1, {%3, %4}], [%2];\n"
+    :: "r"(smem_dst), "l"(tensormap_ptr),
+       "r"(mbar_smem), "r"(coord_x), "r"(coord_y)
+    : "memory");
+}
+
+__device__ __forceinline__
+void tma_load_3d(uint32_t smem_dst, const void* tensormap_ptr,
+                 uint32_t mbar_smem, int c0, int c1, int c2) {
+  asm volatile(
+    "cp.async.bulk.tensor.3d.shared::cluster.global.tile"
+    ".mbarrier::complete_tx::bytes"
+    " [%0], [%1, {%3, %4, %5}], [%2];\n"
+    :: "r"(smem_dst), "l"(tensormap_ptr),
+       "r"(mbar_smem), "r"(c0), "r"(c1), "r"(c2)
+    : "memory");
+}
+
+__device__ __forceinline__
+void tma_load_4d(uint32_t smem_dst, const void* tensormap_ptr,
+                 uint32_t mbar_smem, int c0, int c1, int c2, int c3) {
+  asm volatile(
+    "cp.async.bulk.tensor.4d.shared::cluster.global.tile"
+    ".mbarrier::complete_tx::bytes"
+    " [%0], [%1, {%3, %4, %5, %6}], [%2];\n"
+    :: "r"(smem_dst), "l"(tensormap_ptr),
+       "r"(mbar_smem), "r"(c0), "r"(c1), "r"(c2), "r"(c3)
+    : "memory");
+}
+
+template <int CTA_GROUP>
+__device__ __forceinline__ void tcgen05_dealloc(uint32_t tmem_addr,
+                                                uint32_t n_cols) {
+  static_assert(CTA_GROUP == 1 || CTA_GROUP == 2,
+                "tcgen05_dealloc: CTA_GROUP must be 1 or 2");
+  if constexpr (CTA_GROUP == 1) {
+    asm volatile(
+      "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n"
+      :: "r"(tmem_addr), "r"(n_cols));
+  } else {
+    asm volatile(
+      "tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;\n"
+      :: "r"(tmem_addr), "r"(n_cols));
+  }
+}
+
+__device__ __forceinline__
+void tma_store_2d(const void* tensormap_ptr, int coord_x, int coord_y,
+                  uint32_t smem_src) {
+  asm volatile(
+    "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group"
+    " [%0, {%1, %2}], [%3];\n"
+    :: "l"(tensormap_ptr), "r"(coord_x), "r"(coord_y),
+       "r"(smem_src)
+    : "memory");
+}
+
+__device__ __forceinline__
+void tma_store_3d(const void* tensormap_ptr, int c0, int c1, int c2,
+                  uint32_t smem_src) {
+  asm volatile(
+    "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group"
+    " [%0, {%1, %2, %3}], [%4];\n"
+    :: "l"(tensormap_ptr), "r"(c0), "r"(c1), "r"(c2), "r"(smem_src)
+    : "memory");
+}
+
+__device__ __forceinline__
+void tma_store_4d(const void* tensormap_ptr, int c0, int c1, int c2, int c3,
+                  uint32_t smem_src) {
+  asm volatile(
+    "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group"
+    " [%0, {%1, %2, %3, %4}], [%5];\n"
+    :: "l"(tensormap_ptr), "r"(c0), "r"(c1), "r"(c2), "r"(c3), "r"(smem_src)
+    : "memory");
+}
+
+inline cudaError_t make_tma_2d_tiled(
+    CUtensorMap* out,
+    const void* ptr, int rows, int cols, int box_rows, int box_cols,
+    int elem_bytes, CUtensorMapDataType dtype,
+    CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_128B,
+    CUtensorMapL2promotion l2 = CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
+    CUtensorMapFloatOOBfill oob = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) {
+  uint64_t globalDim[2]     = { (uint64_t)cols, (uint64_t)rows };
+  uint64_t globalStrides[1] = { (uint64_t)cols * (uint64_t)elem_bytes };
+  uint32_t boxDim[2]        = { (uint32_t)box_cols, (uint32_t)box_rows };
+  uint32_t elemStrides[2]   = { 1u, 1u };
+
+  CUresult r = cuTensorMapEncodeTiled(
+      out, dtype, 2,
+      const_cast<void*>(ptr), globalDim, globalStrides,
+      boxDim, elemStrides,
+      CU_TENSOR_MAP_INTERLEAVE_NONE,
+      swizzle, l2, oob);
+  return (r == CUDA_SUCCESS) ? cudaSuccess : cudaErrorInvalidValue;
+}
+
+__device__ __forceinline__
+void cp_async_bulk_commit_group() {
+  asm volatile("cp.async.bulk.commit_group;\n" ::: "memory");
+}
+
+template <int N>
+__device__ __forceinline__
+void cp_async_bulk_wait_group_read() {
+  asm volatile("cp.async.bulk.wait_group.read %0;\n" :: "n"(N) : "memory");
+}
+
+__device__ __forceinline__
+void mbarrier_init(uint32_t mbar_smem, uint32_t arrive_count) {
+  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n"
+               :: "r"(mbar_smem), "r"(arrive_count) : "memory");
+}
+
+template <int CTA_GROUP>
+__device__ __forceinline__ void tcgen05_relinquish_alloc_permit() {
+  static_assert(CTA_GROUP == 1 || CTA_GROUP == 2,
+                "tcgen05_relinquish_alloc_permit: CTA_GROUP must be 1 or 2");
+  if constexpr (CTA_GROUP == 1) {
+    asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::);
+  } else {
+    asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;\n" ::);
+  }
+}
+
+__device__ __forceinline__
+void fence_mbarrier_init_release_cluster() {
+  asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory");
+}
+
+__device__ __forceinline__ void tcgen05_mma_f16_ss_lead(uint32_t lead,
+    uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b, uint32_t idesc,
+    bool enable_input_d) {
+  asm volatile(
+    "{\n\t"
+    ".reg .pred p, q;\n\t"
+    "setp.ne.b32 q, %0, 0;\n\t"
+    "setp.ne.b32 p, %5, 0;\n\t"
+    "@q tcgen05.mma.cta_group::1.kind::f16 [%1], %2, %3, %4, {%6, %7, %8, %9}, p;\n\t"
+    "}\n"
+    :: "r"(lead), "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(idesc),
+       "r"(enable_input_d ? 1u : 0u), "r"(0u), "r"(0u), "r"(0u), "r"(0u));
+}
+
+__device__ __forceinline__ void tcgen05_mma_f16_ts_1sm_lead(uint32_t lead,
+    uint32_t tmem_c, uint32_t tmem_a, uint64_t desc_b, uint32_t idesc,
+    bool enable_input_d) {
+  asm volatile(
+    "{\n\t"
+    ".reg .pred p, q;\n\t"
+    "setp.ne.b32 q, %0, 0;\n\t"
+    "setp.ne.b32 p, %5, 0;\n\t"
+    "@q tcgen05.mma.cta_group::1.kind::f16 [%1], [%2], %3, %4, {%6, %7, %8, %9}, p;\n\t"
+    "}\n"
+    :: "r"(lead), "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(idesc),
+       "r"(enable_input_d ? 1u : 0u), "r"(0u), "r"(0u), "r"(0u), "r"(0u));
+}
+
+enum class SmemSwizzleBlackwell : uint32_t {
+  None = 0,
+  B128_32atom = 1,
+  B128 = 2,
+  B64  = 4,
+  B32  = 6,
+};
+
+__device__ __host__ __forceinline__ uint64_t build_smem_desc_blackwell(
+    uint32_t smem_addr,
+    uint32_t stride_byte_offset,
+    uint32_t leading_byte_offset,
+    SmemSwizzleBlackwell swizzle = SmemSwizzleBlackwell::B128,
+    uint32_t base_offset = 0) {
+  uint64_t d = 0;
+  d |= static_cast<uint64_t>((smem_addr >> 4) & 0x3FFF);
+  d |= static_cast<uint64_t>((leading_byte_offset >> 4) & 0x3FFF) << 16;
+  d |= static_cast<uint64_t>((stride_byte_offset >> 4) & 0x3FFF) << 32;
+  d |= static_cast<uint64_t>(1) << 46;
+  d |= (static_cast<uint64_t>(base_offset) & 0x7) << 49;
+  d |= static_cast<uint64_t>(static_cast<uint32_t>(swizzle) & 0x7) << 61;
+  return d;
+}
+
+__device__ __forceinline__
+uint32_t elect_one_sync() {
+  uint32_t elected;
+  asm volatile(
+    "{\n\t"
+    ".reg .pred p;\n\t"
+    "elect.sync %0|p, 0xffffffff;\n\t"
+    "selp.b32 %0, 1, 0, p;\n\t"
+    "}\n"
+    : "=r"(elected));
+  return elected;
+}
+
+__device__ __forceinline__
+uint32_t elect_one_sync(uint32_t membermask) {
+  uint32_t elected;
+  asm volatile(
+    "{\n\t"
+    ".reg .pred p;\n\t"
+    "elect.sync %0|p, %1;\n\t"
+    "selp.b32 %0, 1, 0, p;\n\t"
+    "}\n"
+    : "=r"(elected) : "r"(membermask));
+  return elected;
+}
+
+template <int N>
+__device__ __forceinline__
+void setmaxnreg_dec() {
+  static_assert(N >= 24 && N <= 256 && N % 8 == 0,
+                "setmaxnreg_dec: N must be in [24, 256] and a multiple of 8");
+  asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n"
+               :: "n"(N) : "memory");
+}
+
+template <int N>
+__device__ __forceinline__
+void setmaxnreg_inc() {
+  static_assert(N >= 24 && N <= 256 && N % 8 == 0,
+                "setmaxnreg_inc: N must be in [24, 256] and a multiple of 8");
+  asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n"
+               :: "n"(N) : "memory");
+}
+
+__device__ __forceinline__
+uint32_t cvt_f32x2_to_bf16x2(float a, float b) {
+  uint32_t r;
+  asm volatile("cvt.rn.bf16x2.f32 %0, %2, %1;\n"
+               : "=r"(r) : "f"(a), "f"(b));
+  return r;
+}
+
+namespace {
+__device__ __forceinline__ uint64_t f32x2_bits(float2 v) {
+  uint64_t b; __builtin_memcpy(&b, &v, 8); return b;
+}
+__device__ __forceinline__ float2 f32x2_make(uint64_t b) {
+  float2 v; __builtin_memcpy(&v, &b, 8); return v;
+}
+}
+
+__device__ __forceinline__ float2 fmul2(float2 a, float2 b) {
+  uint64_t d;
+  asm volatile("mul.f32x2 %0, %1, %2;\n" : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)));
+  return f32x2_make(d);
+}
+
+__device__ __forceinline__ float2 fadd2(float2 a, float2 b) {
+  uint64_t d;
+  asm volatile("add.f32x2 %0, %1, %2;\n" : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)));
+  return f32x2_make(d);
+}
+
+__device__ __forceinline__ float2 ffma2(float2 a, float2 b, float2 c) {
+  uint64_t d;
+  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
+               : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)), "l"(f32x2_bits(c)));
+  return f32x2_make(d);
+}
+
+__device__ __forceinline__ float2 f32x2_splat(float s) { return make_float2(s, s); }
+
+__device__ __forceinline__ float ex2_approx_f32(float z) {
+  float d;
+  asm volatile("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(d) : "f"(z));
+  return d;
+}
+
+__device__ __forceinline__ float2 ex2_emu_f32x2(float x, float y) {
+  uint32_t ox, oy;
+  asm volatile(
+    "{\n\t"
+    ".reg .f32 f1,f2,f3,f4,f5,f6,f7;\n\t"
+    ".reg .b64 l1,l2,l3,l4,l5,l6,l7,l8,l9,l10;\n\t"
+    ".reg .s32 r1,r2,r3,r4,r5,r6,r7,r8;\n\t"
+    "max.f32 f1, %2, 0fC2FE0000;\n\t"
+    "max.f32 f2, %3, 0fC2FE0000;\n\t"
+    "mov.b64 l1, {f1, f2};\n\t"
+    "mov.f32 f3, 0f4B400000;\n\t"
+    "mov.b64 l2, {f3, f3};\n\t"
+    "add.rm.f32x2 l7, l1, l2;\n\t"
+    "sub.rn.f32x2 l8, l7, l2;\n\t"
+    "sub.rn.f32x2 l9, l1, l8;\n\t"
+    "mov.f32 f7, 0f3D9DF09D;\n\t"
+    "mov.b64 l6, {f7, f7};\n\t"
+    "mov.f32 f6, 0f3E6906A4;\n\t"
+    "mov.b64 l5, {f6, f6};\n\t"
+    "mov.f32 f5, 0f3F31F519;\n\t"
+    "mov.b64 l4, {f5, f5};\n\t"
+    "mov.f32 f4, 0f3F800000;\n\t"
+    "mov.b64 l3, {f4, f4};\n\t"
+    "fma.rn.f32x2 l10, l9, l6, l5;\n\t"
+    "fma.rn.f32x2 l10, l10, l9, l4;\n\t"
+    "fma.rn.f32x2 l10, l10, l9, l3;\n\t"
+    "mov.b64 {r1, r2}, l7;\n\t"
+    "mov.b64 {r3, r4}, l10;\n\t"
+    "shl.b32 r5, r1, 23;\n\t"
+    "add.s32 r7, r5, r3;\n\t"
+    "shl.b32 r6, r2, 23;\n\t"
+    "add.s32 r8, r6, r4;\n\t"
+    "mov.b32 %0, r7;\n\t"
+    "mov.b32 %1, r8;\n\t"
+    "}\n"
+    : "=r"(ox), "=r"(oy) : "f"(x), "f"(y));
+  float2 r; __builtin_memcpy(&r.x, &ox, 4); __builtin_memcpy(&r.y, &oy, 4);
+  return r;
+}
+
+__device__ __forceinline__ float rcp_approx_ftz_f32(float x) {
+  float d;
+  asm("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(d) : "f"(x));
+  return d;
+}
+
+__device__ __forceinline__ uint32_t make_idesc_table44(
+    int M, int N,
+    uint32_t dtype, uint32_t atype, uint32_t btype,
+    bool transpose_a = false, bool transpose_b = false,
+    bool negate_a = false, bool negate_b = false) {
+  uint32_t idesc = 0;
+  idesc |= (dtype & 0x3) << 4;
+  idesc |= (atype & 0x7) << 7;
+  idesc |= (btype & 0x7) << 10;
+  idesc |= (negate_a ? 1u : 0u) << 13;
+  idesc |= (negate_b ? 1u : 0u) << 14;
+  idesc |= (transpose_a ? 1u : 0u) << 15;
+  idesc |= (transpose_b ? 1u : 0u) << 16;
+  idesc |= ((static_cast<uint32_t>(N) >> 3) & 0x3F) << 17;
+  idesc |= ((static_cast<uint32_t>(M) >> 4) & 0x1F) << 24;
+  return idesc;
+}
+
+__device__ __forceinline__ uint32_t make_idesc_bf16_f32(
+    int M, int N, bool ta = false, bool tb = false) {
+  return make_idesc_table44(M, N,  1,
+                              1,  1, ta, tb);
+}
+
+__device__ __forceinline__ void tcgen05_ld_32x32b_x16(
+    uint32_t tmem_addr, uint32_t (&r)[16]) {
+  asm volatile("tcgen05.ld.sync.aligned.32x32b.x16.b32 "
+               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
+               "%10,%11,%12,%13,%14,%15}, [%16];\n"
+               : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]),
+                 "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]),
+                 "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]),
+                 "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15])
+               : "r"(tmem_addr));
+}
+
+__device__ __forceinline__ void tcgen05_ld_32x32b_x32(
+    uint32_t tmem_addr, uint32_t (&r)[32]) {
+  asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 "
+               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
+               "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,"
+               "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,"
+               "%30,%31}, [%32];\n"
+               : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]),
+                 "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]),
+                 "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]),
+                 "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]),
+                 "=r"(r[16]),"=r"(r[17]),"=r"(r[18]),"=r"(r[19]),
+                 "=r"(r[20]),"=r"(r[21]),"=r"(r[22]),"=r"(r[23]),
+                 "=r"(r[24]),"=r"(r[25]),"=r"(r[26]),"=r"(r[27]),
+                 "=r"(r[28]),"=r"(r[29]),"=r"(r[30]),"=r"(r[31])
+               : "r"(tmem_addr));
+}
+
+__device__ __forceinline__ void tcgen05_ld_32x32b_x64(
+    uint32_t tmem_addr, uint32_t (&r)[64]) {
+  asm volatile("tcgen05.ld.sync.aligned.32x32b.x64.b32 "
+               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
+               "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,"
+               "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,"
+               "%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,"
+               "%40,%41,%42,%43,%44,%45,%46,%47,%48,%49,"
+               "%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,"
+               "%60,%61,%62,%63}, [%64];\n"
+               : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]),
+                 "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]),
+                 "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]),
+                 "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]),
+                 "=r"(r[16]),"=r"(r[17]),"=r"(r[18]),"=r"(r[19]),
+                 "=r"(r[20]),"=r"(r[21]),"=r"(r[22]),"=r"(r[23]),
+                 "=r"(r[24]),"=r"(r[25]),"=r"(r[26]),"=r"(r[27]),
+                 "=r"(r[28]),"=r"(r[29]),"=r"(r[30]),"=r"(r[31]),
+                 "=r"(r[32]),"=r"(r[33]),"=r"(r[34]),"=r"(r[35]),
+                 "=r"(r[36]),"=r"(r[37]),"=r"(r[38]),"=r"(r[39]),
+                 "=r"(r[40]),"=r"(r[41]),"=r"(r[42]),"=r"(r[43]),
+                 "=r"(r[44]),"=r"(r[45]),"=r"(r[46]),"=r"(r[47]),
+                 "=r"(r[48]),"=r"(r[49]),"=r"(r[50]),"=r"(r[51]),
+                 "=r"(r[52]),"=r"(r[53]),"=r"(r[54]),"=r"(r[55]),
+                 "=r"(r[56]),"=r"(r[57]),"=r"(r[58]),"=r"(r[59]),
+                 "=r"(r[60]),"=r"(r[61]),"=r"(r[62]),"=r"(r[63])
+               : "r"(tmem_addr));
+}
+
+__device__ __forceinline__ void tcgen05_ld_32x32b_x128(
+    uint32_t tmem_addr, uint32_t (&r)[128]) {
+  asm volatile("tcgen05.ld.sync.aligned.32x32b.x128.b32 "
+               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
+               "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,"
+               "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,"
+               "%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,"
+               "%40,%41,%42,%43,%44,%45,%46,%47,%48,%49,"
+               "%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,"
+               "%60,%61,%62,%63,%64,%65,%66,%67,%68,%69,"
+               "%70,%71,%72,%73,%74,%75,%76,%77,%78,%79,"
+               "%80,%81,%82,%83,%84,%85,%86,%87,%88,%89,"
+               "%90,%91,%92,%93,%94,%95,%96,%97,%98,%99,"
+               "%100,%101,%102,%103,%104,%105,%106,%107,%108,%109,"
+               "%110,%111,%112,%113,%114,%115,%116,%117,%118,%119,"
+               "%120,%121,%122,%123,%124,%125,%126,%127}, [%128];\n"
+               : "=r"(r[  0]),"=r"(r[  1]),"=r"(r[  2]),"=r"(r[  3]),
+                 "=r"(r[  4]),"=r"(r[  5]),"=r"(r[  6]),"=r"(r[  7]),
+                 "=r"(r[  8]),"=r"(r[  9]),"=r"(r[ 10]),"=r"(r[ 11]),
+                 "=r"(r[ 12]),"=r"(r[ 13]),"=r"(r[ 14]),"=r"(r[ 15]),
+                 "=r"(r[ 16]),"=r"(r[ 17]),"=r"(r[ 18]),"=r"(r[ 19]),
+                 "=r"(r[ 20]),"=r"(r[ 21]),"=r"(r[ 22]),"=r"(r[ 23]),
+                 "=r"(r[ 24]),"=r"(r[ 25]),"=r"(r[ 26]),"=r"(r[ 27]),
+                 "=r"(r[ 28]),"=r"(r[ 29]),"=r"(r[ 30]),"=r"(r[ 31]),
+                 "=r"(r[ 32]),"=r"(r[ 33]),"=r"(r[ 34]),"=r"(r[ 35]),
+                 "=r"(r[ 36]),"=r"(r[ 37]),"=r"(r[ 38]),"=r"(r[ 39]),
+                 "=r"(r[ 40]),"=r"(r[ 41]),"=r"(r[ 42]),"=r"(r[ 43]),
+                 "=r"(r[ 44]),"=r"(r[ 45]),"=r"(r[ 46]),"=r"(r[ 47]),
+                 "=r"(r[ 48]),"=r"(r[ 49]),"=r"(r[ 50]),"=r"(r[ 51]),
+                 "=r"(r[ 52]),"=r"(r[ 53]),"=r"(r[ 54]),"=r"(r[ 55]),
+                 "=r"(r[ 56]),"=r"(r[ 57]),"=r"(r[ 58]),"=r"(r[ 59]),
+                 "=r"(r[ 60]),"=r"(r[ 61]),"=r"(r[ 62]),"=r"(r[ 63]),
+                 "=r"(r[ 64]),"=r"(r[ 65]),"=r"(r[ 66]),"=r"(r[ 67]),
+                 "=r"(r[ 68]),"=r"(r[ 69]),"=r"(r[ 70]),"=r"(r[ 71]),
+                 "=r"(r[ 72]),"=r"(r[ 73]),"=r"(r[ 74]),"=r"(r[ 75]),
+                 "=r"(r[ 76]),"=r"(r[ 77]),"=r"(r[ 78]),"=r"(r[ 79]),
+                 "=r"(r[ 80]),"=r"(r[ 81]),"=r"(r[ 82]),"=r"(r[ 83]),
+                 "=r"(r[ 84]),"=r"(r[ 85]),"=r"(r[ 86]),"=r"(r[ 87]),
+                 "=r"(r[ 88]),"=r"(r[ 89]),"=r"(r[ 90]),"=r"(r[ 91]),
+                 "=r"(r[ 92]),"=r"(r[ 93]),"=r"(r[ 94]),"=r"(r[ 95]),
+                 "=r"(r[ 96]),"=r"(r[ 97]),"=r"(r[ 98]),"=r"(r[ 99]),
+                 "=r"(r[100]),"=r"(r[101]),"=r"(r[102]),"=r"(r[103]),
+                 "=r"(r[104]),"=r"(r[105]),"=r"(r[106]),"=r"(r[107]),
+                 "=r"(r[108]),"=r"(r[109]),"=r"(r[110]),"=r"(r[111]),
+                 "=r"(r[112]),"=r"(r[113]),"=r"(r[114]),"=r"(r[115]),
+                 "=r"(r[116]),"=r"(r[117]),"=r"(r[118]),"=r"(r[119]),
+                 "=r"(r[120]),"=r"(r[121]),"=r"(r[122]),"=r"(r[123]),
+                 "=r"(r[124]),"=r"(r[125]),"=r"(r[126]),"=r"(r[127])
+               : "r"(tmem_addr));
+}
+
+__device__ __forceinline__
+uint32_t smem_ptr_u32(const void* ptr) {
+  return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+}
+
+__device__ __forceinline__
+void sts_f32(uint32_t smem_addr, float val) {
+  asm volatile("st.shared.f32 [%0], %1;" :: "r"(smem_addr), "f"(val) : "memory");
+}
+

fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu

@@ -0,0 +1,279 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ *
+ * Adapted from flashinfer-ai/flashinfer @ 8a94642d83cba0939035868fb6c309b4474a13d6
+ * (PR #3820), csrc/ulysses_all_to_all.cu.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Torch host bindings for the fused Ulysses all-to-all. Kernel in
+// include/comm/ulysses_all_to_all.cuh.
+//
+// The per-group context is an ncclDevComm plus a registered symmetric window,
+// both created here from the caller's ncclComm_t.
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <nccl.h>
+#include <nccl_device.h>
+
+#include "comm/ulysses_all_to_all.cuh"
+
+namespace fi = fastvideo::comm::ulysses;
+
+namespace {
+
+// A symmetric window every rank can store into, plus the ncclDevComm the
+// kernel opens barrier sessions on.
+struct UlyssesContext {
+  ncclComm_t comm = nullptr;
+  ncclWindow_t win = nullptr;
+  void* buf = nullptr;
+  size_t nbytes = 0;
+  ncclDevComm devComm{};
+  bool dev_comm_created = false;
+  int device = -1;
+  int rank = 0;
+  int world = 0;
+};
+
+#define NCCL_TRY(expr, what)                                                   \
+  do {                                                                         \
+    ncclResult_t _r = (expr);                                                  \
+    TORCH_CHECK(_r == ncclSuccess, what " failed: rc=", static_cast<int>(_r)); \
+  } while (0)
+
+}  // namespace
+
+// Allocate the local half of a context. This is intentionally separate from
+// registration so Python can vote after local allocation: if one rank is OOM,
+// no peer enters a collective window registration alone.
+int64_t allocate_ulysses_a2a(int64_t nbytes, int64_t rank, int64_t world_size,
+                             int64_t device_index) {
+  TORCH_CHECK(world_size == 2 || world_size == 4 || world_size == 6 || world_size == 8,
+              "ulysses a2a only supports world size in (2, 4, 6, 8), got ", world_size);
+  TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank");
+  TORCH_CHECK(nbytes > 0, "nbytes must be positive");
+  TORCH_CHECK(device_index >= 0, "device index must be non-negative");
+
+  const at::cuda::CUDAGuard device_guard(
+      c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(device_index)));
+  auto ctx = std::make_unique<UlyssesContext>();
+  ctx->nbytes = static_cast<size_t>(nbytes);
+  ctx->device = static_cast<int>(device_index);
+  ctx->rank = static_cast<int>(rank);
+  ctx->world = static_cast<int>(world_size);
+
+  // The window must come from NCCL's allocator (4096B aligned per
+  // NCCL_WIN_REQUIRED_ALIGNMENT), which is why this is not a torch tensor.
+  NCCL_TRY(ncclMemAlloc(&ctx->buf, ctx->nbytes), "ncclMemAlloc");
+
+  return reinterpret_cast<int64_t>(ctx.release());
+}
+
+// Register the user window. Collective: every rank must call together.
+void register_ulysses_a2a_window(int64_t handle, int64_t comm_ptr) {
+  auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
+  TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a");
+  TORCH_CHECK(ctx->buf != nullptr && ctx->win == nullptr && !ctx->dev_comm_created,
+              "ulysses a2a context is not in the allocated state");
+
+  const at::cuda::CUDAGuard device_guard(
+      c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(ctx->device)));
+  ctx->comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+  NCCL_TRY(ncclCommWindowRegister(ctx->comm, ctx->buf, ctx->nbytes, &ctx->win,
+                                  NCCL_WIN_COLL_SYMMETRIC),
+           "ncclCommWindowRegister");
+}
+
+// Create the device communicator only after Python has voted that every rank
+// registered its window. This operation is collective as well.
+void create_ulysses_a2a_dev_comm(int64_t handle) {
+  auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
+  TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a");
+  TORCH_CHECK(ctx->comm != nullptr && ctx->win != nullptr && !ctx->dev_comm_created,
+              "ulysses a2a context is not in the window-registered state");
+
+  const at::cuda::CUDAGuard device_guard(
+      c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(ctx->device)));
+
+  ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
+  reqs.lsaBarrierCount = fi::kMaxBlocks;
+  NCCL_TRY(ncclDevCommCreate(ctx->comm, &reqs, &ctx->devComm), "ncclDevCommCreate");
+  ctx->dev_comm_created = true;
+}
+
+static ncclResult_t first_error(ncclResult_t current, ncclResult_t next) {
+  return current == ncclSuccess ? next : current;
+}
+
+void dispose_ulysses_a2a(int64_t handle) {
+  auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
+  if (ctx == nullptr) return;
+
+  const at::cuda::CUDAGuard device_guard(
+      c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(ctx->device)));
+  ncclResult_t result = ncclSuccess;
+  if (ctx->comm != nullptr && ctx->dev_comm_created) {
+    result = first_error(result, ncclDevCommDestroy(ctx->comm, &ctx->devComm));
+    ctx->dev_comm_created = false;
+  }
+  if (ctx->comm != nullptr && ctx->win != nullptr) {
+    result = first_error(result, ncclCommWindowDeregister(ctx->comm, ctx->win));
+    ctx->win = nullptr;
+  }
+  if (ctx->buf != nullptr) {
+    result = first_error(result, ncclMemFree(ctx->buf));
+    ctx->buf = nullptr;
+  }
+  delete ctx;
+  TORCH_CHECK(result == ncclSuccess, "ulysses a2a cleanup failed: rc=", static_cast<int>(result));
+}
+
+// Whether the whole group is load-store accessible. NCCL determined this at
+// ncclCommInitRank.
+bool ulysses_lsa_covers_group(int64_t comm_ptr, int64_t world_size) {
+  auto comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+  ncclCommProperties properties = NCCL_COMM_PROPERTIES_INITIALIZER;
+  NCCL_TRY(ncclCommQueryProperties(comm, &properties), "ncclCommQueryProperties");
+  ncclTeam_t lsa = ncclTeamLsa(comm);
+  return properties.deviceApiSupport && lsa.nRanks == static_cast<int>(world_size);
+}
+
+// Fused-transpose Ulysses all-to-all.
+//   mode == 0: inp [B, S_local, H, D]        -> out [B, S_global, H_local, D]
+//   mode == 1: inp [B, S_global, H_local, D] -> out [B, S_local, H, D]
+// where H is the *global* head count and H_local = H / world_size.
+void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local,
+                 int64_t H, int64_t D, int64_t mode) {
+  auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
+  TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a");
+
+  const at::cuda::CUDAGuard device_guard(inp.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  TORCH_CHECK(inp.is_cuda() && out.is_cuda(), "inp and out must be CUDA tensors");
+  TORCH_CHECK(inp.is_contiguous() && out.is_contiguous(), "inp and out must be contiguous");
+  TORCH_CHECK(inp.device() == out.device(), "inp and out must be on the same device");
+  TORCH_CHECK(inp.scalar_type() == out.scalar_type(), "inp and out must share a dtype");
+  TORCH_CHECK(inp.numel() == out.numel(), "inp and out must have equal element counts");
+  TORCH_CHECK(inp.get_device() == ctx->device, "input is on CUDA device ", inp.get_device(),
+              " but the Ulysses context belongs to device ", ctx->device);
+  TORCH_CHECK(mode == 0 || mode == 1, "mode must be 0 or 1");
+  TORCH_CHECK(inp.dim() == 4 && out.dim() == 4, "inp and out must be 4-D");
+
+  const int W = ctx->world;
+  TORCH_CHECK(H % W == 0, "global head count must be divisible by world size");
+  const int H_local = static_cast<int>(H / W);
+
+  const torch::Tensor& local_op = (mode == 0) ? inp : out;   // [B, S_local, H, D]
+  const torch::Tensor& global_op = (mode == 0) ? out : inp;  // [B, S_global, H_local, D]
+  TORCH_CHECK(local_op.size(0) == B && local_op.size(1) == S_local && local_op.size(2) == H &&
+                  local_op.size(3) == D,
+              "the [B, S_local, H, D] operand of mode ", mode, " has shape (", local_op.size(0),
+              ", ", local_op.size(1), ", ", local_op.size(2), ", ", local_op.size(3),
+              "), expected (", B, ", ", S_local, ", ", H, ", ", D, ")");
+  TORCH_CHECK(global_op.size(0) == B && global_op.size(1) == W * S_local &&
+                  global_op.size(2) == H_local && global_op.size(3) == D,
+              "the [B, S_global, H_local, D] operand of mode ", mode, " has shape (",
+              global_op.size(0), ", ", global_op.size(1), ", ", global_op.size(2), ", ",
+              global_op.size(3), "), expected (", B, ", ", W * S_local, ", ", H_local, ", ", D,
+              ")");
+
+  const size_t out_bytes = out.numel() * out.element_size();
+  TORCH_CHECK(out_bytes <= ctx->nbytes, "operand of ", out_bytes,
+              " bytes exceeds the window capacity ", ctx->nbytes);
+
+  const int64_t num_rows = B * static_cast<int64_t>(W) * S_local;
+  const int blocks =
+      static_cast<int>(std::max<int64_t>(1, std::min<int64_t>(fi::kMaxBlocks, num_rows)));
+  const int threads = fi::kUlyssesThreads;
+
+#define LAUNCH_ULYSSES_A2A(T, NG, MODE)                                             \
+  fi::ulysses_a2a_kernel<T, NG, MODE><<<blocks, threads, 0, stream>>>(              \
+      reinterpret_cast<const T*>(inp.data_ptr()), ctx->devComm, ctx->win, /*off=*/0, \
+      ctx->rank, static_cast<int>(B), static_cast<int>(S_local), H_local, static_cast<int>(D))
+
+#define DISPATCH_NGPUS(T, MODE)                                                \
+  switch (W) {                                                                 \
+    case 2:                                                                    \
+      LAUNCH_ULYSSES_A2A(T, 2, MODE);                                          \
+      break;                                                                   \
+    case 4:                                                                    \
+      LAUNCH_ULYSSES_A2A(T, 4, MODE);                                          \
+      break;                                                                   \
+    case 6:                                                                    \
+      LAUNCH_ULYSSES_A2A(T, 6, MODE);                                          \
+      break;                                                                   \
+    case 8:                                                                    \
+      LAUNCH_ULYSSES_A2A(T, 8, MODE);                                          \
+      break;                                                                   \
+    default:                                                                   \
+      TORCH_CHECK(false, "ulysses_a2a only supports world size in (2,4,6,8)"); \
+  }
+
+#define DISPATCH_DTYPE(MODE)                                                              \
+  switch (out.scalar_type()) {                                                            \
+    case at::ScalarType::Float: {                                                         \
+      DISPATCH_NGPUS(float, MODE);                                                        \
+      break;                                                                              \
+    }                                                                                     \
+    case at::ScalarType::Half: {                                                          \
+      DISPATCH_NGPUS(half, MODE);                                                         \
+      break;                                                                              \
+    }                                                                                     \
+    case at::ScalarType::BFloat16: {                                                      \
+      DISPATCH_NGPUS(nv_bfloat16, MODE);                                                  \
+      break;                                                                              \
+    }                                                                                     \
+    default:                                                                              \
+      TORCH_CHECK(false, "ulysses_a2a only supports float32, float16 and bfloat16, got ", \
+                  out.scalar_type());                                                     \
+  }
+
+  if (mode == 0) {
+    DISPATCH_DTYPE(0);
+  } else {
+    DISPATCH_DTYPE(1);
+  }
+
+#undef DISPATCH_DTYPE
+#undef DISPATCH_NGPUS
+#undef LAUNCH_ULYSSES_A2A
+
+  TORCH_CHECK(cudaGetLastError() == cudaSuccess, "ulysses_a2a kernel launch failed");
+  // Copy this rank's completed result out of the window.
+  auto status = cudaMemcpyAsync(out.data_ptr(), ctx->buf, out_bytes, cudaMemcpyDeviceToDevice,
+                                stream);
+  TORCH_CHECK(status == cudaSuccess, "ulysses_a2a copy-out failed: ", cudaGetErrorString(status));
+}
+
+void register_ulysses_a2a(pybind11::module_& m) {
+  m.def("allocate_ulysses_a2a", &allocate_ulysses_a2a, "allocate a local ulysses a2a window");
+  m.def("register_ulysses_a2a_window", &register_ulysses_a2a_window,
+        "register the ulysses a2a window collectively");
+  m.def("create_ulysses_a2a_dev_comm", &create_ulysses_a2a_dev_comm,
+        "create the ulysses a2a device communicator collectively");
+  m.def("dispose_ulysses_a2a", &dispose_ulysses_a2a, "release a ulysses a2a context");
+  m.def("ulysses_lsa_covers_group", &ulysses_lsa_covers_group,
+        "whether the whole group is load-store accessible");
+  m.def("ulysses_a2a", &ulysses_a2a, "fused-transpose Ulysses all-to-all over NVLink");
+}

fastvideo-kernel/csrc/common_extension.cpp

@@ -22,16 +22,42 @@ extern std::vector<torch::Tensor> block_sparse_attention_backward(
 );
 #endif
 
+#ifdef FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A
+// Ulysses sequence-parallel all-to-all (csrc/comm/)
+void register_ulysses_a2a(pybind11::module_ &);
+#endif
+
 // TurboDiffusion kernels
 void register_quant(pybind11::module_ &);
 void register_rms_norm(pybind11::module_ &);
 void register_layer_norm(pybind11::module_ &);
 void register_gemm(pybind11::module_ &);
 
+#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A
+extern std::vector<torch::Tensor> block_sparse_sm100a_fwd(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<torch::Tensor> v_t,
+    torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes,
+    double sm_scale, bool need_lse);
+extern std::vector<torch::Tensor> block_sparse_sm100a_blk128_fwd(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<torch::Tensor> v_t,
+    torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes,
+    double sm_scale, bool need_lse);
+#endif
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.doc() = "FastVideo CUDA Kernels";
 
+#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A
+    m.def("block_sparse_sm100a_fwd",
+          torch::wrap_pybind_function(block_sparse_sm100a_fwd),
+          "VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a)");
+    m.def("block_sparse_sm100a_blk128_fwd",
+          torch::wrap_pybind_function(block_sparse_sm100a_blk128_fwd),
+          "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a)");
+#endif
+
 #ifdef TK_COMPILE_ST_ATTN
+
     m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)");
 #endif
 
@@ -40,6 +66,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)");
 #endif
 
+#ifdef FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A
+    // Ulysses sequence-parallel all-to-all
+    register_ulysses_a2a(m);
+#endif
+
     // TurboDiffusion
     register_quant(m);
     register_rms_norm(m);

fastvideo-kernel/include/comm/ulysses_all_to_all.cuh

@@ -0,0 +1,165 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ *
+ * Adapted from flashinfer-ai/flashinfer @ 8a94642d83cba0939035868fb6c309b4474a13d6
+ * (PR #3820), which in turn adapted ThunderKittens' NVLink all-to-all:
+ * https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/parallel/all_to_all/all_to_all.cu
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Fused-transpose Ulysses all-to-all over NVLink. Peer addresses come from
+// ncclGetLsaPointer and synchronization from ncclLsaBarrierSession; the index
+// math and slab decomposition below are upstream's.
+//
+// head_dim == 2 layout, uniform sequence splits. With
+//   W        = ulysses world size
+//   H_local  = H / W
+//   S_global = S_local * W
+//
+//   mode == 0 (input  a2a): [B, S_local, H,       D] -> [B, S_global, H_local, D]
+//       y_r[b, j*S_local + s, hl, d] = x_j[b, s, r*H_local + hl, d]
+//   mode == 1 (output a2a): [B, S_global, H_local, D] -> [B, S_local, H,       D]
+//       out_j[b, s, r*H_local + hl, d] = u_r[b, j*S_local + s, hl, d]
+//
+// In both modes the unit of transfer is a contiguous (H_local * D) block, so
+// every cross-GPU store is fully coalesced.
+
+#ifndef FASTVIDEO_COMM_ULYSSES_ALL_TO_ALL_CUH_
+#define FASTVIDEO_COMM_ULYSSES_ALL_TO_ALL_CUH_
+
+#include <cstdint>
+
+#include <nccl.h>
+#include <nccl_device.h>
+
+namespace fastvideo {
+namespace comm {
+namespace ulysses {
+
+constexpr int kUlyssesThreads = 512;
+// Deliberately modest: this is link-bandwidth bound, so a small grid leaves the
+// rest of the GPU free without costing throughput.
+constexpr int kMaxBlocks = 36;
+
+// Shared movement body for the fused-transpose all-to-all (no barriers).
+//
+// Rows are ordered ((b * W + peer) * S_local + s), so consecutive rows share a
+// (batch, peer) and are contiguous on the gather side of the transpose. Each
+// block takes a contiguous slab of rows and flattens its threads over the 16B
+// units in it, so consecutive lanes address one peer buffer back to back and the
+// remote writes coalesce into large bursts rather than (H_local * D)-sized
+// scattered ones.
+template <typename T, int NGPUS, int MODE>
+__device__ __forceinline__ void ulysses_a2a_move(const T* __restrict__ local_in,
+                                                 void* const* peer_ptrs, int rank, int B,
+                                                 int S_local, int H_local, int D) {
+  static_assert(MODE == 0 || MODE == 1, "MODE must be 0 or 1");
+  const int W = NGPUS;
+  const int64_t H = static_cast<int64_t>(H_local) * W;
+  const int64_t S_global = static_cast<int64_t>(S_local) * W;
+  const int64_t block_len = static_cast<int64_t>(H_local) * D;  // elements/row
+  const int64_t num_rows = static_cast<int64_t>(B) * W * S_local;
+
+  // 16B-vectorized fast path when every row is 16B aligned (the common case:
+  // contiguous bf16/fp16/fp32 tensors with block_len * sizeof(T) % 16 == 0).
+  using Vec = int4;
+  constexpr int kVecBytes = sizeof(Vec);
+  const int64_t row_bytes = block_len * static_cast<int64_t>(sizeof(T));
+  const bool vec_ok =
+      (row_bytes % kVecBytes) == 0 && (reinterpret_cast<uintptr_t>(local_in) % kVecBytes) == 0;
+
+  // Contiguous slab of rows for this block.
+  const int64_t rows_per_block = (num_rows + gridDim.x - 1) / gridDim.x;
+  const int64_t row_lo = static_cast<int64_t>(blockIdx.x) * rows_per_block;
+  int64_t row_hi = row_lo + rows_per_block;
+  if (row_hi > num_rows) row_hi = num_rows;
+  if (row_lo >= row_hi) return;
+
+  const int tid = threadIdx.x;
+  const int nthr = blockDim.x;
+
+  // Decode (b, peer, s) and compute src/dst element offsets for a given row.
+  auto offsets = [&](int64_t row, int64_t& src_off, int64_t& dst_off) {
+    const int64_t s = row % S_local;
+    const int64_t tmp = row / S_local;
+    const int64_t peer = tmp % W;
+    const int64_t b = tmp / W;
+    if constexpr (MODE == 0) {
+      src_off = ((b * S_local + s) * H + peer * H_local) * D;
+      dst_off = (b * S_global + static_cast<int64_t>(rank) * S_local + s) * block_len;
+    } else {
+      src_off = (b * S_global + peer * S_local + s) * block_len;
+      dst_off = ((b * S_local + s) * H + static_cast<int64_t>(rank) * H_local) * D;
+    }
+    return peer;
+  };
+
+  if (vec_ok) {
+    const int64_t units_per_row = row_bytes / kVecBytes;
+    const int64_t total_units = (row_hi - row_lo) * units_per_row;
+    for (int64_t u = tid; u < total_units; u += nthr) {
+      const int64_t local_row = u / units_per_row;
+      const int64_t unit = u - local_row * units_per_row;
+      const int64_t row = row_lo + local_row;
+      int64_t src_off, dst_off;
+      const int64_t peer = offsets(row, src_off, dst_off);
+      const Vec* s4 = reinterpret_cast<const Vec*>(local_in + src_off);
+      Vec* d4 = reinterpret_cast<Vec*>((T*)peer_ptrs[peer] + dst_off);
+      d4[unit] = s4[unit];
+    }
+  } else {
+    // Scalar fallback (unaligned / odd shapes).
+    for (int64_t row = row_lo; row < row_hi; ++row) {
+      int64_t src_off, dst_off;
+      const int64_t peer = offsets(row, src_off, dst_off);
+      const T* s_ptr = local_in + src_off;
+      T* d_ptr = (T*)peer_ptrs[peer] + dst_off;
+      for (int64_t i = tid; i < block_len; i += nthr) {
+        d_ptr[i] = s_ptr[i];
+      }
+    }
+  }
+}
+
+// The transfer mode is a compile-time template parameter so the address math
+// specializes and the coalesced slab decomposition is fully unrolled per mode.
+template <typename T, int NGPUS, int MODE>
+__global__ void __launch_bounds__(kUlyssesThreads, 1)
+    ulysses_a2a_kernel(const T* __restrict__ local_in, ncclDevComm devComm, ncclWindow_t win,
+                       size_t win_offset, int rank, int B, int S_local, int H_local, int D) {
+  // Resolved once; the movement loop would otherwise call this per 16B store.
+  void* peer_ptrs[NGPUS];
+#pragma unroll
+  for (int p = 0; p < NGPUS; ++p) {
+    peer_ptrs[p] = ncclGetLsaPointer(win, win_offset, p);
+  }
+
+  // Each CTA owns one barrier generation. Sharing index 0 across independently
+  // scheduled CTAs races the generation counter and is unsupported by NCCL's
+  // device API. The host reserves kMaxBlocks slots when it creates devComm.
+  ncclLsaBarrierSession<ncclCoopCta> bar(
+      ncclCoopCta(), devComm, ncclTeamTagLsa(), /*index=*/blockIdx.x);
+
+  // Every rank must have entered before anyone writes into peer buffers.
+  bar.sync(ncclCoopCta(), cuda::memory_order_relaxed);
+  ulysses_a2a_move<T, NGPUS, MODE>(local_in, peer_ptrs, rank, B, S_local, H_local, D);
+  // Release-acquire: all peer writes visible before a rank reads its window.
+  bar.sync(ncclCoopCta(), cuda::memory_order_acq_rel);
+}
+
+}  // namespace ulysses
+}  // namespace comm
+}  // namespace fastvideo
+
+#endif  // FASTVIDEO_COMM_ULYSSES_ALL_TO_ALL_CUH_

fastvideo-kernel/pyproject.toml

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

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -49,6 +49,36 @@ def _force_tk() -> bool:
     return os.environ.get("FASTVIDEO_VSA_TK", "0") == "1"
 
 
+def _force_sm100a() -> bool:
+    """True iff the sm_100a (Blackwell) forward is explicitly opted into.
+
+    Opt-in only (same env the H3 backend honors): the sm_100a extension is
+    forward-only, so this routing pairs it with the Triton backward -- its lse
+    is already in Triton's M format. Honored only when
+    ``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token
+    metadata falls through to the default selection; unsupported 128-token
+    metadata raises because Triton has no compatible fallback.
+    ``FASTVIDEO_VSA_TRITON`` still wins.
+    """
+    return os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1"
+
+
+def _sm100a_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
+    try:
+        from fastvideo_kernel import block_sparse_attn_sm100a as vsa_sm100a
+    except Exception:
+        return False
+    return vsa_sm100a.is_supported(q, variable_block_sizes)
+
+
+def _infer_block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
+    num_blocks = variable_block_sizes.numel()
+    seq_len = q.shape[2]
+    if num_blocks == 0 or seq_len % num_blocks != 0:
+        return 0
+    return seq_len // num_blocks
+
+
 # ---------------------------------------------------------------------------
 # Index helpers
 # ---------------------------------------------------------------------------
@@ -339,6 +369,76 @@ def _backward_sm90(ctx, grad_o, grad_lse):
 
 block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90)
 
+# ---------------------------------------------------------------------------
+# SM100A backend custom op (index-native)
+#
+# Forward runs the sm_100a CUDA extension; backward reuses the Triton kernels.
+# The sm_100a forward emits lse in exactly Triton's M format (max*log2e +
+# log2(l)), so the pairing needs no conversion. The Triton backward is
+# hardcoded to 64-token blocks, hence the block-size assert below.
+# ---------------------------------------------------------------------------
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_sm100a",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_sm100a_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    from fastvideo_kernel.block_sparse_attn_sm100a import block_sparse_attn_sm100a
+
+    o, M = block_sparse_attn_sm100a(q, k, v, q2k_idx, q2k_num, variable_block_sizes,
+                                    need_lse=True)
+    return o, M
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a")
+def _block_sparse_attn_sm100a_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    o = torch.empty_like(q)
+    M = torch.empty(
+        (q.shape[0], q.shape[1], q.shape[2]),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return o, M
+
+
+def _setup_context_sm100a(ctx, inputs, output):
+    q, k, v, q2k_idx, q2k_num, variable_block_sizes = inputs
+    o, M = output
+    ctx.save_for_backward(q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes)
+
+
+def _backward_sm100a(ctx, grad_o, grad_M):
+    q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors
+    block = q.shape[2] // variable_block_sizes.numel()
+    if block != 64:
+        raise RuntimeError(
+            "block_sparse_attn_sm100a backward pairs the sm_100a forward with the "
+            f"Triton backward, which is hardcoded to 64-token blocks; got {block}. "
+            "Run 128-token-block metadata without grad, or use the Triton forward.")
+    dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx,
+                                                   q2k_num, variable_block_sizes)
+    return dq, dk, dv, None, None, None
+
+
+block_sparse_attn_sm100a_op.register_autograd(_backward_sm100a,
+                                              setup_context=_setup_context_sm100a)
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -364,11 +464,24 @@ def block_sparse_attn_from_indices(
 
     # Backend resolution:
     # - FASTVIDEO_VSA_TRITON forces Triton everywhere.
+    # - FASTVIDEO_VSA_SM100A opts into the sm_100a forward (Triton backward).
+    #   Unsupported 64-token metadata falls through; unsupported 128-token
+    #   metadata raises because Triton cannot consume it.
     # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually
     #   available (else falls through to the default below).
     # - Otherwise: TK on sm_90 if available, else Triton.
     if _force_triton():
         use_sm90 = False
+    elif _force_sm100a():
+        if _sm100a_is_supported(q, variable_block_sizes):
+            return block_sparse_attn_sm100a_op(q, k, v, q2k_idx, q2k_num,
+                                               variable_block_sizes)
+        if _infer_block_size(q, variable_block_sizes) == 128:
+            raise NotImplementedError(
+                "128-token block-sparse attention requires the sm_100a forward; "
+                "the Triton fallback only supports 64-token blocks, and the "
+                "sm_100a route is unavailable for this input.")
+        use_sm90 = sm90_available
     elif _force_tk():
         use_sm90 = sm90_available
     else:

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py

@@ -1,4 +1,4 @@
-"""VSA-256 block-sparse attention wrapper.
+"""VSA-128/256 block-sparse attention wrappers.
 
 The default 256-block path is Triton: it expands the logical 256-block map
 to the existing 64-block Triton kernel via a dense 4x4 expansion per logical
@@ -7,8 +7,8 @@ edge ("route A"), and requires no optional dependencies.
 The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is
 *opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to
 :mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates
-on 128-token KV blocks (this wrapper expands the logical 256-block map /
-sizes into that physical 128-block representation). The CuTe kernel
+on 128-token Q/KV blocks (the 256 wrapper expands its logical KV map and
+sizes into that physical representation). The CuTe kernel
 (``flash_attn.cute`` with block-sparsity) is an optional dependency,
 imported lazily only when this fastpath is selected.
 
@@ -35,7 +35,7 @@ _KV_BLOCK_TRITON = 64  # Existing Triton path uses 64-token KV blocks.
 
 
 def _resolve_backend() -> str:
-    """Pick the backend for the 256-block VSA path.
+    """Pick the backend for the 128/256-block VSA paths.
 
     Default is Triton (no optional deps). The FA4 CuTe fastpath is opt-in
     via ``FASTVIDEO_VSA_CUTEDSL=1`` and requires the optional FA4 CuTe
@@ -49,6 +49,26 @@ def _resolve_backend() -> str:
     return "triton"
 
 
+def _expand_mask_and_sizes_128_to_64(
+    logical_mask_128: torch.Tensor,
+    logical_kv_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Expand a [B, H, Qb128, KVb128] map to 64-token Triton tiles."""
+    expanded_mask = logical_mask_128.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
+    sizes_i32 = logical_kv_sizes_128.to(torch.int32)
+    offsets = torch.tensor(
+        [0, _KV_BLOCK_TRITON],
+        dtype=torch.int32,
+        device=sizes_i32.device,
+    )
+    expanded_sizes = torch.clamp(
+        sizes_i32[:, None] - offsets[None, :],
+        min=0,
+        max=_KV_BLOCK_TRITON,
+    ).reshape(-1)
+    return expanded_mask, expanded_sizes
+
+
 def _expand_mask_and_sizes_256_to_128(
     logical_mask_256: torch.Tensor,
     logical_kv_sizes_256: torch.Tensor,
@@ -112,6 +132,63 @@ def _triton_via_route_a(
     return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
 
 
+def _triton_via_route_a_128(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_mask_128: torch.Tensor,
+    logical_kv_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    from .triton_kernels.index import map_to_index as triton_map_to_index
+
+    mask_64, sizes_64 = _expand_mask_and_sizes_128_to_64(logical_mask_128, logical_kv_sizes_128)
+    q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool))
+    return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
+
+
+def block_sparse_attn_128(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_block_map_128: torch.Tensor,
+    logical_variable_block_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """VSA-128 sparse-branch entrypoint for [B, H, S, D] inputs."""
+    if logical_block_map_128.dim() == 3:
+        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
+
+    if _resolve_backend() == "triton":
+        return _triton_via_route_a_128(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
+    return block_sparse_attn_cute_fwd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+
+def block_sparse_attn_128_bshd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_block_map_128: torch.Tensor,
+    logical_variable_block_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """VSA-128 sparse-branch entrypoint for [B, S, H, D] inputs."""
+    if logical_block_map_128.dim() == 3:
+        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
+
+    if _resolve_backend() == "triton":
+        out_bhsd, aux = _triton_via_route_a_128(
+            q.transpose(1, 2).contiguous(),
+            k.transpose(1, 2).contiguous(),
+            v.transpose(1, 2).contiguous(),
+            logical_block_map_128,
+            logical_variable_block_sizes_128,
+        )
+        return out_bhsd.transpose(1, 2).contiguous(), aux
+
+    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
+    return block_sparse_attn_cute_fwd_bshd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+
 def block_sparse_attn_256(
     q: torch.Tensor,
     k: torch.Tensor,

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py

@@ -1,18 +1,18 @@
-"""CuTe-DSL block-sparse attention forward kernel.
+"""FA4 CuTe-DSL block-sparse attention adapter.
 
-Thin wrapper around `flash_attn.cute.interface._flash_attn_fwd` that adapts
-VSA's `(block_map, variable_block_sizes)` inputs into FA4's
-`BlockSparseTensorsTorch` representation and the per-KV-block validity mask.
+This module adapts VSA's ``(block_map, variable_block_sizes)`` inputs into
+FA4's forward and backward ``BlockSparseTensorsTorch`` representations.
+FA4's public ``flash_attn_func`` owns the forward/backward autograd bridge.
 
 Both [B, H, S, D] (BHSD) and [B, S, H, D] (BSHD) entrypoints are provided.
-The BSHD variant is preferred from VSA-256 callers to avoid layout
+The BSHD variant is preferred from VSA-128/256 callers to avoid layout
 round-trips on the hot path.
 
 The FA4 CuTe block-sparse kernel (``flash_attn.cute`` with
 ``block_sparsity``) is an *optional* dependency: it is imported lazily and
-only exercised when the VSA-256 CuTe fastpath is explicitly selected
-(``FASTVIDEO_VSA_CUTEDSL=1``). The default VSA-256 path is Triton and does
-not require it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
+only exercised when the VSA-128/256 CuTe fastpath is explicitly selected
+(``FASTVIDEO_VSA_CUTEDSL=1``). The default path is Triton and does not require
+it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
 """
 
 from __future__ import annotations
@@ -22,13 +22,14 @@ from typing import Tuple
 
 import torch
 
-_FA4_IMPORT_HINT = ("VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
+_FA4_IMPORT_HINT = ("VSA-128/256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
                     "provides `flash_attn.cute` with block-sparsity support (plus "
                     "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional "
-                    "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe "
+                    "dependency; the default path is Triton. Install the FA4 CuTe "
                     "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.")
 
 
+@functools.lru_cache(maxsize=1)
 def _load_fa4_cute():
     """Lazily import the optional FA4 CuTe block-sparse symbols.
 
@@ -38,14 +39,39 @@ def _load_fa4_cute():
     """
     try:
         from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch
-        from flash_attn.cute.interface import _flash_attn_fwd
+        from flash_attn.cute.interface import (
+            _flash_attn_bwd,
+            _flash_attn_fwd,
+            flash_attn_func,
+        )
     except ImportError as exc:  # pragma: no cover - optional dependency
         raise ImportError(_FA4_IMPORT_HINT) from exc
-    return BlockSparseTensorsTorch, _flash_attn_fwd
+    return BlockSparseTensorsTorch, flash_attn_func, _flash_attn_fwd, _flash_attn_bwd
 
 
-# Q-side tile size; kv_block_size comes from the caller's VSA logical KV block.
-_M_BLOCK_SIZE_DEFAULT = 128
+# FA4's physical Q tile size; KV block size comes from the VSA caller.
+_FA4_Q_BLOCK_SIZE = 128
+
+
+class _SingleQStageLength(int):
+    """Keep the real length while selecting FA4's one-stage Q128 path.
+
+    On sm_100 FA4 derives ``q_stage`` from ``max_seqlen_q > tile_m``. Its
+    kernel supports one 128-token Q stage, but the fixed-length public wrapper
+    does not expose that choice. VSA-128 must select it explicitly; otherwise
+    adjacent logical Q blocks are merged into a 256-token sparse block.
+    """
+
+    def __mul__(self, other):
+        return type(self)(int(self) * int(other))
+
+    def __rmul__(self, other):
+        return type(self)(int(other) * int(self))
+
+    def __gt__(self, other):
+        if int(other) == _FA4_Q_BLOCK_SIZE:
+            return False
+        return int(self) > int(other)
 
 
 def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -64,12 +90,12 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     return triton_map_to_index(block_map)
 
 
-def _choose_q_sparse_block_size(q_len: int, m_block_size: int = _M_BLOCK_SIZE_DEFAULT) -> int:
-    # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > m_block_size.
+def _choose_q_sparse_block_size(q_len: int, q_tile_size: int = _FA4_Q_BLOCK_SIZE) -> int:
+    # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > q_tile_size.
     major, _ = torch.cuda.get_device_capability()
-    if major >= 10 and q_len > m_block_size:
-        return 2 * m_block_size
-    return m_block_size
+    if major >= 10 and q_len > q_tile_size:
+        return 2 * q_tile_size
+    return q_tile_size
 
 
 def _aggregate_q_block_map(
@@ -134,23 +160,35 @@ def _build_vbs_mask_mod(kv_block_size: int):
     return _vbs_mask_mod
 
 
-def _cute_forward(
-    q_bshd: torch.Tensor,
-    k_bshd: torch.Tensor,
-    v_bshd: torch.Tensor,
+def _build_sparse_tensors(
     block_map: torch.Tensor,
     variable_block_sizes: torch.Tensor,
     *,
+    q_len: int,
     q_block_size: int,
     kv_block_size: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Internal: FA4 CuTe BSA fwd with BSHD inputs."""
-    BlockSparseTensorsTorch, _flash_attn_fwd = _load_fa4_cute()
-    q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1])
-    q_sparse_block_size = max(
-        q_block_size,
-        ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
-    )
+    need_backward: bool,
+    force_q_sparse_block_size: int | None = None,
+) -> Tuple[object, object | None]:
+    """Build the Q-owned forward and KV-owned backward sparse metadata.
+
+    ``need_backward`` is False on inference-only calls: the backward metadata
+    is a pair of dense ``[B, H, kv_blocks, q_blocks]`` int32 index tensors that
+    FA4 keeps alive on its autograd ctx until backward runs, so building it
+    when nothing requires grad is pure overhead (~80 MiB per call at Wan-14B
+    720p shape).
+    """
+    BlockSparseTensorsTorch, _, _, _ = _load_fa4_cute()
+    if force_q_sparse_block_size is None:
+        q_sparse_candidate = _choose_q_sparse_block_size(q_len)
+        q_sparse_block_size = max(
+            q_block_size,
+            ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
+        )
+    else:
+        q_sparse_block_size = force_q_sparse_block_size
+        if q_sparse_block_size < q_block_size or q_sparse_block_size % q_block_size != 0:
+            raise ValueError("force_q_sparse_block_size must be a positive multiple of q_block_size")
     sparse_map = _aggregate_q_block_map(
         block_map,
         q_sparse_block_size=q_sparse_block_size,
@@ -158,35 +196,166 @@ def _cute_forward(
     )
     kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1)
     kv_partial = ((variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)).view(1, 1, 1, -1)
-    full_map = sparse_map & kv_full
-    mask_map = sparse_map & kv_partial
 
-    full_block_idx, full_block_cnt = _map_to_index(full_map)
-    mask_block_idx, mask_block_cnt = _map_to_index(mask_map)
+    def from_maps(full_map: torch.Tensor, mask_map: torch.Tensor) -> object:
+        full_block_idx, full_block_cnt = _map_to_index(full_map.contiguous())
+        mask_block_idx, mask_block_cnt = _map_to_index(mask_map.contiguous())
+        return BlockSparseTensorsTorch(
+            full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
+            full_block_idx=full_block_idx.to(torch.int32).contiguous(),
+            mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
+            mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
+            block_size=(q_sparse_block_size, kv_block_size),
+        )
 
-    sparse_tensors = BlockSparseTensorsTorch(
-        full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
-        full_block_idx=full_block_idx.to(torch.int32).contiguous(),
-        mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
-        mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
-        block_size=(q_sparse_block_size, kv_block_size),
+    forward_sparse_tensors = from_maps(
+        sparse_map & kv_full,
+        sparse_map & kv_partial,
     )
 
-    # _flash_attn_fwd returns (out, lse, p, row_max); keep the first two.
-    out, lse = _flash_attn_fwd(
+    if not need_backward:
+        return forward_sparse_tensors, None
+
+    # FA4 backward is KV-owned: for each physical KV tile, list the sparse
+    # query tiles that selected it. Full and partial KV tiles stay separate
+    # so the token-level validity mask only runs for padded tiles.
+    backward_sparse_tensors = from_maps(
+        (sparse_map & kv_full).transpose(2, 3),
+        (sparse_map & kv_partial).transpose(2, 3),
+    )
+    return forward_sparse_tensors, backward_sparse_tensors
+
+
+def _cute_attention_q128_forward(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    *,
+    need_backward: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, object | None]:
+    """Run FA4 with one physical Q stage per logical VSA-128 block."""
+    _, _, flash_attn_fwd, _ = _load_fa4_cute()
+    forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
+        block_map,
+        variable_block_sizes,
+        q_len=q_bshd.shape[1],
+        q_block_size=_FA4_Q_BLOCK_SIZE,
+        kv_block_size=_FA4_Q_BLOCK_SIZE,
+        need_backward=need_backward,
+        force_q_sparse_block_size=_FA4_Q_BLOCK_SIZE,
+    )
+    out, lse = flash_attn_fwd(
         q_bshd,
         k_bshd,
         v_bshd,
-        tile_mn=(_M_BLOCK_SIZE_DEFAULT, kv_block_size),
-        mask_mod=_build_vbs_mask_mod(kv_block_size),
-        block_sparse_tensors=sparse_tensors,
+        tile_mn=(_FA4_Q_BLOCK_SIZE, _FA4_Q_BLOCK_SIZE),
+        max_seqlen_q=_SingleQStageLength(q_bshd.shape[1]),
+        mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
+        block_sparse_tensors=forward_sparse_tensors,
         aux_tensors=[variable_block_sizes],
         causal=False,
         return_lse=True,
     )[:2]
+    return out, lse, backward_sparse_tensors
+
+
+class _CuteAttentionQ128(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes):
+        out, lse, backward_sparse_tensors = _cute_attention_q128_forward(
+            q_bshd,
+            k_bshd,
+            v_bshd,
+            block_map,
+            variable_block_sizes,
+            need_backward=True,
+        )
+        ctx.save_for_backward(q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes)
+        ctx.backward_sparse_tensors = backward_sparse_tensors
+        ctx.mark_non_differentiable(lse)
+        ctx.set_materialize_grads(False)
+        return out, lse
+
+    @staticmethod
+    def backward(ctx, grad_out, grad_lse):
+        del grad_lse
+        q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes = ctx.saved_tensors
+        if grad_out is None:
+            grad_out = torch.zeros_like(out)
+        _, _, _, flash_attn_bwd = _load_fa4_cute()
+        dq, dk, dv = flash_attn_bwd(
+            q_bshd,
+            k_bshd,
+            v_bshd,
+            out,
+            grad_out.contiguous(),
+            lse,
+            softmax_scale=q_bshd.shape[-1]**-0.5,
+            mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
+            aux_tensors=[variable_block_sizes],
+            block_sparse_tensors=ctx.backward_sparse_tensors,
+        )
+        return dq, dk, dv, None, None
+
+
+def _cute_attention_q128(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
+    if need_backward:
+        return _CuteAttentionQ128.apply(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
+    out, lse, _ = _cute_attention_q128_forward(
+        q_bshd,
+        k_bshd,
+        v_bshd,
+        block_map,
+        variable_block_sizes,
+        need_backward=False,
+    )
     return out, lse
 
 
+def _cute_attention(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Run FA4's autograd-enabled block-sparse attention with BSHD inputs."""
+    _, flash_attn_func, _, _ = _load_fa4_cute()
+    q_block_size = q_bshd.shape[1] // block_map.shape[2]
+    kv_block_size = k_bshd.shape[1] // block_map.shape[3]
+    if q_block_size == kv_block_size == _FA4_Q_BLOCK_SIZE:
+        return _cute_attention_q128(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
+    need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
+    forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
+        block_map,
+        variable_block_sizes,
+        q_len=q_bshd.shape[1],
+        q_block_size=q_block_size,
+        kv_block_size=kv_block_size,
+        need_backward=need_backward,
+    )
+    return flash_attn_func(
+        q_bshd,
+        k_bshd,
+        v_bshd,
+        mask_mod=_build_vbs_mask_mod(kv_block_size),
+        aux_tensors=[variable_block_sizes],
+        block_sparse_tensors=forward_sparse_tensors,
+        block_sparse_tensors_bwd=backward_sparse_tensors,
+        return_lse=True,
+    )
+
+
 def block_sparse_attn_cute_fwd(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -194,34 +363,25 @@ def block_sparse_attn_cute_fwd(
     block_map: torch.Tensor,
     variable_block_sizes: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """CuTe forward-only block-sparse attention with [B, H, S, D] inputs."""
+    """Autograd-enabled CuTe block-sparse attention for [B, H, S, D]."""
     if block_map.dim() == 3:
         block_map = block_map.unsqueeze(0)
-    q_block_size = q.shape[2] // block_map.shape[2]
-    kv_block_size = k.shape[2] // block_map.shape[3]
 
     q_bshd = q.transpose(1, 2).contiguous()
     k_bshd = k.transpose(1, 2).contiguous()
     v_bshd = v.transpose(1, 2).contiguous()
-    out_bshd, lse_bshd = _cute_forward(
+    out_bshd, lse = _cute_attention(
         q_bshd,
         k_bshd,
         v_bshd,
         block_map,
         variable_block_sizes,
-        q_block_size=q_block_size,
-        kv_block_size=kv_block_size,
     )
     out = out_bshd.transpose(1, 2).contiguous()
-    if lse_bshd is None:
-        lse = torch.empty(
-            (q.shape[0], q.shape[1], q.shape[2]),
-            dtype=torch.float32,
-            device=q.device,
-        )
-    else:
-        lse = lse_bshd.transpose(1, 2).contiguous()
-    return out, lse
+    # FA4 already returns lse as [B, H, S], matching the Triton path's aux
+    # contract, so it needs no transpose. Detach before any further op: the
+    # value is informational and callers never backprop through it.
+    return out, lse.detach()
 
 
 def block_sparse_attn_cute_fwd_bshd(
@@ -231,27 +391,16 @@ def block_sparse_attn_cute_fwd_bshd(
     block_map: torch.Tensor,
     variable_block_sizes: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """CuTe forward-only block-sparse attention with [B, S, H, D] inputs."""
+    """Autograd-enabled CuTe block-sparse attention for [B, S, H, D]."""
     if block_map.dim() == 3:
         block_map = block_map.unsqueeze(0)
-    q_block_size = q.shape[1] // block_map.shape[2]
-    kv_block_size = k.shape[1] // block_map.shape[3]
 
-    out, lse_bshd = _cute_forward(
+    out, lse = _cute_attention(
         q,
         k,
         v,
         block_map,
         variable_block_sizes,
-        q_block_size=q_block_size,
-        kv_block_size=kv_block_size,
     )
-    if lse_bshd is None:
-        lse = torch.empty(
-            (q.shape[0], q.shape[2], q.shape[1]),
-            dtype=torch.float32,
-            device=q.device,
-        )
-    else:
-        lse = lse_bshd.transpose(1, 2).contiguous()
-    return out, lse
+    # lse is [B, H, S] regardless of the q/k/v layout; see above.
+    return out, lse.detach()

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+"""sm_100a (Blackwell) CUDA block-sparse VSA forward.
+
+A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: it
+returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward``
+writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` fp32 -- so
+``block_sparse_attn_backward_triton`` runs against it unchanged.
+
+The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks
+(tile volumes 64 and 128 in ``build_vsa_metadata``); the block size is inferred from the
+tensors and picks the op. Anything else falls back to Triton via ``is_supported``.
+"""
+
+from typing import Tuple
+
+import torch
+
+try:
+    # The pybind symbols live on fastvideo_kernel_ops, NOT on the _C package that contains it.
+    # `import fastvideo_kernel._C as _C` resolves to the namespace package, whose __init__ is
+    # empty, so hasattr() fails on a wheel install and the caller silently falls back with the
+    # kernel built and present.
+    from fastvideo_kernel._C import fastvideo_kernel_ops as _C
+    _FWD_BY_BLOCK = {
+        64: getattr(_C, "block_sparse_sm100a_fwd", None),
+        128: getattr(_C, "block_sparse_sm100a_blk128_fwd", None),
+    }
+    _HAS_VSA_SM100A = any(_FWD_BY_BLOCK.values())
+except ImportError:  # pragma: no cover - extension not built
+    _C = None
+    _FWD_BY_BLOCK = {}
+    _HAS_VSA_SM100A = False
+
+_SM100 = (10, 0)
+HEAD_DIM = 128
+# Must match the -DVSA_BHSD the extension was compiled with (see CMakeLists).
+BHSD = True
+
+
+def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
+    num_blocks = variable_block_sizes.numel()
+    seqlen = q.shape[2] if BHSD else q.shape[1]
+    return 0 if num_blocks == 0 or seqlen % num_blocks else seqlen // num_blocks
+
+
+def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
+    """True iff this build can run these tensors; otherwise the caller uses Triton.
+
+    Static facts only -- shapes, dtypes, arch, layout. Deliberately NO reads of tensor
+    contents: the previous ``int(variable_block_sizes.min())`` was a GPU->CPU sync on every
+    call, and the kernel no longer needs it (see below). This predicate must stay cheap
+    enough to sit on a per-layer dispatch path.
+
+    What the kernel accepts (and is tested to handle):
+      * q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100 device, head_dim 128, laid out
+        as compiled (BHSD here); seqlen == num_blocks * block with an EVEN num_blocks (a CTA
+        owns an adjacent pair of query blocks) and a 64- or 128-token build present.
+      * q2k_num: any per-row counts in [0, max_kv], NON-uniform across rows included. Rows
+        with count 0 produce exactly-zero output rows (and a finite LSE sentinel) rather
+        than attending anywhere -- so no ``.min()`` floor is required of the caller.
+      * q2k_idx: rows only need valid entries (in [0, num_blocks)) BELOW that row's count;
+        padding past the count (e.g. map_to_index's -1 fill) is never dereferenced. max_kv
+        (= q2k_idx.shape[-1]) must be >= 1, which the host launcher re-checks.
+      * variable_block_sizes: per-KV-block valid-token counts in [0, block]; keys at or past
+        a block's count are masked. Integer metadata is converted to int32/contiguous by
+        ``block_sparse_attn_sm100a`` itself, so int64 inputs merely cost a cast.
+    """
+    if not _HAS_VSA_SM100A or not q.is_cuda:
+        return False
+    if torch.cuda.get_device_capability(q.device) != _SM100:
+        return False
+    if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM:
+        return False
+    if not q.is_contiguous():
+        return False
+    if _FWD_BY_BLOCK.get(_block_size(q, variable_block_sizes)) is None:
+        return False
+    # A CTA owns an adjacent pair of query blocks.
+    if variable_block_sizes.numel() % 2 != 0:
+        return False
+    # Metadata must be integer-typed so the wrapper's int32 conversion is value-preserving.
+    if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32, torch.int64):
+        return False
+    return True
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_sm100a_inference",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _block_sparse_attn_sm100a_inference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    """Opaque no-LSE launch used by the inference-only sm_100a route.
+
+    The extension is exposed as a raw pybind function rather than a dispatcher
+    op. Calling it directly makes Dynamo descend through a Python/C++ boundary
+    that has no fake implementation, so ``torch.compile(fullgraph=True)``
+    cannot capture a sparse H3 block. Keep that boundary inside this custom op;
+    its inputs have already been normalized by the public wrapper below.
+    """
+    fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)]
+    sm_scale = 1.0 / (q.shape[-1]**0.5)
+    res = fwd(q, k, v, None, q2k_idx, q2k_num, variable_block_sizes, sm_scale, False)
+    return res[0]
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_inference")
+def _block_sparse_attn_sm100a_inference_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    # The C++ binding allocates its output with torch::empty_like(q).
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _block_sparse_attn_sm100a_from_mask_inference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    """Opaque mask compaction plus no-LSE sm_100a launch for inference.
+
+    H3 naturally produces a bool block map. Its Triton ``map_to_index`` call
+    must live behind the same opaque boundary as the raw pybind launch;
+    otherwise Dynamo sees that kernel before reaching the index-native custom
+    op and full-graph capture still fails.
+    """
+    from fastvideo_kernel.triton_kernels.index import map_to_index
+
+    q2k_idx, q2k_num = map_to_index(block_map)
+    return _block_sparse_attn_sm100a_inference(
+        q,
+        k,
+        v,
+        q2k_idx.to(torch.int32).contiguous(),
+        q2k_num.to(torch.int32).contiguous(),
+        variable_block_sizes,
+    )
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference")
+def _block_sparse_attn_sm100a_from_mask_inference_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty_like(q)
+
+
+def block_sparse_attn_sm100a_from_mask(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, None]:
+    """Inference forward from a bool block map, with compaction kept opaque."""
+    out = _block_sparse_attn_sm100a_from_mask_inference(
+        q.contiguous(),
+        k.contiguous(),
+        v.contiguous(),
+        block_map.to(torch.bool).contiguous(),
+        variable_block_sizes.to(torch.int32).contiguous(),
+    )
+    return out, None
+
+
+def block_sparse_attn_sm100a(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    need_lse: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor | None]:
+    """Forward pass. Returns ``(out, lse)``; ``out`` has q's layout."""
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    idx = q2k_idx.to(torch.int32).contiguous()
+    num = q2k_num.to(torch.int32).contiguous()
+    vbs = variable_block_sizes.to(torch.int32).contiguous()
+
+    if not need_lse:
+        # This is the production inference path. The custom op keeps the raw
+        # pybind launch opaque to Dynamo while its fake kernel carries output
+        # metadata through full-graph capture.
+        return _block_sparse_attn_sm100a_inference(q, k, v, idx, num, vbs), None
+
+    # Preserve the established LSE-producing path for correctness tests and
+    # any future forward/backward pairing; only inference needs the opaque op.
+    fwd = _FWD_BY_BLOCK[_block_size(q, vbs)]
+    sm_scale = 1.0 / (q.shape[-1]**0.5)
+    res = fwd(q, k, v, None, idx, num, vbs, sm_scale, True)
+    return res[0], res[1]

fastvideo-kernel/python/fastvideo_kernel/comm_ops.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Ulysses sequence-parallel all-to-all ops.
+
+Thin wrappers over csrc/comm/ulysses_all_to_all.cu. The kernel stores directly
+into peers' memory through NCCL's device API, so the caller supplies an
+ncclComm_t for the group.
+"""
+
+import torch
+
+try:
+    from fastvideo_kernel._C import fastvideo_kernel_ops as _ops
+except ImportError:  # pragma: no cover - no compiled extension in this install
+    _ops = None
+
+_SUPPORTED_WORLD_SIZES = (2, 4, 6, 8)
+_REQUIRED_OPS = (
+    "allocate_ulysses_a2a",
+    "register_ulysses_a2a_window",
+    "create_ulysses_a2a_dev_comm",
+    "dispose_ulysses_a2a",
+    "ulysses_lsa_covers_group",
+    "ulysses_a2a",
+)
+
+
+def is_available() -> bool:
+    """Whether this wheel was built with the Ulysses all-to-all kernel."""
+    return _ops is not None and all(hasattr(_ops, name) for name in _REQUIRED_OPS)
+
+
+def _require() -> None:
+    if not is_available():
+        raise RuntimeError(
+            "the Ulysses all-to-all kernel is not present in this fastvideo-kernel build; "
+            "rebuild with ./build.sh or install a wheel that includes csrc/comm/")
+
+
+def lsa_covers_group(comm_ptr: int, world_size: int) -> bool:
+    """Whether every rank in the group is load-store accessible to every other."""
+    _require()
+    return bool(_ops.ulysses_lsa_covers_group(int(comm_ptr), int(world_size)))
+
+
+def allocate(nbytes: int, rank: int, world_size: int, device_index: int) -> int:
+    """Allocate one rank's local symmetric window without a collective."""
+    _require()
+    if world_size not in _SUPPORTED_WORLD_SIZES:
+        raise ValueError(f"ulysses a2a supports world sizes {_SUPPORTED_WORLD_SIZES}, "
+                         f"got {world_size}")
+    return int(_ops.allocate_ulysses_a2a(int(nbytes), int(rank), int(world_size), int(device_index)))
+
+
+def register_window(handle: int, comm_ptr: int) -> None:
+    """Register an allocated window with the supplied communicator.
+
+    Every rank in ``comm_ptr`` must call this together.
+    """
+    _require()
+    _ops.register_ulysses_a2a_window(int(handle), int(comm_ptr))
+
+
+def create_dev_comm(handle: int) -> None:
+    """Create the device communicator for a registered window collectively."""
+    _require()
+    _ops.create_ulysses_a2a_dev_comm(int(handle))
+
+
+def dispose(handle: int) -> None:
+    """Release a handle from :func:`allocate`. It is dangling afterwards."""
+    _require()
+    _ops.dispose_ulysses_a2a(int(handle))
+
+
+def all_to_all(handle: int, inp: torch.Tensor, out: torch.Tensor, B: int, S_local: int, H: int,
+               D: int, mode: int) -> None:
+    """Run one fused all-to-all on the current stream, writing into ``out``.
+
+    ``mode == 0``: ``[B, S_local, H, D] -> [B, S_global, H_local, D]``
+    ``mode == 1``: ``[B, S_global, H_local, D] -> [B, S_local, H, D]``
+
+    ``H`` is the global head count. Every rank must call with consistent
+    geometry in the same order.
+    """
+    _require()
+    _ops.ulysses_a2a(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode))

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -2,6 +2,8 @@ import math
 import torch
 from .block_sparse_attn import block_sparse_attn
 from .block_sparse_attn_256 import (
+    block_sparse_attn_128,
+    block_sparse_attn_128_bshd,
     block_sparse_attn_256,
     block_sparse_attn_256_bshd,
 )
@@ -74,12 +76,13 @@ def video_sparse_attn(
 
     Dispatches the sparse branch by ``block_elements = prod(block_size)``:
     - 64  -> existing TK/Triton path (see ``block_sparse_attn_from_indices``).
+    - 128 -> Triton fallback or CuTe FA4 block-sparse attention.
     - 256 -> CuTe FA4 block-sparse attention (see ``block_sparse_attn_256``).
 
     Backend overrides:
     - ``FASTVIDEO_VSA_TRITON=1`` forces Triton in either path.
     - ``FASTVIDEO_VSA_TK=1`` prefers the sm_90 TK kernel in the 64-block path.
-    - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 256-block path.
+    - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 128/256-block paths.
     """
     if isinstance(block_size, int):
         block_size = (block_size, block_size, block_size)
@@ -119,8 +122,9 @@ def video_sparse_attn(
     # Sparse branch (fused Triton topk mask)
     mask = fused_topk_mask(scores, topk)
 
-    if block_elements == 256:
-        out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0]
+    if block_elements in (128, 256):
+        attention = block_sparse_attn_128 if block_elements == 128 else block_sparse_attn_256
+        out_s = attention(q, k, v, mask, variable_block_sizes)[0]
     else:
         out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
 
@@ -142,14 +146,14 @@ def video_sparse_attn_bshd(
     """VSA entrypoint for [B, S, H, D] tensors.
 
     Avoids the BHSD<->BSHD round-trip that ``video_sparse_attn`` performs on
-    the CuTe 256-block path; the 64-block path still expects BHSD and is not
+    the CuTe 128/256-block paths; the 64-block path still expects BHSD and is not
     supported here.
     """
     if isinstance(block_size, int):
         block_size = (block_size, block_size, block_size)
     block_elements = block_size[0] * block_size[1] * block_size[2]
-    if block_elements != 256:
-        raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 "
+    if block_elements not in (128, 256):
+        raise ValueError("video_sparse_attn_bshd is only defined for block_elements=128 or 256 "
                          f"(got {block_elements}); use video_sparse_attn for the 64-block path.")
 
     batch, q_seq_len, heads, dim = q.shape
@@ -171,19 +175,15 @@ def video_sparse_attn_bshd(
         raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
                          f"got {q_variable_block_sizes.numel()}")
 
-    # Compression branch (BSHD-native: mean over the 256-token axis).
-    token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32)
-    q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1,
-                                                                         1)).view(1, q_num_blocks, block_elements, 1, 1)
-    kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1,
-                                                                        1)).view(1, kv_num_blocks, block_elements, 1, 1)
-
+    # Compression branch (BSHD-native: match fused_block_mean's semantics).
+    # Padding values are expected to be zero; gradients are broadcast across
+    # the full padded block, just like the BHSD fused common path.
     q_c = q.view(batch, q_num_blocks, block_elements, heads, dim)
     k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim)
     v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim)
-    q_c = ((q_c.float() * q_token_valid).sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
-    k_c = ((k_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
-    v_c = ((v_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
+    q_c = (q_c.float().sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
+    k_c = (k_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
+    v_c = (v_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
     q_ch = q_c.permute(0, 2, 1, 3).contiguous()
     k_ch = k_c.permute(0, 2, 1, 3).contiguous()
     v_ch = v_c.permute(0, 2, 1, 3).contiguous()
@@ -195,13 +195,15 @@ def video_sparse_attn_bshd(
 
     # Sparse branch (fused Triton topk mask + CuTe BSHD).
     mask = fused_topk_mask(scores, topk)
-    out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes)
+    attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd
+    out_s, _ = attention(q, k, v, mask, variable_block_sizes)
 
-    out = out_s
-    out_view = out.view(batch, q_num_blocks, block_elements, heads, dim)
+    # Out-of-place: ``out_s`` is the tensor FA4's autograd node saved for its
+    # backward, so mutating it in place invalidates the graph.
+    out_view = out_s.view(batch, q_num_blocks, block_elements, heads, dim)
     if compress_attn_weight is not None:
         gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim)
-        out_view.add_(out_c_blk.unsqueeze(2) * gate_view)
+        out = out_view + out_c_blk.unsqueeze(2) * gate_view
     else:
-        out_view.add_(out_c_blk.unsqueeze(2))
-    return out
+        out = out_view + out_c_blk.unsqueeze(2)
+    return out.view(batch, q_seq_len, heads, dim)

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

@@ -16,14 +16,21 @@ import triton.language as tl
 import math  # small utility needed by the sparse wrapper
 # ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
 
-# We don't run auto-tuning every time to keep the tutorial fast. Keeping
-# the code below and commenting out the equivalent parameters is convenient for
-# re-tuning.
+# BLOCK_M / BLOCK_N are fixed at 64 because they are structural, not tunable:
+# the kernel indexes the top-k list per BLOCK_M q-tile and addresses keys as
+# kv_idx * BLOCK_N, so both must match the granularity q2k_index and
+# variable_block_sizes were built at.
+#
+# num_stages / num_warps ARE free, and the previous {3, 4, 7} was inherited from
+# the upstream tutorial rather than tuned here. It skips 5 and 6; on Blackwell
+# (sm_121) the optimum is num_stages=5, so the search could not reach it. Both
+# block paths independently select 5 once it is available. Autotune still picks
+# per architecture, so other GPUs re-tune rather than inheriting this choice.
 configs = [
     triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
     for BM in [64]\
     for BN in [64]\
-    for s in [3, 4, 7]\
+    for s in [2, 3, 4, 5, 6, 7]\
     for w in [4, 8]\
 ]
 
@@ -237,7 +244,12 @@ def _attn_bwd_dkdv(
         # Load m before computing qk to reduce pipeline stall.
         offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
         m = tl.load(M + offs_m)
-        qkT = tl.dot(k, qT)
+        # Recompute logits exactly as the forward does: raw bf16 operands into
+        # the dot, fp32 scale after accumulation. A bf16 pre-scaled K perturbs
+        # the recomputed logits relative to the saved M by an error
+        # proportional to |logit|, which exp2 amplifies into arbitrarily wrong
+        # probabilities at large activations.
+        qkT = tl.dot(k, qT) * (sm_scale * 1.4426950408889634)
         pT = tl.math.exp2(qkT - m[None, :])
         mask = tl.arange(0, BLOCK_N1) < block_size
         pT = tl.where(mask[:, None], pT, 0.0)
@@ -268,6 +280,7 @@ def _attn_bwd_dq(
         do,
         m,
         D,
+        sm_scale,
         # shared by Q/K/V/DO.
         q2k_index,
         q2k_num,
@@ -315,7 +328,7 @@ def _attn_bwd_dq(
         block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
-        qk = tl.dot(q, kT)
+        qk = tl.dot(q, kT) * (sm_scale * 1.4426950408889634)
         p = tl.math.exp2(qk - m)
         offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
         mask = offs_in_block < block_size
@@ -324,8 +337,7 @@ def _attn_bwd_dq(
         dp = tl.dot(do, vT).to(tl.float32)
         ds = p * (dp - Di[:, None])
         ds = ds.to(tl.bfloat16)
-        # Compute dQ.
-        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        # Compute dQ (kT is raw; the caller applies sm_scale once at the end).
         dq += tl.dot(ds, tl.trans(kT))
         # Increment pointers.
     return dq
@@ -453,6 +465,7 @@ def _attn_bwd(
         do,
         m,
         D,  #
+        sm_scale,
         q2k_index,
         q2k_num,
         max_kv_blks,
@@ -470,7 +483,7 @@ def _attn_bwd(
     )
     # Write back dQ.
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq *= LN2
+    dq *= sm_scale
     tl.store(dq_ptrs, dq)
 
 
@@ -591,6 +604,7 @@ def _attn_bwd_dq_kernel(
         Q,
         K,
         V,
+        sm_scale,
         DO,  #
         DQ,
         M,
@@ -663,6 +677,7 @@ def _attn_bwd_dq_kernel(
         do,
         m,
         D,
+        sm_scale,
         q2k_index,
         q2k_num,
         max_kv_blks,
@@ -680,7 +695,7 @@ def _attn_bwd_dq_kernel(
     )
 
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= LN2
+    dq_acc *= sm_scale
     tl.store(dq_ptrs, dq_acc)
 
 
@@ -748,9 +763,11 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
     dv = torch.empty_like(v)
     BATCH, N_HEAD = q.shape[:2]
     BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
-    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+    # K stays raw: the backward kernels apply sm_scale in fp32 after the dot,
+    # matching the forward's rounding exactly. (A bf16 pre-scaled K perturbs
+    # the recomputed logits vs the saved M; exp2 turns that into unboundedly
+    # wrong probabilities at large activations.)
     arg_k = k
-    arg_k = arg_k * (sm_scale * RCP_LN2)
     PRE_BLOCK = 64
     assert Tq % PRE_BLOCK == 0
     pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
@@ -813,6 +830,7 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
         q,
         arg_k,
         v,
+        sm_scale,
         do,
         dq,
         M,

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.3.2"
+__version__ = "0.3.4"

fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py

@@ -14,7 +14,9 @@ import math
 import torch
 
 VSA_TILE_SIZE = (4, 4, 4)
-_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 256)
+# 128 is served by the sm_100a CUDA backend (block_sparse_attn_sm100a); 64 and 256 by
+# Triton and the CuTe-DSL path. A volume here only needs a backend that accepts it.
+_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 128, 256)
 
 
 def _canonicalize_device(device: torch.device | str) -> torch.device:

fastvideo-kernel/tests/test_vsa128_backward.py

@@ -0,0 +1,146 @@
+"""VSA-128 CuTe/Triton forward and backward parity on Blackwell."""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+from fastvideo_kernel import video_sparse_attn, video_sparse_attn_bshd
+from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128
+
+from .test_vsa256_triton import _metrics, _torch_vsa256_reference
+
+_BLOCK = 128
+_BLOCK_SIZE_3D = (2, 8, 8)
+
+
+def _select_backend(monkeypatch, backend: str) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    if backend == "cute":
+        pytest.importorskip(
+            "flash_attn.cute.block_sparsity",
+            reason="optional FA4 CuTe build (flash_attn.cute) not installed",
+        )
+        monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
+        monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
+        monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
+    else:
+        monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
+        monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False)
+
+
+def _dense_sparse_reference(q, k, v, block_map, variable_block_sizes):
+    token_mask = block_map.repeat_interleave(_BLOCK, dim=2).repeat_interleave(_BLOCK, dim=3)
+    kv_valid = torch.arange(_BLOCK, device=k.device) < variable_block_sizes[:, None]
+    token_mask = token_mask & kv_valid.reshape(1, 1, 1, -1)
+    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) / math.sqrt(q.shape[-1])
+    probabilities = torch.softmax(logits.masked_fill(~token_mask, float("-inf")), dim=-1)
+    return torch.matmul(probabilities, v.float()).to(q.dtype)
+
+
+def _check(tag: str, expected: torch.Tensor, actual: torch.Tensor, avg_tol: float, rel_tol: float) -> None:
+    assert torch.isfinite(actual).all().item(), f"{tag}: non-finite values"
+    avg_abs, max_rel = _metrics(expected, actual)
+    print(f"  {tag}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < avg_tol
+    assert max_rel < rel_tol
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("backend", ["cute", "triton"])
+def test_vsa128_explicit_routes_forward_backward(backend: str, monkeypatch) -> None:
+    """Adjacent Q128 blocks must keep independent routes instead of merging."""
+    _select_backend(monkeypatch, backend)
+    torch.manual_seed(53)
+    shape = (1, 1, 3 * _BLOCK, 128)
+    base = [torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    grad_output = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    variable_block_sizes = torch.tensor([128, 91, 37], device="cuda", dtype=torch.int32)
+    block_map = torch.eye(3, device="cuda", dtype=torch.bool).view(1, 1, 3, 3)
+
+    actual_inputs = [tensor.detach().clone().requires_grad_() for tensor in base]
+    actual, _ = block_sparse_attn_128(*actual_inputs, block_map, variable_block_sizes)
+    (actual * grad_output).sum().backward()
+
+    reference_inputs = [tensor.detach().clone().requires_grad_() for tensor in base]
+    expected = _dense_sparse_reference(*reference_inputs, block_map, variable_block_sizes)
+    (expected * grad_output).sum().backward()
+
+    print(f"[vsa128-explicit-{backend}]")
+    _check("out", expected, actual, 1e-3, 0.2)
+    for name, reference, candidate in zip(("dq", "dk", "dv"), reference_inputs, actual_inputs, strict=True):
+        _check(name, reference.grad, candidate.grad, 2e-2, 0.5)
+
+
+def _zero_kv_tail(x: torch.Tensor, variable_block_sizes: torch.Tensor) -> torch.Tensor:
+    valid = torch.arange(_BLOCK, device=x.device) < variable_block_sizes[:, None]
+    valid = valid.view(1, 1, -1, _BLOCK, 1).expand_as(x.view(1, x.shape[1], -1, _BLOCK, x.shape[-1]))
+    return x * valid.reshape_as(x).to(x.dtype)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("backend", ["cute", "triton"])
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa128_wrapper_forward_backward(backend: str, layout: str, monkeypatch) -> None:
+    _select_backend(monkeypatch, backend)
+    torch.manual_seed(59)
+    batch, heads, dim = 1, 2, 128
+    q_blocks, kv_blocks, topk = 3, 4, 2
+    q_shape = (batch, heads, q_blocks * _BLOCK, dim)
+    kv_shape = (batch, heads, kv_blocks * _BLOCK, dim)
+    q_base = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
+    kv_sizes = torch.tensor([128, 91, 37, 128], device="cuda", dtype=torch.int32)
+    q_sizes = torch.full((q_blocks, ), _BLOCK, device="cuda", dtype=torch.int32)
+    k_base = _zero_kv_tail(torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16), kv_sizes)
+    v_base = _zero_kv_tail(torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16), kv_sizes)
+    gate_base = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16) * 0.1
+    grad_output = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
+
+    if layout == "bhsd":
+        actual_inputs = [tensor.detach().clone().requires_grad_() for tensor in (q_base, k_base, v_base)]
+        actual_gate = gate_base.detach().clone().requires_grad_()
+        actual = video_sparse_attn(
+            *actual_inputs,
+            kv_sizes,
+            q_sizes,
+            topk,
+            block_size=_BLOCK_SIZE_3D,
+            compress_attn_weight=actual_gate,
+        )
+        actual_grads = actual_inputs
+    else:
+        bshd_inputs = [tensor.transpose(1, 2).contiguous().detach().requires_grad_()
+                       for tensor in (q_base, k_base, v_base)]
+        bshd_gate = gate_base.transpose(1, 2).contiguous().detach().requires_grad_()
+        actual = video_sparse_attn_bshd(
+            *bshd_inputs,
+            kv_sizes,
+            q_sizes,
+            topk,
+            block_size=_BLOCK_SIZE_3D,
+            compress_attn_weight=bshd_gate,
+        ).transpose(1, 2)
+        actual_grads = bshd_inputs
+    (actual * grad_output).sum().backward()
+
+    reference_inputs = [tensor.detach().clone().requires_grad_() for tensor in (q_base, k_base, v_base)]
+    reference_gate = gate_base.detach().clone().requires_grad_()
+    expected = _torch_vsa256_reference(
+        *reference_inputs,
+        q_sizes,
+        kv_sizes,
+        topk,
+        compress_attn_weight=reference_gate,
+    )
+    (expected * grad_output).sum().backward()
+
+    print(f"[vsa128-wrapper-{backend}-{layout}]")
+    _check("out", expected, actual, 1e-3, 0.2)
+    for name, reference, candidate in zip(("dq", "dk", "dv"), reference_inputs, actual_grads, strict=True):
+        candidate_grad = candidate.grad if layout == "bhsd" else candidate.grad.transpose(1, 2)
+        _check(name, reference.grad, candidate_grad, 2e-2, 0.5)
+    actual_gate_grad = actual_gate.grad if layout == "bhsd" else bshd_gate.grad.transpose(1, 2)
+    _check("dgate", reference_gate.grad, actual_gate_grad, 1e-3, 0.2)

fastvideo-kernel/tests/test_vsa256_backward.py

@@ -0,0 +1,224 @@
+"""VSA-256 FA4 CuTe forward/backward parity for BHSD and BSHD APIs.
+
+Covers the shapes the CuTe backward actually sees in production: the gated
+compression branch (`compress_attn_weight`), partially filled Q tiles,
+and q_len != kv_len. Also pins the inference fast path, which must skip the
+KV-owned backward metadata without changing the forward result.
+"""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import pytest
+import torch
+
+from fastvideo_kernel import video_sparse_attn, video_sparse_attn_bshd
+
+from .test_vsa256_triton import _metrics, _torch_vsa256_reference
+
+_BLOCK = 256
+_BLOCK_SIZE_3D = (4, 8, 8)  # prod == 256
+
+# Measured on GB200 (sm_100) with bf16 inputs: grads land around 1e-4 avg_abs
+# and <=0.11 max_rel across every case below, so these leave ~10x headroom
+# without being loose enough to hide a real regression.
+_OUT_TOL = (1e-3, 0.2)
+_GRAD_TOL = (1e-3, 0.25)
+
+
+@pytest.fixture(autouse=True)
+def _require_cute_backend(monkeypatch):
+    pytest.importorskip(
+        "flash_attn.cute.block_sparsity",
+        reason="optional FA4 CuTe build (flash_attn.cute) not installed",
+    )
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
+    monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
+    monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
+
+
+def _zero_pad_tail(x: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
+    """Zero the padded tail of every 256-token tile of a [B, H, S, D] tensor.
+
+    VSA callers scatter into a zeroed tile buffer, so padded slots are zero;
+    both the kernel and the reference rely on that.
+    """
+    bsz, heads, _, dim = x.shape
+    blocks = var.numel()
+    token_idx = torch.arange(_BLOCK, device=x.device, dtype=torch.int32)
+    valid = (token_idx.view(1, -1) < var.view(-1, 1)).view(1, 1, blocks, _BLOCK, 1)
+    valid = valid.expand(bsz, heads, blocks, _BLOCK, dim).reshape_as(x)
+    return x * valid.to(x.dtype)
+
+
+def _make_inputs(
+    q_blocks: int,
+    kv_blocks: int,
+    kv_var: torch.Tensor,
+    q_var: torch.Tensor,
+    heads: int = 2,
+    dim: int = 128,
+    seed: int = 42,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    torch.manual_seed(seed)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    sq, skv = q_blocks * _BLOCK, kv_blocks * _BLOCK
+    q = torch.randn(1, heads, sq, dim, device=device, dtype=dtype)
+    k = torch.randn(1, heads, skv, dim, device=device, dtype=dtype)
+    v = torch.randn(1, heads, skv, dim, device=device, dtype=dtype)
+    grad_out = torch.randn_like(q)
+    return _zero_pad_tail(q, q_var), _zero_pad_tail(k, kv_var), _zero_pad_tail(v, kv_var), grad_out
+
+
+def _check(tag: str, ref: torch.Tensor, got: torch.Tensor, tol: Tuple[float, float]) -> None:
+    assert torch.isfinite(got).all().item(), f"{tag}: non-finite values"
+    avg_abs, max_rel = _metrics(ref, got)
+    print(f"  {tag}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < tol[0], f"{tag}: avg_abs {avg_abs:.3e} >= {tol[0]:.3e}"
+    assert max_rel < tol[1], f"{tag}: max_rel {max_rel:.3e} >= {tol[1]:.3e}"
+
+
+def _run_bhsd(q, k, v, kv_var, q_var, topk, gate=None):
+    qg, kg, vg = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out = video_sparse_attn(qg, kg, vg, kv_var, q_var, topk, block_size=_BLOCK_SIZE_3D, compress_attn_weight=gate)
+    return out, (qg, kg, vg)
+
+
+def _run_bshd(q, k, v, kv_var, q_var, topk, gate=None):
+    qg, kg, vg = (t.transpose(1, 2).contiguous().requires_grad_(True) for t in (q, k, v))
+    gate_bshd = None if gate is None else gate.transpose(1, 2).contiguous()
+    out = video_sparse_attn_bshd(qg,
+                                 kg,
+                                 vg,
+                                 kv_var,
+                                 q_var,
+                                 topk,
+                                 block_size=_BLOCK_SIZE_3D,
+                                 compress_attn_weight=gate_bshd)
+    return out.transpose(1, 2), (qg, kg, vg)
+
+
+def _reference(q, k, v, q_var, kv_var, topk, gate=None):
+    qr, kr, vr = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out = _torch_vsa256_reference(qr, kr, vr, q_var, kv_var, topk, compress_attn_weight=gate)
+    return out, (qr, kr, vr)
+
+
+def _compare(tag, layout, q, k, v, kv_var, q_var, topk, grad_out, gate=None):
+    runner = _run_bhsd if layout == "bhsd" else _run_bshd
+    out, (qg, kg, vg) = runner(q, k, v, kv_var, q_var, topk, gate=gate)
+    (out * grad_out).sum().backward()
+    grads = [g.grad if g.grad.dim() == 4 and layout == "bhsd" else g.grad for g in (qg, kg, vg)]
+    if layout == "bshd":
+        grads = [g.transpose(1, 2) for g in grads]
+
+    out_ref, refs = _reference(q, k, v, q_var, kv_var, topk, gate=gate)
+    (out_ref * grad_out).sum().backward()
+
+    print(f"[{tag}-{layout}]")
+    _check("out", out_ref, out, _OUT_TOL)
+    for name, ref, got in zip(("dq", "dk", "dv"), refs, grads):
+        _check(name, ref.grad, got, _GRAD_TOL)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_forward_backward_vs_torch_ref(layout: str) -> None:
+    kv_var = torch.tensor([256, 173, 79, 256], dtype=torch.int32, device="cuda")
+    q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(3, 4, kv_var, q_var)
+    _compare("vsa256-cute", layout, q, k, v, kv_var, q_var, 2, grad_out)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_backward_with_compress_gate(layout: str) -> None:
+    """The gated compression branch is what Wan and MiniMax-H3 actually run.
+
+    It is also the branch that composes the sparse output with the compression
+    output, so it is the one that breaks if that composition mutates FA4's
+    saved output in place.
+    """
+    kv_var = torch.tensor([256, 200, 256, 91], dtype=torch.int32, device="cuda")
+    q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(3, 4, kv_var, q_var, seed=7)
+    gate = torch.randn_like(q) * 0.1
+    _compare("vsa256-cute-gated", layout, q, k, v, kv_var, q_var, 2, grad_out, gate=gate)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_backward_partial_q_blocks(layout: str) -> None:
+    """Q tiles that are not full: only the compression divisor depends on it,
+    but it is the one axis the existing coverage held constant."""
+    kv_var = torch.tensor([256, 128, 256], dtype=torch.int32, device="cuda")
+    q_var = torch.tensor([256, 61, 199], dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(3, 3, kv_var, q_var, seed=11)
+    _compare("vsa256-cute-partial-q", layout, q, k, v, kv_var, q_var, 2, grad_out)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
+def test_vsa256_cute_backward_cross_q_kv(layout: str) -> None:
+    """q_len != kv_len: forward has coverage, backward did not."""
+    kv_var = torch.tensor([256, 143, 256, 256, 88], dtype=torch.int32, device="cuda")
+    q_var = torch.full((2, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, grad_out = _make_inputs(2, 5, kv_var, q_var, seed=13)
+    _compare("vsa256-cute-cross", layout, q, k, v, kv_var, q_var, 3, grad_out)
+
+
+@pytest.mark.cuda
+def test_vsa256_cute_inference_matches_training_forward() -> None:
+    """The KV-owned backward metadata is only built when something requires
+    grad. Skipping it must not perturb the forward result."""
+    kv_var = torch.tensor([256, 173, 79, 256], dtype=torch.int32, device="cuda")
+    q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
+    q, k, v, _ = _make_inputs(3, 4, kv_var, q_var, seed=5)
+
+    with torch.no_grad():
+        out_infer = video_sparse_attn_bshd(
+            q.transpose(1, 2).contiguous(),
+            k.transpose(1, 2).contiguous(),
+            v.transpose(1, 2).contiguous(),
+            kv_var,
+            q_var,
+            2,
+            block_size=_BLOCK_SIZE_3D,
+            compress_attn_weight=None,
+        )
+
+    out_train, _ = _run_bshd(q, k, v, kv_var, q_var, 2)
+    torch.testing.assert_close(out_infer, out_train.transpose(1, 2).detach(), rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+def test_vsa256_cute_lse_is_bhs() -> None:
+    """The aux return is [B, H, S] on both entrypoints, matching the Triton
+    path's contract."""
+    from fastvideo_kernel.block_sparse_attn_256 import (block_sparse_attn_256, block_sparse_attn_256_bshd)
+
+    device = torch.device("cuda")
+    heads, dim, q_blocks, kv_blocks = 2, 128, 3, 4
+    sq, skv = q_blocks * _BLOCK, kv_blocks * _BLOCK
+    q = torch.randn(1, heads, sq, dim, device=device, dtype=torch.bfloat16)
+    k = torch.randn(1, heads, skv, dim, device=device, dtype=torch.bfloat16)
+    v = torch.randn(1, heads, skv, dim, device=device, dtype=torch.bfloat16)
+    vbs = torch.full((kv_blocks, ), _BLOCK, dtype=torch.int32, device=device)
+    mask = torch.zeros(1, heads, q_blocks, kv_blocks, dtype=torch.bool, device=device)
+    mask[..., :2] = True
+
+    _, lse_bhsd = block_sparse_attn_256(q, k, v, mask, vbs)
+    assert lse_bhsd.shape == (1, heads, sq), lse_bhsd.shape
+
+    _, lse_bshd = block_sparse_attn_256_bshd(
+        q.transpose(1, 2).contiguous(),
+        k.transpose(1, 2).contiguous(),
+        v.transpose(1, 2).contiguous(),
+        mask,
+        vbs,
+    )
+    assert lse_bshd.shape == (1, heads, sq), lse_bshd.shape

fastvideo-kernel/tests/test_vsa256_triton.py

@@ -22,6 +22,7 @@ def _torch_vsa256_reference(
     q_var: torch.Tensor,
     kv_var: torch.Tensor,
     topk_logical: int,
+    compress_attn_weight: torch.Tensor | None = None,
 ) -> torch.Tensor:
     bsz, heads, _sq, dim = q.shape
     q_blocks = q_var.numel()
@@ -55,6 +56,8 @@ def _torch_vsa256_reference(
     logits = logits.masked_fill(~token_mask, float("-inf"))
     prob = torch.softmax(logits, dim=-1)
     out_s = torch.matmul(prob, vf).to(q.dtype)
+    if compress_attn_weight is not None:
+        return out_c * compress_attn_weight + out_s
     return out_c + out_s
 
 

fastvideo-kernel/tests/test_vsa_triton_backward_scale.py

@@ -0,0 +1,104 @@
+"""Regression: Triton block-sparse backward gradient parity at realistic activation scale.
+
+The backward used to fold ``sm_scale / ln(2)`` into K in bf16 before the
+exp2-based logit recompute. The bf16 rounding error on the pre-scaled K grows
+proportionally to |logit| and exp2 amplifies it into exponentially wrong
+probabilities, so dQ/dK/dV were correct at unit scale (every pre-existing test)
+but off by orders of magnitude at real activation magnitudes.
+
+This test sweeps the input scale and checks the Triton kernel's gradients
+against an fp32 masked-dense SDPA reference. The unit-scale case is the
+control (it passed even with the broken kernel); the large-scale cases are
+the regression.
+"""
+
+import pytest
+import torch
+
+from fastvideo_kernel.block_sparse_attn import _map_to_index, block_sparse_attn_triton
+
+from .utils import generate_block_sparse_mask_for_function
+
+BLOCK = 64
+
+
+@pytest.fixture(autouse=True)
+def _seed_rng():
+    """Pin the RNG so these cases do not depend on what ran before them.
+
+    Same convention as test_vsa_varlen.py: every tensor here comes from the
+    global torch RNG and the checks use tight thresholds, so an unseeded run
+    would shift inputs whenever an earlier test file draws a different number
+    of randoms.
+    """
+    torch.manual_seed(0)
+    torch.cuda.manual_seed_all(0)
+
+
+def _dense_reference(q, k, v, block_mask):
+    """fp32 masked-dense SDPA over the token-expanded block mask.
+
+    q/k/v: [B, H, S, D]; block_mask: [B, H, S // BLOCK, S // BLOCK] bool.
+    """
+    qf, kf, vf = q.float(), k.float(), v.float()
+    token_mask = block_mask.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)
+    logits = torch.matmul(qf, kf.transpose(-2, -1)) * (q.shape[-1]**-0.5)
+    logits = logits.masked_fill(~token_mask, float("-inf"))
+    return torch.matmul(logits.softmax(dim=-1), vf)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("scale", [1.0, 4.0, 16.0])
+def test_triton_backward_grad_parity_across_input_scales(scale: float) -> None:
+    """Kernel dQ/dK/dV must stay within a few percent of the fp32 reference
+    regardless of input magnitude.
+
+    With the bf16 K pre-scaling bug, scale<=4.0 passes at this geometry while
+    scale=16.0 fails (measured on GB200: dq relative L2 error 5.9e-1 vs 6.9e-3
+    fixed); at larger geometries and real activation magnitudes the broken
+    kernel is off by orders of magnitude. The passing unit-scale case is
+    exactly how the bug survived the original test suite.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    batch, heads, dim = 1, 4, 128
+    num_blocks = 8
+    seq = num_blocks * BLOCK
+
+    q = torch.randn(batch, heads, seq, dim, device=device, dtype=dtype) * scale
+    k = torch.randn(batch, heads, seq, dim, device=device, dtype=dtype) * scale
+    v = torch.randn(batch, heads, seq, dim, device=device, dtype=dtype)
+    grad_out = torch.randn_like(q)
+
+    block_mask = generate_block_sparse_mask_for_function(heads, num_blocks, num_blocks, k=3,
+                                                         device=device).unsqueeze(0)
+    q2k_idx, q2k_num = _map_to_index(block_mask)
+    variable_block_sizes = torch.full((num_blocks, ), BLOCK, dtype=torch.int32, device=device)
+
+    q_ker, k_ker, v_ker = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out_ker, _ = block_sparse_attn_triton(q_ker, k_ker, v_ker, q2k_idx, q2k_num, variable_block_sizes)
+    (out_ker.float() * grad_out.float()).sum().backward()
+
+    q_ref, k_ref, v_ref = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
+    out_ref = _dense_reference(q_ref, k_ref, v_ref, block_mask)
+    (out_ref * grad_out.float()).sum().backward()
+
+    # Forward is exact at any scale; this pins the harness itself.
+    fwd_rel = ((out_ker.float() - out_ref).norm() / out_ref.norm()).item()
+    assert fwd_rel < 2e-2, f"scale={scale}: forward rel err {fwd_rel:.3e}"
+
+    for name, g_ker, g_ref in (
+        ("dq", q_ker.grad, q_ref.grad),
+        ("dk", k_ker.grad, k_ref.grad),
+        ("dv", v_ker.grad, v_ref.grad),
+    ):
+        assert torch.isfinite(g_ker).all().item(), f"scale={scale}: non-finite {name}"
+        ref_norm = g_ref.float().norm()
+        rel = ((g_ker.float() - g_ref.float()).norm() / ref_norm.clamp_min(1e-12)).item()
+        ratio = (g_ker.float().norm() / ref_norm.clamp_min(1e-12)).item()
+        print(f"scale={scale} {name}: rel_l2={rel:.4e} norm_ratio={ratio:.4f}")
+        assert rel < 5e-2, f"scale={scale}: {name} rel l2 err {rel:.3e} >= 5e-2"
+        assert 0.98 < ratio < 1.02, f"scale={scale}: {name} grad-norm ratio {ratio:.4f}"

fastvideo-kernel/tests/test_vsa_varlen.py

@@ -23,6 +23,20 @@ from fastvideo_kernel.block_sparse_attn import (
 from fastvideo_kernel.block_sparse_attn_varlen import block_sparse_attn_varlen
 
 
+@pytest.fixture(autouse=True)
+def _seed_rng():
+    """Pin the RNG so these cases do not depend on what ran before them.
+
+    Every tensor and every variable block size here comes from the global
+    torch RNG, and the gradient checks use a tight max_rel threshold. Without
+    a seed the inputs shift whenever an earlier test file draws a different
+    number of randoms, which surfaces as an unrelated-looking failure in
+    whichever case happens to land on unlucky data.
+    """
+    torch.manual_seed(0)
+    torch.cuda.manual_seed_all(0)
+
+
 def _reference_per_sequence(
     q_list,
     k_list,