[perf] PAN2 sampling: Wan2.2 VAE decode kernel, split decode, fused DiT kernels

- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on
  first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames
  per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output.
- WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32
  accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it);
  decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the
  video, with the ranks agreeing on whether every one of them could build the kernel.
- PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions,
  set_fused_kernels).
- PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels).
- VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame.
- The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.

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/content/cookbook/pan2/pan2.mdx

@@ -0,0 +1,7 @@
+---
+title: PAN 2
+description: Configure text-to-video and first-frame image-to-video generation with PAN 2.
+mode: frame
+---
+
+<Cookbook modelId="pan2" />

docs-site/lib/cookbook-pages.ts

@@ -66,6 +66,15 @@ export const recipePages = [
     description: "Turn a reference image into video with Wan2.1.",
     tags: ["Image to video", "14B", "480p"],
   },
+  {
+    id: "pan2",
+    family: "PAN 2",
+    title: "PAN 2",
+    mark: "PAN",
+    path: "/cookbook/pan2/pan2",
+    description: "Text-to-video and first-frame image-to-video from the 29B PAN 2 model.",
+    tags: ["Text or image to video", "29B"],
+  },
   {
     id: "ltx2-distilled",
     family: "LTX",
@@ -110,7 +119,9 @@ const familyInfo: Record<string, { title: string; path: string; logo: string; ca
   hunyuan: { title: "Hunyuan", path: "hunyuan", logo: "tencent-hunyuan", category: "Video" },
   kandinsky5: { title: "Kandinsky 5", path: "kandinsky5", logo: "kandinsky", category: "Video" },
   kandinsky6: { title: "Kandinsky 6", path: "kandinsky6", logo: "kandinsky", category: "Video" },
+  pan2: { title: "PAN 2", path: "pan2", logo: "ifm", 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" },
@@ -126,12 +137,14 @@ const backendModels: Record<string, string> = {
   "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": "fastwan21",
   "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": "wan21-i2v",
   "FastVideo/LTX2-Distilled-Diffusers": "ltx2-distilled",
+  "IFM/PAN2": "pan2",
 };
 const TASK_GROUPS: Array<[RegExp, string]> = [
   [/world/, "Text to world"],
   [/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 +174,8 @@ export const modelDirectory = [...new Set(sourceRecipes.map((recipe) => recipe.m
       "tencent-hunyuan": "Tencent",
       kandinsky: "Kandinsky Lab",
       "black-forest-labs": "Black Forest Labs",
+      qwen: "Qwen",
+      ifm: "IFM",
       stabilityai: "Stability AI",
       nvidia: "NVIDIA",
     };

docs-site/navigation.json

@@ -127,6 +127,10 @@
         ],
         "display": "group"
       },
+      {
+        "label": "PAN 2 text/image to video",
+        "href": "/inference/pan2"
+      },
       {
         "label": "Examples",
         "href": "/inference/examples/examples_inference_index"
@@ -249,10 +253,33 @@
           }
         ]
       },
+      {
+        "label": "PAN 2",
+        "display": "group",
+        "collapsed": true,
+        "items": [
+          {
+            "label": "PAN 2",
+            "href": "/cookbook/pan2/pan2"
+          },
+          {
+            "label": "Family guide",
+            "href": "/cookbook/pan2"
+          },
+          {
+            "label": "Model reference",
+            "href": "/inference/pan2"
+          }
+        ]
+      },
       {
         "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(38);
   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(38);
 });
 
 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);
@@ -423,6 +443,26 @@ test("source-backed serving pages expose their actual server settings", async ({
   await expect(page.getByRole("button", { name: "Copy config + command", exact: true })).toBeEnabled();
 });
 
+test("PAN 2 has a family card with the IFM logo and an editable serving deployment", async ({ page }) => {
+  await page.goto("/FastVideo/cookbook");
+  await page.getByLabel("Search recipes").fill("PAN 2");
+  const pan2 = page.locator(".recipe-card");
+  await expect(pan2).toHaveCount(1);
+  await expect(pan2.locator("img")).toHaveAttribute("src", /\/ifm\.webp$/);
+  await expect(pan2.locator(".recipe-card-kicker")).toHaveText("IFM");
+  await pan2.click();
+  await expect(page).toHaveURL(/cookbook\/pan2\/pan2$/);
+  await expect(page.locator(".recipe-matrix")).toBeVisible({ timeout: 30000 });
+  await expect(page.getByRole("heading", { level: 1 })).toHaveText("PAN 2");
+  await settings(page, "Serving");
+  await page.getByLabel("Port", { exact: true }).fill("9010");
+  await output(page, "config.yaml");
+  await expect(page.getByLabel("Generated YAML", { exact: true })).toContainText(/"port": 9010/);
+  await expect(page.getByLabel("Generated YAML", { exact: true })).toContainText(/"model_path": "IFM\/PAN2"/);
+  await expect(page.getByLabel("Generated YAML", { exact: true })).toContainText(/"dit_layerwise": false/);
+  await expect(page.getByRole("button", { name: "Download YAML", exact: false })).toBeEnabled();
+});
+
 test("source model hardware preferences survive hydration and reload", async ({ page }) => {
   const errors: string[] = [];
   page.on("pageerror", (error) => errors.push(error.message));
@@ -584,11 +624,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": 20,
   "recipes": [
     {
       "id": "fastwan21-t2v",
@@ -381,6 +381,33 @@
       "expected_artifact": "Upscaled MP4 under outputs_video/kandinsky6_sr/",
       "related": ["kandinsky6-vsr", "kandinsky6-ti2va-piflow"]
     },
+    {
+      "id": "pan2-ti2v",
+      "family": "pan2",
+      "stage": "inference",
+      "task": "Text or image to video",
+      "label": "PAN 2",
+      "summary": "Generate a 13-second, 704 × 1248 video at 24 fps from text, or set IMAGE_PATH in the maintained example to condition on a first frame.",
+      "model": "IFM/PAN2",
+      "source": "examples/inference/basic/basic_pan2.py",
+      "serving": {
+        "source": "examples/serving/openai_pan2.yaml",
+        "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .",
+        "task": "Text to video",
+        "limitations": [
+          "This server walkthrough sends a text prompt only. The server API also accepts an image reference (image_reference or input_reference) and then runs image-to-video from that first frame, but the checked-in clients do not upload one."
+        ]
+      },
+      "command": "python examples/inference/basic/basic_pan2.py",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
+      "evidence": "Source-backed",
+      "expected_artifact": "MP4 video under video_samples_pan2/",
+      "limitations": [
+        "The full IFM/PAN2 checkpoint is not public yet; numerical parity is checked with a public random-weight checkpoint.",
+        "The example keeps the 29B DiT resident in bf16 and offloads only the text encoder. GPU memory requirements have not been established for this recipe; set PAN2_NUM_GPUS to shard video tokens with sequence parallelism, which does not shard the DiT weights."
+      ]
+    },
     {
       "id": "flux2-klein-t2i",
       "family": "flux",
@@ -429,6 +456,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,8 @@ 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 |
+| `ifm.webp` | [Official Institute of Foundation Models (IFM) Hugging Face organization avatar](https://huggingface.co/IFM) | PAN 2 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.                                                                                                                                                                                                                                                                                                               |
@@ -259,6 +259,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_H3_VSA_FP4_K_SMOOTH`                            | bool                                                               | `0`                                                   | attention   | Subtract the per-head key mean before FP4 key quantization, matching the emulator and QAD training numerics. Exact for the softmax; it only changes quantization error.                                                                                                                                                                                                                                                                     |
 | `FASTVIDEO_H3_GATE_CHUNK_ROWS`                             | int                                                                | `0`                                                   | performance | Rows per chunk when the VSA FP4 gate branch expands the pooled compression output back to token rows. 0 keeps the single full-size index_select; a positive value caps the transient at chunk_rows\*heads\*head_dim and is bit-identical.                                                                                                                                                                                                   |
 | `FASTVIDEO_H3_VSA_FINE_ATTENTION`                          | one of , emulator-mix8, emulator-mixmx, emulator-qk, emulator-bf16 | `""`                                                  | attention   | Replace tile-128 fine attention, preserving VSA selection and compression, with the Triton emulator of the sparse FP4 kernel: emulator-mix8 = NVFP4 Q/K video, per-token FP8 Q/K prefix; emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over FASTVIDEO_H3_VSA_FP4. Empty keeps the native route.                 |
+| `FASTVIDEO_WAN22_VAE_DECODE_DIR`                           | path                                                               | `""`                                                  | performance | Directory of the Wan2.2 VAE decoder kernels (fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds on first use, when its config sets decode_kernel. Empty takes the copy in this checkout.                                                                                                                                                                                                                            |
 | `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE`                        | bool                                                               | `0`                                                   | performance | VSA-H3 64/128-token tiles: scatter rows straight into the heads-first layout the sparse kernels read, removing their per-call transpose copy of query, key and value. Bit-identical; no-grad, uncompiled CUDA only.                                                                                                                                                                                                                         |
 | `FASTVIDEO_H3_VSA_TILE_FIRST`                              | bool                                                               | `0`                                                   | attention   | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters.                                                                                                                                                                                                                                                                                                                                 |
 | `FASTVIDEO_H3_VSA_SM89_KERNEL`                             | one of original, bf16, int8                                        | `original`                                            | attention   | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV).                                                                                                                                                                                                                                                                                                                                              |
@@ -328,6 +329,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER`                     | int                                                                | `0`                                                   | test        | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`.                                                                                                                                                                                                                                                                                                                                   |
 | `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT`                     | str                                                                | unset                                                 | test        | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`.                                                                                                                                                                                                                                                                                                                                            |
 | `FASTVIDEO_TEST_SD35_MODEL_DIR`                            | str                                                                | `stabilityai/stable-diffusion-3.5-medium`             | test        | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`.                                                                                                                                                                                                                                                                                                                                                           |
+| `FASTVIDEO_TEST_PAN2_MODEL_DIR`                            | path                                                               | unset                                                 | test        | PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU tests; unset skips them.                                                                                                                                                                                                                                                                                                                                              |
 | `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR`                       | str                                                                | unset                                                 | test        | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`.                                                                                                                                                                                                                                                                                                                                             |
 | `FASTVIDEO_TEST_WAN_ANIMATE_MODEL_DIR`                     | str                                                                | unset                                                 | test        | Local Wan2.2-Animate-14B checkpoint for the Wan-Animate weight tests. Deprecated names: `WAN_ANIMATE_MODEL_PATH`.                                                                                                                                                                                                                                                                                                                           |
 

docs/cookbook/index.md

@@ -6,7 +6,9 @@ The interactive catalog groups models by family and generates YAML for supported
 - [FastH3](minimax-h3.md)
 - [Wan and FastMetal](wan.md)
 - [LTX](ltx.md)
+- [PAN 2](pan2.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/pan2.md

@@ -0,0 +1,76 @@
+# PAN 2 recipes
+
+<RecipeFamily family="pan2" />
+
+Generate a video from text, or from a first frame, with `IFM/PAN2`. One pipeline
+serves both modes: leave out the first frame for text-to-video and supply one
+for image-to-video. The defaults are 704×1248, 313 frames (13 s at 24 fps),
+50 steps and guidance 3.0. See [PAN 2 Text/Image to Video](../inference/pan2.md)
+for the model, sampler and parallelism details.
+
+## Run on a GPU
+
+Follow the [CUDA installation guide](../getting_started/installation/gpu.md), then run
+from the repository root. The full `IFM/PAN2` checkpoint is not public yet; pass a
+local directory or private repo you can access with `PAN2_MODEL_PATH`.
+
+```bash
+python examples/inference/basic/basic_pan2.py
+```
+
+Set `IMAGE_PATH` in the script to a first frame for image-to-video. With more
+GPUs, `PAN2_NUM_GPUS=4` shards the video tokens with sequence parallelism; the
+count must divide PAN 2's 32 attention heads. Each rank still holds the full
+29B DiT, so extra GPUs speed up generation but do not reduce per-device weight
+memory. See [CPU offloading](../inference/offloading.md) for devices that cannot
+hold the DiT.
+
+## Serve with the OpenAI-compatible API
+
+`examples/serving/openai_pan2.yaml` serves PAN 2 through the shared
+`/v1/videos` routes. It mirrors the Python example: one GPU, the text encoder
+offloaded between requests, and the pipeline's own sampling defaults.
+
+```bash
+fastvideo serve --config examples/serving/openai_pan2.yaml --server.host 127.0.0.1
+```
+
+Text-to-video uses the same request as any other served video model, with the
+alias `pan2`:
+
+```bash
+curl --fail-with-body http://127.0.0.1:8000/v1/videos \
+  -H 'Content-Type: application/json' \
+  -d '{"model": "pan2", "prompt": "A red fox trots through fresh snow in a pine forest at dawn."}'
+```
+
+For image-to-video, add a first frame in `image_reference` (or the legacy
+`input_reference`). The server then runs image-to-video; no second deployment or
+workload is needed.
+
+```bash
+curl --fail-with-body http://127.0.0.1:8000/v1/videos \
+  -H 'Content-Type: application/json' \
+  -d '{"model": "pan2", "prompt": "The fox turns toward the camera.",
+       "image_reference": {"image_url": "https://example.com/first-frame.png"}}'
+```
+
+Local image paths must sit under `server.media_roots`, and private or loopback
+URLs need `server.allow_private_media_urls`. Poll `/v1/videos/{id}` and download
+`/v1/videos/{id}/content` as in the [serving guide](openai-api.md). Generation
+time for the default 13-second clip is not measured, so raise the client's
+polling deadline for long jobs instead of treating a slow job as failed.
+
+The cookbook's generated serving clients send a text prompt only. The image
+request above is the supported way to supply a first frame over HTTP, and no
+recorded generation backs it yet.
+
+## Evidence
+
+Both recipes are **Source-backed**: the model ID, example and serving YAML are
+checked against each other, and the serving defaults are tested against the
+registered `pan2_t2v` preset. No GPU generation, memory measurement or HTTP
+client run is recorded, so hardware fit and throughput are unknown. Numerical
+parity with the reference implementation is checked on CPU with a public
+random-weight checkpoint, which verifies the port and produces noise rather than
+useful video.

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/cookbook/recipes/pan2.yaml

@@ -0,0 +1,38 @@
+title: PAN 2
+summary: Text-to-video and first-frame image-to-video from the 29B PAN 2 mixture-of-transformers model.
+
+deployments:
+  cuda-rest:
+    runtime: fastvideo-cuda-rest
+    workload: t2v
+    defaults: examples/serving/openai_pan2.yaml
+    requirements:
+      - >-
+        Requests without a reference run text-to-video. Add image_reference or input_reference
+        to condition on a first frame; the same server and pipeline handle both.
+      - >-
+        To use more than one GPU, set sp_size equal to num_gpus (Ulysses sequence parallelism).
+        It must divide PAN 2's 32 attention heads, and each rank still holds the full DiT weights.
+    options:
+      - server.*
+      - generator.engine.num_gpus
+      - generator.engine.parallelism.*
+      - "!generator.engine.parallelism.hsdp_*"
+      - generator.engine.offload.text_encoder
+      - generator.engine.offload.vae
+      - default_request.sampling.num_frames
+      - default_request.sampling.height
+      - default_request.sampling.width
+      - default_request.sampling.fps
+    overrides:
+      default_request.sampling.num_frames:
+        default: 313
+        description: num_frames - 1 must be a multiple of 4.
+      default_request.sampling.height:
+        default: 704
+        minimum: 32
+        multipleOf: 32
+      default_request.sampling.width:
+        default: 1248
+        minimum: 32
+        multipleOf: 32

docs/cookbook/serving-coverage.md

@@ -49,6 +49,7 @@ never deleted. See `examples/serving/README.md`.
 | `openai_hunyuan15_i2v_480p.yaml` | HunyuanVideo 1.5 I2V 480p | `i2v` | CUDA REST | `hunyuan15-i2v-480p` | Source-backed; preset renders 720x1280 |
 | `openai_kandinsky6_ti2va.yaml` | Kandinsky 6 Pro TI2VA | `t2v` | CUDA REST | `kandinsky6-ti2va` | Source-backed |
 | `openai_kandinsky6_ti2va_distilled.yaml` | Kandinsky 6 Pro-distill TI2VA | `t2v` | CUDA REST | `kandinsky6-ti2va-distilled` | Source-backed |
+| `openai_pan2.yaml` | PAN 2 | `t2v` (a reference switches to image-to-video) | CUDA REST | `pan2/cuda-rest` | Source-backed; checkpoint not public, no recorded run |
 | `openai_ltx2_distilled.yaml` | LTX2 Distilled | `t2v` | CUDA REST | `ltx2-distilled/cuda-rest` | Source-backed |
 | `streaming_demo.yaml` | LTX2 Distilled | `t2v` | CUDA streaming | `ltx2-distilled/cuda-streaming` | Source-backed |
 | `mlx_fasth3.yaml` | FastH3 V1 converted | `t2v` | MLX REST | none (legacy catalog only) | Source runtime; T2VA only |

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:
@@ -282,6 +283,18 @@ surfaces:
           - fastvideo.configs.pipelines.hunyuan15.Hunyuan15T2V720PConfig
           - fastvideo.configs.pipelines.hyworld.HYWorldConfig
           - fastvideo.configs.pipelines.hyworld.Hunyuan15T2V480PConfig
+      text_system_prompt:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+      text_encoder_skip_layers:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+      text_max_sequence_length:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+      i2v_skip_first_latent_cfg:
+        sources:
+          - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
       text_encoder_max_lengths:
         sources:
           - fastvideo.configs.pipelines.hunyuan15.Hunyuan15I2V480PStepDistilledConfig
@@ -442,6 +455,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 +482,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/generate_examples.py

@@ -35,7 +35,9 @@ COOKBOOK_FAMILIES = {
     "cosmos",
     "kandinsky5",
     "kandinsky6",
+    "pan2",
     "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/pan2.md

@@ -0,0 +1,100 @@
+# PAN 2 Text/Image to Video
+
+PAN 2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the
+Wan2.2 VAE. One pipeline serves text-to-video and image-to-video: pass `image_path` to condition on a first frame and
+leave it out for text only.
+
+## Model
+
+| Hub repo | Pipeline | Workloads | Example |
+|---|---|---|---|
+| `IFM/PAN2` | `PAN2ModularPipeline` | T2V, I2V | [`basic_pan2.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pan2.py) |
+
+PAN2 checkpoints use the Modular Diffusers layout (`modular_model_index.json`). Any repo id or local directory whose
+name contains `pan2` (or whose manifest declares `PAN2ModularPipeline`) resolves to this pipeline. The defaults come
+from the `pan2_t2v` and `pan2_i2v` presets:
+
+| Setting | Default |
+|---|---|
+| Resolution | 704x1248 (height and width must be multiples of 32) |
+| Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) |
+| Steps | 50, uniform sigmas from 1 to 0 with the checkpoint's flow shift (9 for `IFM/PAN2-TEST`, the reference sampler's value) |
+| Guidance | 2.5, classifier-free guidance, the reference launcher's default (3 before 2026-10-23). Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) |
+| Negative prompt | None unless the request carries one. The reference sampler always sends the release's structured negative caption (`negative_prompt.json` of the PAN2 repository) with its three `time_range_sec` values set to the clip length, `0 to 20` for 481 frames at 24 fps |
+
+PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
+prompts work, with weaker results.
+
+## Usage
+
+```bash
+python examples/inference/basic/basic_pan2.py
+```
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "IFM/PAN2",
+    num_gpus=1,
+    dit_cpu_offload=False,
+    dit_layerwise_offload=False,
+    text_encoder_cpu_offload=True,
+)
+generator.generate_video(
+    "A red fox trots through fresh snow in a pine forest at dawn.",
+    image_path=None,  # or the path of a first frame for image-to-video
+    output_path="video_samples_pan2",
+)
+```
+
+For image-to-video, the first frame is scaled to cover the requested resolution and center cropped, then encoded by the
+VAE into the first latent frame of the conditioning input.
+
+## Parallelism
+
+With `num_gpus > 1`, FastVideo runs Ulysses sequence parallelism: the video tokens are sharded across ranks and the
+text tokens are replicated. The number of GPUs must divide the 32 attention heads. Sequence-parallel padding of the
+video tokens is trimmed before attention, so no attention mask is needed.
+
+Padding is always excluded, without an opt-in setting. At 704x1248 with 313 frames, the 67,782 video tokens need
+two padding tokens at SP=4. They are removed before joining the text tokens in every attention block. This provides
+the same exclusion as vLLM-Omni's unconditional padding mask in
+[revision `e0c4c731`](https://github.com/vllm-project/vllm-omni/commit/e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74).
+The trimmed sequence can use either FlashAttention or Torch SDPA.
+
+```bash
+PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py
+```
+
+## Performance switches
+
+PAN2 sampling turns these on by default (`PAN2PipelineConfig`); pass any of them as `False` to
+`VideoGenerator.from_pretrained` to turn it off:
+
+| Switch | What it does |
+|---|---|
+| `vae_decode_kernel` | Decodes with fastvideo-kernel's Wan2.2 decoder on B200 / GB200 / B300 (704x1248, batch 1): bf16 with fp32 accumulation, 7x faster than the fp32 decode on a B300 at PSNR 61-63 dB and SSIM 0.998 against it, at a peak of about 20 GiB for a 5 s video. It builds on first use (minutes) into `WAN22_VAE_DECODE_BUILD_DIR`, against fastvideo-kernel's CUTLASS submodule (`git submodule update --init fastvideo-kernel/include/cutlass`, or set `CUTLASS_DIR`). Other GPUs and sizes keep the fp32 decode. |
+| `vae_decode_kernel_parallel` | With 2, 4 or 8 GPUs, splits that decode over them instead of decoding the whole video on each: a 20 s video takes 3.0 s on one B300, 1.55 s on two, 0.79 s on four and 0.47 s on eight. |
+| `dit_fused_kernels` | Runs the DiT's fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation. Measured with a distilled PAN2 model on eight B300s, they cut the denoising time by 21%. The rounding differs from the unfused path; record by record, fused and unfused videos were equally close to the 50-step reference. |
+
+On GPUs with room for every component (B200, B300), keep the text encoder resident as well:
+`text_encoder_cpu_offload=False`.
+
+## Memory and offloading
+
+The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards
+video tokens; each rank still holds the DiT weights. Increasing `num_gpus` alone does not shard those weights.
+For devices that cannot hold the DiT, see [CPU offloading](offloading.md) for `dit_cpu_offload` and
+`dit_layerwise_offload`. Hardware recipes and timings reported by vLLM-Omni use its own offload implementation;
+they have not been validated for FastVideo.
+
+The full `IFM/PAN2` checkpoint is not public yet. Local parity tests use the public random-weight
+`wuqing157/tiny-pan2-modular-pipe` checkpoint; it verifies the integration and produces noise rather than useful video.
+
+## Differences from the vLLM-Omni integration
+
+- The optional prompt and video guardrails (`pan2-guardrail`) are not included.
+- The Cache-DiT `quality` policy and CFG-parallel execution are vLLM-Omni serving features and are not ported.
+- The text encoder is a FastVideo-native Qwen3.5 port. Its Gated DeltaNet layers use a PyTorch chunked delta rule
+  instead of the `fla` kernels.

docs/inference/support_matrix.md

@@ -68,7 +68,9 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`<br />`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`<br />`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`<br />`FastVideo/Matrix-Game-2.0-Base-Diffusers`<br />`FastVideo/Matrix-Game-2.0-GTA-Diffusers`<br />`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`<br />`mignonjia/mg_longtuning_distilled_zelda`<br />`mignonjia/mg_sf_distilled_zelda_1k_steps`<br />`mignonjia/mg_sf_distilled_zelda`<br />`mignonjia/mg_causal_zelda`<br />`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) |
 | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
 | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)<br />[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)<br />[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
+| pan2 | `IFM/PAN2` | T2V, I2V | [basic_pan2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pan2.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_pan2.py

@@ -0,0 +1,40 @@
+import os
+
+from fastvideo import VideoGenerator
+
+# PAN2 is trained on the structured captions its prompt rewriter produces (serialized as one-line JSON). Plain
+# prompts like this one work, with weaker results.
+PROMPT = ("A red fox trots through fresh snow in a pine forest at dawn, its breath visible in the cold air. "
+          "The camera tracks it from the side at ground level as low golden light filters through the trees.")
+
+OUTPUT_PATH = "video_samples_pan2"
+MODEL_PATH = os.environ.get("PAN2_MODEL_PATH", "IFM/PAN2")
+
+# Set to a first frame (e.g. "assets/images/girl.png") for image-to-video; the same pipeline runs both modes.
+IMAGE_PATH = None
+
+
+def main():
+    # The 29B DiT runs in bf16; with more GPUs, num_gpus enables Ulysses sequence parallelism (the 32 attention
+    # heads must be divisible by it).
+    generator = VideoGenerator.from_pretrained(
+        MODEL_PATH,
+        num_gpus=int(os.environ.get("PAN2_NUM_GPUS", "1")),
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    generator.generate_video(
+        PROMPT,
+        image_path=IMAGE_PATH,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

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

examples/serving/README.md

@@ -90,6 +90,7 @@ uncovered registry families.
 | `openai_hunyuan15_i2v_480p.yaml` | HunyuanVideo 1.5 I2V 480p | image-to-video | 1 |
 | `openai_kandinsky6_ti2va.yaml` | Kandinsky 6 Pro TI2VA | text/image-to-video with audio | 1 |
 | `openai_kandinsky6_ti2va_distilled.yaml` | Kandinsky 6 Pro-distill TI2VA | text/image-to-video with audio | 1 |
+| `openai_pan2.yaml` | PAN 2 | text/image-to-video | 1 |
 | `openai_ltx2_distilled.yaml` | LTX2 Distilled | text-to-video with audio | 4 |
 | `openai_flux1_dev.yaml` | FLUX.1-dev | text-to-image | 1 |
 | `openai_flux2_dev.yaml` | FLUX.2-dev | text-to-image | 1 |

examples/serving/openai_pan2.yaml

@@ -0,0 +1,36 @@
+# OpenAI-compatible server for PAN 2 text-to-video and first-frame image-to-video.
+#   fastvideo serve --config examples/serving/openai_pan2.yaml
+# Omit image_reference for text-to-video; send it (or input_reference) for image-to-video.
+generator:
+  model_path: IFM/PAN2
+  engine:
+    num_gpus: 1
+    use_fsdp_inference: false
+    parallelism:
+      tp_size: 1
+      sp_size: 1
+    offload:
+      dit: false
+      dit_layerwise: false
+      vae: false
+      text_encoder: true
+      pin_cpu_memory: true
+  pipeline:
+    workload_type: t2v
+
+server:
+  host: 0.0.0.0
+  port: 8000
+  output_dir: outputs/openai_pan2
+  served_model_name: pan2
+
+default_request:
+  sampling:
+    height: 704
+    width: 1248
+    num_frames: 313
+    fps: 24
+    num_inference_steps: 50
+    guidance_scale: 3.0
+  output:
+    return_frames: false

fastvideo-kernel/README.md

@@ -21,6 +21,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
 | Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
 | FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
 | VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |
+| Wan2.2 VAE decoder (TI2V-5B, 48 latent channels) at 704 x 1248 on `sm_100a` / `sm_103a`, on 1-8 GPUs of a node | `wan22_vae_decode/`, built from `wan22_vae_decode/solution/` on first use against `include/cutlass` | `wan22_vae_decode.WanVAEDecoder`; FastVideo's `WanVAEConfig.decode_kernel`, on for PAN2 sampling (see `wan22_vae_decode/README.md`) |
 
 ## What gets built where, and when
 

fastvideo-kernel/tests/test_wan22_vae_decode.py

@@ -0,0 +1,97 @@
+"""The Wan2.2 VAE decoder (fastvideo-kernel/wan22_vae_decode) through its wrapper, against the reference.
+
+Needs an sm_100a or sm_103a GPU (B200 / GB200, B300) and the VAE's weights: WAN22_VAE_WEIGHTS names either a diffusers AutoencoderKLWan
+directory (its diffusion_pytorch_model.safetensors) or a .pt file holding the 110 tensors in weight_names() order. The
+first test builds the extension (about a minute). The kernel's own checks are wan22_vae_decode/tools/bench.py (nine
+timed workloads and six short clips at 1, 2 and 4 GPUs, on real latents); these tests cover the wrapper on one GPU:
+the weights it reads from a state dict, and decodes against the fp32 reference, each within three times what the bf16
+reference errs by on the same latents (the bench's rule, calibrated per input instead of per clip).
+"""
+
+import functools
+import os
+import sys
+
+import pytest
+import torch
+
+PACKAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "wan22_vae_decode")
+sys.path.insert(0, PACKAGE)
+sys.path.insert(0, os.path.join(PACKAGE, "tools"))
+
+WEIGHTS = os.environ.get("WAN22_VAE_WEIGHTS", "")
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
+                       reason="the Wan2.2 VAE decoder is built for sm_100a / sm_103a (B200 / GB200, B300)"),
+    pytest.mark.skipif(not WEIGHTS or not os.path.exists(WEIGHTS),
+                       reason="set WAN22_VAE_WEIGHTS to the VAE's diffusers directory or a weights .pt file"),
+]
+
+
+@functools.lru_cache(maxsize=None)
+def _state_dict():
+    """name -> tensor of the VAE's parameters, under diffusers' names."""
+    from wan22_vae_decode import weight_names
+
+    if os.path.isdir(WEIGHTS):
+        from safetensors.torch import load_file
+
+        return load_file(os.path.join(WEIGHTS, "diffusion_pytorch_model.safetensors"))
+    tensors = torch.load(WEIGHTS, map_location="cpu", weights_only=True)
+    return tensors if isinstance(tensors, dict) else dict(zip(weight_names(), tensors))
+
+
+@functools.lru_cache(maxsize=None)
+def _decoder():
+    from wan22_vae_decode import WanVAEDecoder
+
+    return WanVAEDecoder(_state_dict(), device="cuda")
+
+
+def test_weights_follow_diffusers_names():
+    from wan22_vae_decode import decoder_weights, weight_names
+
+    names = weight_names()
+    assert len(names) == 110 and names[0] == "post_quant_conv.weight" and names[-1] == "decoder.conv_out.bias"
+    weights = decoder_weights(_state_dict(), "cpu")
+    assert all(w.dtype == torch.bfloat16 and w.is_contiguous() for w in weights)
+    assert all(torch.equal(w, _state_dict()[n].to(torch.bfloat16)) for n, w in zip(names, weights))
+
+
+@pytest.mark.parametrize("latent_frames", [1, 2, 5])
+def test_decode_against_reference(latent_frames):
+    import reference
+    from wan22_vae_decode import weight_names
+
+    generator = torch.Generator().manual_seed(latent_frames)
+    z = torch.randn((1, 48, latent_frames, 44, 78), generator=generator).to("cuda", torch.bfloat16)
+    video = _decoder().decode(z)
+    assert video.shape == (1, 3, 4 * latent_frames - 3, 704, 1248) and video.dtype == torch.bfloat16
+
+    weights = {n: _state_dict()[n].to("cuda", torch.bfloat16).float() for n in weight_names()}
+    ref = reference.decode(z.float(), weights, torch.float32)
+    release = reference.decode(z, weights, torch.bfloat16).float()
+    d, d_release = video.float() - ref, release - ref
+    assert bool(torch.isfinite(video).all()) and float(video.abs().max()) <= 1.0
+    assert float(d.abs().max()) <= 3 * float(d_release.abs().max())
+    assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_release.pow(2).mean().sqrt())
+
+
+def test_decodes_are_independent():
+    """No state carries from one call to the next: the same latents decode to the same video after another clip."""
+    z = torch.randn((1, 48, 2, 44, 78), generator=torch.Generator().manual_seed(7)).to("cuda", torch.bfloat16)
+    first = _decoder().decode(z)
+    _decoder().decode(torch.randn((1, 48, 3, 44, 78), generator=torch.Generator().manual_seed(8)).to("cuda",
+                                                                                                     torch.bfloat16))
+    assert torch.equal(_decoder().decode(z), first)
+
+
+def test_pass_size_does_not_change_the_video():
+    """The causal history crosses passes, so any number of latent frames per pass decodes the same video."""
+    from wan22_vae_decode import WanVAEDecoder
+
+    z = torch.randn((1, 48, 6, 44, 78), generator=torch.Generator().manual_seed(9)).to("cuda", torch.bfloat16)
+    reference_pass = _decoder().decode(z)
+    for frames in (1, 16):
+        assert torch.equal(WanVAEDecoder(_state_dict(), device="cuda", frames_per_pass=frames).decode(z), reference_pass)

fastvideo-kernel/wan22_vae_decode/README.md

@@ -0,0 +1,98 @@
+# Wan2.2 VAE decoder (sm_100a, sm_103a)
+
+The decoder of the Wan2.2 VAE (the VAE of Wan2.2-TI2V-5B: 48 latent channels, 16 x 16 spatial and 4x temporal
+compression) at 704 x 1248, on one, two or four B200 / GB200 GPUs of a node. It decodes latents [1, 48, T, 44, 78]
+into the video [1, 3, 4T - 3, 704, 1248] for any T >= 1, as diffusers' `AutoencoderKLWan.decode(z).sample` does;
+`reference.py` defines the result exactly, a plain-PyTorch decoder bit-exact with diffusers in fp32 and in bf16.
+
+## Use
+
+```python
+import sys
+sys.path.insert(0, "fastvideo-kernel/wan22_vae_decode")
+from wan22_vae_decode import WanVAEDecoder
+
+decoder = WanVAEDecoder(vae.state_dict(), device="cuda")   # diffusers' AutoencoderKLWan, or its state dict
+video = decoder.decode(latents)   # [1, 48, T, 44, 78], as the VAE's decode takes them -> [1, 3, 4T-3, 704, 1248] bf16
+```
+
+- **Over several GPUs**, every rank builds a decoder with the same process group (`WanVAEDecoder(sd, device, pg)`,
+  default: torch.distributed's default group) and calls `decode` with the full latents; rank 0 gets the video and the
+  other ranks `None`. On one GPU without torch.distributed, the decoder initializes a single-rank NCCL group.
+- The first call builds the extension from `solution/` (about a minute) into `$WAN22_VAE_DECODE_BUILD_DIR` (default
+  `~/.cache/fastvideo/wan22_vae_decode`, one directory per `LOCAL_RANK`), and later processes reuse the build.
+- `WanVAEDecoder(..., frames_per_pass=4)` sets the latent frames each GPU carries through one pass (see below).
+- `WanVAEDecoder` prepares the weights once (layout, phase and Winograd transforms of the weights). The raw entry
+  points are `prepare(weights, rank, world, pg)`, `run(z, weights, rank, world, pg, out)` (4 frames per pass) and
+  `run_with(..., out, frames_per_pass)` in `solution/binding.cpp`, destination-passing, on the current stream;
+  `weights` are the 110 tensors of `weight_names()` in bf16.
+
+## Performance and memory
+
+At PAN2's default 704 x 1248, 24 fps, on GB200 GPUs, against PAN2's decode today (diffusers' `AutoencoderKLWan` in bf16
+under autocast, one GPU, one latent frame at a time). Memory is the peak allocated and, after it, the peak reserved by
+PyTorch's caching allocator (what nvidia-smi shows), per GPU; the default is 4 latent frames per GPU per pass:
+
+| Decoder | 8 s (193 frames) | 12 s (289 frames) | 20 s (481 frames) |
+|---|---|---|---|
+| Release, 1 GPU | 5.88 s, 12.9 / 17.1 GiB | 8.84 s, 13.9 / 19.4 GiB | 14.76 s, 15.8 / 91.5 GiB |
+| **This decoder, 1 GPU** | **1.17 s** (x5.0), 20.7 / 36.3 GiB | **1.74 s** (x5.1), 21.2 / 36.8 GiB | **2.89 s** (x5.1), 22.2 / 37.8 GiB |
+| **This decoder, 2 GPUs** | **619 ms**, 18.7 / 33.0 GiB | **929 ms**, 19.2 / 33.5 GiB | **1.54 s**, 20.2 / 34.5 GiB |
+| **This decoder, 4 GPUs** | **328 ms**, 17.8 / 41.4 GiB | **475 ms**, 18.3 / 34.1 GiB | **787 ms** (x18.8), 19.2 / 42.9 GiB |
+
+`frames_per_pass` trades memory for speed and leaves the video unchanged. On one GPU:
+
+| Latent frames per pass | 8 s | 12 s | 20 s |
+|---|---|---|---|
+| 2 | 1.25 s, 14.4 / 24.4 GiB | 1.88 s, 14.9 / 24.9 GiB | 3.13 s, 15.9 / 25.8 GiB |
+| 4 (default) | 1.17 s, 20.7 / 36.3 GiB | 1.74 s, 21.2 / 36.8 GiB | 2.89 s, 22.2 / 37.8 GiB |
+| 8 | 1.10 s, 33.7 / 62.1 GiB | 1.64 s, 34.2 / 62.6 GiB | 2.72 s, 35.2 / 63.6 GiB |
+| 16 | 1.06 s, 59.8 / 121.0 GiB | 1.58 s, 60.2 / 121.7 GiB | 2.67 s, 61.2 / 122.5 GiB |
+
+- With `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, the default reserves 30.6-32.0 GiB on one GPU and 2 frames
+  per pass 19.6-21.1 GiB, at the same speed.
+- On the optimization run's nine workloads (`tools/bench.py`: 121, 241 and 481 frames at 1, 2 and 4 GPUs, SM clocks
+  1.6-1.95 GHz), the default's geomean is 758 ms, against 1043 ms for the run's first solution (x1.38) and 698 ms at
+  16 frames per pass.
+- The CUTLASS Winograd component products take 60 % of the GPU time, the transforms and epilogues 30 %
+  (`notes/optimization.md`).
+
+## Correctness
+
+Against the fp32 decode of real latents, with `d = video - reference` per clip: every element finite and in [-1, 1],
+max |d| at most three times what the release's bf16 decode errs by, and the RMS of d likewise. On this tree, every
+clip at worlds 1, 2 and 4 and the short clips `f1`, `f5`, `f17` at worlds 1 and 4 pass, using 0.30-0.66 of the max
+allowance and 0.35-0.42 of the RMS one (PSNR 56-60 dB against fp32; the release's bf16 decode has 58-61 dB).
+
+```bash
+cd fastvideo-kernel/wan22_vae_decode
+DATA=/path/to/data python tools/bench.py --solution solution            # every workload this job has the GPUs for
+DATA=/path/to/data python tools/bench.py --workload f5x1 f17x1 --no-time  # correctness of the short clips
+cd .. && WAN22_VAE_WEIGHTS=/path/to/vae python -m pytest tests/test_wan22_vae_decode.py   # the wrapper, one GPU
+```
+
+`tools/bench.py` reads `$DATA`: the latents (`latents_<source>.pt`), the weights in `weight_names()` order
+(`decoder_weights_bf16.pt`, `weight_names.json`), the fp32 decodes (`ref_<clip>.pt`) and the release decode's errors
+(`calibration.json`); `tools/spec.py` documents each. The tests need only the VAE's weights.
+
+## Limits
+
+- sm_100a and sm_103a (B200 / GB200, B300): `tcgen05` with tensor memory, TMA and 2-CTA clusters; the loader builds
+  for the GPU it runs on. CUDA 13, torch 2.12 tested; the tests pass on B300 too, built for sm_103a.
+- 704 x 1248 video (latents 44 x 78) and batch 1; any T >= 1.
+- 1, 2 or 4 GPUs of one node were checked; the code accepts up to 8 ranks.
+- bf16 weights and activations, fp32 accumulation and fp32 RMS-norm statistics.
+- In FastVideo, `WanVAEConfig.decode_kernel` routes `AutoencoderKLWan.decode` here where it applies (on one GPU per
+  rank, `WanVAEDecoder(..., local=True)`), and PAN2 sampling turns it on (`PAN2PipelineConfig.vae_decode_kernel`). On
+  a B300 with PAN2 latents it decodes 7x faster than the fp32 decode, at PSNR 61-63 dB and SSIM 0.998 against it.
+
+## Layout
+
+- `solution/`: the kernels and their binding; `solution.json` holds the build flags. The CUTLASS and CuTe headers
+  come from fastvideo-kernel's CUTLASS submodule (`include/cutlass`, or `$CUTLASS_DIR`); the decoder's structure follows
+  diffusers' (`WAN_LICENSE.txt`).
+- `wan22_vae_decode.py`: `WanVAEDecoder`, `decoder_weights` and `weight_names`.
+- `reference.py`: the definition of the result.
+- `tools/`: `bench.py` (check, audit and time), `loader.py` (the build), `spec.py` (workloads, weights, data),
+  `audit.py`, `bounds.py`.
+- `notes/optimization.md`: the design, what each merge bought, what did not pay, and where to go next.

fastvideo-kernel/wan22_vae_decode/notes/optimization.md

@@ -0,0 +1,99 @@
+# How the decoder got here
+
+The solution is the main of a multi-agent optimization run (bifrost, 2026-10-25/26): 32 DeepSeek V4.1 Flash peers
+and two GPT-6 Astra architects on GB200, starting from a first solution written before the run. Every merge passed the
+run's CI (correct on every workload and short clip at worlds 1, 2 and 4, and at least 1 % faster than main by the
+geomean in the same job, unless waived) and an independent review. The comments in `solution/` cite the run's pull
+requests (PRn), issues (In) and jobs (Jn).
+
+## The design
+
+The decoder runs layer by layer, as kernels of its own and CUTLASS kernels instantiated in `solution/` (the
+headers are fastvideo-kernel's CUTLASS submodule). There is no megakernel: see "What did not pay" below.
+
+- **Frames in passes.** A pass carries several latent frames per GPU instead of the reference's one (4 by default,
+  `frames_per_pass`; the run measured its speed at 16, see "Memory"), so its GEMMs are large, and the causal history
+  (the last two frames of each 3 x 3 x 3 convolution's input) is carried from pass to pass as an offset, never by
+  concatenation.
+- **The 18 convolutions of up blocks 1-3** (most of the FLOPs) are F(4x4, 3x3) Winograd in space, keeping the causal
+  3-tap convolution in time, in three steps:
+  1. `wino_in.cu`: the input transform, two blocks per tile;
+  2. 36 component products, each a CUTLASS SM100 implicit-GEMM 3 x 1 x 1 convolution (`conv.cu`: 2-CTA 256 x 256 x 64
+     tiles, TMA im2col, `tcgen05` MMA, fp32 accumulation);
+  3. `wino_epilogue.cu`: the inverse transform with the bias and the residual add, and the *next* convolution's RMS
+     norm, SiLU and causal padding, so that convolution reads its input ready-made.
+- **The other convolutions** at 44 x 78 are direct CUTLASS implicit GEMMs; a few small ones still go through cuDNN
+  and cuBLAS, under 1 % of the GPU time.
+- **The upsamplers** never materialize the upsampled tensor: nearest x2 followed by a 3 x 3 convolution is four 2 x 2
+  phase convolutions at the low resolution, and `resample.cu` interleaves them with the bias and the DupUp3D shortcut
+  in one pass. The temporal upsampler (`timeconv.cu`) writes its two frames straight into place.
+- **Pointwise and the rest:** `norm.cu` (RMS norm + SiLU + causal pad in one pass, where the epilogue does not already
+  apply it), `attn.cu` (the mid block's single-head attention over each frame's 3,432 positions), and `head.cu` (the
+  output head: `conv_out` as a 3 x 1 x 1 convolution to 112 channels, then a 9-tap stencil with the bias, the clamp
+  and the unpatchify write).
+- **Over GPUs** (`comm.cu`, `prefix.cu`): the decoder's prefix (post_quant_conv, conv_in, the mid block's first
+  residual block and its attention) is split by time, each rank computing its frames with exact causal context from
+  the full latents, then transposed to bands of rows (`prefix.cu`). From there each rank owns 44 / world latent rows;
+  every 3 x 3 convolution exchanges one halo row with the neighbours only (NCCL send/recv in one group), and rank 0
+  receives the other ranks' pixels per chunk.
+
+## The merges
+
+| Merge | Change | CI speedup | Geomean after |
+|---|---|---|---|
+| kernel#0 | the first solution: Winograd up blocks, phase upsamplers, fused pointwise, row bands | | 1043 ms |
+| PR3 | 1 x 1 channel shortcuts as own kernels, their bias folded into the residual | x1.017 | 1025 ms |
+| PR24 | mid-block attention as own kernels instead of ATen SDPA (8.5 -> 2.6 ms of f121x1) | x1.013 | 1012 ms |
+| PR15 | RMS norm + SiLU + pad in one pass | x1.110 | 913 ms |
+| PR204 | halo exchange with the neighbours only, one coalesced group | x1.014 | 901 ms |
+| PR392 | the fused Winograd input transform and the fused inverse epilogue (with the next norm) | x1.189 | 757 ms |
+| PR90 | output head as a temporal convolution and a 9-tap stencil (waived at parity) | x1.003 | 755 ms |
+| PR438 | the prefix split by time before the row bands | x1.019 | 740 ms |
+| PR452 | world 8 uses all-gather, not the ring (worlds 1/2/4 unchanged; waived) | x1.001 | 740 ms |
+| PR488 | the upsampler's phases interleaved in one pass, and the temporal upsampler in place | x1.017 | 728 ms |
+| PR513 | chunks of min(16 x world, 61) latent frames instead of 8 x world | x1.023 | 711 ms |
+| PR490 | a check that no fused input leaks across chunks (waived) | x0.999 | 712 ms |
+| PR592 | no dead store of y at each residual block's first convolution | x1.010 | 704 ms |
+| PR517 | the input transform at two blocks per tile (waived at x1.007) | x1.007 | 699 ms |
+
+## Memory
+
+The run optimized time alone, and its main decoded 16 latent frames per GPU per pass at a peak of 100-102 GiB on one
+GPU, against 13-16 GiB for the release's decode, which carries one latent frame. After the run:
+
+- **Tensors are freed at their last use.** At the peak, up block 3's first convolution (512 -> 256 channels at
+  352 x 624) held its padded input after the Winograd transform had read it, the transform (2.25x the input) after
+  every product was enqueued, and the block's input after its shortcut and first norm had read it. Releasing each of
+  them then takes the peak from 101 to 60 GiB at 16 frames per pass, at the same speed, with a bit-identical video.
+- **The pass is 4 latent frames per GPU by default** (`frames_per_pass`), since memory grows with it at 3.5-4.2 GiB per
+  latent frame at 704 x 1248: 21-22 GiB on one GPU and 18-20 GiB per GPU at 2 and 4 GPUs, for 8-9 % more time than 16
+  frames per pass. The time a smaller pass costs is in the low resolutions (44 x 78 and 88 x 156), whose GEMMs do not
+  fill the GPU with few frames, while the memory is in the two high ones (176 x 312 and 352 x 624). Passes of 16 frames
+  through up block 1 and of 2-4 through up blocks 2 and 3 should keep most of the speed at this memory; it is not built.
+
+## What did not pay
+
+- **Megakernels.** Before the run, a persistent kernel per residual block measured 4.1-7.5x slower than two fused
+  per-convolution kernels. During the run, a cooperative kernel for one whole Winograd convolution (forward transform,
+  products and inverse with the next norm, in one launch of 152 CTAs in 2-CTA clusters) was correct but took
+  6.2-7.2 ms at T 8, 352 x 624, 256 channels, against 3.3 ms for the same work in three launches (0.46-0.53x). The gaps
+  between the kernels of a decode are about 5 us each, so fewer launches buy little; what a single kernel loses is
+  the occupancy and the tiling each stage picks for itself.
+- **FP8.** e4m3 operands in the Winograd component products cut the products from 262 to 162 ms at f121x1 (10-11 %
+  of the decode), but fail the error budget on all three clips (RMS at 2.3-3.3x the allowance, even with a power-of-two
+  scale per component): the transformed operands need about one more mantissa bit. Before the run, FP8 in the direct
+  convolutions passed in up blocks 1-3 but bought only 1 % of the decode, since quantization was a pass of its own.
+- **Splitting by time throughout.** A chunk needs 38 latent frames of causal context, so recomputing it is a poor fit;
+  only the prefix is split by time (PR438).
+
+## Where to go next
+
+- **The products** are about 60 % of the GPU time (59.6 % at f121x1: 1,383 CUTLASS launches, 384 ms of 644), at
+  65-68 % of the tensor cores' peak, bound by operand fetch. The transforms and the epilogues are another 30 %. A grouped products kernel (all 36 components of a convolution, persistent, in
+  2-CTA clusters) passed correctness during the run; the step after it is to compute the input transform inside the
+  products' operand producer, which removes a write and a read of the transformed tensor (2.25x the activation) per
+  convolution.
+- **Eight GPUs across two trays:** the code accepts up to 8 ranks (PR452 fixed world 8's all-gather), but world 8
+  was never in CI. A frame pipeline that passes each
+  layer's two-frame context beat row bands at eight GPUs before the run (1.40-1.55 s against 2.12 s on f241, with
+  library convolutions).

fastvideo-kernel/wan22_vae_decode/reference.py

@@ -0,0 +1,251 @@
+"""Wan2.2 VAE decoder: a standalone PyTorch reference of diffusers' AutoencoderKLWan._decode (Wan2.2-TI2V-5B VAE).
+
+decode(z, weights, dtype) maps latents z [B, 48, T, H/16, W/16] to a video [B, 3, 4T-3, H, W] with values in [-1, 1],
+as AutoencoderKLWan.decode(z).sample does with tiling and slicing off (its defaults). It needs torch and nothing else:
+`weights` is a plain dict name -> tensor of the decoder's and post_quant_conv's parameters under diffusers' names
+(the 110 of wan22_vae_decode.weight_names()), such as the VAE's state dict:
+
+    weights = {k: v.cuda().float() for k, v in vae.state_dict().items()}
+    video = decode(latents.cuda(), weights)                    # fp32
+    video = decode(latents.cuda(), weights, torch.bfloat16)    # a pipeline's bf16 decode
+
+What is computed
+----------------
+The decoder of the TI2V-5B VAE (z_dim 48, decoder_base_dim 256, dim_mult [1, 2, 4, 4], num_res_blocks 2,
+is_residual, patch_size 2, temporal upsampling in the first two up blocks, as the VAE's config.json says):
+
+1. post_quant_conv: a 1x1x1 conv 48 -> 48, over all T latent frames at once.
+2. The decoder, run on one latent frame at a time, in order (chunk i is latent frame i). A chunk sees the earlier
+   ones only through the causal caches (below). With h x w = H/16 x W/16:
+     conv_in     causal 3x3x3 conv 48 -> 1024                                                    at h x w
+     mid block   residual block 1024, attention 1024, residual block 1024                        at h x w
+     up block 0  3 residual blocks 1024; upsample3d 1024 (t x2, h w x2); + DupUp3D 1024 -> 1024     to 2h x 2w
+     up block 1  3 residual blocks 1024; upsample3d 1024 (t x2, h w x2); + DupUp3D 1024 -> 1024     to 4h x 4w
+     up block 2  residual blocks 1024 -> 512, 512, 512; upsample2d 512 (h w x2); + DupUp3D 1024 -> 512  to 8h x 8w
+     up block 3  residual blocks 512 -> 256, 256, 256; no resampler, no shortcut                  at 8h x 8w
+     norm_out    RMS norm 256, then SiLU
+     conv_out    causal 3x3x3 conv 256 -> 12                                                     at 8h x 8w
+   Up blocks 0 to 2 add their shortcut, DupUp3D of the block's input, to the output of their resampler.
+   The chunks' outputs are concatenated along time.
+3. unpatchify, [B, 12, F, 8h, 8w] -> [B, 3, F, H, W]: the 12 channels are (colour k, column offset j, row offset i)
+   of 2 x 2 pixel patches, video[b, k, f, 2y + i, 2x + j] = out[b, 4k + 2j + i, f, y, x].
+4. clamp to [-1, 1].
+
+The parts, with diffusers' class names:
+- causal conv (WanCausalConv3d), kernel kt x kh x kw, stride 1: zero padding of kh//2 and kw//2 on both sides in
+  space; in time, 2 * (kt//2) frames before the chunk and none after: the frames of the conv's cache, and zeros for
+  those the cache lacks.
+- RMS norm (WanRMS_norm) over channels: F.normalize, x / max(||x||_2, 1e-12), times sqrt(C), times gamma.
+- residual block (WanResidualBlock), in -> out channels: x1 = conv1(SiLU(norm1(x))), x2 = conv2(SiLU(norm2(x1))),
+  causal 3x3x3 convs; it returns x2 + shortcut, the shortcut being x, or a 1x1x1 conv in -> out of x when in != out.
+  (Dropout is 0, and diffusers decodes in eval mode.)
+- attention (WanAttentionBlock), per frame: n = RMS norm of x; q, k and v are the channels [0, C), [C, 2C) and
+  [2C, 3C) of a 1x1 conv C -> 3C of n; single-head attention softmax(q k^T / sqrt(C)) v over the h*w positions of the
+  frame (head dim 1024, no mask: diffusers' docstring calls it causal, it is not); a 1x1 conv C -> C; plus x.
+- upsample3d (WanResample "upsample3d"): the time conv, a causal 3x1x1 conv C -> 2C whose output channel c becomes
+  frame 2t and channel C + c frame 2t + 1, then upsample2d. upsample2d (WanResample "upsample2d"): per frame, nearest
+  x2 in h and w (mode nearest-exact, WanUpsample) and a 3x3 conv2d C -> C with zero padding 1.
+- DupUp3D, in -> out channels, factor ft in time (2 in up blocks 0 and 1, 1 in up block 2) and 2 in h and w: each
+  input channel repeated r = out * ft * 4 / in times, the copies laid out over the (ft, 2, 2) sub-positions:
+  dup[c, ft*t + a, 2y + i, 2x + j] = x[(4*ft*c + 4a + 2i + j) // r, t, y, x]. In up blocks 0 and 1 (r = 8) that is
+  nearest x2 in t, h and w; in up block 2 (r = 2) it is dup[c, t, 2y + i, 2x + j] = x[2c + i, t, y, x]: channel pairs
+  become row pairs, each value repeated in two columns.
+
+The first latent frame yields one frame
+---------------------------------------
+On the first chunk (diffusers' first_chunk):
+- upsample3d skips the time conv: the chunk's frame passes through as one frame (diffusers marks the conv's cache
+  "Rep"). The time conv's input sequence therefore begins at the second chunk, after zeros: it never sees latent
+  frame 0. (Diffusers runs the second chunk's time conv without history and caches a zero frame in front of that
+  chunk's frame; giving it no history, as here, and zero padding yields the same frames.)
+- DupUp3D keeps only the last of its ft copies in time (x[:, :, ft - 1:]): one frame, as on the main path.
+So chunk 0 gives one video frame and every later chunk four (each upsample3d doubles them): 1 + 4(T - 1) = 4T - 3.
+
+Causal caches
+-------------
+Each causal conv with a 3-frame kernel (conv_in, conv1 and conv2 of every residual block, the two time convs,
+conv_out) keeps the last two frames of its input, the activations it convolves; after a one-frame chunk, the previous
+chunk's last frame and that one. The next chunk's conv sees them before its own frames, zero padding standing in for
+what is missing (the first chunk has no history, a one-frame history lacks a frame). The loop thus computes, chunk by
+chunk, what one causal conv over the whole sequence would, with zeros before the first frame. The 1x1x1 convs
+(post_quant_conv, the residual shortcuts) need no cache.
+
+Over the whole sequence at once, then, the decoder runs every causal conv over all the frames of its level, two zero
+frames before the first; upsample3d's time conv over the frames after the first only, the first passing through; and
+DupUp3D doubling every frame but the first. That form equals the frame loop up to rounding, not bit for bit, its convs
+having other shapes (fp32: max |diff| 1.7e-5, mean 6e-7 on scripts/check_reference.py's latents).
+
+Precision
+---------
+dtype=torch.float32 computes everything in fp32, as diffusers does in fp32 (both under the same torch.backends
+settings; scripts/check_reference.py turns TF32 off).
+dtype=torch.bfloat16 computes what diffusers does with the VAE in bf16 under torch.autocast("cuda", dtype=bfloat16),
+which is how a video pipeline decodes with this VAE in bf16: the weights and the latents are rounded to bf16 (autocast
+casts the latents at post_quant_conv); every conv and the attention take bf16 operands and return bf16 (autocast's
+lower-precision ops; the kernels accumulate in fp32; for head dim 1024 PyTorch runs SDPA's memory-efficient kernel);
+the RMS norm normalizes in fp32 and rounds to bf16 before the products with sqrt(C) and gamma, each of which rounds
+again (diffusers' upcast for 16-bit inputs); the nearest upsampling runs in fp32 (WanUpsample's upcast; it only copies
+values); SiLU, the sums, DupUp3D, unpatchify and clamp run in bf16. The output is bf16. decode turns autocast off
+itself: the casts are explicit.
+
+The operations and their order are diffusers', so that the same kernels see the same operands: on a GB200 with torch
+2.12 and cuDNN 9.20, scripts/check_reference.py finds the reference bit-exact against diffusers in both modes. (The
+bf16 decode itself is max 0.07, mean 0.003 from the fp32 one there.)
+
+Credits
+-------
+Derived from the decoder of diffusers' AutoencoderKLWan (src/diffusers/models/autoencoders/autoencoder_kl_wan.py,
+diffusers 0.39.0), Copyright 2025 The Wan Team and The HuggingFace Team, licensed under the Apache License, Version 2.0
+(http://www.apache.org/licenses/LICENSE-2.0); diffusers' port of the Wan2.2 VAE of Wan-AI. Its structure, parameter
+names and caching rules are theirs. Changes: the modules are rewritten as functions of a weight dict, the cache is
+keyed by parameter name instead of a running index, the "Rep" marker is expressed as the time conv starting at the
+second chunk, and the encoder, tiling and slicing are left out.
+"""
+
+import torch
+import torch.nn.functional as F
+
+CACHE_T = 2  # frames of history kept by a causal conv with a 3-frame kernel
+PATCH_SIZE = 2
+RESNETS_PER_UP_BLOCK = 3  # num_res_blocks + 1
+# decoder.up_blocks.{i}: (in channels, out channels, temporal upsampling, spatial upsampling)
+UP_BLOCKS = ((1024, 1024, True, True), (1024, 1024, True, True), (1024, 512, False, True), (512, 256, False, False))
+PREFIXES = ("decoder.", "post_quant_conv.")
+
+
+@torch.no_grad()
+def decode(z, weights, dtype=torch.float32):
+    """Decodes latents z [B, 48, T, H/16, W/16] into a video [B, 3, 4T-3, H, W] in [-1, 1], of `dtype`.
+
+    `weights` maps diffusers' parameter names to tensors; those of the decoder and post_quant_conv are cast to `dtype`
+    on z's device (no copy when they already are). dtype is torch.float32 or torch.bfloat16 (see the module docstring).
+    """
+    if dtype not in (torch.float32, torch.bfloat16):
+        raise ValueError(f"dtype must be torch.float32 or torch.bfloat16, not {dtype}")
+    weights = {name: t.to(z.device, dtype) for name, t in weights.items() if name.startswith(PREFIXES)}
+
+    with torch.autocast(z.device.type, enabled=False):
+        x = causal_conv(z.to(dtype), weights, "post_quant_conv")
+        cache = {}
+        chunks = [decode_chunk(x[:, :, i : i + 1], weights, cache, first=i == 0) for i in range(x.shape[2])]
+        return unpatchify(torch.cat(chunks, dim=2), PATCH_SIZE).clamp(-1.0, 1.0)
+
+
+def decode_chunk(x, weights, cache, first):
+    """The decoder on one latent frame x [B, 48, 1, h, w]: [B, 12, 1, 8h, 8w] on the first chunk, [B, 12, 4, 8h, 8w]
+    on the others. `cache` holds the causal convs' history, by conv name, and is updated."""
+    x = cached_conv(x, weights, "decoder.conv_in", cache)
+    x = residual_block(x, weights, "decoder.mid_block.resnets.0", cache)
+    x = attention(x, weights, "decoder.mid_block.attentions.0")
+    x = residual_block(x, weights, "decoder.mid_block.resnets.1", cache)
+    for i in range(len(UP_BLOCKS)):
+        x = up_block(x, weights, i, cache, first)
+    x = F.silu(rms_norm(x, weights["decoder.norm_out.gamma"]))
+    return cached_conv(x, weights, "decoder.conv_out", cache)
+
+
+def up_block(x, weights, i, cache, first):
+    """WanResidualUpBlock i: the residual blocks, then (up blocks 0 to 2) the resampler plus DupUp3D of x."""
+    _, out_channels, temporal, spatial = UP_BLOCKS[i]
+    name = f"decoder.up_blocks.{i}"
+    block_input = x
+    for j in range(RESNETS_PER_UP_BLOCK):
+        x = residual_block(x, weights, f"{name}.resnets.{j}", cache)
+    if not spatial:
+        return x
+
+    # The first chunk skips the time conv, whose history therefore starts empty at the second chunk.
+    if temporal and not first:
+        x = time_upsample(x, weights, f"{name}.upsampler.time_conv", cache)
+    x = spatial_upsample(x, weights, f"{name}.upsampler.resample.1")
+    return x + dup_up(block_input, out_channels, 2 if temporal else 1, 2, first)
+
+
+def residual_block(x, weights, name, cache):
+    """WanResidualBlock: conv2(SiLU(norm2(conv1(SiLU(norm1(x)))))) plus x, or plus a 1x1x1 conv of x."""
+    shortcut = causal_conv(x, weights, f"{name}.conv_shortcut") if f"{name}.conv_shortcut.weight" in weights else x
+    x = cached_conv(F.silu(rms_norm(x, weights[f"{name}.norm1.gamma"])), weights, f"{name}.conv1", cache)
+    x = cached_conv(F.silu(rms_norm(x, weights[f"{name}.norm2.gamma"])), weights, f"{name}.conv2", cache)
+    return x + shortcut
+
+
+def attention(x, weights, name):
+    """WanAttentionBlock: single-head attention over the h*w positions of each frame, plus x."""
+    identity = x
+    b, c, t, h, w = x.shape
+    x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+    x = rms_norm(x, weights[f"{name}.norm.gamma"])
+    qkv = F.conv2d(x, weights[f"{name}.to_qkv.weight"], weights[f"{name}.to_qkv.bias"])
+    qkv = qkv.reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous()  # [b*t, 1 head, h*w, 3C]
+    q, k, v = qkv.chunk(3, dim=-1)
+    x = F.scaled_dot_product_attention(q, k, v)
+    x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+    x = F.conv2d(x, weights[f"{name}.proj.weight"], weights[f"{name}.proj.bias"])
+    return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4) + identity
+
+
+def time_upsample(x, weights, name, cache):
+    """upsample3d's time conv: a causal 3x1x1 conv C -> 2C; channel c becomes frame 2t, channel C + c frame 2t + 1."""
+    b, c, t, h, w = x.shape
+    x = cached_conv(x, weights, name, cache)
+    x = x.reshape(b, 2, c, t, h, w)
+    x = torch.stack((x[:, 0], x[:, 1]), 3)
+    return x.reshape(b, c, t * 2, h, w)
+
+
+def spatial_upsample(x, weights, name):
+    """Per frame: nearest x2 in h and w (in fp32, as WanUpsample does; it only copies values), then a 3x3 conv2d."""
+    b, c, t, h, w = x.shape
+    x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+    x = F.interpolate(x.float(), scale_factor=(2.0, 2.0), mode="nearest-exact").type_as(x)
+    x = F.conv2d(x, weights[f"{name}.weight"], weights[f"{name}.bias"], padding=1)
+    return x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
+
+
+def dup_up(x, out_channels, factor_t, factor_s, first):
+    """DupUp3D: each channel repeated r = out * ft * fs * fs / in times, the copies spread over the (ft, fs, fs)
+    sub-positions of time, rows and columns; on the first chunk only the last of the ft copies in time is kept."""
+    b, c, t, h, w = x.shape
+    x = x.repeat_interleave(out_channels * factor_t * factor_s * factor_s // c, dim=1)
+    x = x.view(b, out_channels, factor_t, factor_s, factor_s, t, h, w)
+    x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    x = x.view(b, out_channels, t * factor_t, h * factor_s, w * factor_s)
+    return x[:, :, factor_t - 1 :] if first else x
+
+
+def rms_norm(x, gamma):
+    """WanRMS_norm over channels (dim 1): F.normalize, in fp32 for a 16-bit x and rounded back, times sqrt(C), times
+    gamma."""
+    x_hat = F.normalize(x.float() if x.dtype in (torch.float16, torch.bfloat16) else x, dim=1).to(x.dtype)
+    return x_hat * x.shape[1] ** 0.5 * gamma
+
+
+def cached_conv(x, weights, name, cache):
+    """causal_conv of one chunk with the history in cache[name], which it replaces by the last CACHE_T frames of its
+    input: after a one-frame chunk, the previous chunk's last frame and this one."""
+    history = cache.get(name)
+    keep = x[:, :, -CACHE_T:].clone()
+    if keep.shape[2] < CACHE_T and history is not None:
+        keep = torch.cat([history[:, :, -1:], keep], dim=2)
+    cache[name] = keep
+    return causal_conv(x, weights, name, history)
+
+
+def causal_conv(x, weights, name, history=None):
+    """WanCausalConv3d `name` on x [B, C, t, h, w]: zero padding of kh//2 and kw//2 on both sides in space and, in
+    time, 2 * (kt//2) frames before x: `history` (the frames before x, at most that many), then zeros."""
+    weight, bias = weights[f"{name}.weight"], weights[f"{name}.bias"]
+    kt, kh, kw = weight.shape[2:]
+    pad = [kw // 2, kw // 2, kh // 2, kh // 2, 2 * (kt // 2), 0]
+    if history is not None and pad[4] > 0:
+        x = torch.cat([history, x], dim=2)
+        pad[4] -= history.shape[2]
+    return F.conv3d(F.pad(x, pad), weight, bias)
+
+
+def unpatchify(x, p):
+    """[B, 3*p*p, F, h, w] -> [B, 3, F, p*h, p*w]: channel (k*p + j)*p + i is colour k at row offset i, column j."""
+    b, cp, f, h, w = x.shape
+    x = x.view(b, cp // (p * p), p, p, f, h, w)
+    x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous()
+    return x.view(b, cp // (p * p), f, h * p, w * p)

fastvideo-kernel/wan22_vae_decode/solution/attn.cu

@@ -0,0 +1,217 @@
+// Mid-block self-attention with our own kernels, replacing ATen's
+// scaled_dot_product_attention (fmha_cutlassF, 8.5 ms of f121x1's 954 ms).
+//
+// The reference computes, per latent frame, one head of softmax(q k^T/32) v over
+// the H*W=3432 positions with C=1024 channels and no mask. ATen's path cannot use
+// flash attention here (head dim 1024 > 256) and falls back to the xformers
+// cutlass kernel, which reads and writes far more per frame than the products
+// need: the decode's whole attention is 1.5 TFLOP (0.63 ms at 16-bit peak for
+// f121x1) against 8.5 ms measured, ~13x off.
+//
+// Three stages, each at its own bound:
+//   1. pack: q,k into contiguous [HW,C] per frame and v transposed into the
+//      [C,PAD] weight layout the convolution GEMM wants (~21 MB/frame).
+//   2. S = q k^T per frame (fp32 out) and O = p v per frame (bf16 out), both as
+//      1x1x1 implicit GEMMs through the vendored CUTLASS SM100 convolution
+//      collective (2-CTA tcgen05 MMA, TMA, fp32 accumulation) -- the same
+//      machinery conv.cu uses, so no new GEMM path is introduced.
+//   3. row softmax: S row -> exp(S/sqrt(C) - max)/sum in bf16, one CTA per row,
+//      the row held in shared memory (one read of S, one write of P).
+// P's rows are padded to a multiple of 64 columns and v's transpose is zero
+// there, so the GEMM's tail contributes nothing.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include "attn.h"
+
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/convnd_problem_shape.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/conv/dispatch_policy.hpp"
+#include "cutlass/conv/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/conv/device/conv_universal_adapter.hpp"
+#include "cutlass/conv/kernel/conv_universal.hpp"
+
+using namespace cute;
+using ElementBF16 = cutlass::bfloat16_t;
+using LayoutNDHWC = cutlass::layout::TensorNDHWC;
+using ArchTag = cutlass::arch::Sm100;
+using OperatorClass = cutlass::arch::OpClassTensorOp;
+constexpr auto Fprop = cutlass::conv::Operator::kFprop;
+
+// One 1x1x1 implicit GEMM, y = x (per position) times w, fp32 accumulation and
+// ElementD out. x: channels-last3d [1,C,D,1,1] (may be a strided view); w:
+// [K,C,1,1,1] with explicit strides; y: packed channels-last3d [1,K,D,1,1].
+template <class ElementD, int M, int N, int CL>
+static void gemm1x1(const at::Tensor& x, const at::Tensor& w, const at::Tensor& y) {
+  using Tile = Shape<Int<M>, Int<N>, Shape<_64>>;
+  using Cluster = Shape<Int<CL>, _1, _1>;
+  using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+      float, float, ElementD, LayoutNDHWC, 8, ElementD, LayoutNDHWC, 8,
+      cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+  using Mainloop = typename cutlass::conv::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Fprop, ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, float,
+      Tile, Cluster,
+      cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
+      cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+  using Problem = cutlass::conv::ConvProblemShape<Fprop, 3>;
+  using Kernel = cutlass::conv::kernel::ConvUniversal<Problem, Mainloop, Epilogue>;
+  using Conv = cutlass::conv::device::ConvUniversalAdapter<Kernel>;
+  int C = x.size(1), D = x.size(2);
+  int K = w.size(0);
+  TORCH_CHECK(y.size(1) == K && y.size(2) == D, "attn gemm: output shape mismatch");
+  Problem problem(cutlass::conv::Mode::kCrossCorrelation,
+                  {1, D, 1, 1, C}, {x.stride(0), x.stride(2), x.stride(3), x.stride(4), x.stride(1)},
+                  {K, 1, 1, 1, C}, {w.stride(0), w.stride(2), w.stride(3), w.stride(4), w.stride(1)},
+                  {0, 0, 0}, {0, 0, 0}, {1, 1, 1}, {1, 1, 1}, 1);
+  typename Kernel::StrideC sc{};
+  typename Kernel::StrideD sd{};
+  cute::for_each(cute::make_seq<cute::rank<0>(sc)>{}, [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
+  cute::for_each(cute::make_seq<cute::rank<0>(sd)>{}, [&](auto i) { cute::get<0, i>(sd) = problem.stride_C[Problem::RankT - 2 - i]; });
+  typename Conv::Arguments args{problem, {(const ElementBF16*)x.data_ptr(), (const ElementBF16*)w.data_ptr()},
+                                {{1.f, 0.f}, (const ElementD*)y.data_ptr(), sc, (ElementD*)y.data_ptr(), sd}};
+  Conv conv;
+  TORCH_CHECK(conv.can_implement(args) == cutlass::Status::kSuccess,
+              "attn gemm: CUTLASS cannot implement D=", D, " N=", K, " C=", C);
+  size_t size = Conv::get_workspace_size(args);
+  auto workspace = at::empty({(int64_t)size}, x.options().dtype(at::kByte));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  TORCH_CHECK(conv.initialize(args, workspace.data_ptr(), stream) == cutlass::Status::kSuccess,
+              "attn gemm: initialization failed");
+  TORCH_CHECK(conv.run(stream) == cutlass::Status::kSuccess, "attn gemm: run failed");
+}
+
+// Channels-last-3d views with explicit strides: a size-1 dim built by unsqueeze
+// carries stride 1, which the im2col TMA descriptor rejects.
+static at::Tensor nhwc_in(const at::Tensor& m, int64_t C, int64_t D) {  // [D,C] -> [1,C,D,1,1]
+  return m.as_strided({1, C, D, 1, 1}, {C * D, 1, C, C, C});
+}
+static at::Tensor nhwc_w(const at::Tensor& m, int64_t K, int64_t C) {  // [K,C] -> [K,C,1,1,1]
+  return m.as_strided({K, C, 1, 1, 1}, {C, 1, C, C, C});
+}
+static at::Tensor nhwc_out(const at::Tensor& m, int64_t K, int64_t D) {  // [D,K] -> [1,K,D,1,1]
+  return m.as_strided({1, K, D, 1, 1}, {K * D, 1, K, K, K});
+}
+
+// v [T,1,HW,3C] at offset 2C -> vt [T,C,PAD] (transposed, [T,C,PAD] zeroed first).
+static __global__ void transpose_v_kernel(const __nv_bfloat16* __restrict__ qkv, int T, int HW, int C, int PAD,
+                                   __nv_bfloat16* __restrict__ vt) {
+  __shared__ __nv_bfloat16 tile[32][33];
+  int i0 = blockIdx.x * 32, c0 = blockIdx.y * 32, t = blockIdx.z;
+  {
+    int i = i0 + threadIdx.y, c = c0 + threadIdx.x;
+    if (i < HW && c < C) tile[threadIdx.y][threadIdx.x] = qkv[((size_t)t * HW + i) * (3 * C) + 2 * C + c];
+  }
+  __syncthreads();
+  {
+    int i = i0 + threadIdx.x, c = c0 + threadIdx.y;
+    if (i < HW && c < C) vt[((size_t)t * C + c) * PAD + i] = tile[threadIdx.x][threadIdx.y];
+  }
+}
+
+// One CTA per row: p = exp(s*scale - max)/sum in bf16, then zero the padded tail.
+static __global__ void row_softmax_kernel(const float* __restrict__ s, int cols, int PAD, float scale,
+                                   __nv_bfloat16* __restrict__ p) {
+  extern __shared__ float row[];
+  const int n4 = cols / 4;
+  const float4* src = (const float4*)(s + (size_t)blockIdx.x * cols);
+  __nv_bfloat16* dst = p + (size_t)blockIdx.x * PAD;
+  float4* r4 = (float4*)row;
+  for (int j = threadIdx.x; j < n4; j += blockDim.x) {
+    float4 v = src[j];
+    v.x *= scale;
+    v.y *= scale;
+    v.z *= scale;
+    v.w *= scale;
+    r4[j] = v;
+  }
+  __syncthreads();
+  float m = -INFINITY;
+  for (int j = threadIdx.x; j < n4; j += blockDim.x) {
+    float4 v = r4[j];
+    m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w)));
+  }
+  for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off));
+  __shared__ float red[32];
+  if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = m;
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    m = threadIdx.x < blockDim.x / 32 ? red[threadIdx.x] : -INFINITY;
+    for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off));
+    if (threadIdx.x == 0) red[0] = m;
+  }
+  __syncthreads();
+  m = red[0];
+  __syncthreads();
+  float sum = 0.f;
+  for (int j = threadIdx.x; j < n4; j += blockDim.x) {
+    float4 v = r4[j];
+    v.x = __expf(v.x - m);
+    v.y = __expf(v.y - m);
+    v.z = __expf(v.z - m);
+    v.w = __expf(v.w - m);
+    r4[j] = v;
+    sum += (v.x + v.y) + (v.z + v.w);
+  }
+  for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, off);
+  if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = sum;
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    sum = threadIdx.x < blockDim.x / 32 ? red[threadIdx.x] : 0.f;
+    for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, off);
+    if (threadIdx.x == 0) red[0] = sum;
+  }
+  __syncthreads();
+  const float inv = 1.f / red[0];
+  const int npair = cols / 2;
+  for (int j = threadIdx.x; j < npair; j += blockDim.x)
+    *(__nv_bfloat162*)(dst + 2 * j) = __floats2bfloat162_rn(row[2 * j] * inv, row[2 * j + 1] * inv);
+  for (int j = cols + threadIdx.x; j < PAD; j += blockDim.x) dst[j] = __float2bfloat16(0.f);
+}
+
+at::Tensor attn_forward(const at::Tensor& qkv) {
+  TORCH_CHECK(qkv.dim() == 4 && qkv.size(1) == 1 && qkv.size(3) % 3 == 0, "attn: expected [T,1,HW,3C]");
+  int T = qkv.size(0), HW = qkv.size(2), C = qkv.size(3) / 3;
+  TORCH_CHECK(C % 64 == 0, "attn: channels must be a multiple of 64");
+  int PAD = (HW + 63) / 64 * 64;
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto bf = qkv.options();
+  auto vt = at::zeros({T, C, PAD}, bf);
+  auto s = at::empty({T, HW, HW}, bf.dtype(at::kFloat));
+  auto p = at::empty({T, HW, PAD}, bf);
+  auto o = at::empty({T, 1, HW, C}, bf);
+  const __nv_bfloat16* qkv_p = (const __nv_bfloat16*)qkv.data_ptr();
+
+  transpose_v_kernel<<<dim3((HW + 31) / 32, (C + 31) / 32, T), dim3(32, 32), 0, stream>>>(
+      qkv_p, T, HW, C, PAD, (__nv_bfloat16*)vt.data_ptr());
+
+  // S = q k^T : M = positions, N = keys, K = channels. q and k are read in place:
+  // their channel stride is 1 and their position stride 3C, which TMA accepts.
+  for (int t = 0; t < T; ++t) {
+    auto qv = qkv.select(0, t).select(0, 0);  // [HW, 3C]
+    auto x = qv.narrow(1, 0, C).as_strided({1, C, HW, 1, 1}, {(int64_t)3 * C * HW, 1, 3 * C, 3 * C, 3 * C});
+    auto w = qv.narrow(1, C, C).as_strided({HW, C, 1, 1, 1}, {3 * C, 1, 3 * C, 3 * C, 3 * C});
+    auto y = nhwc_out(s.select(0, t), HW, HW);
+    gemm1x1<float, 256, 256, 2>(x, w, y);
+  }
+
+  float* s_p = (float*)s.data_ptr();
+  float scale = 1.f / sqrtf((float)C);
+  row_softmax_kernel<<<T * HW, 256, HW * sizeof(float), stream>>>(
+      s_p, HW, PAD, scale, (__nv_bfloat16*)p.data_ptr());
+
+  // O = p v : M = positions, N = channels, K = keys (padded).
+  for (int t = 0; t < T; ++t) {
+    auto x = nhwc_in(p.select(0, t), PAD, HW);
+    auto w = nhwc_w(vt.select(0, t), C, PAD);
+    auto y = nhwc_out(o.select(0, t).select(0, 0), C, HW);
+    gemm1x1<ElementBF16, 256, 256, 2>(x, w, y);
+  }
+  return o;
+}

fastvideo-kernel/wan22_vae_decode/solution/attn.h

@@ -0,0 +1,6 @@
+#pragma once
+#include <ATen/ATen.h>
+// One head of self-attention over the H*W positions of every latent frame of
+// one decoder chunk: qkv is [T,1,HW,3C] bf16 (q, k, v concatenated on the last
+// dim, as binding.cpp's attention() builds it); the result is [T,1,HW,C] bf16.
+at::Tensor attn_forward(const at::Tensor& qkv);

fastvideo-kernel/wan22_vae_decode/solution/binding.cpp

@@ -0,0 +1,459 @@
+// Wan2.2 VAE decoder seed. Derived from the reference architecture and
+// indexing in diffusers' AutoencoderKLWan (Wan Team / Hugging Face Team,
+// Apache-2.0; see WAN_LICENSE.txt). Native CUDA execution is implemented here.
+// All invocation state, including temporal history, belongs to Decoder below.
+#include <torch/extension.h>
+#include <ATen/ops/cudnn_convolution.h>
+#include <ATen/ops/scaled_dot_product_attention.h>
+#include <ATen/record_function.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include "kernels.h"
+#include "prefix.h"
+#include "comm.h"
+#include "resample.h"
+#include "attn.h"
+#include "head.h"
+#include <array>
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+
+at::Tensor custom_conv(const at::Tensor&,const at::Tensor&,int64_t,int64_t,int64_t);
+
+// The temporal upsampler, writing the frame-interleaved video directly (solution/timeconv.cu).
+at::Tensor time_upsample(const at::Tensor&,const at::Tensor&,const at::Tensor&,const at::Tensor&,int64_t,int64_t,bool);
+
+using PG = c10::intrusive_ptr<c10d::ProcessGroup>;
+using at::Tensor;
+using Weights = std::vector<Tensor>;
+
+Weights prepare(const Weights& weights,int64_t rank,int64_t world,const PG& pg) {
+  TORCH_CHECK(weights.size()==110,"Expected the 110 decoder tensors");
+  at::NoGradGuard guard;
+  Weights result;
+  result.reserve(weights.size());
+  for(const auto& w:weights) {
+    // A 1x1x1 tensor is "trivially channels-last", so contiguous() leaves it with unit strides on
+    // its size-1 spatial dims; CUTLASS's im2col TMA needs those strides to be 16-byte multiples.
+    // The same bytes with dense strides: as_strided keeps the NDHWC order but gives every spatial
+    // dim the channel-count stride.
+    if(w.dim()==5 && w.size(2)==1 && w.size(3)==1 && w.size(4)==1) {
+      auto wc=w.contiguous();int64_t C=wc.size(1);
+      result.push_back(wc.as_strided({wc.size(0),C,1,1,1},{C,1,C,C,C}));
+    }
+    else if(w.dim()==5) result.push_back(w.contiguous(at::MemoryFormat::ChannelsLast3d));
+    else if(w.dim()==4 && w.size(2)>0 && w.size(1)>1)
+      result.push_back(w.unsqueeze(2).contiguous(at::MemoryFormat::ChannelsLast3d));
+    else result.push_back(w.contiguous());
+  }
+  // The output head: conv_out's nine spatial taps become the output channels of a 3x1x1
+  // (temporal-only) convolution, so its 112 outputs are (tap i,j) x (channel c) and head.cu's
+  // epilogue sums the nine shifted copies. W'[(i*3+j)*12+c, ch, t] = W[c, ch, t, i, j]; the four
+  // channels past 108 are zero and are never read. N=112 rather than 108 keeps the epilogue's
+  // channel count a multiple of 16.
+  {
+    auto orig=result[108];
+    auto t=orig.permute({3,4,0,1,2}).contiguous();          // [i,j,c,ch,t] = [3,3,12,256,3]
+    auto filled=at::zeros({112,256,3},orig.options());      // [n,ch,t], n=(i*3+j)*12+c
+    filled.narrow(0,0,108).copy_(t.reshape({108,256,3}));
+    auto kt=filled.transpose(1,2).contiguous();             // [n,t,ch] contiguous
+    // The layout a [K,C,3,3,3] weight has under ChannelsLast3d: C innermost with stride 1, every
+    // other stride a multiple of 16 bytes. torch normalises a size-1 H/W away, hence as_strided.
+    result[108]=kt.transpose(1,2).as_strided({112,256,3,1,1},{768,1,256,256,256});
+  }
+  // Nearest-x2 + 3x3 convolution equals four 2x2 low-resolution phases.
+  // Form the sums in fp32 and store the transformed weights in bf16.
+  for(int spatial : {39,61,85}) {
+    auto orig=result[spatial].to(at::kFloat);
+    int C=orig.size(0);
+    for(int a=0;a<2;++a)for(int b=0;b<2;++b) {
+      auto transformed=at::zeros({C,C,1,2,2},orig.options());
+      for(int u=0;u<2;++u)for(int v=0;v<2;++v) {
+        auto sum=at::zeros({C,C},orig.options());
+        for(int r=0;r<3;++r)for(int c=0;c<3;++c) {
+          int dy=a+r-1,dx=b+c-1;
+          int iy=dy<0?-1:dy/2,ix=dx<0?-1:dx/2;
+          if(iy==u+a-1&&ix==v+b-1)sum.add_(orig.select(2,0).select(2,r).select(2,c));
+        }
+        transformed.select(2,0).select(2,u).select(2,v).copy_(sum);
+      }
+      result.push_back(transformed.to(at::kBFloat16).contiguous(at::MemoryFormat::ChannelsLast3d));
+    }
+  }
+  // F(4x4,3x3) Winograd weights for up blocks 1, 2, and 3.
+  // The 36 components each retain the causal 3-tap temporal convolution.
+  const float G[6][3]={{.25f,0,0},{-1.f/6,-1.f/6,-1.f/6},{-1.f/6,1.f/6,-1.f/6},{1.f/24,1.f/12,1.f/6},{1.f/24,-1.f/12,1.f/6},{0,0,1}};
+  for(int idx : {44,47,50,53,56,59,66,69,74,77,80,83,88,91,96,99,102,105}) {
+    auto orig=result[idx].to(at::kFloat);
+    for(int a=0;a<6;a++)for(int b=0;b<6;b++) {
+      auto sum=at::zeros({orig.size(0),orig.size(1),3},orig.options());
+      for(int r=0;r<3;r++)for(int c=0;c<3;c++)if(G[a][r]*G[b][c]!=0)
+        sum.add_(orig.select(3,r).select(3,c),G[a][r]*G[b][c]);
+      result.push_back(sum.unsqueeze(3).unsqueeze(4).to(at::kBFloat16).contiguous(at::MemoryFormat::ChannelsLast3d));
+    }
+  }
+  // The 1x1x1 shortcut of a channel-changing resnet is added to the second convolution's
+  // bias + output, out = round(acc2 + b2) + round(acc1 + b1). Fold b1 into b2 instead (one more
+  // bf16 rounding of the bias sum and none of the operand): the shortcut tensor then needs no
+  // separate bias pass, which was a full read+write of it.
+  for(auto pair : {std::pair<int,int>{72,70},std::pair<int,int>{94,92}}) {
+    auto sum=result[pair.second].to(at::kFloat)+result[pair.first].to(at::kFloat);
+    result[pair.second]=sum.to(at::kBFloat16).contiguous();
+  }
+  return result;
+}
+
+struct Decoder {
+  // This must not implicitly bind to a fused-epilogue weight index when
+  // independently extended conv() interfaces are merged.
+  struct TemporalTrim {
+    int frames;
+    explicit constexpr TemporalTrim(int n) : frames(n) {}
+  };
+
+  const Weights& w;
+  int rank,world;
+  PG pg;
+  std::array<Tensor,110> cache;
+  bool sharded=false;
+  // A convolution whose input was already normalized, activated and padded by the previous
+  // Winograd convolution's fused epilogue: (tensor, the conv index it is prepared for, and the
+  // sharding it was prepared under — a mismatch there would be silently wrong, not a crash).
+  Tensor prepared; int prepared_for=-1; bool prepared_sharded=false;
+  // The conv index whose y store was suppressed because this conv reserved it for `prepared`.
+  // It is the RUNTIME half of the y_dead contract: the producer states an intent and this records
+  // it, so a later edit that stops the consumer taking the `pre` edge aborts instead of handing a
+  // never-written tensor to norm_pad. Cleared when the edge is taken, and asserted closed at the
+  // end of every chunk beside prepared_for, because a conv index restarts at the same values.
+  int y_suppressed_for=-1;
+  Decoder(const Weights& weights,int r,int n,const PG& p):w(weights),rank(r),world(n),pg(p){}
+
+  Tensor conv(Tensor x,int idx,int gamma=-1,bool activate=true,bool add_bias=true,const Tensor& residual={},TemporalTrim trim=TemporalTrim{0},int fg=-1,int fnext=-1,bool y_dead=false) {
+    RECORD_FUNCTION("wan.conv."+std::to_string(idx), std::vector<c10::IValue>());
+    const int temporal_trim=trim.frames;
+    int kt=w[idx].size(2),kh=w[idx].size(3),kw=w[idx].size(4);
+    TORCH_CHECK(temporal_trim>=0 && temporal_trim<kt,"Invalid causal prefix trim");
+    int pt=kt-1-temporal_trim,ph=(sharded&&world>1&&kh==3)?1:0;
+    Tensor p=x;
+    bool pre=(prepared_for==idx);
+    // The consumer asserts the producer's y_dead contract AT RUNTIME, not in prose. If a conv had
+    // its y store suppressed for this idx and this conv does not take the fused `prepared` edge,
+    // it is about to read a tensor that was never written -- so abort instead of producing a wrong
+    // video. The existing metadata checks at `pre` cannot catch this, because they live inside the
+    // branch that is not being taken.
+    TORCH_CHECK(pre || y_suppressed_for!=idx,
+                "conv ",idx,": its y store was suppressed FOR THIS CONV, which then did not take "
+                "the prepared edge; the tensor it is about to read is unwritten");
+    // Prefix pruning never uses the fused, full-length prepared-input edge.
+    // A future caller must implement a trimmed prepared layout before mixing them.
+    TORCH_CHECK(temporal_trim==0 || (!pre && fg<0 && fnext<0),
+                "Causal prefix trim cannot use or produce a fused prepared input");
+    if(pre) {
+      TORCH_CHECK(prepared.defined(),"conv ",idx," is marked prepared but no tensor was stored");
+      TORCH_CHECK(prepared_sharded==sharded,"prepared tensor crosses a change of sharding");
+      // The graph edge checks itself. pt/ph here are the consumer's own, read from this conv's
+      // weight, and the producer computed the same two from w[fnext] with fnext==idx -- so these
+      // equalities are true by construction today and are a proof obligation on any later edit:
+      // a wrong fnext would otherwise hand a wrong-shaped tensor to the convolution, and a
+      // wrong-but-same-shaped one would pass every check the kernel makes.
+      TORCH_CHECK(prepared.size(1)==x.size(1) && prepared.size(2)==x.size(2)+pt &&
+                  prepared.size(3)==x.size(3)+2*ph && prepared.size(4)==x.size(4),
+                  "conv ",idx,": prepared ",prepared.sizes()," does not fit input ",x.sizes(),
+                  " (pt ",pt,", ph ",ph,", sharded ",(int)sharded,")");
+      p=prepared; prepared=Tensor(); prepared_for=-1; prepared_sharded=false;
+      y_suppressed_for=-1;   // the suppressed y was reserved for this edge and this edge took it
+    }
+    else if(pt||ph||gamma>=0) p=norm_pad(x,gamma>=0?w[gamma]:Tensor(),cache[idx],pt,ph,activate);
+    // Memory: the unpadded input is dead once the padded copy exists (p holds it when no copy was made).
+    x=Tensor();
+    // The causal cache carries this convolution's input history to the next chunk. `p` may be the
+    // tensor the previous fused epilogue produced and this conv has just consumed, so the copy is
+    // ordered after that kernel by the stream they share -- a side stream here would race it.
+    if(pt) cache[idx]=last_frames(p,pt,ph);
+    if(ph) {
+      RECORD_FUNCTION("wan.halo",std::vector<c10::IValue>());
+      auto recv=halo_exchange(halo_pack(p),pg,rank,world);
+      halo_unpack(p,recv,rank,world);
+    }
+    const int Ho=p.size(3)-2*ph,Wo=p.size(4);
+    std::vector<int64_t> padding={0,ph?0:kh/2,kw/2};
+    bool use_custom=(p.size(4)==78 && kt==3 && kh==3 && w[idx].size(0)==1024 && w[idx].size(1)==1024) || idx==61 || idx==85
+                    || idx==71 || idx==93;
+    // Only the expensive spatial 3x3 convolutions use Winograd. Low-resolution
+    // convolutions remain direct implicit GEMMs; 1x1 and temporal-only layers
+    // use cuDNN. Each transformed GEMM accumulates bf16 operands in fp32.
+    int wi=-1;
+    const int wids[18]={44,47,50,53,56,59,66,69,74,77,80,83,88,91,96,99,102,105};
+    for(int i=0;i<18;i++)if(idx==wids[i])wi=i;
+    Tensor y;
+    if(wi>=0) {
+      auto transformed=winograd_input_fast(p,ph);
+      // Memory: the padded input is dead once transformed (its history went to the cache above), and the transform
+      // (2.25x the input) once every component product is enqueued; the caching allocator reuses both in stream order.
+      p=Tensor();
+      std::vector<Tensor> products;
+      for(int a=0;a<36;a++) {
+        auto input=transformed.select(0,a).unsqueeze(0).permute({0,4,1,2,3});
+        products.push_back(custom_conv(input,w[122+36*wi+a],0,0,3));
+      }
+      transformed=Tensor();
+      // The SiLU this epilogue applies is the CONSUMER's norm activation, not this conv's
+      // `activate`: fusing is valid only when the consumer would have activated. Every fused
+      // consumer in the decoder (the two convs of a residual block) has activate=true, so the
+      // test below holds by construction; a non-activating consumer would have to pass its own.
+      bool fuse=(fg>=0&&fnext>=0&&activate);
+      // write_y is false exactly where the caller asserts nobody reads this launch's y AND the
+      // fused `prepared` edge exists to carry the activation instead. It is recorded so the
+      // consumer can ASSERT the edge rather than trust it; see y_suppressed_for above.
+      const bool write_y=!(y_dead && fuse);
+      int pt2=pt,ph2=ph;
+      if(fuse) { pt2=w[fnext].size(2)-1; ph2=(sharded&&world>1&&w[fnext].size(3)==3)?1:0; }
+      auto out=wino_epilogue(products,Ho,Wo,add_bias?w[idx+1]:Tensor(),residual,
+                             fuse?w[fg]:Tensor(),fuse?cache[fnext]:Tensor(),pt2,ph2,write_y);
+      y=out.y;
+      if(fuse) { prepared=out.n; prepared_for=fnext; prepared_sharded=sharded;
+                 if(!write_y) y_suppressed_for=fnext; }
+    } else y=use_custom?custom_conv(p,w[idx],padding[1],padding[2],3):at::cudnn_convolution(p,w[idx],padding,{1,1,1},{1,1,1},1,false,false,false);
+    if(add_bias && wi<0) finish(y,w[idx+1],residual);
+    return y;
+  }
+
+  Tensor resnet(Tensor x,int base,bool change=false,int nx=-1,int nxg=-1) {
+    RECORD_FUNCTION("wan.resnet."+std::to_string(base),std::vector<c10::IValue>());
+    Tensor skip=change?conv(x,base+6,-1,true,false):x;
+    // Memory: moved into the first convolution, the block's input is released there once its norm has read it
+    // (skip still holds it when it is the residual).
+    auto y=conv(std::move(x),base+1,base,true,true,Tensor(),TemporalTrim{0},base+3,base+4, /*y_dead=*/true);
+    return conv(y,base+4,base+3,true,true,skip,TemporalTrim{0},nxg,nx);
+  }
+
+  Tensor attention(const Tensor& x) {
+    RECORD_FUNCTION("wan.attention",std::vector<c10::IValue>());
+    int T=x.size(2),H=x.size(3),W=x.size(4),C=x.size(1);
+    Tensor qkv=conv(x,5,4,false);
+    auto p=qkv.permute({0,2,3,4,1}).view({T,1,H*W,3*C});
+    Tensor y=attn_forward(p);
+    y=y.view({1,T,H,W,C}).permute({0,4,1,2,3});
+    return conv(y,7,-1,true,true,x);
+  }
+
+  // The output head: RMS norm + SiLU + the causal padding, the 3x1x1 convolution to 112 channels
+  // (108 real) and the 3x3 stencil with the bias and clamp (head.cu). Replaces conv(x,108,107),
+  // whose 12 output channels waste 244 of a 256-wide tile's columns.
+  Tensor output_head(const Tensor& x) {
+    RECORD_FUNCTION("wan.head",std::vector<c10::IValue>());
+    const int pt=2;
+    int ph=world>1?1:0;
+    Tensor p=norm_pad(x,w[107],cache[108],pt,ph,true);
+    if(ph) {
+      RECORD_FUNCTION("wan.halo",std::vector<c10::IValue>());
+      // The third halo site. PR204 deleted pack_edges/unpack_edges, and the resolution that
+      // compiles-and-passes while silently dropping the vectorised pack is to re-add them:
+      // this is the same two pad rows through comm.cu's halo_pack/halo_unpack, which is what
+      // the other two sites (conv/resnet) already use. Census: post_edges == wait_edges == 3.
+      auto recv=halo_exchange(halo_pack(p),pg,rank,world);
+      halo_unpack(p,recv,rank,world);
+    }
+    cache[108]=last_frames(p,pt,ph);
+    Tensor g=conv_n128(p,w[108]);
+    int rows=x.size(3);
+    auto out=at::empty({1,12,g.size(2),rows,g.size(4)},x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+    // 352 is the pixel-plane height (44 latent rows x 8) and lo this rank's first pixel-plane row;
+    // both are the fixed 704x1248 geometry, and lo must stay 8*(44*rank/world) because run()
+    // passes lo*16 - the same row on the 704-tall video grid - to write_video.
+    head_stencil(g,w[109],out,352,ph,8*(44*rank/world));
+    return out;
+  }
+
+  Tensor up(Tensor input,int block,bool first) {
+    RECORD_FUNCTION("wan.up_block."+std::to_string(block),std::vector<c10::IValue>());
+    const int bases[4]={21,43,65,87};
+    int base=bases[block];
+    bool change=block>=2;
+    int b2=base+(change?8:6),b3=base+(change?14:12);
+    // Memory: only the resampler's DupUp3D shortcut reads the block's input again, and block 3 has none.
+    auto x=resnet(block==3?std::move(input):input,base,change,b2+1,b2);
+    x=resnet(x,b2,false,b3+1,b3);
+    x=resnet(x,b3,false);
+    if(block==3) return x;
+    int spatial=base+(change?20:18);
+    if(block<2) {
+      // The 3x1x1 causal time convolution writes the frame-interleaved tensor directly (a stride-2 D epilogue
+      // in solution/timeconv.cu), so `up_time` and `finish` are gone from this seam.  `activate=false` is
+      // INERT: with no gain the dispatcher takes norm.cu's copy-only mode and never reads the flag.
+      int start=first?1:0;
+      Tensor head=first?x.narrow(2,0,1):Tensor();
+      int nf=x.size(2)-start;
+      Tensor p;
+      if(nf>0) {
+        p=norm_pad(x.narrow(2,start,nf),Tensor(),cache[spatial+2],2,0,false);
+        cache[spatial+2]=last_frames(p,2,0);
+      }
+      x=time_upsample(p,head,w[spatial+2],w[spatial+3],x.size(2),nf,first);
+    }
+    if(world>1) {
+      x=norm_pad(x,Tensor(),Tensor(),0,1,false);
+      auto recv=halo_exchange(halo_pack(x),pg,rank,world);
+      halo_unpack(x,recv,rank,world);
+    }
+    // Four-phase spatial upsampler: the four 2x2 phase convolutions, then one bandwidth-bound pass that
+    // interleaves them with the bias and the DupUp3D shortcut (resample.cu).
+    std::vector<Tensor> pw;
+    for(int a=0;a<2;++a)for(int b=0;b<2;++b) pw.push_back(w[110+4*block+2*a+b]);
+    return resample_upsample(x,pw,w[spatial+1],input,world>1,block<2,first);
+  }
+
+  // Stateless exact prefix interval. Three causal convolutions require six
+  // preceding latent frames, but the support is pruned by two after each conv.
+  Tensor prefix_interval(const Tensor& z,int begin,int count) {
+    sharded=false;
+    const int end=begin+count;
+    const int lo6=std::max(0,begin-6),lo4=std::max(0,begin-4),lo2=std::max(0,begin-2);
+    auto raw=z.narrow(2,lo6,end-lo6).contiguous(at::MemoryFormat::ChannelsLast3d);
+    auto x=conv(raw,0);
+    // Reducing left temporal padding prunes the invalid outputs without ever
+    // computing them. Near global frame zero, the remaining padding is real.
+    x=conv(x,2,-1,true,true,Tensor(),TemporalTrim{lo4-lo6});
+    auto y=conv(x,10,9,true,true,Tensor(),TemporalTrim{lo2-lo4});
+    // Batch is one and T is the outer physical NDHWC dimension, so this
+    // frame slice remains dense for finish()'s linear residual indexing.
+    auto skip=x.narrow(2,begin-lo4,count);
+    TORCH_CHECK(skip.is_contiguous(at::MemoryFormat::ChannelsLast3d),"Prefix residual must be dense NDHWC");
+    y=conv(y,13,12,true,true,skip,TemporalTrim{begin-lo2});
+    return attention(y);
+  }
+
+  Tensor distributed_prefix(const Tensor& z,int start,int count) {
+    // No persistent prefix history: temporal ownership changes from batch to
+    // batch. The full latents on each rank provide exact predecessor context.
+    Decoder prefix(w,rank,world,pg);
+    Tensor x;
+    if(count>=2*world) {
+      const int lo=count*rank/world,hi=count*(rank+1)/world;
+      x=prefix.prefix_interval(z,start+lo,hi-lo);
+      x=prefix_temporal_to_rows(x,count,rank,world,pg);
+    } else {
+      // Tiny clips/tails do not amortize the transpose, and may have empty
+      // temporal owners. A stateless replicated prefix remains exact.
+      x=prefix.prefix_interval(z,start,count);
+      const int lo=44*rank/world,hi=44*(rank+1)/world;
+      x=x.narrow(3,lo,hi-lo).contiguous(at::MemoryFormat::ChannelsLast3d);
+    }
+    return x;
+  }
+
+  Tensor tail(Tensor x,bool first) {
+    // Every prepared edge is consumed within its residual/up-block chain.
+    // Tail history survives chunks, but no prepared current-chunk tensor may.
+    TORCH_CHECK(prepared_for<0 && !prepared.defined() && y_suppressed_for<0,
+                "Prepared input crossed a tail-chunk boundary");
+    sharded=true;
+    x=resnet(x,15);
+    for(int i=0;i<4;++i) x=up(std::move(x),i,first);
+    return output_head(x);
+  }
+
+  Tensor chunk(const Tensor& z,bool first) {
+    sharded=false;
+    // The invariant tail() checks must be CHECKED here too, not cleared. Clearing it first makes
+    // tail()'s check vacuous on the only path world 1 takes, and a leaked fused tensor is not
+    // merely retained: prepared_for is a conv INDEX and every chunk restarts at the same indices,
+    // so the next chunk's conv of that index CONSUMES it (`pre`) and the video is silently wrong
+    // at world 1 while world 4 errors out.
+    //
+    // The reset was dead on this tree, so the check costs one compare per chunk and nothing else,
+    // but the closure argument is NOT "every fused pair lives inside one resnet()" -- that is
+    // false. resnet(b)'s second conv is called with (nxg=b2, nx=b2+1) and produces for
+    // fnext=nx=b2+1, which is the FIRST conv of the next resnet; the same edge appears at every
+    // up block (47->50, 53->56 at block 1; 69->74, 77->80 at block 2). The chain closes at the
+    // END OF THE UP BLOCK instead: up() runs resnet(base, nxg=b2, nx=b2+1), then
+    // resnet(b2, nxg=b3, nx=b3+1), then resnet(b3) with the DEFAULT nx=-1, so the third resnet's
+    // last conv has fg=fnext=-1, produces nothing, and no fused tensor leaves up(). up() is the
+    // only Winograd caller and nothing after it is Winograd (output_head goes through conv_n128),
+    // so the triple is clear at every chunk entry.
+    // The predicate is tail()'s, EXACTLY, and that identity is the point rather than a
+    // coincidence: chunk() must accept everything tail() accepts, or the subset property this
+    // change exists to establish is not established. A third conjunct would also be unreachable --
+    // all three of prepared / prepared_for / prepared_sharded are written together at the single
+    // producer site (conv()'s `if(fuse)`) and cleared together at the single consumption site, so
+    // `prepared_sharded` can never differ from `prepared_for` and cannot fail on its own.
+    TORCH_CHECK(prepared_for<0 && !prepared.defined(),
+                "Prepared input crossed a chunk boundary");
+    auto x=conv(z,0);
+    x=conv(x,2);
+    x=resnet(x,9);
+    x=attention(x);
+    if(world>1) {
+      int lo=44*rank/world,hi=44*(rank+1)/world;
+      x=x.narrow(3,lo,hi-lo).contiguous(at::MemoryFormat::ChannelsLast3d);
+    }
+    return tail(x,first);
+  }
+};
+
+void decode_video(const Tensor& z,const Weights& weights,int64_t rank,int64_t world,const PG& pg,Tensor out,
+                  int64_t frames_per_pass) {
+  at::NoGradGuard no_grad;
+  c10::cuda::CUDAGuard device_guard(z.device());
+  TORCH_CHECK(z.scalar_type()==at::kBFloat16&&z.dim()==5&&z.size(0)==1&&z.size(1)==48&&z.size(2)>=1&&z.size(3)==44&&z.size(4)==78,"Expected bf16 [1,48,T,44,78]");
+  TORCH_CHECK(world>=1&&world<=8&&rank>=0&&rank<world,"Invalid rank/world");
+  TORCH_CHECK(weights.size()==770,"Call prepare() once before run()");
+  TORCH_CHECK(frames_per_pass>=1,"frames_per_pass must be at least 1, not ",frames_per_pass);
+  halo_stats()=HaloStats{};  // per call: the debug tail below prints THIS run's exchange, not the process's
+  Decoder decoder(weights,rank,world,pg);
+  // A pass carries frames_per_pass latent frames per GPU, frames_per_pass * world in all. The working memory of a GPU
+  // grows with it (3.5-4.2 GiB per latent frame at 704 x 1248), and so does the speed: a pass's GEMMs at 44 x 78 and
+  // 88 x 156 fill the GPU only with enough frames, and above world 1 every pass also pays a fixed cost (its prefix
+  // and its halo exchanges). The causal history crosses passes in `cache`, so the video does not depend on the size.
+  // At most 61 latent frames: the largest pass this tree was measured with (wino_epilogue.cu's indices are 64-bit).
+  static const int64_t kMaxChunk=61;
+  const int chunk_size=(int)std::min<int64_t>(frames_per_pass*world,kMaxChunk);
+  int T=z.size(2);
+  for(int start=0;start<T;start+=chunk_size) {
+    int count=std::min((int)chunk_size,T-start);
+    Tensor pixels;
+    if(world>1) pixels=decoder.tail(decoder.distributed_prefix(z,start,count),start==0);
+    else {
+      auto input=z.narrow(2,start,count).contiguous(at::MemoryFormat::ChannelsLast3d);
+      pixels=decoder.chunk(input,start==0);
+    }
+    int frame_offset=start==0?0:4*start-3;
+    if(world==1) write_video(pixels,out,frame_offset,0);
+    else if(rank==0) {
+      write_video(pixels,out,frame_offset,0);
+      for(int peer=1;peer<world;++peer) {
+        int lo=44*peer/world,hi=44*(peer+1)/world;
+        auto part=nhwc_empty(12,pixels.size(2),(hi-lo)*8,624,pixels);
+        std::vector<Tensor> buffers={part};
+        pg->recv(buffers,peer,0)->wait();
+        write_video(part,out,frame_offset,lo*16);
+      }
+    } else {
+      std::vector<Tensor> buffers={pixels};
+      pg->send(buffers,0,0)->wait();
+    }
+  }
+  // Debug only, and never set by CI: cached once per process so the measured span runs no getenv.
+  static const bool show_halo_stats=std::getenv("WAN_HALO_STATS")!=nullptr;
+  if(show_halo_stats) {
+    const HaloStats& s=halo_stats();
+    printf("[halo rank %d/%d] exchanges %lld  sent %.1f MB  received %.1f MB\n",(int)rank,(int)world,
+           (long long)s.exchanges,1e-6*s.bytes_sent,1e-6*s.bytes_received);
+    fflush(stdout);
+  }
+}
+
+// The default pass: 4 latent frames per GPU. On one GB200 at 704 x 1248 it peaks at 21-22 GiB allocated and takes
+// 1.09x the time of 16 frames per GPU, which peaks at 60-61 GiB (8 to 20 s of video).
+constexpr int64_t kFramesPerPass=4;
+
+void run(const Tensor& z,const Weights& weights,int64_t rank,int64_t world,const PG& pg,Tensor out) {
+  decode_video(z,weights,rank,world,pg,out,kFramesPerPass);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME,m) {
+  m.def("prepare",&prepare);
+  m.def("run",&run,py::call_guard<py::gil_scoped_release>());
+  m.def("run_with",&decode_video,py::call_guard<py::gil_scoped_release>());
+}

fastvideo-kernel/wan22_vae_decode/solution/comm.cu

@@ -0,0 +1,213 @@
+// Halo edge transport: neighbour-only point to point and the seed's ring all-gather, over an
+// eight-channels-per-thread edge pack/scatter. Which arm runs is use_allgather's world rule below --
+// the ring only at world 2, neighbour-only at world 4 and at every world above it -- which WAN_HALO
+// can force either way for a measurement.
+// THIS POLICY IS STATED IN THREE PLACES AND THEY MUST AGREE: this file header, the block comment on
+// use_allgather, and the WAN_HALO contract in comm.h. An earlier revision of this change corrected one
+// of the three and was rightly blocked on the other two, so the cross-reference is the fix for the class
+// and not for the instance.
+//
+// Why the point to point ops go in one coalesced group: ProcessGroupNCCL launches an unbatched p2p op as a
+// group of its own on the group's stream, so a rank's p2p ops are stream-ordered against each other, and two
+// ranks each waiting to send into the other's not-yet-issued receive is a cycle. (NCCL documents that
+// serialization; w24's M1387 reaches the same deadlock from a direct all-to-all of unbatched p2p ops.)
+// ProcessGroup::startCoalescing/endCoalescing puts all of this rank's neighbour ops into one NCCL group, which
+// removes that ordering: the exchange is one hop, posted and closed with no host block.
+#include "comm.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <cstdlib>
+#include <cstring>
+
+namespace {
+
+cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+
+// One uint4 is eight bf16 channels; C is a multiple of 8 for every layer that sharded space (256/512/1024).
+__global__ void pack_edge8(const uint4* __restrict__ x, uint4* __restrict__ y, int64_t n8, int CV, int T, int H,
+                           int W) {
+  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t)gridDim.x * blockDim.x) {
+    int64_t row = i / (W * CV), rv = i % (W * CV);
+    int t = (int)(row % T), edge = (int)(row / T);
+    y[i] = x[((int64_t)t * H + (edge ? H - 2 : 1)) * W * CV + rv];
+  }
+}
+
+// Scatters the two received edges into the pad rows, eight channels per thread.
+__global__ void unpack_edge8(uint4* __restrict__ x, const uint4* __restrict__ e, int64_t n8, int CV, int T, int H,
+                             int W, int rank, int world) {
+  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t)gridDim.x * blockDim.x) {
+    int64_t row = i / (W * CV), rv = i % (W * CV);
+    int t = (int)(row % T), edge = (int)(row / T);
+    int peer = rank + (edge ? 1 : -1);  // pad row 0 comes from rank-1, pad row H-1 from rank+1
+    uint4 v = make_uint4(0u, 0u, 0u, 0u);
+    if (peer >= 0 && peer < world) v = e[((int64_t)peer * 2 * T + (int64_t)(1 - edge) * T + t) * W * CV + rv];
+    x[((int64_t)t * H + (edge ? H - 1 : 0)) * W * CV + rv] = v;
+  }
+}
+
+int grid_for(int64_t work) {
+  int64_t blocks = (work + 255) / 256;
+  return (int)(blocks > 2048 ? 2048 : (blocks < 1 ? 1 : blocks));
+}
+
+// Where each transport is used, when WAN_HALO does not force one. With P the bytes of one edge (half the packed
+// buffer): the ring sends (world-1)*2P and receives world*2P, of which this exchange reads 2P (one edge from each
+// neighbour); the neighbour path sends and receives one edge per neighbour -- 2P in and out for an interior rank,
+// P at the ends. So it is 6P -> 2P each way at world 4 and 2P/4P -> P/P at world 2, i.e. the received bytes fall
+// by 3x and 4x. So the policy is the ring only where it has been measured to cost nothing, and the
+// neighbour path everywhere else:
+//   world 2   the two arms measured the same, and the copy counts are equal there too. J817, three
+//             alternating rounds in one job at f121x2, minima 698.6 ms (ring) against 696.5 ms
+//             (neighbour), 0.3% apart, J570's 2.3% the other way sitting inside its own run-to-run
+//             spread; and J3057 ran none / ag / p2p at f481x2 in one job, 2305.6 / 2309.7 / 2293.2.
+//   world 4   neighbour. J2211's own `for A in ag p2p ag p2p` at f481x4: ag 2462.9 / 2244.5 against
+//             p2p 1010.6 / 1017.5, the p2p arm's A/A 0.68 %. The ring moves (world-1) copies each
+//             way, the neighbour path one.
+//   above 4   neighbour, and this rung is chosen on COPY COUNTS rather than on a measurement: the
+//             ring moves (world-1) copies each way against the neighbour path's 1, so the ratio grows
+//             with the world size, and no world above 4 has been measured on either arm here -- a
+//             cross-tray one has never run. The world-4 gap (2.32x) is the only scaled evidence and
+//             it is the smallest world where the two differ. Whoever first gets eight GPUs should
+//             treat this rung as the thing to measure, not as settled.
+// What that leaves: every world the harness runs is unchanged -- 1 by the early return, 2 true, 4 false
+// -- and only world 3 and world 5 and up change.
+// WAN_HALO forces either arm for a measurement.
+bool use_allgather(int world) {
+  if (world < 2) return true;  // one rank has no neighbour to talk to
+  const int forced = halo_transport();
+  if (forced == kAllGather) return true;
+  if (forced == kP2P) return false;
+  return world == 2;
+}
+
+}  // namespace
+
+HaloStats& halo_stats() {
+  static HaloStats stats;
+  return stats;
+}
+
+int halo_transport() {
+  static const int mode = [] {
+    const char* s = std::getenv("WAN_HALO");
+    if (!s) return (int)kPolicy;
+    if (std::strcmp(s, "ag") == 0 || std::strcmp(s, "allgather") == 0) return (int)kAllGather;
+    if (std::strcmp(s, "p2p") == 0 || std::strcmp(s, "neighbour") == 0) return (int)kP2P;
+    return (int)kPolicy;
+  }();
+  return mode;
+}
+
+at::Tensor halo_pack(const at::Tensor& padded) {
+  const int C = (int)padded.size(1), T = (int)padded.size(2), H = (int)padded.size(3), W = (int)padded.size(4);
+  TORCH_CHECK(C % 8 == 0, "halo_pack needs C a multiple of 8");
+  at::Tensor e = at::empty({2, T, W, C}, padded.options());
+  const int CV = C / 8;
+  int64_t n8 = (int64_t)2 * T * W * CV;
+  pack_edge8<<<grid_for(n8), 256, 0, stream()>>>((const uint4*)padded.data_ptr(), (uint4*)e.data_ptr(), n8, CV, T, H,
+                                                 W);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return e;
+}
+
+void halo_unpack(at::Tensor& padded, const at::Tensor& edges, int rank, int world) {
+  const int C = (int)padded.size(1), T = (int)padded.size(2), H = (int)padded.size(3), W = (int)padded.size(4);
+  TORCH_CHECK(C % 8 == 0, "halo_unpack needs C a multiple of 8");
+  const int CV = C / 8;
+  int64_t n8 = (int64_t)2 * T * W * CV;
+  unpack_edge8<<<grid_for(n8), 256, 0, stream()>>>((uint4*)padded.data_ptr(), (const uint4*)edges.data_ptr(), n8, CV,
+                                                   T, H, W, rank, world);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+HaloHandle post_edges(const PG& pg, const at::Tensor& send, at::Tensor& recv, int rank, int world,
+                      cudaStream_t comm) {
+  HaloHandle h;
+  h.started = true;
+  // Both buffers are flat here: the p2p slices are element offsets (narrowing the packed [2,T,W,C] on dim 0 would
+  // index a dimension of size 2, which is what round 1 of this PR did).
+  TORCH_CHECK(send.is_contiguous(), "halo edges must be contiguous");
+  TORCH_CHECK(recv.is_contiguous(), "the halo receive buffer must be contiguous");
+  h.send = send.reshape({-1});
+  h.recv = recv.reshape({-1});
+  HaloStats& s = halo_stats();
+  ++s.exchanges;
+  const int64_t bytes = send.numel() * (int64_t)send.element_size();
+  // The transport is ordered on `comm`: ProcessGroupNCCL makes its own stream wait for the current stream and
+  // the current stream wait for its end event, so this guard is what puts the collective after the pack.
+  c10::cuda::CUDAStreamGuard guard(c10::cuda::getStreamFromExternal(comm, c10::cuda::current_device()));
+  if (use_allgather(world)) {
+    // A ring all-gather: each rank's buffer goes round the ring, (world-1) copies each way.
+    at::Tensor in = send;
+    h.keep.push_back(in);
+    h.keep.push_back(recv);
+    h.works.push_back(pg->_allgather_base(recv, in));
+    s.bytes_sent += bytes * (world - 1);
+    s.bytes_received += bytes * (world - 1);
+    return h;
+  }
+  // The packed edges are [2,T,W,C]; the p2p slices below are element offsets, so they are taken flat (above).
+  h.keep.push_back(send);
+  h.keep.push_back(h.send);
+  h.keep.push_back(recv);
+  // Split on rank 0 = rows [0, H/world), so rank+1 is the band BELOW me and rank-1 the band above. Each rank
+  // sends its first real row up to rank-1 and its last real row down to rank+1 -- one edge per neighbour, not
+  // both to each -- and reads back only the two neighbour slots of `recv`; its other world-2 slots are never
+  // written and stay uninitialised, which halo_unpack is careful not to read (it touches slot rank-1's edge 1
+  // and slot rank+1's edge 0 and nothing else). On the all-gather path every slot is filled. Both leave the
+  // same two edges in the same places.
+  const int64_t half = send.numel() / 2;  // one edge: [T,W,C]; edge 0 first, edge 1 second
+  const int neighbours = (rank > 0 ? 1 : 0) + (rank + 1 < world ? 1 : 0);
+  s.bytes_sent += (int64_t)neighbours * half * send.element_size();
+  s.bytes_received += (int64_t)neighbours * half * send.element_size();
+  // Every tensor operation happens before the group opens: a throw between startCoalescing and endCoalescing
+  // would leave the NCCL group open and hang every rank instead of failing this one.
+  const bool hasLower = rank + 1 < world, hasUpper = rank > 0;
+  std::vector<at::Tensor> recvLower, sendUpper, recvUpper, sendLower;
+  if (hasLower) {
+    at::Tensor v = h.recv.narrow(0, (int64_t)(rank + 1) * 2 * half, half);  // below: its edge 0
+    h.keep.push_back(v);
+    recvLower.push_back(v);
+  }
+  if (hasUpper) {
+    at::Tensor v = h.send.narrow(0, 0, half);  // above: my first real row
+    h.keep.push_back(v);
+    sendUpper.push_back(v);
+  }
+  if (hasUpper) {
+    at::Tensor v = h.recv.narrow(0, (int64_t)(rank - 1) * 2 * half + half, half);  // above: its edge 1
+    h.keep.push_back(v);
+    recvUpper.push_back(v);
+  }
+  if (hasLower) {
+    at::Tensor v = h.send.narrow(0, half, half);  // below: my last real row
+    h.keep.push_back(v);
+    sendLower.push_back(v);
+  }
+  pg->startCoalescing(c10::DeviceType::CUDA);
+  if (hasLower) pg->recv(recvLower, rank + 1, 0);
+  if (hasUpper) pg->send(sendUpper, rank - 1, 0);
+  if (hasUpper) pg->recv(recvUpper, rank - 1, 0);
+  if (hasLower) pg->send(sendLower, rank + 1, 0);
+  // The ops inside a coalesced group carry no work of their own: the group's work is the one to wait on.
+  h.works.push_back(pg->endCoalescing(c10::DeviceType::CUDA));
+  return h;
+}
+
+void wait_edges(HaloHandle& h) {
+  if (!h.started) return;
+  // Each wait installs the collective's completion as an event wait on the current stream and does not block
+  // the host, so the caller's next kernel on that stream is ordered after the transfer.
+  for (auto& w : h.works) w->wait();
+  h.works.clear();
+  h.started = false;
+}
+
+at::Tensor halo_exchange(const at::Tensor& send, const PG& pg, int rank, int world) {
+  at::Tensor recv = at::empty({world * send.numel()}, send.options());
+  HaloHandle h = post_edges(pg, send, recv, rank, world, at::cuda::getCurrentCUDAStream());
+  wait_edges(h);
+  return recv;
+}

fastvideo-kernel/wan22_vae_decode/solution/comm.h

@@ -0,0 +1,83 @@
+// Halo edge transport for the row-band split: the two boundary rows a 3x3 convolution needs from its
+// neighbours, moved with the cheapest mechanism available. Representation-agnostic by design: it carries
+// whatever `halo_pack` packs (today the normalized NDHWC bf16 edge rows; the fusion work may change that),
+// and it leaves the packed layout and `halo_unpack` untouched, so the decoded video is bit-identical to the
+// all-gather it replaces.
+//
+// Two ways in:
+//   halo_exchange(...)                 one call: allocates the recv buffer, transports and closes. What
+//                                      Decoder::conv uses today.
+//   post_edges(...) / wait_edges(...)   the split seam for overlapping the transport with interior compute:
+//                                      post enqueues the exchange and returns a handle without a host block,
+//                                      wait closes it. `halo_pack` may run on one stream and the transport on
+//                                      another: post_edges takes the stream the transport is ordered on.
+#pragma once
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <c10/util/intrusive_ptr.h>
+#include <cuda_runtime.h>
+#include <vector>
+
+using PG = c10::intrusive_ptr<c10d::ProcessGroup>;
+
+// Which transport the halo exchange uses, from WAN_HALO once per process. With nothing set the world-size rule
+// in comm.cu decides; the variable is for forcing an arm in a measurement (both must be the same on every rank
+// of a world: one rank in a collective and one in a p2p group would never meet).
+//   unset: the rule -- the ring all-gather only at world 2, where both arms have been measured equal, and the
+//          neighbour path at world 4 and at every world above it. (World 1 has no exchange at all: ph is 0
+//          there, so nothing calls this.)
+//   "ag" or "allgather": the ring collective, whatever the world size.
+//   "p2p" or "neighbour": the neighbour path, whatever the world size. This is what the rule selects above
+//          world 4, and the neighbour path across trays has never been measured: nothing at that world size
+//          has, on either arm.
+enum HaloTransport { kPolicy = 0, kAllGather = 1, kP2P = 2 };
+int halo_transport();
+
+// The exchange sites and the bytes they moved. Debug instrumentation only: nothing on the run path reads them,
+// and with WAN_HALO_STATS unset they cost one increment per exchange.
+struct HaloStats {
+  int64_t exchanges = 0, bytes_sent = 0, bytes_received = 0;
+};
+HaloStats& halo_stats();
+
+// Gather the two edges a rank must send: the first and the last real row of the padded NDHWC tensor
+// [1,C,T,H,W] (H includes exactly one pad row per side), as [2,T,W,C] contiguous, edge-major: edge 0 is the
+// first real row (H index 1), edge 1 the last (H index H-2). Same layout as the seed's pack_edges, eight
+// channels per thread.
+at::Tensor halo_pack(const at::Tensor& padded);
+
+// Scatter the received edges into the pad rows: row 0 takes rank-1's edge 1, row H-1 takes rank+1's edge 0,
+// and a rank with no neighbour on that side gets zeros. `edges` is the transport's `recv` buffer, and rank+1
+// is the band *below* (higher rows), since the split gives rank 0 the rows [0, H/world).
+void halo_unpack(at::Tensor& padded, const at::Tensor& edges, int rank, int world);
+
+// The transport, split in two. `send` is halo_pack's [2,T,W,C]; `recv` is the caller's buffer of
+// world*2*T*W*C elements in the all-gather's layout (slot r is rank r's two edges), which is what
+// halo_unpack reads. The handle owns the works and keeps every send and recv view alive until it is waited.
+struct HaloHandle {
+  std::vector<c10::intrusive_ptr<c10d::Work>> works;
+  std::vector<at::Tensor> keep;
+  at::Tensor send, recv;  // flat views the p2p slices index
+  bool started = false;
+  HaloHandle() = default;
+  HaloHandle(HaloHandle&&) = default;
+  HaloHandle& operator=(HaloHandle&&) = default;
+  HaloHandle(const HaloHandle&) = delete;
+  HaloHandle& operator=(const HaloHandle&) = delete;
+};
+
+// Enqueue only: no host block, and nothing added to the compute stream. `comm` is the stream the transport is
+// ordered on; it must already carry whatever ordering the caller needs (a wait on the pack, say), and the
+// unpack must not run on a different stream before wait_edges has been called on it.
+HaloHandle post_edges(const PG& pg, const at::Tensor& send, at::Tensor& recv, int rank, int world,
+                      cudaStream_t comm);
+
+// Close it: this orders the CALLER'S CURRENT STREAM after the transfer and returns at once. WorkNCCL::wait is
+// an event wait on the stream that is current at the call, not a CPU block. So the next kernel enqueued on
+// that stream sees the edges -- and a kernel enqueued on any other stream does not.
+void wait_edges(HaloHandle& h);
+
+// Allocate the recv buffer, post, close and return the edges in halo_unpack's layout.
+at::Tensor halo_exchange(const at::Tensor& send, const PG& pg, int rank, int world);

fastvideo-kernel/wan22_vae_decode/solution/conv.cu

@@ -0,0 +1,78 @@
+// CUTLASS 4.2 Blackwell implicit convolution: 2-CTA 256x256x64 tiles,
+// TMA im2col loads, tcgen05 MMA, fp32 accumulation, bf16 epilogue.
+// Every included CUTLASS/CuTe dependency is vendored beside this source.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/convnd_problem_shape.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/conv/dispatch_policy.hpp"
+#include "cutlass/conv/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/conv/device/conv_universal_adapter.hpp"
+#include "cutlass/conv/kernel/conv_universal.hpp"
+using namespace cute;
+using ElementBF16=cutlass::bfloat16_t;
+using LayoutNDHWC=cutlass::layout::TensorNDHWC;
+using ArchTag=cutlass::arch::Sm100;
+using OperatorClass=cutlass::arch::OpClassTensorOp;
+constexpr auto Fprop=cutlass::conv::Operator::kFprop;
+
+template<int M,int N,int CL>
+at::Tensor launch(const at::Tensor& x,const at::Tensor& w,int ph,int pw,int phu=-1,int pwu=-1) {
+  if(phu<0)phu=ph;if(pwu<0)pwu=pw;
+  using Tile=Shape<Int<M>,Int<N>,Shape<_64>>;
+  using Cluster=Shape<Int<CL>,_1,_1>;
+  using Epilogue=typename cutlass::epilogue::collective::CollectiveBuilder<
+    ArchTag,OperatorClass,Tile,Cluster,cutlass::epilogue::collective::EpilogueTileAuto,
+    float,float,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+  using Mainloop=typename cutlass::conv::collective::CollectiveBuilder<
+    ArchTag,OperatorClass,Fprop,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,float,Tile,Cluster,
+    cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
+    cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+  using Problem=cutlass::conv::ConvProblemShape<Fprop,3>;
+  using Kernel=cutlass::conv::kernel::ConvUniversal<Problem,Mainloop,Epilogue>;
+  using Conv=cutlass::conv::device::ConvUniversalAdapter<Kernel>;
+  int C=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4);
+  int K=w.size(0),KT=w.size(2),KH=w.size(3),KW=w.size(4);
+  auto y=at::empty({1,K,D-KT+1,H+ph+phu-KH+1,W+pw+pwu-KW+1},x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+  Problem problem(cutlass::conv::Mode::kCrossCorrelation,
+                  {1,D,H,W,C},{x.stride(0),x.stride(2),x.stride(3),x.stride(4),x.stride(1)},
+                  {K,KT,KH,KW,C},{w.stride(0),w.stride(2),w.stride(3),w.stride(4),w.stride(1)},
+                  {0,ph,pw},{0,phu,pwu},{1,1,1},{1,1,1},1);
+  typename Kernel::StrideC sc{};
+  typename Kernel::StrideD sd{};
+  cute::for_each(cute::make_seq<cute::rank<0>(sc)>{},[&](auto i){cute::get<0,i>(sc)=problem.stride_C[Problem::RankT-2-i];});
+  cute::for_each(cute::make_seq<cute::rank<0>(sd)>{},[&](auto i){cute::get<0,i>(sd)=problem.stride_C[Problem::RankT-2-i];});
+  typename Conv::Arguments args{problem,{(const ElementBF16*)x.data_ptr(),(const ElementBF16*)w.data_ptr()},
+    {{1.f,0.f},(const ElementBF16*)y.data_ptr(),sc,(ElementBF16*)y.data_ptr(),sd}};
+  Conv conv;
+  TORCH_CHECK(conv.can_implement(args)==cutlass::Status::kSuccess,"CUTLASS cannot implement shape");
+  size_t size=Conv::get_workspace_size(args);
+  auto workspace=at::empty({(int64_t)size},x.options().dtype(at::kByte));
+  auto stream=at::cuda::getCurrentCUDAStream();
+  auto status=conv.initialize(args,workspace.data_ptr(),stream);
+  TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS initialization failed ",int(status));
+  status=conv.run(stream);
+  TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS convolution failed ",int(status));
+  return y;
+}
+at::Tensor custom_conv(const at::Tensor& x,const at::Tensor& w,int64_t ph,int64_t pw,int64_t variant) {
+  return launch<256,256,2>(x,w,ph,pw);
+}
+
+// The output head's 3x1x1 convolution: N=128 because the head has 112 output channels, and a
+// 256-wide tile would spend 144 of its 256 columns on nothing. Measured at one clock (1951-2022
+// MHz, both arms in J1044 and J1343): N=128 is 0.83 PFLOP/s = 34% of peak, N=256 is 0.69 = 28%.
+// The time is what decides it and it is not close: 6.30 ms against 15.2 ms.
+at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w) {
+  return launch<256,128,2>(x,w,0,0);
+}
+
+at::Tensor phase_conv(const at::Tensor& x,const at::Tensor& w,int a,int b,bool halo) {
+  at::Tensor v=halo?x.narrow(3,a,x.size(3)-1):x;
+  return launch<256,256,2>(v,w,halo?0:1-a,1-b,halo?0:a,b);
+}

fastvideo-kernel/wan22_vae_decode/solution/head.cu

@@ -0,0 +1,82 @@
+// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to
+// [-1,1], from the 112-channel plane conv.cu's `conv_n128` (a 3x1x1 convolution) writes to the
+// 12-channel plane write_video consumes.
+//
+// Why the convolution writes 108 channels and this pass sums 9 of them. conv_out is 256x27 MACs
+// and its per-position INPUT WINDOW is 6912*2 B (the weight is 12x that and stays resident), i.e.
+// 12 FLOP/byte against the machine's ~375, so any form that reads the 27 taps per position is
+// operand-bound. Moving the nine spatial taps into the output channel dimension of a 3x1x1
+// (temporal-only) convolution makes the operand 768*2 B for 108 outputs instead of 13.8 KB for
+// 12, and the nine shifted copies are summed here.
+//
+// The value path, so that the rounding is on the record: the three temporal taps are accumulated
+// in the convolution's fp32 accumulator and rounded to bf16 in g; the nine spatial taps are then
+// summed here in fp32 from nine bf16 g values, the bias is added in fp32, and the result is
+// rounded once to bf16 on the store. Main's path is also two roundings (cuDNN's bf16 conv_out,
+// then finish_scalar's bf16 of x+bias); the count is the same and what differs is WHERE the first
+// one sits - after three taps instead of after twenty-seven - which is why max|d| moves and rms
+// does not. The reference accumulates all 27 taps in one fp32 accumulator.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_bf16.h>
+#include "head.h"
+
+// One thread per output position. Its 12 accumulators take the 12 channels of the same (i,j)
+// tap from nine neighbouring positions; the nine shifted reads of each 216-byte channel vector
+// are served from L1, so the pass moves each element of g about once.
+static __global__ void head_stencil_kernel(const __nv_bfloat16* __restrict__ g,
+                                    const __nv_bfloat16* __restrict__ bias,
+                                    __nv_bfloat16* __restrict__ out,
+                                    int D,int rows,int W,int H,int g_off,int lo) {
+  int idx=blockIdx.x*blockDim.x+threadIdx.x;
+  if(idx>=D*rows*W)return;
+  int w=idx%W, t=idx/W, r=t%rows, d=t/rows;
+  int gh0=lo+r;
+  const int gH=rows+2*g_off;
+  float acc[12];
+#pragma unroll
+  for(int c=0;c<12;++c)acc[c]=__bfloat162float(bias[c]);
+  for(int i=0;i<3;++i) {
+    int gh=gh0+i-1;
+    if(gh<0||gh>=H)continue;                       // conv3d's zero padding, in the plane
+    int gr=r+i-1+g_off;
+    if(gr<0||gr>=gH)continue;
+    for(int j=0;j<3;++j) {
+      int ww=w+j-1;
+      if(ww<0||ww>=W)continue;                     // conv3d's zero padding, across the row
+      const __nv_bfloat162* p=(const __nv_bfloat162*)(g+((((size_t)d*gH+gr)*W)+ww)*112+(i*3+j)*12);
+#pragma unroll
+      for(int k=0;k<6;++k) {
+        float2 f=__bfloat1622float2(p[k]);
+        acc[2*k]+=f.x;
+        acc[2*k+1]+=f.y;
+      }
+    }
+  }
+  __nv_bfloat162* o=(__nv_bfloat162*)(out+((((size_t)d*rows+r)*W)+w)*12);
+#pragma unroll
+  for(int k=0;k<6;++k)
+    o[k]=__floats2bfloat162_rn(fminf(1.f,fmaxf(-1.f,acc[2*k])),fminf(1.f,fmaxf(-1.f,acc[2*k+1])));
+}
+
+void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out,
+                  int64_t H_global,int64_t g_off,int64_t lo) {
+  int D=out.size(2),rows=out.size(3),W=out.size(4);
+  TORCH_CHECK(g.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&g.size(1)==112,
+              "head_stencil: g must be [1,112,D,rows+2*g_off,W] channels-last");
+  TORCH_CHECK(out.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&out.size(1)==12,
+              "head_stencil: out must be [1,12,D,rows,W] channels-last");
+  // The kernel strides g by gH = rows+2*g_off taken from `out` and reads D*rows*W positions of it,
+  // so a g of the wrong depth or height reads out of bounds rather than failing here.
+  TORCH_CHECK(g.size(2)==out.size(2)&&g.size(3)==out.size(3)+2*g_off&&g.size(4)==out.size(4),
+              "head_stencil: g must be [1,112,D,rows+2*g_off,W] of out's [1,12,D,rows,W]");
+  c10::cuda::CUDAGuard guard(g.device());
+  int64_t total=(int64_t)D*rows*W;
+  int threads=256;
+  int blocks=(int)((total+threads-1)/threads);
+  head_stencil_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(
+      (const __nv_bfloat16*)g.data_ptr(),(const __nv_bfloat16*)bias.data_ptr(),
+      (__nv_bfloat16*)out.data_ptr(),D,rows,W,(int)H_global,(int)g_off,(int)lo);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}

fastvideo-kernel/wan22_vae_decode/solution/head.h

@@ -0,0 +1,19 @@
+#pragma once
+#include <ATen/ATen.h>
+
+// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to
+// [-1,1], from the 112-channel plane the 3x1x1 convolution writes (conv.cu's `conv_n128`) to the
+// 12-channel plane write_video consumes.
+//   g:   [1,112,D,rows+2*g_off,W] channels-last-3d bf16, n = (i*3+j)*12 + c for the tap (i,j);
+//        channels 108..111 are the zero-weight lanes and are never read.
+//   out: [1,12,D,rows,W] channels-last-3d bf16, written whole.
+// H_global is the height of the plane this stencil filters (352 = 44 latent rows x 8; the pixel
+// plane is 704 tall and run()'s video row is twice this one). lo is the first row of this rank's
+// band IN THAT PLANE, 8*(44*rank/world), and it must stay that: run() passes lo*16 to write_video,
+// the same row on the 704-tall video grid, and the two derivations have to agree.
+void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out,
+                  int64_t H_global,int64_t g_off,int64_t lo);
+
+// conv.cu: the head's 3x1x1 convolution, at the tile its 112 output channels want. Declared here
+// so the two-sided declaration of a CUTLASS tile cannot drift from its definition in conv.cu.
+at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w);

fastvideo-kernel/wan22_vae_decode/solution/kernels.cu

@@ -0,0 +1,99 @@
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+using B = __nv_bfloat16;
+union Pack { uint4 v; B b[8]; __nv_bfloat162 h[4]; };
+static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+static const B* ptr(const Tensor& t) { return t.defined() ? (const B*)t.data_ptr<at::BFloat16>() : nullptr; }
+static B* ptr(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
+Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like) {
+  return at::empty({1,c,t,h,w},like.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+}
+Tensor last_frames(const Tensor& x,int n,int ph) {
+  int C=x.size(1),T=x.size(2),H=x.size(3)-2*ph,W=x.size(4);
+  Tensor y=nhwc_empty(C,n,H,W,x);
+  size_t row=(size_t)H*W*C*sizeof(B),src_pitch=(size_t)(H+2*ph)*W*C*sizeof(B);
+  const B* source=ptr(x)+((int64_t)(T-n)*(H+2*ph)+ph)*W*C;
+  C10_CUDA_CHECK(cudaMemcpy2DAsync(ptr(y),row,source,src_pitch,row,n,cudaMemcpyDeviceToDevice,stream()));
+  return y;
+}
+// The halo edge pack/scatter lives in comm.cu (halo_pack/halo_unpack, vectorised), with the transport.
+__global__ void finish_vector(B* x,const B* b,const B* s,int64_t n,int C) {
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n/8;i+=(int64_t)gridDim.x*blockDim.x) {
+    Pack v, bias, skip;
+    v.v=((uint4*)x)[i];bias.v=((const uint4*)b)[i%(C/8)];
+    if(s) skip.v=((const uint4*)s)[i];
+    #pragma unroll
+    for(int j=0;j<4;j++) { v.h[j]=__hadd2(v.h[j],bias.h[j]); if(s) v.h[j]=__hadd2(v.h[j],skip.h[j]); }
+    ((uint4*)x)[i]=v.v;
+  }
+}
+__global__ void finish_scalar(B* x,const B* b,const B* s,int64_t n,int C) {
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n;i+=(int64_t)gridDim.x*blockDim.x) {
+    B v=__float2bfloat16(__bfloat162float(x[i])+__bfloat162float(b[i%C]));
+    if(s) v=__float2bfloat16(__bfloat162float(v)+__bfloat162float(s[i])); x[i]=v;
+  }
+}
+void finish(Tensor& x,const Tensor& b,const Tensor& s) {
+  int C=x.size(1);int64_t n=x.numel();
+  if(C%8==0) finish_vector<<<std::min((int64_t)2048,(n+2047)/2048),256,0,stream()>>>(ptr(x),ptr(b),ptr(s),n,C);
+  else finish_scalar<<<std::min((int64_t)2048,(n+255)/256),256,0,stream()>>>(ptr(x),ptr(b),ptr(s),n,C);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+__global__ void write_kernel(const B* x,B* out,int64_t n,int T,int H,int W,int fullT,int fullH,int offsetT,int offsetH) {
+  int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;
+  if(i<n) {
+    int c=i%12;int64_t p=i/12;int w=p%W,h=(p/W)%H,t=p/((int64_t)W*H);
+    int color=c/4,dy=c%2,dx=(c%4)/2;
+    int64_t o=(((int64_t)color*fullT+t+offsetT)*fullH+offsetH+2*h+dy)*(2*W)+2*w+dx;
+    out[o]=__float2bfloat16(fminf(1.f,fmaxf(-1.f,__bfloat162float(x[i]))));
+  }
+}
+void write_video(const Tensor& x,Tensor& out,int frame_offset,int row_offset) {
+  write_kernel<<<(x.numel()+255)/256,256,0,stream()>>>(ptr(x),ptr(out),x.numel(),x.size(2),x.size(3),x.size(4),out.size(2),out.size(3),frame_offset,row_offset);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// phases_kernel used to sit here: PR17 replaces it, together with finish_phases below, with one
+// bandwidth-bound pass in resample.cu that writes the interleaved tensor directly.
+// F(4x4,3x3), with fp32 spatial transforms and bf16 transformed operands.
+__device__ __forceinline__ void bt6(float v0,float v1,float v2,float v3,float v4,float v5,float* y) {
+  float a=v4-4*v2,b=v3-4*v1,c=v4-v2,d=2*(v3-v1);
+  y[0]=4*v0-5*v2+v4;y[1]=a+b;y[2]=a-b;y[3]=c+d;y[4]=c-d;y[5]=4*v1-5*v3+v5;
+}
+__global__ void winograd_input_kernel(const B* x,B* v,int C,int D,int H,int W,int PH,int TH,int TW) {
+  int HP=H+2*PH,CV=C/2;int64_t n=(int64_t)D*TH*TW*CV;
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n;i+=(int64_t)gridDim.x*blockDim.x) {
+    int c=i%CV,p=i/CV,tx=p%TW,ty=(p/TW)%TH,t=p/(TW*TH);
+    float d[2][6][6];
+    #pragma unroll
+    for(int a=0;a<6;a++) {
+      int y=4*ty-1+a+PH;
+      #pragma unroll
+      for(int b=0;b<6;b++) {int xx=4*tx-1+b;float2 f=(y>=0&&y<HP&&xx>=0&&xx<W)?__bfloat1622float2(((const __nv_bfloat162*)x)[(((int64_t)t*HP+y)*W+xx)*CV+c]):make_float2(0,0);d[0][a][b]=f.x;d[1][a][b]=f.y;}
+    }
+    #pragma unroll
+    for(int lane=0;lane<2;lane++) {
+      #pragma unroll
+      for(int b=0;b<6;b++) {float q[6];bt6(d[lane][0][b],d[lane][1][b],d[lane][2][b],d[lane][3][b],d[lane][4][b],d[lane][5][b],q);
+        #pragma unroll
+        for(int a=0;a<6;a++)d[lane][a][b]=q[a];}
+    }
+    #pragma unroll
+    for(int a=0;a<6;a++) {float q0[6],q1[6];bt6(d[0][a][0],d[0][a][1],d[0][a][2],d[0][a][3],d[0][a][4],d[0][a][5],q0);bt6(d[1][a][0],d[1][a][1],d[1][a][2],d[1][a][3],d[1][a][4],d[1][a][5],q1);
+      #pragma unroll
+      for(int b=0;b<6;b++)((__nv_bfloat162*)v)[(6*a+b)*n+i]=__floats2bfloat162_rn(q0[b],q1[b]);}
+  }
+}
+Tensor winograd_input(const Tensor& x,int ph) {
+  int C=x.size(1),D=x.size(2),H=x.size(3)-2*ph,W=x.size(4),TH=(H+3)/4,TW=(W+3)/4;
+  auto v=at::empty({36,D,TH,TW,C},x.options());int64_t n=(int64_t)D*TH*TW*C;
+  winograd_input_kernel<<<std::min((int64_t)2048,(n+255)/256),256,0,stream()>>>(ptr(x),ptr(v),C,D,H,W,ph,TH,TW);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();return v;
+}
+// (WinoPtrs, at6, winograd_output_kernel and winograd_output are deleted: with the fused inverse
+//  in wino_epilogue.cu nothing calls them. That file's a6v reads the same A^T out of shared memory
+//  instead of 36 separate global pointers, and it is the only copy of the transform left there.)
+

fastvideo-kernel/wan22_vae_decode/solution/kernels.h

@@ -0,0 +1,21 @@
+#pragma once
+#include <ATen/ATen.h>
+using at::Tensor;
+Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like);
+Tensor norm_pad(const Tensor& x, const Tensor& gamma, const Tensor& history, int pt, int ph, bool silu);
+Tensor last_frames(const Tensor& padded, int n, int ph);
+void finish(Tensor& x, const Tensor& bias, const Tensor& residual);
+void write_video(const Tensor& x, Tensor& out, int frame_offset, int row_offset);
+
+Tensor winograd_input(const Tensor& x,int ph);
+Tensor winograd_input_fast(const Tensor& x,int ph);
+// winograd_output / winograd_output_kernel are deleted: the fused inverse in wino_epilogue.cu
+// replaced their only caller, and they had no live user left (PR172 round-1 review).
+
+// Fused inverse transform + bias/residual + RMS norm/SiLU.  y = A^T M A + bias + residual;
+// n (when gamma is given) = silu(rmsnorm(y)*sqrt(C)*gain) in the (D+pt, H+2ph, W, C) padded
+// layout the next convolution consumes, with its first pt frames taken from `hist`.
+struct WinoOut { Tensor y; Tensor n; };
+WinoOut wino_epilogue(const std::vector<Tensor>& m,int H,int W,const Tensor& bias,
+                      const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph,
+                      bool write_y=true);

fastvideo-kernel/wan22_vae_decode/solution/norm.cu

@@ -0,0 +1,198 @@
+// norm.cu — RMS norm + SiLU + pad for the Wan2.2 VAE decoder, one warp per position.
+//
+// This replaces the seed's norm_vector<C> / pad_vector pair in kernels.cu with the same
+// call contract (norm_pad(x, gamma, hist, pt, ph, silu)) and the same layout (NDHWC,
+// [T+pt][H+2ph][W][C], history frames copied through, pad rows zeroed).
+//
+// Why it is faster (measured in probe/norm_probe2.cu and norm_probe3.cu, one GB200):
+//  - the seed is XU-pipe bound (w08's and w14's ncu: 78% SM busy, 16% DRAM, XU 75%):
+//    per element it pays cvt bf16->f32 for the sum, five cvts for the reference's three
+//    bf16 roundings, __expf and an IEEE divide. This does
+//      * bf16 -> f32 by an integer shift (ALU pipe, exact, no XU op),
+//      * one rounding for the whole normalize*sqrt(C)*gain product instead of three,
+//        which is also closer to the fp32 reference the decode is graded against,
+//      * SiLU as u + u*tanh(u) with u = 0.5*pre-activation (one tanh.approx.f32),
+//      * bf16 stores two elements at a time (cvt.rn.bf16x2.f32),
+//    i.e. ~1.5 XU ops per element instead of ~11.
+//  - rows are blocked so the per-position integer divisions of the seed's flat-pixel loop
+//    happen once per row, not once per position (at C=256 a position is only 8 elements).
+//  - P positions per warp in flight, the number that measured best per channel count.
+//
+// Numerics: one bf16 rounding for the whole normalize*gain product instead of three, and
+// SiLU in fp32 through tanh.approx.f32. That approximation's documented worst case is an
+// absolute error of about 2^-11 (5e-4) in tanh, which is harmless here because it enters
+// the result as u*2^-11 and the stored value is bf16, so the perturbation stays below the
+// bf16 quantum of u + u*tanh(u) over the range these pre-activations take. Measured, and
+// the place to look if that ever changes (bench.py, f121x1 vs the fp32 reference):
+//   seed   max 0.1709 (0.34x budget) rms 0.002242 (0.38x budget)
+//   this   max 0.1992 (0.40x budget) rms 0.002187 (0.37x budget)
+// i.e. the max moves up a little, the rms moves down, both far inside the budget.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+#include <cstdlib>
+#include <type_traits>
+
+using B = __nv_bfloat16;
+
+namespace {
+
+__device__ __forceinline__ void split2(uint32_t w, float& lo, float& hi) {
+  lo = __uint_as_float(w << 16);      // exact bf16 -> f32, ALU pipe
+  hi = __uint_as_float(w & 0xffff0000u);
+}
+__device__ __forceinline__ float tanh_approx(float x) {
+  float r; asm("tanh.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); return r;
+}
+__device__ __forceinline__ uint32_t pack2(float a, float b) {
+  uint32_t r; asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(b), "f"(a)); return r;
+}
+
+// MODE 0: RMS norm * sqrt(C) * gain, then SiLU.  MODE 1: the same, no SiLU.  MODE 2: copy only.
+template <int C, int P, int MODE>
+__global__ void norm_pad_kernel(const B* __restrict__ x, const B* __restrict__ gain,
+                                const B* __restrict__ hist, B* __restrict__ y,
+                                int T, int H, int W, int PT, int PH) {
+  constexpr int CV = C / 8;      // uint4 per position
+  constexpr int NP = C / 256;    // uint4 per lane per position
+  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5;
+  const int HP = H + 2 * PH;
+  const int64_t rows = (int64_t)(T + PT) * HP;
+  const float sqrtC = sqrtf((float)C);
+  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+    const int h = (int)(row % HP), t = (int)(row / HP);
+    const bool interior = (h >= PH && h < H + PH);
+    const bool current = (t >= PT);       // a history frame is copied through, never re-normalized
+    const int ph = h - PH;
+    const uint4* src = nullptr;
+    if (interior) {
+      if (current) src = (const uint4*)x + ((int64_t)(t - PT) * H + ph) * W * CV;
+      else if (hist) src = (const uint4*)hist + ((int64_t)t * H + ph) * W * CV;
+    }
+    for (int w0 = warp * P; w0 < W; w0 += warps * P) {
+      uint4 v[P][NP];
+      #pragma unroll
+      for (int q = 0; q < P; q++) {
+        const int w = w0 + q;
+        #pragma unroll
+        for (int j = 0; j < NP; j++) {
+          v[q][j] = make_uint4(0, 0, 0, 0);
+          if (src && w < W) v[q][j] = src[(int64_t)w * CV + lane + j * 32];
+        }
+      }
+      if (MODE < 2 && interior && current) {
+        float s[P][2];
+        #pragma unroll
+        for (int q = 0; q < P; q++) { s[q][0] = 0.f; s[q][1] = 0.f; }
+        #pragma unroll
+        for (int q = 0; q < P; q++) {
+          #pragma unroll
+          for (int j = 0; j < NP; j++) {
+            float f[8];
+            split2(v[q][j].x, f[0], f[1]); split2(v[q][j].y, f[2], f[3]);
+            split2(v[q][j].z, f[4], f[5]); split2(v[q][j].w, f[6], f[7]);
+            #pragma unroll
+            for (int k = 0; k < 8; k += 2) {
+              s[q][0] = fmaf(f[k], f[k], s[q][0]);
+              s[q][1] = fmaf(f[k + 1], f[k + 1], s[q][1]);
+            }
+          }
+        }
+        #pragma unroll
+        for (int q = 0; q < P; q++) {
+          float t0 = s[q][0] + s[q][1];
+          #pragma unroll
+          for (int d = 16; d; d >>= 1) t0 += __shfl_xor_sync(0xffffffff, t0, d);
+          // 0.5 * sqrt(C) / max(sqrt(sum), 1e-12): one multiply per element, u = 0.5*pre-activation
+          const float hs = 0.5f * sqrtC * rsqrtf(fmaxf(t0, 1e-24f));
+          #pragma unroll
+          for (int j = 0; j < NP; j++) {
+            const uint4 gn = ((const uint4*)gain)[lane + j * 32];
+            float g[8], xf[8], yf[8];
+            split2(gn.x, g[0], g[1]); split2(gn.y, g[2], g[3]);
+            split2(gn.z, g[4], g[5]); split2(gn.w, g[6], g[7]);
+            split2(v[q][j].x, xf[0], xf[1]); split2(v[q][j].y, xf[2], xf[3]);
+            split2(v[q][j].z, xf[4], xf[5]); split2(v[q][j].w, xf[6], xf[7]);
+            #pragma unroll
+            for (int k = 0; k < 8; k++) {
+              const float u = xf[k] * hs * g[k];
+              yf[k] = (MODE == 0) ? fmaf(u, tanh_approx(u), u) : (2.f * u);
+            }
+            v[q][j].x = pack2(yf[0], yf[1]); v[q][j].y = pack2(yf[2], yf[3]);
+            v[q][j].z = pack2(yf[4], yf[5]); v[q][j].w = pack2(yf[6], yf[7]);
+          }
+        }
+      }
+      #pragma unroll
+      for (int q = 0; q < P; q++) {
+        const int w = w0 + q;
+        if (w < W) {
+          uint4* dst = (uint4*)y + row * W * CV + (int64_t)w * CV + lane;
+          #pragma unroll
+          for (int j = 0; j < NP; j++) dst[j * 32] = v[q][j];
+        }
+      }
+    }
+  }
+}
+
+// Copy + pad for any C divisible by 8: the calls with no gain (post_quant_conv's input, the temporal convs).
+__global__ void copy_row_kernel(const uint4* __restrict__ x, const uint4* __restrict__ hist,
+                                uint4* __restrict__ y, int T, int H, int W, int PT, int PH, int CV) {
+  const int HP = H + 2 * PH;
+  const int64_t rows = (int64_t)(T + PT) * HP;
+  const int64_t nrow = (int64_t)W * CV;
+  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+    const int h = (int)(row % HP), t = (int)(row / HP);
+    const bool interior = (h >= PH && h < H + PH);
+    const uint4* src = nullptr;
+    if (interior) {
+      const int tt = t - PT, ph = h - PH;
+      if (tt >= 0) src = x + ((int64_t)tt * H + ph) * nrow;
+      else if (hist) src = hist + ((int64_t)t * H + ph) * nrow;
+    }
+    uint4* dst = y + row * nrow;
+    const uint4 zz = make_uint4(0, 0, 0, 0);
+    for (int64_t i = threadIdx.x; i < nrow; i += blockDim.x) dst[i] = src ? src[i] : zz;
+  }
+}
+
+}  // namespace
+
+at::Tensor norm_pad(const at::Tensor& x, const at::Tensor& gamma, const at::Tensor& hist,
+                    int pt, int ph, bool silu) {
+  const int C = (int)x.size(1), T = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);
+  auto y = at::empty({1, C, T + pt, H + ph * 2, W},
+                     x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+  const B* xp = (const B*)x.data_ptr<at::BFloat16>();
+  const B* gp = gamma.defined() ? (const B*)gamma.data_ptr<at::BFloat16>() : nullptr;
+  const B* hp = hist.defined() ? (const B*)hist.data_ptr<at::BFloat16>() : nullptr;
+  B* yp = (B*)y.data_ptr<at::BFloat16>();
+  const int64_t rows = (int64_t)(T + pt) * (H + ph * 2);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  const int blocks = (int)std::min<int64_t>(rows, 1024);
+  const int threads = 256;
+  auto go = [&](auto tag) {
+    constexpr int CC = decltype(tag)::value;
+    // P positions in flight per warp: measured best at 4 for C<=512, 2 at C=1024
+    constexpr int P = (CC == 1024) ? 2 : 4;
+    if (!gp) norm_pad_kernel<CC, P, 2><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
+    else if (silu) norm_pad_kernel<CC, P, 0><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
+    else norm_pad_kernel<CC, P, 1><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
+  };
+  const bool fast = (C == 1024 || C == 512 || C == 256);
+  TORCH_CHECK(!gp || fast, "Unsupported RMS width ", C);
+  if (fast) {
+    if (C == 1024) go(std::integral_constant<int, 1024>{});
+    else if (C == 512) go(std::integral_constant<int, 512>{});
+    else go(std::integral_constant<int, 256>{});
+  } else {
+    TORCH_CHECK(C % 8 == 0, "norm_pad copy needs C divisible by 8, got ", C);
+    copy_row_kernel<<<blocks, threads, 0, stream>>>((const uint4*)xp, (const uint4*)hp, (uint4*)yp, T, H, W,
+                                                    pt, ph, C / 8);
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}

fastvideo-kernel/wan22_vae_decode/solution/prefix.cu

@@ -0,0 +1,56 @@
+// Prefix ownership transpose: no arithmetic, only bf16 NDHWC packets.
+#include "prefix.h"
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_runtime.h>
+
+namespace {
+// Blocks for destination bands are disjoint; each packet is [T,H_peer,W,C].
+__global__ void pack_prefix_rows(const uint4* x,uint4* send,int T,int H,int WC8,int world) {
+  const int peer=blockIdx.y;
+  const int lo=H*peer/world, hi=H*(peer+1)/world;
+  const int64_t n=(int64_t)T*(hi-lo)*WC8;
+  const int64_t offset=(int64_t)T*lo*WC8;
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n;i+=(int64_t)gridDim.x*blockDim.x) {
+    const int64_t t=i/((int64_t)(hi-lo)*WC8);
+    const int64_t rowcol=i%((int64_t)(hi-lo)*WC8);
+    send[offset+i]=x[(t*H+lo)*WC8+rowcol];
+  }
+}
+}
+
+at::Tensor prefix_temporal_to_rows(const at::Tensor& local,int64_t total_frames,
+    int rank,int world,const c10::intrusive_ptr<c10d::ProcessGroup>& pg) {
+  const int C=local.size(1),T=local.size(2),H=local.size(3),W=local.size(4);
+  TORCH_CHECK(C%8==0 && local.is_contiguous(at::MemoryFormat::ChannelsLast3d),"Prefix needs aligned NDHWC");
+  TORCH_CHECK(T==total_frames*(rank+1)/world-total_frames*rank/world,"Prefix ownership mismatch");
+  auto send=at::empty({local.numel()},local.options());
+  const int band_height=H*(rank+1)/world-H*rank/world;
+  auto output=nhwc_empty(C,total_frames,band_height,W,local);
+  // ProcessGroup splits count bf16 elements, not bytes. Receive into the
+  // physical NDHWC order: the source ranks own chronological frame intervals.
+  auto physical=output.permute({0,2,3,4,1});
+  TORCH_CHECK(physical.is_contiguous(),"Prefix receive must be dense NDHWC");
+  auto recv=physical.view({-1});
+  if(T) {
+    dim3 grid(256,world);
+    pack_prefix_rows<<<grid,256,0,at::cuda::getCurrentCUDAStream()>>>(
+        (const uint4*)local.data_ptr(),(uint4*)send.data_ptr(),T,H,W*(C/8),world);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+  std::vector<int64_t> in_split(world),out_split(world);
+  int64_t sent=0,received=0;
+  for(int p=0;p<world;++p) {
+    in_split[p]=(int64_t)T*(H*(p+1)/world-H*p/world)*W*C;
+    out_split[p]=(total_frames*(p+1)/world-total_frames*p/world)*band_height*W*C;
+    sent+=in_split[p]; received+=out_split[p];
+  }
+  TORCH_CHECK(sent==send.numel() && received==recv.numel(),"Prefix packet size mismatch");
+  // WorkNCCL::wait orders the caller's stream after the collective; it need
+  // not block the host. ProcessGroup's tensor/allocator stream tracking keeps
+  // send/recv storage live through that dependency. Do not drop this wait
+  // when tail kernels consume output on the caller's current stream.
+  pg->alltoall_base(recv,send,out_split,in_split)->wait();
+  return output;
+}

fastvideo-kernel/wan22_vae_decode/solution/prefix.h

@@ -0,0 +1,7 @@
+#pragma once
+#include <ATen/ATen.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+// Exchange chronological full-frame prefix slices for chronological row bands.
+// Source rank r owns [floor(T*r/world), floor(T*(r+1)/world)).
+at::Tensor prefix_temporal_to_rows(const at::Tensor& local,int64_t total_frames,
+    int rank,int world,const c10::intrusive_ptr<c10d::ProcessGroup>& pg);

fastvideo-kernel/wan22_vae_decode/solution/resample.cu

@@ -0,0 +1,96 @@
+// Resampler stage: the four-phase spatial upsamplers of up_blocks 0,1,2 of the Wan2.2 VAE decoder.
+// See resample.h for why the phase convolutions stay direct and what this pass replaces.
+#include "resample.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+using B = __nv_bfloat16;
+// conv.cu: a 2x2 low-resolution convolution of one output parity, padded per phase; `halo` is set when
+// the input carries a one-row halo from the neighbour rank, so the input is narrowed instead of padded.
+at::Tensor phase_conv(const at::Tensor&, const at::Tensor&, int, int, bool);
+
+static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+static const B* pc(const Tensor& t) { return (const B*)t.data_ptr<at::BFloat16>(); }
+static B* pm(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
+
+union P8 { uint4 v; B b[8]; __nv_bfloat162 h[4]; };
+
+// One thread: eight channels of one 2x2 low-resolution tile, all four phases.  The grid carries the
+// coordinates (tx, t, ty) so there is no integer division on the hot path, and every access is a
+// 16-byte vector: four phase outputs and the DupUp3D source in, four interleaved positions out.
+__global__ void upsampler_kernel(const B* __restrict__ p0, const B* __restrict__ p1,
+                                 const B* __restrict__ p2, const B* __restrict__ p3,
+                                 B* __restrict__ y, const B* __restrict__ bias,
+                                 const B* __restrict__ skip, int C, int HL, int WL, int TC,
+                                 bool temporal, bool first) {
+  const int cg = threadIdx.x;                       // 8-channel group
+  const int tx = blockIdx.x, t = blockIdx.y, ty = blockIdx.z;
+  const int CP = C / 8, SP = TC / 8;
+  const P8 b = {((const uint4*)bias)[cg]};
+  const int st = temporal ? (first ? (t + 1) / 2 : t / 2) : t;
+  const int64_t wC = 2 * (int64_t)WL * C;            // one row of the interleaved output
+  const int64_t tC = 2 * (int64_t)HL * wC;           // one frame of it
+#pragma unroll
+  for (int i = 0; i < 2; ++i) {
+    const int yy = 2 * ty + i;
+    if (yy >= HL) continue;
+#pragma unroll
+    for (int j = 0; j < 2; ++j) {
+      const int xx = 2 * tx + j;
+      if (xx >= WL) continue;
+      const int64_t q0 = ((int64_t)t * HL + yy) * WL + xx;
+      P8 ph[4];
+      ph[0].v = ((const uint4*)p0)[q0 * CP + cg];
+      ph[1].v = ((const uint4*)p1)[q0 * CP + cg];
+      ph[2].v = ((const uint4*)p2)[q0 * CP + cg];
+      ph[3].v = ((const uint4*)p3)[q0 * CP + cg];
+      const int64_t s = ((int64_t)st * HL + yy) * WL + xx;
+      P8 s0, s1;
+      if (temporal) {
+        s0.v = ((const uint4*)skip)[s * SP + cg];
+        s1.v = s0.v;                                 // nearest x2 in time: both row parities alike
+      } else {
+        // up block 2: output channel J takes source channel 2*J + (row parity), so read both parities.
+        s0.v = ((const uint4*)skip)[s * SP + 2 * cg];
+        s1.v = ((const uint4*)skip)[s * SP + 2 * cg + 1];
+      }
+      B* out = y + (int64_t)t * tC + (int64_t)(2 * yy) * wC + (int64_t)(2 * xx) * C + 8 * cg;
+      if (!temporal) {
+        // channel J of the output takes source channel 2J + a: even channels from s0, odd from s1.
+        P8 e, o;
+#pragma unroll
+        for (int k = 0; k < 4; ++k) {
+          e.b[k] = s0.b[2 * k];
+          e.b[k + 4] = s1.b[2 * k];
+          o.b[k] = s0.b[2 * k + 1];
+          o.b[k + 4] = s1.b[2 * k + 1];
+        }
+        s0.v = e.v;
+        s1.v = o.v;
+      }
+      P8 r;
+#pragma unroll
+      for (int p = 0; p < 4; ++p) {
+        const P8& sv = (p & 2) ? s1 : s0;            // row parity of this phase
+#pragma unroll
+        for (int k = 0; k < 4; ++k) r.h[k] = __hadd2(__hadd2(ph[p].h[k], b.h[k]), sv.h[k]);
+        const int a = p >> 1, bcol = p & 1;
+        ((uint4*)(out + (int64_t)a * wC + (int64_t)bcol * C))[0] = r.v;
+      }
+    }
+  }
+}
+
+Tensor resample_upsample(const Tensor& x, const std::vector<Tensor>& pw, const Tensor& bias,
+                         const Tensor& skip, bool halo, bool temporal, bool first) {
+  TORCH_CHECK(pw.size() == 4, "resample_upsample needs the four phase kernels");
+  std::vector<Tensor> q(4);
+  for (int a = 0; a < 2; ++a) for (int b = 0; b < 2; ++b) q[2 * a + b] = phase_conv(x, pw[2 * a + b], a, b, halo);
+  const int C = q[0].size(1), D = q[0].size(2), HL = q[0].size(3), WL = q[0].size(4);
+  Tensor y = at::empty({1, C, D, 2 * HL, 2 * WL}, x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+  dim3 block(C / 8, 1), grid((WL + 1) / 2, D, (HL + 1) / 2);
+  upsampler_kernel<<<grid, block, 0, stream()>>>(pc(q[0]), pc(q[1]), pc(q[2]), pc(q[3]), pm(y), pc(bias),
+                                                 pc(skip), C, HL, WL, (int)skip.size(1), temporal, first);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}

fastvideo-kernel/wan22_vae_decode/solution/resample.h

@@ -0,0 +1,51 @@
+#pragma once
+// Resampler stage of the Wan2.2 VAE decoder: the four-phase spatial upsamplers of up_blocks 0,1,2.
+//
+// A resampler is "nearest x2 in space, then a 3x3x3 convolution", which equals four 2x2 low-resolution
+// convolutions, one per output parity (their repeated taps already summed into per-phase kernels by
+// prepare()).  Those convolutions stay direct (4 taps, no transform: F(2x2,2x2) would cut the MACs 1.78x
+// but materialise 4 x 9/4 = 9 x the low-resolution tensor of transformed input and 2.25x the output of
+// products, which prices out as a wash - measured and posted).
+//
+// What is fixed here is the pass after them: the seed's interleave kernel walked the four phase tensors
+// and the DupUp3D source with four 64-bit division/modulo chains per 16 bytes and ran at 3.9 TB/s, half
+// of HBM.  This one walks tiles and channel groups with the grid carrying the coordinates, so the pass
+// is bandwidth bound: it reads the four phase outputs and the DupUp3D source once and writes the
+// interleaved high-resolution tensor once, with bias and the shortcut added in bf16 at the same
+// rounding points the seed used.
+#include <ATen/ATen.h>
+#include <vector>
+using at::Tensor;
+
+// One resampler over one latent chunk.  x is the low-resolution input, already padded by one row top
+// and bottom ([1,C,D,H+2,W] channels-last) with the halo installed at world > 1 (`halo`), and by the two
+// causal history frames in time; pw are the four per-phase 2x2 kernels ([C_out,C_in,1,2,2], phase =
+// 2*row_parity + col_parity), bias [C_out], and skip the DupUp3D source ([1,TC,Ds,HL,WL]).  `temporal`
+// is true for up_blocks 0 and 1 (the shortcut and the phase kernels are nearest-x2 in time as well),
+// false for up block 2, whose shortcut is the channel-pair-to-row rearrangement.  `first` marks the
+// chunk that holds latent frame 0.
+// Returns the interleaved high-resolution tensor [1,C_out,D,2*HL,2*WL] channels-last.
+//
+// The pass moves 100.5 GB of logical traffic for f121x1 (the four phase tensors read 42.5, the DupUp3D
+// source 15.5, the output written 42.5) in 14.8 ms = 6.8 TB/s.
+//
+// THAT 6.8 IS NOT A BOUND AND MUST NOT BE TRACED TO ONE.  It is solved from 100.5 / 14.8, so quoting it
+// against a ceiling is an identity - w06 M18769/M18811, w22 M18463.  The bounds on this board that could
+// apply are: the MIX ceiling for a read+write stage, 5.07-5.59 TB/s at this stage's own f_r (w08's
+// f_r-form, J2654/J2944); w14's per-thread store-loop write arm, 4.86-5.15 (J3082); and w12's 4-byte-
+// per-lane burst arm, 7.35-7.50 (J3029), which is BURST-FLATTERED at 2 KB against the tree's 128 B and
+// which w12 withdrew.  6.9 TB/s was a deflation of that withdrawn arm, so it inherits the arm's defect:
+// it has no f_r, it belongs to a limb with no read, and this is a read+write stage.
+//
+// So the honest statement is that the stage moves 100.5 GB at 6.8 TB/s and that NO CEILING FOR ITS OWN
+// STREAM HAS BEEN MEASURED.  The counters on this pass would supply one; that run has not been made.
+// Do not write "at its byte bound", and do not write "the model is conservative" (106.7 GB / 14.8 ms is
+// 7.21 TB/s, which is ABOVE the mixed ceiling - the same fact as "the model is 6 % low", stated with the
+// sign that does not flatter it).
+//
+// J2574 also counts 1.37x the sectors of its logical read.  That is sector waste on the access pattern
+// and not work the algorithm needs: this kernel reads each phase tensor and the shortcut at the
+// position's own address, with no neighbour re-read and no stencil anywhere in it.
+
+Tensor resample_upsample(const Tensor& x, const std::vector<Tensor>& pw, const Tensor& bias,
+                         const Tensor& skip, bool halo, bool temporal, bool first);

fastvideo-kernel/wan22_vae_decode/solution/solution.json

@@ -0,0 +1 @@
+{"languages":["cuda_cpp"],"entry_point":"binding.cpp::run","destination_passing_style":true,"compile_options":{"cflags":["-O3","-std=c++17"],"cuda_cflags":["-O3","-std=c++17","-lineinfo"],"ld_flags":["-lcuda"]},"dependencies":[]}

fastvideo-kernel/wan22_vae_decode/solution/timeconv.cu

@@ -0,0 +1,122 @@
+// w11: a CUTLASS implicit convolution whose epilogue stores into a caller-given destination with arbitrary
+// (N,D,H,W) strides, and can accumulate on what the destination already holds (beta=1, source == destination).
+// One user: the temporal upsampler, which writes the frame-interleaved video straight out (stride 2 in D).
+// The resampler's phase tensors come from conv.cu's phase_conv, not from this file.
+// Nothing else changes: same tile (256x256x64), same 2-CTA cluster, same TMA im2col, same bf16 epilogue.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/convnd_problem_shape.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/conv/dispatch_policy.hpp"
+#include "cutlass/conv/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/conv/device/conv_universal_adapter.hpp"
+#include "cutlass/conv/kernel/conv_universal.hpp"
+at::Tensor nhwc_empty(int64_t c, int64_t t, int64_t h, int64_t w, const at::Tensor& like);
+using namespace cute;
+using B = __nv_bfloat16;
+using ElementBF16 = cutlass::bfloat16_t;
+using LayoutNDHWC = cutlass::layout::TensorNDHWC;
+using ArchTag = cutlass::arch::Sm100;
+using OperatorClass = cutlass::arch::OpClassTensorOp;
+constexpr auto Fprop = cutlass::conv::Operator::kFprop;
+
+// x: activation [1,C,D,H,W] channels-last; w: [K,C,KT,KH,KW] channels-last; y: the destination, with its own strides.
+template <int M, int N, int CL>
+static void launch_into(const at::Tensor& x, const at::Tensor& w, at::Tensor& y, int ph, int pw, float beta,
+                        int phu = -1, int pwu = -1) {
+  if (phu < 0) phu = ph;
+  if (pwu < 0) pwu = pw;
+  using Tile = Shape<Int<M>, Int<N>, Shape<_64>>;
+  using Cluster = Shape<Int<CL>, _1, _1>;
+  using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto, float, float,
+      ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8,
+      cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+  using Mainloop = typename cutlass::conv::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Fprop, ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, float, Tile, Cluster,
+      cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
+      cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+  using Problem = cutlass::conv::ConvProblemShape<Fprop, 3>;
+  using Kernel = cutlass::conv::kernel::ConvUniversal<Problem, Mainloop, Epilogue>;
+  using Conv = cutlass::conv::device::ConvUniversalAdapter<Kernel>;
+  int C = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4);
+  int K = w.size(0), KT = w.size(2), KH = w.size(3), KW = w.size(4);
+  int OD = y.size(2), OH = y.size(3), OW = y.size(4);
+  TORCH_CHECK(OD == D - KT + 1 && OH == H + ph + phu - KH + 1 && OW == W + pw + pwu - KW + 1,
+              "launch_into: destination shape does not match the convolution");
+  TORCH_CHECK(y.size(1) == K && y.stride(1) == 1, "launch_into: destination must be channel-innermost");
+  Problem problem(cutlass::conv::Mode::kCrossCorrelation,
+                  {1, D, H, W, C}, {x.stride(0), x.stride(2), x.stride(3), x.stride(4), x.stride(1)},
+                  {K, KT, KH, KW, C}, {w.stride(0), w.stride(2), w.stride(3), w.stride(4), w.stride(1)},
+                  {0, ph, pw}, {0, phu, pwu}, {1, 1, 1}, {1, 1, 1}, 1);
+  // The output's own strides: index 0 N, 1 D, 2 H, 3 W, 4 C (the layout the C strides already use).
+  problem.stride_C[0] = y.stride(0);
+  problem.stride_C[1] = y.stride(2);
+  problem.stride_C[2] = y.stride(3);
+  problem.stride_C[3] = y.stride(4);
+  problem.stride_C[4] = 1;
+  typename Kernel::StrideC sc{};
+  typename Kernel::StrideD sd{};
+  cute::for_each(cute::make_seq<cute::rank<0>(sc)>{},
+                 [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
+  cute::for_each(cute::make_seq<cute::rank<0>(sd)>{},
+                 [&](auto i) { cute::get<0, i>(sd) = problem.stride_C[Problem::RankT - 2 - i]; });
+  const ElementBF16* dst = (const ElementBF16*)y.data_ptr();
+  typename Conv::Arguments args{problem, {(const ElementBF16*)x.data_ptr(), (const ElementBF16*)w.data_ptr()},
+                                {{1.f, beta}, dst, sc, (ElementBF16*)y.data_ptr(), sd}};
+  Conv conv;
+  TORCH_CHECK(conv.can_implement(args) == cutlass::Status::kSuccess, "CUTLASS cannot implement shape");
+  size_t size = Conv::get_workspace_size(args);
+  auto workspace = at::empty({(int64_t)size}, x.options().dtype(at::kByte));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto status = conv.initialize(args, workspace.data_ptr(), stream);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialization failed ", int(status));
+  status = conv.run(stream);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS convolution failed ", int(status));
+}
+
+// The temporal upsampler: y[t] = w[:,:,2] x[t-2] + w[:,:,1] x[t-1] + w[:,:,0] x[t] + bias, with frame 0 passed
+// through unchanged when `first`, the two temporal halves of the 2C channels going to consecutive frames.
+// p is the activation already padded with its two causal history frames ([1,C,nf+2,H,W], channels-last).
+// INVARIANT: every element of y is written by head_bias_kernel, and the two parity arms below cover the
+// non-head frames disjointly, each exactly once -- which is what makes beta=1 onto an at::empty safe.
+static __global__ void head_bias_kernel(const B* head, const B* bias, B* y, int C, int T2, int64_t HW, int off) {
+  int64_t n = (int64_t)T2 * HW * (C / 8);
+  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t)gridDim.x * blockDim.x) {
+    int cg = (int)(i % (C / 8)); int64_t r = i / (C / 8);
+    int64_t pos = r % HW; int t = (int)(r / HW);
+    uint4 v;
+    if (t < off) v = ((const uint4*)head)[pos * (C / 8) + cg];
+    else v = ((const uint4*)(bias + ((t - off) & 1) * C))[cg];
+    ((uint4*)y)[i] = v;
+  }
+}
+
+at::Tensor time_upsample(const at::Tensor& p, const at::Tensor& head, const at::Tensor& w, const at::Tensor& bias,
+                         int64_t T, int64_t nf, bool first) {
+  const at::Tensor& ref = p.defined() ? p : head;
+  int C = bias.size(0) / 2, H = ref.size(3), W = ref.size(4);
+  int T2 = (int)(2 * T - (first ? 1 : 0));
+  at::Tensor y = nhwc_empty(C, T2, H, W, ref);
+  int off = first ? 1 : 0;
+  head_bias_kernel<<<1024, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
+      first ? (const B*)head.data_ptr() : nullptr, (const B*)bias.data_ptr(), (B*)y.data_ptr(), C, T2,
+      (int64_t)H * W, off);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  if (nf > 0) {
+    int64_t sN = y.stride(0), sD = 2 * y.stride(2), sH = y.stride(3), sW = y.stride(4);
+    for (int par = 0; par < 2; par++) {
+      auto yv = y.as_strided({1, (int64_t)C, nf, (int64_t)H, (int64_t)W}, {sN, 1, sD, sH, sW},
+                             (int64_t)(off + par) * y.stride(2));
+      launch_into<256, 256, 2>(p, w.narrow(0, (int64_t)par * C, C), yv, 0, 0, 1.f);
+    }
+  }
+  return y;
+}

fastvideo-kernel/wan22_vae_decode/solution/wino_epilogue.cu

@@ -0,0 +1,400 @@
+// Fused F(4x4,3x3) inverse transform for the causal 3x3x3 Winograd convolutions.
+// One pass over the 36 component tensors M produces BOTH
+//   y = A^T M A + bias + residual         (raw conv output: the resnet's residual operand)
+//   n = silu(rmsnorm(y) * sqrt(C) * gain) (the next convolution's normalized, activated,
+//                                          temporal+spatial padded input)
+// so the normalized tensor is never written and read again by a separate norm pass.
+// The RMS statistic needs every output channel of a position, so one block holds all of C
+// for one 4x4 output tile: blockDim.x == C/2, one channel pair per thread.
+//
+// M layout: 36 separate allocations, each [D][TH][TW][C] bf16, C innermost, so a plane's
+// position stride is 1 in C and its plane offset is (((t*TH+ty)*TW+tx)*C + ch).
+// y is [D][H][W][C]; n is [D+PT][H+2PH][W][C].
+//
+// Two lessons from the guild, both measured elsewhere and both applied here:
+//  - w01's J576 (input transform): the head carries no divisions, no predicates and no 64-bit
+//    address chain. 3-D grid (tx,ty,t), a separate EXACT instantiation with the bounds test
+//    removed, and a 32-bit plane byte offset. EXACT is dispatched per call: it holds at
+//    88x156, 176x312 and 352x624, whose widths are multiples of 4 -- NOT at 44x78, where
+//    W=78 leaves the last tile two columns short and the checked path is taken.
+//  - w05's J627 (inverse transform, 5.5 TB/s at 64 registers): no thread loads 4 B of one plane.
+//    The block stages the 36 components of its tile into shared memory cooperatively -- one
+//    plane's C channels is C*2 contiguous bytes, so a warp pass is 16 B per lane and one fully
+//    coalesced 512 B request -- and the consumer reads smem[comp][2c] as a 4-byte pair, which is
+//    128 contiguous bytes per warp and therefore conflict-free. No transpose on either side.
+//
+// Per-element work follows w02's recipe (M137, PR15): bf16->f32 by an integer shift so the
+// conversions leave the XU pipe, one bf16 rounding for the whole normalize*sqrt(C)*gain product
+// instead of the reference's three, SiLU as u + u*tanh(u) with u = pre-activation/2 and
+// tanh.approx.f32, and 0.5*sqrt(C)*rsqrt(sum) as one per-position scalar.
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+using B=__nv_bfloat16;
+static cudaStream_t e_stream(){return at::cuda::getCurrentCUDAStream();}
+static const B* e_ptr(const Tensor& t){return t.defined()?(const B*)t.data_ptr<at::BFloat16>():nullptr;}
+static B* e_mut(Tensor& t){return t.defined()?(B*)t.data_ptr<at::BFloat16>():nullptr;}
+
+struct EpiPtrs{const B* p[36];};
+
+// A of F(4x4,3x3) applied along one axis: out[i] = sum_v A[i][v] v[v], two fp32 lanes per
+// instruction (a6v below is the only form this file uses). add/sub/fma.f32x2 need sm_100+; this
+// tree is built for sm_100a only (solution.json), so no scalar fallback is kept beside it -- a
+// scalar copy would be dead source here, which is why the scalar `a6` was removed rather than left.
+// f32x2: PTX 8.6 gives add/sub/mul/fma on a .b64 operand holding two f32 lanes (sm_100a).
+// This kernel already carries two independent fp32 lanes per thread -- the channel pair a thread
+// owns -- so `ch` IS the SIMD lane and nothing crosses a lane. Each instruction performs exactly
+// the arithmetic the scalar instruction performs, in the same association.
+struct V2{
+  float x,y;
+  __device__ V2(){} __device__ V2(float a,float b):x(a),y(b){}
+};
+__device__ __forceinline__ unsigned long long pk(V2 a){
+  return (unsigned long long)__float_as_uint(a.y)<<32|__float_as_uint(a.x);
+}
+__device__ __forceinline__ V2 upk(unsigned long long u){
+  V2 r; r.x=__uint_as_float((unsigned)u); r.y=__uint_as_float((unsigned)(u>>32)); return r;
+}
+__device__ __forceinline__ V2 vadd(V2 a,V2 b){
+  unsigned long long r; asm("add.rn.f32x2 %0,%1,%2;":"=l"(r):"l"(pk(a)),"l"(pk(b))); return upk(r);
+}
+__device__ __forceinline__ V2 vsub(V2 a,V2 b){
+  unsigned long long r; asm("sub.rn.f32x2 %0,%1,%2;":"=l"(r):"l"(pk(a)),"l"(pk(b))); return upk(r);
+}
+__device__ __forceinline__ V2 vfmas(V2 a,float c,V2 acc){
+  V2 k; k.x=k.y=c; unsigned long long r;
+  asm("fma.rn.f32x2 %0,%1,%2,%3;":"=l"(r):"l"(pk(a)),"l"(pk(k)),"l"(pk(acc))); return upk(r);
+}
+__device__ __forceinline__ V2 b2v(uint32_t w){
+  V2 r; r.x=__uint_as_float(w<<16); r.y=__uint_as_float(w&0xffff0000u); return r;
+}
+// The transform of A above, two lanes per instruction.
+__device__ __forceinline__ void a6v(V2 v0,V2 v1,V2 v2,V2 v3,V2 v4,V2 v5,V2* y){
+  V2 a=vadd(v1,v2), b=vsub(v1,v2), c=vadd(v3,v4), d=vsub(v3,v4);
+  y[0]=vadd(vadd(v0,a),c); y[1]=vfmas(d,2.f,b); y[2]=vfmas(c,4.f,a); y[3]=vadd(vfmas(d,8.f,b),v5);
+}
+// Exact bf16 pair -> two floats without the XU pipe: a bf16 is the top half of an fp32.
+__device__ __forceinline__ float2 b2f(uint32_t w){
+  return make_float2(__uint_as_float(w<<16),__uint_as_float(w&0xffff0000u));
+}
+__device__ __forceinline__ float tanh_ap(float x){float r;asm("tanh.approx.f32 %0, %1;":"=f"(r):"f"(x));return r;}
+// silu(x) with u = x/2: x*sigmoid(x) = u + u*tanh(u).
+__device__ __forceinline__ float silu2(float u){return __fmaf_rn(u,tanh_ap(u),u);}
+__device__ __forceinline__ uint32_t add2(uint32_t a,uint32_t b){
+  __nv_bfloat162 x=*(const __nv_bfloat162*)&a,y=*(const __nv_bfloat162*)&b;
+  __nv_bfloat162 r=__hadd2(x,y); return *(const uint32_t*)&r;
+}
+__device__ __forceinline__ uint32_t pack2(float x,float y){
+  __nv_bfloat162 r=__floats2bfloat162_rn(x,y); return *(const uint32_t*)&r;
+}
+__device__ __forceinline__ float sq8(const uint4& v){
+  float2 a=b2f(v.x),b=b2f(v.y),c=b2f(v.z),d=b2f(v.w);
+  return (__fmaf_rn(a.x,a.x,a.y*a.y)+__fmaf_rn(b.x,b.x,b.y*b.y))
+        +(__fmaf_rn(c.x,c.x,c.y*c.y)+__fmaf_rn(d.x,d.x,d.y*d.y));
+}
+
+// Stage the 36 components of one 4x4 tile into shared memory: 36*C/8 tasks of 16 bytes,
+// 9 per thread at any C. Each warp's 32 lanes take 32 consecutive 16-byte slots, i.e. 512
+// contiguous bytes of one plane.
+//
+// The copy is `cp.async.cg` (16 B, global->shared, L2-cached, NO register round trip) rather than
+// the `uint4` load-then-store it replaces.  Measured reason (J3871, wino_epi_kernel<1024,1>, ncu,
+// 3 passes): the binding FIXABLE stall is `mio_throttle 2.11` cycles per issued instruction --
+// second only to `not_selected 2.73`, which is not a defect (it means the SM already has eligible
+// warps to pass over) -- and `l1tex__throughput` is the highest pipe at 62.7-64.7 %.  The old form
+// is a global load AND a shared store per 16 B, so it spends two MIO slots and a writeback
+// register per slot; this form spends one slot and no register.  Per block at C=1024 that is
+// 9*512 = 4,608 pairs, i.e. 9,216 MIO instructions becoming 4,608.
+//
+// Bytes moved, alignment and the resulting V are IDENTICAL: same comp/off decomposition, same
+// 16-byte slot, same plane offset inside smem, so the transform reads what it read before and the
+// output is bit-for-bit the old output.  This is a scheduling change, not a numerical one.
+//
+// The commit/wait pair sits here rather than at the call site because the transform consumes the
+// staging immediately -- there is no independent work to overlap the transfer with inside one
+// tile.  Overlap across tiles needs a persistent loop and is a separate change, not this one.
+// The form above is measured, not preferred: `cp.async.cg` beats the `uint4` load+store it
+// replaced by 12.7 ms at f121x1 (J4745, one job, one clock, A/A 0.2 ms).
+template<int C>
+__device__ __forceinline__ void stage36(B* sm,const EpiPtrs& mp,unsigned src_tile){
+  constexpr int CV=C/2;
+  constexpr int PER=C/8;                 // 16-byte slots per plane
+  const unsigned smbase=(unsigned)__cvta_generic_to_shared(sm);
+  #pragma unroll
+  for(int i=0;i<9;i++){
+    const int task=threadIdx.x+i*CV;
+    const int comp=task/PER, off=task%PER;
+    const unsigned dst=smbase+(unsigned)((comp*C+off*8)*2);
+    const char* src=(const char*)mp.p[comp]+src_tile+off*16;
+    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"::"r"(dst),"l"(src));
+  }
+  asm volatile("cp.async.commit_group;\n"::);
+  asm volatile("cp.async.wait_group 0;\n"::);
+}
+
+template<int C,bool EXACT>
+__global__ __launch_bounds__(C/2) void wino_epi_kernel(
+    EpiPtrs mp,B* __restrict__ y,B* __restrict__ n,
+    const B* __restrict__ bias,const B* __restrict__ residual,const B* __restrict__ gain,
+    int H,int W,int TH,int TW,int PT,int PH,bool wy){
+  // `wy` is the CALLER's assertion that nobody reads this launch's y. It is only ever false where
+  // the very next conv has idx == fnext and therefore takes the fused `prepared` edge (see
+  // binding.cpp: resnet's first conv, whose consumer is its own second conv). The n tensor is
+  // written either way, so this can never be false without a n to carry the activation.
+  constexpr int CV=C/2;          // channel pairs per position
+  constexpr int WP=CV/32;        // warps per block
+  const int c=threadIdx.x;
+  const int tx=blockIdx.x, ty=blockIdx.y, t=blockIdx.z;
+  extern __shared__ __align__(16) B smem[];
+  B* sm=smem;                                              // 36 planes x C, this tile
+  // 32-bit plane byte offset: D*TH*TW*C*2 bytes, where D is the PIXEL frame count at this level
+  // (4*chunk-3 at up block 1, not the latent chunk size -- the earlier "34" was the latent count).
+  // chunk_size is NOW min(16*world, 61) -- this PR is what moved the rule off `8*world`, so chunk
+  // 32 at world 4 no longer exists and the worst case moves with it. Over the shipping configs the
+  // largest src_tile is world 2 / chunk 32 = 4.39e8 bytes and world 1 / chunk 16 = 4.29e8, both
+  // ~10x inside UINT32_MAX and both enforced at run time by the guard further down this file.
+  // The `8.8e8` that stood here was 4x too large and it described a configuration this PR retires.
+  // A chunk size that more than doubles again, or a wider clip, must widen this to unsigned long long.
+  const unsigned src_tile=(unsigned)((((t*TH+ty)*TW+tx)*C)*2u);
+
+  stage36<C>(sm,mp,src_tile);
+  __syncthreads();
+
+  V2 acc[4][4];
+  #pragma unroll
+  for(int i=0;i<4;i++)
+    #pragma unroll
+    for(int b=0;b<4;b++) acc[i][b]=V2(0.f,0.f);
+
+  const B* q0=sm+2*c;                     // this thread's channel pair, in each plane
+  #pragma unroll
+  for(int u=0;u<6;u++){
+    V2 f[6];
+    #pragma unroll
+    for(int v=0;v<6;v++) f[v]=b2v(*(const uint32_t*)(q0+(6*u+v)*C));
+    V2 q[4];
+    a6v(f[0],f[1],f[2],f[3],f[4],f[5],q);
+    #pragma unroll
+    for(int b=0;b<4;b++){
+      const V2 v=q[b];
+      if(u==0) acc[0][b]=vadd(acc[0][b],v);
+      else if(u==1){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vadd(acc[1][b],v);acc[2][b]=vadd(acc[2][b],v);acc[3][b]=vadd(acc[3][b],v);}
+      else if(u==2){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vsub(acc[1][b],v);acc[2][b]=vadd(acc[2][b],v);acc[3][b]=vsub(acc[3][b],v);}
+      else if(u==3){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vfmas(v,2.f,acc[1][b]);acc[2][b]=vfmas(v,4.f,acc[2][b]);acc[3][b]=vfmas(v,8.f,acc[3][b]);}
+      else if(u==4){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vfmas(v,-2.f,acc[1][b]);acc[2][b]=vfmas(v,4.f,acc[2][b]);acc[3][b]=vfmas(v,-8.f,acc[3][b]);}
+      else {acc[3][b]=vadd(acc[3][b],v);}
+    }
+  }
+
+  const int Hp=H+2*PH;
+  float2 biasv=make_float2(0.f,0.f),gainv=make_float2(0.f,0.f);
+  if(bias) biasv=b2f(*(const uint32_t*)&((const __nv_bfloat162*)bias)[c]);
+  if(gain) gainv=b2f(*(const uint32_t*)&((const __nv_bfloat162*)gain)[c]);
+  const long yb=(((long)t*H+4*ty)*W+4*tx)*CV+c;
+
+  const __nv_bfloat162 bv=__floats2bfloat162_rn(biasv.x,biasv.y);
+  if(!n){
+  // No prepared tensor: write the raw conv output in the natural per-thread layout.
+  #pragma unroll
+  for(int i=0;i<4;i++){
+    const int yy=4*ty+i;
+    #pragma unroll
+    for(int b=0;b<4;b++){
+      const int xx=4*tx+b;
+      const long off=(long)(i*W+b)*CV;
+      const bool ok=EXACT?true:(yy<H&&xx<W);
+      __nv_bfloat162 v=__floats2bfloat162_rn(acc[i][b].x,acc[i][b].y);
+      if(bias) v=__hadd2(v,bv);
+      if(residual&&ok) v=__hadd2(v,((const __nv_bfloat162*)residual)[yb+off]);
+      if(ok) ((__nv_bfloat162*)y)[yb+off]=v;
+    }
+  }
+  }
+
+  if(n){
+    // Remapped epilogue (w01's M655, measured 1.58x on the standalone inverse): the thread that
+    // computes the transform owns a channel pair, so it would write 4 bytes per output, which
+    // caps the pass at the 4-byte row of the stream table. Instead the 16 outputs go to shared
+    // memory with the bias already applied, one barrier, and the epilogue is re-laid so a warp
+    // owns one position and each lane an octet of channels: the residual read and the y write are
+    // one 512-byte transaction per warp instruction, and the sum of squares becomes a single-warp
+    // butterfly because a warp now holds every channel of its position.
+    __syncthreads();                       // the components are dead; their buffer is reused
+    B* so=smem;                            // so[p*C + ch], bias applied
+    #pragma unroll
+    for(int i=0;i<4;i++)
+      #pragma unroll
+      for(int b=0;b<4;b++){
+        __nv_bfloat162 v=__floats2bfloat162_rn(acc[i][b].x,acc[i][b].y);
+        if(bias) v=__hadd2(v,bv);
+        ((__nv_bfloat162*)so)[(i*4+b)*CV+c]=v;
+      }
+    __syncthreads();
+    constexpr int NPW=16/WP;               // positions per warp (4, 2 or 1)
+    constexpr int CW=C/32;                 // channels per lane (8, 16 or 32)
+    constexpr int NU=CW/8;                 // 16-byte chunks per lane (1, 2 or 4)
+    const int wp=c>>5, ln=c&31;
+    const float hs=0.5f*sqrtf((float)C);
+    uint4 gv[NU];
+    // Chunk map: lane ln takes chunk (k*32+ln), i.e. channel (k*32+ln)*8, so the 32 lanes of one
+    // warp instruction cover 32 consecutive chunks = 512 contiguous bytes. The previous map gave
+    // lane ln the block [ln*CW, ln*CW+CW), which puts neighbours CW*2 bytes apart: contiguous at
+    // C=256 (16 B), half a line at C=512 (32 B), a quarter at C=1024 (64 B) -- and at C=1024 that
+    // is a 4x sector over-fetch on the residual read and on both stores. w21's J1096 gives it in
+    // counters (L1 hit 5.32%, L2 hit 0.20%, DRAM 65.6%): no reuse, so every over-fetched sector
+    // is paid at DRAM. Measured at D=8, min of three rounds: C=256 1.000x (identical expression at
+    // NU=1), C=512 1.244x, C=1024 1.431x. It permutes which channel each k reads, so the 4-term
+    // sum of squares accumulates in a different order and s rounds differently in its last bits --
+    // f121x1 max|d| 0.1605 -> 0.147, rms unchanged to five digits.
+    #pragma unroll
+    for(int k=0;k<NU;k++)
+      gv[k]=gain?*(const uint4*)((const char*)gain+((long)(k*32+ln)*8)*2)
+               :make_uint4(0x3f803f80u,0x3f803f80u,0x3f803f80u,0x3f803f80u);
+    #pragma unroll
+    for(int j=0;j<NPW;j++){
+      const int p=wp*NPW+j, yy=4*ty+(p>>2), xx=4*tx+(p&3);
+      const bool ok=EXACT?true:(yy<H&&xx<W);
+      uint4 yv4[NU]; float sq=0.f;
+      #pragma unroll
+      for(int k=0;k<NU;k++){
+        const int off=(k*32+ln)*8;
+        uint4 v=*(const uint4*)((const char*)so+((long)p*C+off)*2);
+        if(residual&&ok){
+          // (long)t first: the whole offset must be 64-bit, not the last factor. `H` here is this
+          // RANK'S BAND height, not the full 352 -- chunk() narrows dim 3 of the latent grid (44
+          // rows) to [44r/W,44(r+1)/W) before resnet15 and every up block, so block 3 carries 88
+          // rows at world 4 and 44 at world 8. The band shrinks as the temporal chunk grows, which
+          // is why a rank's working set is near-constant -- BUT THE CENSUS THAT STOOD HERE WAS
+          // TRUE ONLY OF `chunk_size = 8*world`, AND THIS PR MOVES THE RULE TO `min(16*world, 61)`.
+          // Post-change THREE cells pass INT_MAX and THE LARGEST IS WORLD 2, so it is listed
+          // first: narrowing against world 4 because "chunk 61 looks biggest" is exactly the
+          // mislabelled-maximum this paragraph exists to prevent.
+          //   world 2 / chunk 32   D=125, rows 176, 624x256   = 3.51e9   <-- THE MAXIMUM
+          //   world 1 / chunk 16   D=61,  rows 352            = 3.43e9
+          //   world 4 / chunk 61   D=241, rows 88             = 3.39e9
+          // World 2's chunk 32 is THIS PR's OWN value at world 2 (min(16*2, 61) = 32), not a
+          // foreign or historical form. ALL THREE ARE PAST INT_MAX. This site is
+          // correct because it chains `(long)` here and again on `nbase` below, so the 64-bit chain
+          // is now LOAD-BEARING at those three cells rather than merely belt-and-braces. Do not
+          // narrow it back on the strength of an `inside INT_MAX` reading from the old rule.
+          const uint4 r=*(const uint4*)((const char*)residual
+                        +((long)((((long)t*H+yy)*W+xx)*C+off))*2);
+          v.x=add2(v.x,r.x); v.y=add2(v.y,r.y); v.z=add2(v.z,r.z); v.w=add2(v.w,r.w);
+        }
+        yv4[k]=v; sq+=sq8(v);
+      }
+      #pragma unroll
+      for(int d=16;d;d>>=1) sq+=__shfl_xor_sync(0xffffffffu,sq,d);
+      if(ok){
+        const float s=hs*rsqrtf(fmaxf(sq,1e-24f));
+        const long ybase=((((long)t*H+yy)*W+xx)*C)*2;
+        const long nbase=(((((long)(t+PT)*Hp+PH+yy)*W+xx)*C))*2;
+        #pragma unroll
+        for(int k=0;k<NU;k++){
+          const long cof=(long)((k*32+ln)*8)*2;
+          if(wy) *(uint4*)((char*)y+ybase+cof)=yv4[k];
+          uint4 u=yv4[k], g=gv[k], o;
+          float2 f0=b2f(u.x), f1=b2f(u.y), f2=b2f(u.z), f3=b2f(u.w);
+          float2 g0=b2f(g.x), g1=b2f(g.y), g2=b2f(g.z), g3=b2f(g.w);
+          o.x=pack2(silu2(f0.x*s*g0.x),silu2(f0.y*s*g0.y));
+          o.y=pack2(silu2(f1.x*s*g1.x),silu2(f1.y*s*g1.y));
+          o.z=pack2(silu2(f2.x*s*g2.x),silu2(f2.y*s*g2.y));
+          o.w=pack2(silu2(f3.x*s*g3.x),silu2(f3.y*s*g3.y));
+          *(uint4*)((char*)n+nbase+cof)=o;
+        }
+      }
+    }
+  }
+}
+
+// Zero the spatial padding rows of every frame of the prepared tensor.
+__global__ void epi_border(B* n,int C,int T,int Hp,int W){
+  const int CV=C/2; const long per=(long)2*W*CV;
+  const long total=(long)T*per;
+  for(long i=(long)blockIdx.x*blockDim.x+threadIdx.x;i<total;i+=(long)gridDim.x*blockDim.x){
+    const int t=(int)(i/per); const long r=i%per; const int edge=(int)(r/((long)W*CV));
+    const long r2=r%((long)W*CV); const int w=(int)(r2/CV),cc=(int)(r2%CV);
+    const long idx=(((long)t*Hp+(edge?Hp-1:0))*W+w)*CV+cc;
+    ((__nv_bfloat162*)n)[idx]=__floats2bfloat162_rn(0.f,0.f);
+  }
+}
+__global__ void epi_zero(B* n,long pairs){
+  const __nv_bfloat162 z=__floats2bfloat162_rn(0.f,0.f);
+  for(long i=(long)blockIdx.x*blockDim.x+threadIdx.x;i<pairs;i+=(long)gridDim.x*blockDim.x)
+    ((__nv_bfloat162*)n)[i]=z;
+}
+
+template<int C,bool EXACT>
+static void launch_epi(const dim3& grid,const dim3& block,size_t smem,
+                       const EpiPtrs& p,B* y,B* n,
+                       const B* b,const B* r,const B* g,
+                       int H,int W,int TH,int TW,int PT,int PH,bool wy,cudaStream_t s){
+  auto k=wino_epi_kernel<C,EXACT>;
+  if(smem>48*1024){
+    static bool done[2]={false,false};
+    if(!done[EXACT?1:0]){
+      C10_CUDA_CHECK(cudaFuncSetAttribute(k,cudaFuncAttributeMaxDynamicSharedMemorySize,36*C*2));
+      done[EXACT?1:0]=true;
+    }
+  }
+  k<<<grid,block,smem,s>>>(p,y,n,b,r,g,H,W,TH,TW,PT,PH,wy);
+}
+
+WinoOut wino_epilogue(const std::vector<Tensor>& m,int H,int W,const Tensor& bias,
+                      const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph,
+                      bool write_y){
+  // `write_y=false` is admissible ONLY when the caller's next conv has idx == the `fnext` this call
+  // prepares for, so that next conv takes the `prepared` edge and never reads the returned y. Out.y
+  // is still a distinct allocation: its unread contents are undefined, not zero.
+  TORCH_CHECK(m.size()==36&&m[0].dim()==5,"wino_epilogue: M list ",m.size()," rank ",m[0].dim());
+  const int C=m[0].size(1),D=m[0].size(2),TH=m[0].size(3),TW=m[0].size(4);
+  TORCH_CHECK(C==256||C==512||C==1024,"wino_epilogue: unsupported channel count");
+  auto s=e_stream();
+  EpiPtrs p; for(int i=0;i<36;i++) p.p[i]=(const B*)m[i].data_ptr<at::BFloat16>();
+  WinoOut out;
+  out.y=nhwc_empty(C,D,H,W,m[0]);
+  const int Hp=H+2*ph;
+  if(gamma.defined()){
+    out.n=nhwc_empty(C,D+pt,Hp,W,m[0]);
+    B* np=e_mut(out.n);
+    const long pairs=(long)(D+pt)*Hp*W*(C/2);
+    if(pt>0&&!hist.defined())
+      epi_zero<<<(int)std::min((long)2048,((long)pt*Hp*W*(C/2)+4095)/4096),256,0,s>>>(np,(long)pt*Hp*W*(C/2));
+    if(ph>0)
+      epi_border<<<(int)std::min((long)2048,(pairs+4095)/4096),256,0,s>>>(np,C,D+pt,Hp,W);
+    if(pt>0&&hist.defined()){
+      // `hist` is a `last_frames` output: a COMPACT [C,pt,H,W] channels-last-3d tensor, so its
+      // per-frame SOURCE pitch is H*W*C -- unpadded. Only the DESTINATION is padded: out.n is
+      // [C,D+pt,Hp,W] with pitch Hp*W*C, and these frames land in its interior rows [ph,ph+H).
+      // The two pitches are one of each and neither is a typo: passing Hp*W*C as the source pitch
+      // would read pt frames of garbage, and passing H*W*C as the destination pitch would write
+      // the pad rows away. If `last_frames` changes its shape or its pitch, this copy changes.
+      const size_t row=(size_t)H*W*C*sizeof(B);
+      C10_CUDA_CHECK(cudaMemcpy2DAsync(np+(long)ph*W*C,(size_t)Hp*W*C*sizeof(B),
+                                       e_ptr(hist),row,row,pt,cudaMemcpyDeviceToDevice,s));
+    }
+  }
+  // src_tile in the kernel is the only 32-bit index in this file, so assert here what it relies on
+  // rather than trusting the kernel's comment. (The y and n offsets in the kernel body are formed in
+  // 64-bit and are not bounded by this.) Checked against UINT32_MAX, not INT32_MAX: the index is
+  // unsigned, and a check against the signed bound would refuse shapes that are in fact safe.
+  { const unsigned long long off=(((unsigned long long)(D-1)*TH+(TH-1))*TW+(TW-1))*(unsigned long long)C*2ull;
+    TORCH_CHECK(off<=0xFFFFFFFFull,"wino_epilogue: 32-bit tile offset ",off,
+                " exceeds UINT32_MAX; widen src_tile in wino_epilogue.cu"); }
+  const bool exact=(TH*4==H&&TW*4==W);
+  const dim3 grid(TW,TH,D);
+  const dim3 block(C/2);
+  const size_t shmem=(size_t)36*C*2;
+  B* yp=e_mut(out.y); B* np=e_mut(out.n);
+  const B* bp=e_ptr(bias); const B* rp=e_ptr(residual); const B* gp=e_ptr(gamma);
+  if(C==256){ if(exact) launch_epi<256,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s);
+              else launch_epi<256,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); }
+  else if(C==512){ if(exact) launch_epi<512,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s);
+                   else launch_epi<512,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); }
+  else { if(exact) launch_epi<1024,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s);
+         else launch_epi<1024,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return out;
+}

fastvideo-kernel/wan22_vae_decode/solution/wino_in.cu

@@ -0,0 +1,247 @@
+// F(4x4,3x3) spatial input transform of the Winograd path of the Wan2.2 VAE decoder.
+//
+// ONE PASS: `winograd_input_fast` replaces `kernels.cu`'s `winograd_input` for the 18 Winograd
+// convolutions (1024/512/256 channels; any other channel count keeps the general kernel), and
+// its output V is BIT-IDENTICAL to it -- same patch origin, same block-uniform fast-path
+// condition, same B^T-then-B order with the in-place column update, same 6a+b plane index, same
+// single bf16 rounding at the store, same plane stride.  Nothing else in the tree moves but the
+// one call site in binding.cpp and one declaration in kernels.h.
+//
+// Why it is faster, from the profile of the kernel it replaces (w01's J549): 581.8 M
+// warp-instructions in 1.031 ms at 59.5% issue-active, of which 43.5% are ALU/integer against
+// 47.6% FMA and 4.9% LSU; DRAM 23.6%, FMA pipe 32.4%.  It sits at 82% of its own instruction
+// issue bound and the integer address math is the only large removable block:
+//   1. a 3-D grid, blockIdx = (tx, ty, t) and threadIdx = the channel pair, so nothing divides;
+//   2. one block-uniform fast path for the interior tiles, where the 36 loads carry no
+//      predicate at all -- only the first and last tile of a row and the first and last tile
+//      row can touch the time pad or the frame edge -- and the patch row pointer advances by
+//      the row stride, so the 36 addresses are 36 LSU instructions with folded offsets;
+//   3. the 36 plane stores walk one pointer by V's plane stride instead of doing an int64
+//      multiply per store.
+//
+// MORE BLOCKS PER TILE INSTEAD OF ONE: SPLIT = 4 AT C = 1024, 2 AT C = 512 AND 256.
+// A tile's C/2 channel pairs used to be covered by one block of C/2 threads, so `blockDim` and
+// `C` were the same variable and the grid held exactly one block per tile.  A block is now the
+// tile's width divided by SPLIT, and SPLIT blocks of the grid cover one tile.  At C = 1024 that
+// moves the level from one block of 16 warps (32.08 ms) to five blocks of 4 warps (22.42 ms,
+// 20 warps resident); at C = 512 and C = 256 the same move is worth 1.05 % and -0.4 % -- i.e.
+// nothing -- so those two stay at SPLIT = 2.  See the table below for why.
+//
+// `__launch_bounds__(CP, 512 / CP)` IS LOAD-BEARING AND IS HALF THE MEASUREMENT.  It pins the
+// register budget at 65536/512 = 128 at EVERY width -- exactly what main's
+// `__launch_bounds__(512, 1)` grants at its one wide instantiation.  Without it ptxas may spend
+// up to the architectural cap of 255 registers per thread (65536/CP is itself 512 at CP = 128,
+// so it is the 255 that binds there), floor(65536/(CP * REG)) falls, the resident warps follow,
+// and the same arm runs 2.25x SLOWER.  The counterfactual is measured at ONE instantiation:
+// J3705 ran `wino_in_kernel_st<256,2>` both ways -- `__launch_bounds__(256, 1)` at 59.11 ms
+// against `(256, 2)` at 26.26 ms, one integer and one body.  At the shipped CP = 128 the no-pin
+// case gives floor(65536/(128*255)) = 2 blocks of 4 warps against the pinned 5, i.e. 20 resident
+// warps down to 8 -- a 2.5x cut, not a halving, and the reason C = 1024 needs the pin most.
+// The declaration and the launch are one change; neither is the arm without the other.
+//
+// Measured (J3868, one build, one process, A B B B B A over SPLIT = 1 2 4 8 2 1, f121x1
+// --kernels, 1132 MHz on every leg, exit 0).  Both control legs are shown so the drift is visible:
+//                     SPLIT=1      SPLIT=2          SPLIT=4        SPLIT=8      2 -> 4
+//   C = 1024 (shipped)  32.08/32.08  26.26/26.53      22.42          22.27       1.1775
+//   C =  512 (2 kept)   64.06/63.92  54.60/55.44      54.45          55.45       1.0105
+//   C =  256 (2 kept)   60.47/60.66  59.56/60.05      60.04        (n/a)         0.9962
+//   CALL               864.7        847.6            844.2          845.2
+//   A/A 0.3 %, B/B 1.0-1.5 %; max |d| 0.1543 of 0.4984, rms 0.002093 of 0.005847, PSNR 59.6 dB,
+//   IDENTICAL to every printed digit on all six legs.
+//   The C = 256 SPLIT = 8 cell is `n/a`, not a decline: J3868's probe only instantiated the CP
+//   values it needed and SPLIT = 8 at C = 256 would ask for CP = 16, which that build did not
+//   contain.  The head's launcher is an exhaustive switch with no fall-back path of any kind.
+//
+// ONLY THE 2 -> 4 MOVE IS A C = 1024-ONLY GAIN, and that is why the other two keep SPLIT = 2:
+//   the shipped SPLIT = 2 already takes the 1 -> 2 step on BOTH of them -- 64.06 -> 54.60 at
+//   C = 512 is 1.1733x and 60.47 -> 59.56 at C = 256 is 1.0183x, the +15.5 % and +1.5 % below --
+//   so neither is a rung that gains nothing.  What is declined on them is the FOURTH block:
+//   C = 512's 1.0105 is inside the B/B drift (its own legs read 54.60 and 55.44, a 1.5 % spread);
+//   C = 256's 0.9962 REVERSES -- its legs read 59.56, 60.05 and 60.66, a 1.8 % spread, and 60.04
+//   sits inside it, so the option is declined rather than merely unpriced;
+//   SPLIT = 8 at C = 1024 is 22.27 against 22.42 -- 0.7 % at a 1.0 % drift, a tie, not a win.
+//
+// WHY C = 1024 AND NOT THE OTHERS.  With REG 96 the realized resident warps are
+// floor(65536/(96*blockDim)) * blockDim/32, i.e. 16 at blockDim 512 or 256, 20 at 128 or 64, and
+// 21 at 32.  Main's C = 1024 case used blockDim 512, which rounds to ONE block and never reaches
+// past 16 warps although 2048/96 = 21.3 is its ceiling; `__launch_bounds__` cannot fix that,
+// because 2 * 512 * 96 = 98,304 > 65,536.  Only a narrower block can, and narrowing it is SPLIT.
+// C = 512 and C = 256 already reach 20 warps at SPLIT = 2, so there is nothing there to buy.
+// The three levels' rates were 2.609 / 3.340 / 3.760 TB/s (J3195's counted bytes).  DERIVED, not
+// re-measured: at fixed byte counts rate_after/rate_before = t_before/t_after, so from J3868's own
+// legs the after-rates are 2.609*32.08/22.42 = 3.733, 3.340*64.06/55.44 = 3.859 and
+// 3.760*60.47/59.56 = 3.817 TB/s -- a 44.1 % spread closed to 3.4 %.  The largest move is the
+// slow level's own (+43 %), and the fast two barely move (+15.5 % and +1.5 %), so what happened is
+// that the slow level was walked up to where the other two already sat, not that a ceiling moved.
+// An earlier draft of this comment stated 3.952 for C = 256, which needs 60.47 -> 57.5 ms and no
+// leg of J3868 is below 59.56; it was a number carried over from a withdrawn row and it is wrong.
+//
+// `J3761`'s rebuild of the round-1 source: BUILD_RC=0, REG:96 STACK:0 SHARED:0 LOCAL:0.
+// `J3889`'s build of this source: the same, with `wino_in_kernel<128,4>` at REG:95 (one register
+// lower, which changes no block count: 65536/(128*95) = 5.39 -> 5, the same as at 96).
+// Correctness at worlds 1 and 4 on all three short clips: J3889 f1x1/f5x1/f17x1 ok at
+// x0.41/x0.41/x0.65 and J3908 f1x4/f5x4/f17x4 ok at x0.45/x0.45/x0.66, every leg inside main's
+// own margins.
+//
+// The share of the call is flat in clip length (17.6 % at f121x1, 17.9 % at f481x1) and this
+// function takes no rank, no world and no process group, so a world-1 cell is the geomean cell.
+//
+// Measured by w04 on main's own tree (kernel#1, this change alone, f121x1 at 1849 MHz): the
+// three instantiations total 114.4 ms against 221.8 ms of `winograd_input_kernel`, and the call
+// 969.9 -> 862.5 ms (-11.1%), with V bit-identical -- f1/f5/f17 keep the seed's max |d| and rms
+// to the digit -- and f481x1 ok at x0.38 of the max and x0.41 of the rms allowance.  w21 wrote
+// the kernel; w08 owns this file.
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+
+using B = __nv_bfloat16;
+static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+
+__device__ __forceinline__ void wbt6(float v0, float v1, float v2, float v3, float v4, float v5, float* y) {
+  float a = v4 - 4 * v2, b = v3 - 4 * v1, c = v4 - v2, d = 2 * (v3 - v1);
+  y[0] = 4 * v0 - 5 * v2 + v4; y[1] = a + b; y[2] = a - b; y[3] = c + d; y[4] = c - d; y[5] = 4 * v1 - 5 * v3 + v5;
+}
+
+// One block per (part-tile, frame); threadIdx.x = the channel pair within this block's share.
+// CP = C/2/SPLIT is the block size.  The shipped dispatch is SPLIT = 4 at C = 1024 and SPLIT = 2
+// at C = 512 and 256, so blockDim is 128 / 128 / 64 at C = 1024 / 512 / 256 -- but SPLIT is NOT
+// uniform and no comment here may claim it is.  The invariant that actually guarantees coverage is
+//     CP * SPLIT == C / 2                                    (128*4 = 128*2 = 64*2 = C/2)
+// enforced by the `static_assert` in `winograd_input_fast` below, together with `grid.x == SPLIT*TW`
+// from the same `constexpr`.  `(blockIdx.x % SPLIT) * CP + threadIdx.x` then walks [0, SPLIT*CP) =
+// [0, C/2) exactly once -- every channel pair of a tile is written once and no pair twice -- and
+// `blockIdx.x / SPLIT` over a grid of SPLIT*TW hits each tile exactly SPLIT times, for any SPLIT:
+// the remainder cycles 0..SPLIT-1 within each group of SPLIT consecutive blockIdx, so TW need not be
+// a multiple of SPLIT.
+//
+// SPLIT = 1 with CP = C/2 gives back main's one-block-per-tile BODY, but not its attribute (this
+// one is `__launch_bounds__(CP, 512/CP)`, which equals main's `(CP, 1)` only at CP = 512) and not
+// its grid shape, so do not use it as an exact revert; revert by checking out the parent commit.
+//
+// The launched blockDim must equal CP.  It does at all three call sites below, from the same
+// template value, so the two cannot drift apart.
+template <int CP, int SPLIT>
+__global__ void __launch_bounds__(CP, 512 / CP)
+wino_in_kernel(const B* __restrict__ x, B* __restrict__ v, int C, int D, int H, int W, int PH, int TH, int TW) {
+  const int cp = (blockIdx.x % SPLIT) * CP + threadIdx.x;
+  const int tx = blockIdx.x / SPLIT;
+  const int ty = blockIdx.y, t = blockIdx.z;
+  const int HP = H + 2 * PH;
+  const int y0 = 4 * ty - 1 + PH, x0 = 4 * tx - 1;
+  const int64_t rowStep = (int64_t)W * C;
+  const B* base = x + (int64_t)t * HP * rowStep + (int64_t)y0 * rowStep + (int64_t)x0 * C + 2 * cp;
+  float d[2][6][6];
+  // Uniform across the block, so it is one branch and not 36 predicates.
+  if ((y0 >= 0) && (y0 + 5 < HP) && (x0 >= 0) && (x0 + 5 < W)) {
+    const B* prow = base;
+#pragma unroll
+    for (int a = 0; a < 6; a++) {
+#pragma unroll
+      for (int b = 0; b < 6; b++) {
+        float2 f = __bfloat1622float2(*(const __nv_bfloat162*)(prow + b * C));
+        d[0][a][b] = f.x; d[1][a][b] = f.y;
+      }
+      prow += rowStep;
+    }
+  } else {
+    const B* prow = base;
+#pragma unroll
+    for (int a = 0; a < 6; a++) {
+      const int y = y0 + a;
+      const bool ok = (y >= 0) && (y < HP);
+#pragma unroll
+      for (int b = 0; b < 6; b++) {
+        const int xx = x0 + b;
+        if (ok && xx >= 0 && xx < W) {
+          float2 f = __bfloat1622float2(*(const __nv_bfloat162*)(prow + b * C));
+          d[0][a][b] = f.x; d[1][a][b] = f.y;
+        } else {
+          d[0][a][b] = 0.f; d[1][a][b] = 0.f;
+        }
+      }
+      prow += rowStep;
+    }
+  }
+  // B^T d down the patch's 6 columns, then B along its 6 rows, one bf16 rounding at the end.
+#pragma unroll
+  for (int k = 0; k < 2; k++)
+#pragma unroll
+    for (int b = 0; b < 6; b++) {
+      float q[6];
+      wbt6(d[k][0][b], d[k][1][b], d[k][2][b], d[k][3][b], d[k][4][b], d[k][5][b], q);
+#pragma unroll
+      for (int a = 0; a < 6; a++) d[k][a][b] = q[a];
+    }
+  B* dst = v + (int64_t)(((t * TH + ty) * TW + tx) * C + 2 * cp);
+  const int64_t plane = (int64_t)D * TH * TW * C;
+#pragma unroll
+  for (int a = 0; a < 6; a++) {
+    float q[2][6];
+#pragma unroll
+    for (int k = 0; k < 2; k++)
+      wbt6(d[k][a][0], d[k][a][1], d[k][a][2], d[k][a][3], d[k][a][4], d[k][a][5], q[k]);
+#pragma unroll
+    for (int b = 0; b < 6; b++, dst += plane) *(__nv_bfloat162*)dst = __floats2bfloat162_rn(q[0][b], q[1][b]);
+  }
+}
+
+Tensor winograd_input_fast(const Tensor& x, int ph) {
+  const int C = x.size(1), D = x.size(2), H = x.size(3) - 2 * ph, W = x.size(4);
+  const int TH = (H + 3) / 4, TW = (W + 3) / 4;
+  // The 18 Winograd convolutions are the 1024-, 512- and 256-channel ones; anything else
+  // keeps the general kernel.
+  if (C != 1024 && C != 512 && C != 256) return winograd_input(x, ph);
+  auto v = at::empty({36, D, TH, TW, C}, x.options());
+  // SPLIT = 4 at C = 1024, 2 at C = 512 and 256.  CP * SPLIT = C/2 at every case, so a tile's
+  // channel pairs are covered exactly once; grid.x is SPLIT * TW, a block count and not a tile
+  // count.  blockDim is CP, taken from the same template value, so the two cannot drift.
+  //
+  // Why 4 only at C = 1024 (J3868: 2 -> 4 on the rung at each C, one build, one process):
+  //   C=1024  <256,2> 26.26/26.53 -> <128,4> 22.42   1.1775   warps 16 -> 20
+  //   C= 512  <128,2> 54.60/55.44 -> < 64,4> 54.45   1.0105   warps 20 -> 20   inside the 1.5 % drift
+  //   C= 256  < 64,2> 59.56/60.05 -> < 32,4> 60.04   0.9962   warps 20 -> 21   it REVERSES
+  // So the resolved gain is C = 1024 alone, and it is the only rung whose split buys warps:
+  // with REG 96 the realized resident warps are floor(65536/(96*blockDim))*blockDim/32, i.e.
+  // 16 at blockDim 512 or 256, 20 at 128 or 64, 21 at 32.  The step off 16 is the whole of it.
+  // SPLIT = 8 at C = 1024 measured 22.27 against 22.42 -- 0.7 % at a 1.0 % B/B drift, not
+  // resolvable, so it is not taken.  A tie is not a win.
+  //
+  // ONE `constexpr` PER CASE BINDS THE TEMPLATE ARGUMENT, THE BLOCK WIDTH, THE GRID MULTIPLIER AND
+  // THE COVERAGE INVARIANT.  Three literals in three places is a desync waiting to become a bug:
+  // raise SPLIT without raising grid.x and `tx = blockIdx.x / SPLIT` stays inside TW, the kernel
+  // silently computes one block's share of every tile and the rest of V is never written; raise
+  // grid.x without SPLIT and `tx` runs past TW, computing x0 = 4*tx-1 for tiles that do not exist,
+  // which is an out-of-bounds read AND write rather than a slowdown.  The `static_assert` makes the
+  // first a compile error and the single `constexpr` makes the second unrepresentable.
+  const B* xp = (const B*)x.data_ptr<at::BFloat16>();
+  B* vp = (B*)v.data_ptr<at::BFloat16>();
+  auto st = stream();
+  switch (C) {
+    case 1024: {
+      constexpr int SPLIT = 4, CP = 1024 / 2 / SPLIT;
+      static_assert(CP * SPLIT == 1024 / 2, "CP * SPLIT must equal C/2 or a tile is not covered");
+      const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D);
+      wino_in_kernel<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
+      break;
+    }
+    case 512: {
+      constexpr int SPLIT = 2, CP = 512 / 2 / SPLIT;
+      static_assert(CP * SPLIT == 512 / 2, "CP * SPLIT must equal C/2 or a tile is not covered");
+      const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D);
+      wino_in_kernel<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
+      break;
+    }
+    case 256: {
+      constexpr int SPLIT = 2, CP = 256 / 2 / SPLIT;
+      static_assert(CP * SPLIT == 256 / 2, "CP * SPLIT must equal C/2 or a tile is not covered");
+      const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D);
+      wino_in_kernel<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
+      break;
+    }
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return v;
+}

fastvideo-kernel/wan22_vae_decode/tools/audit.py

@@ -0,0 +1,359 @@
+"""The stack rules of a solution, checked on its sources; and who launched each kernel of a call, found at run time.
+
+The stack is CUDA C++ (with inline PTX), compiled ahead of time into one extension module (see loader.py), with
+communication through the ProcessGroup it is given (NCCL) or CUDA peer access and IPC between the ranks.
+
+Static rules (check): solution.json declares the languages in LANGUAGES, and RULES maps a rule's name to what it
+looks for in the C/C++/CUDA sources and why. Kinds of match:
+
+    include   the target of an #include
+    text      the source with comments removed
+    flag      solution.json: a compile option
+    vendored  an #include target that must resolve to a file of the solution itself
+
+A task's spec.py may extend the table: RULES = {name: rule, or None to drop a default rule}.
+
+Run time (trace): one call runs under CUPTI. Every kernel is attributed to the library that launched it, from the
+native call stack of its launch: kernels launched from inside cuBLAS/cuBLASLt, cuDNN or ATen (libtorch) are library
+kernels, and so are Triton's (a stack the rules exclude); NCCL's are communication, which the rules allow, counted
+apart; kernels launched from the solution's own module or by driver calls are the solution's, whatever their names,
+so kernels built from CUTLASS/CuTe templates copied into the solution count as its own. Kernels whose launch cannot be
+traced (graph replays, device-side launches) are classified by name. Library kernels are allowed, as placeholders:
+their count and their share of the call's GPU time show where the solution still depends on them. Only kernel
+launches call back into this module (enabling the whole driver API would call back on NCCL's watchdog thread, which
+queries events while it holds a lock the launching thread needs: a deadlock under the GIL).
+
+    python tools/audit.py [solution]    the static rules
+"""
+
+import argparse
+import collections
+import ctypes
+import json
+import os
+import re
+import sys
+
+CPP_FILES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh", ".hh", ".hxx", ".inl", ".inc")
+LANGUAGES = ["cuda_cpp"]
+
+RULES = {
+    "vendored-cutlass": {
+        "vendored": r"(cutlass|cute)/.*",
+        "why": "CUTLASS/CuTe headers are copied into the solution and maintained there (cutlass/..., cute/... next to "
+               "the entry file), not taken from the installed CUTLASS"},
+    "runtime-build": {
+        "include": r"(.*/)?(nvrtc|nvJitLink|nvPTXCompiler)\.h",
+        "text": r"\b(nvrtc[A-Z]\w*|nvJitLink[A-Z]\w*|nvPTXCompiler\w*|cuModuleLoad\w*|"
+                r"cuLink(Create|AddData|AddFile|Complete)\w*|cuLibraryLoad\w*|cudaLibraryLoad\w*)\s*\(",
+        "flag": r"-l(nvrtc|nvJitLink|nvptxcompiler)\w*",
+        "why": "nothing is compiled or loaded at run time (NVRTC, nvJitLink, PTX or cubins loaded through the driver): "
+               "device code, inline PTX included, is compiled by nvcc with the solution"},
+    "python-from-cpp": {
+        "text": r"\b(py(bind11)?::module_?::import|py(bind11)?::(exec|eval)(_file)?|PyImport_\w+|PyRun_\w+)\b",
+        "why": "run() stays in C++ and runs no Python: Triton, torch.compile, torch.utils.cpp_extension and other "
+               "kernel toolkits would come in that way"},
+}
+
+
+def rule_table(extra=None):
+    """The rule table in force: the defaults, changed by a task's RULES (None drops a rule)."""
+    table = dict(RULES)
+    for name, rule in (extra or {}).items():
+        if rule is None:
+            table.pop(name, None)
+        else:
+            table[name] = rule
+    return table
+
+
+def _match(pattern, s):
+    return re.fullmatch(pattern, s) is not None
+
+
+def _strip_cpp_comments(text):
+    """Removes // and /* */ comments (string literals are not parsed; good enough for these rules)."""
+    text = re.sub(r"/\*.*?\*/", " ", text, flags=re.S)
+    return re.sub(r"//[^\n]*", "", text)
+
+
+def _includes(text):
+    return [(m.group(1), text.count("\n", 0, m.start()) + 1)
+            for m in re.finditer(r'^\s*#\s*include\s*[<"]([^>"]+)[>"]', text, flags=re.M)]
+
+
+def check(solution_dir, rules, languages=None):
+    """Violations of `rules` by a solution's sources: a list of "file:line: what (rule: why)" strings.
+
+    languages: the languages the solution may declare in solution.json (default LANGUAGES)."""
+    languages = LANGUAGES if languages is None else languages
+    out = []
+    try:
+        with open(os.path.join(solution_dir, "solution.json")) as f:
+            spec = json.load(f)
+    except (OSError, ValueError) as e:
+        return [f"solution.json: unreadable ({e})"]
+    if not isinstance(spec, dict):
+        return ["solution.json: must hold an object"]
+    declared = spec.get("languages")
+    if not isinstance(declared, list) or not declared or any(lang not in languages for lang in declared):
+        out.append(f"solution.json: languages {declared} are not allowed here (allowed: {languages})")
+    options = spec.get("compile_options") if isinstance(spec.get("compile_options"), dict) else {}
+    flags = [f for v in options.values() if isinstance(v, list) for f in v]
+    for name, rule in rules.items():
+        for flag in flags:
+            if "flag" in rule and _match(rule["flag"], str(flag)):
+                out.append(f"solution.json: compile option {flag!r} ({name}: {rule['why']})")
+    for root, dirs, files in os.walk(solution_dir):
+        dirs[:] = sorted(d for d in dirs if d != "__pycache__")
+        for file in sorted(files):
+            if not file.endswith(CPP_FILES):
+                continue
+            path = os.path.join(root, file)
+            try:
+                with open(path, encoding="utf-8") as f:
+                    text = f.read()
+            except (OSError, UnicodeDecodeError):
+                continue
+            out += _check_cpp(solution_dir, os.path.relpath(path, solution_dir), text, rules)
+    return out
+
+
+def _check_cpp(solution_dir, rel, text, rules):
+    out = []
+    code = _strip_cpp_comments(text)
+    includes = _includes(code)
+    for name, rule in rules.items():
+        for target, line in includes:
+            if "include" in rule and _match(rule["include"], target):
+                out.append(f"{rel}:{line}: #include <{target}> ({name}: {rule['why']})")
+            if "vendored" in rule and _match(rule["vendored"], target):
+                here = os.path.dirname(os.path.join(solution_dir, rel))
+                if not any(os.path.isfile(os.path.join(d, target)) for d in (here, solution_dir)):
+                    out.append(f"{rel}:{line}: #include <{target}> is not in the solution ({name}: {rule['why']})")
+        if "text" in rule:
+            for m in re.finditer(rule["text"], code):
+                line = code.count("\n", 0, m.start()) + 1
+                out.append(f"{rel}:{line}: {m.group(0)!r} ({name}: {rule['why']})")
+                break
+    return out
+
+
+# Run-time audit.
+
+# The library a kernel was launched from, by the path of the innermost library on the launch's native stack.
+LAUNCHERS = [
+    (r".*/libcublas(Lt)?[^/]*", "cuBLAS"),
+    (r".*/libcudnn[^/]*", "cuDNN"),
+    (r".*/lib(torch|c10)[^/]*", "ATen"),
+    (r".*triton.*", "Triton"),
+    (r".*/libnccl[^/]*", "NCCL"),
+    (r".*/lib(cufft|curand|cusparse|cusolver|cutensor|nvshmem|npp)[^/]*", "CUDA library"),
+]
+# Frames skipped on the way out: CUPTI and this module's callback, the Python interpreter, the CUDA runtime/driver.
+SKIP = r"(.*/)?(libffi[^/]*|_ctypes[^/]*|python[^/]*|libpython[^/]*|callback_wrappers[^/]*|cupti[^/]*|libcupti[^/]*|" \
+       r"libcuda\.so[^/]*|libcudart[^/]*|libc\.so[^/]*|ld-linux[^/]*|\?)"
+# Fallback by kernel name, for launches that could not be traced.
+NAMES = [
+    (r"(void )?at::.*|.*at::native::.*|_ZN2at.*|.*at_cuda_detail.*|_ZN3c10.*", "ATen"),
+    (r".*(nvjet|cublas|gemv2?[NT]?_|gemmSN|splitKreduce|xmma|sgemm|hgemm|dgemm).*|cutlass3x_.*|cutlass_\d.*", "cuBLAS"),
+    (r".*cudnn.*", "cuDNN"),
+    (r"triton_.*", "Triton"),
+    (r"(void )?nccl.*|_Z\d+nccl.*", "NCCL"),
+]
+# Libraries whose kernels are communication: allowed outright, and counted apart from the placeholders.
+COMMUNICATION = ("NCCL",)
+# The CUPTI activity buffers: few and large, since a decode launches tens of thousands of kernels.
+BUFFER_BYTES = 32 << 20
+
+_libc = None
+_demangler = None
+
+
+class _DlInfo(ctypes.Structure):
+    _fields_ = [("fname", ctypes.c_char_p), ("fbase", ctypes.c_void_p), ("sname", ctypes.c_char_p),
+                ("saddr", ctypes.c_void_p)]
+
+
+def _launcher():
+    """Path of the innermost library on the current native stack that is not CUPTI, Python or the CUDA runtime."""
+    global _libc
+    if _libc is None:
+        _libc = ctypes.CDLL("libc.so.6")
+        _libc.backtrace.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_int]
+        _libc.dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(_DlInfo)]
+    buf = (ctypes.c_void_p * 96)()
+    n = _libc.backtrace(buf, 96)
+    for i in range(n):
+        info = _DlInfo()
+        path = info.fname.decode() if buf[i] and _libc.dladdr(buf[i], ctypes.byref(info)) and info.fname else "?"
+        if not re.fullmatch(SKIP, path):
+            return path
+    return ""
+
+
+def demangle(name):
+    global _demangler
+    if _demangler is None:
+        try:
+            lib = ctypes.CDLL("libstdc++.so.6")
+            lib.__cxa_demangle.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_size_t),
+                                           ctypes.POINTER(ctypes.c_int)]
+            lib.__cxa_demangle.restype = ctypes.c_char_p
+            _demangler = lib
+        except OSError:
+            _demangler = False
+    if not _demangler or not name.startswith("_Z"):
+        return name
+    status = ctypes.c_int()
+    out = _demangler.__cxa_demangle(name.encode(), None, None, ctypes.byref(status))
+    return out.decode() if status.value == 0 and out else name
+
+
+def short_name(name, limit=110):
+    """A kernel's demangled name without its return type, parameter list and the at::native/anonymous namespaces."""
+    text = demangle(name).replace("(anonymous namespace)::", "")
+    depth, start, end = 0, 0, len(text)
+    for i, c in enumerate(text):
+        if c in "<(":
+            if c == "(" and depth == 0 and i > 0:
+                end = i
+                break
+            depth += 1
+        elif c in ">)":
+            depth -= 1
+        elif c == " " and depth == 0:
+            start = i + 1
+    text = text[start:end] or text
+    for noise in ("at::native::", "at::", "std::"):
+        text = text.replace(noise, "")
+    return text if len(text) <= limit else text[:limit - 3] + "..."
+
+
+def classify(name, launcher, traced):
+    """The library a kernel belongs to (None: the solution's own)."""
+    if traced:
+        for pattern, label in LAUNCHERS:
+            if re.fullmatch(pattern, launcher):
+                return label
+        return None
+    for pattern, label in NAMES:
+        if re.fullmatch(pattern, name) or re.fullmatch(pattern, demangle(name)):
+            return label
+    return None
+
+
+Op = collections.namedtuple("Op", "kind name bytes launcher library ns")
+
+
+def launch_callbacks(cupti):
+    """The driver API callbacks of kernel launches (cuLaunchKernel, cuLaunchKernelEx, cuGraphLaunch, ...)."""
+    return [int(getattr(cbid, "value", cbid)) for name, cbid in cupti.driver_api_trace_cbid.__members__.items()
+            if "Launch" in name and "HostFunc" not in name]
+
+
+def trace(fn, *args):
+    """Calls fn(*args) once under CUPTI; returns its GPU operations in start order, as Op(kind, name, bytes,
+    launcher, library, ns): kind is kernel, memcpy or memset; library is None for the solution's own kernels; ns is
+    the operation's GPU time."""
+    import torch
+    from cupti import cupti
+
+    launches, records, errors = {}, [], []
+
+    def on_api(userdata, domain, cbid, data):
+        try:
+            if data.callback_site == cupti.ApiCallbackSite.API_ENTER:
+                graph = "Graph" in data.function_name
+                launches[data.correlation_id] = "" if graph else _launcher()
+        except Exception as e:  # never let an exception cross into the driver
+            errors.append(repr(e))
+
+    def requested():
+        return BUFFER_BYTES, 0
+
+    def completed(activities):
+        for a in activities:
+            if a.kind == cupti.ActivityKind.CONCURRENT_KERNEL:
+                records.append((a.start, a.end, "kernel", a.name, 0, a.correlation_id,
+                                bool(a.graph_id or a.is_device_launched)))
+            elif a.kind == cupti.ActivityKind.MEMCPY:
+                records.append((a.start, a.end, "memcpy", "", a.bytes, a.correlation_id, False))
+            elif a.kind == cupti.ActivityKind.MEMSET:
+                records.append((a.start, a.end, "memset", "", a.bytes, a.correlation_id, False))
+
+    kinds = (cupti.ActivityKind.CONCURRENT_KERNEL, cupti.ActivityKind.MEMCPY, cupti.ActivityKind.MEMSET)
+    torch.cuda.synchronize()
+    sub = cupti.subscribe(on_api, 0)
+    try:
+        for cbid in launch_callbacks(cupti):
+            cupti.enable_callback(1, sub, cupti.CallbackDomain.DRIVER_API, cbid)
+        for k in kinds:
+            cupti.activity_enable(k)
+        cupti.activity_register_callbacks(requested, completed)
+        fn(*args)
+        torch.cuda.synchronize()
+        cupti.activity_flush_all(1)
+    finally:
+        for k in kinds:
+            cupti.activity_disable(k)
+        cupti.unsubscribe(sub)
+        cupti.finalize()
+    if errors:
+        raise RuntimeError(f"the launch trace failed: {errors[0]}")
+    ops = []
+    for start, end, kind, name, nbytes, corr, untraceable in sorted(records):
+        launcher = launches.get(corr)
+        traced = kind == "kernel" and not untraceable and bool(launcher)
+        library = classify(name, launcher or "", traced) if kind == "kernel" else None
+        ops.append(Op(kind, name, nbytes, os.path.basename(launcher or ""), library, end - start))
+    return ops
+
+
+def placeholder(library):
+    """Whether a library's kernels are placeholders (cuBLAS, cuDNN, ATen, ...), as opposed to the solution's own or
+    communication."""
+    return bool(library) and library not in COMMUNICATION
+
+
+def kernels(ops):
+    """What a call launched, by kernel: [[library or "", name, launches, ns]] by GPU time, the largest first; copies
+    and sets as the kernels "(memcpy)" and "(memset)"."""
+    table = collections.defaultdict(lambda: [0, 0])
+    for op in ops:
+        key = (op.library or "", op.name if op.kind == "kernel" else f"({op.kind})")
+        table[key][0] += 1
+        table[key][1] += op.ns
+    return sorted(([lib, name, n, ns] for (lib, name), (n, ns) in table.items()), key=lambda r: -r[3])
+
+
+def library_share(ops):
+    """The share of a call's GPU time (the summed time of its operations) spent in placeholder library kernels."""
+    total = sum(op.ns for op in ops)
+    return sum(op.ns for op in ops if placeholder(op.library)) / total if total else 0.0
+
+
+def comm_share(ops):
+    """The share of a call's GPU time spent in communication kernels (NCCL)."""
+    total = sum(op.ns for op in ops)
+    return sum(op.ns for op in ops if op.library in COMMUNICATION) / total if total else 0.0
+
+
+def main():
+    here = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, here)
+    import spec  # noqa: E402
+
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("solution", nargs="?", default="solution")
+    a = p.parse_args()
+    table = rule_table(getattr(spec, "RULES", None))
+    found = check(a.solution, table)
+    for v in found:
+        print(v)
+    print(f"stack rules: {'ok' if not found else f'{len(found)} violations'} ({len(table)} rules)")
+    sys.exit(1 if found else 0)
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/wan22_vae_decode/tools/bench.py

@@ -0,0 +1,771 @@
+"""Correctness, timing and library-kernel audit of a solution: the Wan2.2 VAE decode of latents into video, on one, two
+or four GB200 GPUs of one node. The task's CI runs this same file (see ci.py).
+
+The task (spec.py has the details)
+--------
+A solution decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T-3, 704, 1248] (bf16, in [-1, 1]) with
+the Wan2.2 VAE's decoder, as diffusers' AutoencoderKLWan decodes them (reference.py in the kernel repo). It is a
+directory (default solution/, see loader.py) compiled into one extension module that exports
+
+    run(z, weights, rank, world, pg, out)                    positional, destination passing
+    prepare(weights, rank, world, pg) -> list[Tensor]       optional
+
+called in one process per GPU: z the full latents on every rank; weights the decoder's 110 bf16 tensors on the rank's
+GPU, in the order of spec.WEIGHTS (or what prepare returned); pg torch.distributed's NCCL ProcessGroup over the ranks;
+out the video on rank 0, from torch.empty and never zeroed (every element must be written), and an empty tensor on the
+other ranks. When run has returned on every rank, rank 0's out holds the video. prepare, if exported, runs once per
+process before any timed call and is not timed; its result replaces the weights in every run.
+
+The workloads (spec.WORKLOADS) are clips decoded at a world size, named <clip>x<world>. A clip is the first T latent
+frames of one of three real 20 s clips, named by its video frames: f1, f5, f17 (T = 1, 2, 5) and f121, f241, f481
+(T = 31, 61, 121). The long clips at 1, 2 and 4 GPUs are the nine timed workloads (f121x1 ... f481x4), whose geomean
+is the objective; the short clips at 1 and 4 GPUs (f1x1 ... f17x4) are checked and never timed.
+
+Running it
+----------
+In a job with as many GPUs as the largest world size it runs, with the task's data at $DATA (default /data), from the
+repo root:
+
+    python tools/bench.py                             every workload this job has the GPUs for: check, audit, time
+    python tools/bench.py --workload f5x1 f17x1       a subset; the short clips take seconds
+    python tools/bench.py --workload f121x1 --no-time correctness only
+    python tools/bench.py --kernels                   also list the audited calls' kernels by GPU time
+    python tools/bench.py --solution other            another solution directory
+
+A run
+-----
+The solution is built once, by this process, exactly as it is deployed (loader.py; a development copy, tools/, takes a
+build its node already made of the same sources, flags and toolchain from the node's cache, unless --no-cache). Then,
+per world size in turn (1, 2, 4), torch.multiprocessing spawns one process per rank, rank r on GPU r of the job's
+CUDA_VISIBLE_DEVICES. Each initializes NCCL over 127.0.0.1 with its device (the default group, the solution's pg) and
+a gloo group for the harness's own barriers, imports the build (in its working directory, as deployed), moves the
+weights ($DATA/decoder_weights_bf16.pt, checked against spec.WEIGHTS) to its GPU and calls prepare, if the module
+exports it. The ranks then run that world size's workloads in spec.WORKLOADS' order. A rank that raises, dies or
+hangs (WORKLOAD_SEC per workload, SESSION_START_SEC before the first) fails the workload it was at: all the ranks'
+processes are stopped, and the workloads after it run in new ones. This process runs no solution code.
+
+Every call gets z at a fresh address (a copy of the clip's latents) and out at a fresh address, 256-byte aligned and
+filled with NaN on rank 0, so that an element run leaves unwritten fails the check. run must not write z, nor the
+weights (compared after each workload with what prepare left), nor start threads: each fails the workload.
+
+Correctness
+-----------
+The first call of a workload, its warm-up, is checked on rank 0 against the clip's fp32 decode ($DATA/ref_<clip>.pt):
+every element finite and within [-1, 1]; and with d = out - ref in fp32, max |d| <= min(2, 3 x bf16_max_abs) and
+rms(d) <= 3 x bf16_rms, the error of the release's bf16 decode of that clip ($DATA/calibration.json; spec.py,
+Correctness). The line shows max |d| and rms(d), each with its allowance and the share of it used, and the PSNR,
+10 log10(4 / mse). Every later call's output (the audited call's, the trials') must agree with the checked one:
+lose none of the elements it wrote (none left NaN or overflowed), and be within the same allowance of it everywhere
+else. A race, or a state carried from call to call, fails the workload.
+
+Timing (the timed workloads)
+------
+After the warm-up come the trials: 3 when the warm-up took under 1.5 s, else 2. A trial is one call timed by wall
+clock between barriers on all ranks: synchronize the device, barrier, start, run, synchronize, barrier, stop. The
+time of record is rank 0's; the latency of record the mean of the trials, and the fastest trial is shown too. While
+the trials run, this process reads the SM clock of the world's GPUs through NVML every CLOCK_PERIOD_S (a rank inside
+run may hold its GIL for the whole call, so the ranks do not read it themselves); the mean of rank 0's GPU is shown
+with the latency, and the range over the world's GPUs where they differ by 2% or more: the clocks are not locked
+alike on every GPU, and this says at what clock the times were taken. Which processes hold the world's GPUs is
+asked (nvidia-smi) before every trial and after the last: a workload timed while another process held one of them is
+marked SHARED, since its times do not compare with those of another job.
+
+Audit
+-----
+At world 1, one more call per timed clip runs under CUPTI, and each kernel is attributed to the library that launched
+it (see audit.py). Library kernels (cuBLAS/cuBLASLt, cuDNN, ATen) are allowed as placeholders: their count per call
+and their share of the call's GPU time show where the solution still depends on them. NCCL's kernels are
+communication, reported apart and not counted as placeholders. --kernels lists the audited calls' kernels.
+
+Bound
+-----
+Each timed workload's latency is shown with its bound (bounds.py with spec.bounds: the decode's FLOPs spread evenly
+over the world's GPUs at the 16-bit tensor-core peak, or the video and the weights at the memory bandwidth, whichever
+is slower) and the gap, latency / bound. The peak assumes the SM clock GUILD_SM_CLOCK_MHZ where the environment sets
+it (the guild's: 2062 MHz, the GB200's highest), else the GPU's highest.
+
+In CI, and from tools/
+----------------------
+ci.py runs this file in child processes: --save DIR has rank 0 save each workload's checked output there (node-local
+scratch) instead of checking it, so that CI judges it in a process that runs no solution code; --no-check skips the
+check (the timing runs), --no-audit the audit, and --json writes the report. The kernel repo holds copies of this
+file, spec.py, loader.py, audit.py and bounds.py in tools/, which run as they do here: they read the data at $DATA
+(default /data) and the solution from solution/ under the current directory.
+"""
+
+import argparse
+import collections
+import ctypes
+import functools
+import json
+import math
+import os
+import random
+import socket
+import statistics
+import subprocess
+import sys
+import threading
+import time
+import traceback
+from datetime import timedelta
+from multiprocessing import connection
+
+# The caching allocator stays as torch configures it by default: without expandable segments, whose memory
+# cudaIpcGetMemHandle cannot export, since a solution may share its buffers between ranks by CUDA IPC.
+import torch
+import torch.distributed as dist
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import audit  # noqa: E402
+import bounds  # noqa: E402
+import loader  # noqa: E402
+import spec  # noqa: E402
+
+# Trials per timed workload, fewer when one call takes long.
+TRIALS, LONG_TRIALS, LONG_CALL_S = 3, 2, 1.5
+# How often this process reads rank 0's SM clock while ranks run.
+CLOCK_PERIOD_S = 0.02
+NVML_CLOCK_SM = 1  # the nvmlClockType_t of the SM clock
+# Limits that stop a runaway: a session's start (NCCL, the module, the weights, prepare), and one workload (its calls,
+# its check or its save); a call of the release's bf16 decode takes 15 s at the longest clip.
+SESSION_START_SEC = 900
+WORKLOAD_SEC = 1800
+STOP_GRACE_SEC = 15
+# Elements of a video compared at once, in fp32 (1 GiB per operand).
+CHUNK_ELEMENTS = 1 << 28
+
+
+def gpu_holders():
+    """{GPU UUID (lower case, without "GPU-"): ids of the processes that hold it} now, as nvidia-smi gives them, or
+    None when it cannot say."""
+    try:
+        got = subprocess.run(["nvidia-smi", "--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader"],
+                             capture_output=True, text=True, timeout=30)
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    if got.returncode != 0:
+        return None
+    out = collections.defaultdict(set)
+    for line in got.stdout.splitlines():
+        gpu, _, pid = line.partition(",")
+        if pid.strip().isdigit():
+            out[gpu.strip().lower().removeprefix("gpu-")].add(int(pid))
+    return out
+
+
+def others(questions, uuids, ours):
+    """How many processes besides this run's held one of these GPUs at any of the questions (gpu_holders' answers), or
+    None when one could not be asked. This run's processes are known by their ids where nvidia-smi gives the ids this
+    process sees (one PID namespace, the usual case: then a rank that also holds another rank's GPU is still one of
+    ours); otherwise each GPU is taken to hold one of ours, the rank on it."""
+    if not questions or any(q is None for q in questions):
+        return None
+    total = 0
+    for uuid in uuids:
+        seen = set().union(*(q.get(uuid, set()) for q in questions))
+        total += len(seen - ours) if seen & ours else max(0, len(seen) - 1)
+    return total
+
+
+@functools.cache
+def sm_clock(uuid):
+    """A function that reads the SM clock of the GPU with this UUID now, in MHz (None when a read fails), or None where
+    NVML cannot be loaded or does not know the GPU. NVML in this process takes microseconds per read."""
+    try:
+        nvml = ctypes.CDLL("libnvidia-ml.so.1")
+        if nvml.nvmlInit_v2() != 0:
+            return None
+        handle, mhz = ctypes.c_void_p(), ctypes.c_uint()
+        if not uuid or nvml.nvmlDeviceGetHandleByUUID(f"GPU-{uuid}".encode(), ctypes.byref(handle)) != 0:
+            return None
+    except (OSError, AttributeError):
+        return None
+
+    def read():
+        ok = nvml.nvmlDeviceGetClockInfo(handle, NVML_CLOCK_SM, ctypes.byref(mhz)) == 0
+        return mhz.value if ok and mhz.value else None
+
+    return read if read() else None
+
+
+class Clock:
+    """The SM clocks of the world's GPUs (rank order), read every CLOCK_PERIOD_S by a thread of this process, as
+    (monotonic time, [MHz or None per GPU])."""
+
+    def __init__(self, uuids):
+        self.reads, self.samples, self.done = [sm_clock(u) for u in uuids], [], threading.Event()
+        self.thread = threading.Thread(target=self._run, daemon=True) if any(self.reads) else None
+        if self.thread:
+            self.thread.start()
+
+    def _run(self):
+        while not self.done.wait(CLOCK_PERIOD_S):
+            self.samples.append((time.monotonic(), [read() if read else None for read in self.reads]))
+
+    def stop(self):
+        self.done.set()
+        if self.thread:
+            self.thread.join()
+
+    def means(self, windows):
+        """Per GPU, the mean of its reads within these (start, stop) windows of time.monotonic(), or None without
+        any."""
+        inside = [mhz for t, mhz in list(self.samples) if any(a <= t <= b for a, b in windows)]
+        out = []
+        for i in range(len(self.reads)):
+            got = [m[i] for m in inside if m[i]]
+            out.append(statistics.mean(got) if got else None)
+        return out
+
+
+class Pool:
+    """A workload's buffers for one tensor: each call's at a fresh address, shifted by multiples of 256 bytes."""
+
+    ALIGN, MAX_SHIFT = 256, 2048
+
+    def __init__(self, shape, dtype, device, calls, seed):
+        rng = random.Random(seed)
+        self.offsets = [0]
+        for _ in range(calls - 1):
+            self.offsets.append(self.offsets[-1] + rng.randint(1, self.MAX_SHIFT // self.ALIGN))
+        self.shape, self.numel = tuple(shape), math.prod(shape)
+        self.block = max(1, self.ALIGN // torch.empty(0, dtype=dtype).element_size())
+        self.buffer = torch.empty(self.numel + self.offsets[-1] * self.block, dtype=dtype, device=device)
+        self.calls = 0
+
+    def next(self):
+        start = self.offsets[self.calls] * self.block
+        self.calls += 1
+        return self.buffer[start:start + self.numel].view(self.shape)
+
+
+def fingerprint(tensors):
+    """Each tensor's bytes summed as integers: whether run() wrote the weights."""
+    out = []
+    for t in tensors:
+        raw = t.detach().contiguous().view(-1).view(torch.uint8)
+        try:
+            raw = raw.view(torch.int32)
+        except RuntimeError:
+            pass
+        out.append(int(raw.sum(dtype=torch.int64)))
+    return out
+
+
+def judge(out, ref, allowed):
+    """A video against the reference, or a later call's video against the checked one's: {"passed", "nonfinite",
+    "outside" (elements beyond [-1, 1]), "max_abs", "rms", "psnr_db", "max_allowed", "rms_allowed"}, d = out - ref where
+    both are finite (rms and PSNR over those elements, None where there is none), in fp32 on out's GPU (the current one
+    for a video on the CPU), in chunks of frames, the squares summed in fp64."""
+    device = out.device if out.is_cuda else torch.device("cuda", torch.cuda.current_device())
+    step = max(1, CHUNK_ELEMENTS // max(1, out[:, :, :1].numel()))
+    nonfinite = outside = compared = unmatched = 0
+    max_abs = squares = 0.0
+    for f in range(0, out.shape[2], step):
+        x = out[:, :, f:f + step].to(device, torch.float32)
+        y = ref[:, :, f:f + step].to(device, torch.float32)
+        finite = torch.isfinite(x)
+        nonfinite += int((~finite).sum())
+        outside += int((x.abs() > 1).sum())
+        # Where the checked call itself left an element unfinite, a later call is not compared there (the check said).
+        both = finite & torch.isfinite(y)
+        compared += int(both.sum())
+        unmatched += int((~finite & torch.isfinite(y)).sum())
+        d = torch.where(both, x - y, 0.0)
+        max_abs = max(max_abs, float(d.abs().max()))
+        squares += float(d.square().sum(dtype=torch.float64))
+        del x, y, finite, both, d
+    mse = squares / compared if compared else None
+    rms = math.sqrt(mse) if compared else None
+    passed = (nonfinite == 0 and outside == 0 and compared == out.numel() and max_abs <= allowed["max_abs"]
+              and rms <= allowed["rms"])
+    return {"passed": passed, "nonfinite": nonfinite, "unmatched": unmatched, "outside": outside,
+            "max_abs": max_abs if compared else None, "rms": rms, "psnr_db": 10 * math.log10(4 / mse) if mse else None,
+            "max_allowed": allowed["max_abs"], "rms_allowed": allowed["rms"]}
+
+
+def verdict(c):
+    """A check (judge's result) in a few words."""
+    parts = []
+    if c["nonfinite"]:
+        parts.append(f"{c['nonfinite']} elements not finite (NaN: never written, or overflowed)")
+    if c["outside"]:
+        parts.append(f"{c['outside']} elements outside [-1, 1]")
+    if c["max_abs"] is None:
+        parts.append("no element to compare")
+    else:
+        parts.append(f"max |d| {c['max_abs']:.4g} of {c['max_allowed']:.4g} (x{c['max_abs'] / c['max_allowed']:.2f})")
+        parts.append(f"rms {c['rms']:.4g} of {c['rms_allowed']:.4g} (x{c['rms'] / c['rms_allowed']:.2f})")
+        parts.append(f"PSNR {c['psnr_db']:.1f} dB" if c["psnr_db"] is not None else "identical")
+    return f"{'ok  ' if c['passed'] else 'FAIL'} ({', '.join(parts)})"
+
+
+def run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg):
+    """One workload on one rank, in step with the other ranks (see the module docstring); rank 0's result."""
+    entry = spec.WORKLOADS[w]
+    clip, device = entry["clip"], torch.device("cuda", rank)
+    timed = entry["timed"] and cfg["time"]
+    audited = entry["timed"] and cfg["audit"] and world == 1
+    calls = 1 + audited + (TRIALS if timed else 0)
+    seed = list(spec.WORKLOADS).index(w)
+    z0 = spec.make_inputs(w, device)[0]
+    zs = Pool(z0.shape, z0.dtype, device, calls, seed)
+    outs = Pool(spec.output_shape(w), torch.bfloat16, device, calls, seed + 1) if rank == 0 else None
+    allowed = spec.allowance(clip)
+    res = {"clip": clip, "world": world, "timed": entry["timed"]}
+    problems = []
+    threads = threading.active_count()
+
+    def arguments():
+        z = zs.next()
+        z.copy_(z0)
+        out = outs.next().fill_(math.nan) if outs else torch.empty(0, dtype=torch.bfloat16, device=device)
+        return z, out
+
+    def unchanged(z):
+        if not torch.equal(z, z0) and "run() wrote its input z" not in problems:
+            problems.append("run() wrote its input z")
+
+    def call():
+        z, out = arguments()
+        torch.cuda.synchronize(device)
+        dist.barrier(group=harness)
+        start = time.monotonic()
+        sol.fn(z, weights, rank, world, pg, out)
+        torch.cuda.synchronize(device)
+        dist.barrier(group=harness)
+        stop = time.monotonic()
+        unchanged(z)
+        return out, start, stop
+
+    # The warm-up, which is the call checked.
+    out, start, stop = call()
+    warm = stop - start
+    first, worst = None, 0.0
+    if rank == 0:
+        if cfg["save"]:
+            path = os.path.join(cfg["save"], f"{w}.pt")
+            torch.save(out.cpu(), path)
+            res["saved"] = path
+        elif cfg["check"]:
+            res["check"] = judge(out, spec.load_reference(clip), allowed)
+        if calls > 1:
+            first = out.clone()
+
+    def agree(out, what):
+        nonlocal worst
+        # It agrees when it loses no element the checked call wrote, and is within the allowance of it elsewhere;
+        # the checked call's own faults are its check's to report.
+        c = judge(out, first, allowed)
+        worst = max(worst, c["max_abs"] or 0.0)
+        close = c["max_abs"] is None or (c["max_abs"] <= allowed["max_abs"] and c["rms"] <= allowed["rms"])
+        if c["unmatched"] or not close:
+            lost = f"{c['unmatched']} elements not finite where it was; " if c["unmatched"] else ""
+            problems.append(f"{what}'s output does not agree with the checked call's: {lost}"
+                            f"{verdict(c).removeprefix('FAIL').removeprefix('ok').strip()}")
+
+    # World 1 alone: one call under CUPTI, which attributes every kernel.
+    if audited:
+        z, out = arguments()
+        ops = audit.trace(sol.fn, z, weights, rank, world, pg, out)
+        unchanged(z)
+        agree(out, "the audited call")
+        res |= {"kernels": audit.kernels(ops), "launches": sum(1 for op in ops if op.kind == "kernel"),
+                "lib": sum(1 for op in ops if audit.placeholder(op.library)), "lib_share": audit.library_share(ops),
+                "comm": sum(1 for op in ops if op.library in audit.COMMUNICATION), "comm_share": audit.comm_share(ops)}
+        del ops
+
+    if timed:
+        n = [TRIALS if warm < LONG_CALL_S else LONG_TRIALS]
+        dist.broadcast_object_list(n, src=0, group=harness)
+        questions, trials, windows = [], [], []
+        for k in range(n[0]):
+            if rank == 0:
+                questions.append(gpu_holders())
+            out, start, stop = call()
+            trials.append(1e3 * (stop - start))
+            windows.append((start, stop))
+            if rank == 0:
+                agree(out, f"trial {k + 1}")
+        if rank == 0:
+            questions.append(gpu_holders())
+            res |= {"trials": trials, "ms": statistics.mean(trials), "min_ms": min(trials), "windows": windows}
+            held = others(questions, [uuid for _, uuid in ranks], {pid for pid, _ in ranks} | {os.getppid()})
+            if held:
+                res["shared"] = held
+    if first is not None:
+        res["agree_max_abs"] = worst
+    if fingerprint(weights) != marks:
+        problems.append("run() wrote its weights (what prepare returned is read-only to run: keep scratch elsewhere)")
+    if threading.active_count() > threads:
+        problems.append(f"the solution started threads ({threading.active_count() - threads} more)")
+
+    # What every rank found, so that rank 0 reports it.
+    found = [None] * world
+    dist.all_gather_object(found, problems, group=harness)
+    text = [p if world == 1 else f"rank {r}: {p}" for r, ps in enumerate(found) for p in ps]
+    if text and rank == 0:
+        res["error"] = "; ".join(text)
+    return res
+
+
+def worker(rank, world, port, cfg, conn):
+    """The process of one rank (see A run in the module docstring): it reports to the process that spawned it through
+    conn, ("ready", rank, uuid, prepare seconds), ("start", workload) and ("result", workload, result) on rank 0,
+    ("done", rank), or ("error", rank, workload or None, phase, traceback), after which it exits at once."""
+    phase, current = "start", None
+    try:
+        torch.cuda.set_device(rank)
+        device = torch.device("cuda", rank)
+        timeout = timedelta(seconds=WORKLOAD_SEC)
+        dist.init_process_group("nccl", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world,
+                                device_id=device, timeout=timeout)
+        harness = dist.new_group(backend="gloo", timeout=timeout)
+        pg = dist.group.WORLD
+        phase = "load"
+        sol = loader.import_built(cfg["workdir"], cfg["spec"])
+        phase = "weights"
+        weights = spec.load_weights(device)
+        phase = "prepare"
+        prepare_s = None
+        if sol.prepare is not None:
+            torch.cuda.synchronize(device)
+            start = time.monotonic()
+            got = sol.prepare(weights, rank, world, pg)
+            torch.cuda.synchronize(device)
+            prepare_s = time.monotonic() - start
+            if not isinstance(got, (list, tuple)) or not all(isinstance(t, torch.Tensor) for t in got):
+                raise TypeError(f"prepare returned {type(got).__name__}, not a list of tensors")
+            weights = list(got)
+        uuid = str(torch.cuda.get_device_properties(device).uuid).lower().removeprefix("gpu-")
+        ranks = [None] * world
+        dist.all_gather_object(ranks, (os.getpid(), uuid), group=harness)
+        conn.send(("ready", rank, uuid, prepare_s))
+        phase = "workload"
+        marks = fingerprint(weights)
+        for w in cfg["workloads"]:
+            current = w
+            if rank == 0:
+                conn.send(("start", w))
+            res = run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg)
+            if rank == 0:
+                conn.send(("result", w, res))
+            current = None
+            torch.cuda.empty_cache()
+        phase = "end"
+        dist.barrier(group=harness)
+        conn.send(("done", rank))
+        dist.destroy_process_group()
+    except BaseException:
+        try:
+            conn.send(("error", rank, current, phase, traceback.format_exc(limit=8)))
+        finally:
+            sys.stdout.flush()
+            sys.stderr.flush()
+            # At once: NCCL's teardown would wait for the other ranks, which may wait for this one.
+            os._exit(1)
+
+
+def free_port():
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def stop(procs, grace=0.0):
+    """Ends rank processes: those that do not exit within `grace` seconds get SIGTERM, then SIGKILL after
+    STOP_GRACE_SEC."""
+    deadline = time.monotonic() + grace
+    for p in procs:
+        p.join(max(0.0, deadline - time.monotonic()))
+    for p in procs:
+        if p.is_alive():
+            p.terminate()
+    deadline = time.monotonic() + STOP_GRACE_SEC
+    for p in procs:
+        p.join(max(0.1, deadline - time.monotonic()))
+        if p.is_alive():
+            p.kill()
+            p.join(10)
+
+
+def exit_text(code):
+    if code is None:
+        return "did not exit"
+    if code < 0:
+        import signal
+
+        try:
+            return f"was killed by {signal.Signals(-code).name}"
+        except ValueError:
+            return f"was killed by signal {-code}"
+    return f"exited with code {code}"
+
+
+def session(world, workloads, cfg, report, show):
+    """One set of rank processes running these workloads at one world size, in order; results go into report as they
+    arrive, and show(w) prints each. Returns the workloads it did not reach: after a failure at one workload (which gets
+    the error), those after it; [] when every one ran, or when the failure came before any (every one gets it)."""
+    ctx = torch.multiprocessing.get_context("spawn")
+    port = free_port()
+    pipes = [ctx.Pipe(duplex=False) for _ in range(world)]
+    procs = [ctx.Process(target=worker, args=(r, world, port, dict(cfg, workloads=workloads), pipes[r][1]),
+                         daemon=True) for r in range(world)]
+    for p in procs:
+        p.start()
+    for _, send in pipes:
+        send.close()
+    receivers = {pipes[r][0]: r for r in range(world)}
+    done, errors, finished, uuids = set(), [], [], {}
+    current, clock, failure = None, None, None
+    deadline = time.monotonic() + SESSION_START_SEC
+    started = time.monotonic()
+
+    def handle(msg):
+        nonlocal current, clock, deadline
+        kind = msg[0]
+        if kind == "ready":
+            _, rank, uuid, prepare_s = msg
+            uuids[rank] = uuid
+            if rank == 0:
+                report.setdefault("prepare_s", {})[str(world)] = prepare_s
+            if len(uuids) == world:
+                clock = Clock([uuids[r] for r in range(world)])
+                ready = f", prepare {report['prepare_s'][str(world)]:.1f} s" if prepare_s is not None else ""
+                print(f"world {world}: {world} rank{'s' if world > 1 else ''} up in {time.monotonic() - started:.0f} s"
+                      f"{ready}", flush=True)
+        elif kind == "start":
+            current = msg[1]
+            deadline = time.monotonic() + WORKLOAD_SEC
+        elif kind == "result":
+            _, w, res = msg
+            windows = res.pop("windows", None)
+            means = clock.means(windows) if clock and windows else []
+            if means and means[0]:
+                res["clock_mhz"] = means[0]
+            if world > 1 and means and all(means):
+                res["world_clock_mhz"] = [min(means), max(means)]
+            report["workloads"][w] = res
+            finished.append(w)
+            current = None
+            deadline = time.monotonic() + WORKLOAD_SEC
+            show(w)
+        elif kind == "error":
+            errors.append(msg[1:])
+        elif kind == "done":
+            done.add(msg[1])
+
+    try:
+        while len(done) < world:
+            # A rank that exited cleanly may still have its last messages in its pipe: it fails once that is read.
+            bad = [(r, p.exitcode) for r, p in enumerate(procs) if r not in done and p.exitcode is not None
+                   and (p.exitcode != 0 or pipes[r][0] not in receivers)]
+            if bad:
+                failure = "; ".join(f"rank {r} {exit_text(code)}" for r, code in bad)
+                break
+            left = deadline - time.monotonic()
+            if left <= 0:
+                failure = (f"timed out after {WORKLOAD_SEC} s" if current else
+                           f"the ranks did not start within {SESSION_START_SEC} s")
+                break
+            waits = list(receivers) + [p.sentinel for p in procs if p.exitcode is None]
+            for obj in connection.wait(waits, timeout=min(left, 5.0)):
+                if obj in receivers:
+                    try:
+                        handle(obj.recv())
+                    except (EOFError, OSError):
+                        del receivers[obj]
+        # What the other ranks still say, such as the error that made one exit.
+        drain = time.monotonic() + 2.0
+        while failure and receivers and time.monotonic() < drain:
+            for obj in connection.wait(list(receivers), timeout=max(0.0, drain - time.monotonic())):
+                try:
+                    handle(obj.recv())
+                except (EOFError, OSError):
+                    del receivers[obj]
+    finally:
+        stop(procs, 0.0 if failure else STOP_GRACE_SEC)
+        if clock:
+            clock.stop()
+    if not failure:
+        return []
+
+    # The first error a rank reported says what failed; else how a rank ended.
+    if errors:
+        rank, at, phase, text = errors[0]
+        if phase == "load":
+            report["load_error"] = text.strip().splitlines()[-1]
+        text = (f"rank {rank}: " if world > 1 else "") + text.strip()
+    else:
+        text = failure
+    rest = [w for w in workloads if w not in finished]
+    at = current or (errors and errors[0][1])
+    if at not in rest:
+        at = rest[0] if rest and finished else None
+    if not at or report.get("load_error"):
+        for w in rest:
+            report["workloads"][w] = {"error": text, "world": world, "clip": spec.WORKLOADS[w]["clip"],
+                                      "timed": spec.WORKLOADS[w]["timed"]}
+            show(w)
+        return []
+    report["workloads"][at] = {"error": text, "world": world, "clip": spec.WORKLOADS[at]["clip"],
+                               "timed": spec.WORKLOADS[at]["timed"]}
+    show(at)
+    return rest[rest.index(at) + 1:]
+
+
+def kernel_lines(res, limit=25):
+    """The audited call's kernels by GPU time, one line each."""
+    rows = res.get("kernels", [])
+    total = sum(r[3] for r in rows) or 1
+    lines = [f"    {'[' + lib + ']' if lib else '':<9} {ns / 1e6:10.2f} ms {100 * ns / total:5.1f}%  x{n:<6} "
+             f"{name if name.startswith('(') else audit.short_name(name, 110)}" for lib, name, n, ns in rows[:limit]]
+    if len(rows) > limit:
+        rest = rows[limit:]
+        lines.append(f"    {'':<9} {sum(r[3] for r in rest) / 1e6:10.2f} ms "
+                     f"{100 * sum(r[3] for r in rest) / total:5.1f}%  {len(rest)} more kinds")
+    return lines
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--solution", default="solution", help="the solution directory (default: solution)")
+    p.add_argument("--workload", nargs="*", help="workloads to run (default: every one this job has the GPUs for)")
+    p.add_argument("--no-time", action="store_true", help="correctness only: no audit, no trials")
+    p.add_argument("--no-check", action="store_true", help="do not check the outputs (CI's timing runs)")
+    p.add_argument("--no-audit", action="store_true", help="no call under CUPTI")
+    p.add_argument("--kernels", action="store_true", help="list the audited calls' kernels by GPU time")
+    p.add_argument("--build-dir", help="write and build the solution here; reused when it holds the same build")
+    p.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds")
+    p.add_argument("--save", help="CI: save each workload's checked output here instead of checking it")
+    p.add_argument("--json", help="write the results to this file")
+    a = p.parse_args()
+    unknown = [w for w in a.workload or [] if w not in spec.WORKLOADS]
+    if unknown:
+        p.error(f"unknown workloads {unknown}; known: {list(spec.WORKLOADS)}")
+    # Absolute paths, the data's too: the ranks run in the solution's working directory.
+    solution = os.path.abspath(a.solution)
+    if os.path.isfile(solution):
+        solution = os.path.dirname(solution)
+    a.build_dir, a.save, a.json = (x and os.path.abspath(x) for x in (a.build_dir, a.save, a.json))
+    os.environ["DATA"] = os.path.abspath(spec.data_dir())
+    report = {"workloads": {}}
+
+    def finish():
+        if a.json:
+            with open(a.json, "w") as f:
+                json.dump(report, f)
+
+    rules = audit.rule_table(getattr(spec, "RULES", None))
+    report["rules"] = audit.check(solution, rules)
+    for v in report["rules"]:
+        print(f"rule violated: {v}")
+    print(f"stack rules: {'ok' if not report['rules'] else 'FAIL'} ({len(rules)} rules)")
+
+    # The workloads: those named, or every one this job has the GPUs for.
+    gpus = torch.cuda.device_count()
+    chosen = a.workload or [w for w, e in spec.WORKLOADS.items() if e["world"] <= gpus]
+    skipped = [w for w in spec.WORKLOADS if w not in chosen and not a.workload]
+    if skipped:
+        print(f"skipped {', '.join(skipped)}: this job has {gpus} GPU{'s' if gpus != 1 else ''}")
+    order = [w for w in spec.WORKLOADS if w in chosen]
+    timed = [w for w in order if spec.WORKLOADS[w]["timed"] and not a.no_time]
+    limits, dev, peak = bounds.workload_bounds(spec, timed)
+    if limits:
+        print(bounds.describe(dev, peak))
+
+    try:
+        workdir, build_spec, seconds = loader.stage(solution, a.build_dir,
+                                                    cache=None if a.no_cache else loader.cache_dir())
+    except Exception as e:
+        report["load_error"] = str(e) if isinstance(e, loader.SolutionError) else traceback.format_exc(limit=6)
+        print(f"FAILED to build {solution}:\n{report['load_error']}", flush=True)
+        finish()
+        sys.exit(1)
+    report |= {"language": build_spec["languages"][0], "build_s": seconds}
+    built = f"built in {seconds:.1f} s" if seconds is not None else "a build reused"
+    print(f"{build_spec['entry_point']} ({', '.join(build_spec['languages'])}): {built}", flush=True)
+
+    width = max(len(w) for w in spec.WORKLOADS)
+
+    def show(w):
+        res = report["workloads"][w]
+        line = f"{w:<{width}}"
+        if "ms" in res:
+            at = f", SM clock {res['clock_mhz']:.0f} MHz" if "clock_mhz" in res else ""
+            low, high = res.get("world_clock_mhz", (None, None))
+            if low and "clock_mhz" in res and high - low >= 0.02 * res["clock_mhz"]:
+                at += f", the world's GPUs {low:.0f}-{high:.0f} MHz"
+            line += f" {res['ms']:9.1f} ms (fastest trial {res['min_ms']:.1f}{at})"
+            if w in limits:
+                res |= {"bound_ms": limits[w], "gap": res["ms"] / limits[w]}
+                line += f"  bound {limits[w]:.1f} ms, gap x{res['gap']:.2f}"
+        if "check" in res:
+            line += "  " + verdict(res["check"])
+        elif "saved" in res:
+            line += "  saved"
+        if not res.get("timed", True):
+            line += "  (correctness only)"
+        if "lib" in res:
+            line += f"  library kernels {res['lib']} of {res['launches']} ({100 * res['lib_share']:.0f}% of GPU time)"
+        if res.get("comm"):
+            line += f", NCCL {res['comm']} ({100 * res['comm_share']:.0f}%)"
+        if "shared" in res:
+            line += f"  SHARED: {res['shared']} other process(es) held a GPU while it was timed"
+        if "error" in res:
+            line += "  ERROR\n" + res["error"]
+        print(line, flush=True)
+        if a.kernels and "kernels" in res:
+            print("\n".join(kernel_lines(res)), flush=True)
+        finish()
+
+    cfg = {"workdir": workdir, "spec": build_spec, "time": not a.no_time, "check": not a.no_check,
+           "audit": not a.no_audit and (not a.no_time or a.kernels), "save": a.save}
+    for world in sorted({spec.WORKLOADS[w]["world"] for w in order}):
+        todo = [w for w in order if spec.WORKLOADS[w]["world"] == world]
+        if world > gpus or report.get("load_error"):
+            why = report.get("load_error") or f"needs {world} GPUs; this job has {gpus}"
+            for w in todo:
+                report["workloads"][w] = {"error": why, "world": world, "clip": spec.WORKLOADS[w]["clip"],
+                                          "timed": spec.WORKLOADS[w]["timed"]}
+                show(w)
+            continue
+        while todo:
+            todo = session(world, todo, cfg, report, show)
+
+    results = [report["workloads"][w] for w in order if w in report["workloads"]]
+    times = [report["workloads"].get(w, {}).get("ms") for w in timed]
+    if timed and all(times):
+        gaps = [report["workloads"][w]["gap"] for w in timed if "gap" in report["workloads"][w]]
+        gap = f", gap to the bound x{bounds.geomean(gaps):.2f}" if len(gaps) == len(timed) else ""
+        clocks = [report["workloads"][w].get("clock_mhz") for w in timed]
+        at = f", mean SM clock {statistics.mean(clocks):.0f} MHz" if all(clocks) else ""
+        print(f"geomean {bounds.geomean(times):.1f} ms over {len(timed)} timed workloads (mean of trials){gap}{at}")
+    shared = [w for w in order if "shared" in report["workloads"].get(w, {})]
+    if shared:
+        print(f"SHARED GPU: another process held one while {', '.join(shared)} "
+              f"{'was' if len(shared) == 1 else 'were'} timed, so those times do not compare with another job's")
+    checked = [r for r in results if "check" in r]
+    if checked:
+        ok = sum(1 for r in checked if r["check"]["passed"] and "error" not in r)
+        print(f"correctness: ok at {ok} of {len(checked)} workloads checked")
+    failed = [w for w in order if "error" in report["workloads"].get(w, {})]
+    if failed:
+        print(f"errors at {', '.join(failed)}")
+    audited = [r for r in results if "kernels" in r]
+    if audited:
+        libs = collections.Counter()
+        for r in audited:
+            for lib, _, n, _ in r["kernels"]:
+                if audit.placeholder(lib):
+                    libs[lib] += n
+        detail = ", ".join(f"{k} {v}" for k, v in libs.most_common())
+        share = statistics.mean(r["lib_share"] for r in audited)
+        print(f"library kernels: {sum(libs.values())} launches in {len(audited)} audited calls"
+              f"{' (' + detail + ')' if detail else ''}, {100 * share:.0f}% of GPU time on average")
+    finish()
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/wan22_vae_decode/tools/bounds.py

@@ -0,0 +1,156 @@
+"""The physical lower bound of a call's time, from the work it must do and what the GPUs can do per second.
+
+A task's spec.py may define
+
+    bounds(workload) -> {"flops": float, "bytes": float, "gpus": int}    optionally also "math": one of PEAKS ("f16")
+
+lower bounds for one call of run() at that workload: the floating-point operations any correct algorithm performs
+(its convolutions and matrix products) and the bytes the busiest GPU must move to or from its memory at least once,
+and the GPUs the call runs on (default 1). The bound time is then
+
+    bound = max(flops / (gpus x peak FLOP/s of the chosen math), bytes / memory bandwidth)
+
+with the peaks of the GPU it runs on at its SM clock: a call cannot be faster than the slower of the two, even with
+its FLOPs spread evenly over its GPUs and nothing else to do, so the gap, measured / bound, says how much of the way
+remains. The peaks are estimates, derived as SM count x SM clock x per-SM rate per clock (DEVICES), not measured; they
+are the same for every solution, so gaps compare. The SM clock is the one GUILD_SM_CLOCK_MHZ names, where a task's
+environment sets it, else the one DEVICES gives the GPU, else the GPU's highest SM clock: GB200 clocks are not locked,
+and read low while they idle.
+
+    python tools/bounds.py [--spec spec.py]      the bound of every workload on this GPU (or the reference device)
+"""
+
+import argparse
+import math
+import os
+import subprocess
+
+# Per compute capability: dense per-SM throughput per clock, and the memory bandwidth. Estimates from the
+# architecture's published figures; the memory clock does not follow the SM clock. `sm_clock_mhz`, where given, is the
+# clock to assume instead of the GPU's.
+DEVICES = {
+    (10, 0): {
+        "name": "Blackwell, compute capability 10.0",
+        # FLOP per SM per clock, dense (no structured sparsity), fp32 accumulation.
+        "per_sm_clk": {
+            "f16": 8192,  # FP16/BF16 tensor cores (tcgen05.mma kind::f16)
+            "tf32": 4096,  # TF32 tensor cores (kind::tf32)
+            "f32": 256,  # FP32 FMA on the CUDA cores: 128 lanes x 2
+        },
+        "bytes_per_s": 8.0e12,  # HBM3e, about 8 TB/s per GPU
+    },
+}
+PEAKS = {"f16": "16-bit tensor cores", "tf32": "TF32 tensor cores", "f32": "FP32 CUDA cores"}
+# Where no GPU is visible (on a laptop, say): a GB200 GPU at its highest SM clock.
+REFERENCE_DEVICE = {"capability": (10, 0), "sms": 152, "clock_mhz": 2062}
+
+
+def sm_clock_mhz(uuid=None):
+    """The highest SM clock of the GPU with this UUID, or the reference device's."""
+    try:
+        cmd = ["nvidia-smi", "--query-gpu=clocks.max.sm", "--format=csv,noheader,nounits"]
+        if uuid:
+            cmd.insert(1, f"--id={uuid}")
+        out = subprocess.run(cmd, capture_output=True, text=True, timeout=20).stdout.split()
+        return float(out[0]) if out else REFERENCE_DEVICE["clock_mhz"]
+    except (OSError, ValueError, subprocess.SubprocessError):
+        return REFERENCE_DEVICE["clock_mhz"]
+
+
+def assumed_clock_mhz():
+    """The SM clock a task's environment says the bound assumes (GUILD_SM_CLOCK_MHZ), or None where it says none."""
+    try:
+        value = float(os.environ.get("GUILD_SM_CLOCK_MHZ", ""))
+    except ValueError:
+        return None
+    return value if value > 0 else None
+
+
+def device():
+    """{"capability", "sms", "clock_mhz"} of CUDA device 0, or REFERENCE_DEVICE where there is none. It asks the GPU's
+    properties only, which makes no CUDA context: a process that times others must hold none."""
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            p = torch.cuda.get_device_properties(0)
+            uuid = str(getattr(p, "uuid", "") or "")
+            uuid = uuid if not uuid or uuid.startswith(("GPU-", "MIG-")) else f"GPU-{uuid}"
+            capability = (p.major, p.minor)
+            clock = assumed_clock_mhz() or DEVICES.get(capability, {}).get("sm_clock_mhz") or sm_clock_mhz(uuid)
+            return {"capability": capability, "sms": p.multi_processor_count, "clock_mhz": clock}
+    except ImportError:
+        pass
+    return dict(REFERENCE_DEVICE)
+
+
+def peaks(dev):
+    """{"f16": FLOP/s, "tf32": ..., "f32": ..., "bytes_per_s": B/s} of one GPU, or None when DEVICES lacks it."""
+    table = DEVICES.get(tuple(dev["capability"]))
+    if table is None:
+        return None
+    hz = dev["clock_mhz"] * 1e6
+    out = {k: v * dev["sms"] * hz for k, v in table["per_sm_clk"].items()}
+    out["bytes_per_s"] = table["bytes_per_s"]
+    return out
+
+
+def bound_ms(work, peak):
+    """The bound time in ms of one call doing `work` ({"flops", "bytes"[, "gpus", "math"]}) on GPUs with these peaks."""
+    math_kind = work.get("math", "f16")
+    if math_kind not in PEAKS:
+        raise ValueError(f"bounds(): math must be one of {sorted(PEAKS)}, not {math_kind!r}")
+    flops, nbytes = float(work.get("flops", 0.0)), float(work.get("bytes", 0.0))
+    gpus = int(work.get("gpus", 1))
+    return 1e3 * max(flops / (gpus * peak[math_kind]), nbytes / peak["bytes_per_s"])
+
+
+def describe(dev, peak):
+    """One line on how the bound is computed here."""
+    return (f"bound = max(FLOPs / (GPUs x {peak['f16'] / 1e15:.2f} PFLOP/s) ({PEAKS['f16']}, dense), bytes / "
+            f"{peak['bytes_per_s'] / 1e12:.1f} TB/s), estimated for {dev['sms']} SMs per GPU at "
+            f"{dev['clock_mhz']:.0f} MHz")
+
+
+def workload_bounds(spec, workloads, dev=None):
+    """{workload: bound ms} for a spec that defines bounds(), with the device and its peaks; ({}, dev, None) when the
+    spec has no bounds() or the device no peak figures."""
+    dev = dev or device()
+    peak = peaks(dev)
+    fn = getattr(spec, "bounds", None)
+    if fn is None or peak is None:
+        return {}, dev, peak
+    return {w: bound_ms(fn(w), peak) for w in workloads}, dev, peak
+
+
+def geomean(xs):
+    xs = list(xs)
+    return math.exp(sum(math.log(x) for x in xs) / len(xs)) if xs else math.nan
+
+
+def main():
+    import importlib.util
+
+    here = os.path.dirname(os.path.abspath(__file__))
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--spec", default=os.path.join(os.environ.get("CHECKS", here), "spec.py"))
+    a = p.parse_args()
+    module_spec = importlib.util.spec_from_file_location("spec", a.spec)
+    spec = importlib.util.module_from_spec(module_spec)
+    module_spec.loader.exec_module(spec)
+    workloads = [w for w, e in spec.WORKLOADS.items() if e.get("timed", True)]
+    bounds, dev, peak = workload_bounds(spec, workloads)
+    if peak is None:
+        raise SystemExit(f"no peak figures for compute capability {dev['capability']}")
+    if not bounds:
+        raise SystemExit(f"{a.spec} defines no bounds(workload)")
+    print(describe(dev, peak))
+    for w, ms in bounds.items():
+        work = spec.bounds(w)
+        print(f"{w:<10} {ms:10.1f} ms   {work['flops'] / 1e15:8.4f} PFLOP over {work.get('gpus', 1)} GPU(s)   "
+              f"{work['bytes'] / 1e9:6.2f} GB")
+    print(f"geomean {geomean(bounds.values()):.1f} ms over {len(bounds)} workloads")
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/wan22_vae_decode/tools/loader.py

@@ -0,0 +1,542 @@
+"""Builds and loads a solution exactly as it is deployed, so that what runs here runs there.
+
+A solution is a directory holding exactly the files that are deployed, plus `solution.json`, the deployment's build
+spec, with its field names and meaning:
+
+    {"languages": ["cuda_cpp"], "entry_point": "binding.cpp::run", "destination_passing_style": true,
+     "compile_options": {"cflags": [...], "cuda_cflags": [...], "ld_flags": [...]}, "dependencies": []}
+
+- `languages`: always ["cuda_cpp"].
+- `entry_point`: "<file>::<function>", the function the deployment calls, in a C++/CUDA file.
+- `destination_passing_style`: required, and true for this task: run() writes its output into `out`, its last
+  argument, and returns nothing. The calling convention must never be left to inference.
+- `compile_options` (optional): extra flags; see the build below. `dependencies` (optional): a list of names, which
+  does not change the build. `binding` (optional): null or "torch".
+Any other field is refused: the deployment would ignore it, so a misspelled field would silently fall back to its
+default. Every other file in the directory is a source: non-empty UTF-8 text at a relative path, a plain file (no
+links; `__pycache__` and `.pyc` files are ignored).
+
+The module exports the entry, called positionally as the deployment calls it (argument names play no part):
+
+    run(z, weights, rank, world, pg, out)                    every call; writes out on rank 0
+    prepare(weights, rank, world, pg) -> list[Tensor]       optional, under that name: once per process, before run
+
+(spec.py says what each argument is.) Load one solution per process.
+
+stage(solution_dir) writes the sources into a fresh working directory and compiles every .cpp/.cc/.cxx/.c/.cu file at
+its top into one extension module, in a separate process, with torch.utils.cpp_extension.load, build_directory = the
+working directory, include paths [working directory, $CUTLASS_DIR/include, $CUTLASS_DIR/tools/util/include] ($CUTLASS_DIR
+defaults to fastvideo-kernel/include/cutlass, the CUTLASS submodule), and these
+flags on top of torch's own:
+  - no compile_options: cuda_cflags [], cflags [], ld_flags ["-lcuda"];
+  - compile_options given: its lists, where a missing one defaults to cflags [], cuda_cflags ["-O3",
+    "--use_fast_math"], ld_flags ["-lcuda"];
+  - unless a cuda_cflag names an architecture (-gencode, -arch), a -gencode flag for the GPU's own architecture with
+    its arch-specific features goes first (sm_100a on compute capability 10.0), as the deployment does for its GPU.
+Files in subdirectories are not compiled, but can be included. import_built(working directory) then makes it the
+current directory, as the deployment does, imports the module and takes run (and prepare) from it; load() does both.
+The module's name is the build's, not the solution's: PYBIND11_MODULE must be given TORCH_EXTENSION_NAME, as the
+deployment builds the module under a name of its own. bench.py stages a build once and its ranks import that build.
+
+Development builds are kept on each node, so that a job does not compile again what an earlier job on its node built:
+a content-addressed cache in a directory of the node that setup/node.sh makes for the agent user (CACHE_DIR, or
+$GUILD_BUILD_CACHE), keyed by everything that determines the build: the sources, solution.json, the flags, the target
+architecture, this loader, and the toolchain (the machine, Python, torch with its CUDA, nvcc, the host compiler, and
+the variables they read). A build found there is copied into the working directory in place of compiling it; one made
+is put there with its sources and module, without what the compiler left beside them. Only the commands of a
+development copy of the harness use it (tools/ in a repo: `loader.py check` and `build`, `bench.py`), unless they are
+given --no-cache. The copy in the task's checks/, which CI, the hooks and the task's commands run, never reads or
+writes it and builds fresh: the jobs of agents write that directory, so what it holds is theirs to forge. Past a
+quarter of its filesystem, a limit that only stops a runaway, the least recently used builds are removed.
+
+    python tools/loader.py check [solution] [--build-dir DIR]   build and load it, print what was loaded
+    python tools/loader.py build [solution] [DIR]               only write and build it in DIR (default build/): with
+                                                                --arch (e.g. sm_100a), on a machine without a GPU,
+                                                                such as a CPU builder, where the module is not loaded
+"""
+
+import argparse
+import atexit
+import functools
+import hashlib
+import importlib.util
+import json
+import os
+import platform
+import re
+import shutil
+import stat
+import subprocess
+import sys
+import tempfile
+import time
+
+SPEC_FILE = "solution.json"
+# The name the build gives the module; solutions never name it themselves (see the module docstring). It is this
+# kernel's own: another kernel's loader may build an extension of a generic name (solution_ext), and in one process the
+# second extension of a name loads as the first one.
+MODULE = "wan22_vae_decode_ext"
+# The optional function a solution exports beside its entry, called once per process before any call of the entry.
+PREPARE = "prepare"
+BINDINGS = ("torch",)
+CPP_ENTRY_SUFFIXES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh")
+# The fields solution.json may hold.
+SPEC_FIELDS = ("languages", "entry_point", "dependencies", "destination_passing_style", "binding", "compile_options")
+OPTION_FIELDS = ("cflags", "cuda_cflags", "ld_flags")
+# The languages a solution may declare, and the files compiled.
+LANGUAGES = ("cuda_cpp",)
+CPP_SOURCES = (".cu", ".cpp", ".cc", ".cxx", ".c")
+DEFAULT_OPTIONS = {"cflags": [], "cuda_cflags": ["-O3", "--use_fast_math"], "ld_flags": ["-lcuda"]}
+BUILD_TIMEOUT_SEC = 1800
+# The first argument of a PYBIND11_MODULE in a source.
+MODULE_NAME = re.compile(r"\bPYBIND11_MODULE\s*\(\s*([A-Za-z_]\w*)")
+HERE = os.path.dirname(os.path.abspath(__file__))
+# The node's cache of development builds (see the module docstring), and what each build there holds beside its sources.
+CACHE_DIR = "/scratch/kdev-builds"
+STAMP, SIZE = ".stamp", ".size"
+# The share of its filesystem past which the cache removes its least recently used builds: a limit that stops a runaway.
+CACHE_SHARE = 0.25
+# What a job that stopped while it stored a build left in the cache is removed after this long.
+STALE_SEC = 3600
+# The variables the compilers read, which change what they make.
+BUILD_VARIABLES = ("CC", "CXX", "CUDA_HOME", "CUDA_PATH", "CUTLASS_DIR", "NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS",
+                   "CPATH", "C_INCLUDE_PATH", "CPLUS_INCLUDE_PATH", "LIBRARY_PATH")
+
+
+class SolutionError(Exception):
+    """The solution is malformed, or fails to build or load."""
+
+
+def _nonempty_str(x):
+    return isinstance(x, str) and len(x) >= 1
+
+
+def _str_list(x):
+    return isinstance(x, list) and all(isinstance(s, str) for s in x)
+
+
+def spec_problems(spec):
+    """Problems of a build spec, as a list of strings (empty: valid); see the module docstring.
+
+    Stricter than the deployment only where it would coerce a value (a "true" string as a boolean) or fall back to a
+    default (an unknown field, a missing destination_passing_style)."""
+    if not isinstance(spec, dict):
+        return [f"{SPEC_FILE} must hold an object"]
+    out = []
+    unknown = sorted(set(spec) - set(SPEC_FIELDS))
+    if unknown:
+        out.append(f"unknown fields {unknown} (known: {list(SPEC_FIELDS)})")
+    langs = spec.get("languages")
+    if not isinstance(langs, list) or not langs or any(lang not in LANGUAGES for lang in langs):
+        out.append(f"languages: required, {list(LANGUAGES)}")
+    entry = spec.get("entry_point")
+    if not _nonempty_str(entry):
+        out.append("entry_point: required, a non-empty string")
+    elif entry.count("::") != 1:
+        out.append(f'entry_point: {entry!r} is not "<file_path>::<function_name>"')
+    elif os.path.splitext(entry.split("::")[0])[1] not in CPP_ENTRY_SUFFIXES:
+        out.append(f"entry_point: needs a C++/CUDA entry file, not {entry.split('::')[0]!r}")
+    deps = spec.get("dependencies", [])
+    if not isinstance(deps, list) or not all(_nonempty_str(d) for d in deps):
+        out.append("dependencies: a list of non-empty strings")
+    if "destination_passing_style" not in spec:
+        out.append("destination_passing_style: must be given, and true here (run writes out, its last argument)")
+    elif spec["destination_passing_style"] is not True:
+        out.append("destination_passing_style: true: this task's run() writes the video into out, its last argument, "
+                   "and returns nothing")
+    if spec.get("binding") not in (None,) + BINDINGS:
+        out.append(f"binding: null or one of {list(BINDINGS)}")
+    options = spec.get("compile_options")
+    if options is not None:
+        if not isinstance(options, dict):
+            out.append("compile_options: null or an object")
+        else:
+            for k in OPTION_FIELDS:
+                if k in options and not _str_list(options[k]):
+                    out.append(f"compile_options.{k}: a list of strings")
+    return out
+
+
+def read_spec(solution_dir):
+    """solution.json, checked against the build spec's rules (see the module docstring)."""
+    path = os.path.join(solution_dir, SPEC_FILE)
+    try:
+        with open(path) as f:
+            spec = json.load(f)
+    except (OSError, ValueError) as e:
+        raise SolutionError(f"cannot read {path}: {e}") from e
+    problems = spec_problems(spec)
+    if problems:
+        raise SolutionError(f"{SPEC_FILE}: " + "; ".join(problems))
+    return spec
+
+
+def sources(solution_dir):
+    """{relative path: text} of every file of the solution except solution.json."""
+    out = {}
+    for root, dirs, files in os.walk(solution_dir):
+        dirs[:] = sorted(d for d in dirs if d != "__pycache__")
+        for name in sorted(files):
+            path = os.path.join(root, name)
+            rel = os.path.relpath(path, solution_dir).replace(os.sep, "/")
+            if rel == SPEC_FILE or name.endswith(".pyc"):
+                continue
+            if os.path.islink(path):
+                raise SolutionError(f"{rel}: sources must be plain files, not links")
+            try:
+                with open(path, encoding="utf-8") as f:
+                    text = f.read()
+            except (OSError, UnicodeDecodeError) as e:
+                raise SolutionError(f"{rel}: sources must be UTF-8 text ({e})") from e
+            if not text:
+                raise SolutionError(f"{rel}: sources must not be empty")
+            out[rel] = text
+    if not out:
+        raise SolutionError(f"{solution_dir} holds no sources")
+    named = sorted({m for text in out.values() for m in MODULE_NAME.findall(text) if m != "TORCH_EXTENSION_NAME"})
+    if named:
+        raise SolutionError(f"PYBIND11_MODULE names its module {', '.join(named)}: give it TORCH_EXTENSION_NAME, the "
+                            f"name the build chooses (the deployment builds the module under a name of its own)")
+    return out
+
+
+def read(solution_dir):
+    """(spec, {relative path: text}) of a solution, checked as stage() checks it before building: solution.json against
+    the build spec's rules, the sources as the deployment takes them, the entry file among them."""
+    spec = read_spec(solution_dir)
+    files = sources(solution_dir)
+    entry = spec["entry_point"].split("::")[0]
+    if entry not in files:
+        raise SolutionError(f"entry file {entry} is not among the sources")
+    return spec, files
+
+
+def gpu_arch():
+    """The architecture solutions are built for, e.g. sm_100a: the GPU's own, with its arch-specific features. It asks
+    the GPU's properties only, which makes no CUDA context."""
+    import torch
+
+    if not torch.cuda.is_available():
+        raise SolutionError("there is no GPU here to take the architecture from: name it (--arch sm_100a, or the "
+                            "architecture of the GPUs the build is for)")
+    major, minor = torch.cuda.get_device_capability()
+    return f"sm_{major}{minor}" + ("a" if major >= 9 else "")
+
+
+def compile_flags(spec, arch):
+    """(cflags, cuda_cflags, ld_flags) added to torch's own flags (see the module docstring)."""
+    given = spec.get("compile_options")
+    options = {"cflags": [], "cuda_cflags": [], "ld_flags": ["-lcuda"]} if given is None else {
+        k: list(given.get(k, v)) for k, v in DEFAULT_OPTIONS.items()}
+    if not any("-gencode" in f or "-arch" in f for f in options["cuda_cflags"]):
+        options["cuda_cflags"].insert(0, f"-gencode=arch=compute_{arch[3:]},code={arch}")
+    return options["cflags"], options["cuda_cflags"], options["ld_flags"]
+
+
+def cache_dir():
+    """The node's cache of development builds, or None where a build is made fresh: in the copy of the harness in the
+    task's checks/ (CI, hooks and the task's commands run it, which must not load what a job could have put there), and
+    on a node without a cache that this user may write."""
+    checks = os.environ.get("CHECKS")
+    if checks and os.path.realpath(checks) == os.path.realpath(HERE):
+        return None
+    cache = os.environ.get("GUILD_BUILD_CACHE", CACHE_DIR)
+    return cache if os.path.isdir(cache) and os.access(cache, os.W_OK | os.X_OK) else None
+
+
+@functools.lru_cache(maxsize=None)
+def toolchain():
+    """What determines a build beside the solution, its flags and the architecture: this loader, the machine, Python,
+    torch with its CUDA, nvcc, the host compiler, and the variables the compilers read."""
+    import torch
+    from torch.utils import cpp_extension
+
+    def version(cmd):
+        try:
+            p = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
+            return (p.stdout + p.stderr).strip()
+        except (OSError, subprocess.TimeoutExpired) as e:
+            return f"unavailable: {e}"
+
+    with open(os.path.abspath(__file__), "rb") as f:
+        loader = hashlib.sha256(f.read()).hexdigest()
+    return {"loader": loader, "machine": platform.machine(), "python": sys.version,
+            "torch": [torch.__version__, torch.version.cuda, os.path.dirname(torch.__file__)],
+            "nvcc": version([os.path.join(cpp_extension.CUDA_HOME or "/usr/local/cuda", "bin", "nvcc"), "--version"]),
+            "cxx": version([os.environ.get("CXX", "c++"), "--version"]),
+            "variables": {k: os.environ[k] for k in BUILD_VARIABLES if k in os.environ}}
+
+
+def cache_key(spec, files, flags, arch, tools):
+    """The key of a build in the cache: a hash of everything that determines it (see the module docstring)."""
+    what = {"spec": spec, "files": sorted(files.items()), "flags": flags, "arch": arch, "toolchain": tools}
+    return hashlib.sha256(json.dumps(what, sort_keys=True).encode()).hexdigest()
+
+
+def disk_usage(path):
+    """The bytes a directory takes on disk."""
+    total = 0
+    for root, _, names in os.walk(path):
+        for name in names:
+            try:
+                st = os.lstat(os.path.join(root, name))
+            except OSError:
+                continue
+            total += getattr(st, "st_blocks", 0) * 512 or st.st_size
+    return total
+
+
+def fetch(cache, key, build_dir):
+    """Copies the build of a key from the cache into build_dir, in its place, and marks the build used; False when the
+    cache has none, or it went while it was copied."""
+    entry = os.path.join(cache, key)
+    if not os.path.exists(os.path.join(entry, MODULE + ".so")):
+        return False
+    shutil.rmtree(build_dir, ignore_errors=True)
+    try:
+        shutil.copytree(entry, build_dir, symlinks=True, ignore=lambda d, _: [SIZE] if d == entry else [])
+        os.utime(entry)
+    except OSError:
+        shutil.rmtree(build_dir, ignore_errors=True)
+        return False
+    return True
+
+
+def store(cache, key, build_dir, files):
+    """Puts the build in build_dir into the cache under its key: its sources, module and stamp, without what the
+    compiler left beside them. A build of the key that another job put there first stays. Then the cache is kept within
+    its limit. A cache that cannot take it (a full disk, say) is left as it is: the build stands without it."""
+    entry = os.path.join(cache, key)
+    if os.path.exists(entry):
+        return
+    try:
+        tmp = tempfile.mkdtemp(prefix=f".{key[:16]}-", dir=cache)
+    except OSError:
+        return
+    try:
+        for rel in [*files, MODULE + ".so", STAMP]:
+            dest = os.path.join(tmp, rel)
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+            shutil.copy2(os.path.join(build_dir, rel), dest)
+        with open(os.path.join(tmp, SIZE), "w") as f:
+            f.write(str(disk_usage(tmp)))
+        os.rename(tmp, entry)
+    except OSError:
+        shutil.rmtree(tmp, ignore_errors=True)
+        return
+    try:
+        st = os.statvfs(cache)
+        evict(cache, CACHE_SHARE * st.f_blocks * st.f_frsize)
+    except OSError:
+        pass
+
+
+def evict(cache, limit):
+    """Removes the least recently used builds of the cache until it holds at most `limit` bytes, and what jobs that
+    stopped while they stored a build left there (after STALE_SEC). Anything but a directory is left alone."""
+    now = time.time()
+    builds = []
+    for name in os.listdir(cache):
+        path = os.path.join(cache, name)
+        try:
+            st = os.lstat(path)
+        except OSError:
+            continue
+        if not stat.S_ISDIR(st.st_mode):
+            continue
+        if name.startswith("."):
+            if now - st.st_mtime > STALE_SEC:
+                shutil.rmtree(path, ignore_errors=True)
+            continue
+        try:
+            with open(os.path.join(path, SIZE)) as f:
+                size = int(f.read())
+        except (OSError, ValueError):
+            size = disk_usage(path)
+        builds.append((st.st_mtime, path, size))
+    total = sum(size for _, _, size in builds)
+    for _, path, size in sorted(builds):
+        if total <= limit:
+            break
+        shutil.rmtree(path, ignore_errors=True)
+        total -= size
+
+
+def stage(solution_dir, build_dir=None, arch=None, extra_cuda_cflags=(), cache=None):
+    """Writes the sources into build_dir (a fresh temporary directory if None) and builds them there.
+
+    Returns (workdir, spec, build seconds or None when an earlier build was reused). A build_dir that already holds
+    a build of the same sources and flags is reused, so CI builds each solution once for its rounds. With a cache
+    (cache_dir(): a development build), a build of the same key is copied from it, and one made is put there."""
+    spec, files = read(solution_dir)
+    arch = arch or gpu_arch()
+    flags = compile_flags(spec, arch)
+    # Flags a tool adds that do not change the code, such as -lineinfo for source-level profiles.
+    flags[1].extend(extra_cuda_cflags)
+    digest = hashlib.sha256(json.dumps([spec, flags, sorted(files.items())]).encode()).hexdigest()
+    if build_dir is None:
+        build_dir = tempfile.mkdtemp(prefix="solution-")
+        atexit.register(shutil.rmtree, build_dir, True)
+    build_dir = os.path.abspath(build_dir)
+    stamp = os.path.join(build_dir, STAMP)
+    try:
+        with open(stamp) as f:
+            if f.read() == digest and os.path.exists(os.path.join(build_dir, MODULE + ".so")):
+                return build_dir, spec, None
+    except OSError:
+        pass
+    key = cache and cache_key(spec, files, flags, arch, toolchain())
+    if key and fetch(cache, key, build_dir):
+        return build_dir, spec, None
+    shutil.rmtree(build_dir, ignore_errors=True)
+    os.makedirs(build_dir)
+    for rel, text in files.items():
+        dest = os.path.join(build_dir, rel)
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        with open(dest, "w", encoding="utf-8") as f:
+            f.write(text)
+    start = time.monotonic()
+    env = dict(os.environ, PYTORCH_ALLOC_CONF="expandable_segments:True")
+    cmd = [sys.executable, "-s", os.path.abspath(__file__), "_build", json.dumps(flags)]
+    try:
+        p = subprocess.run(cmd, cwd=build_dir, env=env, capture_output=True, text=True, timeout=BUILD_TIMEOUT_SEC)
+    except subprocess.TimeoutExpired as e:
+        raise SolutionError(f"the build timed out after {BUILD_TIMEOUT_SEC} s") from e
+    seconds = time.monotonic() - start
+    if p.returncode != 0 or not os.path.exists(os.path.join(build_dir, MODULE + ".so")):
+        log = (p.stdout + p.stderr).strip()
+        raise SolutionError(f"the build failed after {seconds:.0f} s:\n{_errors(log)}")
+    with open(stamp, "w") as f:
+        f.write(digest)
+    if key:
+        store(cache, key, build_dir, files)
+    return build_dir, spec, seconds
+
+
+def _errors(log, limit=6000):
+    """The compiler's error lines with some context, else the log's tail."""
+    lines = log.splitlines()
+    keep = [i for i, line in enumerate(lines) if "error" in line.lower()]
+    if not keep:
+        return log[-limit:]
+    picked = sorted({j for i in keep for j in range(max(0, i - 2), min(len(lines), i + 3))})
+    return "\n".join(lines[j] for j in picked)[-limit:]
+
+
+# fastvideo-kernel's CUTLASS submodule (git submodule update --init fastvideo-kernel/include/cutlass).
+_CUTLASS_SUBMODULE = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "include",
+                                  "cutlass")
+
+
+def _build(flags):
+    """The compile step, run in the working directory by stage() in its own process."""
+    import torch.utils.cpp_extension as ext
+
+    cflags, cuda_cflags, ld_flags = json.loads(flags)
+    here = os.getcwd()
+    srcs = sorted(os.path.join(here, p) for p in os.listdir(here)
+                  if p.endswith(CPP_SOURCES) and os.path.isfile(os.path.join(here, p)))
+    if not srcs:
+        raise RuntimeError("No CUDA/C++ source files found in working directory")
+    cutlass = os.environ.get("CUTLASS_DIR") or _CUTLASS_SUBMODULE
+    try:
+        ext.load(name=MODULE, sources=srcs, extra_cuda_cflags=cuda_cflags, extra_cflags=cflags, extra_ldflags=ld_flags,
+                 extra_include_paths=[here, f"{cutlass}/include", f"{cutlass}/tools/util/include"],
+                 build_directory=here, verbose=True)
+    except ImportError as e:
+        # Built, and not loadable here: a machine without a GPU has no driver to load it against, and whoever loads it
+        # says whether it loads where it runs. Where a GPU is, a module that does not load (an undefined symbol, a
+        # missing library) is a build that failed, and is neither stamped nor kept in the cache.
+        import torch
+
+        if torch.cuda.is_available():
+            raise
+        print(f"built; not loaded on this machine: {e}", flush=True)
+    built = [p for p in os.listdir(here) if p.startswith(MODULE) and p.endswith(".so") and p != MODULE + ".so"]
+    if built:
+        os.rename(built[0], MODULE + ".so")
+
+
+class Solution:
+    """A loaded solution: fn, the entry (run), called as the deployment calls it, fn(*inputs, *outputs); prepare, the
+    module's prepare function or None; spec (solution.json), entry, languages, language, module, workdir, and
+    build_seconds (None when an earlier build was reused, or when the build was staged elsewhere)."""
+
+    def __init__(self, spec, fn, prepare, module, workdir, build_seconds):
+        self.spec = spec
+        self.languages = list(spec["languages"])
+        self.language = self.languages[0]
+        self.entry = spec["entry_point"]
+        self.fn, self.prepare, self.module = fn, prepare, module
+        self.workdir, self.build_seconds = workdir, build_seconds
+
+
+def import_built(workdir, spec, chdir=True, build_seconds=None):
+    """Imports the module that stage() built in workdir and takes its entry and prepare from it. chdir: make the
+    working directory the current one, as in deployment."""
+    import torch  # noqa: F401  (the module needs libtorch loaded first)
+
+    if chdir:
+        os.chdir(workdir)
+    file, func = spec["entry_point"].split("::")
+    module_spec = importlib.util.spec_from_file_location(MODULE, os.path.join(workdir, MODULE + ".so"))
+    module = importlib.util.module_from_spec(module_spec)
+    try:
+        module_spec.loader.exec_module(module)
+    except Exception as e:
+        raise SolutionError(f"the built module does not load: {e}") from e
+    fn = getattr(module, func, None)
+    if not callable(fn):
+        raise SolutionError(f"{file} defines no function {func} in its PYBIND11_MODULE")
+    prepare = getattr(module, PREPARE, None)
+    if prepare is not None and not callable(prepare):
+        raise SolutionError(f"the module's {PREPARE} is not a function")
+    return Solution(spec, fn, prepare, module, workdir, build_seconds)
+
+
+def load(solution_dir, build_dir=None, chdir=True, extra_cuda_cflags=(), cache=None):
+    """Builds and loads a solution as it is deployed; see the module docstring.
+
+    build_dir: where to write and build it (a fresh temporary directory if None; reused if it holds the same build).
+    chdir: make the working directory the current one, as in deployment. extra_cuda_cflags: flags a tool adds to the
+    deployed ones, that must not change the code (e.g. -lineinfo). cache: the node's cache of development builds
+    (cache_dir()), or None to build fresh."""
+    workdir, spec, seconds = stage(solution_dir, build_dir, extra_cuda_cflags=extra_cuda_cflags, cache=cache)
+    return import_built(workdir, spec, chdir=chdir, build_seconds=seconds)
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    c = sub.add_parser("check", help="build and load a solution")
+    c.add_argument("solution", nargs="?", default="solution")
+    c.add_argument("--build-dir")
+    d = sub.add_parser("build", help="write and build a solution in a directory, without loading it")
+    d.add_argument("solution", nargs="?", default="solution")
+    d.add_argument("build_dir", nargs="?", default="build")
+    d.add_argument("--arch", help="the architecture to build for, e.g. sm_100a (default: this machine's GPU's)")
+    for q in (c, d):
+        q.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds")
+    b = sub.add_parser("_build")
+    b.add_argument("flags")
+    a = p.parse_args()
+    if a.cmd == "_build":
+        _build(a.flags)
+        return
+    cache = None if a.no_cache else cache_dir()
+    try:
+        if a.cmd == "build":
+            where, _, seconds = stage(os.path.abspath(a.solution), os.path.abspath(a.build_dir), a.arch, cache=cache)
+            module = os.path.join(where, MODULE + ".so")
+            print(f"built in {seconds:.1f} s: {module}" if seconds is not None else f"reused a build: {module}")
+            return
+        sol = load(os.path.abspath(a.solution), a.build_dir and os.path.abspath(a.build_dir), cache=cache)
+    except SolutionError as e:
+        sys.exit(f"FAILED: {e}")
+    built = f", built in {sol.build_seconds:.1f} s" if sol.build_seconds is not None else ", a build reused"
+    prepare = f", with {PREPARE}" if sol.prepare is not None else f", without {PREPARE}"
+    print(f"loaded {sol.entry} ({', '.join(sol.languages)}){prepare}{built}")
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/wan22_vae_decode/tools/spec.py

@@ -0,0 +1,428 @@
+"""The task: decoding Wan2.2 VAE latents into video on one, two or four GB200 GPUs of one node, as bench.py, ci.py and
+bounds.py read it. It defines the clips and workloads, the weights' order, the data, the error allowance and the work
+of one decode.
+
+What is computed
+----------------
+reference.py (in the kernel repo) defines the result: diffusers' AutoencoderKLWan._decode of the Wan2.2-TI2V-5B VAE,
+latents z [1, 48, T, 44, 78] to a video [1, 3, 4T-3, 704, 1248] in [-1, 1]: post_quant_conv, then the decoder one
+latent frame at a time with its causal caches (the first latent frame yields one frame, every later one four),
+unpatchify 2 x 2, clamp. The judge compares a solution with the fp32 decode of the same latents, cached as videos in
+$DATA (see Correctness).
+
+The interface (as loader.py builds it and bench.py calls it, positional, destination passing)
+-------------
+    run(z, weights, rank, world, pg, out)
+    prepare(weights, rank, world, pg) -> list[Tensor]       optional
+
+- z: [1, 48, T, 44, 78] bf16, contiguous, the full latents on every rank: what a pipeline passes to
+  vae.decode(latents), rounded to bf16.
+- weights: the 110 tensors of WEIGHTS below, bf16, on this rank's GPU; or what prepare returned.
+- rank, world: this process's rank and the number of ranks (1, 2 or 4), rank r on GPU r.
+- pg: torch.distributed's NCCL ProcessGroup over the ranks; a C++ entry takes it as
+  `const c10::intrusive_ptr<c10d::ProcessGroup>&` (from <torch/csrc/distributed/c10d/ProcessGroup.hpp>).
+- out: on rank 0, the video [1, 3, 4T-3, 704, 1248] bf16, contiguous, from torch.empty and not zeroed (every element
+  must be written); on the other ranks an empty tensor. When run has returned on every rank, rank 0's out holds the
+  video.
+- prepare, when the module exports it, is called once per process before any timed call, and is not timed. It sees
+  the weights, not the latents; its result replaces `weights` in every run.
+
+Numerics: bf16 weights and activations with fp32 accumulation, and fp32 statistics in the RMS norms; nothing narrower
+(no FP8, FP6 or FP4). This is a rule for the reviewers: no check reads it from the code.
+
+Workloads
+---------
+A workload is a clip decoded at a world size, named <clip>x<world>. A clip is the first T latent frames of one of
+three real 20 s clips encoded by the VAE (latents_<source>.pt in $DATA):
+
+    clip   source       latent frames T   video frames 4T-3   worlds    timed
+    f1     row_000168          1                   1           1, 4      no
+    f5     row_000168          2                   5           1, 4      no
+    f17    row_000168          5                  17           1, 4      no
+    f121   row_000168         31                 121           1, 2, 4   yes
+    f241   row_000129         61                 241           1, 2, 4   yes
+    f481   row_000090        121                 481           1, 2, 4   yes
+
+The nine timed workloads (f121x1 ... f481x4) make the objective, their geomean; the six short ones (f1x1 ... f17x4)
+are checked for correctness and never timed: a decoder sees every clip length, and the first latent frame is a
+special case.
+
+Correctness
+-----------
+On rank 0's out, against ref_<clip>.pt, the fp32 decode of the clip's latents (fp32 latents, TF32 off), stored as fp16:
+every element finite and within [-1, 1]; and with d = out - ref in fp32,
+
+    max |d|  <= min(MAX_ABS_CAP, ERROR_FACTOR x bf16_max_abs)
+    rms(d)   <= ERROR_FACTOR x bf16_rms
+
+bf16_max_abs and bf16_rms being the error of the release's bf16 decode (bf16 weights under bf16 autocast, the
+latents rounded to bf16) against that fp32 decode on the same clip, from calibration.json. PSNR, 10 log10(4 / mse), is
+shown beside them.
+
+The weights
+-----------
+run() gets the decoder's parameters and post_quant_conv's, 110 tensors, as a list in diffusers' registration order
+(that of AutoencoderKLWan.named_parameters(), post_quant_conv first, then the decoder from conv_in to conv_out;
+$DATA/weight_names.json lists the same), bf16, in PyTorch's layouts: a conv3d's weight is [out, in, kt, kh, kw], a
+conv2d's [out, in, kh, kw], a bias [out], an RMS norm's gain [C, 1, 1, 1] (the attention's [C, 1, 1]).
+$DATA/decoder_weights_bf16.pt is that list: the release's fp32 weights rounded to bf16 to nearest even, as diffusers'
+from_pretrained(torch_dtype=torch.bfloat16) casts them. The order is not the order of execution: the mid block
+registers its attention before its two residual blocks, though it runs resnets.0, the attention, then resnets.1; an
+upsampler registers its 3x3 conv (resample.1) before its time conv, though the time conv runs first. The attention's
+to_qkv gives q, k and v as its output channels 0-1023, 1024-2047 and 2048-3071. A time conv C -> 2C gives frame 2t of
+its output as channels 0..C-1 and frame 2t+1 as channels C..2C-1. conv_out's 12 channels are (colour k, column offset
+j, row offset i) of 2 x 2 pixel patches: video[k, f, 2y + i, 2x + j] = channel 4k + 2j + i at (f, y, x).
+
+      #  name                                                shape                  role
+      0  post_quant_conv.weight                              [48, 48, 1, 1, 1]      1x1x1 conv 48 -> 48, on the latents
+      1  post_quant_conv.bias                                [48]
+      2  decoder.conv_in.weight                              [1024, 48, 3, 3, 3]    causal 3x3x3 conv 48 -> 1024
+      3  decoder.conv_in.bias                                [1024]
+      4  decoder.mid_block.attentions.0.norm.gamma           [1024, 1, 1]           RMS norm gain, before to_qkv
+      5  decoder.mid_block.attentions.0.to_qkv.weight        [3072, 1024, 1, 1]     1x1 conv2d 1024 -> 3072 (q, k, v)
+      6  decoder.mid_block.attentions.0.to_qkv.bias          [3072]
+      7  decoder.mid_block.attentions.0.proj.weight          [1024, 1024, 1, 1]     1x1 conv2d 1024 -> 1024
+      8  decoder.mid_block.attentions.0.proj.bias            [1024]
+      9  decoder.mid_block.resnets.0.norm1.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv1
+     10  decoder.mid_block.resnets.0.conv1.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     11  decoder.mid_block.resnets.0.conv1.bias              [1024]
+     12  decoder.mid_block.resnets.0.norm2.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv2
+     13  decoder.mid_block.resnets.0.conv2.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     14  decoder.mid_block.resnets.0.conv2.bias              [1024]
+     15  decoder.mid_block.resnets.1.norm1.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv1
+     16  decoder.mid_block.resnets.1.conv1.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     17  decoder.mid_block.resnets.1.conv1.bias              [1024]
+     18  decoder.mid_block.resnets.1.norm2.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv2
+     19  decoder.mid_block.resnets.1.conv2.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     20  decoder.mid_block.resnets.1.conv2.bias              [1024]
+     21  decoder.up_blocks.0.resnets.0.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     22  decoder.up_blocks.0.resnets.0.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     23  decoder.up_blocks.0.resnets.0.conv1.bias            [1024]
+     24  decoder.up_blocks.0.resnets.0.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     25  decoder.up_blocks.0.resnets.0.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     26  decoder.up_blocks.0.resnets.0.conv2.bias            [1024]
+     27  decoder.up_blocks.0.resnets.1.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     28  decoder.up_blocks.0.resnets.1.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     29  decoder.up_blocks.0.resnets.1.conv1.bias            [1024]
+     30  decoder.up_blocks.0.resnets.1.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     31  decoder.up_blocks.0.resnets.1.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     32  decoder.up_blocks.0.resnets.1.conv2.bias            [1024]
+     33  decoder.up_blocks.0.resnets.2.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     34  decoder.up_blocks.0.resnets.2.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     35  decoder.up_blocks.0.resnets.2.conv1.bias            [1024]
+     36  decoder.up_blocks.0.resnets.2.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     37  decoder.up_blocks.0.resnets.2.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     38  decoder.up_blocks.0.resnets.2.conv2.bias            [1024]
+     39  decoder.up_blocks.0.upsampler.resample.1.weight     [1024, 1024, 3, 3]     3x3 conv2d 1024 -> 1024, after x2
+     40  decoder.up_blocks.0.upsampler.resample.1.bias       [1024]
+     41  decoder.up_blocks.0.upsampler.time_conv.weight      [2048, 1024, 3, 1, 1]  causal 3x1x1 conv 1024 -> 2048
+     42  decoder.up_blocks.0.upsampler.time_conv.bias        [2048]
+     43  decoder.up_blocks.1.resnets.0.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     44  decoder.up_blocks.1.resnets.0.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     45  decoder.up_blocks.1.resnets.0.conv1.bias            [1024]
+     46  decoder.up_blocks.1.resnets.0.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     47  decoder.up_blocks.1.resnets.0.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     48  decoder.up_blocks.1.resnets.0.conv2.bias            [1024]
+     49  decoder.up_blocks.1.resnets.1.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     50  decoder.up_blocks.1.resnets.1.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     51  decoder.up_blocks.1.resnets.1.conv1.bias            [1024]
+     52  decoder.up_blocks.1.resnets.1.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     53  decoder.up_blocks.1.resnets.1.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     54  decoder.up_blocks.1.resnets.1.conv2.bias            [1024]
+     55  decoder.up_blocks.1.resnets.2.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     56  decoder.up_blocks.1.resnets.2.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     57  decoder.up_blocks.1.resnets.2.conv1.bias            [1024]
+     58  decoder.up_blocks.1.resnets.2.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     59  decoder.up_blocks.1.resnets.2.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     60  decoder.up_blocks.1.resnets.2.conv2.bias            [1024]
+     61  decoder.up_blocks.1.upsampler.resample.1.weight     [1024, 1024, 3, 3]     3x3 conv2d 1024 -> 1024, after x2
+     62  decoder.up_blocks.1.upsampler.resample.1.bias       [1024]
+     63  decoder.up_blocks.1.upsampler.time_conv.weight      [2048, 1024, 3, 1, 1]  causal 3x1x1 conv 1024 -> 2048
+     64  decoder.up_blocks.1.upsampler.time_conv.bias        [2048]
+     65  decoder.up_blocks.2.resnets.0.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     66  decoder.up_blocks.2.resnets.0.conv1.weight          [512, 1024, 3, 3, 3]   causal 3x3x3 conv 1024 -> 512
+     67  decoder.up_blocks.2.resnets.0.conv1.bias            [512]
+     68  decoder.up_blocks.2.resnets.0.norm2.gamma           [512, 1, 1, 1]         RMS norm gain, before conv2
+     69  decoder.up_blocks.2.resnets.0.conv2.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     70  decoder.up_blocks.2.resnets.0.conv2.bias            [512]
+     71  decoder.up_blocks.2.resnets.0.conv_shortcut.weight  [512, 1024, 1, 1, 1]   1x1x1 conv 1024 -> 512, shortcut
+     72  decoder.up_blocks.2.resnets.0.conv_shortcut.bias    [512]
+     73  decoder.up_blocks.2.resnets.1.norm1.gamma           [512, 1, 1, 1]         RMS norm gain, before conv1
+     74  decoder.up_blocks.2.resnets.1.conv1.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     75  decoder.up_blocks.2.resnets.1.conv1.bias            [512]
+     76  decoder.up_blocks.2.resnets.1.norm2.gamma           [512, 1, 1, 1]         RMS norm gain, before conv2
+     77  decoder.up_blocks.2.resnets.1.conv2.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     78  decoder.up_blocks.2.resnets.1.conv2.bias            [512]
+     79  decoder.up_blocks.2.resnets.2.norm1.gamma           [512, 1, 1, 1]         RMS norm gain, before conv1
+     80  decoder.up_blocks.2.resnets.2.conv1.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     81  decoder.up_blocks.2.resnets.2.conv1.bias            [512]
+     82  decoder.up_blocks.2.resnets.2.norm2.gamma           [512, 1, 1, 1]         RMS norm gain, before conv2
+     83  decoder.up_blocks.2.resnets.2.conv2.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     84  decoder.up_blocks.2.resnets.2.conv2.bias            [512]
+     85  decoder.up_blocks.2.upsampler.resample.1.weight     [512, 512, 3, 3]       3x3 conv2d 512 -> 512, after x2
+     86  decoder.up_blocks.2.upsampler.resample.1.bias       [512]
+     87  decoder.up_blocks.3.resnets.0.norm1.gamma           [512, 1, 1, 1]         RMS norm gain, before conv1
+     88  decoder.up_blocks.3.resnets.0.conv1.weight          [256, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 256
+     89  decoder.up_blocks.3.resnets.0.conv1.bias            [256]
+     90  decoder.up_blocks.3.resnets.0.norm2.gamma           [256, 1, 1, 1]         RMS norm gain, before conv2
+     91  decoder.up_blocks.3.resnets.0.conv2.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+     92  decoder.up_blocks.3.resnets.0.conv2.bias            [256]
+     93  decoder.up_blocks.3.resnets.0.conv_shortcut.weight  [256, 512, 1, 1, 1]    1x1x1 conv 512 -> 256, shortcut
+     94  decoder.up_blocks.3.resnets.0.conv_shortcut.bias    [256]
+     95  decoder.up_blocks.3.resnets.1.norm1.gamma           [256, 1, 1, 1]         RMS norm gain, before conv1
+     96  decoder.up_blocks.3.resnets.1.conv1.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+     97  decoder.up_blocks.3.resnets.1.conv1.bias            [256]
+     98  decoder.up_blocks.3.resnets.1.norm2.gamma           [256, 1, 1, 1]         RMS norm gain, before conv2
+     99  decoder.up_blocks.3.resnets.1.conv2.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+    100  decoder.up_blocks.3.resnets.1.conv2.bias            [256]
+    101  decoder.up_blocks.3.resnets.2.norm1.gamma           [256, 1, 1, 1]         RMS norm gain, before conv1
+    102  decoder.up_blocks.3.resnets.2.conv1.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+    103  decoder.up_blocks.3.resnets.2.conv1.bias            [256]
+    104  decoder.up_blocks.3.resnets.2.norm2.gamma           [256, 1, 1, 1]         RMS norm gain, before conv2
+    105  decoder.up_blocks.3.resnets.2.conv2.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+    106  decoder.up_blocks.3.resnets.2.conv2.bias            [256]
+    107  decoder.norm_out.gamma                              [256, 1, 1, 1]         RMS norm gain, before conv_out
+    108  decoder.conv_out.weight                             [12, 256, 3, 3, 3]     causal 3x3x3 conv 256 -> 12
+    109  decoder.conv_out.bias                               [12]
+
+Up blocks 0 and 1 run at 1024 channels, up block 2 from 1024 to 512, up block 3 from 512 to 256; a residual block's
+first conv changes the channel count, and its 1x1x1 conv_shortcut does so on the shortcut.
+
+Data
+----
+$DATA (default /data), read-only: latents_<source>.pt ({"latents": fp32 [1, 48, 121, 44, 78]}), ref_<clip>.pt
+({"video": fp16 [1, 3, 4T-3, 704, 1248], "source", "latent_frames"}), calibration.json (per clip: bf16_max_abs,
+bf16_mean_abs, bf16_rms, bf16_psnr_db, fp32_decode_s, bf16_decode_s), weight_names.json and decoder_weights_bf16.pt.
+
+The work of one decode (bounds)
+-------------------------------
+FLOPs, a multiply-add counting 2: every convolution at its full kernel over its output positions (2 x the weight's
+elements x frames x rows x columns, zero padding included, which adds a few percent at 44 x 78 and less than one
+beyond), plus the attention's two products per latent frame (2 x 2 x 3432^2 x 1024). The frames at each level, for T
+latent frames: T at 44 x 78 (post_quant_conv, conv_in, the mid block, up block 0's residual blocks); T - 1 for up
+block 0's time conv, which skips the first latent frame; 2T - 1 at 88 x 156 (up block 0's 3x3 conv, up block 1's
+residual blocks); 2T - 2 for up block 1's time conv; 4T - 3 at 176 x 312 (up block 1's 3x3 conv, up block 2's
+residual blocks) and at 352 x 624 (up block 2's 3x3 conv, up block 3, conv_out). The FLOPs are divided over the
+workload's GPUs. Bytes: the video written (on rank 0) and the weights read once.
+"""
+
+import json
+import math
+import os
+
+LATENT_CHANNELS = 48
+LATENT_HW = (44, 78)
+VIDEO_HW = (704, 1248)
+# How far a solution may err, against what the release's bf16 decode errs by on the same clip (see Correctness).
+ERROR_FACTOR = 3.0
+MAX_ABS_CAP = 2.0
+WEIGHTS_FILE = "decoder_weights_bf16.pt"
+NAMES_FILE = "weight_names.json"
+
+CLIPS = {
+    "f1": {"source": "row_000168", "latent_frames": 1},
+    "f5": {"source": "row_000168", "latent_frames": 2},
+    "f17": {"source": "row_000168", "latent_frames": 5},
+    "f121": {"source": "row_000168", "latent_frames": 31},
+    "f241": {"source": "row_000129", "latent_frames": 61},
+    "f481": {"source": "row_000090", "latent_frames": 121},
+}
+# The short clips are checked at one GPU and at four, never timed; the long ones are checked and timed at 1, 2 and 4.
+SHORT_CLIPS, LONG_CLIPS = ("f1", "f5", "f17"), ("f121", "f241", "f481")
+WORKLOADS = {
+    **{f"{c}x{n}": {"clip": c, "world": n, "timed": False} for c in SHORT_CLIPS for n in (1, 4)},
+    **{f"{c}x{n}": {"clip": c, "world": n, "timed": True} for c in LONG_CLIPS for n in (1, 2, 4)},
+}
+
+# decoder.up_blocks.{i}: (in channels, out channels, temporal upsampling, spatial upsampling)
+UP_BLOCKS = ((1024, 1024, True, True), (1024, 1024, True, True), (1024, 512, False, True), (512, 256, False, False))
+RESNETS_PER_UP_BLOCK = 3
+
+
+def _weights():
+    """(name, shape) of run()'s weights in diffusers' registration order (see The weights)."""
+    out = []
+
+    def conv(name, shape):
+        out.extend([(f"{name}.weight", tuple(shape)), (f"{name}.bias", (shape[0],))])
+
+    def resnet(name, cin, cout):
+        out.append((f"{name}.norm1.gamma", (cin, 1, 1, 1)))
+        conv(f"{name}.conv1", (cout, cin, 3, 3, 3))
+        out.append((f"{name}.norm2.gamma", (cout, 1, 1, 1)))
+        conv(f"{name}.conv2", (cout, cout, 3, 3, 3))
+        if cin != cout:
+            conv(f"{name}.conv_shortcut", (cout, cin, 1, 1, 1))
+
+    conv("post_quant_conv", (48, 48, 1, 1, 1))
+    conv("decoder.conv_in", (1024, 48, 3, 3, 3))
+    out.append(("decoder.mid_block.attentions.0.norm.gamma", (1024, 1, 1)))
+    conv("decoder.mid_block.attentions.0.to_qkv", (3072, 1024, 1, 1))
+    conv("decoder.mid_block.attentions.0.proj", (1024, 1024, 1, 1))
+    for j in range(2):
+        resnet(f"decoder.mid_block.resnets.{j}", 1024, 1024)
+    for i, (cin, cout, temporal, spatial) in enumerate(UP_BLOCKS):
+        for j in range(RESNETS_PER_UP_BLOCK):
+            resnet(f"decoder.up_blocks.{i}.resnets.{j}", cin if j == 0 else cout, cout)
+        if spatial:
+            conv(f"decoder.up_blocks.{i}.upsampler.resample.1", (cout, cout, 3, 3))
+        if temporal:
+            conv(f"decoder.up_blocks.{i}.upsampler.time_conv", (2 * cout, cout, 3, 1, 1))
+    out.append(("decoder.norm_out.gamma", (256, 1, 1, 1)))
+    conv("decoder.conv_out", (12, 256, 3, 3, 3))
+    return tuple(out)
+
+
+WEIGHTS = _weights()
+WEIGHT_NAMES = [name for name, _ in WEIGHTS]
+
+
+def data_dir():
+    """The task's data: $DATA, else /data, where the jobs mount it."""
+    return os.environ.get("DATA") or "/data"
+
+
+def frames(latent_frames):
+    """The video frames of T latent frames: one for the first, four for every later one."""
+    return 4 * latent_frames - 3
+
+
+def output_shape(workload):
+    """The shape of rank 0's out at a workload."""
+    return (1, 3, frames(CLIPS[WORKLOADS[workload]["clip"]]["latent_frames"]), *VIDEO_HW)
+
+
+def _load(path):
+    import torch
+
+    try:
+        return torch.load(path, map_location="cpu", weights_only=True, mmap=True)
+    except (OSError, RuntimeError) as e:
+        raise ValueError(f"cannot read {path}: {e}") from e
+
+
+def make_inputs(workload, device):
+    """[z]: the clip's latents, the first T latent frames of its source, rounded to bf16, contiguous on `device`."""
+    import torch
+
+    clip = CLIPS[WORKLOADS[workload]["clip"]]
+    path = os.path.join(data_dir(), f"latents_{clip['source']}.pt")
+    latents = _load(path)["latents"]
+    t = clip["latent_frames"]
+    if latents.dim() != 5 or tuple(latents.shape[:2]) != (1, LATENT_CHANNELS) or tuple(latents.shape[3:]) != LATENT_HW \
+            or latents.shape[2] < t:
+        raise ValueError(f"{path}: latents {tuple(latents.shape)}, expected [1, 48, >= {t}, 44, 78]")
+    return [latents[:, :, :t].to(device=device, dtype=torch.bfloat16).contiguous()]
+
+
+def load_weights(device):
+    """run()'s weights on `device`: decoder_weights_bf16.pt, checked against WEIGHTS (and weight_names.json where it
+    is), so that a solution indexes what this file documents."""
+    import torch
+
+    listed = os.path.join(data_dir(), NAMES_FILE)
+    if os.path.exists(listed):
+        with open(listed) as f:
+            names = [(e["name"], tuple(e["shape"])) for e in json.load(f)]
+        if names != list(WEIGHTS):
+            raise ValueError(f"{listed} does not list spec.WEIGHTS: the weights' order or shapes changed")
+    path = os.path.join(data_dir(), WEIGHTS_FILE)
+    weights = _load(path)
+    if not isinstance(weights, list) or len(weights) != len(WEIGHTS):
+        raise ValueError(f"{path}: expected a list of {len(WEIGHTS)} tensors")
+    for i, (t, (name, shape)) in enumerate(zip(weights, WEIGHTS)):
+        if not isinstance(t, torch.Tensor) or t.dtype != torch.bfloat16 or tuple(t.shape) != shape:
+            got = f"{t.dtype} {tuple(t.shape)}" if isinstance(t, torch.Tensor) else type(t).__name__
+            raise ValueError(f"{path}: tensor {i} is {got}, expected {name}, bf16 {shape}")
+    return [t.to(device).contiguous() for t in weights]
+
+
+def load_reference(clip):
+    """The clip's fp32 decode, as stored: a CPU fp16 tensor [1, 3, 4T-3, 704, 1248] (memory-mapped)."""
+    import torch
+
+    path = os.path.join(data_dir(), f"ref_{clip}.pt")
+    ref = _load(path)
+    t, source = CLIPS[clip]["latent_frames"], CLIPS[clip]["source"]
+    video = ref.get("video") if isinstance(ref, dict) else None
+    shape = (1, 3, frames(t), *VIDEO_HW)
+    if not isinstance(video, torch.Tensor) or tuple(video.shape) != shape or video.dtype != torch.float16:
+        raise ValueError(f"{path}: expected video, fp16 {list(shape)}")
+    if ref.get("source") != source or ref.get("latent_frames") != t:
+        raise ValueError(f"{path} decodes {ref.get('source')}'s first {ref.get('latent_frames')} latent frames, not "
+                         f"{source}'s first {t}")
+    return video
+
+
+def calibration(clip):
+    """The release's bf16 decode of the clip against its fp32 decode (calibration.json's entry)."""
+    path = os.path.join(data_dir(), "calibration.json")
+    try:
+        with open(path) as f:
+            entry = json.load(f)[clip]
+    except (OSError, ValueError, KeyError) as e:
+        raise ValueError(f"{path} has no entry for {clip}: {e}") from e
+    if entry.get("source") != CLIPS[clip]["source"] or entry.get("latent_frames") != CLIPS[clip]["latent_frames"]:
+        raise ValueError(f"{path}: {clip} is calibrated on another clip ({entry.get('source')}, "
+                         f"{entry.get('latent_frames')} latent frames)")
+    return entry
+
+
+def allowance(clip):
+    """{"max_abs", "rms"}: how far a solution's video may be from the fp32 decode of the clip (see Correctness)."""
+    entry = calibration(clip)
+    return {"max_abs": min(MAX_ABS_CAP, ERROR_FACTOR * entry["bf16_max_abs"]), "rms": ERROR_FACTOR * entry["bf16_rms"]}
+
+
+def _conv_level(name, t):
+    """(output frames, upsampling factor in space) of the conv whose weight is `name`, decoding t latent frames."""
+    if name.startswith(("post_quant_conv.", "decoder.conv_in.", "decoder.mid_block.", "decoder.up_blocks.0.resnets.")):
+        return t, 1
+    if name.startswith("decoder.up_blocks.0.upsampler.time_conv."):
+        return t - 1, 1
+    if name.startswith(("decoder.up_blocks.0.upsampler.resample.", "decoder.up_blocks.1.resnets.")):
+        return 2 * t - 1, 2
+    if name.startswith("decoder.up_blocks.1.upsampler.time_conv."):
+        return 2 * t - 2, 2
+    if name.startswith(("decoder.up_blocks.1.upsampler.resample.", "decoder.up_blocks.2.resnets.")):
+        return 4 * t - 3, 4
+    return 4 * t - 3, 8
+
+
+def decode_flops(t):
+    """FLOPs of one decode of t latent frames (see The work of one decode)."""
+    h, w = LATENT_HW
+    total = 0
+    for name, shape in WEIGHTS:
+        if name.endswith(".weight") and len(shape) >= 4:
+            count, scale = _conv_level(name, t)
+            total += 2 * math.prod(shape) * count * (h * scale) * (w * scale)
+    return total + 2 * 2 * (h * w) ** 2 * 1024 * t
+
+
+def bounds(workload):
+    """{"flops", "bytes", "gpus"} of one call at the workload, for bounds.py: the decode's FLOPs (spread over the
+    world's GPUs), and the bytes that must move: the video written and the weights read."""
+    entry = WORKLOADS[workload]
+    t = CLIPS[entry["clip"]]["latent_frames"]
+    video = math.prod(output_shape(workload)) * 2
+    weights = sum(math.prod(shape) for _, shape in WEIGHTS) * 2
+    return {"flops": decode_flops(t), "bytes": video + weights, "gpus": entry["world"]}
+
+
+# The stack rules beside audit.py's: run() reads nothing but its arguments.
+RULES = {
+    "file-io": {
+        "include": r"(fstream|fcntl\.h|sys/mman\.h|filesystem|experimental/filesystem)",
+        "text": r"\b(fopen|fopen64|freopen|fdopen|fread|fwrite|ifstream|ofstream|fstream|filebuf|mmap|mmap64|shm_open|"
+                r"torch::load|torch::save|torch::pickle_load|torch::pickle_save|torch::jit::load)\b|"
+                r"::\s*(open|openat|read|pread|write|pwrite)\s*\(",
+        "why": "run() and prepare() read nothing but their arguments and write nothing but out: no files (the "
+               "reference videos lie on the job's disk) and nothing that outlives a call"},
+}
+
+
+if __name__ == "__main__":
+    for w in WORKLOADS:
+        b = bounds(w)
+        print(f"{w:<8} {b['flops'] / 1e15:8.4f} PFLOP over {b['gpus']} GPU(s), {b['bytes'] / 1e9:6.2f} GB")

fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py

@@ -0,0 +1,143 @@
+"""The decoder of the Wan2.2 VAE (Wan2.2-TI2V-5B, 48 latent channels) on sm_100a (B200 / GB200), on one or several GPUs
+of a node, built from solution/ on first use.
+
+It decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T - 3, 704, 1248] (bf16, in [-1, 1]), as
+diffusers' AutoencoderKLWan.decode(z).sample computes it (reference.py defines the result), for any T >= 1. Other
+latent sizes and batches above one are refused: the kernels are specialised to 704 x 1248.
+
+    decoder = WanVAEDecoder(vae.state_dict(), device="cuda")    # diffusers' AutoencoderKLWan, or its state dict
+    video = decoder.decode(latents)                              # [1, 3, 4T - 3, 704, 1248] bf16
+
+Over several GPUs, every rank builds a decoder with the same process group and calls decode with the full latents;
+rank 0 gets the video and the other ranks None. Without a process group, decode uses torch.distributed's default
+group, and initializes a single-rank NCCL group when there is none (one GPU only). With local=True each process decodes
+on its own GPU alone, whatever torch.distributed holds: one rank of one, so the kernels never communicate (comm.cu
+returns before any exchange below two ranks) and no group has to be created collectively.
+
+frames_per_pass sets how many latent frames each GPU carries through one pass of the decoder, which trades memory for
+speed; the video is the same at any value. On one GB200 at 8 to 20 s of video: 2 frames peak at 14-16 GiB allocated
+(x0.85 the speed of 16 frames), 4 (the default) at 21-22 GiB (x0.92), 8 at 34-35 GiB (x0.97), 16 at 60-61 GiB. The
+allocator reserves more than it allocates; PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True keeps the difference small.
+"""
+
+import atexit
+import functools
+import importlib.util
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+LATENT_CHANNELS, LATENT_HW, VIDEO_HW = 48, (44, 78), (704, 1248)
+
+
+def _own(path, name):
+    """The module at path, imported under a name of this kernel's own. Other kernels built the same way
+    keep tools and references of the same file names; a bare import would hand whichever loaded first to both."""
+    name = f"wan22_vae_decode_{name}"
+    module = sys.modules.get(name)
+    if module is None:
+        module_spec = importlib.util.spec_from_file_location(name, path)
+        module = importlib.util.module_from_spec(module_spec)
+        sys.modules[name] = module
+        try:
+            module_spec.loader.exec_module(module)
+        except BaseException:
+            del sys.modules[name]
+            raise
+    return module
+
+
+@functools.lru_cache(maxsize=None)
+def _tools():
+    """tools/loader.py and tools/spec.py."""
+    return (_own(os.path.join(HERE, "tools", "loader.py"), "loader"),
+            _own(os.path.join(HERE, "tools", "spec.py"), "spec"))
+
+
+@functools.lru_cache(maxsize=None)
+def _module():
+    """The extension built from solution/, once per process. Builds are kept in WAN22_VAE_DECODE_BUILD_DIR (default
+    ~/.cache/fastvideo/wan22_vae_decode), one directory per local rank, and reused while the sources and flags match."""
+    loader, _ = _tools()
+    base = os.environ.get("WAN22_VAE_DECODE_BUILD_DIR", os.path.expanduser("~/.cache/fastvideo/wan22_vae_decode"))
+    build_dir = os.path.join(base, f"rank{os.environ.get('LOCAL_RANK', '0')}")
+    os.makedirs(base, exist_ok=True)
+    return loader.load(os.path.join(HERE, "solution"), build_dir=build_dir, chdir=False).module
+
+
+def weight_names() -> list[str]:
+    """The 110 parameters the decoder takes, under diffusers' AutoencoderKLWan names, in the order it takes them."""
+    return list(_tools()[1].WEIGHT_NAMES)
+
+
+def decoder_weights(state_dict, device=None) -> list[Tensor]:
+    """The decoder's weights from an AutoencoderKLWan state dict (or the module itself): its post_quant_conv and
+    decoder parameters, in weight_names() order, as contiguous bf16 tensors on `device`."""
+    if isinstance(state_dict, torch.nn.Module):
+        state_dict = state_dict.state_dict()
+    _, spec = _tools()
+    weights = []
+    for name, shape in spec.WEIGHTS:
+        if name not in state_dict:
+            raise KeyError(f"the state dict has no {name}: is it the Wan2.2-TI2V-5B VAE's?")
+        w = state_dict[name]
+        if tuple(w.shape) != shape:
+            raise ValueError(f"{name} is {tuple(w.shape)}, expected {shape}")
+        weights.append(w.detach().to(device=device, dtype=torch.bfloat16).contiguous())
+    return weights
+
+
+def _destroy_if_initialized():
+    # The group may be gone already (a caller destroyed it), and destroying it twice raises.
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def _default_group():
+    if not dist.is_initialized():
+        dist.init_process_group("nccl", store=dist.HashStore(), rank=0, world_size=1)
+        atexit.register(_destroy_if_initialized)
+    return dist.group.WORLD
+
+
+class WanVAEDecoder:
+    """The decoder with its weights prepared once (the layout and the Winograd transforms the kernels read)."""
+
+    def __init__(self, state_dict, device=None, pg=None, local=False, frames_per_pass=4):
+        if frames_per_pass < 1:
+            raise ValueError(f"frames_per_pass must be at least 1, not {frames_per_pass}")
+        self.frames_per_pass = int(frames_per_pass)
+        self.device = torch.device(device if device is not None else "cuda")
+        if self.device.index is None:
+            self.device = torch.device("cuda", torch.cuda.current_device())
+        with torch.cuda.device(self.device):
+            if local:
+                if pg is not None:
+                    raise ValueError("local=True decodes on this GPU alone; it takes no process group")
+                # At world 1 the group is only carried, never used: any initialized group serves.
+                self.pg = dist.group.WORLD if dist.is_initialized() else _default_group()
+                self.rank, self.world = 0, 1
+            else:
+                self.pg = pg if pg is not None else _default_group()
+                self.rank, self.world = dist.get_rank(self.pg), dist.get_world_size(self.pg)
+            if self.world > 8:
+                raise ValueError(f"the decoder splits over at most 8 GPUs, not {self.world}")
+            self.weights = _module().prepare(decoder_weights(state_dict, self.device), self.rank, self.world,
+                                             self.pg)
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor | None:
+        """The video of latents z [1, 48, T, 44, 78] on rank 0, None on the other ranks."""
+        if z.dim() != 5 or z.shape[0] != 1 or z.shape[1] != LATENT_CHANNELS or tuple(z.shape[3:]) != LATENT_HW:
+            raise ValueError(f"latents {tuple(z.shape)}: expected [1, 48, T, 44, 78]")
+        z = z.to(device=self.device, dtype=torch.bfloat16).contiguous()
+        frames = 4 * z.shape[2] - 3
+        with torch.cuda.device(self.device):
+            out = torch.empty((1, 3, frames, *VIDEO_HW) if self.rank == 0 else (0,), dtype=torch.bfloat16,
+                              device=self.device)
+            _module().run_with(z, self.weights, self.rank, self.world, self.pg, out, self.frames_per_pass)
+        return out if self.rank == 0 else None

fastvideo/configs/models/dits/__init__.py

@@ -8,6 +8,7 @@ from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
+from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
 from fastvideo.configs.models.dits.mmaudio import (
     MMAUDIO_44K_TRAINING_VARIANTS,
     MMAUDIO_TRAINING_VARIANTS,
@@ -18,6 +19,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
@@ -31,7 +33,7 @@ __all__ = [
     "FluxDiTConfig", "Flux2Config", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig",
     "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig",
     "StableAudioConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotWorldFastVideoConfig",
-    "LingBotVideoConfig", "MiniMaxH3Config", "MMAudioArchConfig", "MMAudioTransformerConfig",
+    "LingBotVideoConfig", "MiniMaxH3Config", "PAN2VideoConfig", "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/pan2.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def is_transformer_block(n: str, m) -> bool:
+    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_refiner_block(n: str, m) -> bool:
+    return "context_refiner" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class PAN2ArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_block, is_refiner_block])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    # Diffusers PAN2Transformer3DModel keys -> native keys. The per-stream Q/K/V projections are fused on load.
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^x_embedder\.(.*)$": r"x_embedder.proj.\1",
+            r"^time_embedder\.linear_1\.(.*)$": r"time_embedder.mlp.fc_in.\1",
+            r"^time_embedder\.linear_2\.(.*)$": r"time_embedder.mlp.fc_out.\1",
+            r"^context_embedder\.0\.(.*)$": r"context_norm.\1",
+            r"^context_embedder\.1\.(.*)$": r"context_proj.\1",
+            r"^context_refiner\.(\d+)\.attn\.to_q\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 0, 3),
+            r"^context_refiner\.(\d+)\.attn\.to_k\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 1, 3),
+            r"^context_refiner\.(\d+)\.attn\.to_v\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 2, 3),
+            r"^context_refiner\.(\d+)\.attn\.to_out\.0\.(.*)$": r"context_refiner.\1.attn.to_out.\2",
+            r"^context_refiner\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"context_refiner.\1.ff.fc_in.\2",
+            r"^context_refiner\.(\d+)\.ff\.net\.2\.(.*)$": r"context_refiner.\1.ff.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
+            (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
+            (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
+            (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.(ff|ff_context)\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.\2.fc_in.\3",
+            r"^transformer_blocks\.(\d+)\.(ff|ff_context)\.net\.2\.(.*)$": r"transformer_blocks.\1.\2.fc_out.\3",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    patch_size: tuple[int, ...] = (1, 2, 2)
+    in_channels: int = 97
+    out_channels: int = 48
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    num_layers: int = 44
+    num_refiner_layers: int = 4
+    mlp_ratio: float = 8.0
+    context_mlp_ratio: float = 4.0
+    refiner_mlp_ratio: float = 4.0
+    text_embed_dim: int = 4096
+    rope_axes_dim: tuple[int, ...] = (16, 56, 56)
+    rope_theta: float = 256.0
+    # Frame positions are rescaled to the frame rate the rotary embedding was trained at (16 fps of 24 fps).
+    rope_temporal_scale: float = 16 / 24
+    # PAN2 runs its text stream in the DiT dtype.
+    cast_prompt_embeds_to_dit_dtype: bool = True
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["x_embedder", "time_embedder", "context_norm", "context_proj", "context_refiner"])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.patch_size = tuple(self.patch_size)
+        self.rope_axes_dim = tuple(self.rope_axes_dim)
+        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
+        self.num_channels_latents: int = self.out_channels
+
+
+@dataclass
+class PAN2VideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=PAN2ArchConfig)
+
+    prefix: str = "PAN2"

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,8 @@ 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.qwen3_5 import Qwen3_5TextArchConfig, Qwen3_5TextConfig
 from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (
     StableAudioConditionerArchConfig,
@@ -49,6 +51,10 @@ __all__ = [
     "StableAudioConditionerConfig",
     "T5GemmaEncoderConfig",
     "Qwen3TextConfig",
+    "QwenImage21Qwen3VLArchConfig",
+    "QwenImage21Qwen3VLConfig",
+    "Qwen3_5TextArchConfig",
+    "Qwen3_5TextConfig",
     "Mistral3TextConfig",
     "LingBotWorld2UMT5ArchConfig",
     "LingBotWorld2UMT5Config",

fastvideo/configs/models/encoders/qwen3_5.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.5 text encoder configuration (hybrid Gated DeltaNet + gated attention)."""
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+_LEGACY_LAYER_TYPES = {
+    "mamba": "linear_attention",
+    "conv": "linear_attention",
+    "attention": "full_attention",
+}
+_DEFAULT_ROPE_THETA = 10000.0
+_DEFAULT_PARTIAL_ROTARY_FACTOR = 0.25
+
+
+def _is_transformer_layer(n: str, m: Any) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m: Any) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m: Any) -> bool:
+    return n == "norm"
+
+
+@dataclass
+class Qwen3_5TextArchConfig(TextEncoderArchConfig):
+    """Architecture config for the Qwen3.5 text model (``Qwen3_5TextModel``).
+
+    Layers alternate between Gated DeltaNet linear attention and gated full
+    attention according to ``layer_types``. The defaults describe the 9B text
+    tower; checkpoint ``config.json`` values override them at load time.
+    """
+
+    vocab_size: int = 248320
+    hidden_size: int = 4096
+    intermediate_size: int = 12288
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 4
+    head_dim: int = 256
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 32768
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    linear_conv_kernel_dim: int = 4
+    linear_key_head_dim: int = 128
+    linear_value_head_dim: int = 128
+    linear_num_key_heads: int = 16
+    linear_num_value_heads: int = 32
+    layer_types: list[str] | None = None
+    full_attention_interval: int = 4
+    # RoPE settings as written by Transformers 5 (``rope_parameters``) or by
+    # older configs (top-level ``rope_theta`` / ``partial_rotary_factor`` /
+    # ``rope_scaling``). Values inside ``rope_parameters`` take precedence.
+    rope_parameters: dict[str, Any] | None = None
+    rope_scaling: dict[str, Any] | None = None
+    rope_theta: float | None = None
+    partial_rotary_factor: float | None = None
+    pad_token_id: int | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # Validate eagerly; the resolved views are recomputed on demand because
+        # ``update_model_arch`` re-runs this hook after replacing raw fields.
+        self.resolved_layer_types()
+        self.resolved_rope_parameters()
+
+    def resolved_layer_types(self) -> list[str]:
+        """Return one ``linear_attention`` / ``full_attention`` entry per layer."""
+        if self.layer_types is None:
+            interval = self.full_attention_interval
+            layer_types = [
+                "linear_attention" if (index + 1) % interval else "full_attention"
+                for index in range(self.num_hidden_layers)
+            ]
+        else:
+            layer_types = [_LEGACY_LAYER_TYPES.get(kind, kind) for kind in self.layer_types]
+        if len(layer_types) < self.num_hidden_layers:
+            raise ValueError(f"Qwen3.5 layer_types has {len(layer_types)} entries for "
+                             f"{self.num_hidden_layers} layers")
+        unknown = set(layer_types) - {"linear_attention", "full_attention"}
+        if unknown:
+            raise ValueError(f"Unsupported Qwen3.5 layer types: {sorted(unknown)}")
+        return layer_types[:self.num_hidden_layers]
+
+    def resolved_rope_parameters(self) -> dict[str, Any]:
+        """Return the effective text RoPE settings with Transformers' precedence."""
+        rope = dict(self.rope_scaling or self.rope_parameters or {})
+        rope.setdefault("rope_type", rope.get("type", "default"))
+        rope.setdefault("rope_theta", self.rope_theta if self.rope_theta is not None else _DEFAULT_ROPE_THETA)
+        rope.setdefault(
+            "partial_rotary_factor",
+            self.partial_rotary_factor if self.partial_rotary_factor is not None else _DEFAULT_PARTIAL_ROTARY_FACTOR,
+        )
+        return rope
+
+
+@dataclass
+class Qwen3_5TextConfig(TextEncoderConfig):
+    """Top-level config for the Qwen3.5 text encoder."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3_5TextArchConfig)
+    prefix: str = "qwen3_5"

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,8 @@ 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.configs.pipelines.pan2 import PAN2PipelineConfig
 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 +28,5 @@ __all__ = [
     "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
     "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
     "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
-    "get_pipeline_config_cls_from_name"
+    "QwenImage21PipelineConfig", "PAN2PipelineConfig", "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/pan2.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 pipeline configuration: text-to-video, and image-to-video from a first frame."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import PAN2VideoConfig
+from fastvideo.configs.models.encoders import Qwen3_5TextConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.models.wan.vae_config import WanVAEConfig
+
+PAN2_SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions."
+
+
+@dataclass
+class PAN2PipelineConfig(PipelineConfig):
+    """One configuration serves text-to-video and, given a first frame, image-to-video."""
+
+    dit_config: DiTConfig = field(default_factory=PAN2VideoConfig)
+    # The Wan2.2 48-channel VAE; its architecture and latent statistics come from vae/config.json at load time.
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Qwen3_5TextConfig(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "fp32"
+    # Decode the video with fastvideo-kernel's Wan2.2 decoder (B200 / GB200 / B300; a 704 x 1248 video, batch 1): bf16
+    # with fp32 accumulation, about 7x faster than the fp32 decode on a B300, at PSNR 61-63 dB and SSIM 0.998 against
+    # it. Other devices and sizes, and False, keep the fp32 decode.
+    vae_decode_kernel: bool = True
+    # With vae_decode_kernel, split the decode over the sequence-parallel GPUs (2, 4 or 8) rather than decoding the
+    # whole video on each of them. On B300s a 20 s video decodes in 3.0 s on one GPU, 1.55 s on 2 (the same video, bit
+    # for bit), 0.79 s on 4 and 0.47 s on 8 (rounded differently, at the same PSNR against the fp32 decode).
+    vae_decode_kernel_parallel: bool = True
+    # Run the DiT's fused Triton kernels (Q/K RMSNorm + RoPE, LayerNorm + modulation): 21% less denoising time on 8
+    # B300s, as far from the eager path's samples as bf16 rounding differences carry them.
+    dit_fused_kernels: bool = True
+    # Unset: the checkpoint's scheduler config decides, which is shift 9 for IFM/PAN2-TEST, the reference sampler's
+    # fixed value. A fixed default here would override the released shift with whatever a test checkpoint declares.
+    flow_shift: float | None = None
+    # Image-to-video: the first latent frame (the conditioning image) takes the positive prediction instead of the
+    # guided one, as the reference sampler does by default (``--i2v_skip_first_latent_cfg``).
+    i2v_skip_first_latent_cfg: bool = True
+    # PAN2 samples on fp32 latents; keep the Euler update out of autocast.
+    scheduler_step_in_fp32: bool = True
+
+    # Prompt encoding: PAN2 conditions on the tokens after the chat template's user header, taken from the text
+    # encoder's hidden states with its last two layers skipped (``hidden_states[-3]``).
+    text_system_prompt: str = PAN2_SYSTEM_PROMPT
+    text_encoder_skip_layers: int = 2
+    text_max_sequence_length: int = 1536
+
+    def __post_init__(self) -> None:
+        # The encoder is needed for the image-to-video first frame.
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.text_encoder_configs[0].arch_config.output_hidden_states = True

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/video_generator.py

@@ -162,6 +162,31 @@ def _quantize_video_frames_to_uint8(src: torch.Tensor) -> torch.Tensor:
         return rearrange(vid_u8, "b c t h w -> t b c h w")
 
 
+def _video_frames(src: torch.Tensor) -> list[np.ndarray]:
+    """`[b, c, t, h, w]` decoded output -> one host uint8 frame `[h, w, 3]` per time step, the batch tiled six to a
+    row by `torchvision.utils.make_grid` (which also repeats a single channel three times).
+
+    One RGB video, every request of most pipelines, is quantized and laid out frames-first and channels-last on
+    ``src``'s device, copied to the host once, and its frames are views of that one array. Building each frame on the
+    CPU instead (a channels-last copy per frame) took 0.48 s against 0.38 s for 241 frames of 704 x 1248 on a B300
+    node, and 1-2 s in a job pinned to two CPUs. Same values either way; the out-of-memory fallback of
+    `_quantize_video_frames_to_uint8` applies to both.
+    """
+    if src.ndim == 5 and src.shape[0] == 1 and src.shape[1] == 3:
+        try:
+            host = pixels_to_uint8(src)[0].permute(1, 2, 3, 0).contiguous().cpu()
+        except torch.cuda.OutOfMemoryError:
+            logger.warning("GPU out of memory laying out the decoded video's frames; doing it on the CPU instead.")
+            src_cpu = src.detach().to("cpu")
+            host = pixels_to_uint8(src_cpu if src_cpu.dtype == torch.uint8 else src_cpu.float())[0].permute(
+                1, 2, 3, 0).contiguous()
+        return list(host.numpy())
+    return [
+        torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy()
+        for x in _quantize_video_frames_to_uint8(src)
+    ]
+
+
 def _validate_request_stage_overrides(model_path: str, request: GenerationRequest) -> None:
     """Validate typed stage overrides against the model's registered preset."""
     if not request.stage_overrides:
@@ -949,7 +974,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()
@@ -1040,10 +1066,7 @@ class VideoGenerator:
             # `samples` above is just the pinned-CPU mirror of `output_batch.output` (`samples.copy_(output)` or
             # `output.cpu()`) with no intervening preprocessing, so reading `output_batch.output` here is the same
             # data. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback.
-            vid_u8 = _quantize_video_frames_to_uint8(output_batch.output)
-            frames = [
-                torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8
-            ]
+            frames = _video_frames(output_batch.output)
         postprocess_time = time.perf_counter() - postprocess_start
         logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
         if logging_info is not None:

fastvideo/envs.py

@@ -635,6 +635,12 @@ FASTVIDEO_H3_VSA_FINE_ATTENTION = EnvChoice(
     "emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized "
     "BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over "
     "FASTVIDEO_H3_VSA_FP4. Empty keeps the native route.")
+FASTVIDEO_WAN22_VAE_DECODE_DIR = EnvPath("",
+                                         category="performance",
+                                         doc="Directory of the Wan2.2 VAE decoder kernels "
+                                         "(fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds "
+                                         "on first use, when its config sets decode_kernel. Empty takes the copy in "
+                                         "this checkout.")
 FASTVIDEO_H3_VSA_HEADS_FIRST_TILE = EnvBool(False,
                                             category="performance",
                                             doc="VSA-H3 64/128-token tiles: scatter rows straight into the heads-first "
@@ -924,6 +930,10 @@ FASTVIDEO_TEST_SD35_MODEL_DIR = EnvStr("stabilityai/stable-diffusion-3.5-medium"
                                        category="test",
                                        doc="Model for the Stable Diffusion 3.5 SSIM test.",
                                        deprecated_names=("SD35_MODEL_DIR", ))
+FASTVIDEO_TEST_PAN2_MODEL_DIR = EnvPath(None,
+                                        category="test",
+                                        doc="PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU "
+                                        "tests; unset skips them.")
 FASTVIDEO_TEST_TAEH3_REFERENCE_DIR = EnvStr(None,
                                             category="test",
                                             doc="Upstream taehv checkout for the MLX TAEH3 parity test.",

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/pan2.py

@@ -0,0 +1,482 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 video transformer.
+
+A mixture-of-transformers DiT: the video and text streams keep separate weights and meet only in one joint
+attention per block, over ``[video; text]``. The text stream is modulated by a learned, timestep-free table.
+Under sequence parallelism the video tokens are sharded and the text tokens are replicated on every rank.
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.attention import DistributedAttention, LocalAttention
+from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
+from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad,
+                                                    sequence_model_parallel_shard)
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.mlp import MLP
+from fastvideo.layers.rotary_embedding import get_1d_rotary_pos_embed
+from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.pan2_fusions import layer_norm_modulate, qk_rmsnorm_rope
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _fusable_layer_norm(x: torch.Tensor) -> bool:
+    """Whether ``layer_norm_modulate`` computes what the eager ``norm(x) * (1 + scale) + shift`` does here.
+
+    The fused op rounds once from FP32, which is the eager result under autocast (FP32 LayerNorm and modulation, cast
+    by the next linear layer) and in FP32; a low-precision tensor outside autocast rounds after every eager op instead.
+    """
+    return x.dtype == torch.float32 or torch.is_autocast_enabled(x.device.type)
+
+
+class PAN2SwiGLU(nn.Module):
+    """SwiGLU MLP of the text refiner; the fused input projection stores ``[value; gate]``."""
+
+    def __init__(self, dim: int, inner_dim: int, prefix: str = ""):
+        super().__init__()
+        self.fc_in = ReplicatedLinear(dim, 2 * inner_dim, bias=False, prefix=f"{prefix}.fc_in")
+        self.fc_out = ReplicatedLinear(inner_dim, dim, bias=False, prefix=f"{prefix}.fc_out")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.fc_in(x)
+        value, gate = x.chunk(2, dim=-1)
+        x, _ = self.fc_out(value * F.silu(gate))
+        return x
+
+
+class PAN2RefinerAttention(nn.Module):
+    """Text-only self-attention of the refiner; the text is replicated on every sequence-parallel rank."""
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 eps: float,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.num_heads = num_heads
+        self.to_qkv = ReplicatedLinear(dim, 3 * num_heads * head_dim, bias=False, prefix=f"{prefix}.to_qkv")
+        self.norm_q = RMSNorm(head_dim, eps=eps)
+        self.norm_k = RMSNorm(head_dim, eps=eps)
+        self.to_out = ReplicatedLinear(num_heads * head_dim, dim, bias=False, prefix=f"{prefix}.to_out")
+        self.attn = LocalAttention(num_heads=num_heads,
+                                   head_size=head_dim,
+                                   causal=False,
+                                   supported_attention_backends=supported_attention_backends)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.unflatten(-1, (3, self.num_heads, -1)).unbind(2)
+        query = self.norm_q(query).to(value.dtype)
+        key = self.norm_k(key).to(value.dtype)
+        hidden_states = self.attn(query, key, value)
+        hidden_states, _ = self.to_out(hidden_states.flatten(2))
+        return hidden_states
+
+
+class PAN2TokenRefinerBlock(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 mlp_ratio: float,
+                 eps: float,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, eps=eps)
+        self.attn = PAN2RefinerAttention(dim,
+                                         num_heads,
+                                         head_dim,
+                                         eps,
+                                         supported_attention_backends=supported_attention_backends,
+                                         prefix=f"{prefix}.attn")
+        self.norm2 = RMSNorm(dim, eps=eps)
+        self.ff = PAN2SwiGLU(dim, int(dim * mlp_ratio), prefix=f"{prefix}.ff")
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states))
+        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
+        return hidden_states
+
+
+class PAN2AdaLayerNormZero(nn.Module):
+    """Timestep modulation of the video stream: modulated attention input plus the gate/shift/scale for the rest."""
+
+    # Set by PAN2Transformer3DModel.set_fused_kernels.
+    fused_kernels = False
+
+    def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""):
+        super().__init__()
+        self.linear = ReplicatedLinear(dim, 6 * dim, bias=True, prefix=f"{prefix}.linear")
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x: torch.Tensor,
+                temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb, _ = self.linear(F.silu(temb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
+        if self.fused_kernels and _fusable_layer_norm(x):
+            x = layer_norm_modulate(x, 1 + scale_msa, shift_msa, self.norm.eps)
+        else:
+            x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class PAN2AdaLayerNormOut(nn.Module):
+
+    # Set by PAN2Transformer3DModel.set_fused_kernels.
+    fused_kernels = False
+
+    def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""):
+        super().__init__()
+        self.linear = ReplicatedLinear(dim, 2 * dim, bias=True, prefix=f"{prefix}.linear")
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        emb, _ = self.linear(F.silu(temb).to(x.dtype))
+        scale, shift = emb.chunk(2, dim=-1)
+        if self.fused_kernels and _fusable_layer_norm(x):
+            return layer_norm_modulate(x, 1 + scale, shift, self.norm.eps)
+        return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+
+
+class PAN2Attention(nn.Module):
+    """Joint attention over ``[video; text]`` with per-stream projections and Q/K RMSNorm.
+
+    RoPE applies to the video tokens only; ``DistributedAttention`` applies it after its all-to-all, where each rank
+    holds the full (unpadded) video sequence for its heads. With fused kernels the Q/K RMSNorm and RoPE run as one op
+    before the all-to-all, on this rank's video shard with ``local_freqs_cis``.
+    """
+
+    # Set by PAN2Transformer3DModel.set_fused_kernels.
+    fused_kernels = False
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 eps: float,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.num_heads = num_heads
+        inner_dim = num_heads * head_dim
+
+        self.to_qkv = ReplicatedLinear(dim, 3 * inner_dim, bias=True, prefix=f"{prefix}.to_qkv")
+        self.norm_q = RMSNorm(head_dim, eps=eps)
+        self.norm_k = RMSNorm(head_dim, eps=eps)
+        self.to_out = ReplicatedLinear(inner_dim, dim, bias=True, prefix=f"{prefix}.to_out")
+
+        self.add_qkv_proj = ReplicatedLinear(dim, 3 * inner_dim, bias=True, prefix=f"{prefix}.add_qkv_proj")
+        self.norm_added_q = RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.to_add_out = ReplicatedLinear(inner_dim, dim, bias=True, prefix=f"{prefix}.to_add_out")
+
+        self.attn = DistributedAttention(num_heads=num_heads,
+                                         head_size=head_dim,
+                                         causal=False,
+                                         supported_attention_backends=supported_attention_backends,
+                                         prefix=f"{prefix}.attn")
+
+    def _project(self, projection: ReplicatedLinear,
+                 hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        qkv, _ = projection(hidden_states)
+        query, key, value = qkv.unflatten(-1, (3, self.num_heads, -1)).unbind(2)
+        return query, key, value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor] | None,
+        original_seq_len: int,
+        local_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        query, key, value = self._project(self.to_qkv, hidden_states)
+        encoder_query, encoder_key, encoder_value = self._project(self.add_qkv_proj, encoder_hidden_states)
+        if self.fused_kernels and local_freqs_cis is not None:
+            # RoPE is per token and per head, so it applies to this rank's shard before the all-to-all as well as to
+            # the gathered sequence after it.
+            query, key = qk_rmsnorm_rope(query, key, self.norm_q.weight, self.norm_k.weight,
+                                         self.norm_q.variance_epsilon, *local_freqs_cis)
+            encoder_query, encoder_key = qk_rmsnorm_rope(encoder_query, encoder_key, self.norm_added_q.weight,
+                                                         self.norm_added_k.weight, self.norm_added_q.variance_epsilon)
+            freqs_cis = None
+        else:
+            query = self.norm_q(query).to(value.dtype)
+            key = self.norm_k(key).to(value.dtype)
+            encoder_query = self.norm_added_q(encoder_query).to(encoder_value.dtype)
+            encoder_key = self.norm_added_k(encoder_key).to(encoder_value.dtype)
+
+        hidden_states, encoder_hidden_states = self.attn(query,
+                                                         key,
+                                                         value,
+                                                         original_seq_len,
+                                                         replicated_q=encoder_query,
+                                                         replicated_k=encoder_key,
+                                                         replicated_v=encoder_value,
+                                                         freqs_cis=freqs_cis)
+        hidden_states, _ = self.to_out(hidden_states.flatten(2))
+        encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states.flatten(2))
+        return hidden_states, encoder_hidden_states
+
+
+class PAN2TransformerBlock(nn.Module):
+
+    # Set by PAN2Transformer3DModel.set_fused_kernels.
+    fused_kernels = False
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 mlp_ratio: float,
+                 context_mlp_ratio: float,
+                 eps: float = 1e-6,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.norm1 = PAN2AdaLayerNormZero(dim, eps=eps, prefix=f"{prefix}.norm1")
+        self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim))
+
+        self.attn = PAN2Attention(dim,
+                                  num_heads,
+                                  head_dim,
+                                  eps,
+                                  supported_attention_backends=supported_attention_backends,
+                                  prefix=f"{prefix}.attn")
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff = MLP(dim, int(dim * mlp_ratio), act_type="gelu_pytorch_tanh", prefix=f"{prefix}.ff")
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.ff_context = MLP(dim,
+                              int(dim * context_mlp_ratio),
+                              act_type="gelu_pytorch_tanh",
+                              prefix=f"{prefix}.ff_context")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        original_seq_len: int,
+        local_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+        (
+            context_shift_msa,
+            context_scale_msa,
+            context_gate_msa,
+            context_shift_mlp,
+            context_scale_mlp,
+            context_gate_mlp,
+        ) = self.scale_shift_table_context.to(temb.dtype).chunk(6, dim=-1)
+        fused = self.fused_kernels and _fusable_layer_norm(encoder_hidden_states)
+        if fused:
+            norm_encoder_hidden_states = layer_norm_modulate(encoder_hidden_states, 1 + context_scale_msa,
+                                                             context_shift_msa, self.norm1_context.eps)
+        else:
+            norm_encoder_hidden_states = (self.norm1_context(encoder_hidden_states) *
+                                          (1 + context_scale_msa) + context_shift_msa)
+
+        attn_output, context_attn_output = self.attn(norm_hidden_states, norm_encoder_hidden_states, freqs_cis,
+                                                     original_seq_len, local_freqs_cis)
+        if fused:
+            # One rounding per gated residual, as PAN2's fused training kernels compute it.
+            hidden_states = torch.addcmul(hidden_states, attn_output, gate_msa.unsqueeze(1))
+            encoder_hidden_states = torch.addcmul(encoder_hidden_states, context_attn_output, context_gate_msa)
+            norm_hidden_states = layer_norm_modulate(hidden_states, 1 + scale_mlp, shift_mlp, self.norm2.eps)
+            norm_encoder_hidden_states = layer_norm_modulate(encoder_hidden_states, 1 + context_scale_mlp,
+                                                             context_shift_mlp, self.norm2_context.eps)
+            hidden_states = torch.addcmul(hidden_states, self.ff(norm_hidden_states), gate_mlp.unsqueeze(1))
+            encoder_hidden_states = torch.addcmul(encoder_hidden_states, self.ff_context(norm_encoder_hidden_states),
+                                                  context_gate_mlp)
+            return hidden_states, encoder_hidden_states
+
+        hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + context_attn_output * context_gate_msa
+
+        norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm_encoder_hidden_states = (self.norm2_context(encoder_hidden_states) *
+                                      (1 + context_scale_mlp) + context_shift_mlp)
+
+        hidden_states = hidden_states + self.ff(norm_hidden_states) * gate_mlp.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + self.ff_context(norm_encoder_hidden_states) * context_gate_mlp
+        return hidden_states, encoder_hidden_states
+
+
+class PAN2Transformer3DModel(BaseDiT):
+    """PAN2 DiT. The input stacks the noisy latents, the conditioning latents and a one-channel conditioning mask
+    along channels, so one model serves text-to-video and first-frame image-to-video."""
+
+    _fsdp_shard_conditions = PAN2VideoConfig()._fsdp_shard_conditions
+    _compile_conditions = PAN2VideoConfig()._compile_conditions
+    _supported_attention_backends = PAN2VideoConfig()._supported_attention_backends
+    param_names_mapping = PAN2VideoConfig().param_names_mapping
+    reverse_param_names_mapping = PAN2VideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = PAN2VideoConfig().lora_param_names_mapping
+
+    def __init__(self, config: PAN2VideoConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = inner_dim
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.num_channels_latents
+        self.out_channels = config.out_channels
+        self.patch_size = tuple(config.patch_size)
+        self.rope_axes_dim = tuple(config.rope_axes_dim)
+        if sum(self.rope_axes_dim) != config.attention_head_dim:
+            raise ValueError(f"rope_axes_dim {self.rope_axes_dim} must sum to the head dim {config.attention_head_dim}")
+        if config.num_attention_heads % get_sp_world_size() != 0:
+            raise ValueError(f"The number of attention heads ({config.num_attention_heads}) must be divisible by the "
+                             f"sequence parallel size ({get_sp_world_size()})")
+
+        self.x_embedder = PatchEmbed(self.patch_size,
+                                     config.in_channels,
+                                     inner_dim,
+                                     prefix=f"{config.prefix}.x_embedder")
+        self.time_embedder = TimestepEmbedder(inner_dim, act_layer="silu", frequency_embedding_size=256)
+
+        self.context_norm = RMSNorm(config.text_embed_dim, eps=1e-5)
+        self.context_proj = ReplicatedLinear(config.text_embed_dim,
+                                             inner_dim,
+                                             bias=True,
+                                             prefix=f"{config.prefix}.context_proj")
+        # Layerwise offload hooks the first nn.ModuleList child, which must be the main blocks.
+        self.transformer_blocks = nn.ModuleList([
+            PAN2TransformerBlock(inner_dim,
+                                 config.num_attention_heads,
+                                 config.attention_head_dim,
+                                 config.mlp_ratio,
+                                 config.context_mlp_ratio,
+                                 supported_attention_backends=self._supported_attention_backends,
+                                 prefix=f"{config.prefix}.transformer_blocks.{i}") for i in range(config.num_layers)
+        ])
+
+        self.context_refiner = nn.ModuleList([
+            PAN2TokenRefinerBlock(inner_dim,
+                                  config.num_attention_heads,
+                                  config.attention_head_dim,
+                                  config.refiner_mlp_ratio,
+                                  eps=1e-5,
+                                  supported_attention_backends=self._supported_attention_backends,
+                                  prefix=f"{config.prefix}.context_refiner.{i}")
+            for i in range(config.num_refiner_layers)
+        ])
+        self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim))
+
+        self.norm_out = PAN2AdaLayerNormOut(inner_dim, eps=1e-6, prefix=f"{config.prefix}.norm_out")
+        self.proj_out = ReplicatedLinear(inner_dim,
+                                         self.patch_size[0] * self.patch_size[1] * self.patch_size[2] *
+                                         self.out_channels,
+                                         bias=True,
+                                         prefix=f"{config.prefix}.proj_out")
+
+        self.gradient_checkpointing = False
+        # See set_fused_kernels; the pipeline turns it on with PAN2PipelineConfig.dit_fused_kernels.
+        self.fused_kernels = False
+        self.__post_init__()
+
+    def set_fused_kernels(self, enabled: bool) -> None:
+        """Run the blocks' fused Q/K RMSNorm + RoPE and LayerNorm + modulation ops (``pan2_fusions``).
+
+        The ops launch Triton kernels on CUDA and run their PyTorch reference elsewhere. The forward reproduces the
+        eager arithmetic up to FP32 summation order, except that each gated residual rounds once instead of twice.
+        """
+        enabled = bool(enabled)
+        for module in self.modules():
+            if isinstance(module, (PAN2TransformerBlock, PAN2Attention, PAN2AdaLayerNormZero, PAN2AdaLayerNormOut)):
+                module.fused_kernels = enabled
+        self.fused_kernels = enabled
+
+    @staticmethod
+    def _local_rotary_pos_embed(freqs_cis: tuple[torch.Tensor, torch.Tensor],
+                                local_len: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """The RoPE rows of this rank's video shard (identity rotation on the sequence-parallel padding)."""
+        cos, sin = freqs_cis
+        start = get_sp_parallel_rank() * local_len
+        pad = start + local_len - cos.shape[0]
+        if pad > 0:
+            cos = F.pad(cos, (0, 0, 0, pad), value=1.0)
+            sin = F.pad(sin, (0, 0, 0, pad), value=0.0)
+        return cos[start:start + local_len], sin[start:start + local_len]
+
+    def _rotary_pos_embed(self, num_frames: int, height: int, width: int,
+                          device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        """Half-dim cos/sin of the 3-axis ``(t, h, w)`` RoPE over the full (unsharded) video token grid."""
+        axes = [torch.arange(size, device=device, dtype=torch.float32) for size in (num_frames, height, width)]
+        grid = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=0)
+        grid[0] = grid[0] * self.config.rope_temporal_scale
+        cos_sin = [
+            get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), theta=self.config.rope_theta, use_real=False)
+            for i, dim in enumerate(self.rope_axes_dim)
+        ]
+        freqs_cos = torch.cat([cos for cos, _ in cos_sin], dim=1).float()
+        freqs_sin = torch.cat([sin for _, sin in cos_sin], dim=1).float()
+        return freqs_cos, freqs_sin
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+                timestep: torch.LongTensor,
+                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+                guidance=None,
+                **kwargs) -> torch.Tensor:
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        batch_size, _, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        freqs_cis = self._rotary_pos_embed(post_patch_num_frames, post_patch_height, post_patch_width,
+                                           hidden_states.device)
+
+        hidden_states = self.x_embedder(hidden_states)
+        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
+        local_freqs_cis = None
+        if self.fused_kernels:
+            local_freqs_cis = self._local_rotary_pos_embed(freqs_cis, hidden_states.shape[1])
+
+        if timestep.dim() == 0:
+            timestep = timestep.expand(batch_size)
+        temb = self.time_embedder(timestep).to(hidden_states.dtype)
+
+        encoder_hidden_states, _ = self.context_proj(self.context_norm(encoder_hidden_states))
+        for refiner_block in self.context_refiner:
+            encoder_hidden_states = refiner_block(encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + self.context_type_embedding
+
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states, encoder_hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len, local_freqs_cis)
+            else:
+                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis,
+                                                             original_seq_len, local_freqs_cis)
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
+        hidden_states, _ = self.proj_out(hidden_states)
+
+        # proj_out emits channels before the patch offsets: (C, p_t, p_h, p_w).
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
+                                              -1, p_t, p_h, p_w)
+        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+
+EntryClass = PAN2Transformer3DModel

fastvideo/models/dits/pan2_fusions.py

@@ -0,0 +1,512 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused PAN2 block kernels with backward passes, ported from the PAN2 reference trainer.
+
+The PAN2 reference fuses its memory-bound per-token work into Triton kernels with hand-written backwards
+(``pan/models/transformers/modules/triton_kernels.py``). Two of them map onto FastVideo's unpacked layout:
+
+- ``qk_rmsnorm_rope``: the per-head query/key RMSNorm and, for the video stream, the interleaved 3-axis RoPE, in one
+  pass over the projection output (the reference's ``qkv_permute_rms`` + ``qkv_rstd_rope`` without its
+  concat/interleave permute). The eager path runs about twenty kernels with FP32 temporaries for this.
+- ``layer_norm_modulate``: the affine-free LayerNorm and its adaptive ``* (1 + scale) + shift`` (the reference's
+  ``packed_modulate`` with ``ln_stats``), computed row-wise in one pass instead of from cached statistics.
+
+Each op is a ``torch.library.custom_op`` with a registered autograd formula, so activation checkpointing and
+``torch.compile`` see one node. The forwards reproduce the eager path's roundings (the normalized value and the
+weighted value round to the input dtype before the next step, as ``RMSNorm.forward_native`` and ``_apply_rotary_emb``
+do), so the fused and eager forwards differ only by FP32 summation order. The backwards keep FP32 throughout, where
+eager autograd rounds some intermediate gradients to BF16.
+
+On CUDA the ops launch Triton; elsewhere they run the PyTorch reference of the same arithmetic.
+"""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:  # pragma: no cover - CPU-only installs
+    triton = None
+    tl = None
+
+HAVE_TRITON = triton is not None
+# Tests set this to drive CPU tensors through the Triton kernels under ``TRITON_INTERPRET=1``.
+_FORCE_TRITON = False
+
+
+def fusions_available(x: torch.Tensor) -> bool:
+    """Whether the fused ops run their Triton kernels for ``x``."""
+    return HAVE_TRITON and (x.is_cuda or _FORCE_TRITON)
+
+
+def _use_triton(x: torch.Tensor) -> bool:
+    return fusions_available(x)
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# Query/key RMSNorm + RoPE
+# --------------------------------------------------------------------------------------------------------------------
+
+_QK_BLOCK_ROWS = 16
+
+if HAVE_TRITON:
+
+    @triton.jit
+    def _qk_rmsnorm_rope_fwd_kernel(X, W, COS, SIN, OUT, RSTD, n_rows, seq_len, num_heads, stride_xr, stride_xh,
+                                    stride_or, stride_oh, stride_cs, eps, HEAD_DIM: tl.constexpr,
+                                    BLOCK_ROWS: tl.constexpr, HAS_ROPE: tl.constexpr):
+        HALF: tl.constexpr = HEAD_DIM // 2
+        rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+        head = tl.program_id(1)
+        pairs = tl.arange(0, HALF)
+        keep = (rows < n_rows)[:, None]
+        rows64 = rows.to(tl.int64)
+        x_base = X + rows64[:, None] * stride_xr + head * stride_xh
+        x_even = tl.load(x_base + 2 * pairs[None, :], mask=keep, other=0.0).to(tl.float32)
+        x_odd = tl.load(x_base + 2 * pairs[None, :] + 1, mask=keep, other=0.0).to(tl.float32)
+        mean_square = (tl.sum(x_even * x_even, axis=1) + tl.sum(x_odd * x_odd, axis=1)) / HEAD_DIM
+        rstd = tl.rsqrt(mean_square + eps)
+        # The eager RMSNorm rounds the normalized value, then the weighted value, to the input dtype.
+        n_even = (x_even * rstd[:, None]).to(OUT.dtype.element_ty).to(tl.float32)
+        n_odd = (x_odd * rstd[:, None]).to(OUT.dtype.element_ty).to(tl.float32)
+        w_even = tl.load(W + 2 * pairs).to(tl.float32)
+        w_odd = tl.load(W + 2 * pairs + 1).to(tl.float32)
+        y_even = (n_even * w_even[None, :]).to(OUT.dtype.element_ty).to(tl.float32)
+        y_odd = (n_odd * w_odd[None, :]).to(OUT.dtype.element_ty).to(tl.float32)
+        if HAS_ROPE:
+            positions = (rows % seq_len).to(tl.int64)
+            cs_offsets = positions[:, None] * stride_cs + pairs[None, :]
+            cos = tl.load(COS + cs_offsets, mask=keep, other=1.0)
+            sin = tl.load(SIN + cs_offsets, mask=keep, other=0.0)
+            out_even = y_even * cos - y_odd * sin
+            out_odd = y_odd * cos + y_even * sin
+        else:
+            out_even = y_even
+            out_odd = y_odd
+        out_base = OUT + rows64[:, None] * stride_or + head * stride_oh
+        tl.store(out_base + 2 * pairs[None, :], out_even.to(OUT.dtype.element_ty), mask=keep)
+        tl.store(out_base + 2 * pairs[None, :] + 1, out_odd.to(OUT.dtype.element_ty), mask=keep)
+        tl.store(RSTD + rows64 * num_heads + head, rstd, mask=rows < n_rows)
+
+    @triton.jit
+    def _qk_rmsnorm_rope_bwd_kernel(X, W, COS, SIN, GRAD_OUT, RSTD, GRAD_X, GRAD_W_PARTIAL, n_rows, seq_len, num_heads,
+                                    stride_xr, stride_xh, stride_gr, stride_gh, stride_cs, HEAD_DIM: tl.constexpr,
+                                    BLOCK_ROWS: tl.constexpr, HAS_ROPE: tl.constexpr):
+        HALF: tl.constexpr = HEAD_DIM // 2
+        block = tl.program_id(0)
+        head = tl.program_id(1)
+        rows = block * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
+        pairs = tl.arange(0, HALF)
+        keep = (rows < n_rows)[:, None]
+        rows64 = rows.to(tl.int64)
+        x_base = X + rows64[:, None] * stride_xr + head * stride_xh
+        x_even = tl.load(x_base + 2 * pairs[None, :], mask=keep, other=0.0).to(tl.float32)
+        x_odd = tl.load(x_base + 2 * pairs[None, :] + 1, mask=keep, other=0.0).to(tl.float32)
+        rstd = tl.load(RSTD + rows64 * num_heads + head, mask=rows < n_rows, other=0.0)
+        n_even = x_even * rstd[:, None]
+        n_odd = x_odd * rstd[:, None]
+        g_base = GRAD_OUT + rows64[:, None] * stride_gr + head * stride_gh
+        g_even = tl.load(g_base + 2 * pairs[None, :], mask=keep, other=0.0).to(tl.float32)
+        g_odd = tl.load(g_base + 2 * pairs[None, :] + 1, mask=keep, other=0.0).to(tl.float32)
+        if HAS_ROPE:
+            positions = (rows % seq_len).to(tl.int64)
+            cs_offsets = positions[:, None] * stride_cs + pairs[None, :]
+            cos = tl.load(COS + cs_offsets, mask=keep, other=1.0)
+            sin = tl.load(SIN + cs_offsets, mask=keep, other=0.0)
+            # The transpose of each pair's rotation.
+            gy_even = g_even * cos + g_odd * sin
+            gy_odd = g_odd * cos - g_even * sin
+        else:
+            gy_even = g_even
+            gy_odd = g_odd
+        # The weight multiplies the rounded normalized value; its gradient sums over this block's rows.
+        rounded_even = n_even.to(GRAD_X.dtype.element_ty).to(tl.float32)
+        rounded_odd = n_odd.to(GRAD_X.dtype.element_ty).to(tl.float32)
+        partial = GRAD_W_PARTIAL + (block * num_heads + head).to(tl.int64) * HEAD_DIM
+        tl.store(partial + 2 * pairs, tl.sum(gy_even * rounded_even, axis=0))
+        tl.store(partial + 2 * pairs + 1, tl.sum(gy_odd * rounded_odd, axis=0))
+        w_even = tl.load(W + 2 * pairs).to(tl.float32)
+        w_odd = tl.load(W + 2 * pairs + 1).to(tl.float32)
+        gn_even = gy_even * w_even[None, :]
+        gn_odd = gy_odd * w_odd[None, :]
+        projection = (tl.sum(gn_even * n_even, axis=1) + tl.sum(gn_odd * n_odd, axis=1)) / HEAD_DIM
+        gx_even = rstd[:, None] * (gn_even - n_even * projection[:, None])
+        gx_odd = rstd[:, None] * (gn_odd - n_odd * projection[:, None])
+        gx_base = GRAD_X + rows64[:, None] * (num_heads * HEAD_DIM) + head * HEAD_DIM
+        tl.store(gx_base + 2 * pairs[None, :], gx_even.to(GRAD_X.dtype.element_ty), mask=keep)
+        tl.store(gx_base + 2 * pairs[None, :] + 1, gx_odd.to(GRAD_X.dtype.element_ty), mask=keep)
+
+
+def _rows_view(x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
+    """``[B, S, H, D]`` (last dim contiguous) as row/head strides over ``B * S`` rows, copying only when it must."""
+    if x.stride(-1) != 1 or x.stride(0) != x.shape[1] * x.stride(1):
+        x = x.contiguous()
+    return x, x.stride(1), x.stride(2)
+
+
+def _check_rope(x: torch.Tensor, cos: torch.Tensor | None, sin: torch.Tensor | None) -> None:
+    if (cos is None) != (sin is None):
+        raise ValueError("cos and sin go together")
+    if cos is not None and (cos.shape != (x.shape[1], x.shape[-1] // 2) or sin.shape != cos.shape):
+        raise ValueError(f"RoPE tables must be [seq_len, head_dim / 2] = [{x.shape[1]}, {x.shape[-1] // 2}], got "
+                         f"{tuple(cos.shape)} and {tuple(sin.shape)}")
+
+
+def _qk_forward_triton(x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None, sin: torch.Tensor | None,
+                       eps: float) -> tuple[torch.Tensor, torch.Tensor]:
+    batch, seq_len, num_heads, head_dim = x.shape
+    x, stride_r, stride_h = _rows_view(x)
+    out = torch.empty((batch, seq_len, num_heads, head_dim), device=x.device, dtype=x.dtype)
+    rstd = torch.empty((batch * seq_len, num_heads), device=x.device, dtype=torch.float32)
+    n_rows = batch * seq_len
+    has_rope = cos is not None
+    tables = (cos.float().contiguous(), sin.float().contiguous()) if has_rope else (rstd, rstd)
+    grid = (triton.cdiv(n_rows, _QK_BLOCK_ROWS), num_heads)
+    _qk_rmsnorm_rope_fwd_kernel[grid](x, weight.contiguous(), tables[0], tables[1], out, rstd, n_rows, seq_len,
+                                      num_heads, stride_r, stride_h, out.stride(1), out.stride(2),
+                                      tables[0].stride(0), eps, HEAD_DIM=head_dim, BLOCK_ROWS=_QK_BLOCK_ROWS,
+                                      HAS_ROPE=has_rope)
+    return out, rstd
+
+
+def _qk_backward_triton(grad_out: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None,
+                        sin: torch.Tensor | None, rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    batch, seq_len, num_heads, head_dim = x.shape
+    x, stride_r, stride_h = _rows_view(x)
+    grad_out, stride_gr, stride_gh = _rows_view(grad_out)
+    n_rows = batch * seq_len
+    num_blocks = triton.cdiv(n_rows, _QK_BLOCK_ROWS)
+    grad_x = torch.empty((batch, seq_len, num_heads, head_dim), device=x.device, dtype=x.dtype)
+    partial = torch.empty((num_blocks * num_heads, head_dim), device=x.device, dtype=torch.float32)
+    has_rope = cos is not None
+    tables = (cos.float().contiguous(), sin.float().contiguous()) if has_rope else (rstd, rstd)
+    _qk_rmsnorm_rope_bwd_kernel[(num_blocks, num_heads)](x, weight.contiguous(), tables[0], tables[1], grad_out, rstd,
+                                                        grad_x, partial, n_rows, seq_len, num_heads, stride_r,
+                                                        stride_h, stride_gr, stride_gh, tables[0].stride(0),
+                                                        HEAD_DIM=head_dim, BLOCK_ROWS=_QK_BLOCK_ROWS,
+                                                        HAS_ROPE=has_rope)
+    return grad_x, partial.sum(0).to(weight.dtype)
+
+
+def _compute_dtype(x: torch.Tensor) -> torch.dtype:
+    """FP32, the eager kernels' arithmetic, except for FP64 inputs (gradient checks)."""
+    return torch.float64 if x.dtype == torch.float64 else torch.float32
+
+
+def _qk_forward_reference(x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None, sin: torch.Tensor | None,
+                          eps: float) -> tuple[torch.Tensor, torch.Tensor]:
+    """The eager arithmetic: ``RMSNorm.forward_native``, ``.to(dtype)``, then ``_apply_rotary_emb`` (GPT-J pairs)."""
+    compute = _compute_dtype(x)
+    x32 = x.to(compute)
+    rstd = torch.rsqrt(x32.pow(2).mean(-1) + eps)
+    normalized = (x32 * rstd[..., None]).to(x.dtype)
+    y = (normalized.to(compute) * weight.to(compute)).to(x.dtype)
+    if cos is not None:
+        y_even, y_odd = y[..., ::2].to(compute), y[..., 1::2].to(compute)
+        cos, sin = cos.to(compute)[None, :, None, :], sin.to(compute)[None, :, None, :]
+        out_even = (y_even * cos - y_odd * sin).to(x.dtype)
+        out_odd = (y_odd * cos + y_even * sin).to(x.dtype)
+        y = torch.stack((out_even, out_odd), dim=-1).flatten(-2)
+    return y.contiguous(), rstd.reshape(-1, x.shape[2])
+
+
+def _qk_backward_reference(grad_out: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor | None,
+                           sin: torch.Tensor | None, rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP32 formulas of the Triton backward (rounding steps count as identities, as eager autograd treats casts)."""
+    compute = _compute_dtype(x)
+    x32 = x.to(compute)
+    rstd = rstd.reshape(x.shape[:3])[..., None].to(compute)
+    normalized = x32 * rstd
+    grad = grad_out.to(compute)
+    if cos is not None:
+        g_even, g_odd = grad[..., ::2], grad[..., 1::2]
+        cos, sin = cos.to(compute)[None, :, None, :], sin.to(compute)[None, :, None, :]
+        grad = torch.stack((g_even * cos + g_odd * sin, g_odd * cos - g_even * sin), dim=-1).flatten(-2)
+    grad_weight = (grad * normalized.to(x.dtype).to(compute)).sum(dim=(0, 1, 2)).to(weight.dtype)
+    grad_normalized = grad * weight.to(compute)
+    projection = (grad_normalized * normalized).mean(-1, keepdim=True)
+    grad_x = rstd * (grad_normalized - normalized * projection)
+    return grad_x.to(x.dtype), grad_weight
+
+
+@torch.library.custom_op("fastvideo::pan2_qk_rmsnorm_rope", mutates_args=())
+def _qk_rmsnorm_rope_op(q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor,
+                        cos: torch.Tensor | None, sin: torch.Tensor | None,
+                        eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    forward = _qk_forward_triton if _use_triton(q) else _qk_forward_reference
+    q_out, q_rstd = forward(q, q_weight, cos, sin, eps)
+    k_out, k_rstd = forward(k, k_weight, cos, sin, eps)
+    return q_out, k_out, q_rstd, k_rstd
+
+
+@_qk_rmsnorm_rope_op.register_fake
+def _qk_rmsnorm_rope_fake(q, k, q_weight, k_weight, cos, sin, eps):
+    del q_weight, k_weight, cos, sin, eps
+    q_rows = q.shape[0] * q.shape[1]
+    k_rows = k.shape[0] * k.shape[1]
+    return (q.new_empty(q.shape), k.new_empty(k.shape), q.new_empty((q_rows, q.shape[2]), dtype=_compute_dtype(q)),
+            k.new_empty((k_rows, k.shape[2]), dtype=_compute_dtype(k)))
+
+
+def _qk_setup_context(ctx, inputs, output) -> None:
+    q, k, q_weight, k_weight, cos, sin, _eps = inputs
+    _q_out, _k_out, q_rstd, k_rstd = output
+    ctx.save_for_backward(q, k, q_weight, k_weight, cos, sin, q_rstd, k_rstd)
+    ctx.mark_non_differentiable(q_rstd, k_rstd)
+
+
+@torch.library.custom_op("fastvideo::pan2_qk_rmsnorm_rope_backward", mutates_args=())
+def _qk_rmsnorm_rope_backward_op(grad_q_out: torch.Tensor, grad_k_out: torch.Tensor, q: torch.Tensor, k: torch.Tensor,
+                                 q_weight: torch.Tensor, k_weight: torch.Tensor, cos: torch.Tensor | None,
+                                 sin: torch.Tensor | None, q_rstd: torch.Tensor,
+                                 k_rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The backward as its own op: compiled graphs (AOTAutograd traces the autograd formula) hold one opaque node
+    instead of tracing into the Triton launch."""
+    backward = _qk_backward_triton if _use_triton(q) else _qk_backward_reference
+    grad_q, grad_q_weight = backward(grad_q_out, q, q_weight, cos, sin, q_rstd)
+    grad_k, grad_k_weight = backward(grad_k_out, k, k_weight, cos, sin, k_rstd)
+    return grad_q, grad_k, grad_q_weight, grad_k_weight
+
+
+@_qk_rmsnorm_rope_backward_op.register_fake
+def _qk_rmsnorm_rope_backward_fake(grad_q_out, grad_k_out, q, k, q_weight, k_weight, cos, sin, q_rstd, k_rstd):
+    del grad_q_out, grad_k_out, cos, sin, q_rstd, k_rstd
+    return q.new_empty(q.shape), k.new_empty(k.shape), q_weight.new_empty(q_weight.shape), k_weight.new_empty(
+        k_weight.shape)
+
+
+def _qk_backward(ctx, grad_q_out, grad_k_out, _grad_q_rstd, _grad_k_rstd):
+    q, k, q_weight, k_weight, cos, sin, q_rstd, k_rstd = ctx.saved_tensors
+    grad_q, grad_k, grad_q_weight, grad_k_weight = _qk_rmsnorm_rope_backward_op(grad_q_out, grad_k_out, q, k,
+                                                                                q_weight, k_weight, cos, sin, q_rstd,
+                                                                                k_rstd)
+    return grad_q, grad_k, grad_q_weight, grad_k_weight, None, None, None
+
+
+_qk_rmsnorm_rope_op.register_autograd(_qk_backward, setup_context=_qk_setup_context)
+
+
+def qk_rmsnorm_rope(q: torch.Tensor,
+                    k: torch.Tensor,
+                    q_weight: torch.Tensor,
+                    k_weight: torch.Tensor,
+                    eps: float,
+                    cos: torch.Tensor | None = None,
+                    sin: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-head RMSNorm of ``q`` and ``k`` (``[B, S, H, D]``, e.g. views of a fused projection), then RoPE.
+
+    ``cos``/``sin`` are the half-dimension tables ``[S, D / 2]`` of ``get_1d_rotary_pos_embed(use_real=False)``,
+    applied to the interleaved ``(x[2i], x[2i + 1])`` pairs as ``_apply_rotary_emb(is_neox_style=False)`` does; without
+    them only the normalization runs. Returns contiguous ``[B, S, H, D]`` tensors in the input dtype.
+    """
+    if q.dim() != 4 or k.shape[:2] != q.shape[:2] or q.shape[-1] % 2:
+        raise ValueError(f"q and k must be [B, S, H, D] with an even D, got {tuple(q.shape)} and {tuple(k.shape)}")
+    _check_rope(q, cos, sin)
+    q_out, k_out, _, _ = _qk_rmsnorm_rope_op(q, k, q_weight, k_weight, cos, sin, float(eps))
+    return q_out, k_out
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# LayerNorm + adaptive modulation
+# --------------------------------------------------------------------------------------------------------------------
+
+_LN_BWD_ROWS_PER_PROGRAM = 16
+
+if HAVE_TRITON:
+
+    @triton.jit
+    def _layer_norm_modulate_fwd_kernel(X, SCALE, SHIFT, OUT, MEAN, RSTD, seq_len, n_cols, stride_xr, stride_or,
+                                        stride_scale_b, stride_shift_b, eps, BLOCK_C: tl.constexpr):
+        row = tl.program_id(0).to(tl.int64)
+        batch = row // seq_len
+        cols = tl.arange(0, BLOCK_C)
+        valid = cols < n_cols
+        x = tl.load(X + row * stride_xr + cols, mask=valid, other=0.0).to(tl.float32)
+        mean = tl.sum(x, axis=0) / n_cols
+        centered = tl.where(valid, x - mean, 0.0)
+        rstd = tl.rsqrt(tl.sum(centered * centered, axis=0) / n_cols + eps)
+        scale = tl.load(SCALE + batch * stride_scale_b + cols, mask=valid, other=0.0).to(tl.float32)
+        shift = tl.load(SHIFT + batch * stride_shift_b + cols, mask=valid, other=0.0).to(tl.float32)
+        out = centered * rstd * scale + shift
+        tl.store(OUT + row * stride_or + cols, out.to(OUT.dtype.element_ty), mask=valid)
+        tl.store(MEAN + row, mean)
+        tl.store(RSTD + row, rstd)
+
+    @triton.jit
+    def _layer_norm_modulate_bwd_kernel(X, SCALE, GRAD_OUT, MEAN, RSTD, GRAD_X, GRAD_SCALE_PARTIAL,
+                                        GRAD_SHIFT_PARTIAL, seq_len, n_cols, rows_per_program, num_chunks, stride_xr,
+                                        stride_gr, stride_scale_b, BLOCK_C: tl.constexpr):
+        chunk = tl.program_id(0)
+        batch = tl.program_id(1).to(tl.int64)
+        cols = tl.arange(0, BLOCK_C)
+        valid = cols < n_cols
+        scale = tl.load(SCALE + batch * stride_scale_b + cols, mask=valid, other=0.0).to(tl.float32)
+        grad_scale = tl.zeros([BLOCK_C], dtype=tl.float32)
+        grad_shift = tl.zeros([BLOCK_C], dtype=tl.float32)
+        first = chunk * rows_per_program
+        last = tl.minimum(first + rows_per_program, seq_len)
+        for offset in range(first, last):
+            row = batch * seq_len + offset
+            x = tl.load(X + row * stride_xr + cols, mask=valid, other=0.0).to(tl.float32)
+            grad = tl.load(GRAD_OUT + row * stride_gr + cols, mask=valid, other=0.0).to(tl.float32)
+            normalized = tl.where(valid, (x - tl.load(MEAN + row)) * tl.load(RSTD + row), 0.0)
+            grad_scale += grad * normalized
+            grad_shift += grad
+            grad_normalized = grad * scale
+            mean_grad = tl.sum(grad_normalized, axis=0) / n_cols
+            mean_projection = tl.sum(grad_normalized * normalized, axis=0) / n_cols
+            grad_x = tl.load(RSTD + row) * (grad_normalized - mean_grad - normalized * mean_projection)
+            tl.store(GRAD_X + row * n_cols + cols, grad_x.to(GRAD_X.dtype.element_ty), mask=valid)
+        partial = (batch * num_chunks + chunk) * n_cols
+        tl.store(GRAD_SCALE_PARTIAL + partial + cols, grad_scale, mask=valid)
+        tl.store(GRAD_SHIFT_PARTIAL + partial + cols, grad_shift, mask=valid)
+
+
+def _batch_rows(table: torch.Tensor, batch: int, n_cols: int) -> torch.Tensor:
+    """A ``[B, C]`` (or broadcastable ``[C]``) table with a contiguous last dimension."""
+    table = table.expand(batch, n_cols)
+    return table if table.stride(-1) == 1 else table.contiguous()
+
+
+def _ln_forward_triton(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor,
+                       eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch, seq_len, n_cols = x.shape
+    rows = x.reshape(-1, n_cols)
+    if rows.stride(-1) != 1:
+        rows = rows.contiguous()
+    out = torch.empty((batch, seq_len, n_cols), device=x.device, dtype=x.dtype)
+    mean = torch.empty(batch * seq_len, device=x.device, dtype=torch.float32)
+    rstd = torch.empty_like(mean)
+    scale, shift = _batch_rows(scale, batch, n_cols), _batch_rows(shift, batch, n_cols)
+    block = triton.next_power_of_2(n_cols)
+    _layer_norm_modulate_fwd_kernel[(batch * seq_len, )](rows, scale, shift, out, mean, rstd, seq_len, n_cols,
+                                                         rows.stride(0), n_cols, scale.stride(0), shift.stride(0), eps,
+                                                         BLOCK_C=block, num_warps=8 if block >= 2048 else 4)
+    return out, mean, rstd
+
+
+def _ln_backward_triton(grad_out: torch.Tensor, x: torch.Tensor, scale: torch.Tensor, mean: torch.Tensor,
+                        rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch, seq_len, n_cols = x.shape
+    rows = x.reshape(-1, n_cols)
+    if rows.stride(-1) != 1:
+        rows = rows.contiguous()
+    grad_rows = grad_out.reshape(-1, n_cols)
+    if grad_rows.stride(-1) != 1:
+        grad_rows = grad_rows.contiguous()
+    scale_rows = _batch_rows(scale, batch, n_cols)
+    num_chunks = triton.cdiv(seq_len, _LN_BWD_ROWS_PER_PROGRAM)
+    grad_x = torch.empty((batch, seq_len, n_cols), device=x.device, dtype=x.dtype)
+    grad_scale = torch.empty((batch, num_chunks, n_cols), device=x.device, dtype=torch.float32)
+    grad_shift = torch.empty_like(grad_scale)
+    block = triton.next_power_of_2(n_cols)
+    _layer_norm_modulate_bwd_kernel[(num_chunks, batch)](rows, scale_rows, grad_rows, mean, rstd, grad_x, grad_scale,
+                                                         grad_shift, seq_len, n_cols, _LN_BWD_ROWS_PER_PROGRAM,
+                                                         num_chunks, rows.stride(0), grad_rows.stride(0),
+                                                         scale_rows.stride(0), BLOCK_C=block,
+                                                         num_warps=8 if block >= 2048 else 4)
+    return grad_x, grad_scale.sum(1), grad_shift.sum(1)
+
+
+def _ln_forward_reference(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor,
+                          eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The autocast arithmetic: FP32 LayerNorm, ``* (1 + scale) + shift`` in FP32, one rounding at the end."""
+    compute = _compute_dtype(x)
+    x32 = x.to(compute)
+    mean = x32.mean(-1)
+    rstd = torch.rsqrt(x32.var(-1, unbiased=False) + eps)
+    normalized = (x32 - mean[..., None]) * rstd[..., None]
+    batch, _, n_cols = x.shape
+    out = (normalized * _batch_rows(scale, batch, n_cols).to(compute)[:, None] +
+           _batch_rows(shift, batch, n_cols).to(compute)[:, None])
+    return out.to(x.dtype), mean.reshape(-1), rstd.reshape(-1)
+
+
+def _ln_backward_reference(grad_out: torch.Tensor, x: torch.Tensor, scale: torch.Tensor, mean: torch.Tensor,
+                           rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    compute = _compute_dtype(x)
+    batch, seq_len, n_cols = x.shape
+    mean, rstd = mean.reshape(batch, seq_len, 1).to(compute), rstd.reshape(batch, seq_len, 1).to(compute)
+    normalized = (x.to(compute) - mean) * rstd
+    grad = grad_out.to(compute)
+    grad_normalized = grad * _batch_rows(scale, batch, n_cols).to(compute)[:, None]
+    grad_x = rstd * (grad_normalized - grad_normalized.mean(-1, keepdim=True) -
+                     normalized * (grad_normalized * normalized).mean(-1, keepdim=True))
+    return grad_x.to(x.dtype), (grad * normalized).sum(1), grad.sum(1)
+
+
+@torch.library.custom_op("fastvideo::pan2_layer_norm_modulate", mutates_args=())
+def _layer_norm_modulate_op(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor,
+                            eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    forward = _ln_forward_triton if _use_triton(x) else _ln_forward_reference
+    return forward(x, scale, shift, eps)
+
+
+@_layer_norm_modulate_op.register_fake
+def _layer_norm_modulate_fake(x, scale, shift, eps):
+    del scale, shift, eps
+    rows = x.shape[0] * x.shape[1]
+    stats = _compute_dtype(x)
+    return x.new_empty(x.shape), x.new_empty(rows, dtype=stats), x.new_empty(rows, dtype=stats)
+
+
+def _ln_setup_context(ctx, inputs, output) -> None:
+    x, scale, shift, _eps = inputs
+    _out, mean, rstd = output
+    ctx.save_for_backward(x, scale, mean, rstd)
+    ctx.scale_shape, ctx.shift_shape = scale.shape, shift.shape
+    ctx.scale_dtype, ctx.shift_dtype = scale.dtype, shift.dtype
+    ctx.mark_non_differentiable(mean, rstd)
+
+
+@torch.library.custom_op("fastvideo::pan2_layer_norm_modulate_backward", mutates_args=())
+def _layer_norm_modulate_backward_op(grad_out: torch.Tensor, x: torch.Tensor, scale: torch.Tensor, mean: torch.Tensor,
+                                     rstd: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """(grad x, grad scale, grad shift), the tables' gradients per sample ``[B, C]``; an op of its own for the same
+    reason as ``_qk_rmsnorm_rope_backward_op``."""
+    backward = _ln_backward_triton if _use_triton(x) else _ln_backward_reference
+    return backward(grad_out, x, scale, mean, rstd)
+
+
+@_layer_norm_modulate_backward_op.register_fake
+def _layer_norm_modulate_backward_fake(grad_out, x, scale, mean, rstd):
+    del grad_out, scale, mean, rstd
+    tables = (x.shape[0], x.shape[2])
+    stats = _compute_dtype(x)
+    return x.new_empty(x.shape), x.new_empty(tables, dtype=stats), x.new_empty(tables, dtype=stats)
+
+
+def _ln_backward(ctx, grad_out, _grad_mean, _grad_rstd):
+    x, scale, mean, rstd = ctx.saved_tensors
+    grad_x, grad_scale, grad_shift = _layer_norm_modulate_backward_op(grad_out, x, scale, mean, rstd)
+    # Tables shared by the batch ([C]) receive the batch's summed gradient.
+    grad_scale = grad_scale.sum_to_size(ctx.scale_shape).to(ctx.scale_dtype)
+    grad_shift = grad_shift.sum_to_size(ctx.shift_shape).to(ctx.shift_dtype)
+    return grad_x, grad_scale, grad_shift, None
+
+
+_layer_norm_modulate_op.register_autograd(_ln_backward, setup_context=_ln_setup_context)
+
+
+def layer_norm_modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float) -> torch.Tensor:
+    """``LayerNorm(x) * scale + shift`` for ``x`` ``[B, S, C]``, affine-free LayerNorm, in one pass.
+
+    ``scale`` (PAN2's ``1 + scale``) and ``shift`` are ``[B, C]`` per sample or ``[C]`` shared. The arithmetic is the
+    autocast path's: an FP32 LayerNorm and FP32 modulation, rounded once to ``x``'s dtype. Under autocast that is
+    exactly the value the next linear layer would have cast its FP32 input to.
+    """
+    if x.dim() != 3 or scale.shape[-1] != x.shape[-1] or shift.shape[-1] != x.shape[-1]:
+        raise ValueError(f"x must be [B, S, C] with [B, C] or [C] tables, got {tuple(x.shape)}, "
+                         f"{tuple(scale.shape)} and {tuple(shift.shape)}")
+    out, _, _ = _layer_norm_modulate_op(x, scale, shift, float(eps))
+    return out
+
+
+
+__all__ = [
+    "HAVE_TRITON",
+    "fusions_available",
+    "layer_norm_modulate",
+    "qk_rmsnorm_rope",
+]

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/qwen3_5.py

@@ -0,0 +1,574 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.5 text encoder: hybrid Gated DeltaNet linear attention + gated full attention.
+
+The encoder runs replicated (no tensor or sequence parallelism) and
+inference-only, so the linear-attention layers use a pure-PyTorch chunked gated
+delta rule instead of the fused ``fla`` / ``causal_conv1d`` kernels. Full
+attention calls ``scaled_dot_product_attention`` directly with head-major
+``[B, H, L, D]`` inputs: that is the layout the reference text model feeds SDPA,
+and keeping it makes bf16 hidden states reproducible against that reference.
+"""
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextArchConfig, Qwen3_5TextConfig
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.logger import init_logger
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+logger = init_logger(__name__)
+
+_GATED_DELTA_CHUNK_SIZE = 64
+# Checkpoints saved from the multimodal wrapper nest the text tower.
+_CHECKPOINT_PREFIXES = ("model.language_model.", "language_model.", "model.")
+_SKIPPED_CHECKPOINT_PREFIXES = ("lm_head.", "mtp.", "visual.", "model.visual.")
+
+
+class Qwen3_5RMSNorm(nn.Module):
+    """Zero-centered RMSNorm: the stored weight is an offset from one."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = x.float()
+        output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + self.eps)
+        # The scale is applied in float32 before the downcast.
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class Qwen3_5RMSNormGated(nn.Module):
+    """Per-head RMSNorm followed by a SiLU output gate (plain, not zero-centered, weight)."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        hidden_states = self.weight * hidden_states.to(input_dtype)
+        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+        return hidden_states.to(input_dtype)
+
+
+def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
+
+
+def chunk_gated_delta_rule(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = _GATED_DELTA_CHUNK_SIZE,
+) -> torch.Tensor:
+    """Chunked gated delta rule over a full prefix with a zero initial state.
+
+    Inputs are ``[batch, seq, heads, dim]`` (``g`` / ``beta`` are
+    ``[batch, seq, heads]``); q/k must already be L2-normalized. The recurrence
+    runs in float32 and the output is returned in the query dtype.
+    """
+    output_dtype = query.dtype
+    query, key, value, beta, g = (x.transpose(1, 2).contiguous().to(torch.float32)
+                                  for x in (query, key, value, beta, g))
+
+    batch_size, num_heads, sequence_length, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+    query = F.pad(query, (0, 0, 0, pad_size))
+    key = F.pad(key, (0, 0, 0, pad_size))
+    value = F.pad(value, (0, 0, 0, pad_size))
+    beta = F.pad(beta, (0, pad_size))
+    g = F.pad(g, (0, pad_size))
+    total_sequence_length = sequence_length + pad_size
+    query = query * (1 / (query.shape[-1]**0.5))
+
+    v_beta = value * beta.unsqueeze(-1)
+    k_beta = key * beta.unsqueeze(-1)
+    query, key, value, k_beta, v_beta = (x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
+                                         for x in (query, key, value, k_beta, v_beta))
+    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
+    upper_with_diagonal = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
+                                     diagonal=0)
+
+    # Intra-chunk decay and the (I + strictly-lower)^-1 WY transform, solved by
+    # forward substitution row by row.
+    g = g.cumsum(dim=-1)
+    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(upper_with_diagonal, 0)
+    for i in range(1, chunk_size):
+        row = attn[..., i, :i].clone()
+        sub = attn[..., :i, :i].clone()
+        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+    value = attn @ v_beta
+    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
+
+    recurrent_state = torch.zeros(batch_size,
+                                  num_heads,
+                                  k_head_dim,
+                                  v_head_dim,
+                                  dtype=value.dtype,
+                                  device=value.device)
+    core_attn_out = torch.zeros_like(value)
+    for i in range(total_sequence_length // chunk_size):
+        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
+        # decay_mask is already zero above the diagonal, which keeps this causal.
+        attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]
+        v_prime = k_cumdecay[:, :, i] @ recurrent_state
+        v_new = v_i - v_prime
+        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ recurrent_state
+        core_attn_out[:, :, i] = attn_inter + attn @ v_new
+        recurrent_state = (recurrent_state * g[:, :, i, -1, None, None].exp() +
+                           (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new)
+
+    core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
+    core_attn_out = core_attn_out[:, :, :sequence_length]
+    return core_attn_out.transpose(1, 2).contiguous().to(output_dtype)
+
+
+class Qwen3_5GatedDeltaNet(nn.Module):
+    """Gated DeltaNet token mixer used by the ``linear_attention`` layers."""
+
+    def __init__(
+        self,
+        config: Qwen3_5TextArchConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_v_heads = config.linear_num_value_heads
+        self.num_k_heads = config.linear_num_key_heads
+        if self.num_v_heads % self.num_k_heads:
+            raise ValueError(f"Qwen3.5 linear value heads ({self.num_v_heads}) must be a multiple of key heads "
+                             f"({self.num_k_heads})")
+        self.head_k_dim = config.linear_key_head_dim
+        self.head_v_dim = config.linear_value_head_dim
+        self.key_dim = self.head_k_dim * self.num_k_heads
+        self.value_dim = self.head_v_dim * self.num_v_heads
+        self.conv_kernel_size = config.linear_conv_kernel_dim
+        if config.hidden_act != "silu":
+            raise ValueError(f"Unsupported Qwen3.5 activation: {config.hidden_act}")
+
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=False,
+            kernel_size=self.conv_kernel_size,
+            groups=self.conv_dim,
+            padding=self.conv_kernel_size - 1,
+        )
+        self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
+        self.A_log = nn.Parameter(torch.zeros(self.num_v_heads))
+        self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=config.rms_norm_eps)
+        self.out_proj = ReplicatedLinear(self.value_dim,
+                                         self.hidden_size,
+                                         bias=False,
+                                         quant_config=quant_config,
+                                         prefix=f"{prefix}.out_proj")
+        self.in_proj_qkv = ReplicatedLinear(self.hidden_size,
+                                            self.conv_dim,
+                                            bias=False,
+                                            quant_config=quant_config,
+                                            prefix=f"{prefix}.in_proj_qkv")
+        self.in_proj_z = ReplicatedLinear(self.hidden_size,
+                                          self.value_dim,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.in_proj_z")
+        self.in_proj_b = ReplicatedLinear(self.hidden_size,
+                                          self.num_v_heads,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.in_proj_b")
+        self.in_proj_a = ReplicatedLinear(self.hidden_size,
+                                          self.num_v_heads,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.in_proj_a")
+
+    def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor:
+        if padding_mask is not None:
+            # Zero padded tokens so they cannot leak into the recurrent state.
+            hidden_states = (hidden_states * padding_mask[:, :, None]).to(hidden_states.dtype)
+        batch_size, seq_len, _ = hidden_states.shape
+
+        mixed_qkv, _ = self.in_proj_qkv(hidden_states)
+        z, _ = self.in_proj_z(hidden_states)
+        b, _ = self.in_proj_b(hidden_states)
+        a, _ = self.in_proj_a(hidden_states)
+
+        # Causal depthwise conv: left padding of kernel-1, keep the first seq_len outputs.
+        mixed_qkv = mixed_qkv.transpose(1, 2)
+        mixed_qkv = F.silu(self.conv1d(mixed_qkv.to(self.conv1d.weight.dtype))[:, :, :seq_len]).to(mixed_qkv.dtype)
+        mixed_qkv = mixed_qkv.transpose(1, 2)
+        query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
+        query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)
+        key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)
+        value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)
+
+        beta = b.sigmoid()
+        # float32 keeps exp(A_log) finite for half-precision weights.
+        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+        if self.num_v_heads > self.num_k_heads:
+            repeats = self.num_v_heads // self.num_k_heads
+            query = query.repeat_interleave(repeats, dim=2)
+            key = key.repeat_interleave(repeats, dim=2)
+
+        core_attn_out = chunk_gated_delta_rule(_l2norm(query), _l2norm(key), value, g=g, beta=beta)
+
+        core_attn_out = self.norm(core_attn_out.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim))
+        core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)
+        output, _ = self.out_proj(core_attn_out)
+        return output
+
+
+class Qwen3_5RotaryEmbedding(nn.Module):
+    """Text-only partial RoPE.
+
+    Qwen3.5's interleaved M-RoPE collapses to standard RoPE when the temporal,
+    height and width positions are identical, which holds for text tokens.
+    """
+
+    def __init__(self, config: Qwen3_5TextArchConfig) -> None:
+        super().__init__()
+        rope = config.resolved_rope_parameters()
+        if rope["rope_type"] != "default":
+            raise NotImplementedError(f"Qwen3.5 rope_type {rope['rope_type']!r} is not supported")
+        rotary_dim = int(config.head_dim * rope["partial_rotary_factor"])
+        inv_freq = 1.0 / (rope["rope_theta"]**(
+            torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cpu") / rotary_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq = self.inv_freq.to(device=hidden_states.device, dtype=torch.float32)
+        device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            inv_freq = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
+            freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+        return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _repeat_kv(x: torch.Tensor, groups: int) -> torch.Tensor:
+    batch_size, num_kv_heads, seq_len, head_dim = x.shape
+    x = x[:, :, None].expand(batch_size, num_kv_heads, groups, seq_len, head_dim)
+    return x.reshape(batch_size, num_kv_heads * groups, seq_len, head_dim)
+
+
+def _apply_partial_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Rotate the leading ``cos.shape[-1]`` channels of ``[B, H, L, D]`` and pass the rest through."""
+    rotary_dim = cos.shape[-1]
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
+    x_rot = x_rot * cos + _rotate_half(x_rot) * sin
+    return torch.cat([x_rot, x_pass], dim=-1)
+
+
+class Qwen3_5Attention(nn.Module):
+    """Gated full attention with per-head QK-norm, partial RoPE and GQA."""
+
+    def __init__(
+        self,
+        config: Qwen3_5TextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        arch = config.arch_config
+        self.num_heads = arch.num_attention_heads
+        self.num_kv_heads = arch.num_key_value_heads
+        if self.num_heads % self.num_kv_heads:
+            raise ValueError(f"Qwen3.5 attention heads ({self.num_heads}) must be a multiple of KV heads "
+                             f"({self.num_kv_heads})")
+        self.head_dim = arch.head_dim
+        self.scaling = self.head_dim**-0.5
+        # q_proj packs, per head, the query followed by its output gate.
+        self.q_proj = ReplicatedLinear(arch.hidden_size,
+                                       self.num_heads * self.head_dim * 2,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.q_proj")
+        self.k_proj = ReplicatedLinear(arch.hidden_size,
+                                       self.num_kv_heads * self.head_dim,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.k_proj")
+        self.v_proj = ReplicatedLinear(arch.hidden_size,
+                                       self.num_kv_heads * self.head_dim,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.v_proj")
+        self.o_proj = ReplicatedLinear(self.num_heads * self.head_dim,
+                                       arch.hidden_size,
+                                       bias=arch.attention_bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.o_proj")
+        self.q_norm = Qwen3_5RMSNorm(self.head_dim, eps=arch.rms_norm_eps)
+        self.k_norm = Qwen3_5RMSNorm(self.head_dim, eps=arch.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        packed, _ = self.q_proj(hidden_states)
+        query, gate = torch.chunk(packed.view(batch_size, seq_len, -1, self.head_dim * 2), 2, dim=-1)
+        gate = gate.reshape(batch_size, seq_len, -1)
+        key, _ = self.k_proj(hidden_states)
+        value, _ = self.v_proj(hidden_states)
+
+        query = self.q_norm(query).transpose(1, 2)
+        key = self.k_norm(key.view(batch_size, seq_len, -1, self.head_dim)).transpose(1, 2)
+        value = value.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
+        cos, sin = position_embeddings
+        query = _apply_partial_rotary(query, cos, sin)
+        key = _apply_partial_rotary(key, cos, sin)
+
+        groups = self.num_heads // self.num_kv_heads
+        # Native GQA is only used where SDPA keeps a fused kernel (no mask, head_dim <= 256).
+        native_gqa = padding_mask is None and self.head_dim <= 256
+        if groups > 1 and not native_gqa:
+            key = _repeat_kv(key, groups)
+            value = _repeat_kv(value, groups)
+        if padding_mask is None:
+            attn_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                dropout_p=0.0,
+                is_causal=seq_len > 1,
+                scale=self.scaling,
+                enable_gqa=groups > 1 and native_gqa,
+            )
+        else:
+            causal = torch.ones(seq_len, seq_len, device=query.device, dtype=torch.bool).tril()
+            mask = causal[None, None] & padding_mask.to(device=query.device, dtype=torch.bool)[:, None, None, :]
+            attn_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scaling,
+            )
+
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, -1)
+        attn_output = attn_output * torch.sigmoid(gate)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen3_5MLP(nn.Module):
+
+    def __init__(
+        self,
+        config: Qwen3_5TextArchConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if config.hidden_act != "silu":
+            raise ValueError(f"Unsupported Qwen3.5 activation: {config.hidden_act}")
+        self.gate_proj = ReplicatedLinear(config.hidden_size,
+                                          config.intermediate_size,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.gate_proj")
+        self.up_proj = ReplicatedLinear(config.hidden_size,
+                                        config.intermediate_size,
+                                        bias=False,
+                                        quant_config=quant_config,
+                                        prefix=f"{prefix}.up_proj")
+        self.down_proj = ReplicatedLinear(config.intermediate_size,
+                                          config.hidden_size,
+                                          bias=False,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.down_proj")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, _ = self.gate_proj(x)
+        up, _ = self.up_proj(x)
+        output, _ = self.down_proj(F.silu(gate) * up)
+        return output
+
+
+class Qwen3_5DecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: Qwen3_5TextConfig,
+        layer_type: str,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        arch = config.arch_config
+        self.layer_type = layer_type
+        if layer_type == "linear_attention":
+            self.linear_attn = Qwen3_5GatedDeltaNet(arch, quant_config=quant_config, prefix=f"{prefix}.linear_attn")
+        else:
+            self.self_attn = Qwen3_5Attention(config, quant_config=quant_config, prefix=f"{prefix}.self_attn")
+        self.mlp = Qwen3_5MLP(arch, quant_config=quant_config, prefix=f"{prefix}.mlp")
+        self.input_layernorm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        if self.layer_type == "linear_attention":
+            hidden_states = self.linear_attn(hidden_states, padding_mask)
+        else:
+            hidden_states = self.self_attn(hidden_states, position_embeddings, padding_mask)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
+
+class Qwen3_5TextModel(TextEncoder):
+    """Qwen3.5 text tower used as a diffusion text encoder (e.g. PAN2).
+
+    ``hidden_states`` follows the Transformers contract: entry 0 is the token
+    embeddings, entry ``i`` is the output of layer ``i - 1``, and the last entry
+    is the final-normed output (it replaces the raw last-layer output).
+    """
+
+    def __init__(self, config: Qwen3_5TextConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        arch = config.arch_config
+        self.quant_config = getattr(config, "quant_config", None)
+        self.layer_types = arch.resolved_layer_types()
+
+        self.embed_tokens = VocabParallelEmbedding(
+            arch.vocab_size,
+            arch.hidden_size,
+            org_num_embeddings=arch.vocab_size,
+            quant_config=self.quant_config,
+        )
+        self.layers = nn.ModuleList([
+            Qwen3_5DecoderLayer(
+                config,
+                layer_type,
+                quant_config=self.quant_config,
+                prefix=f"{config.prefix}.layers.{index}",
+            ) for index, layer_type in enumerate(self.layer_types)
+        ])
+        self.norm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+        self.rotary_emb = Qwen3_5RotaryEmbedding(arch)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs: Any,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else self.config.arch_config.output_hidden_states)
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds")
+        if inputs_embeds is None:
+            assert input_ids is not None
+            if input_ids.ndim == 1:
+                input_ids = input_ids.unsqueeze(0)
+            inputs_embeds = self.get_input_embeddings(input_ids)
+        hidden_states = inputs_embeds
+        batch_size, seq_len = hidden_states.shape[:2]
+
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0).expand(batch_size, -1)
+        elif position_ids.ndim == 1:
+            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        elif position_ids.ndim != 2:
+            raise ValueError("Qwen3.5 text encoder expects text position_ids shaped [batch, seq]")
+
+        # An all-ones mask is equivalent to no mask; only real padding takes the masked path.
+        padding_mask = None
+        if attention_mask is not None:
+            if attention_mask.ndim != 2:
+                raise ValueError("Qwen3.5 text encoder expects a 2D [batch, seq] padding mask")
+            if not bool(attention_mask.to(torch.bool).all()):
+                padding_mask = attention_mask.to(hidden_states.device)
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (hidden_states, )
+            hidden_states = layer(hidden_states, position_embeddings, padding_mask)
+
+        hidden_states = self.norm(hidden_states)
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states, )
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attention_mask=attention_mask,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if name.startswith(_SKIPPED_CHECKPOINT_PREFIXES) or "rotary_emb.inv_freq" in name:
+                continue
+            for checkpoint_prefix in _CHECKPOINT_PREFIXES:
+                if name.startswith(checkpoint_prefix):
+                    name = name[len(checkpoint_prefix):]
+                    break
+            if name not in params_dict:
+                logger.debug("Skipping unexpected Qwen3.5 checkpoint tensor %s", name)
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+EntryClass = Qwen3_5TextModel

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

@@ -44,6 +44,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"),
     "Kandinsky6Transformer3DModel": ("dits", "kandinsky6", "Kandinsky6Transformer3DModel"),
     "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"),
+    "PAN2Transformer3DModel": ("dits", "pan2", "PAN2Transformer3DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -71,7 +72,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"),
@@ -86,6 +92,7 @@ _TEXT_ENCODER_MODELS = {
     "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
     "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
     "Mistral3ForConditionalGeneration": ("encoders", "mistral3", "Mistral3ForConditionalGeneration"),
+    "Qwen3_5TextModel": ("encoders", "qwen3_5", "Qwen3_5TextModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -108,6 +115,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 +162,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/models/wan/vae.py

@@ -15,18 +15,61 @@
 # limitations under the License.
 
 import contextvars
+import functools
+import importlib.util
+import os
+import sys
 from contextlib import contextmanager
+from types import ModuleType
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
 
+import fastvideo.envs as envs
+from fastvideo.distributed.parallel_state import GroupCoordinator, get_sp_group, model_parallel_is_initialized
 from fastvideo.layers.activation import get_act_fn
+from fastvideo.logger import init_logger
 from fastvideo.models.vaes.common import (DiagonalGaussianDistribution, ParallelTiledVAE)
 from fastvideo.models.wan.vae_config import WanVAEConfig
 
+logger = init_logger(__name__)
+
 CACHE_T = 2
+# fastvideo-kernel/wan22_vae_decode decodes the Wan2.2-TI2V-5B VAE's latents [1, 48, T, 44, 78] (704 x 1248) on
+# sm_100a (B200 / GB200) and sm_103a (B300).
+_DECODE_KERNEL_LATENT = (48, 44, 78)
+_DECODE_KERNEL_VIDEO = (704, 1248)
+_DECODE_KERNEL_CAPABILITIES = ((10, 0), (10, 3))
+# The group sizes it splits a decode over (decode_kernel_parallel), each checked to decode the same video as one GPU.
+_DECODE_KERNEL_WORLDS = (2, 4, 8)
+
+
+@functools.cache
+def _wan22_vae_decode() -> ModuleType:
+    """The wrapper of fastvideo-kernel/wan22_vae_decode (``FASTVIDEO_WAN22_VAE_DECODE_DIR``, by default this
+    checkout's), loaded once per process as ``wan22_vae_decode``, as its own tests import it. Its extension builds on
+    first use (into ``WAN22_VAE_DECODE_BUILD_DIR``, see the wrapper)."""
+    module = sys.modules.get("wan22_vae_decode")
+    if module is None:
+        root = envs.FASTVIDEO_WAN22_VAE_DECODE_DIR.get() or os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))),
+            "fastvideo-kernel", "wan22_vae_decode")
+        path = os.path.join(root, "wan22_vae_decode.py")
+        if not os.path.isfile(path):
+            raise RuntimeError(f"decode_kernel needs fastvideo-kernel/wan22_vae_decode; {path} does not exist "
+                               "(FASTVIDEO_WAN22_VAE_DECODE_DIR).")
+        spec = importlib.util.spec_from_file_location("wan22_vae_decode", path)
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        sys.modules["wan22_vae_decode"] = module
+        try:
+            spec.loader.exec_module(module)
+        except BaseException:
+            del sys.modules["wan22_vae_decode"]
+            raise
+    return module
 
 is_first_frame = contextvars.ContextVar("is_first_frame", default=False)
 feat_cache = contextvars.ContextVar("feat_cache", default=None)
@@ -1096,6 +1139,10 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
             )
 
         self.use_feature_cache = config.use_feature_cache
+        # The kernel decoders of decode_kernel, one per device and process group (their weights prepared once; None
+        # where the group's ranks could not all build one), and why the kernel cannot run.
+        self._kernel_decoders: dict[tuple[torch.device, tuple[int, ...] | None], object | None] = {}
+        self._kernel_unavailable: str | None = None
 
     def clear_cache(self) -> None:
 
@@ -1187,6 +1234,10 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
         return enc
 
     def decode(self, z: torch.Tensor) -> torch.Tensor:
+        if getattr(self.config, "decode_kernel", False):
+            video = self._kernel_decode(z)
+            if video is not None:
+                return video
         if self.use_feature_cache:
             self.clear_cache()
             iter_ = z.shape[2]
@@ -1215,6 +1266,80 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
 
         return out
 
+    def _kernel_decode(self, z: torch.Tensor) -> torch.Tensor | None:
+        """``decode`` by fastvideo-kernel's Wan2.2 decoder (``decode_kernel``): bf16 with fp32 accumulation, within the
+        error of a bf16 decode of the fp32 one, or None where it does not apply (other devices, geometries, batches or
+        VAEs), and the decode below runs as before. Every rank of a multi-GPU run returns the whole video, as the
+        decode below does: decoded on its GPU alone, or with ``decode_kernel_parallel`` split over the
+        sequence-parallel group and broadcast from the group's first rank."""
+        # Only what is the same on every rank decides here, so that the ranks of a group take the split together; a
+        # rank whose own build failed still joins the group's agreement (see _kernel_decoder).
+        reason = None
+        if not z.is_cuda:
+            reason = "the latents are not on a GPU"
+        elif torch.cuda.get_device_capability(z.device) not in _DECODE_KERNEL_CAPABILITIES:
+            major, minor = torch.cuda.get_device_capability(z.device)
+            reason = f"the GPU is sm_{major}{minor}, not sm_100 or sm_103"
+        elif z.dim() != 5 or z.shape[0] != 1 or (z.shape[1], *z.shape[3:]) != _DECODE_KERNEL_LATENT:
+            reason = f"latents {tuple(z.shape)} are not [1, 48, T, 44, 78] (a 704 x 1248 video, batch 1)"
+        if reason is not None:
+            logger.warning_once(f"decode_kernel: {reason}; decoding with AutoencoderKLWan.")
+            return None
+        group = self._kernel_group(z)
+        if group is not None and (decoder := self._kernel_decoder(z.device, group)) is not None:
+            logger.info_once(f"decode_kernel: decoding with fastvideo-kernel's Wan2.2 decoder over the "
+                             f"{group.world_size} GPUs of the sequence-parallel group.")
+            video = decoder.decode(z)  # the video on the group's first rank, None on the others
+            if video is None:
+                video = torch.empty((1, 3, 4 * z.shape[2] - 3, *_DECODE_KERNEL_VIDEO), dtype=torch.bfloat16,
+                                    device=z.device)
+            return group.broadcast(video, src=0).float()
+        decoder = self._kernel_decoder(z.device, None)
+        if decoder is None:
+            logger.warning_once(f"decode_kernel: {self._kernel_unavailable}; decoding with AutoencoderKLWan.")
+            return None
+        logger.info_once("decode_kernel: decoding with fastvideo-kernel's Wan2.2 decoder on this GPU.")
+        return decoder.decode(z).float()
+
+    def _kernel_group(self, z: torch.Tensor) -> GroupCoordinator | None:
+        """The sequence-parallel group ``decode_kernel_parallel`` splits the decode of ``z`` over, or None (decode
+        alone): the switch is off, there is no group of a size the kernel splits over, or the video has fewer latent
+        frames than the group has GPUs, where the split's fixed cost (about 50 ms over 4 B300s, 100 ms over 8) exceeds
+        the decode."""
+        if not getattr(self.config, "decode_kernel_parallel", False) or not model_parallel_is_initialized():
+            return None
+        group = get_sp_group()
+        return group if group.world_size in _DECODE_KERNEL_WORLDS and z.shape[2] >= group.world_size else None
+
+    def _kernel_decoder(self, device: torch.device, group: GroupCoordinator | None) -> object | None:
+        """The kernel's decoder on ``device``, alone (``group`` None) or as this rank of ``group``, built on first use
+        (its extension built or loaded, its weights prepared), or None where it cannot be. Over a group every rank
+        builds its own without communicating, then the ranks agree: the split decode needs every one of them, so a
+        rank that could not build makes the whole group decode alone."""
+        key = (device, None if group is None else tuple(group.ranks))
+        if key in self._kernel_decoders:
+            return self._kernel_decoders[key]
+        decoder = None
+        if self._kernel_unavailable is None:
+            try:
+                decoder = _wan22_vae_decode().WanVAEDecoder(self.state_dict(), device=device,
+                                                            pg=None if group is None else group.device_group,
+                                                            local=group is None)
+            except (KeyError, ValueError) as error:  # not the Wan2.2-TI2V-5B VAE's decoder
+                self._kernel_unavailable = f"this VAE's decoder does not fit the kernel ({error})"
+            except Exception as error:  # the kernel failed to load or build; generation goes on without it
+                self._kernel_unavailable = f"the kernel did not load ({type(error).__name__}: {error})"
+        if group is not None:
+            ready = torch.tensor([int(decoder is not None)], dtype=torch.int32)
+            torch.distributed.all_reduce(ready, op=torch.distributed.ReduceOp.MIN, group=group.cpu_group)
+            if not ready.item():
+                if decoder is not None:
+                    logger.warning("decode_kernel_parallel: a rank of the sequence-parallel group could not build "
+                                   "the kernel's decoder; every rank decodes alone.")
+                decoder = None
+        self._kernel_decoders[key] = decoder
+        return decoder
+
     def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor:
         x = self.post_quant_conv(z)
         with forward_context(first_frame_arg=first_frame, use_light_vae_arg=self.config.use_light_vae):

fastvideo/models/wan/vae_config.py

@@ -79,6 +79,13 @@ class WanVAEConfig(VAEConfig):
     use_parallel_tiling: bool = False
 
     use_light_vae: bool = False
+    # Decode with fastvideo-kernel's Wan2.2 decoder (wan22_vae_decode: B200 / GB200 / B300, the Wan2.2-TI2V-5B VAE,
+    # 704 x 1248 video, batch 1) in bf16 where it applies; other inputs and devices take the decode here.
+    decode_kernel: bool = False
+    # With decode_kernel, split each decode over the sequence-parallel group (2, 4 or 8 GPUs) instead of decoding on
+    # every GPU alone. Every rank of the group must decode the same latents at the same time, as a pipeline's decoding
+    # stage does; the group's first rank assembles the video and broadcasts it, so every rank returns it as before.
+    decode_kernel_parallel: bool = False
 
     def __post_init__(self):
         self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2

fastvideo/pipelines/basic/pan2/__init__.py

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

fastvideo/pipelines/basic/pan2/pan2_pipeline.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 video generation: text-to-video, and image-to-video from a first frame."""
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.basic.pan2.stages import (
+    PAN2ConditioningStage,
+    PAN2DenoisingStage,
+    PAN2InputValidationStage,
+    PAN2TextEncodingStage,
+    PAN2TimestepPreparationStage,
+)
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import DecodingStage, LatentPreparationStage
+
+
+class PAN2Pipeline(ComposedPipelineBase):
+    """One stage chain serves both modes: a request with ``image_path`` (or a PIL ``pil_image``) runs image-to-video
+    from that first frame, and one without runs text-to-video."""
+
+    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage", stage=PAN2InputValidationStage())
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=PAN2TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=PAN2TimestepPreparationStage(scheduler=self.get_module("scheduler")))
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(scheduler=self.get_module("scheduler"),
+                                                    transformer=self.get_module("transformer"),
+                                                    latent_dtype=torch.float32))
+        self.add_stage(stage_name="conditioning_stage", stage=PAN2ConditioningStage(vae=self.get_module("vae")))
+        self.add_stage(stage_name="denoising_stage",
+                       stage=PAN2DenoisingStage(transformer=self.get_module("transformer"),
+                                                scheduler=self.get_module("scheduler"),
+                                                pipeline=self))
+        # The performance switches of PAN2PipelineConfig: the DiT's fused kernels and the VAE's decode kernel.
+        pipeline_config = getattr(fastvideo_args, "pipeline_config", None)
+        transformer = self.get_module("transformer")
+        if hasattr(transformer, "set_fused_kernels"):
+            transformer.set_fused_kernels(bool(getattr(pipeline_config, "dit_fused_kernels", False)))
+        vae = self.get_module("vae")
+        vae_config = getattr(vae, "config", None)
+        if hasattr(vae_config, "decode_kernel"):
+            vae_config.decode_kernel = bool(getattr(pipeline_config, "vae_decode_kernel", False))
+            vae_config.decode_kernel_parallel = bool(getattr(pipeline_config, "vae_decode_kernel_parallel", False))
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=vae, pipeline=self))
+
+
+class PAN2ModularPipeline(PAN2Pipeline):
+    """The pipeline class PAN2 checkpoints declare in their ``modular_model_index.json``."""
+
+
+EntryClass = [PAN2Pipeline, PAN2ModularPipeline]

fastvideo/pipelines/basic/pan2/presets.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inference presets for PAN2 text-to-video and image-to-video."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="PAN2 flow-matching denoising",
+    allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}),
+)
+
+# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 2.5 (the release launcher's default since ifm-ai/PAN2
+# 86d5023, 2026-10-23; 3 before). The preset sends no negative prompt (an empty one is encoded
+# as a single space); the reference sampler sends the release's structured negative caption, so pass it as the
+# request's ``negative_prompt`` to follow that recipe.
+_SHARED_DEFAULTS = {
+    "height": 704,
+    "width": 1248,
+    "num_frames": 313,
+    "fps": 24,
+    "num_inference_steps": 50,
+    "guidance_scale": 2.5,
+    "negative_prompt": "",
+}
+
+PAN2_T2V = InferencePreset(
+    name="pan2_t2v",
+    version=1,
+    model_family="pan2",
+    description="PAN2 text-to-video",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+# The same pipeline runs image-to-video when the request carries a first frame (image_path).
+PAN2_I2V = InferencePreset(
+    name="pan2_i2v",
+    version=1,
+    model_family="pan2",
+    description="PAN2 image-to-video from a first frame",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+ALL_PRESETS = (PAN2_T2V, PAN2_I2V)
+
+__all__ = ["ALL_PRESETS", "PAN2_I2V", "PAN2_T2V"]

fastvideo/pipelines/basic/pan2/stages.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2-specific pipeline stages: prompt encoding, first-frame conditioning, sigma schedule and fp32 sampling."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+import PIL.Image
+import torch
+import torchvision.transforms.functional as TF
+from torch.distributed.tensor import DTensor
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.utils import PRECISION_TO_TYPE
+
+_USER_HEADER = "<|im_start|>user\n"
+
+
+class PAN2InputValidationStage(InputValidationStage):
+    """Rejects geometry the PAN2 VAE and patch embedding cannot represent."""
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        batch = super().forward(batch, fastvideo_args)
+        pipeline_config = fastvideo_args.pipeline_config
+        vae_arch = pipeline_config.vae_config.arch_config
+        spatial_multiple = vae_arch.spatial_compression_ratio * pipeline_config.dit_config.arch_config.patch_size[1]
+        if batch.height % spatial_multiple != 0 or batch.width % spatial_multiple != 0:
+            raise ValueError(f"PAN2 height and width must be divisible by {spatial_multiple}, "
+                             f"got {batch.height} and {batch.width}.")
+        if (batch.num_frames - 1) % vae_arch.temporal_compression_ratio != 0:
+            raise ValueError(f"PAN2 num_frames - 1 must be divisible by {vae_arch.temporal_compression_ratio}, "
+                             f"got {batch.num_frames}.")
+        if batch.pil_image is not None and not isinstance(batch.pil_image, PIL.Image.Image):
+            raise ValueError("PAN2 image-to-video expects the first frame as a PIL image or an image_path.")
+        return batch
+
+
+class PAN2TextEncodingStage(TextEncodingStage):
+    """Encodes a prompt with PAN2's chat template.
+
+    PAN2 conditions on the tokens after the ``<|im_start|>user\\n`` header (the prompt and the assistant header),
+    read from the text encoder's hidden states with its last ``text_encoder_skip_layers`` layers skipped.
+    """
+
+    @staticmethod
+    def tokenize(tokenizer: Any, prompt: str, system_prompt: str, max_length: int) -> tuple[torch.Tensor, int]:
+        """Token ids truncated to ``max_length`` tokens after the user header, and where that header ends."""
+        # An empty prompt is encoded as a single space, as in training.
+        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt or " "}]
+        input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+                                                  return_dict=True)["input_ids"]
+        if isinstance(input_ids, torch.Tensor):
+            input_ids = input_ids.tolist()
+        if input_ids and isinstance(input_ids[0], list):
+            input_ids = input_ids[0]
+        header_ids = tokenizer(_USER_HEADER, add_special_tokens=False)["input_ids"]
+        crop_start = next((i + len(header_ids)
+                           for i in range(len(input_ids)) if input_ids[i:i + len(header_ids)] == header_ids), None)
+        if crop_start is None:
+            raise ValueError(f"The PAN2 tokenizer's chat template has no {_USER_HEADER!r} header.")
+        return torch.tensor([input_ids[:crop_start + max_length]], dtype=torch.long), crop_start
+
+    @torch.no_grad()
+    def encode_text(
+        self,
+        text: str | list[str],
+        fastvideo_args: FastVideoArgs,
+        encoder_index: int | list[int] | None = None,
+        return_attention_mask: bool = False,
+        return_type: str = "list",
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+        max_length: int | Sequence[int | None] | None = None,
+        truncation: bool | None = None,
+        padding: bool | str | None = None,
+    ):
+        pipeline_config = fastvideo_args.pipeline_config
+        texts = [text] if isinstance(text, str) else list(text)
+        if len(texts) != 1:
+            raise ValueError(f"PAN2 encodes one prompt per request, got {len(texts)}.")
+        if isinstance(max_length, Sequence):
+            max_length = max_length[0]
+        max_length = max_length or pipeline_config.text_max_sequence_length
+
+        tokenizer = self.tokenizers[0]
+        text_encoder = self.text_encoders[0]
+        target_device = torch.device(device) if device is not None else get_local_torch_device()
+        # A module parked on CPU by text_encoder_cpu_offload computes on the target device and goes back afterwards;
+        # FSDP-wrapped (DTensor) encoders move their own shards.
+        first_param = next(text_encoder.parameters(), None)
+        input_device = first_param.device if first_param is not None else target_device
+        moved_for_forward = False
+        if (first_param is not None and not isinstance(first_param, DTensor)
+                and input_device.type != target_device.type):
+            text_encoder = text_encoder.to(target_device)
+            input_device = target_device
+            moved_for_forward = True
+
+        input_ids, crop_start = self.tokenize(tokenizer, texts[0], pipeline_config.text_system_prompt, max_length)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            outputs = text_encoder(input_ids=input_ids.to(input_device), output_hidden_states=True)
+        prompt_embeds = outputs.hidden_states[-(pipeline_config.text_encoder_skip_layers + 1)][:, crop_start:]
+        # fp32 is exact for the bf16 encoder output; the denoiser casts the embeddings to the DiT dtype.
+        prompt_embeds = prompt_embeds.to(device=target_device, dtype=dtype or torch.float32, copy=True)
+        attention_mask = torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=target_device)
+        if moved_for_forward and fastvideo_args.text_encoder_cpu_offload:
+            text_encoder.to("cpu")
+        return self.return_embeds([prompt_embeds], [attention_mask], return_type, return_attention_mask, [0])
+
+
+class PAN2TimestepPreparationStage(TimestepPreparationStage):
+    """PAN2 samples on uniformly spaced sigmas from 1 down to (excluding) 0; the scheduler applies the shift."""
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.timesteps is None and batch.sigmas is None:
+            batch.sigmas = np.linspace(1.0, 0.0, batch.num_inference_steps + 1)[:-1].tolist()
+        return super().forward(batch, fastvideo_args)
+
+
+class PAN2ConditioningStage(PipelineStage):
+    """Builds the conditioning latents and mask that PAN2 stacks after the noisy latents along channels.
+
+    Text-to-video conditions on zeros. Image-to-video fills the first latent frame with the normalized VAE latents
+    of the first frame, scaled to cover the output and center cropped, and sets that frame's mask channel to one.
+    """
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    @staticmethod
+    def preprocess_image(image: PIL.Image.Image, height: int, width: int) -> torch.Tensor:
+        """``[1, 3, 1, H, W]`` pixels in ``[-1, 1]``: scale to cover the canvas, round the size, center crop."""
+        image = image.convert("RGB")
+        scale = max(width / image.width, height / image.height)
+        image = image.resize((round(image.width * scale), round(image.height * scale)), resample=PIL.Image.LANCZOS)
+        left, top = round((image.width - width) / 2), round((image.height - height) / 2)
+        image = image.crop((left, top, left + width, top + height))
+        return (TF.to_tensor(image) * 2.0 - 1.0)[None, :, None]
+
+    def _encode_first_frame(self, image: PIL.Image.Image, batch: ForwardBatch,
+                            fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        device = get_local_torch_device()
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        self.vae = self.vae.to(device)
+        pixels = self.preprocess_image(image, batch.height, batch.width).to(device=device, dtype=vae_dtype)
+        latents = self.vae.encode(pixels).mode().float()
+        shift = self.vae.shift_factor.to(latents.device, latents.dtype)
+        scale = self.vae.scaling_factor.to(latents.device, latents.dtype)
+        if fastvideo_args.vae_cpu_offload:
+            self.vae.to("cpu")
+        return (latents - shift) * scale
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        latents = batch.latents
+        batch_size, num_channels, num_frames, height, width = latents.shape
+        condition = torch.zeros(batch_size,
+                                num_channels + 1,
+                                num_frames,
+                                height,
+                                width,
+                                device=latents.device,
+                                dtype=torch.float32)
+        if batch.pil_image is not None:
+            image_latents = self._encode_first_frame(batch.pil_image, batch, fastvideo_args)
+            condition[:, :num_channels, :1] = image_latents.to(condition.device)
+            condition[:, num_channels:, :1] = 1.0
+        batch.image_latent = condition
+        return batch
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)])
+        return result
+
+
+class PAN2DenoisingStage(DenoisingStage):
+    """Keeps the latents in fp32 across the loop; they are cast to the DiT dtype for each forward.
+
+    Image-to-video leaves the conditioning image's latent frame out of classifier-free guidance: that frame takes the
+    positive prediction, as in the reference sampler.
+    """
+
+    def finish_guidance(self, noise_pred, noise_pred_text, batch, fastvideo_args) -> torch.Tensor:
+        if batch.pil_image is not None and fastvideo_args.pipeline_config.i2v_skip_first_latent_cfg:
+            noise_pred[:, :, :1] = noise_pred_text[:, :, :1]
+        return noise_pred
+
+    def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor:
+        return noise_pred.float()

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)