[feat] OpenAI data URLs: decode in bounded chunks, test cancelled decoding

Review follow-ups for the inline video/audio reference decoder:

- decode_data_url now undoes percent escapes and strips whitespace in
  64 KiB chunks and stops as soon as the payload must decode to more than
  server.media_reference_max_bytes. An oversized request costs memory in
  proportion to the limit, not to its own size. Escapes are never split
  across chunks.
- Data URL checks lower-case only the scheme instead of copying the whole
  URL.
- Tests: the wrapped/percent-escaped payload decodes the same at chunk
  sizes that split escapes and line breaks; an oversized payload is refused
  after a few chunks; and a request cancelled before its decoding thread
  writes the file still leaves no upload once the thread finishes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

@@ -1,15 +1,15 @@
 ---
 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.
+description: Work on FastVideo's Spark 4 and Slinky Slurm GPU CI lanes, static Buildkite graph, trusted ci-runner policy, lane scripts, and hardware validation.
 ---
 
-# Slinky Slurm CI lanes
+# GPU 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.
+GPU CI. The trusted uploader runs on the Slinky login plane. Slinky agents
+dispatch cluster lanes through Slurm into isolated Enroot containers. The
+Spark 4 agent dispatches eligible single-GPU lanes through local Slurm into
+isolated Docker containers. Both execute the immutable PR SHA.
 
 The files under `fastvideo/tests/modal/` and `.buildkite/scripts/pr_test.sh`
 are dormant rollback code. Never add an active Buildkite or slash-command
@@ -33,6 +33,14 @@ and the coordination contract with that bundle.
   A Spark runner sets `FASTVIDEO_CI_BACKEND=spark4`; gate cluster-only tests on
   that flag, never on memory heuristics. Keep
   `SPARK_CAPABLE_LANES` in `test_ci_test_collection.py` in sync.
+- The trusted uploader fills `FV_SPARK_BACKEND`: `spark4` for `fastcheck`,
+  `merge` and `direct` builds, including unset scope treated as Fastcheck,
+  while an agent tagged `queue=ci-runner,backend=spark4` is connected and
+  not paused, `slurm`
+  otherwise. Selection happens at upload time. Legacy graphs without backend
+  tags need explicit normalization; setting the variable alone does not pin
+  those jobs to Slurm. It rejects any other `backend` value. The upload step
+  in each entry pipeline requires `backend=slurm` so it never lands on the Spark 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
@@ -64,7 +72,7 @@ and the coordination contract with that bundle.
 - 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
+- The Slinky 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.

.github/mergify.yml

@@ -9,6 +9,28 @@ merge_protections:
       - check-success=fastcheck-passed
       - check-success=full-suite-passed
 
+  # Mirrors the `paths` filter of .github/workflows/ci-macos-mlx.yml so the
+  # required checks always exist whenever this protection applies.
+  - name: MLX smoke requirements
+    if:
+      - base = main
+      - or:
+          - files~=^\.github/workflows/ci-macos-mlx\.yml$
+          - files~=^fastvideo/mlx_runtime/
+          - files~=^fastvideo/tests/mlx/
+          - files~=^fastvideo/tests/platforms/test_mps_vsa_error\.py$
+          - files~=^fastvideo/tests/platforms/test_cpu_sdpa\.py$
+          - files~=^fastvideo/platforms/cpu\.py$
+          - files~=^fastvideo/platforms/mps\.py$
+          - files~=^fastvideo/platforms/__init__\.py$
+          - files~=^fastvideo/__init__\.py$
+          - files~=^examples/inference/basic/mlx_[^/]*\.py$
+          - files~=^fastvideo/benchmarks/mlx_[^/]*\.py$
+          - files~=^pyproject\.toml$
+    success_conditions:
+      - check-success=mlx-smoke
+      - check-success=mlx-smoke-linux-cpu
+
 pull_request_rules:
 
   # ============================================================

.github/scripts/plan_merge_ci.py

@@ -131,6 +131,11 @@ class FamilyCoverage:
 
 
 FAMILY_COVERAGE = (
+    FamilyCoverage(
+        re.compile(r"(^|[/_.-])qwen[_-]?image[_-]?2[._-]?1([/_.-]|$)"),
+        ("test_qwen_image21.py", ),
+        ("test_qwen_image21_similarity.py", ),
+    ),
     FamilyCoverage(
         re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"),
         ("test_flux2_klein.py", ),

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

@@ -2,6 +2,9 @@ name: macOS MLX Smoke
 
 on:
   pull_request:
+    # The defaults plus ready_for_review: both jobs skip drafts, and Mergify requires them
+    # for MLX changes, so marking a draft ready must run them.
+    types: [opened, synchronize, reopened, ready_for_review]
     branches: [main]
     paths:
       - ".github/workflows/ci-macos-mlx.yml"
@@ -85,7 +88,10 @@ jobs:
             fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \
+            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py \
+            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \
+            fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_ref_layout.py \
@@ -157,7 +163,10 @@ jobs:
             fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \
+            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py \
+            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \
+            fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_ref_layout.py \

docker/Dockerfile.spark-ci

@@ -0,0 +1,11 @@
+# syntax=docker/dockerfile:1.7
+# CI-only GB10 image. The regular Spark image keeps its native SDPA defaults.
+ARG SPARK_BASE_IMAGE=ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0-latest
+FROM ${SPARK_BASE_IMAGE}
+
+# VMoBA and the Wan causal gate require FA2. Install only this wheel so the
+# validated torch/flashinfer/CUTLASS stack from the base image stays intact.
+ARG FLASH_ATTN_WHEEL=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22/flash_attn-2.8.3+cu130torch2.12-cp312-cp312-linux_aarch64.whl
+RUN /root/.local/bin/uv pip install --python /opt/venv/bin/python --no-deps "${FLASH_ATTN_WHEEL}" && \
+    /opt/venv/bin/python -c "import glob, shutil; [shutil.rmtree(p) for p in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
+    /opt/venv/bin/python -c "import flash_attn, flashinfer; assert hasattr(flash_attn, 'flash_attn_varlen_func'); print('FA2', flash_attn.__version__, 'flashinfer', flashinfer.__version__)"

docs-site/lib/cookbook-pages.ts

@@ -111,6 +111,7 @@ const familyInfo: Record<string, { title: string; path: string; logo: string; ca
   kandinsky5: { title: "Kandinsky 5", path: "kandinsky5", logo: "kandinsky", category: "Video" },
   kandinsky6: { title: "Kandinsky 6", path: "kandinsky6", logo: "kandinsky", category: "Video" },
   flux: { title: "FLUX", path: "flux", logo: "black-forest-labs", category: "Image" },
+  qwen_image21: { title: "Qwen-Image-2.1", path: "qwen-image21", logo: "qwen", category: "Image" },
   sd35: { title: "Stable Diffusion", path: "stable-diffusion", logo: "stabilityai", category: "Image" },
   stable_audio: { title: "Stable Audio", path: "stable-audio", logo: "stabilityai", category: "Audio" },
   mmaudio: { title: "MMAudio", path: "mmaudio", logo: "fastvideo", category: "Audio" },
@@ -132,6 +133,7 @@ const TASK_GROUPS: Array<[RegExp, string]> = [
   [/video\/text to audio|video.*to audio/, "Video to audio"],
   [/text to audio/, "Text to audio"],
   [/text to image/, "Text to image"],
+  [/image to image/, "Image to image"],
   [/reference media to video/, "Reference to video"],
   [/first\/last frame to video/, "Image to video"],
   [/text or image to video/, "Text/image to video"],
@@ -161,6 +163,7 @@ export const modelDirectory = [...new Set(sourceRecipes.map((recipe) => recipe.m
       "tencent-hunyuan": "Tencent",
       kandinsky: "Kandinsky Lab",
       "black-forest-labs": "Black Forest Labs",
+      qwen: "Qwen",
       stabilityai: "Stability AI",
       nvidia: "NVIDIA",
     };

docs-site/navigation.json

@@ -253,6 +253,10 @@
         "label": "FLUX",
         "href": "/cookbook/flux"
       },
+      {
+        "label": "Qwen-Image-2.1",
+        "href": "/cookbook/qwen-image21"
+      },
       {
         "label": "Stable Diffusion",
         "href": "/cookbook/stable-diffusion"

docs-site/tests/cookbook.spec.ts

@@ -223,7 +223,7 @@ test("inherited values, explicit false, zero and null retain their meanings", as
 test("catalog discovery opens dedicated model recipes and filters by family and task", async ({ page }) => {
   await page.goto("/FastVideo/cookbook");
   const cards = page.locator(".recipe-card");
-  await expect(cards).toHaveCount(36);
+  await expect(cards).toHaveCount(37);
   await expect(cards.first()).toContainText("FastH3 OmniRef");
   await page.getByLabel("Model family", { exact: true }).selectOption("Wan");
   await page.getByLabel("Search recipes").fill("Wan2.1 I2V");
@@ -277,7 +277,7 @@ test("output tabs support keyboard navigation and precision chips use the backen
   await page.getByLabel("Search recipes").fill("does not exist");
   await expect(page.getByRole("heading", { name: "No matching recipes" })).toBeVisible();
   await page.getByRole("button", { name: "Clear filters" }).click();
-  await expect(page.locator(".recipe-card")).toHaveCount(36);
+  await expect(page.locator(".recipe-card")).toHaveCount(37);
 });
 
 test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => {
@@ -297,7 +297,7 @@ test("cookbook cards and controls remain readable in dark mode and reflow at 320
 test("image and audio models have logos and dedicated source recipe pages", async ({ page }) => {
   await page.goto("/FastVideo/cookbook");
   await page.getByRole("button", { name: "Image", exact: true }).click();
-  await expect(page.locator(".recipe-card")).toHaveCount(4);
+  await expect(page.locator(".recipe-card")).toHaveCount(5);
   const card = page.locator(".recipe-card").filter({ hasText: "FLUX.2 Klein" });
   await expect(card.locator("img")).toBeVisible();
   await card.click();
@@ -309,6 +309,26 @@ test("image and audio models have logos and dedicated source recipe pages", asyn
   await expect(page.getByRole("button", { name: "Serving", exact: true })).toHaveCount(0);
   await page.getByLabel("Seed", { exact: true }).fill("0");
   await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
+  await page.goto("/FastVideo/cookbook/qwen-image21");
+  const qwen = page.locator(".recipe-card").filter({ hasText: "Qwen-Image-2.1" });
+  await expect(qwen.locator("img")).toHaveAttribute("src", /\/qwen\.webp$/);
+  await qwen.click();
+  await expect(page.getByRole("heading", { level: 1 })).toHaveText("Qwen-Image-2.1");
+  await expect(page.getByRole("button", { name: "Serving", exact: true })).toHaveCount(0);
+  await page.getByLabel("Height", { exact: true }).fill("512");
+  await page.getByRole("tab", { name: "config.yaml", exact: true }).click();
+  const qwenConfig = page.getByLabel("Source example", { exact: true });
+  await expect(qwenConfig).toContainText("model_path: Qwen/Qwen-Image-2.1");
+  await expect(qwenConfig).toContainText("height: 512");
+  await expect(qwenConfig).toContainText("dit_layerwise: true");
+  await expect(page.getByRole("button", { name: "Download YAML", exact: true })).toBeEnabled();
+  await page
+    .getByRole("group", { name: "Recipe variant", exact: true })
+    .getByRole("button", { name: "Qwen-Image-2.1 references and editing", exact: true })
+    .click();
+  await page.getByRole("tab", { name: "Commands", exact: true }).click();
+  await expect(page.getByLabel("Inference command", { exact: true })).toContainText("--reference source.png");
+  await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
   await page.goto("/FastVideo/cookbook");
   await page.getByRole("button", { name: "Audio", exact: true }).click();
   await expect(page.locator(".recipe-card")).toHaveCount(3);
@@ -584,11 +604,11 @@ test("directory filters round-trip through the URL", async ({ page }) => {
   await page.goto("/FastVideo/cookbook?category=Image&runtime=cuda");
   await expect(page.getByRole("button", { name: "Image", exact: true })).toHaveAttribute("aria-pressed", "true");
   await expect(page.getByLabel("Runtime", { exact: true })).toHaveValue("cuda");
-  await expect(page.locator(".recipe-card")).toHaveCount(4);
+  await expect(page.locator(".recipe-card")).toHaveCount(5);
   await page.reload();
   await expect(page.getByRole("button", { name: "Image", exact: true })).toHaveAttribute("aria-pressed", "true");
   await expect(page.getByLabel("Runtime", { exact: true })).toHaveValue("cuda");
-  await expect(page.locator(".recipe-card")).toHaveCount(4);
+  await expect(page.locator(".recipe-card")).toHaveCount(5);
 });
 
 test("a saved target missing from the catalog stays visible and clearable", async ({ page }) => {

docs/assets/cookbook-recipes.json

@@ -1,5 +1,5 @@
 {
-  "version": 18,
+  "version": 19,
   "recipes": [
     {
       "id": "fastwan21-t2v",
@@ -429,6 +429,44 @@
       "expected_artifact": "PNG images under outputs/flux_dev/samples/",
       "limitations": ["FLUX.1 is loadable by ID but registers no model_family in fastvideo/registry.py; it is grouped under FLUX for documentation only."]
     },
+    {
+      "id": "qwen-image21-t2i",
+      "family": "qwen_image21",
+      "stage": "inference",
+      "task": "Text to image",
+      "label": "Qwen-Image-2.1",
+      "summary": "Generate a 1024 × 1024 image in 40 steps on one CUDA GPU, with the CPU offload settings recorded in the maintained YAML.",
+      "model": "Qwen/Qwen-Image-2.1",
+      "source": "examples/inference/basic/qwen_image21_t2i.yaml",
+      "command": "fastvideo generate --config examples/inference/basic/qwen_image21_t2i.yaml",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
+      "evidence": "Source-backed",
+      "expected_artifact": "RGBA PNG image at outputs/qwen_image21/teapot.png",
+      "limitations": ["The source uses DiT layerwise offload, text-encoder and VAE offload, and lazy loading. GPU memory requirements have not been established for this recipe."],
+      "related": ["qwen-image21-i2i"]
+    },
+    {
+      "id": "qwen-image21-i2i",
+      "family": "qwen_image21",
+      "stage": "inference",
+      "task": "Image to image",
+      "label": "Qwen-Image-2.1 references and editing",
+      "summary": "Generate or edit an image from a prompt and up to ten ordered reference images, using the maintained single-GPU Python example.",
+      "model": "Qwen/Qwen-Image-2.1",
+      "source": "examples/inference/basic/basic_qwen_image21.py",
+      "command": "python examples/inference/basic/basic_qwen_image21.py --reference source.png --prompt 'Change the shirt to blue; retain the subject and composition.' --output outputs/qwen_image21/edit.png",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
+      "evidence": "Source-backed",
+      "expected_artifact": "RGBA PNG image at outputs/qwen_image21/edit.png",
+      "limitations": [
+        "Replace source.png with an existing image. Repeat --reference for additional images and describe their order in the prompt.",
+        "Mask images use the same ordered-reference path; they do not impose hard pixel-preservation constraints.",
+        "The source enables CPU offload and lazy loading. This entry does not establish a minimum GPU memory requirement."
+      ],
+      "related": ["qwen-image21-t2i"]
+    },
     {
       "id": "sd35-medium-t2i",
       "family": "sd35",

docs/assets/logos/SOURCES.md

@@ -13,6 +13,7 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo.
 | `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header |
 | `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers |
 | `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header |
+| `qwen.webp` | [Official Qwen Hugging Face organization avatar](https://huggingface.co/Qwen) | Qwen-Image-2.1 family card and page header |
 | `minimax.webp` | [Official MiniMax Hugging Face organization avatar](https://huggingface.co/MiniMaxAI) | MiniMax H3 family card and page header |
 | `stabilityai.webp` | [Official Stability AI Hugging Face organization avatar](https://huggingface.co/stabilityai) | Stable Diffusion and Stable Audio cards and page headers |
 | `fastvideo.webp` | [Official FastVideo Hugging Face organization avatar](https://huggingface.co/FastVideo) | Matrix Game and MMAudio cards (converted weights published by this org) |

docs/contributing/ci_architecture.md

@@ -17,7 +17,7 @@ PR opened or updated
   |     style, lint, type, spelling, Markdown, workflow syntax, filenames
   |
   |-- Tier 2: Fastcheck
-  |     Buildkite schedules six lanes on the Slinky Slurm cluster
+  |     Buildkite schedules six lanes across Spark 4 and Slinky Slurm
   |     encoder, VAE, transformer, kernel, unit, DreamVerse
   |
   |-- /merge or ready label
@@ -25,7 +25,7 @@ PR opened or updated
         `-- 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
+              API, or performance lanes on Spark 4 or Slinky Slurm
               |
               pass -> Mergify squash-merges when all merge conditions pass
               fail -> fix, push, and re-run
@@ -43,11 +43,12 @@ CI is not one monolithic job:
   updates, docs deployment, image builds, package publishing, and community
   automations.
 - 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.
+  plane. Slinky agents dispatch tests from the login plane; the Spark 4 agent
+  dispatches local Slurm jobs.
+- Slinky Slurm runs cluster lanes in isolated Enroot containers using a
+  persistent four-GPU allocation. Spark 4 runs eligible single-GPU lanes in
+  isolated Docker containers through its own Slurm queue. Both test 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
@@ -99,7 +100,7 @@ debugging a hook implementation.
 | Attribute | Value |
 |---|---|
 | Triggered by | Buildkite PR builds with `TEST_SCOPE=fastcheck` or unset |
-| Compute | Slinky Slurm (`ci-runner` queue) |
+| Compute | Spark 4 when connected and unpaused at upload time; otherwise Slinky Slurm. Both use `ci-runner`. |
 | Definition | `.buildkite/pipeline.yml` |
 | Entrypoint | Trusted host driver -> `.buildkite/scripts/unit_test.sh` or `.buildkite/scripts/lanes/*.sh` |
 
@@ -122,15 +123,61 @@ Setting it to `spark4` moves only the Spark-capable lanes to a GB10 runner.
 That runner exports `FASTVIDEO_CI_BACKEND=spark4`, which the lane scripts and
 a few tests honor to skip GB200-only coverage (for example the LTX-2 finetune
 test, which needs more than 128 GB). GB10 golden references are keyed by
-`NVIDIA_GB10`; see `fastvideo/tests/golden_gate`. The private uploader policy
-must allow the `backend` agent key before this routing is active.
+`NVIDIA_GB10`; see `fastvideo/tests/golden_gate`.
+
+The trusted uploader on the Slurm login plane picks the value at upload time.
+It sets `FV_SPARK_BACKEND=spark4` only when both hold:
+
+- the build's `TEST_SCOPE` is `fastcheck`, `merge` or `direct`, or is unset;
+- the Buildkite agents API reports an agent tagged
+  `queue=ci-runner,backend=spark4` that is connected and not paused.
+
+Otherwise it sets `slurm`. For pipelines with the backend tags above,
+`/test full`, the weekly Full Suite, a paused or disconnected Spark agent, and
+any failure to read the agents API select Slurm for every lane. This choice
+happens at upload time; it does not move jobs already queued for Spark.
+Older graphs without backend tags need the normalization described below.
+
+The uploader rejects any `backend` value other than the two forms in the
+table, so a pull request cannot pin a lane to the Spark runner directly. The upload step in each entry pipeline requires `backend=slurm`, so
+the Spark runner never takes it and this policy always runs on the login plane.
+
+The regular GB10 inference image intentionally omits FlashAttention to match
+native Spark output. CI uses `docker/Dockerfile.spark-ci` to add the pinned
+FA2 wheel required by VMoBA and Wan causal tests. This image keeps the base
+Torch, FlashInfer, and CUTLASS packages and removes the wheel's unsupported
+FA4 CuTe implementation. Build it separately from the inference image:
+
+```bash
+docker build -f docker/Dockerfile.spark-ci \
+  --build-arg SPARK_BASE_IMAGE=<validated-GB10-image> \
+  -t fastvideo-ci-gb10:<version> .
+```
+
+Golden references must match this CI image, including its FA2 version. Seed
+only from unchanged upstream main, verify in separate processes, and review
+before publishing new `golden_gates/NVIDIA_GB10/` files. Preserve existing
+references for other GPUs. A successful import alone is insufficient; validate
+an actual FA2 forward, the VMoBA tests, and each lane's complete payload before
+selecting the new image in the host-owned runner configuration.
+
+The host runner must allowlist Spark-capable lane types before allocating the
+GPU. Label each Docker container with its job token and remove it from the
+outer runner on completion or cancellation. Slurm killing the Docker client
+does not stop a container owned by the Docker daemon.
+
+The trusted uploader must preserve each literal `backend: "slurm"` even when
+it selects `FV_SPARK_BACKEND=spark4`. Older PR revisions without backend tags
+must be normalized to `backend: "slurm"` before upload; `queue: "ci-runner"`
+alone also matches the Spark agent. Update the private uploader's lane policy
+with routing changes and verify the resolved agent queries on a canary build.
 
 ### Tier 3: Change-Aware Merge Gate
 
 | Attribute | Value |
 |---|---|
 | Triggered by | `/merge`, adding `ready`, or a new push to a PR that already has `ready` |
-| Compute | Slinky Slurm only (`ci-runner` queue) |
+| Compute | Spark 4 or Slinky Slurm, selected per lane (`ci-runner` queue) |
 | Definition | `.buildkite/pipeline.yml` |
 | Entrypoint | `/opt/fastvideo-ci-runner/run-ci` (`run-unit` is a compatibility wrapper) |
 
@@ -396,7 +443,7 @@ Valid direct test names:
 | `/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.
+they select the same test 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
@@ -476,7 +523,7 @@ 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.
+internal suffix is an implementation detail; backend tags select the compute host.
 
 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

docs/contributing/env_vars.md

@@ -215,7 +215,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_MLX_COMPILE`                                    | bool                                                               | `0`                                                   | performance | Compile the MLX DiT forward with mx.compile.                                                                                                                                                                                                                                                                                                                                                                                                |
 | `FASTVIDEO_MLX_FAST_NORM`                                  | bool                                                               | `0`                                                   | performance | Use MLX fast normalization kernels.                                                                                                                                                                                                                                                                                                                                                                                                         |
 | `FASTVIDEO_MLX_DQ_GEMM`                                    | str                                                                | `1`                                                   | performance | MLX dequantized GEMM for affine-quantized weights: 0 turns it off, 1 uses the measured minimum row count, and an integer sets the minimum row count.                                                                                                                                                                                                                                                                                        |
-| `FASTVIDEO_MLX_QUANT_GEMM`                                 | str                                                                | `auto`                                                | performance | MLX H3 affine-quantized GEMM path: `auto` keeps `FASTVIDEO_MLX_DQ_GEMM` behaviour, `dequant` always uses dequantize + dense GEMM for wide inputs, `qmm` always uses `mx.quantized_matmul`.                                                                                                                                                                                                                                                  |
+| `FASTVIDEO_MLX_QUANT_GEMM`                                 | str                                                                | `auto`                                                | performance | MLX H3 affine-quantized GEMM path: auto keeps FASTVIDEO_MLX_DQ_GEMM behaviour, dequant always uses dequantize + dense GEMM for wide inputs, qmm always uses mx.quantized_matmul.                                                                                                                                                                                                                                                            |
 | `FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D`                      | bool                                                               | `1`                                                   | performance | LTX-2 VAE uses the channels_last_3d memory format.                                                                                                                                                                                                                                                                                                                                                                                          |
 | `FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST`                    | bool                                                               | `1`                                                   | performance | LTX-2 audio decoding runs without CUDA autocast. Deprecated names: `LTX2_DISABLE_AUDIO_AUTOCAST`.                                                                                                                                                                                                                                                                                                                                           |
 | `FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION` | bool                                                               | `0`                                                   | performance | Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for the 4-step Klein model.                                                                                                                                                                                                                                                                                                               |

docs/cookbook/index.md

@@ -7,6 +7,7 @@ The interactive catalog groups models by family and generates YAML for supported
 - [Wan and FastMetal](wan.md)
 - [LTX](ltx.md)
 - [FLUX](flux.md)
+- [Qwen-Image-2.1](qwen-image21.md)
 - [Stable Diffusion](stable-diffusion.md)
 - [Stable Audio](stable-audio.md)
 - [MMAudio](mmaudio.md)

docs/cookbook/qwen-image21.md

@@ -0,0 +1,108 @@
+# Qwen-Image-2.1 recipes
+
+<RecipeFamily family="qwen_image21" />
+
+Generate images, edit an image, or combine up to ten ordered references with
+`Qwen/Qwen-Image-2.1`. Outputs are RGBA PNGs. Request a transparent background
+in the prompt when needed.
+
+## Run on a GPU
+
+Follow the [CUDA installation guide](../getting_started/installation/gpu.md), then
+run from the repository root on one CUDA GPU (`tp_size=sp_size=1`). The examples
+use CPU offload and VAE tiling. The loader downloads the complete checkpoint;
+use `--model-path /path/to/Qwen-Image-2.1` for local weights. No conversion is needed.
+
+```bash
+python examples/inference/basic/basic_qwen_image21.py \
+  --prompt 'A red ceramic teapot, isolated on a transparent background.' \
+  --output outputs/teapot.png
+```
+
+This saves a 1024×1024 PNG after 40 denoising steps. For instruction editing,
+replace `source.png` with your image:
+
+```bash
+python examples/inference/basic/basic_qwen_image21.py \
+  --reference source.png \
+  --prompt 'Change the shirt to blue; retain the subject and composition.' \
+  --output outputs/edit.png
+```
+
+Repeat `--reference` for reference-to-image generation, annotations or masks:
+
+```bash
+python examples/inference/basic/basic_qwen_image21.py \
+  --reference subject.png --reference style.png \
+  --prompt 'Paint the subject from image 1 in the style of image 2.'
+```
+
+A mask is another reference interpreted through the prompt; pixels outside it
+are not guaranteed to remain identical.
+
+## Configuration
+
+- Set `--height`, `--width`, `--steps` and `--seed` as needed. Dimensions must be
+  positive multiples of 32; use `--height 2048 --width 2048` for 2K output.
+- `--reference-resolution` defaults to 1024 and must be a positive multiple of
+  32. References retain their aspect ratio and are resized to approximately
+  that resolution squared in pixels, independently of output dimensions.
+- Prefix KV caching avoids recomputing text/reference tokens each step. The
+  Python example caches on CPU; `--no-use-kv-cache` disables reuse.
+- For two-branch CFG, set `--true-cfg-scale` above 1 and supply
+  `--negative-prompt` (an empty string is valid). Embedded guidance is not used.
+
+For the CLI, use the checked-in [T2I YAML](../../examples/inference/basic/qwen_image21_t2i.yaml)
+or [editing YAML](../../examples/inference/basic/qwen_image21_edit.yaml) with
+`fastvideo generate --config <file>`. The T2I YAML keeps prefix KV caches on CUDA;
+the Python example and editing YAML default to CPU caches. See
+[Configuration](../inference/configuration.md) for the Python API and YAML options.
+
+## Attention
+
+Dense image attention uses FastVideo's shared automatic backend selector, which
+prefers FlashAttention on supported CUDA hardware. Masked text attention uses
+SDPA. To request FA4 with its dependencies installed on supported hardware:
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+export FASTVIDEO_FA4=1
+```
+
+Use `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` and `FASTVIDEO_FA4=0` for upstream
+parity comparisons. Different attention kernels can produce different pixels.
+See [Optimizations](../inference/optimizations.md) for backend and offload settings.
+
+## Measured memory
+
+Peaks below are in **GiB**, measured on one **NVIDIA GB200** with BF16,
+1024×1024 output, seed 42, compilation disabled and two 40-step requests per
+process. Image cases use one reference. These measurements do not establish
+minimum hardware capacity; an actual 24 GB GPU has not been verified.
+
+| Configuration | KV cache | GPU allocated | GPU reserved | Device VRAM including driver | Host process-tree RSS |
+| --- | --- | ---: | ---: | ---: | ---: |
+| T2I YAML, SDPA, offloaded, tiled | CUDA | 13.28 | 22.54 | 24.19 | 43.87 |
+| Plain I2I Python, SDPA, offloaded, tiled | CPU | 13.29 | 16.11 | 17.76 | 48.29 |
+| Editing Python, SDPA, offloaded, tiled | CPU | 13.29 | 16.11 | 17.76 | 49.25 |
+| T2I, FA4, resident models, untiled VAE | CUDA | 32.77 | 35.50 | 37.09 | 13.12 |
+
+Measurements cover loading, cold/warm requests, PNG export and shutdown.
+Allocated/reserved values use PyTorch counters; device VRAM and simultaneous
+process-tree RSS were sampled every 20 ms. Device VRAM includes the GB200's
+0.72 GiB driver reservation. Allow host RAM headroom beyond process RSS.
+Offloaded rows explicitly use SDPA; the resident FA4 row disables offload and
+tiling and is a separate configuration, not a controlled latency comparison.
+
+## Validation and scope
+
+On GB200, 40-step upstream SDPA parity passed for eleven 1024×1024 cases:
+T2I, plain I2I, editing, reference-to-image, RGBA generation/editing, annotations,
+mask images, true CFG, uncached execution and ten references. Separate 2048×2048
+T2I also passed. Denoising trajectories and output pixels matched exactly.
+FA4 block and T2I/edit image regressions passed against their matching references.
+See the [validation runbook](../../tests/local_tests/qwen_image21/README.md)
+for reproducible commands, pinned revisions and numerical diagnostics.
+
+The supported scope is single GPU, batch one and unquantized inference.
+SP/TP, LoRA, prompt rewriting and HTTP serving are not included.

docs/design/inference_schema_parity_inventory.yaml

@@ -149,6 +149,7 @@ surfaces:
       preprocess_text_funcs: "Internal text preprocessing hooks."
       postprocess_text_funcs: "Internal text postprocessing hooks."
       scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
+      output_channels: "Decoded pixel channel count fixed by the model; used to allocate the output buffer."
 
   pipeline_config_extensions:
     preset_owned:
@@ -442,6 +443,15 @@ surfaces:
         sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig]
       sample_fps:
         sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig]
+      reference_resolution:
+        target: generator.pipeline.experimental.reference_resolution
+        sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
+      use_kv_cache:
+        target: generator.pipeline.experimental.use_kv_cache
+        sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
+      kv_cache_device:
+        target: generator.pipeline.experimental.kv_cache_device
+        sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
     compatibility_only:
       batch_size: "Gen3C inference-only tuning field pending typed batching design."
       gradient_checkpointing: "Gen3C inference-only compatibility field pending typed batching design."
@@ -460,6 +470,8 @@ surfaces:
       text_offset: "MagiHuman internal data-proxy text alignment offset."
       vae_stride: "MagiHuman internal VAE/data-proxy stride setting."
       vsa_ref_keep_rate: "MiniMax-H3 PDD reference-video VSA keep rate; defaults to the checkpoint's trained value."
+      video_scheduler_shift: "MiniMax-H3 training-time video noise clock; inference samples with the checkpoint's scheduler config."
+      audio_scheduler_shift: "MiniMax-H3 training-time audio noise clock; inference samples with the checkpoint's scheduler config."
       z_dim: "MagiHuman internal VAE latent channel setting."
       vocoder_config: "Legacy internal component config object."
       vocoder_precision: "Precision override pending dedicated component precision design."

docs/design/server_contracts/openai.md

@@ -88,12 +88,27 @@ block. If `num_frames` is absent, `seconds * fps` is used. FastVideo also keeps
 the legacy `input_reference`, `reference_url`, `video_path`, and `video_url`
 spellings.
 
-Reference objects support URL or local-path strings through `image_url`,
-`video_url`, and `audio_url`. `file_id` references are schema-compatible but
-return HTTP 400 because FastVideo does not provide an OpenAI Files store.
-Image URLs, data URLs, local paths, and multipart `input_reference` uploads are
-materialized and decoded under the configured output directory during
-admission. Invalid media returns HTTP 400 before a job is created.
+Reference objects carry one string each through `image_url`, `video_url`, and
+`audio_url`. Each field accepts these forms:
+
+| Field | `http(s)` URL | `data:` URL | Local path |
+| --- | --- | --- | --- |
+| `image_reference[].image_url`, `input_reference`, `reference_url` | yes | `data:image/...;base64,` | under `server.media_roots` |
+| `video_reference[].video_url`, `video_path`, `video_url` | yes | `data:video/...;base64,` | under `server.media_roots` |
+| `audio_reference[].audio_url` | yes | `data:audio/...;base64,` | under `server.media_roots` |
+
+`input_reference` also accepts a multipart upload. `file_id` references are
+schema-compatible but return HTTP 400 because FastVideo does not provide an
+OpenAI Files store.
+
+URLs must resolve to public addresses unless
+`server.allow_private_media_urls` is set. A `data:` URL must be base64-encoded,
+and its media type must match its field: a `data:audio/...` URL in
+`video_reference` is refused. Each downloaded file, and each decoded video or
+audio file, is limited to `server.media_reference_max_bytes`. URLs and `data:` URLs are
+written under `server.output_dir/uploads` during admission, keep their order
+within each field, and are deleted when the job ends. Invalid, oversized, or
+mismatched media returns HTTP 400 before a job is created.
 
 ## Jobs and synchronous responses
 

docs/generate_examples.py

@@ -36,6 +36,7 @@ COOKBOOK_FAMILIES = {
     "kandinsky5",
     "kandinsky6",
     "flux",
+    "qwen_image21",
     "sd35",
     "minimax_h3",
     "stable_audio",

docs/inference/fasth3_omniref_mlx.md

@@ -73,6 +73,34 @@ Quantized checkpoints pick a linear path per layer. `--quant-gemm` (or
 
 The same flag is on the example: `--quant-gemm {auto,dequant,qmm}`.
 
+### Opt-in M5 attention backend
+
+`--vsa-impl sdpa` evaluates the trained VSA tile mask through fused MLX SDPA,
+sharing each attention head's K/V across its query tiles. On M5, the built-in
+matrix operations can be faster than the custom SIMD kernel even though they
+visit masked keys. The default `auto` selection is unchanged.
+
+This backend keeps the reference resolution, selected tiles, ragged-row masks,
+and compression gate. Accumulation differs from the SIMD backend, so outputs
+are not bit-identical. It supports tiles 64/128 with head dimension 128 and
+raises if fused SDPA is unavailable. Prefix K/V reuse is not supported; do not
+combine it with `--reuse-prefix-kv`. Query tiles are processed in chunks of at
+most 128, capped at 256 MiB per boolean-mask chunk with two chunks in flight.
+
+`--vsa-impl nax` is a separate opt-in backend for Apple M5 GPUs. It uses Metal
+TensorOps to visit only the selected key blocks and their valid rows, avoiding
+the dense mask and masked-key traversal of the SDPA backend. Dense prefix
+queries and the trained compression branch keep their existing implementations.
+The selected reference tokens, output resolution and denoising steps are unchanged.
+
+The NAX backend uses private helper headers from the installed MLX package;
+it accepts only the header hashes validated with MLX 0.32.4. Changed headers need
+numerical and generation validation before their hashes are accepted. Missing headers, unsupported GPUs
+and compilation failures raise errors rather than silently selecting another
+backend. Tiles 64/128 and head dimension 128 are supported. Prefix K/V reuse is
+not supported. Floating-point accumulation differs from both SIMD and SDPA,
+so the generated clip can differ even with the same seed. `auto` remains unchanged.
+
 ## Generate
 
 ```bash

docs/inference/support_matrix.md

@@ -69,6 +69,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | 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) |
 | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) |
+| qwen_image21 | `Qwen/Qwen-Image-2.1` | T2I, I2I | [basic_qwen_image21.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_qwen_image21.py) |
 | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) |
 | stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) |
 | wan | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | T2V | [basic.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic.py) |

examples/inference/basic/basic_qwen_image21.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen-Image-2.1 generation, editing, references and transparent PNG output."""
+
+import argparse
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OffloadConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-path", default="Qwen/Qwen-Image-2.1")
+    parser.add_argument("--revision")
+    parser.add_argument("--prompt", default="A red ceramic teapot, isolated on a transparent background.")
+    parser.add_argument("--reference", action="append", default=[], help="Ordered reference image; repeat up to ten times")
+    parser.add_argument("--negative-prompt", default=None)
+    parser.add_argument("--true-cfg-scale", type=float, default=1.0)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--reference-resolution", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=40)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--use-kv-cache", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--output", default="outputs/qwen_image21/image.png")
+    args = parser.parse_args()
+    if len(args.reference) > 10:
+        parser.error("At most ten ordered reference images are supported")
+    workload = "i2i" if args.reference else "t2i"
+    generator = VideoGenerator.from_config(
+        GeneratorConfig(
+            model_path=args.model_path,
+            revision=args.revision,
+            engine=EngineConfig(
+                num_gpus=1,
+                parallelism=ParallelismConfig(tp_size=1, sp_size=1),
+                offload=OffloadConfig(dit=True, dit_layerwise=True, text_encoder=True, vae=True,
+                                      lazy_module_load=True),
+            ),
+            pipeline=PipelineSelection(
+                workload_type=workload,
+                experimental={
+                    "reference_resolution": args.reference_resolution,
+                    "use_kv_cache": args.use_kv_cache,
+                },
+            ),
+        ))
+    try:
+        generator.generate(
+            GenerationRequest(
+                prompt=args.prompt,
+                negative_prompt=args.negative_prompt,
+                inputs=InputConfig(references=args.reference or None),
+                sampling=SamplingConfig(height=args.height, width=args.width, num_frames=1, fps=1,
+                                        num_inference_steps=args.steps, guidance_scale=1.0,
+                                        true_cfg_scale=args.true_cfg_scale, seed=args.seed),
+                output=OutputConfig(output_path=args.output, save_video=True, return_frames=False),
+            ))
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/mlx_fasth3_omniref.py

@@ -114,9 +114,10 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--vsa-ref-keep-rate", type=float, default=None,
                         help="fraction of each reference video's tiles a query keeps (p2_multi_region); "
                         "default is the contract's trained value")
-    parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd"), default="auto",
+    parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd", "sdpa", "nax"), default="auto",
                         help="sparse attention backend; auto uses the SIMD Metal kernel for 128-token tiles and "
-                        "gather+SDPA otherwise, reference forces gather+SDPA")
+                        "gather+SDPA otherwise, reference forces gather+SDPA; sdpa shares K/V with a tile mask "
+                        "and nax uses M5 TensorOps over selected blocks (both opt-in, incompatible with --reuse-prefix-kv)")
     parser.add_argument("--reuse-prefix-kv", action="store_true",
                         help="opt-in, lossy: record the text/reference rows' per-block K/V on step 1 and run later "
                         "steps on the target audio/video rows only (faster, slightly different output)")
@@ -145,6 +146,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         parser.error("Ref2VA needs at least one image or video reference")
     if args.model_root is None and (args.export is None or args.base_model_path is None):
         parser.error("give --model-root, or --export together with --base-model-path")
+    if args.vsa_impl in ("sdpa", "nax") and args.reuse_prefix_kv:
+        parser.error(f"--vsa-impl {args.vsa_impl} does not support --reuse-prefix-kv")
     return args
 
 

examples/inference/basic/qwen_image21_edit.yaml

@@ -0,0 +1,39 @@
+# Replace source.png and mask.png with your own images. The mask is an ordered
+# reference interpreted by the model, not a hard pixel-preservation constraint.
+generator:
+  model_path: Qwen/Qwen-Image-2.1
+  engine:
+    num_gpus: 1
+    parallelism:
+      tp_size: 1
+      sp_size: 1
+    offload:
+      dit: true
+      dit_layerwise: true
+      text_encoder: true
+      vae: true
+      lazy_module_load: true
+  pipeline:
+    workload_type: i2i
+    experimental:
+      reference_resolution: 1024
+      use_kv_cache: true
+request:
+  prompt: In image 1, turn the object highlighted by the mask in image 2 blue.
+  inputs:
+    references:
+      - source.png
+      - mask.png
+  sampling:
+    height: 1024
+    width: 1024
+    num_frames: 1
+    fps: 1
+    num_inference_steps: 40
+    guidance_scale: 1.0
+    true_cfg_scale: 1.0
+    seed: 42
+  output:
+    output_path: outputs/qwen_image21/edit.png
+    save_video: true
+    return_frames: false

examples/inference/basic/qwen_image21_t2i.yaml

@@ -0,0 +1,40 @@
+# Run on an allocated CUDA GPU after installing FastVideo and providing weights.
+# fastvideo generate --config examples/inference/basic/qwen_image21_t2i.yaml
+generator:
+  model_path: Qwen/Qwen-Image-2.1
+  engine:
+    num_gpus: 1
+    compile:
+      enabled: false
+      backend: inductor
+      fullgraph: false
+    parallelism:
+      tp_size: 1
+      sp_size: 1
+    offload:
+      dit: true
+      dit_layerwise: true
+      text_encoder: true
+      vae: true
+      lazy_module_load: true
+  pipeline:
+    workload_type: t2i
+    experimental:
+      kv_cache_device: cuda
+      reference_resolution: 1024
+      use_kv_cache: true
+request:
+  prompt: A red ceramic teapot, isolated on a transparent background.
+  sampling:
+    height: 1024
+    width: 1024
+    num_frames: 1
+    fps: 1
+    num_inference_steps: 40
+    guidance_scale: 1.0
+    true_cfg_scale: 1.0
+    seed: 42
+  output:
+    output_path: outputs/qwen_image21/teapot.png
+    save_video: true
+    return_frames: false

fastvideo/configs/models/dits/__init__.py

@@ -18,6 +18,7 @@ from fastvideo.configs.models.dits.mmaudio import (
 )
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.models.wan.config import WanVideoConfig
+from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig, Kandinsky6VideoAudioConfig
@@ -33,5 +34,5 @@ __all__ = [
     "StableAudioConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotWorldFastVideoConfig",
     "LingBotVideoConfig", "MiniMaxH3Config", "MMAudioArchConfig", "MMAudioTransformerConfig",
     "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS", "MMAUDIO_44K_TRAINING_VARIANTS",
-    "get_mmaudio_transformer_config"
+    "get_mmaudio_transformer_config", "QwenImage21ArchConfig", "QwenImage21Config"
 ]

fastvideo/configs/models/dits/qwen_image21.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Architecture and checkpoint loading configuration for Qwen-Image-2.1."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _is_transformer_block(name: str, module) -> bool:
+    parts = name.split(".")
+    return len(parts) >= 2 and parts[-2] == "transformer_blocks" and parts[-1].isdigit()
+
+
+@dataclass
+class QwenImage21ArchConfig(DiTArchConfig):
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int = 64
+    num_layers: int = 32
+    attention_head_dim: int = 128
+    num_attention_heads: int = 32
+    context_in_dim: int = 4096
+    mlp_ratio: int = 3
+    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
+    eps: float = 1e-6
+    causal_condition: bool = True
+    cast_prompt_embeds_to_dit_dtype: bool = True
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_block])
+    _supported_attention_backends: tuple[AttentionBackendEnum,
+                                         ...] = (AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN,
+                                                 AttentionBackendEnum.SAGE_ATTN)
+    param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
+
+    def __post_init__(self) -> None:
+        axes_dims_rope = tuple(self.axes_dims_rope)
+        if len(axes_dims_rope) != 3 or any(dim <= 0 or dim % 2 for dim in axes_dims_rope):
+            raise ValueError("axes_dims_rope must contain three positive even dimensions")
+        self.axes_dims_rope = (axes_dims_rope[0], axes_dims_rope[1], axes_dims_rope[2])
+        if sum(self.axes_dims_rope) != self.attention_head_dim:
+            raise ValueError("axes_dims_rope must sum to attention_head_dim")
+        if self.num_layers <= 0 or self.num_attention_heads <= 0:
+            raise ValueError("num_layers and num_attention_heads must be positive")
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class QwenImage21Config(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=QwenImage21ArchConfig)
+    prefix: str = "QwenImage21"

fastvideo/configs/models/encoders/__init__.py

@@ -15,6 +15,7 @@ from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig)
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLArchConfig, QwenImage21Qwen3VLConfig
 from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (
     StableAudioConditionerArchConfig,
@@ -49,6 +50,8 @@ __all__ = [
     "StableAudioConditionerConfig",
     "T5GemmaEncoderConfig",
     "Qwen3TextConfig",
+    "QwenImage21Qwen3VLArchConfig",
+    "QwenImage21Qwen3VLConfig",
     "Mistral3TextConfig",
     "LingBotWorld2UMT5ArchConfig",
     "LingBotWorld2UMT5Config",

fastvideo/configs/models/encoders/qwen_image21.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3-VL architecture used for Qwen-Image-2.1 conditioning."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.base import ModelConfig
+from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (
+    MiniMaxH3Qwen3VLArchConfig,
+    MiniMaxH3Qwen3VLConfig,
+    _OFFICIAL_ARCHITECTURES,
+    _VISION_CONFIG_MAPPING,
+)
+
+
+@dataclass
+class QwenImage21Qwen3VLArchConfig(MiniMaxH3Qwen3VLArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["QwenImage21Qwen3VLConditioner"])
+    hidden_size: int = 4096
+    intermediate_size: int = 12288
+    num_hidden_layers: int = 36
+    output_hidden_state_index: int = 36
+    num_hidden_layers_override: int | None = None
+    num_attention_heads: int = 32
+    vision_out_hidden_size: int = 4096
+    # The image pipeline's processor owns prompt lengths, including reference
+    # image tokens; the H3 text-only truncation limit is inappropriate here.
+    text_len: int = 262144
+    require_processor: bool = True
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.output_hidden_state_index != self.num_hidden_layers:
+            raise ValueError("Qwen-Image-2.1 conditioning requires the last decoder layer before final normalization")
+        if self.num_hidden_layers_override is not None:
+            raise ValueError("Qwen-Image-2.1 conditioning requires all text decoder layers")
+        if self.hidden_act != "silu" or self.vision_hidden_act != "gelu_pytorch_tanh":
+            raise ValueError("Qwen-Image-2.1 requires SiLU text MLPs and tanh-approximate GELU vision MLPs")
+        if (self.rope_scaling or {}).get("rope_type", "default") != "default":
+            raise ValueError("Qwen-Image-2.1 conditioning supports default multimodal RoPE")
+        self.tokenizer_kwargs = {
+            "padding": True,
+            "padding_side": "left",
+            "truncation": False,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class QwenImage21Qwen3VLConfig(MiniMaxH3Qwen3VLConfig):
+    arch_config: QwenImage21Qwen3VLArchConfig = field(default_factory=QwenImage21Qwen3VLArchConfig)
+    prefix: str = "qwen_image21_qwen3_vl"
+
+    def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
+        flattened = dict(source_model_dict)
+        flattened.update(dict(flattened.pop("text_config", {})))
+        vision = flattened.pop("vision_config", {})
+        for source_name, target_name in _VISION_CONFIG_MAPPING.items():
+            if source_name in vision:
+                flattened[target_name] = vision[source_name]
+
+        architectures = flattened.get("architectures", [])
+        if isinstance(architectures, str):
+            architectures = [architectures]
+        if any(architecture in _OFFICIAL_ARCHITECTURES for architecture in architectures):
+            flattened["architectures"] = ["QwenImage21Qwen3VLConditioner"]
+        if "num_hidden_layers" in flattened:
+            flattened["output_hidden_state_index"] = flattened["num_hidden_layers"]
+        for name in ("mrope_section", "vision_deepstack_visual_indexes"):
+            if isinstance(flattened.get(name), list):
+                flattened[name] = tuple(flattened[name])
+        # Transformers 5 serializes the same default RoPE parameters under a
+        # renamed field. Keep the checkpoint's frequencies in either format.
+        if "rope_parameters" in flattened:
+            parameters = dict(flattened["rope_parameters"])
+            flattened["rope_scaling"] = parameters
+            if "rope_theta" in parameters:
+                flattened["rope_theta"] = parameters["rope_theta"]
+        ModelConfig.update_model_arch(self, flattened)

fastvideo/configs/models/vaes/__init__.py

@@ -12,6 +12,7 @@ from fastvideo.configs.models.vaes.minimax_h3_video import (
 )
 from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig
 from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
+from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEArchConfig, QwenImage21VAEConfig
 from fastvideo.models.wan.vae_config import WanVAEConfig
 
 __all__ = [
@@ -31,4 +32,6 @@ __all__ = [
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
     "Flux2VAEConfig",
+    "QwenImage21VAEArchConfig",
+    "QwenImage21VAEConfig",
 ]

fastvideo/configs/models/vaes/qwen_image21.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen-Image-2.1 RGBA autoencoder configuration."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class QwenImage21VAEArchConfig(VAEArchConfig):
+    _class_name: str = "AutoencoderKLQwenImage21"
+    base_dim: int = 96
+    decoder_base_dim: int | None = 144
+    z_dim: int = 64
+    dim_mult: tuple[int, ...] = (1, 2, 4, 8, 8)
+    num_res_blocks: int = 2
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True, True)
+    dropout: float = 0.0
+    is_residual: bool = True
+    in_channels: int = 4
+    out_channels: int = 4
+    patch_size: int | None = None
+    scale_factor_spatial: int = 16
+    scale_factor_temporal: int = 8
+    latents_mean: tuple[float, ...] = (
+        0.5126,
+        0.7721,
+        -0.0631,
+        1.3506,
+        -0.7855,
+        -2.1025,
+        -0.3458,
+        1.3722,
+        1.8873,
+        -1.7177,
+        -0.6510,
+        0.2732,
+        0.7562,
+        -0.6163,
+        -1.0277,
+        3.8363,
+        2.0210,
+        0.0472,
+        0.9320,
+        2.0087,
+        2.4954,
+        -0.1391,
+        -1.4249,
+        1.8464,
+        -0.5236,
+        1.2826,
+        3.7046,
+        -1.3035,
+        2.7286,
+        -1.4518,
+        -1.9036,
+        -1.9955,
+        -0.0342,
+        -1.0265,
+        -0.7636,
+        3.0555,
+        0.0746,
+        -3.0751,
+        -0.1076,
+        1.7376,
+        -1.0914,
+        -1.9435,
+        -0.2784,
+        -1.3680,
+        0.4809,
+        -0.4433,
+        0.3764,
+        0.5729,
+        -2.0595,
+        1.0960,
+        -1.3260,
+        -2.0211,
+        -5.0179,
+        0.5275,
+        4.0162,
+        1.8505,
+        0.3026,
+        1.9373,
+        1.4937,
+        0.2632,
+        0.5547,
+        -1.7121,
+        -0.1562,
+        0.0304,
+    )
+    latents_std: tuple[float, ...] = (
+        3.2001,
+        3.2936,
+        3.4321,
+        3.0091,
+        3.1061,
+        4.0379,
+        4.0705,
+        3.7910,
+        3.0785,
+        3.6500,
+        3.9308,
+        3.0904,
+        2.8778,
+        3.7675,
+        3.7320,
+        5.0756,
+        3.2864,
+        4.0397,
+        3.1317,
+        4.0443,
+        2.9249,
+        3.9454,
+        3.0988,
+        4.2489,
+        3.4896,
+        3.8513,
+        3.9323,
+        3.4719,
+        3.7498,
+        4.2830,
+        3.5694,
+        4.2467,
+        3.9037,
+        3.2947,
+        5.0770,
+        3.5075,
+        3.2700,
+        3.4767,
+        2.8063,
+        5.1125,
+        3.5327,
+        4.7833,
+        3.1286,
+        4.1819,
+        3.8527,
+        3.8312,
+        3.5605,
+        4.3875,
+        3.9624,
+        4.0168,
+        3.5643,
+        4.0550,
+        5.5614,
+        4.2963,
+        4.4080,
+        3.4959,
+        3.8747,
+        3.7608,
+        3.5735,
+        3.1490,
+        3.7662,
+        3.6746,
+        3.4563,
+        3.8161,
+    )
+    scaling_factor: float = 1.0
+    temporal_compression_ratio: int = 1
+    spatial_compression_ratio: int = 16
+    param_names_mapping: dict[str, str] = field(default_factory=lambda: {r"^(.*)$": r"\1"})
+
+    @property
+    def latent_channels(self) -> int:
+        return self.z_dim
+
+    def __post_init__(self) -> None:
+        if len(self.dim_mult) < 2 or len(self.temperal_downsample) != len(self.dim_mult) - 1:
+            raise ValueError("`temperal_downsample` must have one entry per downsampling stage.")
+        if len(self.latents_mean) != self.z_dim or len(self.latents_std) != self.z_dim:
+            raise ValueError("Qwen-Image-2.1 latent mean/std must have one entry per latent channel.")
+        if any(value <= 0 for value in self.latents_std):
+            raise ValueError("Qwen-Image-2.1 latent standard deviations must be positive.")
+        if self.patch_size is not None and self.patch_size < 1:
+            raise ValueError("`patch_size` must be positive or None.")
+        actual_spatial_scale = 2**(len(self.dim_mult) - 1) * (self.patch_size or 1)
+        if self.scale_factor_spatial != actual_spatial_scale:
+            raise ValueError("`scale_factor_spatial` must match the VAE downsampling stages and patch size.")
+        self.spatial_compression_ratio = self.scale_factor_spatial
+
+
+@dataclass
+class QwenImage21VAEConfig(VAEConfig):
+    arch_config: QwenImage21VAEArchConfig = field(default_factory=QwenImage21VAEArchConfig)
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False

fastvideo/configs/pipelines/__init__.py

@@ -14,6 +14,7 @@ from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
 from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
@@ -26,5 +27,5 @@ __all__ = [
     "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
     "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
     "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
-    "get_pipeline_config_cls_from_name"
+    "QwenImage21PipelineConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -64,6 +64,7 @@ class PipelineConfig:
     vae_decode_precision: str | None = None
     vae_tiling: bool = True
     vae_sp: bool = True
+    output_channels: int = 3
 
     # Image encoder configuration
     image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)

fastvideo/configs/pipelines/qwen_image21.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native Qwen-Image-2.1 inference configuration."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.qwen_image21 import QwenImage21Config
+from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
+from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class QwenImage21PipelineConfig(PipelineConfig):
+    dit_config: QwenImage21Config = field(default_factory=QwenImage21Config)
+    vae_config: QwenImage21VAEConfig = field(default_factory=QwenImage21VAEConfig)
+    text_encoder_configs: tuple[QwenImage21Qwen3VLConfig,
+                                ...] = field(default_factory=lambda: (QwenImage21Qwen3VLConfig(), ))
+    text_encoder_precisions: tuple[str, ...] = ("bf16", )
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    vae_tiling: bool = True
+    vae_sp: bool = False
+    # Four output channels are required to preserve transparency through PNG export.
+    output_channels: int = 4
+    flow_shift: float | None = None
+    scheduler_step_in_fp32: bool = True
+    reference_resolution: int = 1024
+    use_kv_cache: bool = True
+    # Runtime placement policy. CPU keeps ten reference prefixes out of VRAM.
+    kv_cache_device: str = "cpu"

fastvideo/entrypoints/openai/media_policy.py

@@ -5,23 +5,28 @@ Local paths are readable only inside operator-listed ``server.media_roots``.
 URLs must use http(s) and, unless ``server.allow_private_media_urls`` is set,
 must resolve to public addresses on every redirect hop. Downloads are capped
 by streamed bytes and leave no partial file on failure or cancellation.
+Inline base64 ``data:`` URLs are decoded under the same byte cap.
 """
 
 from __future__ import annotations
 
 import asyncio
+import base64
+import binascii
 import contextlib
 import ipaddress
 import os
 import socket
 from dataclasses import dataclass
-from urllib.parse import urljoin, urlparse
+from urllib.parse import unquote, urljoin, urlparse
 
 import aiofiles
 import httpx
 
 MAX_REDIRECTS = 5
 READ_TIMEOUT_S = 60.0
+DATA_URL_HEADER_MAX_CHARS = 256
+DATA_URL_CHUNK_CHARS = 1 << 16  # at least 3, so a percent escape fits in one chunk
 
 
 class MediaPolicyError(ValueError):
@@ -74,6 +79,73 @@ def discard(path: str) -> None:
         os.unlink(path)
 
 
+def _too_large(max_bytes: int) -> MediaPolicyError:
+    return MediaPolicyError(f"Media reference is larger than the {max_bytes}-byte limit; raise "
+                            "server.media_reference_max_bytes in the serve config or provide a smaller file.")
+
+
+def data_url_media_type(source: str) -> str:
+    """The lower-cased media type of a base64 ``data:`` URL ("" when the URL names none)."""
+    separator = source.find(",")
+    if source[:5].lower() != "data:" or separator < 0 or separator > DATA_URL_HEADER_MAX_CHARS:
+        raise MediaPolicyError("Data URLs must have the form data:<media type>;base64,<payload>.")
+    parameters = source[5:separator].split(";")
+    if "base64" not in (parameter.strip().lower() for parameter in parameters[1:]):
+        raise MediaPolicyError("Data URLs must be base64-encoded: data:<media type>;base64,<payload>.")
+    return parameters[0].strip().lower()
+
+
+def _base64_payload(source: str, start: int, max_bytes: int) -> str:
+    """``source[start:]`` without percent escapes or whitespace, normalized in bounded chunks.
+
+    Stops as soon as the payload must decode to more than ``max_bytes``, so an oversized request costs memory
+    in proportion to the limit, not to its own size.
+    """
+    chunks: list[str] = []
+    length = 0
+    while start < len(source):
+        end = min(start + DATA_URL_CHUNK_CHARS, len(source))
+        escape = source.find("%", max(start + 1, end - 2), end) if end < len(source) else -1
+        if escape >= 0:
+            end = escape  # an escape that may continue past the chunk starts the next one
+        # Same leniency as SGLang Diffusion: percent escapes and line-wrapped payloads decode.
+        chunk = source[start:end]
+        if "%" in chunk:
+            chunk = unquote(chunk)
+        chunk = "".join(chunk.split())
+        length += len(chunk)
+        # Base64 decodes to 3 bytes per 4 characters, less at most 2 bytes of padding.
+        if length // 4 * 3 - 2 > max_bytes:
+            raise _too_large(max_bytes)
+        chunks.append(chunk)
+        start = end
+    return "".join(chunks)
+
+
+def decode_data_url(source: str, target: str, *, max_bytes: int) -> None:
+    """Decode a base64 ``data:`` URL to ``target``; nothing is written if it is invalid or over ``max_bytes``."""
+    data_url_media_type(source)
+    payload = _base64_payload(source, source.find(",") + 1, max_bytes)
+    padding = len(payload) - len(payload.rstrip("="))
+    if len(payload) // 4 * 3 - padding > max_bytes:
+        raise _too_large(max_bytes)
+    try:
+        data = base64.b64decode(payload, validate=True)
+    except (binascii.Error, ValueError) as error:
+        raise MediaPolicyError(f"Data URL payload is not valid base64: {error}") from error
+    if not data:
+        raise MediaPolicyError("Data URL payload is empty.")
+    if len(data) > max_bytes:
+        raise _too_large(max_bytes)
+    os.makedirs(os.path.dirname(os.path.abspath(target)), exist_ok=True)
+    try:
+        with open(target, "wb") as output:
+            output.write(data)
+    except BaseException:
+        discard(target)
+        raise
+
+
 async def fetch_media(
     url: str,
     target: str,
@@ -101,9 +173,7 @@ async def fetch_media(
                         async for chunk in response.aiter_bytes():
                             written += len(chunk)
                             if written > max_bytes:
-                                raise MediaPolicyError(
-                                    f"Media reference is larger than the {max_bytes}-byte limit; raise "
-                                    "server.media_reference_max_bytes in the serve config or provide a smaller file.")
+                                raise _too_large(max_bytes)
                             await output.write(chunk)
                     return response.headers
             raise MediaPolicyError(f"Media URL redirected more than {MAX_REDIRECTS} times.")

fastvideo/entrypoints/openai/request_adapter.py

@@ -12,6 +12,8 @@ from urllib.parse import urlparse
 from fastvideo.entrypoints.openai.media_policy import (
     MediaPolicy,
     MediaPolicyError,
+    data_url_media_type,
+    decode_data_url,
     discard,
     fetch_media,
     resolve_local_media,
@@ -41,7 +43,23 @@ class RequestAdaptationError(ValueError):
     """The transport request cannot be represented by the loaded pipeline."""
 
 
-_VIDEO_SUFFIXES = frozenset({".avi", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm"})
+_MEDIA_SUFFIXES = {
+    "video": frozenset({".avi", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm"}),
+    "audio": frozenset({".aac", ".flac", ".m4a", ".mp3", ".ogg", ".opus", ".wav", ".webm"}),
+}
+_DEFAULT_SUFFIXES = {"video": ".mp4", "audio": ".wav"}
+# Media types whose subtype is not already the file suffix; decoding probes content, so this is only a hint.
+_DATA_URL_SUFFIXES = {
+    "video/quicktime": ".mov",
+    "video/x-matroska": ".mkv",
+    "video/x-msvideo": ".avi",
+    "audio/mpeg": ".mp3",
+    "audio/mp4": ".m4a",
+    "audio/x-m4a": ".m4a",
+    "audio/wave": ".wav",
+    "audio/x-wav": ".wav",
+    "audio/x-flac": ".flac",
+}
 
 
 def _as_list(value: Any | list[Any] | None) -> list[Any]:
@@ -93,30 +111,63 @@ def _local_media(source: str, policy: MediaPolicy) -> str:
         raise RequestAdaptationError(str(error)) from error
 
 
-async def _materialize_video(source: str,
+async def _decode_data_media(source: str, target_base: str, kind: str, *, max_bytes: int | None) -> str:
+    try:
+        media_type = data_url_media_type(source)
+        if media_type.partition("/")[0] != kind:
+            raise MediaPolicyError(f"{kind.capitalize()} references need a data:{kind}/... URL, got media type "
+                                   f"{media_type or 'none'!r}.")
+        subtype_suffix = f".{media_type.partition('/')[2]}"
+        suffix = _DATA_URL_SUFFIXES.get(media_type, subtype_suffix)
+        target = f"{target_base}{suffix if suffix in _MEDIA_SUFFIXES[kind] else _DEFAULT_SUFFIXES[kind]}"
+        decoding = asyncio.ensure_future(
+            asyncio.to_thread(decode_data_url,
+                              source,
+                              target,
+                              max_bytes=get_media_reference_max_bytes() if max_bytes is None else max_bytes))
+        try:
+            await asyncio.shield(decoding)
+        except asyncio.CancelledError:
+            # The decoding thread cannot be interrupted; remove its file once it has finished writing.
+            def remove(done: asyncio.Future) -> None:
+                if not done.cancelled():
+                    done.exception()  # mark it retrieved; the request is already gone
+                discard(target)
+
+            decoding.add_done_callback(remove)
+            raise
+    except MediaPolicyError as error:
+        raise RequestAdaptationError(str(error)) from error
+    return target
+
+
+async def _materialize_media(source: str,
                              target_base: str,
+                             kind: str,
                              *,
                              max_bytes: int | None = None,
-                             policy: MediaPolicy | None = None) -> str:
+                             policy: MediaPolicy | None = None) -> tuple[str, bool]:
+    """Return a local file for a video or audio reference and whether this request created it."""
     policy = get_media_policy() if policy is None else policy
     source = source.strip()
     if not source:
-        raise RequestAdaptationError("Video references must not be empty.")
-    if source.lower().startswith(("http://", "https://")):
+        raise RequestAdaptationError(f"{kind.capitalize()} references must not be empty.")
+    # Only the scheme is lower-cased: copying a whole data: URL would double its memory.
+    if source[:8].lower().startswith(("http://", "https://")):
         suffix = os.path.splitext(urlparse(source).path)[1].lower()
-        if suffix not in _VIDEO_SUFFIXES:
-            suffix = ".mp4"
+        if suffix not in _MEDIA_SUFFIXES[kind]:
+            suffix = _DEFAULT_SUFFIXES[kind]
         target = f"{target_base}{suffix}"
         try:
             await _download_media(source, target, max_bytes=max_bytes, policy=policy)
         except RequestAdaptationError:
             raise
         except Exception as error:
-            raise RequestAdaptationError(f"Unable to fetch video reference: {error}") from error
-        return target
-    if source.lower().startswith("data:"):
-        raise RequestAdaptationError("Data video references are not supported; provide a URL or local path.")
-    return _local_media(source, policy)
+            raise RequestAdaptationError(f"Unable to fetch {kind} reference: {error}") from error
+        return target, True
+    if source[:5].lower() == "data:":
+        return await _decode_data_media(source, target_base, kind, max_bytes=max_bytes), True
+    return _local_media(source, policy), False
 
 
 async def _materialize_image_url(source: str, target_base: str, *, max_bytes: int | None, policy: MediaPolicy) -> str:
@@ -163,6 +214,8 @@ async def prepare_reference_media(
                 local_path = track(await save_base64_image_to_path(source, target))
             except Exception as error:
                 raise RequestAdaptationError(f"Unable to decode image reference: {error}") from error
+        elif source[:5].lower() == "data:":
+            raise RequestAdaptationError("Image references need a data:image/... URL.")
         elif is_remote(source):
             local_path = track(await _materialize_image_url(source, target, max_bytes=max_bytes, policy=policy))
         else:
@@ -173,6 +226,14 @@ async def prepare_reference_media(
             raise RequestAdaptationError(f"Unable to decode image reference: {error}") from error
         return local_path
 
+    async def materialize_clip(source: str, name: str, kind: str) -> str:
+        local_path, created_here = await _materialize_media(source,
+                                                            os.path.join(uploads_dir, f"{request_id}_{name}"),
+                                                            kind,
+                                                            max_bytes=max_bytes,
+                                                            policy=policy)
+        return track(local_path) if created_here else local_path
+
     try:
         index = 0
         for reference in _as_list(request.image_reference):
@@ -191,29 +252,13 @@ async def prepare_reference_media(
 
         for index, reference in enumerate(_as_list(request.video_reference)):
             if isinstance(reference, UrlVideoReference):
-                remote = is_remote(reference.video_url)
-                reference.video_url = await _materialize_video(reference.video_url,
-                                                               os.path.join(uploads_dir, f"{request_id}_video_{index}"),
-                                                               max_bytes=max_bytes,
-                                                               policy=policy)
-                if remote:
-                    track(reference.video_url)
+                reference.video_url = await materialize_clip(reference.video_url, f"video_{index}", "video")
         if request.video_path:
-            remote = is_remote(request.video_path)
-            request.video_path = await _materialize_video(request.video_path,
-                                                          os.path.join(uploads_dir, f"{request_id}_video_path"),
-                                                          max_bytes=max_bytes,
-                                                          policy=policy)
-            if remote:
-                track(request.video_path)
+            request.video_path = await materialize_clip(request.video_path, "video_path", "video")
         if request.video_url:
-            remote = is_remote(request.video_url)
-            request.video_url = await _materialize_video(request.video_url,
-                                                         os.path.join(uploads_dir, f"{request_id}_video_url"),
-                                                         max_bytes=max_bytes,
-                                                         policy=policy)
-            if remote:
-                track(request.video_url)
+            request.video_url = await materialize_clip(request.video_url, "video_url", "video")
+        for index, reference in enumerate(_as_list(request.audio_reference)):
+            reference.audio_url = await materialize_clip(reference.audio_url, f"audio_{index}", "audio")
     except BaseException:
         # Includes cancellation, so a disconnected client leaves no downloaded inputs behind.
         discard_materialized_media(created)

fastvideo/entrypoints/video_generator.py

@@ -949,7 +949,8 @@ class VideoGenerator:
             samples = torch.empty(0, device='cpu')
         else:
             samples = allocate_cpu_tensor_with_pin_fallback(
-                (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
+                (latent_batch_size, getattr(fastvideo_args.pipeline_config, "output_channels",
+                                            3), sampling_param.num_frames, sampling_param.height, sampling_param.width),
                 pin_memory=fastvideo_args.pin_cpu_memory)
         try:
             thread.join()

fastvideo/mlx_runtime/minimax_h3.py

@@ -1603,6 +1603,7 @@ class MLXMiniMaxH3DiT:
         pdd_block: tuple[int, int] | None = None,
         pdd_integration_weights: dict[str, np.ndarray] | None = None,
         record_prefix_kv: Any | None = None,
+        precomputed_text_embeds: Any | None = None,
     ):
         """Denoise-step forward served entirely from the AdaLN cache.
 
@@ -1616,6 +1617,11 @@ class MLXMiniMaxH3DiT:
         ``step_timesteps``. With a PDD checkpoint pass the step's
         ``pdd_block`` and ``pdd_integration_weights`` so the widened output
         heads are fused into the block's weighted mean.
+
+        ``precomputed_text_embeds`` may hold ``refine_text(text_rows)`` for
+        this request. The refiner is independent of the denoising timestep;
+        callers own its lifetime and must recompute it for changed text or
+        weights. No embeddings are retained on the model.
         """
         import mlx.core as mx
 
@@ -1640,7 +1646,16 @@ class MLXMiniMaxH3DiT:
             self.weights["audio_proj_in.weight"],
             self.weights["audio_proj_in.bias"],
         )
-        text_embeds = self.refine_text(text_rows)
+        if precomputed_text_embeds is None:
+            text_embeds = self.refine_text(text_rows)
+        else:
+            expected_shape = (len(layout.text_indices), self.hidden_size)
+            if precomputed_text_embeds.shape != expected_shape:
+                raise ValueError(f"Precomputed text embeddings have shape {precomputed_text_embeds.shape}; "
+                                 f"expected {expected_shape}.")
+            if precomputed_text_embeds.dtype not in (mx.float32, mx.float16, mx.bfloat16):
+                raise ValueError("Precomputed text embeddings must have a floating-point dtype.")
+            text_embeds = precomputed_text_embeds
 
         packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
         packed[mx.array(layout.text_indices)] = text_embeds

fastvideo/mlx_runtime/minimax_h3_pipeline.py

@@ -937,6 +937,10 @@ class MiniMaxH3MLXPipeline:
         elif layout.num_condition_audio_rows:
             raise ValueError("A reference-prefixed layout needs its encoded audio condition rows.")
         text = mx.array(text_rows.astype(np.float32))
+        # The token refiner depends only on this request's conditioning, not
+        # the changing latents or timesteps. A local value also prevents a
+        # resident DiT from reusing embeddings from an earlier request.
+        text_embeds = None
         dit_forward_s = 0.0
         prefix_kv = None
         if reuse_prefix_kv:
@@ -998,6 +1002,9 @@ class MiniMaxH3MLXPipeline:
                 if record_step:
                     pdd_forward_kwargs["record_prefix_kv"] = prefix_kv
                     recorded_steps.append(step_index)
+                if text_embeds is None:
+                    text_embeds = dit.refine_text(text)
+                    mx.eval(text_embeds)
                 video_velocity, audio_velocity = dit.forward_with_cache(
                     x_v,
                     x_a,
@@ -1006,6 +1013,7 @@ class MiniMaxH3MLXPipeline:
                     step_timesteps=unique,
                     row_timestep_inverse=inverse,
                     step_index=step_index,
+                    precomputed_text_embeds=text_embeds,
                     **pdd_forward_kwargs,
                 )
                 video_velocity = video_velocity[layout.num_condition_video_rows:]
@@ -1326,6 +1334,9 @@ class MiniMaxH3MLXPipeline:
         timings: dict[str, float] = {}
         peaks: dict[str, float] = {}
 
+        if vsa_impl in ("sdpa", "nax") and reuse_prefix_kv:
+            raise ValueError(
+                f"{vsa_impl.upper()} VSA does not support prefix K/V reuse; disable reuse or select another backend")
         if fast_sharpen < 0:
             raise ValueError(f"fast_sharpen must be non-negative, got {fast_sharpen}.")
         if inter_step_cooldown_s < 0:

fastvideo/mlx_runtime/minimax_h3_vsa.py

@@ -23,6 +23,10 @@ Execution backends:
 - **simd** — flash-style SIMD-group 8x8 matrix operations over the reference tile
   map for tile sizes 64 and 128 and head dimension 128. Unsupported shapes and kernel
   failures fall back to the reference backend.
+- **sdpa** — opt-in fused SDPA with the same tile mask and shared K/V. Intended
+  for M5; visits masked keys but can use faster matrix operations. ``auto`` is unchanged.
+- **nax** — opt-in M5 TensorOps kernel that visits only selected blocks and valid
+  key rows. Uses the installed MLX NAX helpers; ``auto`` is unchanged.
 
 Dense fused SDPA remains the default when VSA is disabled, the geometry is
 unsupported, or a dense-only checkpoint is loaded.
@@ -50,7 +54,7 @@ VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = {
 VSA_REF_POLICIES = ("p2_multi_region", )
 VSA_GATE_KEY_SUFFIX = "attn.to_gate_compress.weight"
 VSA_PREFIX_MODES = ("exempt", "compete")
-PUBLIC_VSA_IMPLS = ("auto", "reference", "simd")
+PUBLIC_VSA_IMPLS = ("auto", "reference", "simd", "sdpa", "nax")
 VSA_IMPLS = PUBLIC_VSA_IMPLS
 # Tile sizes where ``auto`` picks the SIMD kernel (trained Omni tiles); other tiles stay opt-in.
 _AUTO_SIMD_TILES: tuple[int, ...] = (128, )
@@ -58,7 +62,7 @@ _REFERENCE_FULL_MASK_TILE_LIMIT = 24
 _REFERENCE_FULL_MASK_MAX_ELEMENTS = 256 * 1024**2
 
 PrefixMode = Literal["exempt", "compete"]
-VSAImpl = Literal["auto", "reference", "simd"]
+VSAImpl = Literal["auto", "reference", "simd", "sdpa", "nax"]
 
 
 def validate_ref_policy(ref_policy: str | None, ref_keep_rate: float | None) -> None:
@@ -873,6 +877,20 @@ def _gate_compress_output(scores, v_tiled, gate_tiled, geometry: MiniMaxH3VSAGeo
 
 
 def resolve_impl(requested: VSAImpl, head_dim: int, tile_elems: int) -> str:
+    if requested == "nax":
+        from fastvideo.mlx_runtime.minimax_h3_vsa_nax import NAX_HEAD_DIM, NAX_TILE_SIZES
+
+        if tile_elems not in NAX_TILE_SIZES or head_dim != NAX_HEAD_DIM:
+            raise ValueError("NAX VSA requires tile 64 or 128 and head dimension 128 "
+                             f"(got tile={tile_elems}, head_dim={head_dim})")
+        return "nax"
+    if requested == "sdpa":
+        from fastvideo.mlx_runtime.minimax_h3_vsa_sdpa import SDPA_HEAD_DIM, SDPA_TILE_SIZES
+
+        if tile_elems not in SDPA_TILE_SIZES or head_dim != SDPA_HEAD_DIM:
+            raise ValueError("SDPA VSA requires tile 64 or 128 and head dimension 128 "
+                             f"(got tile={tile_elems}, head_dim={head_dim})")
+        return "sdpa"
     if requested == "reference" or (requested == "auto" and tile_elems not in _AUTO_SIMD_TILES):
         return "reference"
 
@@ -987,6 +1005,16 @@ def h3_vsa_attention(
                 if stats is not None:
                     stats.dense_fallback_reason = f"simd kernel failed: {error}"
                 video_tiled = _reference_gather_sdpa(q_tiled, k_tiled, v_tiled, block_idx, geometry, scale)
+        elif chosen == "nax":
+            from fastvideo.mlx_runtime.minimax_h3_vsa_nax import nax_block_sparse
+
+            video_tiled = nax_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry,
+                                           scale)[n_prefix_pad:]
+        elif chosen == "sdpa":
+            from fastvideo.mlx_runtime.minimax_h3_vsa_sdpa import sdpa_block_sparse
+
+            video_tiled = sdpa_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry,
+                                            scale)[n_prefix_pad:]
         elif _reference_full_mask_fits(geometry, heads):
             scores_np = np.array(scores, dtype=np.float32)
             mask = build_block_mask(

fastvideo/mlx_runtime/minimax_h3_vsa_nax.py

@@ -0,0 +1,388 @@
+# SPDX-License-Identifier: Apache-2.0 AND MIT
+# mypy: disable-error-code=no-untyped-call
+#
+# This file contains a kernel adapted from MLX. Its MIT notice follows.
+# MIT License
+#
+# Copyright © 2023 Apple Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""Opt-in M5 block-sparse TensorOps attention using installed MLX NAX helpers.
+
+Only selected blocks and their valid rows are visited, without K/V duplication
+or a dense mask. Prefix and padded query outputs remain zero. Online softmax
+and matrix accumulation can differ numerically from SIMD and fused SDPA.
+
+This backend depends on private headers from the installed MLX package, tested
+with MLX 0.32.4 and the Metal 4 SDK. Only the validated header hashes are accepted.
+Missing headers, incompatible APIs and
+unsupported hardware raise explicitly; no fallback changes the implementation.
+Importing this module performs no MLX operations. Header copyrights are retained
+and get_nax_provenance() provides hashes for reproducible benchmark receipts.
+"""
+from __future__ import annotations
+
+import hashlib
+import math
+import re
+from functools import lru_cache
+from pathlib import Path
+
+NAX_TILE_SIZES = (64, 128)
+NAX_HEAD_DIM = 128
+DEFAULT_QUERY_ROWS = 32
+_NAX_HEADER = "mlx/backend/metal/kernels/steel/attn/nax.h"
+# Private fragment layouts can change without renaming their APIs. Accept only
+# the headers used by the numerical and full-generation validation on MLX 0.32.4.
+_VALIDATED_HEADER_SHA256 = {
+    _NAX_HEADER:
+    "fb10bcae44095e9fed5c81213d03527dd5fe5ac11be11ec7784bd357c6f317ec",
+    "mlx/backend/metal/kernels/steel/defines.h":
+    "b03cea6a7d5cfe814e6838d8536e2aacdcafd2744cb408e954fda336ce21759a",
+    "mlx/backend/metal/kernels/steel/utils/integral_constant.h":
+    "886eccaac4f8f00eab5a10551fdb5315df7ab7462aba13d69bb95765c1a4735e",
+    "mlx/backend/metal/kernels/steel/utils/type_traits.h":
+    "4766c5a3809d3cf7740e63acc5703ef76b6b4e85d8baf16602c9d517a2fd426c",
+}
+
+_OPERATORS = r"""
+namespace h3_sparse_nax {
+struct MaxOp {
+  template <typename V>
+  METAL_FUNC static constexpr V apply(V x, V y) { return metal::max(x, y); }
+};
+struct SumOp {
+  template <typename V>
+  METAL_FUNC static constexpr V apply(V x, V y) { return x + y; }
+};
+struct MulOp {
+  template <typename V>
+  METAL_FUNC static constexpr V apply(V x, V y) { return x * y; }
+};
+struct ExpSubOp {
+  template <typename V>
+  METAL_FUNC static constexpr V apply(V x, V y) { return metal::fast::exp2(x - y); }
+};
+}
+"""
+
+# Adapted from MLX steel_attention_nax.h (Copyright 2024-26 Apple Inc.).
+# The fragment math/online softmax follows that implementation; this launch
+# addresses the existing H3 tiled layout and walks selected key blocks directly.
+_SOURCE = r"""
+    using namespace mlx::steel;
+    using namespace h3_sparse_nax;
+    constexpr int D = 128;
+    constexpr int BK = 32;
+    constexpr int TK = BK / 16;
+    constexpr int TD = D / 16;
+    constexpr float NEG = -3.402823466e38f;
+
+    const int qsub = int(threadgroup_position_in_grid.x);
+    const int query_tile = int(threadgroup_position_in_grid.y);
+    const int head = int(threadgroup_position_in_grid.z);
+    const int sid = int(simdgroup_index_in_threadgroup);
+    const int k_max = meta_i[0];
+    const int n_prefix = meta_i[1];
+    const int n_video = meta_i[2];
+    const int heads = meta_i[3];
+    const int qt = n_prefix + query_tile;
+    const int q_valid = vbs[qt];
+    const int row_base = qsub * BQ;
+    const int nsel = block_num[head * n_video + query_tile];
+
+    // Uniform across the threadgroup: later barriers cannot be skipped by only
+    // some SIMD groups. Inactive SIMD groups use guarded loads and stores.
+    if (row_base >= q_valid || nsel <= 0) { return; }
+
+    const int row_stride = heads * D;
+    const int qrow = row_base + sid * 16;
+    const short q_remaining = short(q_valid - qrow);
+    const device T* Q = q + (long(qt * TILE + qrow) * heads + head) * D;
+    device T* O = out + (long(qt * TILE + qrow) * heads + head) * D;
+    const float scale2 = meta_f[0] * 1.44269504088896341f;
+
+    using otile_t = NAXTile<float, 1, TD>;
+    using stile_t = NAXTile<float, 1, TK>;
+    otile_t Otile;
+    Otile.clear();
+    constexpr short kRowsPT = otile_t::kRowsPerThread;
+    metal::vec<float, kRowsPT> max_score;
+    metal::vec<float, kRowsPT> sum_score{0};
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) { max_score[i] = NEG; }
+    const short sn = otile_t::NAXFrag_t::get_coord().x;
+
+    for (int selected = 0; selected < nsel; ++selected) {
+      const int kt = block_idx[(head * n_video + query_tile) * k_max + selected];
+      const int k_valid = vbs[kt];
+      const long kv_offset = (long(kt) * TILE * heads + head) * D;
+      for (int krow = 0; krow < k_valid; krow += BK) {
+        const short k_remaining = short(metal::min(BK, k_valid - krow));
+        const device T* K = k + kv_offset + long(krow) * row_stride;
+        const device T* V = v + kv_offset + long(krow) * row_stride;
+        stile_t Stile;
+        Stile.clear();
+
+        // QK: retain MLX's four-way partial unroll for the 128-wide head.
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik += 2) {
+#pragma clang loop unroll_count(4)
+          for (short id = 0; id < TD; ++id) {
+            NAXTile<T, 1, 1> Qtile;
+            NAXTile<T, 2, 1> Ktile;
+            if (q_remaining < 16) {
+              Qtile.load_rows(Q + id * 16, row_stride, q_remaining);
+            } else {
+              Qtile.load(Q + id * 16, row_stride);
+            }
+            if (k_remaining < BK) {
+              Ktile.load_rows(K + ik * 16 * row_stride + id * 16,
+                              row_stride, k_remaining - ik * 16);
+            } else {
+              Ktile.load(K + ik * 16 * row_stride + id * 16, row_stride);
+            }
+            stile_t::NAXFrag_t::mma(
+                Stile.frag_at(0, ik), Stile.frag_at(0, ik + 1),
+                Qtile.frag_at(0, 0), metal::false_type{},
+                Ktile.frag_at(0, 0), Ktile.frag_at(1, 0), metal::true_type{});
+          }
+        }
+        STEEL_PRAGMA_UNROLL
+        for (short i = 0; i < stile_t::kElemsPerTile; ++i) {
+          Stile.elems()[i] *= scale2;
+        }
+        if (k_remaining < BK) {
+          STEEL_PRAGMA_UNROLL
+          for (short ik = 0; ik < TK; ++ik) {
+            thread auto& frag = Stile.frag_at(0, ik);
+            STEEL_PRAGMA_UNROLL
+            for (short i = 0; i < stile_t::kFragThrRows; ++i) {
+              STEEL_PRAGMA_UNROLL
+              for (short j = 0; j < stile_t::kFragThrCols; ++j) {
+                if (ik * 16 + sn + j >= k_remaining) {
+                  frag[i * stile_t::kFragThrCols + j] = NEG;
+                }
+              }
+            }
+          }
+        }
+
+        metal::vec<float, kRowsPT> new_max = max_score;
+        metal::vec<float, kRowsPT> factor;
+        Stile.template row_reduce<MaxOp>(new_max);
+        Stile.template row_bin_op<ExpSubOp>(new_max);
+        STEEL_PRAGMA_UNROLL
+        for (short i = 0; i < kRowsPT; ++i) {
+          factor[i] = metal::fast::exp2(max_score[i] - new_max[i]);
+          max_score[i] = new_max[i];
+          sum_score[i] *= factor[i];
+        }
+        Stile.template row_reduce<SumOp>(sum_score);
+        Otile.template row_bin_op<MulOp>(factor);
+        simdgroup_barrier(mem_flags::mem_none);
+
+        // PV retains float32 probabilities, matching MLX's NAX SDPA kernel.
+        STEEL_PRAGMA_UNROLL
+        for (short id = 0; id < TD; id += 2) {
+          if (id == 4) { threadgroup_barrier(mem_flags::mem_none); }
+          STEEL_PRAGMA_UNROLL
+          for (short ik = 0; ik < TK; ++ik) {
+            NAXTile<T, 1, 2> Vtile;
+            if (k_remaining < BK) {
+              Vtile.load_rows(V + ik * 16 * row_stride + id * 16,
+                              row_stride, k_remaining - ik * 16);
+            } else {
+              Vtile.load(V + ik * 16 * row_stride + id * 16, row_stride);
+            }
+            otile_t::NAXFrag_t::mma(
+                Otile.frag_at(0, id), Otile.frag_at(0, id + 1),
+                Stile.frag_at(0, ik), metal::false_type{},
+                Vtile.frag_at(0, 0), Vtile.frag_at(0, 1), metal::false_type{});
+          }
+        }
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+    metal::vec<float, kRowsPT> rcp;
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      rcp[i] = sum_score[i] > 0.f ? 1.f / sum_score[i] : 0.f;
+    }
+    Otile.template row_bin_op<MulOp>(rcp);
+    if (q_remaining > 0) {
+      if (q_remaining < 16) { Otile.store_rows(O, row_stride, q_remaining); }
+      else { Otile.store(O, row_stride); }
+    }
+"""
+
+
+def _mlx_include_root() -> Path:
+    """Resolve headers beside the imported package, including namespace installs."""
+    import mlx
+
+    for package_path in mlx.__path__:
+        include_root = Path(package_path) / "include"
+        if (include_root / _NAX_HEADER).is_file():
+            return include_root
+    raise RuntimeError("NAX VSA requires installed MLX NAX headers (tested with MLX 0.32.4); "
+                       "the mlx package has no include/mlx/backend/metal/kernels/steel/attn/nax.h")
+
+
+@lru_cache(maxsize=1)
+def _header():
+    """Inline quoted MLX includes; system Metal/MPP includes stay intact."""
+    include_root = _mlx_include_root()
+    visited = {}
+
+    def expand(relative):
+        if relative in visited:
+            return f"// Already included: {relative}\n"
+        path = include_root / relative
+        try:
+            original_bytes = path.read_bytes()
+        except OSError:
+            raise RuntimeError(f"NAX VSA cannot read required installed MLX header: {relative}") from None
+        digest = hashlib.sha256(original_bytes).hexdigest()
+        if digest != _VALIDATED_HEADER_SHA256.get(relative):
+            raise RuntimeError(f"Unvalidated MLX private header: {relative}; NAX VSA requires the headers "
+                               "validated with MLX 0.32.4. Revalidate this backend before accepting new headers.")
+        original = original_bytes.decode("utf-8")
+        visited[relative] = digest
+        content = re.sub(r"^\s*#pragma once\s*$", "", original, flags=re.MULTILINE)
+        content = re.sub(r'^\s*#include "(mlx/[^"\n]+)"\s*$',
+                         lambda match: expand(match.group(1)),
+                         content,
+                         flags=re.MULTILINE)
+        if relative == _NAX_HEADER:
+            old = "transpose_b,\n        true,\n        mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate"
+            expected_api = ("struct BaseNAXFrag", "struct NAXTile", "get_coord()", "row_reduce(", "row_bin_op(",
+                            "load_rows(", "store_rows(")
+            if content.count(old) != 2 or any(marker not in content for marker in expected_api):
+                raise RuntimeError("Installed MLX NAX private API changed; this backend was tested with MLX 0.32.4 "
+                                   "and needs compatibility review")
+        return f"\n// Begin installed MLX header: {relative}\n{content}\n// End: {relative}\n"
+
+    header = "#ifndef METAL_FUNC\n#define METAL_FUNC inline __attribute__((always_inline))\n#endif\n"
+    header += expand(_NAX_HEADER)
+    header += _OPERATORS
+    return header, visited
+
+
+def get_nax_provenance():
+    """CPU-only metadata, containing no machine-specific absolute paths."""
+    header, dependencies = _header()
+    return {
+        "kernel": "h3-true-sparse-nax-v1",
+        "relaxed_precision": True,
+        "math_mode": "safe",
+        "key_rows": 32,
+        "source_sha256": hashlib.sha256(_SOURCE.encode()).hexdigest(),
+        "expanded_header_sha256": hashlib.sha256(header.encode()).hexdigest(),
+        "installed_header_sha256": dependencies
+    }
+
+
+def _check_nax_device() -> None:
+    """The TensorOps backend is deliberately restricted to validated M5 GPUs."""
+    import mlx.core as mx
+
+    if not mx.metal.is_available():
+        raise RuntimeError("NAX VSA requires an Apple M5 GPU and Metal 4")
+    device_name = str(mx.device_info().get("device_name", ""))
+    if re.match(r"^Apple M5(?:\s|$)", device_name) is None:
+        raise RuntimeError(f"NAX VSA supports Apple M5 GPUs only, got {device_name or 'unknown GPU'}")
+
+
+@lru_cache(maxsize=1)
+def _nax_kernel():
+    import mlx.core as mx
+
+    _check_nax_device()
+    header, _ = _header()
+    if not hasattr(mx.fast, "metal_kernel"):
+        raise RuntimeError("NAX VSA requires mx.fast.metal_kernel with safe math compile_options")
+    try:
+        return mx.fast.metal_kernel(name="h3_vsa_sparse_nax",
+                                    input_names=["q", "k", "v", "block_idx", "block_num", "vbs", "meta_i", "meta_f"],
+                                    output_names=["out"],
+                                    source=_SOURCE,
+                                    header=header,
+                                    compile_options={"math_mode": "safe"})
+    except (TypeError, ValueError, RuntimeError) as error:
+        raise RuntimeError("NAX VSA kernel construction failed; this backend requires compatible MLX NAX headers "
+                           "and Metal 4 TensorOps support") from error
+
+
+def nax_block_sparse(q_tiled,
+                     k_tiled,
+                     v_tiled,
+                     block_idx,
+                     block_num,
+                     geometry,
+                     scale: float,
+                     *,
+                     query_rows: int = DEFAULT_QUERY_ROWS):
+    """Same [S_pad,H,128] contract as the existing H3 sparse implementations.
+
+    query_rows controls SIMD groups per threadgroup.
+    Active indices must be valid and unique, as emitted by the H3 router. Values
+    in inactive index slots are never read. Prefix/padding outputs are zero.
+    """
+    import mlx.core as mx
+
+    if query_rows not in (16, 32, 64, 128):
+        raise ValueError("query_rows must be 16, 32, 64 or 128")
+    if q_tiled.ndim != 3 or q_tiled.shape != k_tiled.shape or q_tiled.shape != v_tiled.shape:
+        raise ValueError("NAX sparse attention requires matching [S_pad,H,D] Q/K/V")
+    padded, heads, dim = q_tiled.shape
+    tile = geometry.tile_elems
+    if tile not in NAX_TILE_SIZES or dim != NAX_HEAD_DIM or query_rows > tile:
+        raise ValueError("NAX sparse attention requires tile 64/128, head dim 128, query_rows <= tile")
+    if padded != geometry.padded_length or heads < 1:
+        raise ValueError("Q/K/V shape does not match geometry")
+    if q_tiled.dtype not in (mx.float32, mx.float16, mx.bfloat16):
+        raise ValueError("NAX sparse attention supports float32, float16 or bfloat16")
+    if q_tiled.dtype != k_tiled.dtype or q_tiled.dtype != v_tiled.dtype:
+        raise ValueError("Q/K/V dtypes must match")
+    if (block_idx.ndim != 3 or block_idx.shape[:2] != (heads, geometry.num_video_tiles)
+            or block_num.shape != (heads, geometry.num_video_tiles)):
+        raise ValueError("Routing shapes do not match head/tile geometry")
+    if not math.isfinite(scale):
+        raise ValueError("scale must be finite")
+    if geometry.num_video_tiles == 0:
+        return mx.zeros_like(q_tiled)
+    threads = (query_rows // 16) * 32
+    result, = _nax_kernel()(inputs=[
+        mx.contiguous(q_tiled),
+        mx.contiguous(k_tiled),
+        mx.contiguous(v_tiled),
+        mx.contiguous(block_idx.astype(mx.int32)),
+        mx.contiguous(block_num.astype(mx.int32)),
+        mx.array(geometry.variable_block_sizes, dtype=mx.int32),
+        mx.array([block_idx.shape[-1], geometry.num_prefix_tiles, geometry.num_video_tiles, heads], dtype=mx.int32),
+        mx.array([scale], dtype=mx.float32)
+    ],
+                            template=[("T", q_tiled.dtype), ("TILE", tile), ("BQ", query_rows)],
+                            grid=(threads * (tile // query_rows), geometry.num_video_tiles, heads),
+                            threadgroup=(threads, 1, 1),
+                            output_shapes=[q_tiled.shape],
+                            output_dtypes=[q_tiled.dtype],
+                            init_value=0)
+    return result

fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# mypy: disable-error-code=no-untyped-call
+"""VSA tile masks through fused SDPA, sharing K/V with grouped-query attention.
+
+This opt-in backend preserves selected blocks, ragged-row masks and the trained
+compression branch. Unlike the SIMD kernel, fused SDPA visits masked key blocks;
+on M5 its faster matrix operations can outweigh that additional work. Floating-
+point accumulation differs, so this is not a bit-identical implementation.
+"""
+
+from __future__ import annotations
+
+SDPA_TILE_SIZES = (64, 128)
+SDPA_HEAD_DIM = 128
+DEFAULT_QUERY_CHUNK = 128
+_MASK_CHUNK_MAX_BYTES = 256 * 1024**2
+
+
+def sdpa_block_sparse(q_tiled,
+                      k_tiled,
+                      v_tiled,
+                      block_idx,
+                      block_num,
+                      geometry,
+                      scale: float,
+                      *,
+                      chunk_size: int = DEFAULT_QUERY_CHUNK):
+    """Return tiled [S_pad,H,D], with zero prefix rows and zero padded queries.
+
+    Query heads are ordered (original head, query tile), so GQA maps each group
+    of query tiles to its original K/V head. A tile's key mask broadcasts across
+    its query rows. At most two mask chunks are in flight. Unsupported fused
+    operations raise instead of silently allocating unfused attention scores.
+    """
+    import mlx.core as mx
+
+    if chunk_size < 1:
+        raise ValueError("SDPA VSA chunk_size must be positive")
+    if q_tiled.ndim != 3 or k_tiled.shape != q_tiled.shape or v_tiled.shape != q_tiled.shape:
+        raise ValueError("SDPA VSA requires matching [S_pad,H,D] Q/K/V shapes")
+    padded, heads, dim = q_tiled.shape
+    tile = geometry.tile_elems
+    if tile not in SDPA_TILE_SIZES or dim != SDPA_HEAD_DIM:
+        raise ValueError(f"SDPA VSA requires tile 64 or 128 and head dimension 128, got tile={tile}, dim={dim}")
+    if padded != geometry.padded_length or heads < 1:
+        raise ValueError("SDPA VSA Q/K/V shape does not match the tile geometry")
+    if q_tiled.dtype != k_tiled.dtype or q_tiled.dtype != v_tiled.dtype:
+        raise ValueError("SDPA VSA requires matching Q/K/V dtypes")
+    n_tiles = geometry.num_tiles
+    n_video = geometry.num_video_tiles
+    n_prefix = geometry.num_prefix_tiles
+    if (block_idx.ndim != 3 or block_idx.shape[:2] != (heads, n_video) or block_num.shape != (heads, n_video)):
+        raise ValueError("SDPA VSA routing shape does not match the query head/tile geometry")
+
+    query_chunk = min(chunk_size, max(1, _MASK_CHUNK_MAX_BYTES // max(heads * padded, 1)))
+    q_video = q_tiled.reshape(n_tiles, tile, heads, dim).transpose(2, 0, 1, 3)[:, n_prefix:]
+    # Fused full attention accepts these strided views: the D axis is contiguous.
+    k_shared = k_tiled.transpose(1, 0, 2)[None]
+    v_shared = v_tiled.transpose(1, 0, 2)[None]
+    sizes = mx.array(geometry.variable_block_sizes, dtype=mx.int32)
+    row = mx.arange(tile, dtype=mx.int32)
+    key_valid = row[None, :] < sizes[:, None]
+    head_index = mx.arange(heads, dtype=mx.int32)[:, None, None]
+    selected_slot = mx.arange(block_idx.shape[-1], dtype=mx.int32)[None, None, :]
+    outputs = []
+    for start in range(0, n_video, query_chunk):
+        end = min(start + query_chunk, n_video)
+        count = end - start
+        indices = block_idx[:, start:end].astype(mx.int32)
+        counts = block_num[:, start:end]
+        # Send inactive slots to a sentinel column that is discarded. Every
+        # scatter writes True, so duplicate indices cannot erase a valid key.
+        indices = mx.where(selected_slot < counts[..., None], indices, n_tiles)
+        keep = mx.zeros((heads, count, n_tiles + 1), dtype=mx.bool_)
+        tile_index = mx.arange(count, dtype=mx.int32)[None, :, None]
+        keep[head_index, tile_index, indices] = True
+        mask = (keep[:, :, :n_tiles, None] & key_valid[None, None]).reshape(1, heads * count, 1, padded)
+        query = mx.contiguous(q_video[:, start:end].reshape(1, heads * count, tile, dim))
+        output = mx.fast.scaled_dot_product_attention(query,
+                                                      k_shared,
+                                                      v_shared,
+                                                      scale=scale,
+                                                      mask=mask,
+                                                      force_fused=True)
+        output = output.reshape(heads, count, tile, dim)
+        valid_query = row[None, :] < sizes[n_prefix + start:n_prefix + end, None]
+        output = mx.where(valid_query[None, :, :, None] & (counts[..., None, None] > 0), output, 0)
+        mx.async_eval(output)
+        outputs.append(output)
+        if len(outputs) > 1:
+            mx.eval(outputs[-2])
+    if not outputs:
+        return mx.zeros_like(q_tiled)
+    mx.eval(outputs[-1])
+    video = mx.concatenate(outputs, axis=1) if len(outputs) > 1 else outputs[0]
+    video = video.transpose(1, 2, 0, 3).reshape(n_video * tile, heads, dim)
+    return mx.concatenate((mx.zeros((n_prefix * tile, heads, dim), dtype=q_tiled.dtype), video), axis=0)

fastvideo/mlx_runtime/minimax_h3_vsa_target.py

@@ -43,6 +43,9 @@ def h3_vsa_attention_target(
     ref_keep_rate: float | None = None,
 ):
     """query/gate_compress: [T, H, D] target rows; key/value: [S, H, D] packed. Returns [T, H, D]."""
+    if impl in ("sdpa", "nax"):
+        raise ValueError(
+            f"{impl.upper()} VSA does not support prefix K/V reuse; disable reuse or select another backend")
     import mlx.core as mx
 
     _, heads, dim = query.shape

fastvideo/mlx_runtime/tests/test_minimax_h3_text_refiner_cache.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Request-local text refinement must preserve denoising and refresh per request."""
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+mx = pytest.importorskip("mlx.core")
+
+from fastvideo.mlx_runtime import minimax_h3 as h3
+from fastvideo.mlx_runtime import minimax_h3_pipeline as pipeline_module
+
+
+def _toy_dit():
+    rng = np.random.default_rng(11)
+
+    def array(shape):
+        return mx.array(rng.normal(0, 0.05, shape).astype(np.float32))
+
+    hidden, ffn = 16, 24
+    config = dict(hidden_size=hidden, num_attention_heads=1, attention_head_dim=hidden,
+                  ffn_dim=ffn, in_channels=4, audio_in_channels=4, patch_size=[1, 1, 1],
+                  text_dim=8, freq_dim=hidden, time_embed_dim=hidden, rope_freq_dim=2,
+                  rope_theta=10000., norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5)
+    weights = {}
+    for name, source, target in (("context_embedder", 8, hidden), ("proj_in", 4, hidden),
+                                 ("audio_proj_in", 4, hidden), ("proj_out", hidden, 4),
+                                 ("audio_proj_out", hidden, 4),
+                                 ("time_embedder.linear_1", hidden, hidden),
+                                 ("time_embedder.linear_2", hidden, hidden),
+                                 ("norm_out.linear", hidden, 2 * hidden)):
+        weights[name + ".weight"] = array((target, source))
+        weights[name + ".bias"] = array((target,))
+    weights["token_refiner.final_norm.weight"] = mx.ones((hidden,))
+    weights["norm_out.norm.weight"] = mx.ones((hidden,))
+
+    def block():
+        result = {name: mx.ones((hidden,)) for name in
+                  ("norm1.weight", "norm2.weight", "attn.norm_q.weight", "attn.norm_k.weight")}
+        for name in ("attn.to_q.weight", "attn.to_k.weight", "attn.to_v.weight", "attn.to_out.0.weight"):
+            result[name] = array((hidden, hidden))
+        result["ff.net.0.proj.weight"] = array((2 * ffn, hidden))
+        result["ff.net.2.weight"] = array((hidden, ffn))
+        result["adaln_proj.linear.weight"] = array((18 * hidden, hidden))
+        result["adaln_proj.linear.bias"] = array((18 * hidden,))
+        return result
+
+    return h3.MLXMiniMaxH3DiT(weights, [block()], [block()], config)
+
+
+def test_precomputed_refiner_is_exact_and_runs_once_for_eight_steps(monkeypatch):
+    dit = _toy_dit()
+    layout = h3.build_packed_layout(3, 1, 2, 2, 2, patch_size=dit.patch_size)
+    timesteps, inverse = h3.build_row_timesteps(layout, .7, .4)
+    dit.precompute_adaln(timesteps)
+    text = mx.array(np.arange(24, dtype=np.float32).reshape(3, 8) / 24)
+    video = mx.ones((len(layout.video_indices), 4))
+    audio = mx.ones((len(layout.audio_indices), 4))
+    calls = []
+    original = dit.refine_text
+
+    def counted(rows):
+        calls.append(rows)
+        return original(rows)
+
+    monkeypatch.setattr(dit, "refine_text", counted)
+    kwargs = dict(layout=layout, step_timesteps=timesteps, row_timestep_inverse=inverse)
+    expected = [dit.forward_with_cache(video * (step + 1), audio, text, **kwargs) for step in range(8)]
+    mx.eval(expected)
+    assert len(calls) == 8
+    calls.clear()
+    embeds = dit.refine_text(text)
+    mx.eval(embeds)
+    actual = [dit.forward_with_cache(video * (step + 1), audio, text,
+                                    precomputed_text_embeds=embeds, **kwargs) for step in range(8)]
+    mx.eval(actual)
+    assert len(calls) == 1
+    for reference, cached in zip(expected, actual, strict=True):
+        for left, right in zip(reference, cached, strict=True):
+            np.testing.assert_array_equal(np.asarray(left), np.asarray(right))
+    with pytest.raises(ValueError, match="shape"):
+        dit.forward_with_cache(video, audio, text, precomputed_text_embeds=embeds[:1], **kwargs)
+    with pytest.raises(ValueError, match="floating-point"):
+        dit.forward_with_cache(video, audio, text, precomputed_text_embeds=embeds.astype(mx.int32), **kwargs)
+
+
+def test_pipeline_refines_again_for_next_request(monkeypatch, tmp_path):
+    dit = _toy_dit()
+    video_schedule = h3.MiniMaxH3SchedulerState.create(10, 8)
+    audio_schedule = h3.MiniMaxH3SchedulerState.create(3, 8)
+    union = h3.adaln_timestep_union(video_schedule, audio_schedule)
+    schedule = SimpleNamespace(video=video_schedule, audio=audio_schedule, adaln_timesteps=union,
+                               num_steps=8, pdd_plan=None)
+    monkeypatch.setattr(pipeline_module, "resolve_h3_denoise_schedule", lambda *a: schedule)
+    monkeypatch.setattr(pipeline_module, "_cleanup_mlx", lambda: None)
+    pipeline = object.__new__(pipeline_module.MiniMaxH3MLXPipeline)
+    pipeline.resident = False
+    pipeline.model_root = tmp_path
+    pipeline.resolve_geometry = lambda *a, **k: dict(num_frames=124, latent_frame_count=1,
+                                                   latent_height=2, latent_width=2)
+    calls = []
+    original = dit.refine_text
+
+    def counted(rows):
+        calls.append(np.asarray(rows))
+        return original(rows)
+
+    monkeypatch.setattr(dit, "refine_text", counted)
+    first_text = np.arange(24, dtype=np.float32).reshape(3, 8) / 24
+    second_text = first_text[:, ::-1].copy()
+    kwargs = dict(height=32, width=32, num_frames=124, seed=0, num_steps=8, dit=dit)
+    first = pipeline.denoise(first_text, np.ones(3, dtype=np.int64), **kwargs)
+    second = pipeline.denoise(second_text, np.ones(3, dtype=np.int64), **kwargs)
+    repeated = pipeline.denoise(first_text, np.ones(3, dtype=np.int64), **kwargs)
+    assert len(calls) == 3
+    np.testing.assert_array_equal(calls[0], first_text)
+    np.testing.assert_array_equal(calls[1], second_text)
+    assert not np.array_equal(first[0], second[0])
+    np.testing.assert_array_equal(first[0], repeated[0])
+    np.testing.assert_array_equal(first[1], repeated[1])

fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py

@@ -50,5 +50,10 @@ def test_compiled_rotary_matches_uncompiled_fp32_math():
     rotated = mx.concatenate([-second, first], axis=-1)
     out = (x_rot * cos[:, None, :] + rotated * sin[:, None, :]).astype(x.dtype)
     expected = mx.concatenate([out, x[..., rotary_dim:]], axis=-1)
-    np.testing.assert_array_equal(np.array(apply_h3_rotary(x, cos, sin).astype(mx.float32)),
-                                  np.array(expected.astype(mx.float32)))
+    # The compiled kernel may fuse the multiply-add (FMA) where the eager graph rounds each op,
+    # which can move a result by one bf16 ULP on some CPUs. Allow exactly that: rtol 2**-7 is one
+    # bf16 ULP at the bottom of each binade, so a real math error still fails.
+    np.testing.assert_allclose(np.array(apply_h3_rotary(x, cos, sin).astype(mx.float32)),
+                               np.array(expected.astype(mx.float32)),
+                               rtol=2**-7,
+                               atol=0)

fastvideo/models/dits/qwen_image21.py

@@ -0,0 +1,592 @@
+# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Native Qwen-Image-2.1 single-stream diffusion transformer.
+
+The interleaved text/image block-causal stream has no matching FastVideo
+sequence-parallel attention primitive. Mask-free dense calls use FastVideo's
+shared attention selector, retaining PyTorch SDPA for causal text and padding
+masks. Tensor and sequence parallelism are not implemented.
+Prefix K/V can reside on the host, so decoding restores only one layer's cache
+at a time instead of keeping every layer's reference images on the GPU.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.attention.selector import effective_attention_backend
+from fastvideo.configs.models.dits.qwen_image21 import QwenImage21Config
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.base import BaseDiT
+
+if TYPE_CHECKING:
+    from fastvideo.platforms import AttentionBackendEnum
+
+logger = init_logger(__name__)
+
+_IMG_TOKENS_PER_SLOT = 4
+
+
+class QwenImage21KVCacheAllocationError(RuntimeError):
+    """CPU prefix-cache allocation failed; an inference caller may retry uncached."""
+
+
+class QwenImage21KVLayerCache:
+    """Post-RoPE prefix K/V in ``(batch, tokens, heads, head_dim)`` format."""
+
+    def __init__(self, storage_device: str | torch.device | None = None):
+        self.storage_device = torch.device(storage_device) if storage_device is not None else None
+        self.k: torch.Tensor | None = None
+        self.v: torch.Tensor | None = None
+
+    def store(self, k: torch.Tensor, v: torch.Tensor) -> None:
+        if k.shape != v.shape or k.ndim != 4:
+            raise ValueError("Prefix K/V must have matching (batch, tokens, heads, head_dim) shapes")
+        # A fresh allocation releases the larger prefill tensor after the layer finishes.
+        device = self.storage_device or k.device
+        try:
+            self.k = k.detach().to(device=device, copy=True)
+            self.v = v.detach().to(device=device, copy=True)
+        except (MemoryError, RuntimeError) as exc:
+            self.clear()
+            message = str(exc).lower()
+            allocation_error = isinstance(exc, MemoryError) or any(
+                marker in message for marker in (
+                    "defaultcpuallocator", "cannot allocate memory", "can't allocate memory",
+                    "not enough memory", "out of memory", "std::bad_alloc"
+                ))
+            if device.type == "cpu" and allocation_error and "cuda" not in message and "hip" not in message:
+                raise QwenImage21KVCacheAllocationError(
+                    "CPU memory exhausted while allocating the prefix KV cache") from exc
+            raise
+
+    def get(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        if self.k is None or self.v is None:
+            raise RuntimeError("Prefix KV cache has not been populated; run an extract step first")
+        return self.k.to(device=device, dtype=dtype), self.v.to(device=device, dtype=dtype)
+
+    def clear(self) -> None:
+        self.k = self.v = None
+
+
+class QwenImage21KVCache:
+    """Per-request cache container; never share it between prompts or CFG branches."""
+
+    def __init__(self, num_layers: int, storage_device: str | torch.device | None = None):
+        if num_layers <= 0:
+            raise ValueError("num_layers must be positive")
+        self.layer_caches = [QwenImage21KVLayerCache(storage_device) for _ in range(num_layers)]
+        self.prefix_len: int | None = None
+
+    def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache:
+        return self.layer_caches[layer_idx]
+
+    def clear(self) -> None:
+        for layer in self.layer_caches:
+            layer.clear()
+        self.prefix_len = None
+
+
+def apply_rotary_emb_qwen(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    return torch.view_as_real(rotated * freqs_cis.unsqueeze(1)).flatten(3).to(x.dtype)
+
+
+class QwenImage21TemporalTimesteps(nn.Module):
+
+    def __init__(self, timestep_dim: int = 256, max_period: int = 10000, time_factor: float = 1000.0):
+        super().__init__()
+        self.timestep_dim = timestep_dim
+        self.time_factor = time_factor
+        half = timestep_dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32, device="cpu") / half)
+        self.register_buffer("freqs", freqs, persistent=False)
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        args = (self.time_factor * timestep.float())[:, None] * self.freqs.float()[None].to(timestep.device)
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.timestep_dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding.to(timestep.dtype)
+
+
+class QwenImage21TimestepEmbedding(nn.Module):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.linear_1 = ReplicatedLinear(256, embedding_dim, bias=False)
+        self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(embedding_dim, embedding_dim, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.linear_1(hidden_states)
+        return self.linear_2(self.act(hidden_states))[0]
+
+
+class QwenImage21TimestepProjEmbeddings(nn.Module):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.time_proj = QwenImage21TemporalTimesteps()
+        self.timestep_embedder = QwenImage21TimestepEmbedding(embedding_dim)
+
+    def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.timestep_embedder(self.time_proj(timestep).to(hidden_states.dtype))
+
+
+class QwenImage21ZeroCenterRMSNorm(nn.Module):
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(dim))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        rrms = torch.rsqrt(torch.mean(hidden_states**2, dim=-1, keepdim=True) + self.eps)
+        return (hidden_states * rrms * (self.weight.float() + 1)).to(input_dtype)
+
+
+class QwenImage21RMSNorm(nn.Module):
+
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        variance = hidden_states.float().pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        if self.weight.dtype in (torch.float16, torch.bfloat16):
+            hidden_states = hidden_states.to(self.weight.dtype)
+        return hidden_states * self.weight
+
+
+class QwenImage21TextProjection(nn.Module):
+
+    def __init__(self, context_in_dim: int, hidden_size: int, eps: float):
+        super().__init__()
+        self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps)
+        self.in_layer = ReplicatedLinear(context_in_dim, hidden_size, bias=False)
+        self.act = nn.GELU(approximate="tanh")
+        self.out_layer = ReplicatedLinear(hidden_size, hidden_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.in_layer(self.text_norm(hidden_states))
+        return self.out_layer(self.act(hidden_states))[0]
+
+
+class QwenImage21SwiGLUFeedForward(nn.Module):
+
+    def __init__(self, hidden_size: int, mlp_hidden_size: int):
+        super().__init__()
+        self.proj = ReplicatedLinear(hidden_size, mlp_hidden_size, bias=False)
+        self.out = ReplicatedLinear(mlp_hidden_size, hidden_size, bias=False)
+        self.gate_layer = ReplicatedLinear(hidden_size, mlp_hidden_size, bias=False)
+        self.activation_fn = nn.SiLU()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate = self.activation_fn(self.gate_layer(hidden_states)[0])
+        return self.out(gate * self.proj(hidden_states)[0])[0]
+
+
+def _select_modulation_rows(params: torch.Tensor, target_token_mask: torch.Tensor | None) -> torch.Tensor:
+    if target_token_mask is None:
+        return params.unsqueeze(1)
+    real, zero = params[:-1].unsqueeze(1), params[-1:].unsqueeze(0)
+    return torch.where(target_token_mask.view(1, -1, 1), real, zero)
+
+
+class QwenImage21AdaLayerNormContinuous(nn.Module):
+
+    def __init__(self, embedding_dim: int, eps: float):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(embedding_dim, embedding_dim, bias=False)
+        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=False)
+
+    def forward(self, hidden_states, conditioning_embedding, target_token_mask=None):
+        scale = self.linear(self.silu(conditioning_embedding).to(hidden_states.dtype))[0]
+        return self.norm(hidden_states) * (1 + _select_modulation_rows(scale, target_token_mask))
+
+
+def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]:
+    prefix_ids = image_ids[:prefix_len].tolist()
+    segments = []
+    start = 0
+    for index in range(1, prefix_len + 1):
+        if index == prefix_len or prefix_ids[index] != prefix_ids[start]:
+            segments.append((start, index, prefix_ids[start] < 0))
+            start = index
+    return segments
+
+
+def _attention(query, key, value, mask=None):
+    # SDPA expects (batch, heads, tokens, dim); projected QKV use tokens before heads.
+    output = F.scaled_dot_product_attention(query.transpose(1, 2),
+                                            key.transpose(1, 2),
+                                            value.transpose(1, 2),
+                                            attn_mask=mask,
+                                            dropout_p=0.0)
+    if mask is not None:
+        # Older CPU SDPA returns NaNs for fully masked padding queries.
+        output = output.masked_fill(~mask.any(dim=-1, keepdim=True), 0.0)
+    return output.transpose(1, 2)
+
+
+def _make_dense_attention(heads: int, dim_head: int,
+                          supported_attention_backends: tuple[AttentionBackendEnum, ...],
+                          requested_backend: AttentionBackendEnum | None) -> nn.Module | None:
+    # Keep optional attention kernels and GPU dispatch imports off the SDPA path.
+    from fastvideo.attention.layer import LocalAttention
+    from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
+    from fastvideo.platforms import AttentionBackendEnum
+    from fastvideo.utils import get_compute_dtype
+
+    dtype = get_compute_dtype()
+    selection_kwargs = {} if requested_backend is None else {"requested": requested_backend}
+    backend_cls = get_attn_backend(dim_head, dtype,
+                                   supported_attention_backends=supported_attention_backends,
+                                   **selection_kwargs)
+    backend = backend_name_to_enum(backend_cls.get_name())
+    if backend is AttentionBackendEnum.TORCH_SDPA:
+        return None
+
+    class DenseLocalAttention(LocalAttention):
+        """Qwen owns segment lengths and has already applied its complex RoPE.
+
+        Resolve explicit component requests through the public selector above;
+        direct construction without a request retains scope/env selection.
+        LocalAttention's ordinary forward reads unrelated global metadata,
+        which must not reach our unequal Q/KV calls.
+        """
+
+        def __init__(self):
+            nn.Module.__init__(self)
+            self.softmax_scale = dim_head**-0.5
+            self.attn_impl = backend_cls.get_impl_cls()(num_heads=heads, head_size=dim_head,
+                                                        softmax_scale=self.softmax_scale,
+                                                        num_kv_heads=heads, causal=False)
+            self.num_heads = self.num_kv_heads = heads
+            self.head_size, self.backend, self.dtype = dim_head, backend, dtype
+
+        def forward(self, q, k, v):
+            return self.attn_impl.forward(q, k, v, None)
+
+    return DenseLocalAttention()
+
+
+class QwenImage21Attention(nn.Module):
+
+    def __init__(self, dim: int, heads: int, dim_head: int, eps: float, *,
+                 requested_backend: AttentionBackendEnum | None = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None):
+        super().__init__()
+        self.heads = heads
+        self.inner_dim = heads * dim_head
+        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=False)
+        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=False)
+        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=False)
+        self.to_out = nn.ModuleList([ReplicatedLinear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)])
+        self.norm_q = QwenImage21RMSNorm(dim_head, eps)
+        self.norm_k = QwenImage21RMSNorm(dim_head, eps)
+        self.dense_attn = None
+        if requested_backend is None or requested_backend.name != "TORCH_SDPA":
+            supported = supported_attention_backends or _DEFAULT_CONFIG.arch_config._supported_attention_backends
+            self.dense_attn = _make_dense_attention(heads, dim_head, supported, requested_backend)
+
+    def _attend(self, query, key, value, mask=None):
+        if mask is not None or self.dense_attn is None:
+            return _attention(query, key, value, mask)
+        return self.dense_attn(query, key, value)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                rotary_emb: torch.Tensor | None = None,
+                attention_mask: torch.Tensor | None = None,
+                layer_cache: QwenImage21KVLayerCache | None = None,
+                kv_cache_mode: str | None = None,
+                cache_write_slice: slice | None = None,
+                segments: list[tuple[int, int, bool]] | None = None,
+                key_valid: torch.Tensor | None = None) -> torch.Tensor:
+        query = self.to_q(hidden_states)[0].unflatten(-1, (self.heads, -1))
+        key = self.to_k(hidden_states)[0].unflatten(-1, (self.heads, -1))
+        value = self.to_v(hidden_states)[0].unflatten(-1, (self.heads, -1))
+        query, key = self.norm_q(query).to(value.dtype), self.norm_k(key).to(value.dtype)
+        if rotary_emb is not None:
+            query = apply_rotary_emb_qwen(query, rotary_emb)
+            key = apply_rotary_emb_qwen(key, rotary_emb)
+        if layer_cache is not None:
+            if kv_cache_mode == "extract" and cache_write_slice is not None:
+                layer_cache.store(key[:, cache_write_slice], value[:, cache_write_slice])
+            elif kv_cache_mode == "cached":
+                cached_k, cached_v = layer_cache.get(key.device, key.dtype)
+                if cached_k.shape[0] != key.shape[0] or cached_k.shape[2:] != key.shape[2:]:
+                    raise ValueError("Cached prefix batch/head shape differs from the current request")
+                key = torch.cat([cached_k, key], dim=1)
+                value = torch.cat([cached_v, value], dim=1)
+                del cached_k, cached_v
+        if segments is None:
+            output = self._attend(query, key, value, attention_mask)
+        else:
+            outputs = []
+            for start, end, is_text in segments:
+                mask = None
+                if is_text:
+                    length = end - start
+                    mask = torch.cat([torch.ones(length, start, dtype=torch.bool, device=query.device),
+                                      torch.tril(torch.ones(length, length, dtype=torch.bool, device=query.device))],
+                                     dim=1)[None, None]
+                if key_valid is not None:
+                    valid = key_valid[:, None, None, :end]
+                    mask = valid if mask is None else mask & valid
+                outputs.append(self._attend(query[:, start:end], key[:, :end], value[:, :end], mask))
+            prefix_len = segments[-1][1] if segments else 0
+            mask = None if key_valid is None else key_valid[:, None, None, :]
+            outputs.append(self._attend(query[:, prefix_len:], key, value, mask))
+            output = torch.cat(outputs, dim=1)
+        output = output.flatten(2, 3).to(query.dtype)
+        return self.to_out[1](self.to_out[0](output)[0])
+
+
+class QwenImage21TransformerBlock(nn.Module):
+
+    def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: int, eps: float, *,
+                 requested_backend: AttentionBackendEnum | None = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None):
+        super().__init__()
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = QwenImage21Attention(dim, num_attention_heads, attention_head_dim, eps,
+                                       requested_backend=requested_backend,
+                                       supported_attention_backends=supported_attention_backends)
+        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = QwenImage21SwiGLUFeedForward(dim, dim * mlp_ratio)
+
+    def _modulate(self, hidden_states, mod_params, target_token_mask):
+        scale, gate = mod_params.chunk(2, dim=-1)
+        scale = _select_modulation_rows(scale, target_token_mask)
+        gate = _select_modulation_rows(gate, target_token_mask)
+        return hidden_states * (1 + scale), gate
+
+    def forward(self, hidden_states, modulation, target_token_mask=None, **attention_kwargs):
+        mod1, mod2 = modulation.chunk(2, dim=-1)
+        normalized, gate = self._modulate(self.img_norm1(hidden_states), mod1, target_token_mask)
+        hidden_states = hidden_states + gate.tanh() * self.attn(normalized, **attention_kwargs)
+        normalized, gate = self._modulate(self.img_norm2(hidden_states), mod2, target_token_mask)
+        hidden_states = hidden_states + gate.tanh() * self.img_mlp(normalized)
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+        return hidden_states
+
+
+class QwenImage21Rope(nn.Module):
+
+    def __init__(self, theta: int, axes_dim: tuple[int, int, int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def rope_params(self, index: torch.Tensor, dim: int) -> torch.Tensor:
+        powers = torch.arange(0, dim, 2, dtype=torch.float32, device=index.device).div(dim)
+        freqs = torch.outer(index.float(), 1.0 / torch.pow(self.theta, powers))
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def forward(self, img_shapes, image_pad_mask, device):
+        frame_index, image_height_index, image_width_index = [], [], []
+        cursor, position = 0, 0
+        is_image = image_pad_mask.tolist()
+        for _, height, width in img_shapes:
+            block_start = is_image.index(True, cursor)
+            text_len = block_start - cursor
+            frame_index.extend(range(position, position + text_len))
+            position += text_len
+            cursor = block_start + height * width
+            frame_index.extend([position] * (height * width))
+            position += max(height, width)
+            image_height_index.extend([h for h in range(-(height - height // 2), height // 2) for _ in range(width)])
+            image_width_index.extend([w for _ in range(height) for w in range(-(width - width // 2), width // 2)])
+        if cursor < len(is_image):
+            frame_index.extend(range(position, position + len(is_image) - cursor))
+        # Derive frequencies from signed coordinates, retaining Qwen's CPU fp32/complex64 numerics.
+        frame_index = torch.tensor(frame_index, dtype=torch.long, device="cpu")
+        image_pad_mask = torch.tensor(is_image, dtype=torch.bool, device="cpu")
+        height_index, width_index = frame_index.clone(), frame_index.clone()
+        height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device="cpu")
+        width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device="cpu")
+        return torch.cat([self.rope_params(index, dim) for index, dim in
+                          zip((frame_index, height_index, width_index), self.axes_dim, strict=True)], dim=-1).to(device)
+
+
+_DEFAULT_CONFIG = QwenImage21Config()
+
+
+class QwenImage21Transformer2DModel(BaseDiT):
+    _fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions
+    _compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions
+    _supported_attention_backends = _DEFAULT_CONFIG.arch_config._supported_attention_backends
+    param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
+    reverse_param_names_mapping = _DEFAULT_CONFIG.arch_config.reverse_param_names_mapping
+
+    def __init__(self, config: QwenImage21Config, hf_config: dict[str, Any] | None = None, **kwargs):
+        super().__init__(config, hf_config or {}, **kwargs)
+        if getattr(config, "quant_config", None) is not None:
+            raise NotImplementedError("Qwen-Image-2.1 transformer quantization has not been implemented")
+        self.inner_dim = self.hidden_size = config.hidden_size
+        self.out_channels = config.out_channels
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.num_channels_latents
+        self.pos_embed = QwenImage21Rope(10000, config.axes_dims_rope)
+        self.time_text_embed = QwenImage21TimestepProjEmbeddings(self.inner_dim)
+        self.txt_in = QwenImage21TextProjection(config.context_in_dim, self.inner_dim, config.eps)
+        self.img_in = ReplicatedLinear(config.in_channels * config.patch_size**2, self.inner_dim, bias=False)
+        self.modulation = nn.Sequential(nn.SiLU(), ReplicatedLinear(self.inner_dim, 4 * self.inner_dim, bias=False))
+        requested_backend = effective_attention_backend(config)
+        self.transformer_blocks = nn.ModuleList([
+            QwenImage21TransformerBlock(self.inner_dim, config.num_attention_heads, config.attention_head_dim,
+                                        config.mlp_ratio, config.eps, requested_backend=requested_backend,
+                                        supported_attention_backends=self._supported_attention_backends)
+            for _ in range(config.num_layers)
+        ])
+        self.norm_out = QwenImage21AdaLayerNormContinuous(self.inner_dim, config.eps)
+        self.proj_out = ReplicatedLinear(self.inner_dim, config.patch_size**2 * config.out_channels, bias=False)
+
+    @staticmethod
+    def build_token_metadata(image_pad_mask, img_shapes):
+        positions = image_pad_mask.nonzero(as_tuple=True)[0]
+        lengths = [math.prod(shape) for shape in img_shapes]
+        if not lengths or any(length <= 0 for length in lengths) or sum(lengths) != positions.numel():
+            raise ValueError("img_shapes token counts must match image_pad_mask")
+        image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long)
+        ids = torch.repeat_interleave(torch.arange(len(lengths), device=image_pad_mask.device),
+                                     torch.tensor(lengths, device=image_pad_mask.device))
+        image_ids[positions] = ids
+        target_mask = torch.zeros_like(image_pad_mask)
+        target_mask[positions[-lengths[-1]:]] = True
+        return image_ids, target_mask
+
+    def _validate_inputs(self, hidden_states, encoder_hidden_states, timestep, img_shapes, img_mask, text_mask):
+        batch_size = hidden_states.shape[0]
+        if hidden_states.ndim != 3 or encoder_hidden_states.ndim != 3:
+            raise ValueError("Latents and text embeddings must have (batch, tokens, channels) shapes")
+        if encoder_hidden_states.shape[0] != batch_size:
+            raise ValueError("Latent and text batch sizes must match")
+        if len(img_shapes) != batch_size or not img_shapes or not img_shapes[0]:
+            raise ValueError("img_shapes must provide a nonempty image layout for each sample")
+        if any(shapes != img_shapes[0] for shapes in img_shapes):
+            raise ValueError("All samples must share the same image layout")
+        if any(frame != 1 or height <= 0 or width <= 0 for frame, height, width in img_shapes[0]):
+            raise ValueError("Qwen-Image-2.1 expects positive single-frame image shapes")
+        target_tokens = math.prod(img_shapes[0][-1])
+        if target_tokens % _IMG_TOKENS_PER_SLOT:
+            raise ValueError("The target latent token count must be divisible by four")
+        if hidden_states.shape[1] != sum(math.prod(shape) for shape in img_shapes[0]):
+            raise ValueError("Packed latents must include all references followed by the target image")
+        expected_slots = encoder_hidden_states.shape[1] + target_tokens // _IMG_TOKENS_PER_SLOT
+        if img_mask.shape != (batch_size, expected_slots):
+            raise ValueError("img_mask must include the VLM sequence and appended target-image slots")
+        if not torch.equal(img_mask.bool(), img_mask[0:1].bool().expand_as(img_mask)):
+            raise ValueError("All samples must share the same image-slot layout")
+        if not bool(img_mask[:, -target_tokens // _IMG_TOKENS_PER_SLOT:].all()):
+            raise ValueError("Target-image slots must be appended at the end of img_mask")
+        if int(img_mask[0].count_nonzero()) * _IMG_TOKENS_PER_SLOT != hidden_states.shape[1]:
+            raise ValueError("Image-slot count must match the reference and target latent token count")
+        if timestep.numel() != batch_size:
+            raise ValueError("timestep must contain one value per sample")
+        if text_mask is not None and text_mask.shape != encoder_hidden_states.shape[:2]:
+            raise ValueError("encoder_hidden_states_mask must match the VLM embedding sequence")
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor,
+                timestep: torch.Tensor,
+                img_shapes: list[list[tuple[int, int, int]]],
+                img_mask: torch.Tensor,
+                encoder_hidden_states_mask: torch.Tensor | None = None,
+                attention_kwargs: dict[str, Any] | None = None,
+                kv_cache: QwenImage21KVCache | None = None,
+                kv_cache_mode: str | None = None,
+                return_dict: bool | None = None,
+                **kwargs) -> torch.Tensor | tuple[torch.Tensor]:
+        """Return joint-sequence predictions; the pipeline selects the target tail.
+
+        ``kv_cache_mode='extract'`` stores timestep-independent text/reference
+        K/V. ``'cached'`` returns only the target image tokens. By default the
+        return value is a tensor; ``return_dict=False`` selects a one-item tuple.
+        """
+        self._validate_inputs(hidden_states, encoder_hidden_states, timestep, img_shapes, img_mask,
+                              encoder_hidden_states_mask)
+        if kv_cache is not None and not self.config.causal_condition:
+            raise ValueError("Prefix KV caching requires causal_condition=True")
+        if kv_cache is not None and kv_cache_mode not in ("extract", "cached"):
+            raise ValueError("kv_cache_mode must be extract or cached when a cache is provided")
+        if kv_cache is None and kv_cache_mode is not None:
+            raise ValueError("kv_cache_mode requires a cache container")
+        if kv_cache is not None and len(kv_cache.layer_caches) != len(self.transformer_blocks):
+            raise ValueError("KV cache layer count differs from the transformer")
+        if attention_kwargs:
+            raise ValueError("Custom attention processors and LoRA scaling are not implemented for Qwen-Image-2.1")
+        batch_size = hidden_states.shape[0]
+        hidden_states = self.img_in(hidden_states)[0]
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+        img_mask = img_mask.bool().to(hidden_states.device)
+        repeats = torch.where(img_mask, _IMG_TOKENS_PER_SLOT, 1)[0]
+        image_pad_mask = torch.repeat_interleave(img_mask[0], repeats)
+        target_tokens = math.prod(img_shapes[0][-1])
+        joint = torch.cat([encoder_hidden_states,
+                           encoder_hidden_states.new_zeros(batch_size, target_tokens // 4, self.inner_dim)], dim=1)
+        joint = joint.repeat_interleave(repeats, dim=1)
+        joint[:, image_pad_mask] = hidden_states
+        rotary_emb = self.pos_embed(img_shapes[0], image_pad_mask, hidden_states.device)
+        image_ids, target_mask = self.build_token_metadata(image_pad_mask, img_shapes[0])
+        timestep = timestep.reshape(batch_size).to(device=hidden_states.device, dtype=hidden_states.dtype)
+        if self.config.causal_condition:
+            timestep = torch.cat([timestep, timestep.new_zeros(1)])
+            modulation_mask = target_mask
+        else:
+            modulation_mask = None
+        temb = self.time_text_embed(timestep, hidden_states)
+        modulation = self.modulation[1](self.modulation[0](temb))[0]
+        joint_key_valid = None
+        if encoder_hidden_states_mask is not None:
+            joint_key_valid = torch.ones(batch_size, image_pad_mask.numel(), dtype=torch.bool,
+                                         device=hidden_states.device)
+            text_positions = (~image_pad_mask).nonzero(as_tuple=True)[0]
+            vlm_text_positions = ~img_mask[0, :encoder_hidden_states.shape[1]]
+            valid_text = encoder_hidden_states_mask.bool().to(hidden_states.device)[:, vlm_text_positions]
+            joint_key_valid[:, text_positions] = valid_text
+        prefix_len = image_pad_mask.numel() - target_tokens
+        if kv_cache_mode == "cached":
+            if kv_cache.prefix_len != prefix_len:
+                raise ValueError("Cached prefix layout differs from the current request or has not been extracted")
+            joint, rotary_emb = joint[:, prefix_len:], rotary_emb[prefix_len:]
+            modulation_mask = modulation_mask[prefix_len:]
+            attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :]
+            segments, block_key_valid, cache_slice = None, None, None
+        else:
+            segments = _qwenimage21_prefix_segments(image_ids, prefix_len)
+            attention_mask, block_key_valid = None, joint_key_valid
+            cache_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None
+            if kv_cache is not None:
+                kv_cache.prefix_len = prefix_len
+        for index, block in enumerate(self.transformer_blocks):
+            joint = block(joint, modulation,
+                          target_token_mask=modulation_mask,
+                          rotary_emb=rotary_emb,
+                          attention_mask=attention_mask,
+                          layer_cache=kv_cache.get_layer(index) if kv_cache is not None else None,
+                          kv_cache_mode=kv_cache_mode,
+                          cache_write_slice=cache_slice,
+                          segments=segments,
+                          key_valid=block_key_valid)
+        output = self.proj_out(self.norm_out(joint, temb, modulation_mask))[0]
+        return (output,) if return_dict is False else output
+
+
+EntryClass = QwenImage21Transformer2DModel

fastvideo/models/encoders/qwen_image21.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native Qwen3-VL conditioner returning Qwen-Image-2.1's pre-norm features."""
+
+from collections.abc import Iterable
+
+import torch
+
+from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
+from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+
+class QwenImage21Qwen3VLConditioner(MiniMaxH3Qwen3VLConditioner):
+    """Qwen3-VL text and image features, without autoregressive generation.
+
+    The shared native stack stops immediately after the configured last
+    decoder layer. Its final RMSNorm is loaded for checkpoint completeness
+    but never applied to the image transformer's conditioning features.
+    Sequence parallel execution is not implemented for this conditioner.
+    """
+
+    supported_checkpoint_quantization_methods: frozenset[str] = frozenset()
+    _checkpoint_quantization_configs: dict = {}
+
+    def __init__(self, config: QwenImage21Qwen3VLConfig) -> None:
+        super().__init__(config)
+        # CPU construction of this constant avoids device-dependent power
+        # rounding accumulating across the complete 36-layer decoder stack.
+        rotary = self.language_model.rotary_emb
+        exponents = torch.arange(0, config.head_dim, 2, dtype=torch.float32, device="cpu") / config.head_dim
+        rotary.inv_freq = (1.0 / (config.rope_theta**exponents)).to(rotary.inv_freq.device)
+
+    def _visual_features(
+        self,
+        pixels: torch.Tensor,
+        grid_thw: torch.Tensor,
+    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        if grid_thw.ndim != 2 or grid_thw.shape[1] != 3 or grid_thw.shape[0] == 0:
+            raise ValueError("image_grid_thw must contain one nonempty [time, height, width] row per image")
+        grids = grid_thw.to(device="cpu", dtype=torch.long)
+        merge = self.config.vision_spatial_merge_size
+        if bool((grids <= 0).any()) or bool((grids[:, 1:] % merge != 0).any()):
+            raise ValueError("Image grids must be positive and spatially divisible by the vision merge size")
+        patch_counts = grids.prod(dim=1).tolist()
+        if sum(patch_counts) != pixels.shape[0]:
+            raise ValueError("pixel_values patch count does not match image_grid_thw")
+
+        # Keep packed projection shapes: splitting images changes BF16 GEMM
+        # rounding before the complete text decoder amplifies the difference.
+        # Vision attention still executes separately for each image grid.
+        return self.visual(pixels.to(self.visual.patch_embed.proj.weight.dtype), grid_thw)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        image_grid_thw: torch.Tensor | None = None,
+        mm_token_type_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Return ``[batch, tokens, width]``, or ``[tokens, width]`` for 1-D IDs.
+
+        Padding is zeroed and removed by the prompt stage using the original
+        attention mask before stripping its template.
+        """
+        single_sequence = input_ids.ndim == 1
+        if single_sequence:
+            input_ids = input_ids.unsqueeze(0)
+            if attention_mask is not None and attention_mask.ndim == 1:
+                attention_mask = attention_mask.unsqueeze(0)
+            if mm_token_type_ids is not None and mm_token_type_ids.ndim == 1:
+                mm_token_type_ids = mm_token_type_ids.unsqueeze(0)
+        if input_ids.ndim != 2 or input_ids.shape[1] == 0:
+            raise ValueError("Qwen-Image-2.1 input_ids must have shape [batch, nonempty sequence] or [sequence]")
+        if attention_mask is not None:
+            if attention_mask.shape != input_ids.shape:
+                raise ValueError("attention_mask must have the same shape as input_ids")
+            if bool((attention_mask.to(torch.bool).sum(dim=1) == 0).any()):
+                raise ValueError("Each conditioning sequence must contain an unmasked token")
+        if (pixel_values is None) != (image_grid_thw is None):
+            raise ValueError("pixel_values and image_grid_thw must be provided together")
+        if bool((input_ids == self.config.video_token_id).any()):
+            raise ValueError("Qwen-Image-2.1 conditioning accepts image references, not video tokens")
+        if mm_token_type_ids is not None:
+            expected = (input_ids == self.config.image_token_id).to(mm_token_type_ids.dtype)
+            if mm_token_type_ids.shape != input_ids.shape or not torch.equal(mm_token_type_ids, expected):
+                raise ValueError("mm_token_type_ids must mark image placeholders as 1 and other tokens as 0")
+
+        inputs_embeds = self.language_model.embed_tokens(input_ids)
+        visual_mask = None
+        deepstack = None
+        if pixel_values is not None:
+            features, deepstack = self._visual_features(pixel_values, image_grid_thw)
+            features = features.to(inputs_embeds.device, inputs_embeds.dtype)
+            visual_mask = self._placeholder_mask(input_ids, inputs_embeds, self.config.image_token_id,
+                                                features, "image")
+            inputs_embeds = inputs_embeds.masked_scatter(visual_mask.unsqueeze(-1), features)
+        elif bool((input_ids == self.config.image_token_id).any()):
+            raise ValueError("Image placeholder tokens require pixel_values and image_grid_thw")
+
+        # Text-only prefills retain positions in the padded sequence. A
+        # constant RoPE shift is mathematically equivalent, but changes BF16
+        # rounding; image prefills need each sample's multimodal layout.
+        positions = self._get_rope_index(input_ids, image_grid_thw, None,
+                                         attention_mask if image_grid_thw is not None else None)
+        # Unpadding changes projection and attention shapes in reduced
+        # precision. Preserve the batch through every decoder layer instead.
+        hidden_states = self.language_model(inputs_embeds, positions, attention_mask, visual_mask, deepstack)
+        if attention_mask is not None:
+            hidden_states = hidden_states.masked_fill(~attention_mask.to(torch.bool).unsqueeze(-1), 0)
+        return hidden_states[0] if single_sequence else hidden_states
+
+    @torch.no_grad()
+    def encode_ids(
+        self,
+        input_ids: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        image_grid_thw: torch.Tensor | None = None,
+        mm_token_type_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Encode one sequence without imposing a tokenizer truncation limit."""
+        if input_ids.ndim != 1:
+            raise ValueError("encode_ids requires 1-D input_ids; use forward for a padded batch")
+        return self.forward(
+            input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+            mm_token_type_ids=mm_token_type_ids,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        parameters = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for source_name, tensor in weights:
+            if source_name == "lm_head.weight":
+                continue
+            name = source_name.removeprefix("model.")
+            if name not in parameters:
+                raise ValueError(f"Unexpected Qwen-Image-2.1 Qwen3-VL checkpoint key: {source_name}")
+            if name in loaded:
+                raise ValueError(f"Duplicate Qwen-Image-2.1 Qwen3-VL checkpoint key: {source_name}")
+            parameter = parameters[name]
+            loader = getattr(parameter, "weight_loader", default_weight_loader)
+            loader(parameter, tensor)
+            loaded.add(name)
+        return loaded
+
+
+EntryClass = QwenImage21Qwen3VLConditioner

fastvideo/models/loader/component_loader.py

@@ -1032,7 +1032,7 @@ class VAELoader(ComponentLoader):
         # Diffusers-format AutoencoderKL checkpoints should match exactly; load
         # strictly so missing/unexpected keys are surfaced early.
         # Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well.
-        strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"}
+        strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE", "AutoencoderKLQwenImage21"}
         vae.load_state_dict(loaded, strict=strict_load)
         del loaded
         if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:

fastvideo/models/registry.py

@@ -71,7 +71,12 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
 }
 
+_TEXT_TO_IMAGE_DIT_MODELS = {
+    "QwenImage21Transformer2DModel": ("dits", "qwen_image21", "QwenImage21Transformer2DModel"),
+}
+
 _TEXT_ENCODER_MODELS = {
+    "QwenImage21Qwen3VLConditioner": ("encoders", "qwen_image21", "QwenImage21Qwen3VLConditioner"),
     "MMAudioDFNCLIPTextEncoder": ("encoders", "mmaudio_clip", "MMAudioDFNCLIPTextEncoder"),
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection": ("encoders", "clip", "CLIPTextModelWithProjection"),
@@ -108,6 +113,7 @@ _VAE_MODELS = {
     "AutoencoderKLWan": ("wan", "vae", "AutoencoderKLWan"),
     "LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"),
     "AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"),
+    "AutoencoderKLQwenImage21": ("vaes", "qwen_image21", "AutoencoderKLQwenImage21"),
     "AutoencoderKLGen3CTokenizer": ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
     "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
     "AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"),
@@ -154,6 +160,7 @@ _UPSAMPLERS = {
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
+    **_TEXT_TO_IMAGE_DIT_MODELS,
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,

fastvideo/models/vaes/qwen_image21.py

@@ -0,0 +1,656 @@
+# Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved.
+#
+# 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.
+"""Native single-frame RGBA autoencoder for Qwen-Image-2.1."""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from fastvideo.configs.models.vaes.base import VAEConfig
+from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEConfig
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class DiagonalGaussianDistribution:
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None:
+        self.parameters = parameters
+        self.mean, self.logvar = parameters.chunk(2, dim=1)
+        self.logvar = self.logvar.clamp(-30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if deterministic:
+            self.std = torch.zeros_like(self.mean)
+            self.var = torch.zeros_like(self.mean)
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+    def sample(self, generator: torch.Generator | list[torch.Generator] | None = None) -> torch.Tensor:
+        if isinstance(generator, list):
+            if len(generator) != self.mean.shape[0]:
+                raise ValueError("A generator list must contain one generator per batch element.")
+            noise = torch.cat([
+                self._noise(self.mean[index:index + 1], item) for index, item in enumerate(generator)
+            ])
+        else:
+            noise = self._noise(self.mean, generator)
+        return self.mean + self.std * noise
+
+    @staticmethod
+    def _noise(sample: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
+        noise_device = sample.device if generator is None else generator.device
+        return torch.randn(sample.shape, generator=generator, device=noise_device, dtype=sample.dtype).to(sample.device)
+
+    def kl(self, other: DiagonalGaussianDistribution | None = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.zeros(1, device=self.mean.device, dtype=self.mean.dtype)
+        if other is None:
+            terms = self.mean.square() + self.var - 1.0 - self.logvar
+        else:
+            terms = ((self.mean - other.mean).square() + self.var) / other.var - 1.0 - self.logvar + other.logvar
+        return 0.5 * terms.sum(dim=(1, 2, 3))
+
+    def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor:
+        if self.deterministic:
+            return torch.zeros(1, device=self.mean.device, dtype=self.mean.dtype)
+        return 0.5 * (math.log(2 * math.pi) + self.logvar + (sample - self.mean).square() / self.var).sum(dim=dims)
+
+
+@dataclass
+class AutoencoderKLOutput:
+    latent_dist: DiagonalGaussianDistribution
+
+
+@dataclass
+class DecoderOutput:
+    sample: torch.Tensor
+
+
+class QwenImage21AvgDown3D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = factor_t * factor_s * factor_s
+        if in_channels * self.factor % out_channels:
+            raise ValueError("Average shortcut channels must be divisible by the downsampling factor.")
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        x = F.pad(x, (0, 0, 0, 0, pad_t, 0))
+        batch, channels, frames, height, width = x.shape
+        x = x.view(batch, channels, frames // self.factor_t, self.factor_t,
+                   height // self.factor_s, self.factor_s, width // self.factor_s, self.factor_s)
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(batch, channels * self.factor, frames // self.factor_t,
+                   height // self.factor_s, width // self.factor_s)
+        x = x.view(batch, self.out_channels, self.group_size, frames // self.factor_t,
+                   height // self.factor_s, width // self.factor_s)
+        return x.mean(dim=2)
+
+
+class QwenImage21DupUp3D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1) -> None:
+        super().__init__()
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        factor = factor_t * factor_s * factor_s
+        if out_channels * factor % in_channels:
+            raise ValueError("Duplicate shortcut channels must be divisible by the upsampling factor.")
+        self.repeats = out_channels * factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk: bool = False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(x.shape[0], self.out_channels, self.factor_t, self.factor_s, self.factor_s,
+                   x.shape[2], x.shape[3], x.shape[4])
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(x.shape[0], self.out_channels, x.shape[2] * self.factor_t,
+                   x.shape[4] * self.factor_s, x.shape[6] * self.factor_s)
+        return x[:, :, self.factor_t - 1:] if first_chunk else x
+
+
+class QwenImage21CausalConv3d(nn.Conv2d):
+    """Spatial convolution retaining the image VAE's singleton frame axis."""
+
+    def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0) -> None:
+        super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
+        self._padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0])
+        self.padding = (0, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 5 or x.shape[2] != 1:
+            raise ValueError("Qwen-Image-2.1 spatial convolutions require exactly one frame.")
+        return super().forward(F.pad(x.squeeze(2), self._padding)).unsqueeze(2)
+
+
+class QwenImage21RMSNorm(nn.Module):
+    def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
+        super().__init__()
+        broadcast_dims = (1, 1) if images else (1, 1, 1)
+        shape = (dim, *broadcast_dims) if channel_first else (dim,)
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        upcast = x.dtype in (torch.float16, torch.bfloat16) or any(
+            name in str(x.dtype) for name in ("float4_", "float8_")
+        )
+        normalized = F.normalize(x.float() if upcast else x, dim=1 if self.channel_first else -1).to(x.dtype)
+        return normalized * self.scale * self.gamma + self.bias
+
+
+class QwenImage21Upsample(nn.Upsample):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x.float()).type_as(x)
+
+
+class QwenImage21Resample(nn.Module):
+    def __init__(self, dim: int, mode: str, upsample_out_dim: int | None = None) -> None:
+        super().__init__()
+        self.mode = mode
+        if upsample_out_dim is None:
+            upsample_out_dim = dim // 2
+        if mode in ("upsample2d", "upsample3d"):
+            self.resample = nn.Sequential(
+                QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
+            )
+            if mode == "upsample3d":
+                self.time_conv = QwenImage21CausalConv3d(dim, dim * 2, (1, 1), padding=(0, 0))
+        elif mode in ("downsample2d", "downsample3d"):
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            if mode == "downsample3d":
+                self.time_conv = QwenImage21CausalConv3d(dim, dim, (1, 1), stride=(1, 1), padding=(0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, channels, frames, height, width = x.shape
+        # Temporal convolutions have no preceding frame in this image model.
+        x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
+        x = self.resample(x)
+        return x.view(batch, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4)
+
+
+class QwenImage21ResidualBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm1 = QwenImage21RMSNorm(in_dim, images=False)
+        self.conv1 = QwenImage21CausalConv3d(in_dim, out_dim, 3, padding=1)
+        self.norm2 = QwenImage21RMSNorm(out_dim, images=False)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1)
+        self.conv_shortcut = QwenImage21CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = self.conv_shortcut(x)
+        x = self.conv1(F.silu(self.norm1(x)))
+        x = self.conv2(self.dropout(F.silu(self.norm2(x))))
+        return x + residual
+
+
+class QwenImage21AttentionBlock(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.norm = QwenImage21RMSNorm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        batch, channels, frames, height, width = x.shape
+        x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
+        qkv = self.to_qkv(self.norm(x)).reshape(batch * frames, 1, channels * 3, -1)
+        q, k, v = qkv.permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+        # The VAE's one-head spatial attention has no distributed sequence layout.
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = x.squeeze(1).permute(0, 2, 1).reshape(batch * frames, channels, height, width)
+        x = self.proj(x).view(batch, frames, channels, height, width).permute(0, 2, 1, 3, 4)
+        return x + residual
+
+
+class QwenImage21MidBlock(nn.Module):
+    def __init__(self, dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            QwenImage21ResidualBlock(dim, dim, dropout),
+            QwenImage21ResidualBlock(dim, dim, dropout),
+        ])
+        self.attentions = nn.ModuleList([QwenImage21AttentionBlock(dim)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.resnets[1](self.attentions[0](self.resnets[0](x)))
+
+
+class QwenImage21ResidualDownBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, dropout: float, num_res_blocks: int,
+                 temporal_downsample: bool, down_flag: bool) -> None:
+        super().__init__()
+        self.avg_shortcut = QwenImage21AvgDown3D(in_dim, out_dim, 2 if temporal_downsample else 1,
+                                               2 if down_flag else 1)
+        self.resnets = nn.ModuleList([
+            QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, dropout)
+            for index in range(num_res_blocks)
+        ])
+        self.downsampler = QwenImage21Resample(
+            out_dim, "downsample3d" if temporal_downsample else "downsample2d"
+        ) if down_flag else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x.clone()
+        for resnet in self.resnets:
+            x = resnet(x)
+        if self.downsampler is not None:
+            x = self.downsampler(x)
+        return x + self.avg_shortcut(residual)
+
+
+class QwenImage21Encoder3d(nn.Module):
+    def __init__(self, arch) -> None:
+        super().__init__()
+        dims = [arch.base_dim * factor for factor in [1, *arch.dim_mult]]
+        self.conv_in = QwenImage21CausalConv3d(arch.in_channels, dims[0], 3, padding=1)
+        self.down_blocks = nn.ModuleList()
+        scale = 1.0
+        for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            down_flag = index != len(arch.dim_mult) - 1
+            temporal = arch.temperal_downsample[index] if down_flag else False
+            if arch.is_residual:
+                self.down_blocks.append(QwenImage21ResidualDownBlock(
+                    in_dim, out_dim, arch.dropout, arch.num_res_blocks, temporal, down_flag
+                ))
+            else:
+                for _ in range(arch.num_res_blocks):
+                    self.down_blocks.append(QwenImage21ResidualBlock(in_dim, out_dim, arch.dropout))
+                    if scale in arch.attn_scales:
+                        self.down_blocks.append(QwenImage21AttentionBlock(out_dim))
+                    in_dim = out_dim
+                if down_flag:
+                    mode = "downsample3d" if temporal else "downsample2d"
+                    self.down_blocks.append(QwenImage21Resample(out_dim, mode))
+                    scale /= 2.0
+        self.mid_block = QwenImage21MidBlock(out_dim, arch.dropout)
+        self.norm_out = QwenImage21RMSNorm(out_dim, images=False)
+        self.conv_out = QwenImage21CausalConv3d(out_dim, arch.z_dim * 2, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_in(x)
+        for block in self.down_blocks:
+            x = block(x)
+        x = self.mid_block(x)
+        return self.conv_out(F.silu(self.norm_out(x)))
+
+
+class QwenImage21ResidualUpBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, arch, temporal_upsample: bool, up_flag: bool) -> None:
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, arch.dropout)
+            for index in range(arch.num_res_blocks + 1)
+        ])
+        self.avg_shortcut = QwenImage21DupUp3D(
+            in_dim, out_dim, 2 if temporal_upsample else 1, 2
+        ) if up_flag else None
+        self.upsampler = QwenImage21Resample(
+            out_dim, "upsample3d" if temporal_upsample else "upsample2d", upsample_out_dim=out_dim
+        ) if up_flag else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x.clone()
+        for resnet in self.resnets:
+            x = resnet(x)
+        if self.upsampler is not None:
+            x = self.upsampler(x)
+        if self.avg_shortcut is not None:
+            x = x + self.avg_shortcut(residual, first_chunk=True)
+        return x
+
+
+class QwenImage21UpBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, arch, upsample_mode: str | None) -> None:
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, arch.dropout)
+            for index in range(arch.num_res_blocks + 1)
+        ])
+        self.upsamplers = nn.ModuleList([QwenImage21Resample(out_dim, upsample_mode)]) if upsample_mode else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for resnet in self.resnets:
+            x = resnet(x)
+        return self.upsamplers[0](x) if self.upsamplers is not None else x
+
+
+class QwenImage21Decoder3d(nn.Module):
+    def __init__(self, arch) -> None:
+        super().__init__()
+        dim = arch.decoder_base_dim if arch.decoder_base_dim is not None else arch.base_dim
+        dims = [dim * factor for factor in [arch.dim_mult[-1], *arch.dim_mult[::-1]]]
+        temporal_upsample = arch.temperal_downsample[::-1]
+        self.conv_in = QwenImage21CausalConv3d(arch.z_dim, dims[0], 3, padding=1)
+        self.mid_block = QwenImage21MidBlock(dims[0], arch.dropout)
+        self.up_blocks = nn.ModuleList()
+        for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            up_flag = index != len(arch.dim_mult) - 1
+            temporal = temporal_upsample[index] if up_flag else False
+            if arch.is_residual:
+                self.up_blocks.append(QwenImage21ResidualUpBlock(in_dim, out_dim, arch, temporal, up_flag))
+            else:
+                if index > 0:
+                    in_dim //= 2
+                mode = ("upsample3d" if temporal else "upsample2d") if up_flag else None
+                self.up_blocks.append(QwenImage21UpBlock(in_dim, out_dim, arch, mode))
+        self.norm_out = QwenImage21RMSNorm(out_dim, images=False)
+        self.conv_out = QwenImage21CausalConv3d(out_dim, arch.out_channels, 3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.mid_block(self.conv_in(x))
+        for block in self.up_blocks:
+            x = block(x)
+        return self.conv_out(F.silu(self.norm_out(x)))
+
+
+def _patchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
+    if patch_size == 1:
+        return x
+    batch, channels, frames, height, width = x.shape
+    if height % patch_size or width % patch_size:
+        raise ValueError("Image height and width must be divisible by the VAE patch size.")
+    x = x.view(batch, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size)
+    x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous()
+    return x.view(batch, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size)
+
+
+def _unpatchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
+    if patch_size == 1:
+        return x
+    batch, channels, frames, height, width = x.shape
+    channels //= patch_size * patch_size
+    x = x.view(batch, channels, patch_size, patch_size, frames, height, width)
+    x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous()
+    return x.view(batch, channels, frames, height * patch_size, width * patch_size)
+
+
+class AutoencoderKLQwenImage21(nn.Module):
+    _supports_gradient_checkpointing = False
+    _group_offload_block_modules = ["quant_conv", "post_quant_conv", "encoder", "decoder"]
+
+    def __init__(self, config: VAEConfig) -> None:
+        super().__init__()
+        self.config = config
+        if not config.load_encoder or not config.load_decoder:
+            raise ValueError("Qwen-Image-2.1 requires the complete encoder and decoder checkpoint.")
+        arch = config.arch_config
+        self.z_dim = arch.z_dim
+        self.latent_channels = arch.z_dim
+        self.encoder = QwenImage21Encoder3d(arch)
+        self.quant_conv = QwenImage21CausalConv3d(arch.z_dim * 2, arch.z_dim * 2, 1)
+        self.post_quant_conv = QwenImage21CausalConv3d(arch.z_dim, arch.z_dim, 1)
+        self.decoder = QwenImage21Decoder3d(arch)
+        self.spatial_compression_ratio = arch.scale_factor_spatial
+        self.temporal_compression_ratio = 1
+        self.use_slicing = False
+        self.use_tiling = config.use_tiling
+        self.tile_sample_min_height = config.tile_sample_min_height
+        self.tile_sample_min_width = config.tile_sample_min_width
+        self.tile_sample_stride_height = config.tile_sample_stride_height
+        self.tile_sample_stride_width = config.tile_sample_stride_width
+        self._validate_tiling()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def enable_slicing(self) -> None:
+        self.use_slicing = True
+
+    def disable_slicing(self) -> None:
+        self.use_slicing = False
+
+    def _validate_tiling(self) -> None:
+        ratio = self.spatial_compression_ratio
+        for extent, stride in ((self.tile_sample_min_height, self.tile_sample_stride_height),
+                               (self.tile_sample_min_width, self.tile_sample_stride_width)):
+            if extent < ratio or stride < ratio or extent % ratio or stride % ratio or stride > extent:
+                raise ValueError(
+                    "VAE tile size and stride must be positive spatial-scale multiples, with stride <= size."
+                )
+
+    def enable_tiling(self, tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None,
+                      tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None,
+                      **kwargs) -> None:
+        self.use_tiling = True
+        self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
+        self._validate_tiling()
+
+    def disable_tiling(self) -> None:
+        self.use_tiling = False
+
+    def clear_cache(self) -> None:
+        """Image inference retains no temporal feature cache."""
+
+    @staticmethod
+    def _validate_image(x: torch.Tensor, channels: int, name: str) -> None:
+        if x.ndim != 5 or x.shape[1] != channels or x.shape[2] != 1:
+            raise ValueError(f"`{name}` must have shape [B, {channels}, 1, H, W], got {tuple(x.shape)}.")
+        if x.shape[0] < 1 or min(x.shape[-2:]) < 1:
+            raise ValueError(f"`{name}` must contain at least one nonempty image.")
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.patch_size is not None:
+            x = _patchify(x, self.config.patch_size)
+        if self.use_tiling and (x.shape[-2] > self.tile_sample_min_height or x.shape[-1] > self.tile_sample_min_width):
+            return self.tiled_encode(x)
+        return self.quant_conv(self.encoder(x))
+
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
+        self._validate_image(x, self.config.in_channels, "x")
+        if x.shape[-2] % self.spatial_compression_ratio or x.shape[-1] % self.spatial_compression_ratio:
+            raise ValueError("Image height and width must be divisible by the VAE spatial compression ratio.")
+        if self.use_slicing and x.shape[0] > 1:
+            moments = torch.cat([self._encode(item) for item in x.split(1)])
+        else:
+            moments = self._encode(x)
+        posterior = DiagonalGaussianDistribution(moments)
+        return AutoencoderKLOutput(posterior) if return_dict else (posterior,)
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        if self.use_tiling and (z.shape[-2] > min_height or z.shape[-1] > min_width):
+            return self.tiled_decode(z).sample
+        decoded = self.decoder(self.post_quant_conv(z))
+        if self.config.patch_size is not None:
+            decoded = _unpatchify(decoded, self.config.patch_size)
+        return decoded.clamp(-1.0, 1.0)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+        self._validate_image(z, self.z_dim, "z")
+        if self.use_slicing and z.shape[0] > 1:
+            decoded = torch.cat([self._decode(item) for item in z.split(1)])
+        else:
+            decoded = self._decode(z)
+        return DecoderOutput(decoded) if return_dict else (decoded,)
+
+    @staticmethod
+    def blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = (
+                a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent)
+                + b[:, :, :, y, :] * (y / blend_extent)
+            )
+        return b
+
+    @staticmethod
+    def blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = (
+                a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent)
+                + b[:, :, :, :, x] * (x / blend_extent)
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        height, width = x.shape[-2:]
+        ratio = self.spatial_compression_ratio // (self.config.patch_size or 1)
+        tile_height = self.tile_sample_min_height // ratio
+        tile_width = self.tile_sample_min_width // ratio
+        stride_height = self.tile_sample_stride_height // ratio
+        stride_width = self.tile_sample_stride_width // ratio
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[:, :, :, i:i + self.tile_sample_min_height, j:j + self.tile_sample_min_width]
+                row.append(self.quant_conv(self.encoder(tile)))
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, tile_height - stride_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, tile_width - stride_width)
+                result_row.append(tile[:, :, :, :stride_height, :stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=3)[:, :, :, :height // ratio, :width // ratio]
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+        height, width = z.shape[-2:]
+        ratio = self.spatial_compression_ratio
+        tile_height = self.tile_sample_min_height // ratio
+        tile_width = self.tile_sample_min_width // ratio
+        stride_height = self.tile_sample_stride_height // ratio
+        stride_width = self.tile_sample_stride_width // ratio
+        sample_height, sample_width = height * ratio, width * ratio
+        sample_stride_height = self.tile_sample_stride_height
+        sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            patch_size = self.config.patch_size
+            sample_height //= patch_size
+            sample_width //= patch_size
+            sample_stride_height //= patch_size
+            sample_stride_width //= patch_size
+            blend_height = self.tile_sample_min_height // patch_size - sample_stride_height
+            blend_width = self.tile_sample_min_width // patch_size - sample_stride_width
+        else:
+            blend_height = self.tile_sample_min_height - sample_stride_height
+            blend_width = self.tile_sample_min_width - sample_stride_width
+        rows = []
+        for i in range(0, height, stride_height):
+            row = []
+            for j in range(0, width, stride_width):
+                tile = z[:, :, :, i:i + tile_height, j:j + tile_width]
+                row.append(self.decoder(self.post_quant_conv(tile)))
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :sample_stride_height, :sample_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        decoded = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+        if self.config.patch_size is not None:
+            decoded = _unpatchify(decoded, self.config.patch_size)
+        decoded = decoded.clamp(-1.0, 1.0)
+        return DecoderOutput(decoded) if return_dict else (decoded,)
+
+    def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        mean = torch.tensor(self.config.latents_mean, device=latents.device, dtype=latents.dtype)
+        std = torch.tensor(self.config.latents_std, device=latents.device, dtype=latents.dtype)
+        mean, std = mean.view(1, self.z_dim, 1, 1, 1), std.view(1, self.z_dim, 1, 1, 1)
+        return (latents - mean) / std
+
+    def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        mean = torch.tensor(self.config.latents_mean, device=latents.device, dtype=latents.dtype)
+        std = torch.tensor(self.config.latents_std, device=latents.device, dtype=latents.dtype)
+        mean, std = mean.view(1, self.z_dim, 1, 1, 1), std.view(1, self.z_dim, 1, 1, 1)
+        return latents * std + mean
+
+    def forward(self, sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True,
+                generator: torch.Generator | None = None) -> DecoderOutput | tuple[torch.Tensor]:
+        posterior = self.encode(sample).latent_dist
+        latents = posterior.sample(generator=generator) if sample_posterior else posterior.mode()
+        return self.decode(latents, return_dict=return_dict)
+
+    @classmethod
+    def from_pretrained(cls, model_path: str | Path, *, subfolder: str | None = None,
+                        config: VAEConfig | None = None, torch_dtype: torch.dtype | None = None,
+                        device: torch.device | str = "cpu", local_files_only: bool = True) -> AutoencoderKLQwenImage21:
+        """Strictly load a local Diffusers-layout VAE component without downloading weights."""
+        from safetensors.torch import load_file
+
+        if not local_files_only:
+            raise ValueError("Qwen-Image-2.1 VAE.from_pretrained accepts local checkpoints only.")
+        path = Path(model_path)
+        if subfolder:
+            path /= subfolder
+        if not path.is_dir():
+            raise FileNotFoundError(f"Local Qwen-Image-2.1 VAE directory does not exist: {path}")
+        with (path / "config.json").open() as handle:
+            arch_fields = json.load(handle)
+        class_name = arch_fields.get("_class_name")
+        if class_name != "AutoencoderKLQwenImage21":
+            raise ValueError(f"Expected AutoencoderKLQwenImage21 checkpoint, got {class_name!r}.")
+        config = config or QwenImage21VAEConfig()
+        config.update_model_arch(arch_fields)
+        model = cls(config)
+        files = sorted(path.glob("*.safetensors"))
+        if not files:
+            raise FileNotFoundError(f"No safetensors VAE checkpoint found in {path}.")
+        state = {}
+        for file in files:
+            shard = load_file(str(file))
+            duplicates = state.keys() & shard.keys()
+            if duplicates:
+                raise ValueError(f"Duplicate tensors in VAE checkpoint shards: {sorted(duplicates)[:5]}")
+            state.update(shard)
+        model.load_state_dict(state, strict=True)
+        logger.info("Loaded Qwen-Image-2.1 VAE strictly from %s", path)
+        return model.to(device=device, dtype=torch_dtype).eval()
+
+
+EntryClass = AutoencoderKLQwenImage21

fastvideo/pipelines/basic/qwen_image21/__init__.py

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

fastvideo/pipelines/basic/qwen_image21/presets.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen-Image-2.1 defaults shared by text and reference image generation."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+QWEN_IMAGE21 = InferencePreset(
+    name="qwen_image21",
+    version=1,
+    model_family="qwen_image21",
+    description="Qwen-Image-2.1 image generation, editing and ordered references",
+    workload_type="t2i",
+    stage_schemas=(PresetStageSpec(name="denoise",
+                                   kind="denoising",
+                                   description="Qwen-Image-2.1 Euler flow steps",
+                                   allowed_overrides=frozenset({"num_inference_steps", "true_cfg_scale"})), ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "num_videos_per_prompt": 1,
+        "fps": 1,
+        "seed": 42,
+        "num_inference_steps": 40,
+        "guidance_scale": 1.0,
+        "true_cfg_scale": 1.0,
+    },
+)
+
+QWEN_IMAGE21_EDIT = InferencePreset(
+    name="qwen_image21_edit",
+    version=1,
+    model_family="qwen_image21",
+    description="Qwen-Image-2.1 image edits and reference-to-image generation",
+    workload_type="i2i",
+    stage_schemas=QWEN_IMAGE21.stage_schemas,
+    defaults=dict(QWEN_IMAGE21.defaults),
+)
+
+ALL_PRESETS = (QWEN_IMAGE21, QWEN_IMAGE21_EDIT)

fastvideo/pipelines/basic/qwen_image21/qwen_image21_pipeline.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native Qwen-Image-2.1 pipeline for generation and image-conditioned edits."""
+
+from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+from .stages import (
+    QwenImage21DecodingStage,
+    QwenImage21DenoisingStage,
+    QwenImage21InputStage,
+    QwenImage21LatentStage,
+    QwenImage21ScheduleStage,
+    QwenImage21TextEncodingStage,
+)
+
+
+class QwenImage21Pipeline(ComposedPipelineBase):
+    pipeline_config_cls = QwenImage21PipelineConfig
+    sampling_params_cls = SamplingParam
+    _required_config_modules = ["processor", "text_encoder", "transformer", "vae", "scheduler"]
+    _lazy_module_names = ("text_encoder", "transformer", "vae")
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage("input_preparation", QwenImage21InputStage())
+        self.add_stage("text_encoding",
+                       QwenImage21TextEncodingStage(self.get_module("text_encoder"), self.get_module("processor")))
+        self.add_stage("latent_preparation", QwenImage21LatentStage(self.get_module("vae")))
+        self.add_stage("timestep_preparation", QwenImage21ScheduleStage(self.get_module("scheduler")))
+        self.add_stage("denoising",
+                       QwenImage21DenoisingStage(self.get_module("transformer"), self.get_module("scheduler")))
+        self.add_stage("decoding", QwenImage21DecodingStage(self.get_module("vae")))
+
+
+EntryClass = QwenImage21Pipeline

fastvideo/pipelines/basic/qwen_image21/stages.py

@@ -0,0 +1,510 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved.
+# Adapted from Diffusers' Apache-2.0 QwenImage21Pipeline.
+"""Separate preparation, encoding, scheduling, denoising and decoding stages."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+from torch.distributed.tensor import DTensor
+from tqdm.auto import tqdm
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.hooks.hooks import ModuleHookManager
+from fastvideo.logger import init_logger
+from fastvideo.models import pinned_offload
+from fastvideo.models.dits.qwen_image21 import QwenImage21KVCache, QwenImage21KVCacheAllocationError
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import V, VerificationResult
+
+logger = init_logger(__name__)
+
+STATE_KEY = "qwen_image21"
+SYSTEM_PROMPT = "Comprehend and analyze the provided prompt."
+IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
+
+
+@dataclass
+class QwenImage21State:
+    references: list[Image.Image] = field(default_factory=list)
+    img_shapes: list[list[tuple[int, int, int]]] = field(default_factory=list)
+    image_pad_mask: torch.Tensor | None = None
+    negative_image_pad_mask: torch.Tensor | None = None
+
+
+def load_reference(value: Any) -> Image.Image:
+    """Detach a local image from its file and preserve alpha for the VAE."""
+    if isinstance(value, str | Path):
+        with Image.open(value) as image:
+            return ImageOps.exif_transpose(image).convert("RGBA")
+    if isinstance(value, Image.Image):
+        return ImageOps.exif_transpose(value).convert("RGBA")
+    if isinstance(value, np.ndarray) and value.dtype == np.uint8 and value.ndim in (2, 3):
+        return Image.fromarray(value).convert("RGBA")
+    raise TypeError("Each reference must be a local image path, PIL image, or uint8 image array")
+
+
+def collect_references(references: list[Any] | None, image_path: Any, pil_image: Any) -> list[Image.Image]:
+    if references is not None and not isinstance(references, list | tuple):
+        raise TypeError("references must be an ordered list of images")
+    shortcuts = [value for value in (image_path, pil_image) if value is not None]
+    if len(shortcuts) > 1 or (references and shortcuts):
+        raise ValueError("Supply references, image_path, or pil_image as a single source of reference images")
+    values = list(references or shortcuts)
+    if len(values) > 10:
+        raise ValueError("Qwen-Image-2.1 supports at most 10 ordered reference images")
+    return [load_reference(value) for value in values]
+
+
+def reference_dimensions(width: int, height: int, resolution: int) -> tuple[int, int]:
+    if min(width, height, resolution) <= 0:
+        raise ValueError("Image dimensions and reference_resolution must be positive")
+    ratio = width / height
+    width_float = math.sqrt(resolution**2 * ratio)
+    height_float = width_float / ratio
+    new_width = round(width_float / 32) * 32
+    new_height = round(height_float / 32) * 32
+    if min(new_width, new_height) < 32:
+        raise ValueError("Reference aspect ratio is too extreme for a 32-pixel latent grid")
+    return new_width, new_height
+
+
+def resize_reference(image: Image.Image, resolution: int) -> Image.Image:
+    return image.resize(reference_dimensions(image.width, image.height, resolution), resample=Image.Resampling.LANCZOS)
+
+
+def vision_reference(image: Image.Image) -> Image.Image:
+    white = Image.new("RGB", image.size, (255, 255, 255))
+    white.paste(image, mask=image.getchannel("A"))
+    return white
+
+
+def reference_pixels(image: Image.Image) -> torch.Tensor:
+    # [B, RGBA, one frame, H, W], with all four channels normalized to [-1, 1].
+    # Build the batch axis before permuting, as upstream's numpy_to_pt does.
+    # Even a singleton batch stride affects BF16 convolution kernel rounding.
+    pixels = torch.from_numpy(np.array(image, dtype=np.float32)[None]).permute(0, 3, 1, 2)
+    return (pixels / 127.5 - 1).unsqueeze(2)
+
+
+def prompt_template(prompt: str, num_references: int) -> str:
+    if not 0 <= num_references <= 10:
+        raise ValueError("Qwen-Image-2.1 supports 0–10 references")
+    prefix = " ".join(f"<image{index + 1}>{IMAGE_PLACEHOLDER}" for index in range(num_references))
+    return (f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+            f"<|im_start|>user\n{prefix}{prompt or ' '}<|im_end|>\n<|im_start|>assistant\n")
+
+
+def extract_prompt(hidden: torch.Tensor, input_ids: torch.Tensor, attention_mask: torch.Tensor, drop_tokens: int,
+                   image_token_id: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    rows, image_rows = [], []
+    for embeddings, ids, mask in zip(hidden, input_ids, attention_mask, strict=True):
+        valid = mask.bool()
+        rows.append(embeddings[valid][drop_tokens:])
+        image_rows.append((ids[valid] == image_token_id)[drop_tokens:])
+    if not rows or min(row.shape[0] for row in rows) == 0:
+        raise ValueError("The processor produced an empty prompt after removing system tokens")
+    length = max(row.shape[0] for row in rows)
+    embeds = torch.stack([torch.cat((row, row.new_zeros(length - row.shape[0], row.shape[1]))) for row in rows])
+    masks = torch.stack([torch.arange(length, device=row.device) < row.shape[0] for row in rows])
+    image_masks = torch.stack([torch.cat((row, row.new_zeros(length - row.shape[0]))) for row in image_rows])
+    return embeds, masks, image_masks
+
+
+def pack_latents(latents: torch.Tensor) -> torch.Tensor:
+    """Pack [B, C, 1, H, W] into [B, H*W, C], with width varying fastest."""
+    if latents.ndim != 5 or latents.shape[2] != 1:
+        raise ValueError("Qwen-Image-2.1 latents must have shape [B, C, 1, H, W]")
+    return latents.flatten(2).transpose(1, 2)
+
+
+def unpack_latents(latents: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """Restore [B, C, 1, H/16, W/16] from tokens containing all C channels."""
+    if latents.ndim != 3 or latents.shape[1] != (height // 16) * (width // 16):
+        raise ValueError("Packed latent length does not match the requested image dimensions")
+    return latents.transpose(1, 2).reshape(latents.shape[0], latents.shape[2], 1, height // 16, width // 16)
+
+
+def normalize_latents(latents: torch.Tensor, means: Any, stds: Any, *, inverse: bool = False) -> torch.Tensor:
+    mean = torch.as_tensor(means, device=latents.device, dtype=latents.dtype).reshape(1, -1, 1, 1, 1)
+    std = torch.as_tensor(stds, device=latents.device, dtype=latents.dtype).reshape(1, -1, 1, 1, 1)
+    if mean.shape[1] != latents.shape[1] or std.shape[1] != latents.shape[1] or bool((std <= 0).any()):
+        raise ValueError("VAE latent statistics must match its channels and have positive standard deviations")
+    return latents * std + mean if inverse else (latents - mean) / std
+
+
+def schedule_shift(target_tokens: int, config: Any) -> float:
+    base_len, max_len = config.get("base_image_seq_len", 256), config.get("max_image_seq_len", 8192)
+    base_shift, max_shift = config.get("base_shift", 0.5), config.get("max_shift", 0.9)
+    slope = (max_shift - base_shift) / (max_len - base_len)
+    return base_shift + slope * (target_tokens - base_len)
+
+
+def validate_request(batch: Any, args: Any) -> None:
+    if args.tp_size != 1 or args.sp_size != 1 or args.num_gpus != 1:
+        raise ValueError("The Qwen-Image-2.1 port currently supports one GPU with tp_size=sp_size=1")
+    if batch.num_frames != 1 or batch.num_videos_per_prompt != 1:
+        raise ValueError("Qwen-Image-2.1 requires num_frames=1 and num_videos_per_prompt=1")
+    if not isinstance(batch.prompt, str) and not (isinstance(batch.prompt, list) and len(batch.prompt) == 1
+                                                  and isinstance(batch.prompt[0], str)):
+        raise ValueError("Qwen-Image-2.1 requires one text prompt per request")
+    for name, value in (("height", batch.height), ("width", batch.width), ("reference_resolution",
+                                                                           args.pipeline_config.reference_resolution)):
+        if not isinstance(value, int) or isinstance(value, bool) or value <= 0 or value % 32:
+            raise ValueError(f"{name} must be a positive multiple of 32")
+    if not isinstance(batch.num_inference_steps, int) or batch.num_inference_steps < 1:
+        raise ValueError("num_inference_steps must be positive")
+    if not math.isfinite(batch.true_cfg_scale) or batch.true_cfg_scale < 1:
+        raise ValueError("true_cfg_scale must be finite and at least 1")
+    if batch.guidance_scale != 1:
+        raise ValueError("Qwen-Image-2.1 has no embedded guidance; use true_cfg_scale with negative_prompt")
+    if batch.video_path is not None or getattr(batch, "conditioning_mask", None) is not None:
+        raise ValueError("Pass annotations or mask images as references; this pipeline takes single-frame images")
+    if batch.prompt_embeds or batch.negative_prompt_embeds:
+        raise ValueError("Externally supplied prompt embeddings are not supported by this Qwen-Image-2.1 port")
+    if batch.negative_prompt is not None and not isinstance(batch.negative_prompt, str):
+        raise ValueError("negative_prompt must be a string or None")
+
+
+def _state(batch: ForwardBatch) -> QwenImage21State:
+    state = batch.extra.get(STATE_KEY)
+    if not isinstance(state, QwenImage21State):
+        raise ValueError("Qwen-Image-2.1 input preparation must precede this stage")
+    return state
+
+
+class QwenImage21InputStage(PipelineStage):
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("height", batch.height, V.positive_int_divisible(32))
+        result.add_check("width", batch.width, V.positive_int_divisible(32))
+        result.add_check("num_frames", batch.num_frames, lambda value: value == 1)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check(STATE_KEY, batch.extra.get(STATE_KEY),
+                                              lambda value: isinstance(value, QwenImage21State))
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        validate_request(batch, fastvideo_args)
+        references = collect_references(batch.references, batch.image_path, batch.pil_image)
+        references = [
+            resize_reference(image, fastvideo_args.pipeline_config.reference_resolution) for image in references
+        ]
+        shapes = [(1, image.height // 16, image.width // 16) for image in references]
+        shapes.append((1, batch.height // 16, batch.width // 16))
+        batch.extra[STATE_KEY] = QwenImage21State(references=references, img_shapes=[shapes])
+        batch.batch_size = 1
+        batch.do_classifier_free_guidance = batch.true_cfg_scale > 1 and batch.negative_prompt is not None
+        if batch.true_cfg_scale > 1 and not batch.do_classifier_free_guidance:
+            logger.warning("true_cfg_scale > 1 requires an explicit negative_prompt; guidance is disabled")
+        return batch
+
+
+class QwenImage21TextEncodingStage(PipelineStage):
+    performance_component_metric = "text_encoder_time_s"
+
+    def __init__(self, text_encoder, processor) -> None:
+        self.text_encoder = text_encoder
+        self.processor = processor
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check(STATE_KEY, batch.extra.get(STATE_KEY),
+                                              lambda value: isinstance(value, QwenImage21State))
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("prompt_embeds", batch.prompt_embeds,
+                                              lambda value: len(value) == 1 and value[0].ndim == 3)
+
+    def _encode(self, prompt: str, state: QwenImage21State, device: torch.device):
+        kwargs = dict(text=[prompt_template(prompt, len(state.references))],
+                      padding=True,
+                      padding_side="left",
+                      return_tensors="pt")
+        if state.references:
+            kwargs["images"] = [vision_reference(image) for image in state.references]
+        model_inputs = self.processor(**kwargs)
+        inputs = {key: value.to(device) for key, value in model_inputs.items() if isinstance(value, torch.Tensor)}
+        hidden = self.text_encoder(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            pixel_values=inputs.get("pixel_values"),
+            image_grid_thw=inputs.get("image_grid_thw"),
+            mm_token_type_ids=inputs.get("mm_token_type_ids"),
+        )
+        system = [{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}]
+        system_ids = self.processor.apply_chat_template(system, tokenize=True, return_dict=False)
+        drop_tokens = len(system_ids[0])
+        image_token_id = self.processor.tokenizer.encode("<|image_pad|>")[0]
+        embeds, mask, image_mask = extract_prompt(hidden, inputs["input_ids"], inputs["attention_mask"], drop_tokens,
+                                                  image_token_id)
+        return embeds, None if bool(mask.all()) else mask, image_mask
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        state = _state(batch)
+        device = get_local_torch_device()
+        first_param = next(self.text_encoder.parameters(), None)
+        # The loader's FSDP2 CPU offload already streams individual text layers.
+        move = (fastvideo_args.text_encoder_cpu_offload and first_param is not None
+                and not isinstance(first_param, DTensor))
+        try:
+            if move:
+                pinned_offload.load(self.text_encoder, device, pin=fastvideo_args.pin_cpu_memory)
+            prompt = batch.prompt[0] if isinstance(batch.prompt, list) else batch.prompt
+            if prompt is None:
+                raise ValueError("A text prompt is required")
+            embeds, mask, state.image_pad_mask = self._encode(prompt, state, device)
+            batch.prompt_embeds = [embeds]
+            batch.prompt_attention_mask = [mask]
+            if batch.do_classifier_free_guidance:
+                if batch.negative_prompt is None:
+                    raise ValueError("True CFG requires an explicit negative prompt")
+                embeds, mask, state.negative_image_pad_mask = self._encode(batch.negative_prompt, state, device)
+                batch.negative_prompt_embeds = [embeds]
+                batch.negative_attention_mask = [mask]
+        finally:
+            if move:
+                pinned_offload.unload(self.text_encoder)
+        return batch
+
+
+class QwenImage21LatentStage(PipelineStage):
+    performance_component_metric = "vae_encode_time_s"
+
+    def __init__(self, vae) -> None:
+        self.vae = vae
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("prompt_embeds", batch.prompt_embeds, lambda value: len(value) == 1)
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        state = _state(batch)
+        device, dtype = get_local_torch_device(), batch.prompt_embeds[0].dtype
+        channels = fastvideo_args.pipeline_config.dit_config.arch_config.in_channels
+        shape = (1, channels, 1, batch.height // 16, batch.width // 16)
+        if batch.latents is None:
+            generator = batch.generator
+            if isinstance(generator, list):
+                if len(generator) != 1:
+                    raise ValueError("Qwen-Image-2.1 requires one generator")
+                generator = generator[0]
+            if generator is None:
+                generator = torch.Generator(device="cpu").manual_seed(batch.seed or 0)
+            noise_device = generator.device
+            batch.latents = pack_latents(
+                torch.randn(shape, generator=generator, device=noise_device, dtype=dtype).to(device))
+        else:
+            if batch.latents.ndim == 5:
+                batch.latents = pack_latents(batch.latents)
+            if tuple(batch.latents.shape) != (1, shape[3] * shape[4], channels):
+                raise ValueError("Supplied latents must match [1, target tokens, latent channels]")
+            batch.latents = batch.latents.to(device=device, dtype=dtype)
+        batch.raw_latent_shape = shape
+        if state.references:
+            try:
+                pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory)
+                if fastvideo_args.pipeline_config.vae_tiling:
+                    self.vae.enable_tiling()
+                else:
+                    self.vae.disable_tiling()
+                vae_dtype = next(self.vae.parameters()).dtype
+                rows = []
+                for image in state.references:
+                    pixels = reference_pixels(image).to(device=device, dtype=vae_dtype)
+                    latents = self.vae.encode(pixels).latent_dist.mode()
+                    latents = normalize_latents(latents, self.vae.config.latents_mean, self.vae.config.latents_std)
+                    rows.append(pack_latents(latents).to(dtype))
+                batch.image_latent = torch.cat(rows, dim=1)
+                expected = int(state.image_pad_mask.sum().item()) * 4
+                if expected != batch.image_latent.shape[1]:
+                    raise ValueError("The processor's image slots do not match the VAE reference grid. "
+                                     "Use the checkpoint's processor and a reference_resolution "
+                                     "within its image budget.")
+            finally:
+                if fastvideo_args.vae_cpu_offload:
+                    pinned_offload.unload(self.vae)
+        # One target mask slot stands for four unpatched latent tokens.
+        slots = state.image_pad_mask.new_ones(1, batch.latents.shape[1] // 4)
+        state.image_pad_mask = torch.cat((state.image_pad_mask, slots), dim=1)
+        if state.negative_image_pad_mask is not None:
+            slots = state.negative_image_pad_mask.new_ones(1, batch.latents.shape[1] // 4)
+            state.negative_image_pad_mask = torch.cat((state.negative_image_pad_mask, slots), dim=1)
+        return batch
+
+
+class QwenImage21ScheduleStage(PipelineStage):
+
+    def __init__(self, scheduler) -> None:
+        self.scheduler = scheduler
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("timesteps", batch.timesteps, V.with_dims(1))
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        sigmas = batch.sigmas.tolist() if isinstance(batch.sigmas, torch.Tensor) else batch.sigmas
+        if sigmas is None:
+            sigmas = np.linspace(1.0, 1 / batch.num_inference_steps, batch.num_inference_steps).tolist()
+        if not sigmas or any(not np.isfinite(sigma) or sigma <= 0 or sigma > 1 for sigma in sigmas):
+            raise ValueError("sigmas must be finite values in (0, 1]")
+        if any(first <= second for first, second in zip(sigmas, sigmas[1:], strict=False)):
+            raise ValueError("sigmas must be strictly decreasing")
+        self.scheduler.set_timesteps(sigmas=sigmas,
+                                     device=get_local_torch_device(),
+                                     mu=schedule_shift(batch.latents.shape[1], self.scheduler.config))
+        self.scheduler.set_begin_index(0)
+        batch.timesteps = self.scheduler.timesteps
+        batch.num_inference_steps = len(batch.timesteps)
+        return batch
+
+
+class QwenImage21DenoisingStage(PipelineStage):
+    performance_component_metric = "dit_time_s"
+
+    def __init__(self, transformer, scheduler) -> None:
+        self.transformer, self.scheduler = transformer, scheduler
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("latents", batch.latents,
+                                              V.with_dims(3)).add_check("timesteps", batch.timesteps, V.with_dims(1))
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        state, device = _state(batch), get_local_torch_device()
+        cache_device = fastvideo_args.pipeline_config.kv_cache_device
+        if cache_device not in ("cpu", "cuda"):
+            raise ValueError("kv_cache_device must be cpu or cuda; use_kv_cache=False disables caching")
+        cache_enabled = fastvideo_args.pipeline_config.use_kv_cache and self.transformer.config.causal_condition
+        caches = []
+        if cache_enabled:
+            caches = [QwenImage21KVCache(len(self.transformer.transformer_blocks), storage_device=cache_device)]
+            if batch.do_classifier_free_guidance:
+                caches.append(QwenImage21KVCache(len(self.transformer.transformer_blocks), storage_device=cache_device))
+        move = (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
+                and not fastvideo_args.use_fsdp_inference)
+        try:
+            if move:
+                pinned_offload.load(self.transformer, device, pin=fastvideo_args.pin_cpu_memory)
+            for index, timestep in enumerate(tqdm(batch.timesteps, desc="Denoising", total=len(batch.timesteps))):
+                if getattr(self, "interrupt", False):
+                    break
+                latent_input = batch.latents
+                if batch.image_latent is not None:
+                    latent_input = torch.cat((batch.image_latent, latent_input), dim=1)
+                mode = "extract" if cache_enabled and index == 0 else ("cached" if cache_enabled else None)
+
+                def predict(embeds: torch.Tensor, mask: torch.Tensor | None, image_mask: torch.Tensor,
+                            cache: QwenImage21KVCache | None, latent_input: torch.Tensor, timestep: torch.Tensor,
+                            mode: str | None) -> torch.Tensor:
+                    result = self.transformer(
+                        hidden_states=latent_input,
+                        timestep=timestep.expand(1).to(batch.latents.dtype) / 1000,
+                        encoder_hidden_states=embeds,
+                        encoder_hidden_states_mask=mask,
+                        img_shapes=state.img_shapes,
+                        img_mask=image_mask,
+                        kv_cache=cache,
+                        kv_cache_mode=mode,
+                        return_dict=False,
+                    )[0]
+                    return result[:, -batch.latents.shape[1]:]
+
+                while True:
+                    try:
+                        positive = predict(batch.prompt_embeds[0], batch.prompt_attention_mask[0], state.image_pad_mask,
+                                           caches[0] if caches else None, latent_input, timestep, mode)
+                        noise = positive
+                        if batch.do_classifier_free_guidance:
+                            negative = predict(batch.negative_prompt_embeds[0], batch.negative_attention_mask[0],
+                                               state.negative_image_pad_mask, caches[1] if caches else None,
+                                               latent_input, timestep, mode)
+                            noise = negative + batch.true_cfg_scale * (positive - negative)
+                        break
+                    except QwenImage21KVCacheAllocationError:
+                        if not cache_enabled:
+                            raise
+                        logger.warning("Host prefix-cache allocation failed; retrying this request without KV caching")
+                        for cache in caches:
+                            cache.clear()
+                        caches = []
+                        cache_enabled, mode = False, None
+                        # An aborted layer may retain the next layer's prefetch.
+                        self._release_layerwise()
+                batch.latents = self.scheduler.step(noise, timestep, batch.latents, return_dict=False)[0]
+                batch.step_index, batch.timestep = index, timestep
+        finally:
+            for cache in caches:
+                cache.clear()
+            if fastvideo_args.dit_layerwise_offload:
+                self._release_layerwise()
+            if move:
+                pinned_offload.unload(self.transformer)
+        return batch
+
+    def _release_layerwise(self) -> None:
+        # The shared offloader links the last block back to the first, leaving
+        # its next-step prefetch resident. Also handle an interrupted block.
+        for block in self.transformer.transformer_blocks:
+            manager = ModuleHookManager.get_from(block)
+            hook = manager.get_forward_hook("LayerwiseOffloadHook") if manager is not None else None
+            if hook is not None and hook.state.gpu_named_parameters:
+                torch.cuda.current_stream().wait_stream(hook.state.async_copy_stream)
+                hook.state.release_gpu_params()
+
+
+class QwenImage21DecodingStage(PipelineStage):
+    performance_component_metric = "vae_decode_time_s"
+
+    def __init__(self, vae) -> None:
+        self.vae = vae
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        return VerificationResult().add_check("output", batch.output, V.is_tensor)
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if fastvideo_args.output_type == "latent":
+            batch.output = batch.latents
+            return batch
+        device = get_local_torch_device()
+        try:
+            pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory)
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
+            dtype = next(self.vae.parameters()).dtype
+            latents = unpack_latents(batch.latents, batch.height, batch.width).to(device=device, dtype=dtype)
+            latents = normalize_latents(latents,
+                                        self.vae.config.latents_mean,
+                                        self.vae.config.latents_std,
+                                        inverse=True)
+            # Keep alpha as the fourth channel all the way to the PNG exporter.
+            batch.output = self.vae.decode(latents, return_dict=False)[0].div(2).add(0.5).clamp(0, 1).float()
+        finally:
+            if fastvideo_args.vae_cpu_offload:
+                pinned_offload.unload(self.vae)
+        return batch

fastvideo/pipelines/stages/input_validation.py

@@ -64,8 +64,7 @@ class InputValidationStage(PipelineStage):
             raise ValueError("Either `prompt` or `prompt_embeds` must be provided")
 
         # Ensure negative prompt is properly formatted if using classifier-free guidance
-        if (batch.do_classifier_free_guidance and batch.negative_prompt is None
-                and batch.negative_prompt_embeds is None):
+        if (batch.do_classifier_free_guidance and batch.negative_prompt is None and not batch.negative_prompt_embeds):
             raise ValueError("For classifier-free guidance, either `negative_prompt` or "
                              "`negative_prompt_embeds` must be provided")
 

fastvideo/registry.py

@@ -55,6 +55,7 @@ from fastvideo.models.wan.definition import WAN_MODEL_DEFINITIONS, WanModelDefin
 from fastvideo.configs.pipelines.flux import FluxPipelineConfig
 from fastvideo.configs.pipelines.sd35 import SD35Config
 from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
+from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
 from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
@@ -222,14 +223,21 @@ def _get_config_info(
     revision: str | None = None,
     allow_remote: bool = True,
 ) -> ConfigInfo | None:
+    # Normalize only the lookup key; keep the real path for manifest IO and detectors.
+    lookup_path = model_path.rstrip("/") or model_path
+    parts = lookup_path.split("/")
+    # Standard Hugging Face cache layout: models--{org}--{name}/snapshots/{revision}.
+    if len(parts) >= 3 and parts[-2] == "snapshots" and parts[-3].startswith("models--"):
+        lookup_path = parts[-3][len("models--"):].replace("--", "/", 1)
+
     # 1. Exact match
-    if model_path in _MODEL_HF_PATH_TO_NAME:
-        model_id = _MODEL_HF_PATH_TO_NAME[model_path]
+    if lookup_path in _MODEL_HF_PATH_TO_NAME:
+        model_id = _MODEL_HF_PATH_TO_NAME[lookup_path]
         logger.debug("Resolved model path '%s' from exact path match.", model_path)
         return _CONFIG_REGISTRY.get(model_id)
 
     # 2. Partial match: use short model name.
-    model_name = get_model_short_name(model_path.lower())
+    model_name = get_model_short_name(lookup_path.lower())
     all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)
     for registered_model_hf_id in all_model_hf_paths:
         registered_model_name = get_model_short_name(registered_model_hf_id.lower())
@@ -1204,6 +1212,19 @@ def _register_configs() -> None:
         ],
     )
 
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=QwenImage21PipelineConfig,
+        workload_types=(WorkloadType.T2I, WorkloadType.I2I),
+        hf_model_paths=["Qwen/Qwen-Image-2.1"],
+        model_detectors=[
+            lambda path: "qwenimage21pipeline" in path or "qwen-image-2.1" in path or "qwen_image21" in path,
+        ],
+        model_family="qwen_image21",
+        default_preset="qwen_image21",
+        pipeline_cls_name="QwenImage21Pipeline",
+    )
+
 
 # --- Part 3: Main Resolver ---
 
@@ -1337,6 +1358,8 @@ def _register_presets() -> None:
         ALL_PRESETS as STABLE_AUDIO_PRESETS, )
     from fastvideo.pipelines.basic.wan.presets import (
         ALL_PRESETS as WAN_PRESETS, )
+    from fastvideo.pipelines.basic.qwen_image21.presets import (
+        ALL_PRESETS as QWEN_IMAGE21_PRESETS, )
     from fastvideo.pipelines.basic.flux_2.presets import (
         ALL_PRESETS as FLUX2_PRESETS, )
 
@@ -1364,6 +1387,7 @@ def _register_presets() -> None:
         SD35_PRESETS,
         STABLE_AUDIO_PRESETS,
         WAN_PRESETS,
+        QWEN_IMAGE21_PRESETS,
     )
     for group in all_preset_groups:
         for preset in group:

fastvideo/tests/api/test_wan_definitions.py

@@ -80,6 +80,40 @@ def test_legacy_aliases_and_sampling_defaults(model_path, config_name, preset_na
             assert getattr(sampling, name) == value
 
 
+@pytest.mark.parametrize("hf_id", [
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
+])
+@pytest.mark.parametrize("path_form", ["renamed_parent", "trailing_slash", "snapshot", "snapshot_slash"])
+def test_local_path_forms_resolve_like_hub_id(tmp_path, hf_id, path_form):
+    expected = registry._get_config_info(hf_id)
+    org, name = hf_id.split("/", 1)
+    if path_form.startswith("snapshot"):
+        checkpoint = tmp_path / "hub" / f"models--{org}--{name}" / "snapshots" / ("0" * 40)
+    else:
+        checkpoint = tmp_path / "renamed-parent" / name
+    checkpoint.mkdir(parents=True)
+    (checkpoint / "model_index.json").write_text(json.dumps({
+        "_class_name": "WanPipeline", "_diffusers_version": "0.33.0",
+    }))
+    path = str(checkpoint) + ("/" if path_form.endswith("slash") else "")
+    info = registry._get_config_info(path)
+    assert info.pipeline_config_cls is expected.pipeline_config_cls
+    assert info.default_preset == expected.default_preset
+
+
+@pytest.mark.parametrize("suffix", ["", "/"])
+def test_unknown_snapshot_falls_back_to_real_manifest(tmp_path, suffix):
+    checkpoint = tmp_path / "models--x--y" / "snapshots" / ("0" * 40)
+    checkpoint.mkdir(parents=True)
+    (checkpoint / "model_index.json").write_text(json.dumps({
+        "_class_name": "WanImageToVideoPipeline", "_diffusers_version": "0.33.0",
+    }))
+    info = registry._get_config_info(str(checkpoint) + suffix)
+    assert info.pipeline_config_cls is pipeline_config.WanI2V480PConfig
+    assert info.default_preset == "wan_i2v_14b_480p"
+
+
 @pytest.mark.parametrize("directory,manifest_class,expected_family,expected_preset", [
     ("custom", "WanPipeline", "wan", "wan_t2v_1_3b"),
     ("custom", "WanImageToVideoPipeline", "wan", "wan_i2v_14b_480p"),

fastvideo/tests/contract/test_env_policy.py

@@ -190,7 +190,6 @@ KNOWN_VIOLATIONS: dict[str, int] = {
     'fastvideo/tests/ssim/ci_runner.py: whole-environ': 1,
     'fastvideo/tests/ssim/reference_videos_cli.py: read <dynamic>': 1,
     'fastvideo/tests/train/methods/test_minimax_h3_finetune.py: whole-environ': 2,
-    'fastvideo/tests/train/models/test_load_kandinsky5.py: write FASTVIDEO_ATTENTION_BACKEND': 1,
     'fastvideo/tests/training/distill/test_anyflow_smoke.py: whole-environ': 2,
     'fastvideo/train/entrypoint/train.py: write FASTVIDEO_ATTENTION_BACKEND': 1,
     'fastvideo/utils.py: read <dynamic>': 5,

fastvideo/tests/entrypoints/test_openai_serving_engine.py

@@ -4,6 +4,8 @@
 from __future__ import annotations
 
 import asyncio
+import base64
+import os
 from pathlib import Path
 import threading
 import time
@@ -15,11 +17,13 @@ from fastapi.testclient import TestClient
 
 from fastvideo.api.parser import parse_config
 from fastvideo.api.schema import GenerationRequest
+from fastvideo.entrypoints.openai import media_policy, request_adapter
 from fastvideo.entrypoints.openai.media_policy import MediaPolicy
-from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest
+from fastvideo.entrypoints.openai.protocol import UrlAudioReference, VideoGenerationRequest
 from fastvideo.entrypoints.openai.request_adapter import (
     RequestAdaptationError,
     build_generation_request,
+    discard_materialized_media,
     prepare_reference_media,
     validate_served_model_name,
 )
@@ -241,7 +245,8 @@ def test_wav_media_type_is_detected_from_output_path() -> None:
 def test_missing_video_reference_fails_during_admission(tmp_path: Path) -> None:
     request = VideoGenerationRequest(prompt="waves", video_reference={"video_url": str(tmp_path / "missing.mp4")})
     with pytest.raises(RequestAdaptationError, match="does not exist"):
-        asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy(roots=(str(tmp_path), ))))
+        asyncio.run(
+            prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy(roots=(str(tmp_path), ))))
 
 
 def test_local_media_paths_need_a_configured_root(tmp_path: Path) -> None:
@@ -252,6 +257,166 @@ def test_local_media_paths_need_a_configured_root(tmp_path: Path) -> None:
         asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
 
 
+def _data_url(media_type: str, payload: bytes) -> str:
+    return f"data:{media_type};base64,{base64.b64encode(payload).decode()}"
+
+
+def _ref(kind: str, url: str) -> dict[str, str]:
+    return {f"{kind}_url": url}
+
+
+def _uploads(tmp_path: Path) -> list[Path]:
+    return list((tmp_path / "uploads").glob("*"))
+
+
+def test_inline_video_and_audio_references_are_decoded_in_order(tmp_path: Path) -> None:
+    request = VideoGenerationRequest(
+        prompt="waves",
+        video_reference=[
+            _ref("video", _data_url("video/mp4", b"first-clip")),
+            _ref("video", _data_url("video/quicktime", b"second-clip")),
+        ],
+        audio_reference=[
+            _ref("audio", _data_url("audio/mpeg", b"first-track")),
+            _ref("audio", _data_url("audio/wav", b"second-track")),
+        ],
+    )
+    created = asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
+    assert isinstance(request.video_reference, list) and isinstance(request.audio_reference, list)
+    videos = [reference.video_url for reference in request.video_reference]
+    audios = [reference.audio_url for reference in request.audio_reference]
+    assert created == videos + audios
+    contents = [Path(path).read_bytes() for path in created]
+    assert contents == [b"first-clip", b"second-clip", b"first-track", b"second-track"]
+    assert [Path(path).suffix for path in created] == [".mp4", ".mov", ".mp3", ".wav"]
+    assert {Path(path).parent for path in created} == {tmp_path / "uploads"}
+    discard_materialized_media(created)
+    assert not _uploads(tmp_path)
+
+
+@pytest.mark.parametrize("chunk_chars", [media_policy.DATA_URL_CHUNK_CHARS, 3, 4, 5, 7])
+def test_inline_reference_payload_may_be_wrapped_and_percent_escaped(tmp_path: Path, monkeypatch: pytest.MonkeyPatch,
+                                                                    chunk_chars: int) -> None:
+    # Small chunks put chunk boundaries inside percent escapes and line breaks.
+    monkeypatch.setattr(media_policy, "DATA_URL_CHUNK_CHARS", chunk_chars)
+    track = bytes(range(256)) * 2
+    encoded = base64.b64encode(track).decode()
+    lines = [encoded[start:start + 16] for start in range(0, len(encoded), 16)]
+    wrapped = "".join(line + ("%0A" if number % 2 else "\r\n") for number, line in enumerate(lines))
+    escaped = wrapped.replace("+", "%2B").replace("/", "%2f").replace("=", "%3D")
+    request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": f"data:audio/flac;base64,{escaped}"})
+    created = asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
+    assert Path(created[0]).read_bytes() == track
+    assert created[0].endswith(".flac")
+
+
+@pytest.mark.parametrize(
+    ("kind", "url", "message"),
+    [
+        ("video", "data:video/mp4;base64,not*base64", "not valid base64"),
+        ("audio", "data:audio/mpeg;base64,", "payload is empty"),
+        ("video", "data:video/mp4,raw-bytes", "must be base64-encoded"),
+        ("video", _data_url("audio/mpeg", b"track"), "need a data:video/"),
+        ("audio", _data_url("video/mp4", b"clip"), "need a data:audio/"),
+        ("audio", _data_url("application/octet-stream", b"track"), "need a data:audio/"),
+        ("image", _data_url("video/mp4", b"clip"), "need a data:image/"),
+    ],
+)
+def test_invalid_inline_references_are_refused(tmp_path: Path, kind: str, url: str, message: str) -> None:
+    request = VideoGenerationRequest(prompt="waves", **{f"{kind}_reference": _ref(kind, url)})
+    with pytest.raises(RequestAdaptationError, match=message):
+        asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
+    assert not _uploads(tmp_path)
+
+
+def test_inline_reference_over_the_media_limit_is_refused(tmp_path: Path) -> None:
+    request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": _data_url("audio/wav", b"12345")})
+    with pytest.raises(RequestAdaptationError, match="4-byte limit"):
+        asyncio.run(prepare_reference_media("request", request, str(tmp_path), max_bytes=4, policy=MediaPolicy()))
+    assert not _uploads(tmp_path)
+    at_limit = VideoGenerationRequest(prompt="waves", video_reference={"video_url": _data_url("video/mp4", b"1234")})
+    created = asyncio.run(prepare_reference_media("request", at_limit, str(tmp_path), max_bytes=4,
+                                                  policy=MediaPolicy()))
+    assert Path(created[0]).read_bytes() == b"1234"
+
+
+def test_oversized_inline_payload_is_refused_before_it_is_all_read(tmp_path: Path,
+                                                                   monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(media_policy, "DATA_URL_CHUNK_CHARS", 8)
+    unescaped: list[int] = []
+    original_unquote = media_policy.unquote
+
+    def counting_unquote(chunk: str) -> str:
+        unescaped.append(len(chunk))
+        return original_unquote(chunk)
+
+    monkeypatch.setattr(media_policy, "unquote", counting_unquote)
+    source = "data:video/mp4;base64," + "QU%4AB\n" * 10_000
+    with pytest.raises(media_policy.MediaPolicyError, match="4-byte limit"):
+        media_policy.decode_data_url(source, str(tmp_path / "clip.mp4"), max_bytes=4)
+    assert max(unescaped) <= 8 and sum(unescaped) < 64
+    assert not (tmp_path / "clip.mp4").exists()
+
+
+def test_decoded_references_are_removed_when_a_later_reference_fails(tmp_path: Path) -> None:
+    request = VideoGenerationRequest(
+        prompt="waves",
+        video_reference={"video_url": _data_url("video/mp4", b"clip")},
+        audio_reference={"audio_url": "data:audio/mpeg;base64,%%%"},
+    )
+    with pytest.raises(RequestAdaptationError, match="not valid base64"):
+        asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
+    assert not _uploads(tmp_path)
+
+
+def test_cancelled_decoding_removes_its_file_once_the_thread_finishes(tmp_path: Path,
+                                                                      monkeypatch: pytest.MonkeyPatch) -> None:
+    asyncio.run(_assert_cancelled_decoding_removes_its_file(tmp_path, monkeypatch))
+
+
+async def _assert_cancelled_decoding_removes_its_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    started, release, finished = threading.Event(), threading.Event(), threading.Event()
+    written: list[bool] = []
+
+    def decode_after_release(source: str, target: str, *, max_bytes: int) -> None:
+        started.set()
+        release.wait(timeout=5)
+        try:
+            media_policy.decode_data_url(source, target, max_bytes=max_bytes)
+            written.append(os.path.exists(target))
+        finally:
+            finished.set()
+
+    monkeypatch.setattr(request_adapter, "decode_data_url", decode_after_release)
+    request = VideoGenerationRequest(prompt="waves", video_reference={"video_url": _data_url("video/mp4", b"clip")})
+    preparing = asyncio.create_task(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
+    assert await asyncio.to_thread(started.wait, 5)
+    preparing.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await preparing
+    # The thread writes its file only after the request is gone, so cleanup has to wait for the thread.
+    release.set()
+    assert await asyncio.to_thread(finished.wait, 5)
+    assert written == [True]
+    deadline = time.monotonic() + 5
+    while _uploads(tmp_path) and time.monotonic() < deadline:
+        await asyncio.sleep(0.01)
+    assert not _uploads(tmp_path)
+
+
+def test_local_audio_paths_need_a_configured_root(tmp_path: Path) -> None:
+    track = tmp_path / "track.wav"
+    track.write_bytes(b"x")
+    request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": str(track)})
+    with pytest.raises(RequestAdaptationError, match="server.media_roots"):
+        asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy()))
+    created = asyncio.run(
+        prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy(roots=(str(tmp_path), ))))
+    assert created == []
+    assert isinstance(request.audio_reference, UrlAudioReference)
+    assert request.audio_reference.audio_url == os.path.realpath(track)
+
+
 def test_mmaudio_v2a_requires_video_reference(tmp_path: Path) -> None:
     request = VideoGenerationRequest(prompt="waves")
     with pytest.raises(RequestAdaptationError, match="requires video_reference"):
@@ -267,13 +432,18 @@ def test_mmaudio_v2a_requires_video_reference(tmp_path: Path) -> None:
 def test_missing_reference_fails_during_admission(tmp_path: Path) -> None:
     request = VideoGenerationRequest(prompt="a fox", input_reference=str(tmp_path / "missing.png"))
     with pytest.raises(RequestAdaptationError, match="does not exist"):
-        asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy(roots=(str(tmp_path), ))))
+        asyncio.run(
+            prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy(roots=(str(tmp_path), ))))
 
 
 def test_request_adapter_resolves_vllm_nested_params(tmp_path: Path) -> None:
     request = VideoGenerationRequest(
         prompt="a fox",
-        video_params={"width": 832, "height": 480, "fps": 30},
+        video_params={
+            "width": 832,
+            "height": 480,
+            "fps": 30
+        },
         seconds="2",
         seed=7,
     )
@@ -304,7 +474,11 @@ def test_request_adapter_accepts_matching_startup_lora(tmp_path: Path) -> None:
     request = VideoGenerationRequest(
         prompt="a fox",
         model="fast",
-        lora={"name": "fast", "path": adapter, "scale": 0.75},
+        lora={
+            "name": "fast",
+            "path": adapter,
+            "scale": 0.75
+        },
     )
 
     build_generation_request(
@@ -359,10 +533,22 @@ def test_fasth3_request_uses_the_general_adapter(tmp_path: Path) -> None:
 @pytest.mark.parametrize(
     ("overrides", "message"),
     [
-        ({"aspect_ratio": "16:9", "num_frames": 22}, "generates 5-15 seconds"),
-        ({"size": "100x100", "num_frames": 124}, "positive multiples of 32"),
-        ({"aspect_ratio": "1:9", "num_frames": 124}, "supports aspect ratios"),
-        ({"aspect_ratio": "16:9", "num_frames": 120}, "causal-VAE grid"),
+        ({
+            "aspect_ratio": "16:9",
+            "num_frames": 22
+        }, "generates 5-15 seconds"),
+        ({
+            "size": "100x100",
+            "num_frames": 124
+        }, "positive multiples of 32"),
+        ({
+            "aspect_ratio": "1:9",
+            "num_frames": 124
+        }, "supports aspect ratios"),
+        ({
+            "aspect_ratio": "16:9",
+            "num_frames": 120
+        }, "causal-VAE grid"),
     ],
 )
 def test_fasth3_invalid_geometry_fails_at_admission(tmp_path: Path, overrides, message) -> None:
@@ -396,7 +582,9 @@ def test_fasth3_ref2va_limits_references_at_admission(tmp_path: Path) -> None:
         task="ref2va",
         aspect_ratio="16:9",
         num_frames=124,
-        image_reference=[{"image_url": f"/tmp/reference-{index}.png"} for index in range(10)],
+        image_reference=[{
+            "image_url": f"/tmp/reference-{index}.png"
+        } for index in range(10)],
     )
     with pytest.raises(RequestAdaptationError, match="at most 9 image references"):
         build_generation_request(
@@ -489,7 +677,11 @@ def test_audio_jobs_publish_wav_media_type_and_seconds(tmp_path: Path) -> None:
         with TestClient(app) as client:
             created = client.post(
                 "/v1/videos",
-                json={"model": "stable-audio", "prompt": "rain", "seconds": 3},
+                json={
+                    "model": "stable-audio",
+                    "prompt": "rain",
+                    "seconds": 3
+                },
             )
             assert created.status_code == 200
             audio_id = created.json()["id"]
@@ -522,6 +714,69 @@ def test_audio_jobs_publish_wav_media_type_and_seconds(tmp_path: Path) -> None:
         state.clear_state()
 
 
+class _ReferenceGenerator(_FileGenerator):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.references: list[tuple[str, bytes]] = []
+
+    def generate(self, request):
+        self.references = [(reference.media_type, Path(reference.source).read_bytes())
+                           for reference in request.inputs.references]
+        return super().generate(request)
+
+
+def test_ref2va_route_accepts_inline_references_and_removes_them(tmp_path: Path) -> None:
+    from fastvideo.entrypoints.openai import state
+    from fastvideo.entrypoints.openai.video_api import router
+
+    generator = _ReferenceGenerator()
+    engine = OpenAIServingEngine(generator)  # type: ignore[arg-type]
+    args = _args("MiniMaxAI/MiniMax-H3", override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline")
+    state.set_state(
+        generator,  # type: ignore[arg-type]
+        engine,
+        args,  # type: ignore[arg-type]
+        str(tmp_path),
+        served_model_name="minimax-h3",
+    )
+    asyncio.run(VIDEO_STORE.clear())
+    app = FastAPI()
+    app.include_router(router)
+    body = dict(
+        model="minimax-h3",
+        prompt="a fox sings along",
+        task="ref2va",
+        aspect_ratio="16:9",
+        num_frames=124,
+        video_reference=[_ref("video", _data_url("video/mp4", b"clip-with-sound"))],
+        audio_reference=[
+            _ref("audio", _data_url("audio/mpeg", b"first-voice")),
+            _ref("audio", _data_url("audio/wav", b"second-voice")),
+        ],
+    )
+
+    try:
+        with TestClient(app) as client:
+            sync = client.post("/v1/videos/sync", json=body)
+            assert sync.status_code == 200, sync.text
+            assert generator.references == [
+                ("video", b"clip-with-sound"),
+                ("audio", b"first-voice"),
+                ("audio", b"second-voice"),
+            ]
+            assert not _uploads(tmp_path)
+
+            mismatched = {**body, "audio_reference": [_ref("audio", _data_url("video/mp4", b"clip"))]}
+            refused = client.post("/v1/videos", json=mismatched)
+            assert refused.status_code == 400
+            assert "need a data:audio/" in refused.text
+            assert not _uploads(tmp_path)
+    finally:
+        asyncio.run(VIDEO_STORE.clear())
+        state.clear_state()
+
+
 def test_video_routes_cover_async_sync_list_content_and_delete(tmp_path: Path) -> None:
     from fastvideo.entrypoints.openai import state
     from fastvideo.entrypoints.openai.video_api import router

fastvideo/tests/golden_gate/test_qwen_image21.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Golden-gate: Qwen-Image-2.1 transformer block 0 with real BF16 weights.
+
+The small packed sequence contains text, one reference image, more text, and
+the target image. Causal text segments use masked SDPA; reference and target
+segments use FlashAttention, including their unequal query/key lengths. Complex RoPE and
+the two modulation rows match the causal-conditioning block interface.
+
+Requires a BF16-capable CUDA GPU and FlashAttention. Select FA4 with
+FASTVIDEO_FA4=1 before starting Python. The shared harness downloads
+only checkpoint shards containing block 0 and compares a device/runtime-matched
+golden exactly; it never loads the text encoder or VAE.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate
+
+__all__ = ["distributed_runtime"]
+
+DIM = 4096  # 32 heads x 128, Qwen-Image-2.1 transformer/config.json
+TEXT_BEFORE, TEXT_AFTER = 32, 16
+REFERENCE_H, REFERENCE_W = 4, 4
+TARGET_H, TARGET_W = 8, 16
+CHECKPOINT_REVISION = "d26bb61231c349cf6b7896fa83353113880e1ba3"
+
+
+def _build_block(layer: int) -> torch.nn.Module:
+    from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig
+    from fastvideo.models.dits.qwen_image21 import QwenImage21TransformerBlock
+    from fastvideo.platforms import AttentionBackendEnum
+
+    arch = QwenImage21ArchConfig()
+    # The selector reads the construction dtype before the harness casts the
+    # loaded block. Keep it BF16 here so FLASH_ATTN cannot fall back to SDPA.
+    previous_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.bfloat16)
+    try:
+        block = QwenImage21TransformerBlock(
+            dim=arch.num_attention_heads * arch.attention_head_dim,
+            num_attention_heads=arch.num_attention_heads,
+            attention_head_dim=arch.attention_head_dim,
+            mlp_ratio=arch.mlp_ratio,
+            eps=arch.eps,
+            requested_backend=AttentionBackendEnum.FLASH_ATTN,
+            supported_attention_backends=arch._supported_attention_backends,
+        )
+    finally:
+        torch.set_default_dtype(previous_dtype)
+    assert block.attn.dense_attn is not None, "Qwen golden requires dense FlashAttention"
+    assert block.attn.dense_attn.backend is AttentionBackendEnum.FLASH_ATTN
+    return block
+
+
+def _make_inputs(device: torch.device, seed: int) -> dict:
+    from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig
+    from fastvideo.models.dits.qwen_image21 import (
+        QwenImage21Rope,
+        QwenImage21Transformer2DModel,
+        _qwenimage21_prefix_segments,
+    )
+
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    reference_tokens = REFERENCE_H * REFERENCE_W
+    target_tokens = TARGET_H * TARGET_W
+    image_pad_mask = torch.tensor(
+        [False] * TEXT_BEFORE + [True] * reference_tokens + [False] * TEXT_AFTER + [True] * target_tokens,
+        dtype=torch.bool,
+        device=device,
+    )
+    img_shapes = [(1, REFERENCE_H, REFERENCE_W), (1, TARGET_H, TARGET_W)]
+    image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes)
+    rope = QwenImage21Rope(theta=10000, axes_dim=QwenImage21ArchConfig().axes_dims_rope)
+
+    def rnd(*shape):
+        return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device=device, dtype=torch.bfloat16)
+
+    return {
+        "hidden_states": rnd(1, image_pad_mask.numel(), DIM),
+        # Row 0 conditions target tokens; row 1 is the prefix's zero-timestep row.
+        "modulation": rnd(2, 4 * DIM),
+        "target_token_mask": target_token_mask,
+        "rotary_emb": rope(img_shapes, image_pad_mask, device),
+        "segments": _qwenimage21_prefix_segments(image_ids, image_pad_mask.numel() - target_tokens),
+    }
+
+
+SPEC = GateSpec(
+    name="qwen_image21",
+    repo_id="Qwen/Qwen-Image-2.1",
+    revision=CHECKPOINT_REVISION,
+    build_block=_build_block,
+    make_inputs=_make_inputs,
+    prefix_template="transformer_blocks.{N}.",
+    attention_backend="FLASH_ATTN",
+)
+
+
+def test_qwen_image21_golden_gate(distributed_runtime) -> None:
+    run_gate(SPEC)

fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py

@@ -211,11 +211,11 @@ def test_quant_gemm_env_selector_and_validation(env_overrides, quant_gemm_mode_r
         set_quant_gemm_mode("bogus")
 
 
-def test_unset_pipeline_quant_gemm_keeps_the_environment(monkeypatch) -> None:
+def test_unset_pipeline_quant_gemm_keeps_the_environment(env_overrides) -> None:
     from fastvideo.mlx_runtime import fastwan
     from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline
 
-    monkeypatch.setenv("FASTVIDEO_MLX_QUANT_GEMM", "qmm")
+    env_overrides.enter_context(envs.FASTVIDEO_MLX_QUANT_GEMM.override("qmm"))
     pipeline = object.__new__(MiniMaxH3MLXPipeline)
     pipeline._apply_quant_gemm(None)
     try:

fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py

@@ -80,8 +80,10 @@ def _cache_for(dit, layout, storage="native"):
     return MiniMaxH3PrefixKVCache(len(dit.blocks), frozen_prefix_rows(layout), dit.num_heads, dit.head_dim, storage)
 
 
-def _full(dit, layout, inputs, record=None, video_t=0.7, audio_t=0.4, step_index=0):
+def _full(dit, layout, inputs, record=None, video_t=0.7, audio_t=0.4, step_index=0, text_embeds=None):
     kwargs = {"record_prefix_kv": record} if record is not None else {}
+    if text_embeds is not None:
+        kwargs["precomputed_text_embeds"] = text_embeds
     video, audio = dit.forward_with_cache(inputs["x_v"], inputs["x_a"], inputs["text"], layout=layout,
                                           step_index=step_index, **_step_args(layout, video_t, audio_t), **kwargs)
     mx.eval(video, audio)
@@ -112,6 +114,24 @@ def test_recording_step_matches_the_normal_forward_exactly(tiny_dit) -> None:
     assert cache.nbytes == expected
 
 
+def test_precomputed_text_preserves_recorded_prefix_kv(tiny_dit) -> None:
+    """Cached refinement must preserve both visible outputs and later-step K/V."""
+    dit, layout = tiny_dit
+    inputs = _inputs(layout, seed=13)
+    expected_cache = _cache_for(dit, layout)
+    expected = _full(dit, layout, inputs, record=expected_cache)
+    text_embeds = dit.refine_text(inputs["text"])
+    mx.eval(text_embeds)
+    actual_cache = _cache_for(dit, layout)
+    actual = _full(dit, layout, inputs, record=actual_cache, text_embeds=text_embeds)
+    for left, right in zip(expected, actual, strict=True):
+        np.testing.assert_array_equal(left, right)
+    assert actual_cache.complete
+    for layer in range(len(dit.blocks)):
+        for left, right in zip(expected_cache.load(layer), actual_cache.load(layer), strict=True):
+            np.testing.assert_array_equal(np.asarray(left), np.asarray(right))
+
+
 def test_reuse_with_unmoved_inputs_equals_the_full_forward(tiny_dit) -> None:
     """If nothing the prefix depends on changed, frozen K/V is exactly what the full pass computes."""
     dit, layout = tiny_dit
@@ -290,13 +310,20 @@ class _StubDiT:
 
     def __init__(self) -> None:
         self.calls: list[str] = []
+        self.refiner_inputs: list[np.ndarray] = []
+        self.text_embeddings: list[object] = []
 
     def precompute_adaln(self, timesteps, *, drop_weights: bool = True):
         return None
 
+    def refine_text(self, text):
+        self.refiner_inputs.append(np.asarray(text).copy())
+        return text + 1.0
+
     def forward_with_cache(self, x_v, x_a, text, *, layout, step_timesteps, row_timestep_inverse, step_index=0,
-                           record_prefix_kv=None):
+                           record_prefix_kv=None, precomputed_text_embeds=None):
         self.calls.append("record" if record_prefix_kv is not None else "full")
+        self.text_embeddings.append(precomputed_text_embeds)
         return x_v * 0.0 + 1.0, x_a * 0.0 + 1.0
 
     def forward_target_rows(self, x_v, x_a, *, layout, step_timesteps, row_timestep_inverse, prefix_kv, step_index=0):
@@ -336,10 +363,18 @@ def test_denoise_loop_records_then_reuses_and_refreshes(tmp_path: Path) -> None:
     stats = pipeline.last_prefix_kv_stats
     assert stats["enabled"] and stats["recorded_steps"] == [0] and stats["reused_steps"] == [1, 2, 3]
     assert len(pipeline.last_denoise_step_times) == 4
+    assert len(dit.refiner_inputs) == len(dit.text_embeddings) == 1
+    np.testing.assert_array_equal(np.asarray(dit.text_embeddings[0]), dit.refiner_inputs[0] + 1.0)
 
     dit = _StubDiT()
     pipeline = _denoise(tmp_path, dit, num_steps=4, reuse_prefix_kv=True, prefix_kv_refresh_every=2)
     assert dit.calls == ["record", "reuse", "record", "reuse"]
+    assert pipeline.last_prefix_kv_stats["recorded_steps"] == [0, 2]
+    assert len(dit.refiner_inputs) == 1
+    assert len(dit.text_embeddings) == 2
+    # Refreshing transformer K/V does not invalidate timestep-independent text.
+    for embeds in dit.text_embeddings:
+        np.testing.assert_array_equal(np.asarray(embeds), dit.refiner_inputs[0] + 1.0)
 
 
 def test_denoise_default_never_touches_prefix_kv(tmp_path: Path) -> None:
@@ -347,6 +382,10 @@ def test_denoise_default_never_touches_prefix_kv(tmp_path: Path) -> None:
     pipeline = _denoise(tmp_path, dit, num_steps=2)
     assert dit.calls == ["full", "full"]
     assert pipeline.last_prefix_kv_stats["enabled"] is False
+    assert len(dit.refiner_inputs) == 1
+    assert len(dit.text_embeddings) == 2
+    for embeds in dit.text_embeddings:
+        np.testing.assert_array_equal(np.asarray(embeds), dit.refiner_inputs[0] + 1.0)
 
 
 def test_denoise_refuses_a_cache_over_budget(tmp_path: Path) -> None:
@@ -354,6 +393,7 @@ def test_denoise_refuses_a_cache_over_budget(tmp_path: Path) -> None:
     with pytest.raises(PrefixKVBudgetError, match="prefix_kv_cache_max_gib"):
         _denoise(tmp_path, dit, reuse_prefix_kv=True, prefix_kv_cache_max_gib=1e-9)
     assert dit.calls == []
+    assert dit.refiner_inputs == []
 
 
 def test_reuse_needs_references() -> None:

fastvideo/tests/mlx/test_mlx_minimax_h3_ref2va_pipeline.py

@@ -71,14 +71,22 @@ class _RecordingDiT:
 
     def __init__(self) -> None:
         self.calls: list[dict[str, np.ndarray]] = []
+        self.refiner_inputs: list[np.ndarray] = []
+        self.text_embeddings: list[object] = []
 
     def precompute_adaln(self, timesteps, *, drop_weights: bool = True):
         self.adaln_timesteps = np.asarray(timesteps)
         return None
 
-    def forward_with_cache(self, x_v, x_a, text, *, layout, step_timesteps, row_timestep_inverse, step_index=0):
-        import mlx.core as mx
+    def refine_text(self, text):
+        self.refiner_inputs.append(np.asarray(text).copy())
+        # Distinguishable from the raw conditioning so a missing cache argument
+        # cannot accidentally satisfy the pipeline handoff assertions.
+        return text * 2.0 + 1.0
 
+    def forward_with_cache(self, x_v, x_a, text, *, layout, step_timesteps, row_timestep_inverse, step_index=0,
+                           precomputed_text_embeds=None):
+        self.text_embeddings.append(precomputed_text_embeds)
         self.calls.append({
             "video_rows": np.asarray(x_v),
             "audio_rows": np.asarray(x_a),
@@ -180,9 +188,10 @@ def test_denoise_hands_the_encoder_rows_to_every_forward(tmp_path: Path) -> None
     video_conditions = np.full((layout.num_condition_video_rows, PATCH_DIM), 3.0, dtype=np.float32)
     audio_conditions = np.full((layout.num_condition_audio_rows, AUDIO_IN_CHANNELS), 4.0, dtype=np.float32)
     dit = _RecordingDiT()
+    text_rows = np.arange(NUM_TEXT_TOKENS * 16, dtype=np.float32).reshape(NUM_TEXT_TOKENS, 16) / 16
 
     video_rows, audio_rows = pipeline.denoise(
-        np.zeros((NUM_TEXT_TOKENS, 16), dtype=np.float32),
+        text_rows,
         np.full(NUM_TEXT_TOKENS, 1, dtype=np.int64),
         height=HEIGHT,
         width=WIDTH,
@@ -197,6 +206,11 @@ def test_denoise_hands_the_encoder_rows_to_every_forward(tmp_path: Path) -> None
     )
 
     assert len(dit.calls) == 2
+    assert len(dit.refiner_inputs) == 1
+    np.testing.assert_array_equal(dit.refiner_inputs[0], text_rows)
+    assert len(dit.text_embeddings) == len(dit.calls)
+    for embeds in dit.text_embeddings:
+        np.testing.assert_array_equal(np.asarray(embeds), text_rows * 2.0 + 1.0)
     for call in dit.calls:
         assert call["video_rows"].shape == (layout.video_indices.shape[0], PATCH_DIM)
         assert call["audio_rows"].shape == (layout.audio_indices.shape[0], AUDIO_IN_CHANNELS)

fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py

@@ -387,7 +387,7 @@ def test_untile_hidden_never_consumes_padded_query_rows() -> None:
 
 
 def test_simd_unsupported_tile_falls_back_to_reference() -> None:
-    assert PUBLIC_VSA_IMPLS == ("auto", "reference", "simd")
+    assert PUBLIC_VSA_IMPLS == ("auto", "reference", "simd", "sdpa", "nax")
     assert "metal" not in PUBLIC_VSA_IMPLS
     assert resolve_impl("simd", head_dim=128, tile_elems=256) == "reference"
     assert resolve_impl("simd", head_dim=64, tile_elems=64) == "reference"

fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Independent selected-token checks for the explicitly requested M5 backend."""
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import PurePosixPath
+
+import numpy as np
+import pytest
+
+mx = pytest.importorskip("mlx.core", reason="Apple MLX is required")
+
+from fastvideo.mlx_runtime import minimax_h3_vsa as vsa
+from fastvideo.mlx_runtime import minimax_h3_vsa_nax as nax
+
+
+@pytest.fixture
+def nax_device():
+    if not mx.metal.is_available():
+        pytest.skip("NAX attention requires Metal on an Apple M5 GPU")
+    if not str(mx.device_info().get("device_name", "")).startswith("Apple M5"):
+        pytest.skip("NAX attention is validated only on Apple M5 GPUs")
+
+
+def _selected_rows(indices, counts, geometry, head, query):
+    tile = geometry.tile_elems
+    query_tile = geometry.num_prefix_tiles + query
+    query_rows = np.arange(query_tile * tile, query_tile * tile + geometry.variable_block_sizes[query_tile])
+    selected = indices[head, query, :counts[head, query]]
+    key_rows = (np.concatenate([np.arange(index * tile, index * tile + geometry.variable_block_sizes[index])
+                                for index in selected]) if len(selected) else np.empty(0, dtype=np.int64))
+    return query_rows, key_rows
+
+
+def _numpy_reference(q, k, value, indices, counts, geometry, scale):
+    """Select individual valid tokens independently of all MLX block masks."""
+    result = np.zeros_like(q, dtype=np.float64)
+    for head in range(q.shape[1]):
+        for query in range(geometry.num_video_tiles):
+            query_rows, key_rows = _selected_rows(indices, counts, geometry, head, query)
+            if not len(key_rows):
+                continue
+            scores = q[query_rows, head].astype(np.float64) @ k[key_rows, head].astype(np.float64).T * scale
+            probabilities = np.exp(scores - scores.max(axis=-1, keepdims=True))
+            probabilities /= probabilities.sum(axis=-1, keepdims=True)
+            result[query_rows, head] = probabilities @ value[key_rows, head].astype(np.float64)
+    return result
+
+
+def _native_reference(q, k, value, indices, counts, geometry, scale):
+    """Gather the same valid tokens; no custom kernel, padding, masks or GQA."""
+    result = np.zeros(q.shape, dtype=np.float32)
+    for head in range(q.shape[1]):
+        for query in range(geometry.num_video_tiles):
+            query_rows, key_rows = _selected_rows(indices, counts, geometry, head, query)
+            if not len(key_rows):
+                continue
+            output = mx.fast.scaled_dot_product_attention(
+                q[mx.array(query_rows), head][None, None],
+                k[mx.array(key_rows), head][None, None],
+                value[mx.array(key_rows), head][None, None], scale=scale, force_fused=True)
+            mx.eval(output)
+            result[query_rows, head] = np.asarray(output[0, 0].astype(mx.float32))
+    return result
+
+
+def _metrics(actual, expected):
+    difference = actual.astype(np.float64) - expected
+    return {"max_abs": float(np.abs(difference).max()),
+            "rmse": float(np.sqrt(np.mean(difference**2))),
+            "relative_l2": float(np.linalg.norm(difference) / max(np.linalg.norm(expected), 1e-30))}
+
+
+def _case(tile_size, dtype_name, *, segments=None):
+    # Dense/reference/target regions are interleaved before tiling. Ragged tiles
+    # cross 16/32/64-row boundaries, including inactive SIMD groups in a launch.
+    geometry = vsa.build_h3_segment_geometry(segments or (7, (1, 3, 5), 13, (3, 7, 11)), tile_size)
+    rng = np.random.default_rng(9241)
+    shape = (geometry.padded_length, 3, 128)
+    arrays = [rng.normal(size=shape).astype(np.float32) * 0.6 for _ in range(3)]
+    padding = np.concatenate([np.arange(index * tile_size + size, (index + 1) * tile_size)
+                              for index, size in enumerate(geometry.variable_block_sizes)])
+    for index, array in enumerate(arrays):
+        array[padding] = 30.0 + index
+    width = min(5, geometry.num_tiles)
+    # Intentionally out of bounds: inactive route slots must never be read.
+    indices = np.full((3, geometry.num_video_tiles, width), np.iinfo(np.int32).max, dtype=np.int32)
+    counts = np.empty((3, geometry.num_video_tiles), dtype=np.int32)
+    for head in range(3):
+        for query in range(geometry.num_video_tiles):
+            count = 1 + (head + query) % width
+            counts[head, query] = count
+            indices[head, query, :count] = rng.permutation(geometry.num_tiles)[:count]
+    counts[-1, -1] = 0
+    indices[-1, -1] = np.iinfo(np.int32).max
+    tensors = [mx.array(array).astype(getattr(mx, dtype_name)) for array in arrays]
+    mx.eval(*tensors)
+    return geometry, tensors, indices, counts, padding
+
+
+# MLX's installed steel/attn/nax.h sets MPP relaxed_precision=true even for
+# float32. The native gathered control reproduces its ~7e-4 NumPy discrepancy
+# on M5. Keep the 1e-3 independent FP32 bound and record BOTH numerical controls;
+# changing the selected tokens must not be excused as a precision difference.
+# The probability accumulation schedule differs from the native dense kernel.
+@pytest.mark.parametrize("dtype_name,tolerance", [("float16", 0.003), ("bfloat16", 0.02), ("float32", 0.001)])
+@pytest.mark.parametrize("tile_size,query_rows", [(128, 16), (128, 32), (128, 64), (128, 128),
+                                                  (64, 16), (64, 32), (64, 64)])
+def test_nax_matches_independent_valid_token_oracle(nax_device, record_property, dtype_name,
+                                                   tolerance, tile_size, query_rows):
+    geometry, tensors, indices, counts, padding = _case(tile_size, dtype_name)
+    scale = 128**-0.5
+    rounded = [np.asarray(tensor.astype(mx.float32)) for tensor in tensors]
+    expected = _numpy_reference(*rounded, indices, counts, geometry, scale)
+    actual = nax.nax_block_sparse(*tensors, mx.array(indices), mx.array(counts), geometry,
+                                  scale, query_rows=query_rows)
+    mx.eval(actual)
+    got = np.asarray(actual.astype(mx.float32))
+    record_property("candidate_vs_numpy", json.dumps(_metrics(got, expected)))
+    assert actual.shape == tensors[0].shape
+    assert actual.dtype == tensors[0].dtype
+    assert np.isfinite(got).all()
+    np.testing.assert_allclose(got, expected, atol=tolerance, rtol=tolerance)
+    np.testing.assert_array_equal(got[padding], 0)
+    np.testing.assert_array_equal(got[:geometry.num_prefix_tiles * tile_size], 0)
+    zero_rows, _ = _selected_rows(indices, counts, geometry, 2, geometry.num_video_tiles - 1)
+    np.testing.assert_array_equal(got[zero_rows, 2], 0)
+    if dtype_name == "float32":
+        native = _native_reference(*tensors, indices, counts, geometry, scale)
+        record_property("native_vs_numpy", json.dumps(_metrics(native, expected)))
+        record_property("candidate_vs_native", json.dumps(_metrics(got, native)))
+        np.testing.assert_allclose(native, expected, atol=tolerance, rtol=tolerance)
+        # Both paths have independently bounded NumPy error. This pairwise
+        # budget is their sum, not a claim of bitwise native-kernel equivalence.
+        np.testing.assert_allclose(got, native, atol=2 * tolerance, rtol=2 * tolerance)
+
+
+@pytest.mark.parametrize("query_rows", [16, 32, 64, 128])
+def test_nax_uniform_logits_expose_head_and_tile_selection(nax_device, query_rows):
+    geometry, tensors, indices, counts, padding = _case(128, "float32")
+    shape = tensors[0].shape
+    q, k = mx.zeros(shape), mx.zeros(shape)
+    value = np.zeros(shape, dtype=np.float32)
+    for tile, size in enumerate(geometry.variable_block_sizes):
+        for head in range(shape[1]):
+            # Each tile and head carries a distinct basis vector. Uniform
+            # logits make the answer a valid-row-weighted tile histogram.
+            value[tile * 128:tile * 128 + size, head, tile] = head + 1
+    value[padding] = 32.0
+    expected = _numpy_reference(np.zeros(shape), np.zeros(shape), value, indices, counts, geometry, 0.071)
+    actual = nax.nax_block_sparse(q, k, mx.array(value), mx.array(indices), mx.array(counts),
+                                  geometry, 0.071, query_rows=query_rows)
+    mx.eval(actual)
+    np.testing.assert_allclose(np.asarray(actual), expected, atol=2e-6, rtol=2e-6)
+
+
+def test_nax_without_dense_prefix_and_with_nondefault_scale(nax_device):
+    geometry, tensors, indices, counts, padding = _case(64, "bfloat16", segments=((3, 7, 11),))
+    assert geometry.num_prefix_tiles == 0
+    expected = _numpy_reference(*[np.asarray(t.astype(mx.float32)) for t in tensors],
+                                indices, counts, geometry, 0.17)
+    actual = nax.nax_block_sparse(*tensors, mx.array(indices), mx.array(counts), geometry, 0.17,
+                                  query_rows=32)
+    mx.eval(actual)
+    np.testing.assert_allclose(np.asarray(actual.astype(mx.float32)), expected, atol=0.02, rtol=0.02)
+    np.testing.assert_array_equal(np.asarray(actual.astype(mx.float32))[padding], 0)
+
+
+def _attention_inputs():
+    geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), 128)
+    rng = np.random.default_rng(874)
+    shape = (geometry.total_seq_length, 2, 128)
+    tensors = [mx.array(rng.normal(size=shape).astype(np.float32) * 0.1) for _ in range(4)]
+    return geometry, tensors
+
+
+def test_explicit_nax_preserves_router_dense_prefix_and_compression(nax_device, monkeypatch):
+    geometry, (q, k, value, gate) = _attention_inputs()
+    expected = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5,
+                                     gate_compress=gate, impl="reference")
+    calls = []
+    original = nax.nax_block_sparse
+
+    def record(*args, **kwargs):
+        calls.append((args[3].shape, args[4].shape))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(nax, "nax_block_sparse", record)
+    stats = vsa.MiniMaxH3VSAStats()
+    actual = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5,
+                                  gate_compress=gate, impl="nax", stats=stats)
+    mx.eval(actual, expected)
+    assert calls == [((2, geometry.num_video_tiles, calls[0][0][-1]), (2, geometry.num_video_tiles))]
+    assert stats.impl == "nax"
+    assert stats.dense_fallback_reason is None
+    np.testing.assert_allclose(np.asarray(actual), np.asarray(expected), atol=2e-5, rtol=2e-4)
+
+
+def test_explicit_nax_propagates_kernel_failure(nax_device, monkeypatch):
+    geometry, (q, k, value, _) = _attention_inputs()
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("deliberate NAX kernel failure")
+
+    monkeypatch.setattr(nax, "nax_block_sparse", fail)
+    with pytest.raises(RuntimeError, match="deliberate NAX kernel failure"):
+        actual = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, impl="nax")
+        mx.eval(actual)
+
+
+def test_nax_requires_explicit_selection(monkeypatch):
+    from fastvideo.mlx_runtime import minimax_h3_vsa_simd as simd
+
+    monkeypatch.setattr(simd, "simd_kernel_available", lambda: True)
+    assert "nax" in vsa.PUBLIC_VSA_IMPLS
+    assert vsa.resolve_impl("nax", 128, 128) == "nax"
+    assert vsa.resolve_impl("nax", 128, 64) == "nax"
+    assert vsa.resolve_impl("auto", 128, 128) == "simd"
+    assert vsa.resolve_impl("reference", 128, 128) == "reference"
+    monkeypatch.setattr(simd, "simd_kernel_available", lambda: False)
+    monkeypatch.setattr(simd, "simd_kernel_error", lambda: "test unavailable")
+    assert vsa.resolve_impl("auto", 128, 128) == "reference"
+
+
+@pytest.mark.parametrize("head_dim,tile_size", [(64, 128), (128, 256)])
+def test_nax_rejects_unsupported_public_shapes(head_dim, tile_size):
+    with pytest.raises(ValueError):
+        vsa.resolve_impl("nax", head_dim, tile_size)
+
+
+@pytest.mark.parametrize("device_name", ["Apple M4 Max", "Apple M3", "Apple M50", "AMD Radeon", ""])
+def test_nax_rejects_unsupported_device_before_kernel_construction(monkeypatch, device_name):
+    monkeypatch.setattr(mx.metal, "is_available", lambda: True)
+    monkeypatch.setattr(mx, "device_info", lambda: {"device_name": device_name}, raising=False)
+
+    def unexpected(*args, **kwargs):
+        pytest.fail("An unsupported device must not construct a Metal kernel")
+
+    monkeypatch.setattr(mx.fast, "metal_kernel", unexpected, raising=False)
+    nax._nax_kernel.cache_clear()
+    try:
+        with pytest.raises(RuntimeError, match="M5"):
+            nax._nax_kernel()
+    finally:
+        nax._nax_kernel.cache_clear()
+
+
+def test_nax_rejects_missing_metal(monkeypatch):
+    monkeypatch.setattr(mx.metal, "is_available", lambda: False)
+    with pytest.raises(RuntimeError, match="Metal|M5"):
+        nax._check_nax_device()
+
+
+def test_nax_construction_failure_keeps_original_cause(monkeypatch):
+    original_error = RuntimeError("deliberate compiler error")
+    monkeypatch.setattr(nax, "_check_nax_device", lambda: None)
+    monkeypatch.setattr(nax, "_header", lambda: ("unused", {}))
+
+    def fail(*args, **kwargs):
+        raise original_error
+
+    monkeypatch.setattr(mx.fast, "metal_kernel", fail, raising=False)
+    nax._nax_kernel.cache_clear()
+    try:
+        with pytest.raises(RuntimeError, match="kernel construction") as captured:
+            nax._nax_kernel()
+        assert captured.value.__cause__ is original_error
+    finally:
+        nax._nax_kernel.cache_clear()
+
+
+def test_unvalidated_precision_override_is_not_a_public_option():
+    with pytest.raises(TypeError, match="relaxed_precision"):
+        nax.nax_block_sparse(None, None, None, None, None, None, 1.0,
+                             relaxed_precision=False)
+
+
+def _assert_nax_provenance_schema(provenance):
+    assert provenance["relaxed_precision"] is True
+    assert provenance["math_mode"] == "safe"
+    assert provenance["key_rows"] == 32
+    for name in ("source_sha256", "expanded_header_sha256"):
+        assert len(provenance[name]) == 64
+        int(provenance[name], 16)
+    dependencies = provenance["installed_header_sha256"]
+    assert dependencies
+    for path, digest in dependencies.items():
+        assert path.startswith("mlx/")
+        assert not PurePosixPath(path).is_absolute()
+        assert ".." not in PurePosixPath(path).parts
+        assert len(digest) == 64
+        int(digest, 16)
+
+
+def test_nax_provenance_schema_without_installed_metal_headers(monkeypatch):
+    # The Linux MLX CPU wheel need not ship Apple's Metal helper headers.
+    # Exercise receipt formatting and hashing without consulting that package.
+    header = "// synthetic expanded NAX header\n"
+    dependencies = {"mlx/backend/metal/kernels/steel/attn/nax.h": hashlib.sha256(b"synthetic header").hexdigest()}
+    monkeypatch.setattr(nax, "_header", lambda: (header, dependencies))
+    provenance = nax.get_nax_provenance()
+    _assert_nax_provenance_schema(provenance)
+    assert provenance["expanded_header_sha256"] == hashlib.sha256(header.encode()).hexdigest()
+    assert provenance["source_sha256"] == hashlib.sha256(nax._SOURCE.encode()).hexdigest()
+    assert provenance["installed_header_sha256"] == dependencies
+
+
+def test_nax_installed_header_provenance_on_supported_device(nax_device):
+    # A supported M5 installation must supply the real validated headers;
+    # missing or incompatible headers are a failure here, not a skip.
+    provenance = nax.get_nax_provenance()
+    _assert_nax_provenance_schema(provenance)
+    include_root = nax._mlx_include_root()
+    for relative, digest in provenance["installed_header_sha256"].items():
+        assert hashlib.sha256((include_root / relative).read_bytes()).hexdigest() == digest
+
+
+@pytest.fixture
+def synthetic_nax_headers(monkeypatch, tmp_path):
+    # Parse a tiny include tree without Metal or the installed private headers.
+    # The expected API markers deliberately match the loader's known contract.
+    dependency = "mlx/testing/nax_dependency.h"
+    descriptor = "transpose_b,\n        true,\n        mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate\n"
+    header = (f'#include "{dependency}"\n'
+              "struct BaseNAXFrag; struct NAXTile;\n"
+              "get_coord(); row_reduce(); row_bin_op(); load_rows(); store_rows();\n"
+              + descriptor * 2)
+    files = {nax._NAX_HEADER: header, dependency: "// validated dependency\n"}
+    hashes = {}
+    for relative, text in files.items():
+        path = tmp_path / relative
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(text)
+        hashes[relative] = hashlib.sha256(path.read_bytes()).hexdigest()
+    monkeypatch.setattr(nax, "_mlx_include_root", lambda: tmp_path)
+    monkeypatch.setattr(nax, "_VALIDATED_HEADER_SHA256", hashes)
+    nax._header.cache_clear()
+    try:
+        yield tmp_path, dependency, hashes
+    finally:
+        nax._header.cache_clear()
+
+
+def test_nax_rejects_modified_validated_header(synthetic_nax_headers):
+    include_root, dependency, hashes = synthetic_nax_headers
+    _, actual_hashes = nax._header()
+    assert actual_hashes == hashes
+    (include_root / dependency).write_text("// changed dependency with unchanged API\n")
+    nax._header.cache_clear()
+    with pytest.raises(RuntimeError, match="Unvalidated MLX private header") as captured:
+        nax._header()
+    assert dependency in str(captured.value)
+
+
+def test_nax_rejects_unregistered_header_dependency(synthetic_nax_headers):
+    include_root, _, hashes = synthetic_nax_headers
+    unknown = "mlx/testing/unregistered_dependency.h"
+    (include_root / unknown).write_text("// This dependency is outside the validated set.\n")
+    header_path = include_root / nax._NAX_HEADER
+    header_path.write_text(header_path.read_text() + f'#include "{unknown}"\n')
+    hashes[nax._NAX_HEADER] = hashlib.sha256(header_path.read_bytes()).hexdigest()
+    with pytest.raises(RuntimeError, match="Unvalidated MLX private header") as captured:
+        nax._header()
+    assert unknown in str(captured.value)
+
+
+def test_pipeline_rejects_nax_prefix_reuse_before_loading_components():
+    from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline
+
+    pipeline = object.__new__(MiniMaxH3MLXPipeline)
+    with pytest.raises(ValueError, match="prefix|Prefix"):
+        pipeline.generate("unused", output_path="unused.mp4", vsa_impl="nax", reuse_prefix_kv=True)

fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Masked SDPA keeps H3's selected tiles and ragged-row semantics."""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+mx = pytest.importorskip("mlx.core", reason="Apple MLX is required")
+
+from fastvideo.mlx_runtime import minimax_h3_vsa as vsa
+from fastvideo.mlx_runtime import minimax_h3_vsa_sdpa as sdpa
+
+pytestmark = pytest.mark.skipif(not mx.metal.is_available(), reason="Fused MLX SDPA requires Metal")
+
+
+def _selected_token_reference(q, k, value, indices, counts, geometry):
+    """Independent token selection, followed by NumPy float64 softmax."""
+    tile = geometry.tile_elems
+    sizes = geometry.variable_block_sizes
+    output = np.zeros_like(q, dtype=np.float64)
+    for head in range(q.shape[1]):
+        for query in range(geometry.num_video_tiles):
+            query_tile = geometry.num_prefix_tiles + query
+            query_rows = np.arange(query_tile * tile, query_tile * tile + sizes[query_tile])
+            selected = indices[head, query, :counts[head, query]]
+            if not len(selected):
+                continue
+            key_rows = np.concatenate([np.arange(index * tile, index * tile + sizes[index]) for index in selected])
+            logits = q[query_rows, head].astype(np.float64) @ k[key_rows, head].astype(np.float64).T
+            logits *= q.shape[-1]**-0.5
+            probabilities = np.exp(logits - logits.max(axis=-1, keepdims=True))
+            probabilities /= probabilities.sum(axis=-1, keepdims=True)
+            output[query_rows, head] = probabilities @ value[key_rows, head].astype(np.float64)
+    return output
+
+
+def _gathered_native_reference(q, k, value, indices, counts, geometry):
+    """Native control with independently gathered valid keys and padded queries.
+
+    Match the backend's query length: shorter ragged queries can select a
+    different native arithmetic path. Query rows attend independently, so
+    padding queries does not change the expected values of the valid rows.
+    This control uses no GQA, key padding or attention mask.
+    """
+    output = np.zeros(q.shape, dtype=np.float32)
+    tile = geometry.tile_elems
+    sizes = geometry.variable_block_sizes
+    for head in range(q.shape[1]):
+        for query in range(geometry.num_video_tiles):
+            query_tile = geometry.num_prefix_tiles + query
+            selected = indices[head, query, :counts[head, query]]
+            if not len(selected):
+                continue
+            query_start = query_tile * tile
+            query_rows = np.arange(query_start, query_start + sizes[query_tile])
+            key_rows = np.concatenate([np.arange(index * tile, index * tile + sizes[index]) for index in selected])
+            result = mx.fast.scaled_dot_product_attention(
+                q[query_start:query_start + tile, head][None, None],
+                k[mx.array(key_rows), head][None, None],
+                value[mx.array(key_rows), head][None, None],
+                scale=q.shape[-1]**-0.5, force_fused=True)
+            mx.eval(result)
+            output[query_rows, head] = np.asarray(result[0, 0, :sizes[query_tile]].astype(mx.float32))
+    return output
+
+
+# MLX steel/attn/nax.h uses MPP matmul2d_descriptor(relaxed_precision=true),
+# including float32 inputs. The MPP SDK explicitly permits an accuracy/performance
+# tradeoff. On M5, the independent native gathered control and this backend both
+# differ from NumPy by up to 7.45e-4 across these tile sizes. Matching their query
+# lengths reduces native arithmetic differences without sharing key selection.
+# Bound that native error separately from the stricter mask/layout control below.
+@pytest.mark.parametrize("dtype_name,tolerance", [("float32", 1e-3), ("bfloat16", 0.02)])
+@pytest.mark.parametrize("chunk_size", [1, 3, 8])
+@pytest.mark.parametrize("tile_size", [64, 128])
+def test_sparse_sdpa_matches_selected_valid_tokens(dtype_name, tolerance, chunk_size, tile_size):
+    geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), tile_size)
+    rng = np.random.default_rng(9241)
+    heads, dim = 3, 128
+    shape = (geometry.padded_length, heads, dim)
+    originals = [rng.normal(size=shape).astype(np.float32) * 0.6 for _ in range(3)]
+    tile = geometry.tile_elems
+    sizes = geometry.variable_block_sizes
+    padding = np.concatenate([np.arange(index * tile + count, (index + 1) * tile)
+                              for index, count in enumerate(sizes)])
+    for index, array in enumerate(originals):
+        array[padding] = 30.0 + index  # Missing padding masks cause large errors.
+    width = min(5, geometry.num_tiles)
+    indices = np.zeros((heads, geometry.num_video_tiles, width), dtype=np.int32)
+    counts = np.empty((heads, geometry.num_video_tiles), dtype=np.int32)
+    for head in range(heads):
+        for query in range(geometry.num_video_tiles):
+            count = 1 + (head + query) % width
+            counts[head, query] = count
+            indices[head, query, :count] = rng.permutation(geometry.num_tiles)[:count]
+            unselected = np.setdiff1d(np.arange(geometry.num_tiles), indices[head, query, :count])
+            indices[head, query, count:] = unselected[-1] if len(unselected) else 0
+    counts[-1, -1] = 0
+    q, k, value = [mx.array(item).astype(getattr(mx, dtype_name)) for item in originals]
+    mx.eval(q, k, value)
+    rounded = [np.asarray(item.astype(mx.float32)) for item in (q, k, value)]
+    expected = _selected_token_reference(*rounded, indices, counts, geometry)
+    actual = sdpa.sdpa_block_sparse(q, k, value, mx.array(indices), mx.array(counts), geometry,
+                                    dim**-0.5, chunk_size=chunk_size)
+    mx.eval(actual)
+    got = np.asarray(actual.astype(mx.float32))
+    assert np.isfinite(got).all()
+    np.testing.assert_allclose(got, expected, atol=tolerance, rtol=tolerance)
+    if dtype_name == "float32":
+        native = _gathered_native_reference(q, k, value, indices, counts, geometry)
+        np.testing.assert_allclose(native, expected, atol=tolerance, rtol=tolerance)
+        # This independent control removes all candidate mask construction,
+        # reshaping and GQA, so a wrong selected tile or head cannot hide behind
+        # the NAX kernel's looser float32 arithmetic.
+        np.testing.assert_allclose(got, native, atol=1e-4, rtol=2e-4)
+    np.testing.assert_array_equal(got[padding], 0)
+    np.testing.assert_array_equal(got[:geometry.num_prefix_tiles * tile], 0)
+
+
+def test_mask_memory_cap_splits_dispatch_without_changing_selected_tokens(monkeypatch):
+    geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), 128)
+    rng = np.random.default_rng(116)
+    heads, dim = 3, 128
+    arrays = [rng.normal(size=(geometry.padded_length, heads, dim)).astype(np.float32) * 0.6
+              for _ in range(3)]
+    padding = np.concatenate([np.arange(tile * 128 + count, (tile + 1) * 128)
+                              for tile, count in enumerate(geometry.variable_block_sizes)])
+    for index, array in enumerate(arrays):
+        array[padding] = 30.0 + index
+    indices = np.empty((heads, geometry.num_video_tiles, 3), dtype=np.int32)
+    for head in range(heads):
+        for query in range(geometry.num_video_tiles):
+            indices[head, query] = (head + query + np.array([0, 2, 4])) % geometry.num_tiles
+    counts = np.full((heads, geometry.num_video_tiles), 3, dtype=np.int32)
+    counts[-1, -1] = 0
+    expected = _selected_token_reference(*arrays, indices, counts, geometry)
+    q, k, value = [mx.array(array) for array in arrays]
+
+    # A fixed tiny budget exercises the automatic cap without allocating a
+    # production-sized mask. Observe actual native calls, not the chunk formula.
+    cap_bytes = 8192
+    monkeypatch.setattr(sdpa, "_MASK_CHUNK_MAX_BYTES", cap_bytes)
+    dispatches = []
+    native_sdpa = mx.fast.scaled_dot_product_attention
+
+    def record_dispatch(query, key, value, **kwargs):
+        dispatches.append({"mask_bytes": kwargs["mask"].nbytes,
+                           "query_heads": query.shape[1]})
+        return native_sdpa(query, key, value, **kwargs)
+
+    monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", record_dispatch)
+    actual = sdpa.sdpa_block_sparse(q, k, value, mx.array(indices), mx.array(counts), geometry,
+                                    dim**-0.5, chunk_size=geometry.num_video_tiles)
+    mx.eval(actual)
+    assert len(dispatches) > 1
+    assert all(0 < dispatch["mask_bytes"] <= cap_bytes for dispatch in dispatches)
+    assert sum(dispatch["query_heads"] for dispatch in dispatches) == heads * geometry.num_video_tiles
+    got = np.asarray(actual)
+    assert np.isfinite(got).all()
+    np.testing.assert_allclose(got, expected, atol=1e-3, rtol=1e-3)
+    np.testing.assert_array_equal(got[padding], 0)
+    np.testing.assert_array_equal(got[:geometry.num_prefix_tiles * geometry.tile_elems], 0)
+
+
+def _attention_inputs():
+    geometry = vsa.build_h3_segment_geometry((7, (1, 3, 5), 13, (3, 7, 11)), 128)
+    rng = np.random.default_rng(874)
+    shape = (geometry.total_seq_length, 2, 128)
+    tensors = [mx.array(rng.normal(size=shape).astype(np.float32) * 0.1) for _ in range(4)]
+    return geometry, tensors
+
+
+def test_explicit_sdpa_routes_without_changing_attention_selection(monkeypatch):
+    geometry, (q, k, value, gate) = _attention_inputs()
+    expected = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5,
+                                     gate_compress=gate, impl="reference")
+    calls = []
+    original = sdpa.sdpa_block_sparse
+
+    def record(*args, **kwargs):
+        calls.append((args[3].shape, args[4].shape))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(sdpa, "sdpa_block_sparse", record)
+    stats = vsa.MiniMaxH3VSAStats()
+    actual = vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5,
+                                  gate_compress=gate, impl="sdpa", stats=stats)
+    mx.eval(actual, expected)
+    assert len(calls) == 1
+    assert calls[0][0][:2] == (2, geometry.num_video_tiles)
+    assert calls[0][1] == (2, geometry.num_video_tiles)
+    assert stats.impl == "sdpa"
+    assert stats.dense_fallback_reason is None
+    np.testing.assert_allclose(np.asarray(actual), np.asarray(expected), atol=2e-5, rtol=2e-4)
+
+
+def test_explicit_sdpa_propagates_kernel_failure(monkeypatch):
+    geometry, (q, k, value, _) = _attention_inputs()
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("test SDPA failure")
+
+    monkeypatch.setattr(sdpa, "sdpa_block_sparse", fail)
+    with pytest.raises(RuntimeError, match="test SDPA failure"):
+        vsa.h3_vsa_attention(q, k, value, geometry, sparsity=0.5, ref_keep_rate=0.5, impl="sdpa")
+
+
+def test_sdpa_is_explicit_and_auto_selection_stays_unchanged(monkeypatch):
+    from fastvideo.mlx_runtime import minimax_h3_vsa_simd as simd
+
+    monkeypatch.setattr(simd, "simd_kernel_available", lambda: True)
+    assert vsa.resolve_impl("sdpa", 128, 128) == "sdpa"
+    assert vsa.resolve_impl("auto", 128, 128) == "simd"
+
+
+@pytest.mark.parametrize("head_dim,tile_size", [(64, 128), (128, 256)])
+def test_explicit_sdpa_rejects_unsupported_shapes(head_dim, tile_size):
+    with pytest.raises(ValueError):
+        vsa.resolve_impl("sdpa", head_dim, tile_size)
+
+
+def test_pipeline_rejects_sdpa_prefix_reuse_before_loading_components():
+    from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline
+
+    pipeline = object.__new__(MiniMaxH3MLXPipeline)
+    with pytest.raises(ValueError, match="prefix|Prefix"):
+        pipeline.generate("unused", output_path="unused.mp4", vsa_impl="sdpa", reuse_prefix_kv=True)
+
+
+def test_sparse_sdpa_rejects_invalid_chunk_size():
+    geometry = vsa.build_h3_segment_geometry((3, (1, 3, 5)), 128)
+    array = mx.zeros((geometry.padded_length, 1, 128))
+    indices = mx.zeros((1, geometry.num_video_tiles, 1), dtype=mx.int32)
+    counts = mx.ones((1, geometry.num_video_tiles), dtype=mx.int32)
+    with pytest.raises(ValueError, match="chunk"):
+        sdpa.sdpa_block_sparse(array, array, array, indices, counts, geometry, 128**-0.5, chunk_size=0)

fastvideo/tests/ssim/test_qwen_image21_similarity.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-GPU Qwen T2I and editing regressions, including exported alpha.
+
+Requires the Qwen checkpoint, FlashAttention, and device-matched PNG references.
+The shared SSIM harness compares all four RGBA channels and records metrics.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+from PIL import Image, ImageDraw
+
+from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    DEVICE_MAPPINGS,
+    run_text_to_video_similarity_test,
+)
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 1
+MODEL_ID = "Qwen/Qwen-Image-2.1"
+MODEL_REVISION = "d26bb61231c349cf6b7896fa83353113880e1ba3"
+MIN_ACCEPTABLE_SSIM = 0.98
+
+QWEN_IMAGE21_DEFAULT_PARAMS: dict[str, object] = {
+    "model_path": MODEL_ID,
+    "num_gpus": REQUIRED_GPUS,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": 512,
+    "width": 512,
+    "num_frames": 1,
+    "fps": 1,
+    "num_inference_steps": 4,
+    "guidance_scale": 1.0,
+    "seed": 21,
+}
+
+QWEN_IMAGE21_FULL_QUALITY_PARAMS: dict[str, object] = {
+    **QWEN_IMAGE21_DEFAULT_PARAMS,
+    "height": 1024,
+    "width": 1024,
+    "num_inference_steps": 40,
+}
+
+QWEN_IMAGE21_MODEL_TO_PARAMS = {"Qwen/Qwen-Image-2.1": QWEN_IMAGE21_DEFAULT_PARAMS}
+QWEN_IMAGE21_FULL_QUALITY_MODEL_TO_PARAMS = {"Qwen/Qwen-Image-2.1": QWEN_IMAGE21_FULL_QUALITY_PARAMS}
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Qwen image SSIM requires one CUDA GPU")
+
+
+def _run_similarity(model_id: str, attention_backend_name: str, prompt: str,
+                    references: list[str] | None = None) -> None:
+    device_reference_folder = resolve_device_reference_folder(
+        DEVICE_MAPPINGS + (("RTX 5090", "RTX5090"), ("RTX 4090", "RTX4090")),
+        device_name=get_cuda_device_name(),
+        logger=logger,
+    )
+    if device_reference_folder is None:
+        pytest.skip(f"No Qwen image SSIM reference mapping for {get_cuda_device_name()}")
+
+    script_dir = str(Path(__file__).resolve().parent)
+    selected_params = select_ssim_params(QWEN_IMAGE21_MODEL_TO_PARAMS,
+                                        QWEN_IMAGE21_FULL_QUALITY_MODEL_TO_PARAMS)[model_id]
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=script_dir,
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=QWEN_IMAGE21_MODEL_TO_PARAMS,
+        full_quality_params_map=QWEN_IMAGE21_FULL_QUALITY_MODEL_TO_PARAMS,
+        min_acceptable_ssim=MIN_ACCEPTABLE_SSIM,
+        media_extension=".png",
+        init_kwargs_override={
+            "pipeline_config": QwenImage21PipelineConfig(
+                reference_resolution=selected_params["height"],
+                use_kv_cache=True,
+            ),
+            "revision": MODEL_REVISION,
+            "workload_type": "i2i" if references else "t2i",
+            "use_fsdp_inference": False,
+            "dit_cpu_offload": True,
+            "dit_layerwise_offload": True,
+            "text_encoder_cpu_offload": True,
+            "vae_cpu_offload": True,
+            "lazy_module_load": True,
+            "pin_cpu_memory": False,
+            "vae_tiling": True,
+        },
+        generation_kwargs_override={
+            "references": references,
+            "negative_prompt": None,
+            "true_cfg_scale": 1.0,
+            "return_frames": True,
+            "save_video": True,
+        },
+    )
+
+    output_dir = build_generated_output_dir(script_dir, device_reference_folder, model_id, attention_backend_name)
+    with Image.open(Path(output_dir) / f"{prompt[:100].strip()}.png") as generated:
+        assert generated.mode == "RGBA", "Qwen PNG export must retain its alpha channel"
+        assert generated.size == (selected_params["width"], selected_params["height"])
+
+
+@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
+@pytest.mark.parametrize("model_id", list(QWEN_IMAGE21_MODEL_TO_PARAMS))
+def test_qwen_image21_t2i_similarity(model_id: str, attention_backend_name: str) -> None:
+    _run_similarity(
+        model_id,
+        attention_backend_name,
+        "A studio photograph of a blue ceramic vase on a white background",
+    )
+
+
+@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
+@pytest.mark.parametrize("model_id", list(QWEN_IMAGE21_MODEL_TO_PARAMS))
+def test_qwen_image21_edit_similarity(model_id: str, attention_backend_name: str, tmp_path: Path) -> None:
+    reference = Image.new("RGB", (512, 512), "white")
+    draw = ImageDraw.Draw(reference)
+    draw.ellipse((176, 160, 336, 432), fill=(32, 96, 224))
+    draw.rectangle((224, 96, 288, 192), fill=(32, 96, 224))
+    draw.ellipse((224, 88, 288, 112), fill=(32, 96, 224))
+    reference_path = tmp_path / "blue-vase.png"
+    reference.save(reference_path)
+    _run_similarity(
+        model_id,
+        attention_backend_name,
+        "Change the blue vase to red and keep the white background",
+        references=[str(reference_path)],
+    )

fastvideo/tests/stages/test_input_validation.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Weight-free regression tests for negative conditioning validation."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+
+@pytest.fixture(autouse=True)
+def disable_gpu_stage_timing(monkeypatch):
+    monkeypatch.setattr(
+        "fastvideo.pipelines.stages.base.envs.FASTVIDEO_STAGE_LOGGING",
+        SimpleNamespace(get=lambda: False),
+    )
+
+
+def _batch(**overrides):
+    kwargs = dict(
+        data_type="video",
+        prompt="a red fox",
+        negative_prompt=None,
+        height=16,
+        width=16,
+        num_frames=1,
+        seed=0,
+        num_inference_steps=2,
+        guidance_scale=5.0,
+    )
+    kwargs.update(overrides)
+    return ForwardBatch(**kwargs)
+
+
+def _args(verify):
+    return SimpleNamespace(
+        enable_stage_verification=verify,
+        pipeline_config=SimpleNamespace(ti2v_task=False, is_causal=False),
+    )
+
+
+@pytest.mark.parametrize("verify", [False, True])
+@pytest.mark.parametrize("negative_prompt_embeds", [None, []])
+def test_cfg_rejects_missing_negative_conditioning(verify, negative_prompt_embeds):
+    batch = _batch(negative_prompt_embeds=negative_prompt_embeds)
+    assert batch.do_classifier_free_guidance
+    assert batch.negative_prompt_embeds == []
+    with pytest.raises(ValueError, match="For classifier-free guidance, either `negative_prompt`"):
+        InputValidationStage()(batch, _args(verify))
+
+
+@pytest.mark.parametrize("verify", [False, True])
+def test_cfg_accepts_empty_negative_prompt(verify):
+    batch = _batch(negative_prompt="")
+    assert InputValidationStage()(batch, _args(verify)) is batch
+    assert batch.negative_prompt == ""
+
+
+@pytest.mark.parametrize("verify", [False, True])
+def test_cfg_accepts_precomputed_negative_embeddings(verify):
+    embeddings = [torch.ones(1, 2, 3)]
+    batch = _batch(negative_prompt_embeds=embeddings)
+    assert InputValidationStage()(batch, _args(verify)) is batch
+    assert batch.negative_prompt_embeds is embeddings
+    assert torch.equal(batch.negative_prompt_embeds[0], torch.ones(1, 2, 3))
+
+
+@pytest.mark.parametrize("verify", [False, True])
+@pytest.mark.parametrize("negative_prompt_embeds", [None, []])
+def test_without_cfg_does_not_require_negative_conditioning(verify, negative_prompt_embeds):
+    batch = _batch(guidance_scale=1.0, negative_prompt_embeds=negative_prompt_embeds)
+    assert not batch.do_classifier_free_guidance
+    assert InputValidationStage()(batch, _args(verify)) is batch
+
+
+def test_empty_negative_prompt_reaches_text_encoding(monkeypatch):
+    batch = _batch(negative_prompt="")
+    args = _args(False)
+    args.pipeline_config.text_encoder_configs = [SimpleNamespace()]
+    encoder = TextEncodingStage([object()], [object()])
+    prompts = []
+
+    def encode_text(prompt, *args, **kwargs):
+        prompts.append(prompt)
+        return [torch.ones(1, 2, 3)], [torch.ones(1, 2, dtype=torch.long)]
+
+    monkeypatch.setattr(encoder, "encode_text", encode_text)
+    InputValidationStage()(batch, args)
+    assert encoder(batch, args) is batch
+    assert prompts == ["a red fox", ""]
+    assert len(batch.prompt_embeds) == len(batch.negative_prompt_embeds) == 1

fastvideo/tests/stages/test_qwen_image21_stages.py

@@ -0,0 +1,304 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen request and stage contracts with real batches and weight-free components.
+
+Tensor operations run on CPU. These checks exercise stage integration; independent
+checkpoint parity remains in tests/local_tests/qwen_image21/.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+from fastvideo.models.dits.qwen_image21 import QwenImage21KVCacheAllocationError
+from fastvideo.pipelines.basic.qwen_image21 import stages
+from fastvideo.pipelines.basic.qwen_image21.stages import STATE_KEY
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+@pytest.fixture
+def stage_args(monkeypatch):
+    # Keep the real stages and caches; only replace runtime device/weight moves.
+    monkeypatch.setattr(stages, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(stages.pinned_offload, "load", lambda model, device, pin=False: model.to(device))
+    return SimpleNamespace(
+        tp_size=1,
+        sp_size=1,
+        num_gpus=1,
+        text_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        use_fsdp_inference=False,
+        pin_cpu_memory=False,
+        output_type="pil",
+        pipeline_config=SimpleNamespace(
+            dit_config=SimpleNamespace(arch_config=SimpleNamespace(in_channels=2)),
+            vae_tiling=False,
+            reference_resolution=32,
+            use_kv_cache=True,
+            kv_cache_device="cpu",
+        ),
+    )
+
+
+def _batch(**updates):
+    values = dict(
+        data_type="image",
+        prompt="draw",
+        height=64,
+        width=64,
+        num_frames=1,
+        num_inference_steps=2,
+        guidance_scale=1.0,
+        true_cfg_scale=1.0,
+        seed=42,
+    )
+    values.update(updates)
+    return ForwardBatch(**values)
+
+
+class TinyProcessor:
+    tokenizer = SimpleNamespace(encode=lambda _text: [9])
+
+    def __init__(self):
+        self.images = []
+
+    def apply_chat_template(self, *_args, **_kwargs):
+        return [[1, 2]]
+
+    def __call__(self, text, images=None, **kwargs):
+        assert kwargs["padding_side"] == "left"
+        self.images = list(images or [])
+        last_id = 5 if "bad" in text[0] else 4
+        ids = torch.tensor([[1, 2, 3] + [9] * len(self.images) + [last_id]])
+        return dict(input_ids=ids, attention_mask=torch.ones_like(ids))
+
+
+class TinyEncoder(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(1))
+
+    def forward(self, input_ids, **_kwargs):
+        return input_ids.float().unsqueeze(-1).repeat(1, 1, 2)
+
+
+class TinyVAE(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(1))
+        self.config = SimpleNamespace(latents_mean=(0.0, 0.0), latents_std=(1.0, 1.0))
+        self.encoded_pixels = []
+
+    def disable_tiling(self):
+        pass
+
+    def encode(self, pixels):
+        self.encoded_pixels.append(pixels.clone())
+        latent = torch.nn.functional.avg_pool2d(pixels[:, [0, 3], 0], 16).unsqueeze(2)
+        return SimpleNamespace(latent_dist=SimpleNamespace(mode=lambda: latent))
+
+    def decode(self, latents, return_dict=False):
+        image = latents[:, :1].repeat_interleave(16, 3).repeat_interleave(16, 4)
+        return (torch.cat((image.repeat(1, 3, 1, 1, 1), torch.full_like(image, -0.5)), dim=1),)
+
+
+class TinyScheduler:
+    config = dict(base_image_seq_len=256, max_image_seq_len=8192, base_shift=0.5, max_shift=0.9)
+
+    def set_timesteps(self, sigmas, device, mu):
+        self.timesteps = torch.tensor(sigmas, device=device) * 1000
+
+    def set_begin_index(self, index):
+        assert index == 0
+
+    def step(self, noise, timestep, latents, return_dict=False):
+        return (latents - noise * 0.1,)
+
+
+class TinyTransformer(nn.Module):
+    """Populate real prefix caches so ownership and cleanup checks inspect data."""
+
+    def __init__(self, fail_cache=False, model_error=None):
+        super().__init__()
+        self.config = SimpleNamespace(causal_condition=True)
+        self.transformer_blocks = nn.ModuleList([nn.Identity()])
+        self.calls = []
+        self.fail_cache = fail_cache
+        self.model_error = model_error
+
+    def forward(self, hidden_states, encoder_hidden_states, kv_cache_mode, kv_cache, **_kwargs):
+        self.calls.append(SimpleNamespace(mode=kv_cache_mode, cache=kv_cache))
+        condition = encoder_hidden_states[:, -1:, :]
+        if kv_cache_mode == "extract":
+            prefix = condition.unsqueeze(2)
+            kv_cache.get_layer(0).store(prefix, prefix)
+            kv_cache.prefix_len = 1
+            if self.fail_cache:
+                self.fail_cache = False
+                raise QwenImage21KVCacheAllocationError("simulated host allocation failure")
+        elif kv_cache_mode == "cached":
+            condition = kv_cache.get_layer(0).get()[1].squeeze(2)
+        if self.model_error is not None:
+            raise self.model_error
+        return (condition.expand_as(hidden_states),)
+
+
+def _run(batch, args, transformer=None):
+    transformer = transformer if transformer is not None else TinyTransformer()
+    vae, processor, scheduler = TinyVAE(), TinyProcessor(), TinyScheduler()
+    pipeline = [
+        stages.QwenImage21InputStage(),
+        stages.QwenImage21TextEncodingStage(TinyEncoder(), processor),
+        stages.QwenImage21LatentStage(vae),
+        stages.QwenImage21ScheduleStage(scheduler),
+        stages.QwenImage21DenoisingStage(transformer, scheduler),
+        stages.QwenImage21DecodingStage(vae),
+    ]
+    for stage in pipeline:
+        batch = stage.forward(batch, args)
+    return batch, transformer, vae, processor
+
+
+def _assert_caches_cleared(transformer):
+    caches = [call.cache for call in transformer.calls if call.cache is not None]
+    assert caches
+    for cache in caches:
+        assert cache.prefix_len is None
+        assert all(layer.k is None and layer.v is None for layer in cache.layer_caches)
+
+
+def test_request_validation_accepts_real_forward_batch(stage_args):
+    batch = _batch()
+    result = stages.QwenImage21InputStage().forward(batch, stage_args)
+    assert result is batch
+    assert result.batch_size == 1
+    assert result.extra[STATE_KEY].img_shapes == [[(1, 4, 4)]]
+    assert not result.do_classifier_free_guidance
+    for updates, message in [
+        ({"num_frames": 2}, "num_frames"),
+        ({"guidance_scale": 3}, "true_cfg_scale"),
+        ({"height": 63}, "multiple of 32"),
+        ({"true_cfg_scale": float("nan")}, "finite"),
+        ({"num_videos_per_prompt": 2}, "num_videos_per_prompt"),
+    ]:
+        with pytest.raises(ValueError, match=message):
+            stages.QwenImage21InputStage().forward(_batch(**updates), stage_args)
+    for resolution in (0, 33):
+        stage_args.pipeline_config.reference_resolution = resolution
+        with pytest.raises(ValueError, match="reference_resolution"):
+            stages.QwenImage21InputStage().forward(_batch(), stage_args)
+
+
+def test_ordered_references_preserve_rgba_for_vae_and_composite_for_encoder(stage_args):
+    references = [Image.new("RGBA", (32, 32), color) for color in [(0, 0, 0, 128), (128, 0, 0, 64)]]
+    batch, transformer, vae, processor = _run(_batch(references=references), stage_args)
+    assert batch.extra[STATE_KEY].img_shapes == [[(1, 2, 2), (1, 2, 2), (1, 4, 4)]]
+    assert batch.image_latent.shape == (1, 8, 2)
+    expected_red = torch.tensor([0.0, 128.0]) / 127.5 - 1
+    expected_alpha = torch.tensor([128.0, 64.0]) / 127.5 - 1
+    torch.testing.assert_close(batch.image_latent[0, ::4, 0], expected_red)
+    torch.testing.assert_close(batch.image_latent[0, ::4, 1], expected_alpha)
+    assert all(pixels.shape == (1, 4, 1, 32, 32) for pixels in vae.encoded_pixels)
+    for source, encoded in zip(references, processor.images, strict=True):
+        expected = Image.new("RGB", source.size, "white")
+        expected.paste(source, mask=source.getchannel("A"))
+        assert encoded.mode == "RGB"
+        assert encoded.tobytes() == expected.tobytes()
+    assert batch.output.shape == (1, 4, 1, 64, 64)
+    torch.testing.assert_close(batch.output[:, 3], torch.full((1, 1, 64, 64), 0.25))
+    _assert_caches_cleared(transformer)
+
+
+def test_true_cfg_requires_explicit_negative_prompt_and_separate_caches(stage_args):
+    initial = torch.zeros(1, 16, 2)
+    guided, transformer, _, _ = _run(
+        _batch(latents=initial.clone(), true_cfg_scale=2, negative_prompt="bad"), stage_args)
+    assert guided.do_classifier_free_guidance
+    torch.testing.assert_close(guided.latents, torch.full_like(initial, -0.6))
+    assert [call.mode for call in transformer.calls] == ["extract", "extract", "cached", "cached"]
+    positive_cache, negative_cache = transformer.calls[0].cache, transformer.calls[1].cache
+    assert positive_cache is not negative_cache
+    assert transformer.calls[2].cache is positive_cache
+    assert transformer.calls[3].cache is negative_cache
+    _assert_caches_cleared(transformer)
+
+    unguided, transformer, _, _ = _run(_batch(latents=initial.clone(), true_cfg_scale=2), stage_args)
+    assert not unguided.do_classifier_free_guidance
+    assert [call.mode for call in transformer.calls] == ["extract", "cached"]
+    torch.testing.assert_close(unguided.latents, torch.full_like(initial, -0.8))
+    _assert_caches_cleared(transformer)
+
+
+def test_host_cache_allocation_failure_falls_back_to_uncached_execution(stage_args):
+    transformer = TinyTransformer(fail_cache=True)
+    actual, _, _, _ = _run(_batch(), stage_args, transformer)
+    assert [call.mode for call in transformer.calls] == ["extract", None, None]
+    assert all(call.cache is None for call in transformer.calls[1:])
+    _assert_caches_cleared(transformer)
+    stage_args.pipeline_config.use_kv_cache = False
+    expected, _, _, _ = _run(_batch(), stage_args)
+    torch.testing.assert_close(actual.latents, expected.latents, atol=0, rtol=0)
+    torch.testing.assert_close(actual.output, expected.output, atol=0, rtol=0)
+
+
+def test_model_error_propagates_and_clears_request_cache(stage_args):
+    error = RuntimeError("model failure")
+    transformer = TinyTransformer(model_error=error)
+    with pytest.raises(RuntimeError, match="model failure") as raised:
+        _run(_batch(), stage_args, transformer)
+    assert raised.value is error
+    assert [call.mode for call in transformer.calls] == ["extract"]
+    _assert_caches_cleared(transformer)
+
+
+@pytest.mark.parametrize("offload", ["whole", "layerwise"])
+def test_cancellation_stops_remaining_steps_and_releases_cache_and_weights(stage_args, monkeypatch, offload):
+    batch = _batch(num_inference_steps=3, latents=torch.zeros(1, 16, 2))
+    scheduler = TinyScheduler()
+    for stage in (
+        stages.QwenImage21InputStage(),
+        stages.QwenImage21TextEncodingStage(TinyEncoder(), TinyProcessor()),
+        stages.QwenImage21LatentStage(TinyVAE()),
+        stages.QwenImage21ScheduleStage(scheduler),
+    ):
+        batch = stage.forward(batch, stage_args)
+
+    class InterruptingTransformer(TinyTransformer):
+
+        def forward(self, **kwargs):
+            result = super().forward(**kwargs)
+            denoising.interrupt = True
+            return result
+
+    transformer = InterruptingTransformer()
+    denoising = stages.QwenImage21DenoisingStage(transformer, scheduler)
+    stage_args.dit_cpu_offload = offload == "whole"
+    stage_args.dit_layerwise_offload = offload == "layerwise"
+    loaded, unloaded, layerwise_released, steps = [], [], [], []
+    monkeypatch.setattr(stages.pinned_offload, "load", lambda model, device, pin=False: loaded.append(model))
+    monkeypatch.setattr(stages.pinned_offload, "unload", lambda model: unloaded.append(model))
+    monkeypatch.setattr(denoising, "_release_layerwise", lambda: layerwise_released.append(True))
+    original_step = scheduler.step
+
+    def record_step(*args, **kwargs):
+        steps.append(args[1])
+        return original_step(*args, **kwargs)
+
+    monkeypatch.setattr(scheduler, "step", record_step)
+    result = denoising.forward(batch, stage_args)
+
+    assert result is batch
+    assert denoising.interrupt is True  # The executor owns the cancellation latch.
+    assert [call.mode for call in transformer.calls] == ["extract"]
+    assert len(steps) == 1
+    assert batch.step_index == 0
+    torch.testing.assert_close(batch.latents, torch.full_like(batch.latents, -0.4))
+    _assert_caches_cleared(transformer)
+    assert loaded == unloaded == ([transformer] if offload == "whole" else [])
+    assert layerwise_released == ([True] if offload == "layerwise" else [])

fastvideo/tests/train/models/test_load_kandinsky5.py

@@ -13,8 +13,6 @@ and ``Kandinsky5DenoisingStage``.
 
 from __future__ import annotations
 
-import os
-
 import fastvideo.envs as envs
 
 envs.setdefault_external("MASTER_ADDR", "localhost")
@@ -33,11 +31,11 @@ _FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "kandinsky5
 
 
 @pytest.mark.usefixtures("distributed_setup")
-def test_kandinsky5_model_loads_and_forwards(monkeypatch: pytest.MonkeyPatch):
+def test_kandinsky5_model_loads_and_forwards(env_overrides):
     if not torch.cuda.is_available():
         pytest.skip("requires CUDA")
-    if "FASTVIDEO_ATTENTION_BACKEND" not in os.environ:
-        monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+    if not envs.FASTVIDEO_ATTENTION_BACKEND.is_set():
+        env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
 
     cfg = load_run_config(_FIXTURE)
     model = Kandinsky5Model(

fastvideo/tests/transformers/test_qwen_image21_attention.py

@@ -0,0 +1,410 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen-Image-2.1 attention and prefix-cache regressions with tiny native weights.
+
+RoPE checks run on CPU; attention/cache cases require one CUDA GPU and the
+ordinary FastVideo runtime. No model download or optional FlashAttention
+installation is needed. Independent checkpoint parity remains in
+``tests/local_tests/transformers/``.
+"""
+
+from __future__ import annotations
+
+from functools import cache
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.testing import assert_close
+
+import fastvideo.envs as envs
+import fastvideo.models.dits.qwen_image21 as qwen
+from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.platforms import AttentionBackendEnum
+
+@pytest.fixture
+def qwen_setup(env_overrides, monkeypatch, request):
+    if not torch.cuda.is_available():
+        pytest.skip("Requires one CUDA GPU")
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+    monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
+    monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False)
+    request.getfixturevalue("distributed_setup")
+    with torch.no_grad(), sdpa_kernel(SDPBackend.MATH):
+        yield torch.device("cuda")
+
+
+def _initialize_linears(module):
+    # ReplicatedLinear allocates empty storage instead of initializing weights.
+    for layer in module.modules():
+        if isinstance(layer, ReplicatedLinear):
+            torch.nn.init.xavier_uniform_(layer.weight)
+            if layer.bias is not None:
+                torch.nn.init.zeros_(layer.bias)
+    return module.eval()
+
+
+def _model(device):
+    config = QwenImage21Config(arch_config=QwenImage21ArchConfig(
+        in_channels=4, out_channels=4, num_layers=2, num_attention_heads=2,
+        attention_head_dim=8, context_in_dim=12, axes_dims_rope=(2, 2, 4), mlp_ratio=2))
+    return _initialize_linears(qwen.QwenImage21Transformer2DModel(config).to(device=device, dtype=torch.float32))
+
+
+def _inputs(device):
+    # Adjacent reference slots must still form distinct bidirectional blocks.
+    return dict(
+        hidden_states=torch.randn(2, 12, 4, device=device),
+        encoder_hidden_states=torch.randn(2, 5, 12, device=device),
+        timestep=torch.tensor([0.6, 0.7], device=device),
+        img_shapes=[[(1, 2, 2)] * 3 for _ in range(2)],
+        img_mask=torch.tensor([[False, True, True, False, False, True]] * 2, device=device),
+        encoder_hidden_states_mask=torch.tensor(
+            [[True, True, True, True, False], [False, True, True, True, False]], device=device),
+    )
+
+
+@pytest.mark.parametrize("position", [8191, 8192, 9215, 9216])
+def test_rope_positive_position_does_not_alias_negative_coordinates(position):
+    rope = qwen.QwenImage21Rope(theta=10000, axes_dim=(2, 2, 2))
+    mask = torch.tensor([True] + [False] * position)
+    actual = rope([(1, 1, 1)], mask, torch.device("cpu"))
+
+    # Each two-dimensional axis has frequency 1: its phase is exp(i * position).
+    expected = torch.exp(1j * torch.tensor([position] * 3, dtype=torch.float32))
+    assert_close(actual[-1], expected, atol=1e-6, rtol=1e-6)
+    image_expected = torch.exp(1j * torch.tensor([0, -1, -1], dtype=torch.float32))
+    assert_close(actual[0], image_expected, atol=1e-6, rtol=1e-6)
+    assert actual.shape == (position + 1, 3)
+    assert actual.dtype == torch.complex64 and actual.device.type == "cpu"
+
+
+@pytest.mark.parametrize("device_name", ["cpu", "cuda"])
+def test_rope_long_layout_preserves_multi_image_coordinates_and_short_requests(device_name):
+    if device_name == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA RoPE placement requires one GPU")
+    device = torch.device(device_name)
+    rope = qwen.QwenImage21Rope(theta=10000, axes_dim=(2, 2, 2)).to(device=device, dtype=torch.bfloat16)
+    shapes = [(1, 2049, 1), (1, 1, 2049)]
+
+    def layout(text_tokens):
+        mask = torch.tensor([False] * text_tokens + [True] * 2049 + [False] * 3 + [True] * 2049 + [False] * 2,
+                            device=device)
+        # Odd image axes span -1025..1023, beyond the old negative table.
+        # Each image advances subsequent text by its largest spatial axis.
+        positions = ([[p] * 3 for p in range(text_tokens)]
+                     + [[text_tokens, h, -1] for h in range(-1025, 1024)]
+                     + [[text_tokens + p] * 3 for p in (2049, 2050, 2051)]
+                     + [[text_tokens + 2052, -1, w] for w in range(-1025, 1024)]
+                     + [[text_tokens + p] * 3 for p in (4101, 4102)])
+        expected = torch.exp(1j * torch.tensor(positions, dtype=torch.float32, device=device))
+        return mask, expected
+
+    short_mask, short_expected = layout(3)
+    before = rope(shapes, short_mask, device).clone()
+    assert_close(before, short_expected, atol=1e-6, rtol=1e-6)
+    long_mask, long_expected = layout(9216)
+    actual = rope(shapes, long_mask, device)
+    assert_close(actual, long_expected, atol=1e-6, rtol=1e-6)
+    assert actual.device == long_mask.device and actual.dtype == torch.complex64
+    assert_close(rope(shapes, short_mask, device), before, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("device_name", ["cpu", "cuda"])
+def test_rope_long_layout_keeps_nontrivial_short_frequencies_bitwise(device_name):
+    if device_name == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA RoPE placement requires one GPU")
+    device = torch.device(device_name)
+    rope = qwen.QwenImage21Rope(theta=10000, axes_dim=(2, 2, 4)).to(device=device, dtype=torch.bfloat16)
+    shapes = [(1, 2, 2)]
+    short_mask = torch.tensor([False] * 3 + [True] * 4, device=device)
+    before = rope(shapes, short_mask, device).clone()
+    rope(shapes, torch.tensor([False] * 9216 + [True] * 4, device=device), device)
+    after = rope(shapes, short_mask, device)
+    assert after.dtype == torch.complex64 and after.device == short_mask.device
+    assert_close(after, before, atol=0, rtol=0)
+
+
+def test_segmented_attention_matches_independent_dense_mask(qwen_setup):
+    device = qwen_setup
+    attention = _initialize_linears(qwen.QwenImage21Attention(16, 2, 8, 1e-6).to(device))
+    hidden = torch.randn(2, 15, 16, device=device)
+    image_ids = torch.tensor([-1, -1] + [0] * 4 + [1] * 4 + [-1] + [2] * 4, device=device)
+    valid = torch.ones(2, 15, dtype=torch.bool, device=device)
+    valid[0, 0] = False
+    valid[1, 1] = False
+
+    # Independently encode causality: preceding tokens, or any token in the
+    # same image. In particular, adjacent reference images cannot see ahead.
+    positions = torch.arange(15, device=device)
+    causal = positions[:, None] >= positions[None, :]
+    same_image = (image_ids[:, None] == image_ids[None, :]) & (image_ids[:, None] >= 0)
+    mask = (causal | same_image)[None, None] & valid[:, None, None, :]
+    query = attention.norm_q(attention.to_q(hidden)[0].unflatten(-1, (2, 8)))
+    key = attention.norm_k(attention.to_k(hidden)[0].unflatten(-1, (2, 8)))
+    value = attention.to_v(hidden)[0].unflatten(-1, (2, 8))
+    expected = F.scaled_dot_product_attention(
+        query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=mask, dropout_p=0.0)
+    expected = expected.masked_fill(~mask.any(-1, keepdim=True), 0).transpose(1, 2).flatten(2, 3)
+    expected = attention.to_out[0](expected)[0]
+
+    actual = attention(hidden, segments=qwen._qwenimage21_prefix_segments(image_ids, 11), key_valid=valid)
+
+    assert torch.isfinite(actual).all()
+    assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_cached_decode_matches_recompute_and_owns_prefix_storage(qwen_setup):
+    device = qwen_setup
+    model = _model(device)
+    inputs = _inputs(device)
+    cache = qwen.QwenImage21KVCache(model.config.num_layers, storage_device="cpu")
+    extracted = model(**inputs, kv_cache=cache, kv_cache_mode="extract")
+    assert_close(extracted, model(**inputs), atol=0, rtol=0)
+    assert cache.prefix_len == 11
+    prefix_values = []
+    for layer in cache.layer_caches:
+        assert layer.k.shape == layer.v.shape == (2, 11, 2, 8)
+        assert layer.k.device.type == layer.v.device.type == "cpu"
+        for tensor in (layer.k, layer.v):
+            assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size()
+        prefix_values.append((layer.k.clone(), layer.v.clone()))
+
+    inputs["timestep"] = torch.tensor([0.25, 0.35], device=device)
+    inputs["hidden_states"] = inputs["hidden_states"].clone()
+    inputs["hidden_states"][:, -4:] += 0.7
+    actual = model(**inputs, kv_cache=cache, kv_cache_mode="cached")
+    expected = model(**inputs)
+
+    assert actual.shape == (2, 4, 4)
+    assert_close(actual, expected[:, -4:], atol=1e-5, rtol=1e-5)
+    assert_close(extracted[:, :-4], expected[:, :-4], atol=0, rtol=0)
+    for layer, (key, value) in zip(cache.layer_caches, prefix_values, strict=True):
+        assert_close(layer.k, key, atol=0, rtol=0)
+        assert_close(layer.v, value, atol=0, rtol=0)
+
+    # A prefix view must own its bytes, even when source and cache are both on
+    # the CPU. Otherwise retaining it keeps the entire prefill allocation alive.
+    source = torch.randn(1, 15, 2, 8)
+    owned = qwen.QwenImage21KVLayerCache("cpu")
+    owned.store(source[:, :3], source[:, :3])
+    original = source[:, :3].clone()
+    assert owned.k.untyped_storage().data_ptr() != source.untyped_storage().data_ptr()
+    assert owned.k.untyped_storage().nbytes() == owned.k.numel() * owned.k.element_size()
+    source.zero_()
+    assert_close(owned.k, original, atol=0, rtol=0)
+    assert_close(owned.v, original, atol=0, rtol=0)
+    cache.clear()
+    owned.clear()
+    assert cache.prefix_len is None
+    assert all(layer.k is None and layer.v is None for layer in cache.layer_caches)
+    with pytest.raises(RuntimeError, match="extract"):
+        owned.get()
+
+
+def test_padding_embeddings_cannot_change_target_predictions(qwen_setup):
+    device = qwen_setup
+    model = _model(device)
+    inputs = _inputs(device)
+    expected = model(**inputs)[:, -4:]
+    embeddings = inputs["encoder_hidden_states"].clone()
+    embeddings[~inputs["encoder_hidden_states_mask"]] = 1000
+    inputs["encoder_hidden_states"] = embeddings
+    cache = qwen.QwenImage21KVCache(model.config.num_layers, storage_device="cpu")
+
+    extracted = model(**inputs, kv_cache=cache, kv_cache_mode="extract")[:, -4:]
+    cached = model(**inputs, kv_cache=cache, kv_cache_mode="cached")
+
+    assert torch.isfinite(extracted).all()
+    assert torch.isfinite(cached).all()
+    assert_close(extracted, expected, atol=0, rtol=0)
+    assert_close(cached, expected, atol=1e-5, rtol=1e-5)
+    cache.clear()
+
+
+def test_masked_calls_use_sdpa_and_dense_calls_use_selected_backend(qwen_setup, monkeypatch):
+    import fastvideo.attention.selector as selector
+
+    device = qwen_setup
+    dense_calls, resolutions, masked_calls = [], [], []
+
+    class RecordingImpl:
+
+        def __init__(self, **kwargs):
+            assert kwargs["causal"] is False
+            assert kwargs["num_heads"] == kwargs["num_kv_heads"] == 2
+            assert kwargs["head_size"] == 8
+
+        def forward(self, query, key, value, attn_metadata):
+            assert attn_metadata is None
+            dense_calls.append((query.shape[1], key.shape[1]))
+            return F.scaled_dot_product_attention(
+                query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2),
+                dropout_p=0.0).transpose(1, 2)
+
+    class RecordingBackend:
+
+        @staticmethod
+        def get_name():
+            return "FLASH_ATTN"
+
+        @staticmethod
+        def get_impl_cls():
+            return RecordingImpl
+
+    def resolve_backend(head_size, dtype, **kwargs):
+        assert head_size == 8
+        resolutions.append(kwargs)
+        return RecordingBackend
+
+    monkeypatch.setattr(selector, "get_attn_backend", resolve_backend)
+    baseline = _initialize_linears(qwen.QwenImage21Attention(
+        16, 2, 8, 1e-6, requested_backend=AttentionBackendEnum.TORCH_SDPA).to(device))
+    assert resolutions == []
+    attention = qwen.QwenImage21Attention(16, 2, 8, 1e-6).to(device).eval()
+    explicit_attention = qwen.QwenImage21Attention(
+        16, 2, 8, 1e-6, requested_backend=AttentionBackendEnum.FLASH_ATTN).to(device).eval()
+    attention.load_state_dict(baseline.state_dict())
+    explicit_attention.load_state_dict(baseline.state_dict())
+    assert [resolution.get("requested", selector.NO_REQUEST) for resolution in resolutions] == [
+        selector.NO_REQUEST, AttentionBackendEnum.FLASH_ATTN]
+    assert all("default_backend" not in resolution for resolution in resolutions)
+    assert all(AttentionBackendEnum.FLASH_ATTN in resolution["supported_attention_backends"]
+               for resolution in resolutions)
+
+    hidden = torch.randn(1, 11, 16, device=device)
+    image_ids = torch.tensor([-1, -1] + [0] * 4 + [-1] + [1] * 4, device=device)
+    segments = qwen._qwenimage21_prefix_segments(image_ids, 7)
+    valid = torch.ones(1, 11, dtype=torch.bool, device=device)
+    valid[:, 0] = False
+    expected = baseline(hidden, segments=segments)
+    padded_expected = baseline(hidden, segments=segments, key_valid=valid)
+    assert_close(explicit_attention(hidden, segments=segments), expected, atol=1e-5, rtol=1e-5)
+    assert dense_calls == [(4, 6), (4, 11)]
+    dense_calls.clear()
+    sdpa = qwen._attention
+
+    def record_masked_sdpa(query, key, value, mask=None):
+        assert mask is not None
+        masked_calls.append((query.shape[1], key.shape[1]))
+        return sdpa(query, key, value, mask)
+
+    monkeypatch.setattr(qwen, "_attention", record_masked_sdpa)
+    cache = qwen.QwenImage21KVLayerCache()
+    actual = attention(hidden, segments=segments, layer_cache=cache,
+                       kv_cache_mode="extract", cache_write_slice=slice(0, 7))
+    assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+    assert masked_calls == [(2, 2), (1, 7)]
+    assert dense_calls == [(4, 6), (4, 11)]
+
+    dense_calls.clear()
+    masked_calls.clear()
+    cached = attention(hidden[:, 7:], layer_cache=cache, kv_cache_mode="cached")
+    assert_close(cached, expected[:, 7:], atol=1e-5, rtol=1e-5)
+    assert dense_calls == [(4, 11)]
+    assert masked_calls == []
+
+    dense_calls.clear()
+    padded = attention(hidden, segments=segments, key_valid=valid)
+    assert_close(padded, padded_expected, atol=1e-5, rtol=1e-5)
+    assert dense_calls == []
+    assert masked_calls == [(2, 2), (4, 6), (1, 7), (4, 11)]
+    masked_calls.clear()
+    padded_cached = attention(hidden[:, 7:], layer_cache=cache, kv_cache_mode="cached",
+                              attention_mask=valid[:, None, None, :])
+    assert_close(padded_cached, padded_expected[:, 7:], atol=1e-5, rtol=1e-5)
+    assert dense_calls == []
+    assert masked_calls == [(4, 11)]
+    cache.clear()
+
+
+@pytest.fixture
+def selector_resolution(qwen_setup, monkeypatch):
+    import fastvideo.attention.selector as selector
+    import fastvideo.platforms as platforms
+    from fastvideo.attention.backends.sdpa import SDPABackend
+
+    class AlternativeImpl:
+
+        def __init__(self, **kwargs):
+            pass
+
+    class AlternativeBackend:
+
+        @staticmethod
+        def get_name():
+            return "FLASH_ATTN"
+
+        @staticmethod
+        def get_impl_cls():
+            return AlternativeImpl
+
+    def platform_backend(requested, head_size, dtype):
+        # A dense alternative is available automatically, so losing an SDPA
+        # override cannot silently pass by falling back to SDPA anyway.
+        return "TORCH_SDPA" if requested is AttentionBackendEnum.TORCH_SDPA else "FLASH_ATTN"
+
+    monkeypatch.setattr(platforms.current_platform, "get_attn_backend_cls", platform_backend)
+    monkeypatch.setattr(selector, "resolve_obj_by_qualname",
+                        lambda name: SDPABackend if name == "TORCH_SDPA" else AlternativeBackend)
+    # Exercise the real selector and its cache without leaving fake platform
+    # resolutions in the shared process cache or clearing unrelated entries.
+    monkeypatch.setattr(selector, "_cached_get_attn_backend", cache(selector._cached_get_attn_backend.__wrapped__))
+
+
+def _direct_attention(kind, recorded_backend=None):
+    if kind == "block":
+        return qwen.QwenImage21TransformerBlock(16, 2, 8, 2, 1e-6).attn
+    config = QwenImage21Config(arch_config=QwenImage21ArchConfig(
+        in_channels=4, out_channels=4, num_layers=1, num_attention_heads=2,
+        attention_head_dim=8, context_in_dim=12, axes_dims_rope=(2, 2, 4), mlp_ratio=2))
+    config._resolved_attention_backend = recorded_backend
+    return qwen.QwenImage21Transformer2DModel(config).transformer_blocks[0].attn
+
+
+def _selected_backend(attention):
+    return (AttentionBackendEnum.TORCH_SDPA if attention.dense_attn is None else attention.dense_attn.backend)
+
+
+@pytest.mark.parametrize("kind", ["transformer", "block"])
+@pytest.mark.parametrize("environment", ["TORCH_SDPA", "FLASH_ATTN", None])
+def test_direct_construction_honors_environment_backend(selector_resolution, kind, environment):
+    with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
+        attention = _direct_attention(kind)
+    expected = AttentionBackendEnum.TORCH_SDPA if environment == "TORCH_SDPA" else AttentionBackendEnum.FLASH_ATTN
+    assert _selected_backend(attention) is expected
+
+
+@pytest.mark.parametrize("recorded", [AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN])
+def test_direct_transformer_recorded_backend_overrides_environment(selector_resolution, recorded):
+    from fastvideo.attention.selector import _component_attention_backend_scope
+
+    environment = "FLASH_ATTN" if recorded is AttentionBackendEnum.TORCH_SDPA else "TORCH_SDPA"
+    with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
+        with _component_attention_backend_scope(environment, component="qwen"):
+            attention = _direct_attention("transformer", recorded)
+    assert _selected_backend(attention) is recorded
+
+
+@pytest.mark.parametrize("requested", [AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN])
+def test_direct_block_explicit_backend_overrides_environment(selector_resolution, requested):
+    environment = "FLASH_ATTN" if requested is AttentionBackendEnum.TORCH_SDPA else "TORCH_SDPA"
+    with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
+        block = qwen.QwenImage21TransformerBlock(16, 2, 8, 2, 1e-6, requested_backend=requested)
+    assert _selected_backend(block.attn) is requested
+
+
+@pytest.mark.parametrize("kind", ["transformer", "block"])
+@pytest.mark.parametrize("scoped,consult_env,environment,expected", [
+    (AttentionBackendEnum.TORCH_SDPA, False, "FLASH_ATTN", AttentionBackendEnum.TORCH_SDPA),
+    (None, False, "TORCH_SDPA", AttentionBackendEnum.FLASH_ATTN),
+    (None, True, "TORCH_SDPA", AttentionBackendEnum.TORCH_SDPA),
+])
+def test_direct_construction_honors_component_scope(selector_resolution, kind, scoped, consult_env, environment, expected):
+    from fastvideo.attention.selector import _component_attention_backend_scope
+
+    with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
+        with _component_attention_backend_scope(scoped, component="qwen", consult_env=consult_env):
+            attention = _direct_attention(kind)
+    assert _selected_backend(attention) is expected

pyproject.toml

@@ -305,7 +305,8 @@ skip = "./data,./wandb,apps/fastvideo_studio/package-lock.json,apps/performance_
 # lowercase form silences all case variants.
 # "padd" is the literal value of Wan2.2-S2V's ``framepack_drop_mode`` in the
 # published checkpoint config; it must be spelled exactly to match.
-ignore-words-list = "tread,passt,dout,padd"
+# "nax" is the name of MLX's M5 TensorOps helper and the corresponding H3 backend.
+ignore-words-list = "tread,passt,dout,padd,nax"
 
 [tool.ruff]
 # Allow lines to be as long as 120.

tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Released-weight Qwen-Image-2.1 conditioner parity through the production loader.
+
+Run with QWEN_IMAGE21_RUN_ENCODER_PARITY=1 and QWEN_IMAGE21_MODEL_ROOT set to
+the complete local HF snapshot. This gate requires a BF16 CUDA GPU; opting in
+turns missing assets or dependencies into failures rather than skipped parity.
+Set QWEN_IMAGE21_ENCODER_PARITY_REPORT to save the numerical diagnostics as JSON.
+"""
+
+import gc
+import json
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from torch.testing import assert_close
+
+PARITY_SCOPE = "production_loader"
+ATOL = 1e-2
+RTOL = 1e-2
+
+
+def _numerical_metrics(actual, expected):
+    actual, expected = actual.float(), expected.float()
+    difference = actual - expected
+    reference_rms = expected.square().mean().sqrt()
+    rmse = difference.square().mean().sqrt()
+    actual_flat, expected_flat = actual.flatten().double(), expected.flatten().double()
+    cosine = torch.nn.functional.cosine_similarity(actual_flat, expected_flat, dim=0).clamp(-1, 1)
+    return {
+        "shape": list(actual.shape),
+        "finite": bool(torch.isfinite(actual).all() and torch.isfinite(expected).all()),
+        "max_abs": difference.abs().max().item(),
+        "mean_abs": difference.abs().mean().item(),
+        "rmse": rmse.item(),
+        "reference_rms": reference_rms.item(),
+        "relative_l2": (rmse / reference_rms.clamp_min(torch.finfo(torch.float32).tiny)).item(),
+        "cosine": cosine.item(),
+        "mismatched_elements": int((difference.abs() > ATOL + RTOL * expected.abs()).sum()),
+        "total_elements": actual.numel(),
+    }
+
+
+def _make_cases(processor):
+    prompt = "<|im_start|>user\nA red fox on fresh snow.<|im_end|>\n<|im_start|>assistant\n"
+    cases = {"text": processor(text=[prompt], padding=True, padding_side="left", return_tensors="pt")}
+    long_prompt = "<|im_start|>user\n" + "Small red fox on fresh snow. " * 220 + "<|im_end|>\n"
+    cases["long_text"] = processor(text=[long_prompt], return_tensors="pt")
+    assert cases["long_text"]["input_ids"].shape[-1] > 1024
+    for count in (1, 2, 10):
+        images = []
+        for index in range(count):
+            pixels = np.arange(64 * 96 * 3, dtype=np.uint32).reshape(64, 96, 3)
+            images.append(Image.fromarray(((pixels * (index + 1) + index * 31) % 256).astype(np.uint8)))
+        image_prompt = "<|im_start|>user\n" + " ".join(
+            f"<image{index + 1}><|vision_start|><|image_pad|><|vision_end|>" for index in range(count))
+        image_prompt += " Blend these references into a snowy scene.<|im_end|>\n"
+        cases[f"images_{count}"] = processor(
+            text=[image_prompt],
+            images=images,
+            min_pixels=4096,
+            max_pixels=16384,
+            return_tensors="pt",
+        )
+    cases["left_padding"] = processor(text=[prompt, prompt + " A blue sky."], padding=True,
+                                       padding_side="left", return_tensors="pt")
+    return cases
+
+
+def test_qwen_image21_conditioner_real_weight_parity():
+    if os.environ.get("QWEN_IMAGE21_RUN_ENCODER_PARITY") != "1":
+        pytest.skip("set QWEN_IMAGE21_RUN_ENCODER_PARITY=1 on an allocated CUDA GPU")
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        pytest.fail("Qwen-Image-2.1 conditioner parity requires a BF16 CUDA GPU", pytrace=False)
+    checkpoint = os.environ.get("QWEN_IMAGE21_MODEL_ROOT")
+    if not checkpoint:
+        pytest.fail("set QWEN_IMAGE21_MODEL_ROOT to the downloaded Qwen-Image-2.1 HF snapshot", pytrace=False)
+    root = Path(checkpoint)
+    if not (root / "text_encoder" / "config.json").is_file() or not (root / "processor").is_dir():
+        pytest.fail("The checkpoint must include text_encoder and processor component directories", pytrace=False)
+
+    from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor, __version__ as transformers_version
+
+    from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+    from fastvideo.models.loader.component_loader import TextEncoderLoader
+
+    device = torch.device("cuda", 0)
+    processor = Qwen3VLProcessor.from_pretrained(root / "processor", local_files_only=True)
+    cases = _make_cases(processor)
+    official, loading_info = Qwen3VLForConditionalGeneration.from_pretrained(
+        root / "text_encoder",
+        dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+        local_files_only=True,
+        low_cpu_mem_usage=True,
+        output_loading_info=True,
+    )
+    errors = {name: loading_info[name] for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
+              if loading_info.get(name)}
+    assert not errors, f"Official checkpoint strict loading failed: {errors}"
+    official = official.to(device).eval()
+    norm = official.model.language_model.norm
+    handle = norm.register_forward_hook(lambda _module, args, _output: args[0])
+    expected = {}
+    try:
+        for name, case in cases.items():
+            inputs = {key: value.to(device=device, dtype=torch.bfloat16 if key == "pixel_values" else value.dtype)
+                      for key, value in case.items()}
+            with torch.no_grad():
+                output = official(**inputs, output_hidden_states=True, use_cache=False)
+            expected[name] = output.hidden_states[-1].detach().cpu()
+            del output, inputs
+    finally:
+        handle.remove()
+    del official, norm
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29625")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    config = QwenImage21Qwen3VLConfig()
+    config.update_model_arch(json.loads((root / "text_encoder" / "config.json").read_text()))
+    args = SimpleNamespace(
+        text_encoder_cpu_offload=True,
+        override_text_encoder_quant=None,
+        override_text_encoder_safetensors=None,
+        pin_cpu_memory=False,
+        disable_offload_on_unified_memory=lambda _device_id, offload_flag=None: False,
+    )
+    report = {
+        "component": "qwen_image21_conditioner",
+        "checkpoint_revision": root.name,
+        "torch_version": torch.__version__,
+        "transformers_version": transformers_version,
+        "atol": ATOL,
+        "rtol": RTOL,
+        "cases": {},
+    }
+    failures = []
+    try:
+        native = TextEncoderLoader().load_model(str(root / "text_encoder"), config, device, args, dtype="bf16")
+        for name, case in cases.items():
+            inputs = {key: value.to(device=device, dtype=torch.bfloat16 if key == "pixel_values" else value.dtype)
+                      for key, value in case.items()}
+            with torch.no_grad():
+                actual = native(**inputs).detach().cpu()
+            mask = case["attention_mask"].bool()
+            metrics = _numerical_metrics(actual[mask], expected[name][mask])
+            report["cases"][name] = metrics
+            print(f"QWEN_IMAGE21_ENCODER_PARITY {name} {json.dumps(metrics, sort_keys=True)}", flush=True)
+            try:
+                assert_close(actual[mask], expected[name][mask], rtol=RTOL, atol=ATOL,
+                             msg=lambda message: f"{name}: {message}")
+            except AssertionError as error:
+                failures.append(f"{name}: {error}")
+            del actual, inputs
+    finally:
+        report_path = os.environ.get("QWEN_IMAGE21_ENCODER_PARITY_REPORT")
+        if report_path:
+            destination = Path(report_path)
+            destination.parent.mkdir(parents=True, exist_ok=True)
+            destination.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
+        cleanup_dist_env_and_memory()
+    assert not failures, "\n\n".join(failures)

tests/local_tests/qwen_image21/README.md

@@ -0,0 +1,138 @@
+# Qwen-Image-2.1 validation
+
+Use the package tests for regression coverage and the opt-in local tests to
+compare FastVideo with pinned upstream Diffusers. For generation commands and
+memory measurements, see the [cookbook](../../../docs/cookbook/qwen-image21.md).
+
+## Confirmed coverage
+
+Validation on NVIDIA GB200 at `c3d59459` completed on 2026-10-23: **36 cases
+passed, with no failures or skips**.
+
+| Check | Cases | Result |
+| --- | ---: | --- |
+| Package attention and stage tests | 9 | Backend selection, masks, caches and pipeline stages pass |
+| FA4 transformer-block golden | 1 | Zero drift against the reviewed device/runtime reference |
+| FA4 T2I and editing image regressions | 2 | SSIM 1.0; unchanged threshold 0.98 |
+| Encoder, DiT and VAE component parity | 12 | All component cases pass |
+| SDPA upstream pipeline parity | 12 | All modes below pass at 40 steps |
+
+The pipeline matrix covers T2I, plain I2I, editing, reference-to-image,
+transparent RGBA generation and editing, annotations, mask images, true CFG,
+uncached execution and ten references at 1024 x 1024. A separate T2I case covers
+2048 x 2048. Conditioning, initial/final latents, all 40 noise predictions and
+latent updates, floating-point RGBA output and quantized PNG pixels were exact.
+Raw FP32 scheduler arrays have maximum differences of `1.1920928955078125e-7`
+for sigmas and `0.0001220703125` for timesteps; they are not bitwise equal.
+
+Upstream comparisons explicitly use SDPA with matching cuDNN SDPA settings.
+FA4 is validated separately against its golden and reviewed image references;
+FA4 and SDPA outputs are not expected to be bitwise equal. The automatic
+attention selector was also observed selecting FA4 for all 64 blocks across the
+two image regressions with `FASTVIDEO_FA4=1`.
+
+These results establish the tested GB200 configurations. They do not establish
+other GPU/backend combinations, 2K with ten references, or hard preservation of
+pixels outside a semantic mask.
+
+## Package regression tests
+
+Install FastVideo using the [CUDA installation guide](../../../docs/getting_started/installation/gpu.md)
+and run commands from the repository root on an allocated CUDA GPU:
+
+```bash
+python -m pytest fastvideo/tests/transformers/test_qwen_image21_attention.py \
+  fastvideo/tests/stages/test_qwen_image21_stages.py -xq
+
+FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python -m pytest \
+  fastvideo/tests/golden_gate/test_qwen_image21.py -xq
+
+FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python -m pytest \
+  fastvideo/tests/ssim/test_qwen_image21_similarity.py -xq
+```
+
+These files cover attention, pipeline stages, transformer blocks and generated
+images through the existing package CI lanes. The golden loads one released
+BF16 transformer block. The image tests generate T2I and editing outputs at
+512 x 512, four steps and seed 21, comparing all four RGBA channels.
+
+The shared harnesses download missing references from
+[`FastVideo/ssim-reference-videos`](https://huggingface.co/datasets/FastVideo/ssim-reference-videos).
+Qwen references are available for the GB200 FA4 golden and default image tier.
+The optional `--ssim-full-quality` tier uses 1024 x 1024 and 40 steps and needs
+its own reviewed references; those are not covered by the published default
+tier. See the [SSIM guide](../../../fastvideo/tests/ssim/README.md) for reference
+management.
+
+## Prepare upstream parity
+
+Local parity tests under `tests/local_tests/` are opt-in and run outside CI.
+Use BF16 on a CUDA GPU, a complete Qwen checkpoint and a clean checkout of the
+pinned upstream revision. The pipeline harness uses CPU offload, VAE tiling
+and CPU prefix KV storage; its pytest duration is not a generation benchmark.
+Use the cookbook's measured memory profiles to plan resources.
+
+```bash
+git clone https://github.com/huggingface/diffusers.git ../diffusers-qwen-image21
+git -C ../diffusers-qwen-image21 checkout 578c9b2c6636ab2424a0e56186268b83623656b2
+hf download Qwen/Qwen-Image-2.1 \
+  --revision d26bb61231c349cf6b7896fa83353113880e1ba3 \
+  --local-dir ../Qwen-Image-2.1
+
+export QWEN_IMAGE21_DIFFUSERS_DIR="$PWD/../diffusers-qwen-image21"
+export QWEN_IMAGE21_MODEL_ROOT="$PWD/../Qwen-Image-2.1"
+export PYTHONPATH="$QWEN_IMAGE21_DIFFUSERS_DIR/src${PYTHONPATH:+:$PYTHONPATH}"
+export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+```
+
+An existing checkpoint can be used by changing `QWEN_IMAGE21_MODEL_ROOT`. It
+must include `model_index.json`, `processor`, `text_encoder`, `transformer`,
+`vae`, `scheduler` and all weight shards. No checkpoint conversion is needed.
+Explicitly enabled real-weight tests fail if prerequisites are missing.
+
+### Component parity
+
+```bash
+QWEN_IMAGE21_RUN_ENCODER_PARITY=1 python -m pytest \
+  tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py -xq
+
+QWEN_IMAGE21_RUN_DIT_PARITY=1 python -m pytest \
+  tests/local_tests/transformers/test_qwen_image21_transformer_parity.py -xq
+
+QWEN_IMAGE21_RUN_VAE_PARITY=1 python -m pytest \
+  tests/local_tests/vaes/test_qwen_image21_vae_parity.py -xq
+```
+
+These checks cover text and 1/2/10-image conditioning, every DiT block and
+cached decoding, and VAE posterior/decode behavior including tiling. Optional
+`QWEN_IMAGE21_ENCODER_PARITY_REPORT` and `QWEN_IMAGE21_DIT_PARITY_REPORT` paths
+save component diagnostics as JSON.
+
+### Pipeline parity
+
+Run all eleven 1024-pixel modes and the separate 2K T2I case at 40 steps:
+
+```bash
+QWEN_IMAGE21_RUN_PIPELINE_PARITY=1 QWEN_IMAGE21_RUN_2K_PARITY=1 \
+  QWEN_IMAGE21_PARITY_STEPS=40 \
+  python -m pytest tests/local_tests/qwen_image21/test_pipeline_parity.py -xq \
+  --basetemp outputs/qwen_image21/parity-40step
+```
+
+Choose a fresh `--basetemp` path for each run: pytest replaces an existing
+directory at that path. For one mode, select its node, for example
+`tests/local_tests/qwen_image21/test_pipeline_parity.py::test_qwen_image21_modes_1024_pipeline_parity[edit]`.
+Without `QWEN_IMAGE21_PARITY_STEPS`, the harness defaults to four steps.
+
+Each case runs native and upstream pipelines sequentially in separate
+processes, with identical BF16 initial latents, sigmas and cache settings.
+Synthetic reference images are generated by default.
+`QWEN_IMAGE21_PARITY_CASES_JSON` can supply per-case prompts and ordered image
+paths; relative paths resolve beside that JSON file. For parallel invocations,
+assign distinct `CUDA_VISIBLE_DEVICES`, `MASTER_PORT` and artifact paths.
+
+The tests assert RGBA output parity within the existing BF16 tolerance and
+verify PNG alpha export. They also write final-output error metrics. The stronger
+exact-trajectory results reported above came from separate diagnostic runs;
+a passing pixel assertion alone does not establish bitwise equality at every
+stage.

tests/local_tests/qwen_image21/test_pipeline_parity.py

@@ -0,0 +1,235 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Opt-in, released-weight pipeline parity through the public FastVideo API.
+
+Each implementation runs in its own process so their weights never overlap in
+VRAM. Synthetic references are the default; QWEN_IMAGE21_PARITY_CASES_JSON can
+provide per-case prompts and ordered reference paths from the original source.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image, ImageDraw, ImageOps
+
+ROOT = Path(__file__).resolve().parents[3]
+REFERENCE_REVISION = "578c9b2c6636ab2424a0e56186268b83623656b2"
+PARITY_SCOPE = "pipeline"
+CASES = ("t2i", "i2i", "edit", "ref2img", "ten_refs", "rgba", "rgba_edit", "annotation", "mask_image", "true_cfg", "uncached")
+
+
+def _worker(kind: str, spec_path: str, output_path: str) -> None:
+    # Native CUDA startup disables cuDNN SDPA. Match that policy in the isolated
+    # reference process; its default can change encoder numerics on GB200.
+    torch.backends.cuda.enable_cudnn_sdp(False)
+    spec = json.loads(Path(spec_path).read_text())
+    sys.path.insert(0, str(ROOT))
+    sys.path.insert(0, str(Path(spec["reference"]) / "src"))
+    height, width = spec.get("height", spec["resolution"]), spec.get("width", spec["resolution"])
+    seed = spec.get("seed", 21)
+    reference_resolution = spec.get("reference_resolution", 1024)
+    sigmas = spec.get("sigmas", np.linspace(1.0, 1 / spec["steps"], spec["steps"]).tolist())
+    noise = torch.load(spec["latents"], map_location="cpu", weights_only=True)
+    assert noise.dtype == torch.bfloat16
+    assert noise.shape == (1, (height // 16) * (width // 16), 64)
+    torch.manual_seed(seed)
+    if kind == "native":
+        from fastvideo import VideoGenerator
+        from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig,
+                                   OutputConfig, ParallelismConfig, PipelineSelection, SamplingConfig)
+
+        generator = VideoGenerator.from_config(
+            GeneratorConfig(
+                model_path=spec["model"],
+                engine=EngineConfig(num_gpus=1, parallelism=ParallelismConfig(tp_size=1, sp_size=1),
+                                    offload=OffloadConfig(dit=True, dit_layerwise=True, text_encoder=True,
+                                                          vae=True, lazy_module_load=True)),
+                pipeline=PipelineSelection(workload_type="i2i" if spec["refs"] else "t2i",
+                                           experimental={"attention_backend": "TORCH_SDPA",
+                                                         "reference_resolution": reference_resolution,
+                                                         "use_kv_cache": spec["cache"]}),
+            ))
+        try:
+            result = generator.generate(
+                GenerationRequest(
+                    prompt=spec["prompt"], negative_prompt=spec["negative_prompt"],
+                    inputs=InputConfig(references=spec["refs"] or None, latents=noise.clone()),
+                    sampling=SamplingConfig(height=height, width=width, num_frames=1,
+                                            fps=1, num_inference_steps=spec["steps"], guidance_scale=1,
+                                            true_cfg_scale=spec["true_cfg_scale"], seed=seed, sigmas=sigmas),
+                    output=OutputConfig(output_path=str(Path(output_path).with_suffix(".png")),
+                                        return_frames=True, save_video=True),
+                ))
+            torch.save({"pixels": result.samples.detach().float().cpu(), "png_path": result.video_path,
+                        "peak_memory_mb": result.peak_memory_mb}, output_path)
+        finally:
+            generator.shutdown()
+    else:
+        from diffusers.pipelines.qwenimage21 import pipeline_qwenimage21
+
+        expected = Path(spec["reference"]) / "src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py"
+        assert Path(pipeline_qwenimage21.__file__).resolve() == expected.resolve()
+        pipe = pipeline_qwenimage21.QwenImage21Pipeline.from_pretrained(
+            spec["model"], torch_dtype=torch.bfloat16, local_files_only=True)
+        pipe.vae.enable_tiling()
+        pipe.enable_sequential_cpu_offload()
+        references = []
+        for path in spec["refs"]:
+            with Image.open(path) as image:
+                references.append(ImageOps.exif_transpose(image).convert("RGBA"))
+        with torch.no_grad():
+            output = pipe(prompt=spec["prompt"], image=references or None, height=height, width=width,
+                          output_resolution=reference_resolution, num_inference_steps=spec["steps"], sigmas=sigmas,
+                          negative_prompt=spec["negative_prompt"], true_cfg_scale=spec["true_cfg_scale"],
+                          generator=torch.Generator(device="cpu").manual_seed(seed),
+                          latents=noise.clone(), output_type="pt", use_kv_cache=spec["cache"])
+        torch.save({"pixels": output.images.unsqueeze(2).float().cpu()}, output_path)
+
+
+def _assets(tmp_path: Path) -> list[str]:
+    files = []
+    for index in range(10):
+        image = Image.new("RGBA", (256, 256), (230, 235, 240, 0 if index == 0 else 255))
+        draw = ImageDraw.Draw(image)
+        draw.ellipse((48 + index, 48, 208, 208), fill=(30 + index * 20, 80, 140, 200))
+        if index == 1:
+            draw.rectangle((24, 24, 232, 232), outline="red", width=8)
+        path = tmp_path / f"reference_{index + 1}.png"
+        image.save(path)
+        files.append(str(path))
+    mask = Image.new("RGBA", (256, 256), "black")
+    ImageDraw.Draw(mask).ellipse((48, 48, 208, 208), fill="white")
+    path = tmp_path / "mask.png"
+    mask.save(path)
+    return files + [str(path)]
+
+
+def _case_overrides(case: str) -> dict:
+    path = os.environ.get("QWEN_IMAGE21_PARITY_CASES_JSON")
+    if not path:
+        return {}
+    manifest = Path(path).resolve()
+    overrides = json.loads(manifest.read_text()).get(case, {})
+    allowed = {"prompt", "refs", "height", "width", "steps", "seed", "reference_resolution", "negative_prompt",
+               "true_cfg_scale", "cache", "sigmas"}
+    assert isinstance(overrides, dict) and not (set(overrides) - allowed)
+    if "refs" in overrides:
+        assert isinstance(overrides["refs"], list) and len(overrides["refs"]) <= 10
+        paths = [Path(value) for value in overrides["refs"]]
+        paths = [value if value.is_absolute() else manifest.parent / value for value in paths]
+        assert all(value.is_file() for value in paths), "Manifest references must be existing local files"
+        overrides["refs"] = [str(value.resolve()) for value in paths]
+    return overrides
+
+
+def _difference_metrics(actual: torch.Tensor, expected: torch.Tensor, *, atol=0.04, rtol=0.02) -> dict[str, float]:
+    actual, expected = actual.float(), expected.float()
+    error = (actual - expected).abs()
+    mean_error = error.mean().item()
+    reference_mean = expected.abs().mean().item()
+    rmse = error.square().mean().sqrt().item()
+    return {
+        "abs_mean": mean_error,
+        "abs_max": error.max().item(),
+        "rmse": rmse,
+        "relative_l1": mean_error / max(reference_mean, 1e-12),
+        "relative_l2": rmse / max(expected.square().mean().sqrt().item(), 1e-12),
+        "actual_abs_mean": actual.abs().mean().item(),
+        "reference_abs_mean": reference_mean,
+        "fraction_outside_tolerance": (error > atol + rtol * expected.abs()).float().mean().item(),
+    }
+
+
+def _run_parity(case, resolution, tmp_path):
+    if os.environ.get("QWEN_IMAGE21_RUN_PIPELINE_PARITY") != "1":
+        pytest.skip("Set QWEN_IMAGE21_RUN_PIPELINE_PARITY=1 on an allocated CUDA node")
+    if not torch.cuda.is_available():
+        pytest.fail("Qwen-Image-2.1 pipeline parity requires CUDA", pytrace=False)
+    model = Path(os.environ.get("QWEN_IMAGE21_MODEL_ROOT", ROOT / "official_weights/Qwen-Image-2.1"))
+    reference = Path(os.environ.get("QWEN_IMAGE21_DIFFUSERS_DIR", ROOT / "official_reference/diffusers"))
+    if not (model / "model_index.json").is_file() or not (reference / "src").is_dir():
+        pytest.fail("Provide QWEN_IMAGE21_MODEL_ROOT and the pinned QWEN_IMAGE21_DIFFUSERS_DIR", pytrace=False)
+    revision = subprocess.check_output(["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip()
+    assert revision == REFERENCE_REVISION
+    subprocess.run(["git", "-C", str(reference), "diff", "--exit-code", "HEAD", "--", "src/diffusers"], check=True)
+    files = _assets(tmp_path)
+    counts = dict(i2i=1, edit=1, ref2img=2, ten_refs=10, rgba_edit=1, annotation=2, mask_image=1, true_cfg=1)
+    refs = files[:counts.get(case, 0)]
+    prompt = "An elegant blue ceramic vase in a studio, with the word Qwen written on it."
+    if refs:
+        prompt = "Make a blue sculpture inspired by the shapes in the reference images."
+    if case == "edit":
+        prompt = "Change the object in image 1 into a blue ceramic sphere; retain the background and composition."
+    if case in ("rgba", "rgba_edit"):
+        prompt = "A blue sphere isolated on a transparent background."
+    if case == "annotation":
+        prompt = "Change the object enclosed by the red annotation in image 2 into a blue sphere in image 1."
+    if case == "mask_image":
+        refs.append(files[-1])
+        prompt = "Replace the white masked area of image 2 in image 1 with flowers."
+    noise_path = tmp_path / "initial_latents.pt"
+    spec = dict(model=str(model.resolve()), reference=str(reference.resolve()), prompt=prompt, refs=refs,
+                resolution=resolution, steps=int(os.environ.get("QWEN_IMAGE21_PARITY_STEPS", "4")),
+                negative_prompt="" if case == "true_cfg" else None,
+                true_cfg_scale=2 if case == "true_cfg" else 1, cache=case != "uncached", latents=str(noise_path),
+                seed=21, reference_resolution=1024, height=resolution, width=resolution)
+    spec.update(_case_overrides(case))
+    height, width = spec["height"], spec["width"]
+    assert min(height, width, spec["reference_resolution"]) > 0
+    assert all(value % 32 == 0 for value in (height, width, spec["reference_resolution"]))
+    assert spec["steps"] > 0
+    spec.setdefault("sigmas", np.linspace(1.0, 1 / spec["steps"], spec["steps"]).tolist())
+    torch.save(torch.randn(1, (height // 16) * (width // 16), 64,
+                           generator=torch.Generator(device="cpu").manual_seed(spec["seed"]),
+                           dtype=torch.bfloat16), noise_path)
+    spec_path = tmp_path / "spec.json"
+    spec_path.write_text(json.dumps(spec))
+    outputs = []
+    for kind in ("native", "reference"):
+        output = tmp_path / f"{kind}.pt"
+        subprocess.run([sys.executable, str(Path(__file__).resolve()), "--worker", kind, str(spec_path), str(output)],
+                       check=True, cwd=ROOT)
+        outputs.append(torch.load(output, map_location="cpu", weights_only=True))
+    actual, expected = outputs[0]["pixels"], outputs[1]["pixels"]
+    assert actual.shape == expected.shape == (1, 4, 1, height, width)
+    assert torch.isfinite(actual).all() and bool(((actual >= 0) & (actual <= 1)).all())
+    assert torch.isfinite(expected).all() and bool(((expected >= 0) & (expected <= 1)).all())
+    metrics = {"case": case, "steps": spec["steps"], "cache": spec["cache"], "seed": spec["seed"],
+               "rgb": _difference_metrics(actual[:, :3], expected[:, :3]),
+               "all_channels": _difference_metrics(actual, expected)}
+    (tmp_path / "parity_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    print(json.dumps(metrics), flush=True)
+    # Provisional tolerances: calibrate from measured BF16 CUDA evidence before release.
+    torch.testing.assert_close(actual, expected, atol=0.04, rtol=0.02)
+    native_png = Image.open(outputs[0]["png_path"])
+    assert native_png.mode == "RGBA"
+    expected_u8 = (actual[0, :, 0].permute(1, 2, 0) * 255).clamp(0, 255).to(torch.uint8).numpy()
+    np.testing.assert_array_equal(np.array(native_png), expected_u8)
+    if outputs[0]["peak_memory_mb"] is not None:
+        assert outputs[0]["peak_memory_mb"] <= 24 * 1024
+
+
+@pytest.mark.parametrize("case", CASES)
+def test_qwen_image21_modes_1024_pipeline_parity(case, tmp_path):
+    _run_parity(case, 1024, tmp_path)
+
+
+def test_qwen_image21_t2i_2k_pipeline_parity(tmp_path):
+    if os.environ.get("QWEN_IMAGE21_RUN_2K_PARITY") != "1":
+        pytest.skip("Set QWEN_IMAGE21_RUN_2K_PARITY=1 for the separate 2K gate")
+    if os.environ.get("QWEN_IMAGE21_RUN_PIPELINE_PARITY") != "1":
+        pytest.fail("The 2K gate also requires QWEN_IMAGE21_RUN_PIPELINE_PARITY=1", pytrace=False)
+    _run_parity("t2i", 2048, tmp_path)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 5 or sys.argv[1] != "--worker":
+        raise SystemExit("This helper is launched by pytest")
+    _worker(sys.argv[2], sys.argv[3], sys.argv[4])

tests/local_tests/transformers/test_qwen_image21_transformer_parity.py

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen-Image-2.1 transformer reference and production-loader parity.
+
+Set QWEN_IMAGE21_DIFFUSERS_DIR to Diffusers revision below for random-weight
+CPU parity. Real-weight BF16 parity additionally needs an allocated CUDA GPU,
+QWEN_IMAGE21_MODEL_ROOT, and QWEN_IMAGE21_RUN_DIT_PARITY=1. Explicitly requested
+real-weight tests fail rather than skip when their prerequisites are missing.
+Set QWEN_IMAGE21_DIT_PARITY_REPORT to save block and prediction diagnostics as JSON.
+"""
+
+from __future__ import annotations
+
+import gc
+import importlib
+import json
+import math
+import os
+import subprocess
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+ROOT = Path(__file__).resolve().parents[3]
+REFERENCE_REVISION = "578c9b2c6636ab2424a0e56186268b83623656b2"
+REFERENCE_ROOT = Path(os.environ.get("QWEN_IMAGE21_DIFFUSERS_DIR", ROOT / "official_reference/diffusers"))
+MODEL_ROOT = Path(os.environ.get("QWEN_IMAGE21_MODEL_ROOT", ROOT / "official_weights/Qwen-Image-2.1"))
+RUN_ENV = "QWEN_IMAGE21_RUN_DIT_PARITY"
+PARITY_SCOPE = "both"
+
+
+def _numerical_metrics(actual, expected):
+    actual, expected = actual.detach().float().cpu(), expected.detach().float().cpu()
+    difference = actual - expected
+    reference_rms = expected.square().mean().sqrt()
+    reference_abs_mean = expected.abs().mean()
+    rmse = difference.square().mean().sqrt()
+    mean_abs = difference.abs().mean()
+    floor = torch.finfo(torch.float32).tiny
+    return {
+        "shape": list(actual.shape),
+        "finite": bool(torch.isfinite(actual).all() and torch.isfinite(expected).all()),
+        "max_abs": difference.abs().max().item(),
+        "mean_abs": mean_abs.item(),
+        "rmse": rmse.item(),
+        "reference_rms": reference_rms.item(),
+        "reference_abs_mean": reference_abs_mean.item(),
+        "relative_l2": (rmse / reference_rms.clamp_min(floor)).item(),
+        "relative_abs_mean": (mean_abs / reference_abs_mean.clamp_min(floor)).item(),
+        "total_elements": actual.numel(),
+    }
+
+
+def _modality_summaries(actual, expected, masks):
+    actual, expected = actual.detach().cpu(), expected.detach().cpu()
+    summaries = {"global": _numerical_metrics(actual, expected)}
+    for name, mask in masks.items():
+        assert mask.numel() == actual.shape[1], f"{name} token mask does not match the joint sequence"
+        if bool(mask.any()):
+            summaries[name] = _numerical_metrics(actual[:, mask], expected[:, mask])
+    return summaries
+
+
+def _reference_module(required=False):
+    if not (REFERENCE_ROOT / "src").is_dir():
+        message = f"Set QWEN_IMAGE21_DIFFUSERS_DIR to Diffusers checkout at {REFERENCE_REVISION}"
+        if required:
+            pytest.fail(message, pytrace=False)
+        pytest.skip(message)
+    revision = subprocess.check_output(["git", "-C", str(REFERENCE_ROOT), "rev-parse", "HEAD"], text=True).strip()
+    assert revision == REFERENCE_REVISION, f"Reference revision mismatch: {revision}"
+    sys.path.insert(0, str(REFERENCE_ROOT / "src"))
+    try:
+        module = importlib.import_module("diffusers.models.transformers.transformer_qwenimage21")
+    except ImportError as exc:
+        if required:
+            pytest.fail(f"Qwen-Image-2.1 reference dependencies unavailable: {exc}", pytrace=False)
+        pytest.skip(f"Qwen-Image-2.1 reference dependencies unavailable: {exc}")
+    expected = REFERENCE_ROOT / "src/diffusers/models/transformers"
+    assert Path(module.__file__).resolve().parent == expected.resolve()
+    return module
+
+
+def _inputs(channels=4, text_dim=12, device="cpu", dtype=torch.float32, batch=2):
+    generator = torch.Generator().manual_seed(21)
+    image = torch.randn(batch, 12, channels, generator=generator).to(device=device, dtype=dtype)
+    text = torch.randn(batch, 5, text_dim, generator=generator).to(device=device, dtype=dtype)
+    return dict(hidden_states=image,
+                encoder_hidden_states=text,
+                timestep=torch.full((batch,), 0.65, device=device, dtype=dtype),
+                img_shapes=[[(1, 2, 2), (1, 2, 2), (1, 2, 2)]] * batch,
+                img_mask=torch.tensor([[False, True, True, False, False, True]] * batch, device=device),
+                encoder_hidden_states_mask=torch.tensor([[True, True, True, True, False]] * batch, device=device))
+
+
+@pytest.fixture
+def tiny_pair():
+    reference = _reference_module()
+    from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
+    from fastvideo.models.dits.qwen_image21 import QwenImage21Transformer2DModel
+    from fastvideo.platforms import AttentionBackendEnum
+
+    kwargs = dict(in_channels=4, out_channels=4, num_layers=2, attention_head_dim=8,
+                  num_attention_heads=2, context_in_dim=12, mlp_ratio=3, axes_dims_rope=(2, 2, 4))
+    torch.manual_seed(20)
+    official = reference.QwenImage21Transformer2DModel(**kwargs).eval()
+    config = QwenImage21Config(arch_config=QwenImage21ArchConfig(**kwargs))
+    config._resolved_attention_backend = AttentionBackendEnum.TORCH_SDPA
+    native = QwenImage21Transformer2DModel(config, {}).eval()
+    official_state = official.state_dict()
+    assert {key: value.shape for key, value in native.state_dict().items()} == {
+        key: value.shape for key, value in official_state.items()
+    }
+    native.load_state_dict(official_state, strict=True)
+    return reference, official, native
+
+
+def test_qwen_image21_dit_random_reference_parity(tiny_pair):
+    _, official, native = tiny_pair
+    inputs = _inputs()
+    with torch.inference_mode():
+        expected = official(**inputs, return_dict=False)[0]
+        actual = native(**inputs)
+    assert_close(actual, expected, atol=2e-6, rtol=2e-5)
+
+
+def test_qwen_image21_qk_norm_dtype_reference_parity():
+    reference = _reference_module()
+    from fastvideo.models.dits.qwen_image21 import QwenImage21RMSNorm
+
+    for weight_dtype in (torch.float32, torch.bfloat16, torch.float16):
+        for input_dtype in (torch.float32, torch.bfloat16, torch.float16):
+            official = reference.RMSNorm(8, eps=1e-6).to(weight_dtype)
+            native = QwenImage21RMSNorm(8, eps=1e-6).to(weight_dtype)
+            weight = torch.randn(8, generator=torch.Generator().manual_seed(23)).to(weight_dtype)
+            official.weight.data.copy_(weight)
+            native.weight.data.copy_(weight)
+            hidden = torch.randn(2, 7, 3, 8, generator=torch.Generator().manual_seed(24)).to(input_dtype)
+            assert_close(native(hidden), official(hidden), atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("storage_device", [None, "cpu"])
+def test_qwen_image21_dit_reference_cache_parity(tiny_pair, storage_device):
+    reference, official, native = tiny_pair
+    from fastvideo.models.dits.qwen_image21 import QwenImage21KVCache
+
+    inputs = _inputs()
+    official_cache = reference.QwenImage21KVCache(2)
+    native_cache = QwenImage21KVCache(2, storage_device)
+    with torch.inference_mode():
+        expected = official(**inputs, kv_cache=official_cache, kv_cache_mode="extract", return_dict=False)[0]
+        actual = native(**inputs, kv_cache=native_cache, kv_cache_mode="extract")
+        assert_close(actual, expected, atol=2e-6, rtol=2e-5)
+        for layer_idx in range(2):
+            expected_k, expected_v = official_cache.get_layer(layer_idx).get()
+            actual_k, actual_v = native_cache.get_layer(layer_idx).get()
+            assert_close(actual_k, expected_k, atol=2e-6, rtol=2e-5)
+            assert_close(actual_v, expected_v, atol=2e-6, rtol=2e-5)
+        inputs["timestep"].fill_(0.3)
+        inputs["hidden_states"][:, -4:] *= 0.8
+        expected = official(**inputs, kv_cache=official_cache, kv_cache_mode="cached", return_dict=False)[0]
+        actual = native(**inputs, kv_cache=native_cache, kv_cache_mode="cached")
+    assert_close(actual, expected, atol=2e-6, rtol=2e-5)
+
+
+def test_qwen_image21_dit_production_loader_parity():
+    if os.environ.get(RUN_ENV) != "1":
+        pytest.skip(f"Set {RUN_ENV}=1 on an allocated CUDA node")
+    if not torch.cuda.is_available():
+        pytest.fail("Qwen-Image-2.1 DiT production-loader parity requires CUDA", pytrace=False)
+    reference = _reference_module(required=True)
+    component = MODEL_ROOT / "transformer"
+    if not (component / "config.json").is_file() or not list(component.glob("*.safetensors")):
+        pytest.fail(f"Local transformer checkpoint missing at {component}", pytrace=False)
+
+    from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
+    from fastvideo.models.dits.qwen_image21 import QwenImage21KVCache
+    from fastvideo.models.loader.component_loader import PipelineComponentLoader
+
+    args = SimpleNamespace(pipeline_config=QwenImage21PipelineConfig(), model_paths={},
+                           override_transformer_cls_name=None, hsdp_replicate_dim=1, hsdp_shard_dim=1,
+                           dit_cpu_offload=False, pin_cpu_memory=False, use_fsdp_inference=False,
+                           training_mode=False, enable_torch_compile=False, torch_compile_kwargs={},
+                           inference_torch_compile=False, VSA_tile_size=None, inference_mode=True,
+                           dit_layerwise_offload=False, attention_backend="TORCH_SDPA")
+    native = PipelineComponentLoader.load_module("transformer", str(component), "diffusers", args)
+    inputs = _inputs(channels=native.config.in_channels, text_dim=native.config.context_in_dim,
+                     device="cuda:0", dtype=torch.bfloat16, batch=1)
+    cached_inputs = dict(inputs)
+    cached_inputs["timestep"] = inputs["timestep"] * 0.5
+    actual_layers = []
+    handles = [block.register_forward_hook(lambda module, args, output: actual_layers.append(output.float().cpu()))
+               for block in native.transformer_blocks]
+    cache = QwenImage21KVCache(len(native.transformer_blocks), "cpu")
+    with torch.inference_mode():
+        actual = native(**inputs, kv_cache=cache, kv_cache_mode="extract").float().cpu()
+        for handle in handles:
+            handle.remove()
+        actual_cached = native(**cached_inputs, kv_cache=cache, kv_cache_mode="cached").float().cpu()
+    del native, cache
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    official, info = reference.QwenImage21Transformer2DModel.from_pretrained(component, local_files_only=True,
+                                                                            torch_dtype=torch.bfloat16,
+                                                                            output_loading_info=True)
+    assert not {key: info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
+                if info.get(key)}
+    official = official.to("cuda:0").eval()
+    expected_layers = []
+    handles = [block.register_forward_hook(lambda module, args, output: expected_layers.append(output.float().cpu()))
+               for block in official.transformer_blocks]
+    official_cache = reference.QwenImage21KVCache(len(official.transformer_blocks))
+    with torch.inference_mode():
+        expected = official(**inputs, kv_cache=official_cache, kv_cache_mode="extract", return_dict=False)[0]
+        for handle in handles:
+            handle.remove()
+        expected_cached = official(**cached_inputs, kv_cache=official_cache, kv_cache_mode="cached",
+                                   return_dict=False)[0]
+
+    # Image slots expand into four unpatched latent tokens; the input masks
+    # independently identify text, reference images, target image and padding.
+    slots = inputs["img_mask"][0].bool().cpu()
+    repeats = torch.where(slots, 4, 1)
+    image_tokens = slots.repeat_interleave(repeats)
+    target_tokens = math.prod(inputs["img_shapes"][0][-1])
+    prefix = torch.arange(image_tokens.numel()) < image_tokens.numel() - target_tokens
+    prompt_valid = inputs["encoder_hidden_states_mask"][0].bool().cpu()
+    slot_valid = torch.cat((prompt_valid, torch.ones(slots.numel() - prompt_valid.numel(), dtype=torch.bool)))
+    text_valid = slot_valid.repeat_interleave(repeats)
+    masks = {
+        "text_prefix": ~image_tokens & text_valid & prefix,
+        "image_tokens": image_tokens,
+        "reference_images": image_tokens & prefix,
+        "target_image": image_tokens & ~prefix,
+        "padding": ~image_tokens & ~text_valid & prefix,
+    }
+    report = {
+        "component": "qwen_image21_transformer",
+        "checkpoint_revision": MODEL_ROOT.name,
+        "reference_revision": REFERENCE_REVISION,
+        "torch_version": torch.__version__,
+        "atol": 3e-3,
+        "rtol": 1e-2,
+        "blocks": [],
+    }
+    for index, (actual_layer, expected_layer) in enumerate(zip(actual_layers, expected_layers, strict=True)):
+        summary = {"block": index, **_modality_summaries(actual_layer, expected_layer, masks)}
+        report["blocks"].append(summary)
+        print(f"QWEN_IMAGE21_DIT_PARITY block_{index} {json.dumps(summary, sort_keys=True)}", flush=True)
+    report["final_prediction"] = _modality_summaries(actual, expected, masks)
+    report["cached_prediction"] = _modality_summaries(
+        actual_cached, expected_cached, {"target_image": torch.ones(target_tokens, dtype=torch.bool)})
+    for name in ("final_prediction", "cached_prediction"):
+        print(f"QWEN_IMAGE21_DIT_PARITY {name} {json.dumps(report[name], sort_keys=True)}", flush=True)
+    report_path = os.environ.get("QWEN_IMAGE21_DIT_PARITY_REPORT")
+    if report_path:
+        destination = Path(report_path)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        destination.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
+
+    for index, (actual_layer, expected_layer) in enumerate(zip(actual_layers, expected_layers, strict=True)):
+        assert_close(actual_layer, expected_layer, atol=3e-3, rtol=1e-2, msg=f"Transformer block {index}")
+    assert_close(actual, expected.float().cpu(), atol=3e-3, rtol=1e-2)
+    assert_close(actual_cached, expected_cached.float().cpu(), atol=3e-3, rtol=1e-2)

tests/local_tests/vaes/test_qwen_image21_vae_parity.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen-Image-2.1 VAE parity against a pinned Diffusers checkout.
+
+Random-weight CPU tests cover the image-specific residual paths and tiling.
+The production-loader gate requires one CUDA GPU and user-provided weights.
+"""
+
+from __future__ import annotations
+
+import gc
+import importlib
+import os
+import subprocess
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+from torch.testing import assert_close
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+REFERENCE_REVISION = "578c9b2c6636ab2424a0e56186268b83623656b2"
+REFERENCE_ROOT = Path(os.environ.get("QWEN_IMAGE21_DIFFUSERS_DIR", REPO_ROOT / "official_reference" / "diffusers"))
+MODEL_ROOT = Path(os.environ.get("QWEN_IMAGE21_MODEL_ROOT", REPO_ROOT / "official_weights" / "Qwen-Image-2.1"))
+RUN_ENV = "QWEN_IMAGE21_RUN_VAE_PARITY"
+PARITY_SCOPE = "both"
+
+
+def _reference_class(required: bool = False):
+    if not (REFERENCE_ROOT / "src").is_dir():
+        message = f"set QWEN_IMAGE21_DIFFUSERS_DIR to Diffusers checkout at {REFERENCE_REVISION}"
+        if required:
+            pytest.fail(message, pytrace=False)
+        pytest.skip(message)
+    revision = subprocess.check_output(["git", "-C", str(REFERENCE_ROOT), "rev-parse", "HEAD"], text=True).strip()
+    assert revision == REFERENCE_REVISION, f"reference revision mismatch: {revision}"
+    sys.path.insert(0, str(REFERENCE_ROOT / "src"))
+    try:
+        module = importlib.import_module("diffusers.models.autoencoders.autoencoder_kl_qwenimage21")
+    except ImportError as exc:
+        if required:
+            pytest.fail(f"cannot import Qwen-Image-2.1 reference: {exc}", pytrace=False)
+        pytest.skip(f"reference dependencies unavailable: {exc}")
+    expected_source = REFERENCE_ROOT / "src" / "diffusers" / "models" / "autoencoders"
+    assert Path(module.__file__).resolve().parent == expected_source.resolve()
+    return module.AutoencoderKLQwenImage21
+
+
+def _tiny_config():
+    from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEArchConfig, QwenImage21VAEConfig
+
+    return QwenImage21VAEConfig(
+        arch_config=QwenImage21VAEArchConfig(
+            base_dim=4,
+            decoder_base_dim=4,
+            z_dim=4,
+            dim_mult=(1, 2, 4, 8, 8),
+            num_res_blocks=1,
+            latents_mean=(0.0,) * 4,
+            latents_std=(1.0,) * 4,
+        ),
+        use_tiling=False,
+    )
+
+
+@pytest.fixture
+def _tiny_pair():
+    official_cls = _reference_class()
+    from fastvideo.models.vaes.qwen_image21 import AutoencoderKLQwenImage21
+
+    config = _tiny_config()
+    arch = config.arch_config
+    official = official_cls(
+        base_dim=arch.base_dim,
+        decoder_base_dim=arch.decoder_base_dim,
+        z_dim=arch.z_dim,
+        dim_mult=list(arch.dim_mult),
+        num_res_blocks=arch.num_res_blocks,
+        temperal_downsample=list(arch.temperal_downsample),
+        latents_mean=list(arch.latents_mean),
+        latents_std=list(arch.latents_std),
+    ).eval()
+    native = AutoencoderKLQwenImage21(config).eval()
+    official_state = official.state_dict()
+    assert {key: value.shape for key, value in native.state_dict().items()} == {
+        key: value.shape for key, value in official_state.items()
+    }
+    native.load_state_dict(official_state, strict=True)
+    return official, native
+
+
+@pytest.mark.parametrize("tiled", [False, True])
+def test_qwen_image21_vae_random_cpu_parity(_tiny_pair, tiled):
+    official, native = _tiny_pair
+    if tiled:
+        for model in (official, native):
+            model.enable_tiling(
+                tile_sample_min_height=32,
+                tile_sample_min_width=32,
+                tile_sample_stride_height=16,
+                tile_sample_stride_width=16,
+            )
+    generator = torch.Generator().manual_seed(21)
+    pixels = torch.rand((2, 4, 1, 48, 64), generator=generator).mul_(2).sub_(1)
+    with torch.inference_mode():
+        expected = official.encode(pixels).latent_dist
+        actual = native.encode(pixels).latent_dist
+        assert_close(actual.mean, expected.mean, atol=2e-6, rtol=2e-5)
+        assert_close(actual.logvar, expected.logvar, atol=2e-6, rtol=2e-5)
+        expected_pixels = official.decode(expected.mode()).sample
+        actual_pixels = native.decode(expected.mode()).sample
+    assert_close(actual_pixels, expected_pixels, atol=2e-6, rtol=2e-5)
+    assert actual_pixels.shape == pixels.shape
+
+
+def test_qwen_image21_vae_slicing_and_output_contract(_tiny_pair):
+    official, native = _tiny_pair
+    for model in (official, native):
+        model.enable_slicing()
+    pixels = torch.rand((2, 4, 1, 32, 48), generator=torch.Generator().manual_seed(22)).mul_(2).sub_(1)
+    with torch.inference_mode():
+        expected = official.encode(pixels, return_dict=False)[0].mode()
+        actual = native.encode(pixels, return_dict=False)[0].mode()
+        assert_close(actual, expected, atol=2e-6, rtol=2e-5)
+        assert_close(
+            native.decode(actual, return_dict=False)[0],
+            official.decode(expected, return_dict=False)[0],
+            atol=2e-6,
+            rtol=2e-5,
+        )
+
+
+def test_qwen_image21_vae_production_loader_parity():
+    if os.environ.get(RUN_ENV) != "1":
+        pytest.skip(f"set {RUN_ENV}=1 on an allocated CUDA node")
+    if not torch.cuda.is_available():
+        pytest.fail("Qwen-Image-2.1 VAE production parity requires CUDA", pytrace=False)
+    official_cls = _reference_class(required=True)
+    component_path = MODEL_ROOT / "vae"
+    if not (component_path / "config.json").is_file() or not list(component_path.glob("*.safetensors")):
+        pytest.fail(f"local VAE checkpoint missing at {component_path}", pytrace=False)
+
+    from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
+    from fastvideo.models.loader.component_loader import VAELoader
+
+    device = torch.device("cuda:0")
+    args = SimpleNamespace(
+        pipeline_config=QwenImage21PipelineConfig(),
+        model_paths={},
+        vae_cpu_offload=False,
+    )
+    args.pipeline_config.vae_precision = "fp32"
+    args.pipeline_config.vae_config.use_tiling = False
+    native = VAELoader().load(str(component_path), args)
+    assert all(parameter.dtype == torch.float32 for parameter in native.parameters())
+    generator = torch.Generator().manual_seed(23)
+    pixels = torch.rand((1, 4, 1, 64, 80), generator=generator).mul_(2).sub_(1).to(device)
+    with torch.inference_mode():
+        posterior = native.encode(pixels).latent_dist
+        actual_mean = posterior.mean.cpu()
+        actual_logvar = posterior.logvar.cpu()
+        actual_pixels = native.decode(posterior.mode()).sample.cpu()
+    del native, posterior
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    official, info = official_cls.from_pretrained(
+        component_path,
+        local_files_only=True,
+        torch_dtype=torch.float32,
+        output_loading_info=True,
+    )
+    assert not {key: info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
+                if info.get(key)}
+    official = official.to(device).eval()
+    with torch.inference_mode():
+        expected = official.encode(pixels).latent_dist
+        assert_close(actual_mean, expected.mean.cpu(), atol=1e-5, rtol=1e-4)
+        assert_close(actual_logvar, expected.logvar.cpu(), atol=1e-5, rtol=1e-4)
+        assert_close(actual_pixels, official.decode(expected.mode()).sample.cpu(), atol=1e-5, rtol=1e-4)
+
+
+@pytest.mark.parametrize("height,width,tiled", [(64, 80, False), (288, 320, True)])
+def test_qwen_image21_vae_bf16_reference_layout_parity(height, width, tiled):
+    """Independent preprocessing must preserve singleton strides for BF16 kernels."""
+    if os.environ.get(RUN_ENV) != "1":
+        pytest.skip(f"set {RUN_ENV}=1 on an allocated CUDA node")
+    if not torch.cuda.is_available():
+        pytest.fail("Qwen-Image-2.1 BF16 VAE production parity requires CUDA", pytrace=False)
+    official_cls = _reference_class(required=True)
+    component_path = MODEL_ROOT / "vae"
+    if not (component_path / "config.json").is_file() or not list(component_path.glob("*.safetensors")):
+        pytest.fail(f"local VAE checkpoint missing at {component_path}", pytrace=False)
+
+    from diffusers.image_processor import VaeImageProcessor
+    from diffusers.pipelines.qwenimage21.pipeline_qwenimage21 import QwenImage21Pipeline
+    from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
+    from fastvideo.models.loader.component_loader import VAELoader
+    from fastvideo.pipelines.basic.qwen_image21.stages import normalize_latents, pack_latents, reference_pixels
+
+    rgb = torch.randint(0, 256, (height, width, 3), generator=torch.Generator().manual_seed(23), dtype=torch.uint8)
+    image = Image.fromarray(rgb.numpy()).convert("RGBA")
+    processor = VaeImageProcessor(vae_scale_factor=16, do_convert_rgb=False)
+    native_pixels = reference_pixels(image)
+    official_pixels = processor.preprocess(image, height=height, width=width).unsqueeze(2)
+    assert_close(native_pixels, official_pixels, atol=0, rtol=0)
+    # Size-one strides can select a different BF16 convolution kernel even
+    # though both tensors contain identical values and are channels_last_3d.
+    assert native_pixels.stride() == official_pixels.stride()
+    device = torch.device("cuda:0")
+    native_pixels = native_pixels.to(device=device, dtype=torch.bfloat16)
+    official_pixels = official_pixels.to(device=device, dtype=torch.bfloat16)
+    assert native_pixels.stride() == official_pixels.stride()
+
+    args = SimpleNamespace(pipeline_config=QwenImage21PipelineConfig(), model_paths={}, vae_cpu_offload=False)
+    args.pipeline_config.vae_config.use_tiling = tiled
+    native = VAELoader().load(str(component_path), args)
+    assert all(parameter.dtype == torch.bfloat16 for parameter in native.parameters())
+    with torch.no_grad():
+        posterior = native.encode(native_pixels).latent_dist
+        actual_mean, actual_logvar = posterior.mean.cpu(), posterior.logvar.cpu()
+        normalized = normalize_latents(posterior.mode(), native.config.latents_mean, native.config.latents_std)
+        actual_packed = pack_latents(normalized).cpu()
+    del native, posterior, normalized
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    official, info = official_cls.from_pretrained(
+        component_path, local_files_only=True, torch_dtype=torch.bfloat16, output_loading_info=True)
+    assert not {key: info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
+                if info.get(key)}
+    official = official.to(device).eval()
+    if tiled:
+        official.enable_tiling()
+    with torch.no_grad():
+        expected = official.encode(official_pixels).latent_dist
+        assert_close(actual_mean, expected.mean.cpu(), atol=1e-5, rtol=1e-4)
+        assert_close(actual_logvar, expected.logvar.cpu(), atol=1e-5, rtol=1e-4)
+        mean = torch.tensor(official.config.latents_mean).view(1, official.config.z_dim, 1, 1, 1)
+        std = torch.tensor(official.config.latents_std).view(1, official.config.z_dim, 1, 1, 1)
+        normalized = (expected.mode() - mean.to(device, torch.bfloat16)) / std.to(device, torch.bfloat16)
+        expected_packed = QwenImage21Pipeline._pack_latents(
+            normalized, 1, official.config.z_dim, height // 16, width // 16).cpu()
+        assert_close(actual_packed, expected_packed, atol=1e-5, rtol=1e-4)