diff --git a/docs-site/content/cookbook/pan/pan2.mdx b/docs-site/content/cookbook/pan/pan2.mdx
new file mode 100644
index 000000000..07836529a
--- /dev/null
+++ b/docs-site/content/cookbook/pan/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
+---
+
+
diff --git a/docs-site/lib/cookbook-pages.ts b/docs-site/lib/cookbook-pages.ts
index af4e8e077..9cadf1d6b 100644
--- a/docs-site/lib/cookbook-pages.ts
+++ b/docs-site/lib/cookbook-pages.ts
@@ -75,6 +75,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/pan/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",
@@ -120,6 +129,7 @@ const familyInfo: Record = {
"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 hiddenBackends = new Set(["fasth3-preview"]);
@@ -185,6 +196,7 @@ const mappedModels = [...new Set(sourceRecipes.map((recipe) => recipe.model))].m
kandinsky: "Kandinsky Lab",
"black-forest-labs": "Black Forest Labs",
qwen: "Qwen",
+ ifm: "IFM",
stabilityai: "Stability AI",
nvidia: "NVIDIA",
};
diff --git a/docs-site/navigation.json b/docs-site/navigation.json
index 8885794eb..b3794c9eb 100644
--- a/docs-site/navigation.json
+++ b/docs-site/navigation.json
@@ -131,6 +131,10 @@
],
"display": "group"
},
+ {
+ "label": "PAN 2 text/image to video",
+ "href": "/inference/pan2"
+ },
{
"label": "Examples",
"href": "/inference/examples/examples_inference_index"
@@ -253,6 +257,25 @@
}
]
},
+ {
+ "label": "PAN 2",
+ "display": "group",
+ "collapsed": true,
+ "items": [
+ {
+ "label": "PAN 2",
+ "href": "/cookbook/pan/pan2"
+ },
+ {
+ "label": "Family guide",
+ "href": "/cookbook/pan"
+ },
+ {
+ "label": "Model reference",
+ "href": "/inference/pan2"
+ }
+ ]
+ },
{
"label": "FLUX",
"href": "/cookbook/flux"
diff --git a/docs-site/tests/cookbook.spec.ts b/docs-site/tests/cookbook.spec.ts
index d1906df92..06e1cfe85 100644
--- a/docs-site/tests/cookbook.spec.ts
+++ b/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(38);
+ await expect(cards).toHaveCount(39);
await expect(cards.first()).toContainText("FastH3 Omni v1");
await page.getByLabel("Model family", { exact: true }).selectOption("Wan");
await page.getByLabel("Search recipes").fill("Wan2.1 I2V");
@@ -276,7 +276,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(38);
+ await expect(page.locator(".recipe-card")).toHaveCount(39);
});
test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => {
@@ -471,6 +471,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\/pan\/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));
@@ -789,7 +809,7 @@ test("MLX Omni discovery and its original route expose the MLX launch command",
test("category counts match visible cards and exclude retired Preview", async ({ page }) => {
await page.goto("/FastVideo/cookbook");
- await expect(page.locator(".recipe-card")).toHaveCount(38);
+ await expect(page.locator(".recipe-card")).toHaveCount(39);
for (const category of ["All models", "Video", "Image", "Audio", "World"]) {
const button = page.getByRole("button", { name: category, exact: true });
const count = Number(await button.locator(".recipe-filter-count").innerText());
diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json
index 0c13d83fd..df8be65b2 100644
--- a/docs/assets/cookbook-recipes.json
+++ b/docs/assets/cookbook-recipes.json
@@ -381,6 +381,32 @@
"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 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",
diff --git a/docs/assets/logos/SOURCES.md b/docs/assets/logos/SOURCES.md
index 0a2a1f9c5..9e6f78f82 100644
--- a/docs/assets/logos/SOURCES.md
+++ b/docs/assets/logos/SOURCES.md
@@ -13,6 +13,7 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo.
| `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header |
| `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers |
| `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header |
+| `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 |
| `pika.webp` | [Official Pika-Labs Hugging Face organization avatar](https://huggingface.co/Pika-Labs) | Pika H3 Flex card and page header |
diff --git a/docs/assets/logos/ifm.webp b/docs/assets/logos/ifm.webp
new file mode 100644
index 000000000..fe4602a9f
Binary files /dev/null and b/docs/assets/logos/ifm.webp differ
diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md
index e0dae6729..53f090b82 100644
--- a/docs/contributing/env_vars.md
+++ b/docs/contributing/env_vars.md
@@ -266,6 +266,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). |
@@ -339,6 +340,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`. |
diff --git a/docs/cookbook/index.md b/docs/cookbook/index.md
index 10a75432c..0011ce674 100644
--- a/docs/cookbook/index.md
+++ b/docs/cookbook/index.md
@@ -6,6 +6,7 @@ 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](pan.md)
- [FLUX](flux.md)
- [Qwen-Image-2.1](qwen-image21.md)
- [Stable Diffusion](stable-diffusion.md)
diff --git a/docs/cookbook/pan.md b/docs/cookbook/pan.md
new file mode 100644
index 000000000..4380cf9ef
--- /dev/null
+++ b/docs/cookbook/pan.md
@@ -0,0 +1,76 @@
+# PAN 2 recipes
+
+
+
+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, guidance 2.5 and PAN 2's default negative prompt. 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. Set `PAN2_MODEL_PATH` to run from a local copy of the
+checkpoint. PAN 2 runs with its prompt and video guardrails by default, so install
+them first (see [Guardrails](../inference/pan2.md#guardrails)):
+
+```bash
+pip install git+https://github.com/ifm-ai/pan2-guardrail
+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 by the PAN 2 parity tests.
diff --git a/docs/cookbook/recipes/pan2.yaml b/docs/cookbook/recipes/pan2.yaml
new file mode 100644
index 000000000..bfd1a0540
--- /dev/null
+++ b/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
diff --git a/docs/cookbook/serving-coverage.md b/docs/cookbook/serving-coverage.md
index 23280315f..2ba60a70d 100644
--- a/docs/cookbook/serving-coverage.md
+++ b/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; 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 |
diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml
index 678fff1f8..05d40405c 100644
--- a/docs/design/inference_schema_parity_inventory.yaml
+++ b/docs/design/inference_schema_parity_inventory.yaml
@@ -296,6 +296,33 @@ 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
+ vae_decode_kernel:
+ sources:
+ - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+ vae_decode_kernel_parallel:
+ sources:
+ - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+ dit_fused_kernels:
+ sources:
+ - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+ enable_guardrails:
+ sources:
+ - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
+ offload_guardrail_models:
+ sources:
+ - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig
text_encoder_max_lengths:
sources:
- fastvideo.configs.pipelines.hunyuan15.Hunyuan15I2V480PStepDistilledConfig
diff --git a/docs/generate_examples.py b/docs/generate_examples.py
index 8203bead7..f7ee16b5e 100644
--- a/docs/generate_examples.py
+++ b/docs/generate_examples.py
@@ -35,6 +35,7 @@ COOKBOOK_FAMILIES = {
"cosmos",
"kandinsky5",
"kandinsky6",
+ "pan2",
"flux",
"qwen_image21",
"sd35",
diff --git a/docs/inference/pan2.md b/docs/inference/pan2.md
new file mode 100644
index 000000000..ed3bb47ce
--- /dev/null
+++ b/docs/inference/pan2.md
@@ -0,0 +1,120 @@
+# 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`, 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 | 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, as in the reference sampler. It is used whenever the request's negative prompt is empty; a single space gives the empty unconditional prompt PAN 2 trained with |
+
+PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON: the style,
+lighting and camera, each subject, the first frame, and timed events. `basic_pan2.py` runs the release's own
+text-to-video caption, a 10 s clip. Short raw prompts, like the one in the snippet below, 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. As in the reference sampler, an `image_path` is read as its
+stored pixels: an EXIF orientation tag does not rotate it.
+
+## Guardrails
+
+PAN 2 runs with guardrails by default, as in its vLLM-Omni and Diffusers integrations. The
+[`pan2-guardrail`](https://github.com/ifm-ai/pan2-guardrail) package checks every prompt before any compute and every
+generated video after decoding. A blocked prompt or video raises `PAN2GuardrailViolationError`. Install the package
+before loading the pipeline:
+
+```bash
+pip install git+https://github.com/ifm-ai/pan2-guardrail
+```
+
+| Setting | Default | What it does |
+|---|---|---|
+| `enable_guardrails` | `True` | Pass `False` to `VideoGenerator.from_pretrained` to turn both checks off, at your own risk. Without the package, loading fails unless the guardrails are off. |
+| `offload_guardrail_models` | `False` | Keeps the guard models on CPU and moves each one to the GPU only for its check. |
+
+With sequence parallelism, the first GPU of the group loads the guard models and runs the checks, and every rank
+raises on its verdict. The video check reads the decoded pixels, so `output_type="latent"` is refused while the
+guardrails are on.
+
+## 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) against fastvideo-kernel's CUTLASS submodule (`git submodule update --init fastvideo-kernel/include/cutlass`, or set `CUTLASS_DIR`); each process builds in a directory of its own under `WAN22_VAE_DECODE_BUILD_DIR`, and later processes copy the build from the cache there. 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. Their rounding differs slightly from the unfused path. |
+
+On Blackwell GPUs (sm_100 and newer) with FlashAttention-4 (`flash_attn.cute`) installed, PAN2 also attends with FA4 unless `FASTVIDEO_FA4` is set: 288 s instead of 670 s with the Torch SDPA fallback for 241 frames on one B300. Set `FASTVIDEO_FA4=0` to keep the backend FastVideo picks otherwise. PAN2 sets the variable only while it loads, but the attention modules keep FA4 for the rest of the process, so a later pipeline in the same process also attends with it.
+
+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.
+
+## Differences from the vLLM-Omni integration
+
+- 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.
diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md
index e6b40458f..a6d68eee2 100644
--- a/docs/inference/support_matrix.md
+++ b/docs/inference/support_matrix.md
@@ -70,6 +70,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
| 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) [FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py) [Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
| pika_h3_flex | `Pika-Labs/Pika-H3-Flex` | T2V, I2V | [T2VA/I2VA/L2VA/FL2VA/Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pika_h3_flex.py) |
| open_h3_regenerate_2k | `Efficient-Large-Model/Open-H3-Regenerate-2K` | — | [basic_open_h3_regenerate_2k.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_open_h3_regenerate_2k.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) |
diff --git a/examples/inference/basic/basic_pan2.py b/examples/inference/basic/basic_pan2.py
new file mode 100644
index 000000000..2a2f1ee8a
--- /dev/null
+++ b/examples/inference/basic/basic_pan2.py
@@ -0,0 +1,150 @@
+import json
+import os
+
+from fastvideo import VideoGenerator
+
+# PAN2 is prompted with a structured caption, serialized as one-line JSON: the style, lighting and camera, each
+# subject, the first frame, and timed events with their camera motion and actions. Its prompt rewriter
+# (`inference/pan-prompt-rewriter` in ifm-ai/PAN2) turns a plain description into one. This is the release's own
+# text-to-video example (`inference/prompts/examples.json`), a 10 s clip whose event times span the video.
+CAPTION = json.loads("""
+{
+ "video_style": "Realistic sports cinematography in one uninterrupted continuous shot.",
+ "lighting": "Direct morning sunlight from the pool's far side; bright surface reflections, natural skin color and blue-green underwater light with moving caustics.",
+ "perspective_view": "Low side three-quarter view from an adjacent empty lane, tracking right alongside the swimmer. The waterproof camera descends through the surface after entry, follows underwater, then rises above water while maintaining the same side and direction.",
+ "subjects_lifecycle": [
+ {
+ "subject_id": "<|Subject1|>",
+ "display_name": "Racing swimmer",
+ "subject_attribute": "Athletic man in his early 30s, clean-shaven, wearing a white silicone cap, dark racing goggles and navy knee-length racing jammers."
+ }
+ ],
+ "init_frame_visual_description": "The swimmer crouches on a white starting block at the left of an outdoor competition pool, facing right along his lane. His fingers curl over the block's front edge, hips raised and feet staggered. The low camera sits just above the water in the neighboring lane, framing his entire starting posture and several meters of clear water ahead. Blue-and-white lane ropes converge toward the distant end wall. Sunlight defines his shoulders and the matte fabric of his jammers; small surface ripples break the reflection of the pale pool deck, with a dark floor guide visible beneath them.",
+ "events": [
+ {
+ "event_index": 0,
+ "time_range_sec": "0 to 2.2",
+ "event_description": "He drives off the block, extends his arms beyond his head and enters fingertips first farther down the lane. His torso and pointed feet follow through the opening, leaving a brief upright splash behind him. He lengthens into a submerged glide as the camera follows right and descends through the rippling surface beside his entry trail.",
+ "camera_motion": [
+ {
+ "time_range_sec": "0 to 0.3",
+ "camera_motion": [
+ "static"
+ ]
+ },
+ {
+ "time_range_sec": "0.3 to 1.5",
+ "camera_motion": [
+ "track_right"
+ ]
+ },
+ {
+ "time_range_sec": "1.5 to 2.2",
+ "camera_motion": [
+ "track_right",
+ "dolly_down"
+ ]
+ }
+ ],
+ "contained_actions": [
+ {
+ "time_range_sec": "0.2 to 1.5",
+ "action": "<|Subject1|> pushes from the block, dives forward with arms extended and enters the water hands first, followed by torso and feet."
+ },
+ {
+ "time_range_sec": "1.5 to 2.2",
+ "action": "<|Subject1|> glides underwater in a long streamlined shape, carrying forward momentum from the dive."
+ }
+ ]
+ },
+ {
+ "event_index": 1,
+ "time_range_sec": "2.2 to 5.8",
+ "event_description": "Seen fully underwater, he adds two compact dolphin kicks, then angles toward the surface. One arm separates into the first pull as his legs switch to a fast flutter kick. His head breaks the surface and he establishes front crawl; from below, alternating forearms press water backward while small bubbles trail his hands and shoulders. The floor guide passes beneath him as the camera matches his forward speed.",
+ "camera_motion": [
+ {
+ "time_range_sec": "2.2 to 5.8",
+ "camera_motion": [
+ "track_right"
+ ]
+ }
+ ],
+ "contained_actions": [
+ {
+ "time_range_sec": "2.2 to 3.5",
+ "action": "<|Subject1|> performs two dolphin kicks during the glide and directs his streamlined body upward."
+ },
+ {
+ "time_range_sec": "3.5 to 5.8",
+ "action": "<|Subject1|> breaks out into front crawl, alternating underwater pulls and changing to a continuous flutter kick."
+ }
+ ]
+ },
+ {
+ "event_index": 2,
+ "time_range_sec": "5.8 to 10",
+ "event_description": "The camera rises through the waterline, revealing his recovering arms above the surface while keeping his advancing body at useful scale. He powers onward with alternating reaches and pulls, each hand entry producing a short splash in a regular rhythm. He rolls his head toward the camera for one breath, returns his face to the water and drives into the next strokes without slowing.",
+ "camera_motion": [
+ {
+ "time_range_sec": "5.8 to 6.6",
+ "camera_motion": [
+ "track_right",
+ "dolly_up"
+ ]
+ },
+ {
+ "time_range_sec": "6.6 to 10",
+ "camera_motion": [
+ "track_right"
+ ]
+ }
+ ],
+ "contained_actions": [
+ {
+ "time_range_sec": "5.8 to 10",
+ "action": "<|Subject1|> races forward in front crawl with alternating arm recoveries and rhythmic hand entries, takes one side breath and returns his face to the water while maintaining his kick."
+ }
+ ]
+ }
+ ],
+ "quality": {
+ "video_quality": "Resolved athletic form and transparent water, with brief motion blur on fast hands and splashes.",
+ "text_logos_overlays": null,
+ "artifacts": null
+ }
+}
+""")
+PROMPT = json.dumps(CAPTION, ensure_ascii=False, separators=(",", ":"))
+NUM_FRAMES = 241 # 10 s at 24 fps, the caption's length.
+
+OUTPUT_PATH = "video_samples_pan2"
+MODEL_PATH = os.environ.get("PAN2_MODEL_PATH", "IFM/PAN2")
+
+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,
+ num_frames=NUM_FRAMES,
+ output_path=OUTPUT_PATH,
+ save_video=True,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/serving/README.md b/examples/serving/README.md
index b95910345..7cbd01ec3 100644
--- a/examples/serving/README.md
+++ b/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 |
diff --git a/examples/serving/openai_pan2.yaml b/examples/serving/openai_pan2.yaml
new file mode 100644
index 000000000..9f44b4da5
--- /dev/null
+++ b/examples/serving/openai_pan2.yaml
@@ -0,0 +1,38 @@
+# 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:
+ # Empty: PAN 2's default negative prompt, its time ranges set to the clip's duration.
+ negative_prompt: ""
+ sampling:
+ height: 704
+ width: 1248
+ num_frames: 313
+ fps: 24
+ num_inference_steps: 50
+ guidance_scale: 2.5
+ output:
+ return_frames: false
diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md
index f332e2fed..dd19a5a79 100644
--- a/fastvideo-kernel/README.md
+++ b/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
diff --git a/fastvideo-kernel/tests/test_wan22_vae_decode.py b/fastvideo-kernel/tests/test_wan22_vae_decode.py
new file mode 100644
index 000000000..e1d115bcd
--- /dev/null
+++ b/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)
diff --git a/fastvideo-kernel/wan22_vae_decode/README.md b/fastvideo-kernel/wan22_vae_decode/README.md
new file mode 100644
index 000000000..69480ae3b
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/README.md
@@ -0,0 +1,99 @@
+# 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`). Each process builds in a directory of its own, so concurrent processes do
+ not collide, and puts the finished build in the `cache/` beside it, from which later processes copy it.
+- `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_.pt`), the weights in `weight_names()` order
+(`decoder_weights_bf16.pt`, `weight_names.json`), the fp32 decodes (`ref_.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.
diff --git a/fastvideo-kernel/wan22_vae_decode/notes/optimization.md b/fastvideo-kernel/wan22_vae_decode/notes/optimization.md
new file mode 100644
index 000000000..fcdae9f2b
--- /dev/null
+++ b/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).
diff --git a/fastvideo-kernel/wan22_vae_decode/reference.py b/fastvideo-kernel/wan22_vae_decode/reference.py
new file mode 100644
index 000000000..cc036fa29
--- /dev/null
+++ b/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)
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/attn.cu b/fastvideo-kernel/wan22_vae_decode/solution/attn.cu
new file mode 100644
index 000000000..4559bc4e3
--- /dev/null
+++ b/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
+#include
+#include
+#include
+#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
+static void gemm1x1(const at::Tensor& x, const at::Tensor& w, const at::Tensor& y) {
+ using Tile = Shape, Int, Shape<_64>>;
+ using Cluster = Shape, _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(sizeof(typename Epilogue::SharedStorage))>,
+ cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+ using Problem = cutlass::conv::ConvProblemShape;
+ using Kernel = cutlass::conv::kernel::ConvUniversal;
+ using Conv = cutlass::conv::device::ConvUniversalAdapter;
+ 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(sc)>{}, [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
+ cute::for_each(cute::make_seq(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<<>>(
+ 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(x, w, y);
+ }
+
+ float* s_p = (float*)s.data_ptr();
+ float scale = 1.f / sqrtf((float)C);
+ row_softmax_kernel<<>>(
+ 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(x, w, y);
+ }
+ return o;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/attn.h b/fastvideo-kernel/wan22_vae_decode/solution/attn.h
new file mode 100644
index 000000000..de1df1567
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/solution/attn.h
@@ -0,0 +1,6 @@
+#pragma once
+#include
+// 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);
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/binding.cpp b/fastvideo-kernel/wan22_vae_decode/solution/binding.cpp
new file mode 100644
index 000000000..a8c59c67b
--- /dev/null
+++ b/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
+#include
+#include
+#include
+#include
+#include
+#include "kernels.h"
+#include "prefix.h"
+#include "comm.h"
+#include "resample.h"
+#include "attn.h"
+#include "head.h"
+#include
+#include
+#include
+#include
+
+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;
+using at::Tensor;
+using Weights = std::vector;
+
+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{72,70},std::pair{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 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());
+ 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_trim1&&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());
+ 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 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 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());
+ 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());
+ 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());
+ 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());
+ // 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());
+ 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 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=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(frames_per_pass*world,kMaxChunk);
+ int T=z.size(2);
+ for(int start=0;start1) 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 buffers={part};
+ pg->recv(buffers,peer,0)->wait();
+ write_video(part,out,frame_offset,lo*16);
+ }
+ } else {
+ std::vector 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());
+ m.def("run_with",&decode_video,py::call_guard());
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/comm.cu b/fastvideo-kernel/wan22_vae_decode/solution/comm.cu
new file mode 100644
index 000000000..7abe62072
--- /dev/null
+++ b/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
+#include
+#include
+#include
+#include
+
+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<<>>((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<<>>((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 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;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/comm.h b/fastvideo-kernel/wan22_vae_decode/solution/comm.h
new file mode 100644
index 000000000..d1824eb60
--- /dev/null
+++ b/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
+#include
+#include
+#include
+#include
+#include
+#include
+
+using PG = c10::intrusive_ptr;
+
+// 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> works;
+ std::vector 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);
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/conv.cu b/fastvideo-kernel/wan22_vae_decode/solution/conv.cu
new file mode 100644
index 000000000..874c766e3
--- /dev/null
+++ b/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
+#include
+#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
+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,Shape<_64>>;
+ using Cluster=Shape,_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(sizeof(typename Epilogue::SharedStorage))>,
+ cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+ using Problem=cutlass::conv::ConvProblemShape;
+ using Kernel=cutlass::conv::kernel::ConvUniversal;
+ using Conv=cutlass::conv::device::ConvUniversalAdapter;
+ 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(sc)>{},[&](auto i){cute::get<0,i>(sc)=problem.stride_C[Problem::RankT-2-i];});
+ cute::for_each(cute::make_seq(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);
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/head.cu b/fastvideo-kernel/wan22_vae_decode/solution/head.cu
new file mode 100644
index 000000000..cf9443c31
--- /dev/null
+++ 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
+#include
+#include
+#include
+#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<<>>(
+ (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();
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/head.h b/fastvideo-kernel/wan22_vae_decode/solution/head.h
new file mode 100644
index 000000000..0a76b4efb
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/solution/head.h
@@ -0,0 +1,19 @@
+#pragma once
+#include
+
+// 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);
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/kernels.cu b/fastvideo-kernel/wan22_vae_decode/solution/kernels.cu
new file mode 100644
index 000000000..684ea67a6
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/solution/kernels.cu
@@ -0,0 +1,99 @@
+#include "kernels.h"
+#include
+#include
+#include
+#include
+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() : nullptr; }
+static B* ptr(Tensor& t) { return (B*)t.data_ptr(); }
+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>>(ptr(x),ptr(b),ptr(s),n,C);
+ else finish_scalar<<>>(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>>(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=0&&y=0&&xx>>(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.)
+
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/kernels.h b/fastvideo-kernel/wan22_vae_decode/solution/kernels.h
new file mode 100644
index 000000000..ff1efda7b
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/solution/kernels.h
@@ -0,0 +1,21 @@
+#pragma once
+#include
+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& 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);
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/norm.cu b/fastvideo-kernel/wan22_vae_decode/solution/norm.cu
new file mode 100644
index 000000000..6e0bbd124
--- /dev/null
+++ b/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 / 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
+#include
+#include
+#include
+#include
+#include
+#include
+
+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
+__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();
+ const B* gp = gamma.defined() ? (const B*)gamma.data_ptr() : nullptr;
+ const B* hp = hist.defined() ? (const B*)hist.data_ptr() : nullptr;
+ B* yp = (B*)y.data_ptr();
+ const int64_t rows = (int64_t)(T + pt) * (H + ph * 2);
+ auto stream = at::cuda::getCurrentCUDAStream();
+ const int blocks = (int)std::min(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<<>>(xp, gp, hp, yp, T, H, W, pt, ph);
+ else if (silu) norm_pad_kernel<<>>(xp, gp, hp, yp, T, H, W, pt, ph);
+ else norm_pad_kernel<<>>(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{});
+ else if (C == 512) go(std::integral_constant{});
+ else go(std::integral_constant{});
+ } else {
+ TORCH_CHECK(C % 8 == 0, "norm_pad copy needs C divisible by 8, got ", C);
+ copy_row_kernel<<>>((const uint4*)xp, (const uint4*)hp, (uint4*)yp, T, H, W,
+ pt, ph, C / 8);
+ }
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+ return y;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/prefix.cu b/fastvideo-kernel/wan22_vae_decode/solution/prefix.cu
new file mode 100644
index 000000000..00a9b9630
--- /dev/null
+++ b/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
+#include
+#include
+
+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& 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<<>>(
+ (const uint4*)local.data_ptr(),(uint4*)send.data_ptr(),T,H,W*(C/8),world);
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+ }
+ std::vector in_split(world),out_split(world);
+ int64_t sent=0,received=0;
+ for(int p=0;palltoall_base(recv,send,out_split,in_split)->wait();
+ return output;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/prefix.h b/fastvideo-kernel/wan22_vae_decode/solution/prefix.h
new file mode 100644
index 000000000..5815a99f1
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/solution/prefix.h
@@ -0,0 +1,7 @@
+#pragma once
+#include
+#include
+// 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& pg);
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/resample.cu b/fastvideo-kernel/wan22_vae_decode/solution/resample.cu
new file mode 100644
index 000000000..896b99003
--- /dev/null
+++ b/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
+#include
+#include
+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(); }
+static B* pm(Tensor& t) { return (B*)t.data_ptr(); }
+
+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& 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 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<<>>(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;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/resample.h b/fastvideo-kernel/wan22_vae_decode/solution/resample.h
new file mode 100644
index 000000000..4f9be646c
--- /dev/null
+++ b/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
+#include
+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& pw, const Tensor& bias,
+ const Tensor& skip, bool halo, bool temporal, bool first);
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/solution.json b/fastvideo-kernel/wan22_vae_decode/solution/solution.json
new file mode 100644
index 000000000..9f1850549
--- /dev/null
+++ b/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":[]}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/timeconv.cu b/fastvideo-kernel/wan22_vae_decode/solution/timeconv.cu
new file mode 100644
index 000000000..cb0f2ece8
--- /dev/null
+++ b/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
+#include
+#include
+#include
+#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
+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, Shape<_64>>;
+ using Cluster = Shape, _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(sizeof(typename Epilogue::SharedStorage))>,
+ cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+ using Problem = cutlass::conv::ConvProblemShape;
+ using Kernel = cutlass::conv::kernel::ConvUniversal;
+ using Conv = cutlass::conv::device::ConvUniversalAdapter;
+ 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(sc)>{},
+ [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
+ cute::for_each(cute::make_seq(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;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/wino_epilogue.cu b/fastvideo-kernel/wan22_vae_decode/solution/wino_epilogue.cu
new file mode 100644
index 000000000..c94ce25fb
--- /dev/null
+++ b/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
+#include
+#include
+#include
+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():nullptr;}
+static B* e_mut(Tensor& t){return t.defined()?(B*)t.data_ptr():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
+__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
+__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(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>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>2), xx=4*tx+(p&3);
+ const bool ok=EXACT?true:(yy>=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
+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;
+ 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<<>>(p,y,n,b,r,g,H,W,TH,TW,PT,PH,wy);
+}
+
+WinoOut wino_epilogue(const std::vector& 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();
+ 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;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/solution/wino_in.cu b/fastvideo-kernel/wan22_vae_decode/solution/wino_in.cu
new file mode 100644
index 000000000..0b73fc350
--- /dev/null
+++ b/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
+#include
+#include
+#include
+
+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
+__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();
+ B* vp = (B*)v.data_ptr();
+ 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<<>>(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<<>>(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<<>>(xp, vp, C, D, H, W, ph, TH, TW);
+ break;
+ }
+ }
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+ return v;
+}
diff --git a/fastvideo-kernel/wan22_vae_decode/tools/audit.py b/fastvideo-kernel/wan22_vae_decode/tools/audit.py
new file mode 100644
index 000000000..4d3026938
--- /dev/null
+++ b/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()
diff --git a/fastvideo-kernel/wan22_vae_decode/tools/bench.py b/fastvideo-kernel/wan22_vae_decode/tools/bench.py
new file mode 100644
index 000000000..93826a539
--- /dev/null
+++ b/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 x. 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_.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()
diff --git a/fastvideo-kernel/wan22_vae_decode/tools/bounds.py b/fastvideo-kernel/wan22_vae_decode/tools/bounds.py
new file mode 100644
index 000000000..caad07537
--- /dev/null
+++ b/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()
diff --git a/fastvideo-kernel/wan22_vae_decode/tools/loader.py b/fastvideo-kernel/wan22_vae_decode/tools/loader.py
new file mode 100644
index 000000000..3c33506d1
--- /dev/null
+++ b/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`: "::", 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 "::"')
+ 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()
diff --git a/fastvideo-kernel/wan22_vae_decode/tools/spec.py b/fastvideo-kernel/wan22_vae_decode/tools/spec.py
new file mode 100644
index 000000000..f995bbe5b
--- /dev/null
+++ b/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&` (from ).
+- 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 x. A clip is the first T latent frames of one of
+three real 20 s clips encoded by the VAE (latents_.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_.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_.pt ({"latents": fp32 [1, 48, 121, 44, 78]}), ref_.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")
diff --git a/fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py b/fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py
new file mode 100644
index 000000000..0f17a43e9
--- /dev/null
+++ b/fastvideo-kernel/wan22_vae_decode/wan22_vae_decode.py
@@ -0,0 +1,152 @@
+"""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 shutil
+import socket
+import sys
+import tempfile
+
+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, in a build directory of the process's own under
+ WAN22_VAE_DECODE_BUILD_DIR (default ~/.cache/fastvideo/wan22_vae_decode), removed when the process exits.
+
+ Processes that build at once each write only their own directory. Finished builds go to the cache beside them,
+ keyed by the sources, flags and toolchain and added by an atomic rename, so later processes copy one instead of
+ building again."""
+ loader, _ = _tools()
+ base = os.environ.get("WAN22_VAE_DECODE_BUILD_DIR", os.path.expanduser("~/.cache/fastvideo/wan22_vae_decode"))
+ cache = os.path.join(base, "cache")
+ os.makedirs(cache, exist_ok=True)
+ build_dir = tempfile.mkdtemp(prefix=f"build-{socket.gethostname()}-{os.getpid()}-", dir=base)
+ atexit.register(shutil.rmtree, build_dir, True)
+ return loader.load(os.path.join(HERE, "solution"), build_dir=build_dir, chdir=False, cache=cache).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
diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py
index ab4a4e0f7..4490c2cc5 100644
--- a/fastvideo/configs/models/dits/__init__.py
+++ b/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,
@@ -32,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", "QwenImage21ArchConfig", "QwenImage21Config"
]
diff --git a/fastvideo/configs/models/dits/pan2.py b/fastvideo/configs/models/dits/pan2.py
new file mode 100644
index 000000000..8ec5b6d03
--- /dev/null
+++ b/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"
diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py
index d2797f2db..67f442a91 100644
--- a/fastvideo/configs/models/encoders/__init__.py
+++ b/fastvideo/configs/models/encoders/__init__.py
@@ -16,6 +16,7 @@ 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,
@@ -52,6 +53,8 @@ __all__ = [
"Qwen3TextConfig",
"QwenImage21Qwen3VLArchConfig",
"QwenImage21Qwen3VLConfig",
+ "Qwen3_5TextArchConfig",
+ "Qwen3_5TextConfig",
"Mistral3TextConfig",
"LingBotWorld2UMT5ArchConfig",
"LingBotWorld2UMT5Config",
diff --git a/fastvideo/configs/models/encoders/qwen3_5.py b/fastvideo/configs/models/encoders/qwen3_5.py
new file mode 100644
index 000000000..f469657ab
--- /dev/null
+++ b/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"
diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py
index f42e517b8..a925d5123 100644
--- a/fastvideo/configs/pipelines/__init__.py
+++ b/fastvideo/configs/pipelines/__init__.py
@@ -15,6 +15,7 @@ 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,
@@ -27,5 +28,5 @@ __all__ = [
"Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
"Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
"LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
- "QwenImage21PipelineConfig", "get_pipeline_config_cls_from_name"
+ "QwenImage21PipelineConfig", "PAN2PipelineConfig", "get_pipeline_config_cls_from_name"
]
diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py
index 7faf283b4..c76ff8dac 100644
--- a/fastvideo/configs/pipelines/base.py
+++ b/fastvideo/configs/pipelines/base.py
@@ -2,7 +2,7 @@
import json
from collections.abc import Callable
from dataclasses import asdict, dataclass, field, fields
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, ClassVar, cast
import torch
@@ -76,6 +76,8 @@ class PipelineConfig:
image_encoder_precisions: tuple[str, ...] | None = None
# Text encoder configuration
+ # This is a model-fixed normalization policy, not a serialized request knob.
+ preserve_single_space_negative_prompt: ClassVar[bool] = False
DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32", )
text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (EncoderConfig(), ))
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
diff --git a/fastvideo/configs/pipelines/pan2.py b/fastvideo/configs/pipelines/pan2.py
new file mode 100644
index 000000000..70a8e0e2c
--- /dev/null
+++ b/fastvideo/configs/pipelines/pan2.py
@@ -0,0 +1,66 @@
+# 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 typing import ClassVar
+
+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, the reference sampler's
+ # fixed value. A fixed default here would override the shift the 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
+
+ # Check the prompt and every generated video with the PAN 2 guardrails (the `pan2-guardrail` package), as the
+ # vLLM-Omni and Diffusers integrations do by default. A blocked prompt or video raises
+ # `PAN2GuardrailViolationError`. False turns both checks off, at your own risk.
+ enable_guardrails: bool = True
+ # Keep the guard models on CPU and move each to the GPU only for its check.
+ offload_guardrail_models: bool = False
+
+ # A single space is PAN2's trained unconditional prompt; normalization must not turn it into an empty string.
+ preserve_single_space_negative_prompt: ClassVar[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:
+ self.vae_config.load_encoder = True
+ self.vae_config.load_decoder = True
+ self.text_encoder_configs[0].arch_config.output_hidden_states = True
diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py
index dc52f444c..0676f4e88 100644
--- a/fastvideo/entrypoints/video_generator.py
+++ b/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:
@@ -861,7 +886,10 @@ class VideoGenerator:
sampling_param.prompt = prompt
# Process negative prompt
if sampling_param.negative_prompt is not None:
- sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
+ keep_single_space = (sampling_param.negative_prompt == " " and getattr(
+ fastvideo_args.pipeline_config, "preserve_single_space_negative_prompt", False))
+ if not keep_single_space:
+ sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
# Validate dimensions
if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0):
@@ -1064,11 +1092,7 @@ class VideoGenerator:
and src_frames.shape[0] == 1 and src_frames.shape[2] > 1):
frames = list(src_frames[0].permute(1, 2, 3, 0).contiguous().cpu().numpy())
else:
- vid_u8 = _quantize_video_frames_to_uint8(src_frames)
- frames = [
- torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy()
- for x in vid_u8
- ]
+ frames = _video_frames(src_frames)
postprocess_time = time.perf_counter() - postprocess_start
logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
if logging_info is not None:
diff --git a/fastvideo/envs.py b/fastvideo/envs.py
index 9eea7dffb..ae0bd90f0 100644
--- a/fastvideo/envs.py
+++ b/fastvideo/envs.py
@@ -677,6 +677,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 "
@@ -993,6 +999,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.",
diff --git a/fastvideo/models/dits/pan2.py b/fastvideo/models/dits/pan2.py
new file mode 100644
index 000000000..157ce7e62
--- /dev/null
+++ b/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
diff --git a/fastvideo/models/dits/pan2_fusions.py b/fastvideo/models/dits/pan2_fusions.py
new file mode 100644
index 000000000..2785efe76
--- /dev/null
+++ b/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",
+]
diff --git a/fastvideo/models/encoders/qwen3_5.py b/fastvideo/models/encoders/qwen3_5.py
new file mode 100644
index 000000000..c036620fd
--- /dev/null
+++ b/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
diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py
index 897e3f276..de6464b31 100644
--- a/fastvideo/models/registry.py
+++ b/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 = {
@@ -91,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] = {
diff --git a/fastvideo/models/vision_utils.py b/fastvideo/models/vision_utils.py
index feffe2178..dd10c93d3 100644
--- a/fastvideo/models/vision_utils.py
+++ b/fastvideo/models/vision_utils.py
@@ -124,7 +124,8 @@ def _fetch_image_bytes(url: str, attempts: int = 3, timeout: float = 30.0) -> by
# adapted from diffusers.utils import load_image
def load_image(image: str | PIL.Image.Image,
- convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None) -> PIL.Image.Image:
+ convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None,
+ exif_transpose: bool = True) -> PIL.Image.Image:
"""
Loads `image` to a PIL Image.
@@ -134,6 +135,8 @@ def load_image(image: str | PIL.Image.Image,
convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*):
A conversion method to apply to the image after loading it. When set to `None` the image will be converted
"RGB".
+ exif_transpose (`bool`, *optional*, defaults to `True`):
+ Whether to rotate the image upright by its EXIF orientation tag. `False` keeps the stored pixels.
Returns:
`PIL.Image.Image`:
@@ -154,7 +157,8 @@ def load_image(image: str | PIL.Image.Image,
raise ValueError(
"Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image.")
- image = PIL.ImageOps.exif_transpose(image)
+ if exif_transpose:
+ image = PIL.ImageOps.exif_transpose(image)
image = convert_method(image) if convert_method is not None else image.convert("RGB")
diff --git a/fastvideo/models/wan/vae.py b/fastvideo/models/wan/vae.py
index e26cd8955..46a82b4b5 100644
--- a/fastvideo/models/wan/vae.py
+++ b/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):
diff --git a/fastvideo/models/wan/vae_config.py b/fastvideo/models/wan/vae_config.py
index 8c6e67955..89935b69d 100644
--- a/fastvideo/models/wan/vae_config.py
+++ b/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
diff --git a/fastvideo/pipelines/basic/pan2/__init__.py b/fastvideo/pipelines/basic/pan2/__init__.py
new file mode 100644
index 000000000..988131360
--- /dev/null
+++ b/fastvideo/pipelines/basic/pan2/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/fastvideo/pipelines/basic/pan2/guardrails.py b/fastvideo/pipelines/basic/pan2/guardrails.py
new file mode 100644
index 000000000..a93a5386c
--- /dev/null
+++ b/fastvideo/pipelines/basic/pan2/guardrails.py
@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 guardrails: the prompt and every generated video are checked with the ``pan2_guardrail`` package's
+``PAN2SafetyChecker``, as in the vLLM-Omni and Diffusers integrations.
+
+They are on by default (``PAN2PipelineConfig.enable_guardrails``). Either check blocks the request with a
+``PAN2GuardrailViolationError``.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from typing import Any, TypeVar, cast
+
+import numpy as np
+import torch
+
+from fastvideo.distributed import get_local_torch_device, get_sp_group, model_parallel_is_initialized
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.utils import pixels_to_uint8
+
+logger = init_logger(__name__)
+
+T = TypeVar("T")
+
+_DISABLE_HINT = "or turn the guardrails off with `enable_guardrails=False` at your own risk."
+
+
+class PAN2GuardrailViolationError(ValueError):
+ """A prompt or a generated video was blocked by the PAN 2 guardrails."""
+
+
+def _sp_rank_and_size() -> tuple[int, int]:
+ if not model_parallel_is_initialized():
+ return 0, 1
+ group = get_sp_group()
+ return group.rank_in_group, group.world_size
+
+
+def _on_first_rank(fn: Callable[[], T], what: str) -> T:
+ """Runs ``fn`` on the first rank of the sequence-parallel group and returns its result on every rank.
+
+ A failure on the first rank is broadcast too, so the other ranks raise with it at once instead of waiting at their
+ next collective until it times out.
+ """
+ rank, world_size = _sp_rank_and_size()
+ outcome: tuple[bool, Any] | None = None
+ if rank == 0:
+ try:
+ outcome = (True, fn())
+ except BaseException as e:
+ if world_size > 1:
+ get_sp_group().broadcast_object((False, f"{type(e).__name__}: {e}"), src=0)
+ raise
+ if world_size > 1:
+ outcome = get_sp_group().broadcast_object(outcome, src=0)
+ assert outcome is not None
+ ok, value = outcome
+ if not ok:
+ raise RuntimeError(f"The first sequence-parallel rank failed to {what}: {value}")
+ return cast(T, value)
+
+
+class PAN2Guardrails:
+ """The checker, loaded on the first rank of each sequence-parallel group only.
+
+ That rank runs every check and broadcasts its verdict, so the ranks of a group raise together and the guard models
+ take memory on one GPU, not on all of them.
+ """
+
+ def __init__(self, offload_to_cpu: bool = False) -> None:
+ if importlib.util.find_spec("pan2_guardrail") is None:
+ raise ImportError("PAN 2 runs with guardrails by default, but the pan2-guardrail package is not installed. "
+ "Install it with `pip install git+https://github.com/ifm-ai/pan2-guardrail`, " +
+ _DISABLE_HINT)
+ self.offload_to_cpu = offload_to_cpu
+ self.checker = None
+
+ def load() -> None:
+ from pan2_guardrail import PAN2SafetyChecker
+
+ logger.info("Loading the PAN 2 guardrails (offload_to_cpu=%s)", offload_to_cpu)
+ # The checker loads on CPU. With offloading, each guard moves to the GPU only for its check.
+ self.checker = PAN2SafetyChecker.from_pretrained()
+ if not offload_to_cpu:
+ self.checker.to(get_local_torch_device())
+
+ _on_first_rank(load, "load the guardrails")
+
+ @contextmanager
+ def _on_device(self, module: torch.nn.Module) -> Iterator[None]:
+ if not self.offload_to_cpu:
+ yield
+ return
+ module.to(get_local_torch_device())
+ try:
+ yield
+ finally:
+ module.to("cpu")
+ # Return the guard's cached blocks, so the offloaded guard does not keep GPU memory between checks.
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+ def _passes(self, check: Callable[[], bool]) -> bool:
+ return bool(_on_first_rank(check, "run the guardrail check"))
+
+ def check_prompt(self, prompt: str) -> None:
+
+ def check() -> bool:
+ with self._on_device(self.checker.text_guard):
+ return self.checker.check_text_safety(prompt)
+
+ if not self._passes(check):
+ # PAN2SafetyChecker logs the reason at CRITICAL.
+ raise PAN2GuardrailViolationError("The prompt was blocked by the PAN 2 guardrails.")
+
+ def check_video(self, frames: Callable[[], np.ndarray]) -> None:
+ """``frames`` builds the uint8 ``[T, H, W, 3]`` frames, on the checking rank only."""
+
+ def check() -> bool:
+ with self._on_device(self.checker.video_guard):
+ # `check_video_safety` returns the unchanged frames when they pass and None when the video is blocked.
+ return self.checker.check_video_safety(frames()) is not None
+
+ if not self._passes(check):
+ raise PAN2GuardrailViolationError("The generated video was blocked by the PAN 2 guardrails.")
+
+
+class PAN2PromptGuardrailStage(PipelineStage):
+ """Checks the prompt before any compute. Latent output, which the video check cannot see, is refused here too."""
+
+ def __init__(self, guardrails: PAN2Guardrails) -> None:
+ super().__init__()
+ self.guardrails = guardrails
+
+ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+ if fastvideo_args.output_type == "latent":
+ raise ValueError("The PAN 2 guardrails check the decoded video, so latent output is not available with "
+ "them; " + _DISABLE_HINT)
+ prompts = [batch.prompt] if isinstance(batch.prompt, str) else list(batch.prompt or [])
+ for prompt in prompts:
+ self.guardrails.check_prompt(prompt)
+ return batch
+
+
+class PAN2VideoGuardrailStage(PipelineStage):
+ """Checks every decoded video, as the uint8 frames the generator delivers."""
+
+ def __init__(self, guardrails: PAN2Guardrails) -> None:
+ super().__init__()
+ self.guardrails = guardrails
+
+ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+ # `batch.output` is `[B, 3, T, H, W]` in [0, 1]; `pixels_to_uint8` quantizes as the generator does.
+ for video in batch.output:
+
+ def frames(video: torch.Tensor = video) -> np.ndarray:
+ return pixels_to_uint8(video).permute(1, 2, 3, 0).cpu().numpy()
+
+ self.guardrails.check_video(frames)
+ return batch
diff --git a/fastvideo/pipelines/basic/pan2/negative_prompt.py b/fastvideo/pipelines/basic/pan2/negative_prompt.py
new file mode 100644
index 000000000..de2f68307
--- /dev/null
+++ b/fastvideo/pipelines/basic/pan2/negative_prompt.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The negative prompt PAN2 was released with (``inference/prompts/negative_prompt.json`` in ifm-ai/PAN2).
+
+Its ``time_range_sec`` fields are a ``0 to 10`` template; ``render_pan2_default_negative_prompt`` sets them to span
+the clip, as the reference sampler does for records without a negative prompt of their own.
+"""
+
+import json
+from typing import Any
+
+# yapf: disable
+PAN2_DEFAULT_NEGATIVE_PROMPT = {
+ "video_style": "low quality, blur, softness, reduced detail, noise, compression, clipped highlights, "
+ "blurred sidebar, black border, letterboxing, watermark, subtitle, AIGC text, oily waxy "
+ "plastic texture, oversaturation or color cast, color clipping or chroma bleeding, flicker, "
+ "jitter, ghosting, motion smear, choppy or laggy motion, slow motion, Ken Burns effect, "
+ "unmotivated slow zoom over a frozen scene, pan across a still image, global frame scaling "
+ "or drifting, simulated movement from a static photograph, repeated micro-motion loop, "
+ "discontinuous shot, scene cut, hard cut, crossfade, fade-in, fade-out, fading transition, "
+ "scene replacement, spontaneous object or person pop-in or disappearance, duplicated or "
+ "missing people or props, identity or count drift, rigid-object deformation or morphing, "
+ "bending or melting rigid structures, structural shape or component drift",
+ "lighting": "overexposure, clipped highlights, underexposure, harsh lighting, overly deep skin shadows, "
+ "brightness fading, excessive lighting saturation, chroma-clipped highlights, neon-like skin "
+ "tones, colored spill overwhelming local material colors, oversaturated bloom, unnatural neon "
+ "color casts, hue drift, chroma noise, bad white balance, color pulsing, flickering light, "
+ "unnatural lighting",
+ "perspective_view": "warped geometry, distorted perspective, inconsistent spatial relationships, "
+ "unmotivated layout changes, camera teleport or jump, viewpoint reset, wrong depth or "
+ "occlusion",
+ "subjects_lifecycle": [
+ {
+ "subject_id": "<|Subject1|>",
+ "display_name": "malformed inconsistent subject",
+ "subject_attribute": "broken anatomy, distorted body proportions, reversed pelvis, "
+ "twisted torso, deformed arms or legs, reversed joints, fused "
+ "feet or toes, warped or collapsing faces, transparent or "
+ "floating head and body, missing body parts, detached or "
+ "disembodied body parts, floating hands or limbs, detached "
+ "articulated parts, disconnected joints or supports, hands "
+ "disconnected from wrists or forearms, discontinuous "
+ "limb-to-torso connections, blurred or fused fingers, missing "
+ "thumbs, extra or missing fingers, malformed wrists, extra "
+ "limbs, duplicated hands, mirrored face, identity drift, wrong "
+ "relative sizes, human-object fusion, unintended body-material "
+ "fusion, loss of separate material boundaries, opaque white "
+ "nail polish, gray nail polish, chalky opaque fingernails, "
+ "artificial acrylic nails",
+ }
+ ],
+ "init_frame_visual_description": "blur or softness, reduced detail, noise or compression artifacts, clipped "
+ "highlights, blurred sidebar, black border or letterboxing, oily waxy "
+ "plastic surfaces, oversaturation or color cast, exaggerated local "
+ "contrast, exposure fading, bad white balance, smeared dense textures, "
+ "blurred or distorted face, transparent or missing body parts, detached or "
+ "disembodied body parts, floating hands or limbs, detached articulated "
+ "parts, disconnected joints or supports, disconnected wrists or forearms, "
+ "discontinuous limb-to-torso connections, twisted torso, fused or missing "
+ "fingers, missing thumbs, malformed wrists, duplicated limbs, deformed "
+ "arms or legs, reversed knees, fused toes, warped geometry, deformed or "
+ "asymmetrically warped rigid objects, bent rigid structures, incorrect "
+ "rigid part geometry, topology, proportions, or component count, wrong "
+ "occlusion, floating or hovering objects, unsupported objects, object or "
+ "human-object fusion, unintended fusion of distinct surfaces or layers, "
+ "loss of separate material boundaries, merged objects, hand-object "
+ "penetration, object-object penetration, penetration, collision, "
+ "impossible contact, wrong sizes.",
+ "events": [
+ {
+ "event_index": 0,
+ "time_range_sec": "0 to 10",
+ "event_description": "blur, softness, reduced detail, compression artifacts, clipped "
+ "highlights, flicker, jitter, ghosting, motion smear, choppy or laggy "
+ "motion, slow motion, Ken Burns effect, unmotivated slow zoom over a "
+ "frozen scene, pan across a still image, global frame scaling or drifting, "
+ "simulated movement from a static photograph, repeated micro-motion loop, "
+ "oily oversaturated or color-pulsing surfaces, shimmering surface detail, "
+ "fingers fusing or changing count, hands smearing, facial features "
+ "merging, reversed joints, upper body rotating independently of attached "
+ "hips and feet, corkscrewed waist, limbs changing length or attachment "
+ "after occlusion, transparent body or head, detached or disembodied body "
+ "parts, floating hands or limbs, detached articulated parts, disconnected "
+ "joints or supports, hands disconnected from wrists or forearms, body "
+ "parts moving independently, vanishing limb-to-torso connections, merged "
+ "people or limbs, limbs fused into torso, merged or fused feet and toes, "
+ "floating or hovering objects, unsupported objects, object or human-object "
+ "fusion, unintended fusion of distinct surfaces or layers, loss of "
+ "separate material boundaries, merged objects, hand-object penetration, "
+ "object-object penetration, penetration or clipping, objects colliding or "
+ "morphing, rigid objects tearing, fragmenting, fusing, bending, melting or "
+ "changing dimensions, topology, silhouette, component count or layout, "
+ "subjects changing identity, splitting or duplicating, objects appearing "
+ "or disappearing without entry, exit or occlusion, changed identity or "
+ "count after occlusion, state changes inconsistent with the causing "
+ "action, broken physics, uncaused deformation or motion, state changes "
+ "preceding the causing interaction, loss of coupling during contact, pose "
+ "jumps, discontinuous shot, camera jump, hard scene cut, crossfade, "
+ "fade-in, fade-out, fading transition, scene replacement or viewpoint "
+ "reset.",
+ "camera_motion": [
+ {"time_range_sec": "0 to 10", "camera_motion": ["camera_jump", "hard_cut", "viewpoint_reset"]}
+ ],
+ "contained_actions": [
+ {
+ "time_range_sec": "0 to 10",
+ "action": "<|Subject1|> fingers fusing during grip changes, palms "
+ "merging with props, wrists collapsing, crossed limbs fusing, "
+ "intersecting or swapping attachment, duplicated hands or "
+ "feet, detached or disembodied hands or limbs, hands "
+ "disconnected from wrists or forearms, discontinuous "
+ "limb-to-torso connections, merged or fused feet and toes, "
+ "unstable contact, unintended grip slip, hand-object "
+ "penetration, body-object penetration, passage through solid "
+ "surfaces instead of openings, object-object penetration, "
+ "merged or fused objects, rigid objects fragmenting, fusing, "
+ "smearing or changing dimensions during motion, unintended "
+ "fusion of distinct surfaces or layers, loss of separate "
+ "material boundaries, floating or hovering objects, "
+ "unsupported objects, objects duplicating, splitting or "
+ "disappearing during interaction or occlusion, unchanged "
+ "action result, pose jump",
+ }
+ ],
+ }
+ ],
+ "quality": {
+ "video_quality": "low resolution, blur, softness, reduced detail, noise, compression, clipped "
+ "highlights, blurred sidebar, black border, letterboxing, oily waxy plastic "
+ "surfaces, greasy highlights, oversaturated colors, faded or washed-out "
+ "colors, exposure fading, smeared faces, hands or held objects, surface detail "
+ "smearing or swimming, material or existing markings changing without a "
+ "physical cause, color or shape morphing, object pop-in or pop-out, sudden "
+ "disappearance, missing persistent objects, object duplication or count drift, "
+ "rigid-object warping or deformation, changing rigid proportions, topology, "
+ "silhouette, or component layout",
+ "text_logos_overlays": "watermark, subtitle, captions, AIGC text or logo",
+ "artifacts": [
+ "watermark",
+ "template_or_text_dominated",
+ "talking_head",
+ "visual_quality_issue",
+ "edited_or_discontinuous",
+ "video_game",
+ "animation",
+ "ai_generated_content",
+ "black_border",
+ "ken_burn",
+ "time_lapse",
+ "blurred_sidebar",
+ "hard_scene_cuts",
+ ],
+ },
+}
+# yapf: enable
+
+
+def render_pan2_default_negative_prompt(num_frames: int, fps: float) -> str:
+ """The default negative prompt as compact one-line JSON, with every time range spanning the clip."""
+ span = f"0 to {(num_frames - 1) / fps:g}"
+
+ def fill(node: Any) -> Any:
+ if isinstance(node, dict):
+ return {key: span if key == "time_range_sec" else fill(value) for key, value in node.items()}
+ if isinstance(node, list):
+ return [fill(item) for item in node]
+ return node
+
+ return json.dumps(fill(PAN2_DEFAULT_NEGATIVE_PROMPT), ensure_ascii=False, separators=(",", ":"))
diff --git a/fastvideo/pipelines/basic/pan2/pan2_pipeline.py b/fastvideo/pipelines/basic/pan2/pan2_pipeline.py
new file mode 100644
index 000000000..bc2b7061c
--- /dev/null
+++ b/fastvideo/pipelines/basic/pan2/pan2_pipeline.py
@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 video generation: text-to-video, and image-to-video from a first frame."""
+
+import contextlib
+import importlib
+import sys
+
+import torch
+
+from fastvideo import envs
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.pan2.guardrails import (
+ PAN2Guardrails,
+ PAN2PromptGuardrailStage,
+ PAN2VideoGuardrailStage,
+)
+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
+
+logger = init_logger(__name__)
+
+# The modules that read FASTVIDEO_FA4 when they are imported.
+_FA4_READERS = ("fastvideo.attention.utils.flash_attn_default", "fastvideo.attention.utils.flash_attn_no_pad")
+
+
+def flash_attention_4_by_default() -> bool:
+ """Whether PAN2 turns FlashAttention-4 on, which on Blackwell GPUs is about 2.3x faster than the SDPA fallback.
+
+ Only when FASTVIDEO_FA4 is unset, the GPU is sm_100 or newer, FA4 imports, and no attention module has read the
+ variable yet.
+ """
+ if envs.FASTVIDEO_FA4.is_set() or not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
+ return False
+ try:
+ import fastvideo.attention.utils.flash_attn_cute # noqa: F401
+ except ImportError:
+ logger.info("PAN2 would use FlashAttention-4 on this GPU, but flash_attn.cute is not installed.")
+ return False
+ if any(name in sys.modules for name in _FA4_READERS):
+ logger.warning("PAN2 defaults to FlashAttention-4 on this GPU, but attention was set up before the pipeline; "
+ "set FASTVIDEO_FA4=1 to use it.")
+ return False
+ return True
+
+
+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 __init__(self, *args, **kwargs) -> None:
+ if not flash_attention_4_by_default():
+ super().__init__(*args, **kwargs)
+ return
+ # FASTVIDEO_FA4 is set only while the pipeline loads: the attention modules read it when they are imported
+ # and keep that choice, so other code in the process does not see the variable. Those modules are shared, so a
+ # later pipeline in the same process still runs the attention they chose.
+ logger.info("PAN2 uses FlashAttention-4 on this GPU; set FASTVIDEO_FA4=0 to turn it off.")
+ with envs.FASTVIDEO_FA4.override(True):
+ for name in _FA4_READERS:
+ # flash_attn_no_pad also needs the flash_attn FA2 package.
+ with contextlib.suppress(ImportError):
+ importlib.import_module(name)
+ super().__init__(*args, **kwargs)
+
+ def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+ pipeline_config = getattr(fastvideo_args, "pipeline_config", None)
+ # The guardrails are on unless the config turns them off.
+ guardrails = None
+ if getattr(pipeline_config, "enable_guardrails", True):
+ offload = bool(getattr(pipeline_config, "offload_guardrail_models", False))
+ guardrails = PAN2Guardrails(offload_to_cpu=offload)
+
+ self.add_stage(stage_name="input_validation_stage", stage=PAN2InputValidationStage())
+ if guardrails is not None:
+ self.add_stage(stage_name="prompt_guardrail_stage", stage=PAN2PromptGuardrailStage(guardrails))
+ 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.
+ 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))
+ if guardrails is not None:
+ self.add_stage(stage_name="video_guardrail_stage", stage=PAN2VideoGuardrailStage(guardrails))
+
+
+class PAN2ModularPipeline(PAN2Pipeline):
+ """The pipeline class PAN2 checkpoints declare in their ``modular_model_index.json``."""
+
+
+EntryClass = [PAN2Pipeline, PAN2ModularPipeline]
diff --git a/fastvideo/pipelines/basic/pan2/presets.py b/fastvideo/pipelines/basic/pan2/presets.py
new file mode 100644
index 000000000..87a2904c2
--- /dev/null
+++ b/fastvideo/pipelines/basic/pan2/presets.py
@@ -0,0 +1,49 @@
+# 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 empty negative prompt selects the release's structured negative caption,
+# rendered for the clip's duration by ``PAN2TextEncodingStage`` as in the reference sampler.
+_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"]
diff --git a/fastvideo/pipelines/basic/pan2/stages.py b/fastvideo/pipelines/basic/pan2/stages.py
new file mode 100644
index 000000000..2dec0fc90
--- /dev/null
+++ b/fastvideo/pipelines/basic/pan2/stages.py
@@ -0,0 +1,239 @@
+# 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.models.vision_utils import load_image
+from fastvideo.pipelines.basic.pan2.negative_prompt import render_pan2_default_negative_prompt
+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"
+
+
+def _set_default_negative_prompt(batch: ForwardBatch) -> None:
+ if not getattr(batch, "do_classifier_free_guidance", False) or getattr(batch, "negative_prompt", None):
+ return
+ num_frames = batch.num_frames[0] if isinstance(batch.num_frames, list) else batch.num_frames
+ fps = batch.fps[0] if isinstance(batch.fps, list) else batch.fps
+ batch.negative_prompt = render_pan2_default_negative_prompt(num_frames, fps or 24)
+
+
+class PAN2InputValidationStage(InputValidationStage):
+ """Rejects geometry the PAN2 VAE and patch embedding cannot represent.
+
+ An ``image_path`` first frame is loaded as the reference sampler loads it: the stored pixels, without the rotation
+ of an EXIF orientation tag.
+ """
+
+ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+ _set_default_negative_prompt(batch)
+ image_path = batch.image_path
+ if image_path is not None and not image_path.endswith(".mp4"):
+ batch.pil_image = load_image(image_path, exif_transpose=False)
+ batch.image_path = None
+ try:
+ batch = super().forward(batch, fastvideo_args)
+ finally:
+ batch.image_path = image_path
+ 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.
+
+ An empty or missing negative prompt is replaced by the release's default one, its time ranges set to the clip's
+ duration, as in the reference sampler. A single space encodes the empty unconditional prompt PAN2 trained with.
+ """
+
+ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+ _set_default_negative_prompt(batch)
+ return super().forward(batch, fastvideo_args)
+
+ @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()
diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py
index 4e8faa0ee..8373493ff 100644
--- a/fastvideo/pipelines/stages/denoising.py
+++ b/fastvideo/pipelines/stages/denoising.py
@@ -479,6 +479,9 @@ class DenoisingStage(PipelineStage):
noise_pred_text,
guidance_rescale=batch.guidance_rescale,
)
+
+ noise_pred = self.finish_guidance(noise_pred, noise_pred_text, batch, fastvideo_args)
+ noise_pred = self.prepare_scheduler_model_output(noise_pred, state)
if scheduler_fp32:
# Diffusers-style: fp32 Euler update outside autocast avoids BF16 drift.
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
@@ -612,6 +615,14 @@ class DenoisingStage(PipelineStage):
def prepare_timestep(self, timestep, latent_model_input, batch, state) -> torch.Tensor:
return timestep.repeat(latent_model_input.shape[0])
+ def finish_guidance(self, noise_pred, noise_pred_text, batch, fastvideo_args) -> torch.Tensor:
+ """The guided prediction handed on to the scheduler; a model overrides it to keep part of the output out of CFG."""
+ return noise_pred
+
+ def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor:
+ """The model output handed to the scheduler step, which returns the next latents in its dtype."""
+ return noise_pred
+
def finish_step(self, latents, state) -> torch.Tensor:
return latents
diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py
index e0116af9a..ef0e94df8 100644
--- a/fastvideo/pipelines/stages/latent_preparation.py
+++ b/fastvideo/pipelines/stages/latent_preparation.py
@@ -30,11 +30,16 @@ class LatentPreparationStage(PipelineStage):
denoised during the diffusion process.
"""
- def __init__(self, scheduler, transformer, use_btchw_layout: bool = False) -> None:
+ def __init__(self,
+ scheduler,
+ transformer,
+ use_btchw_layout: bool = False,
+ latent_dtype: torch.dtype | None = None) -> None:
super().__init__()
self.scheduler = scheduler
self.transformer = transformer
self.use_btchw_layout = use_btchw_layout
+ self.latent_dtype = latent_dtype
def forward(
self,
@@ -88,7 +93,7 @@ class LatentPreparationStage(PipelineStage):
batch.prompt_embeds = [dummy_prompt]
batch.negative_prompt_embeds = []
batch.do_classifier_free_guidance = False
- dtype = batch.prompt_embeds[0].dtype
+ dtype = self.latent_dtype or batch.prompt_embeds[0].dtype
device = get_local_torch_device()
generator = batch.generator
latents = batch.latents
diff --git a/fastvideo/registry.py b/fastvideo/registry.py
index 0a0d18361..5efe6b10e 100644
--- a/fastvideo/registry.py
+++ b/fastvideo/registry.py
@@ -10,6 +10,7 @@ from __future__ import annotations
import dataclasses
import os
+import re
from collections.abc import Callable
from functools import lru_cache
from typing import TYPE_CHECKING, Any
@@ -44,6 +45,7 @@ from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
from fastvideo.configs.pipelines.pika_h3_flex import PikaH3FlexPipelineConfig
from fastvideo.configs.pipelines.open_h3_regenerate_2k import OpenH3Regenerate2KPipelineConfig
+from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
from fastvideo.configs.pipelines.mmaudio import (
MMAudioLarge44kV2PipelineConfig,
MMAudioLarge44kV2AConfig,
@@ -1227,6 +1229,22 @@ def _register_configs() -> None:
pipeline_cls_name="MiniMaxH3ModularPipeline",
)
+ # PAN2: one pipeline runs text-to-video, and image-to-video when the request carries a first frame.
+ def _is_pan2(name: str) -> bool:
+ name = name.lower()
+ return "pan2modularpipeline" in name or re.search(r"(^|[/_-])pan-?2([/_.-]|$)", name) is not None
+
+ register_configs(
+ sampling_param_cls=None,
+ pipeline_config_cls=PAN2PipelineConfig,
+ workload_types=(WorkloadType.T2V, WorkloadType.I2V),
+ hf_model_paths=["IFM/PAN2"],
+ model_detectors=[_is_pan2],
+ model_family="pan2",
+ default_preset="pan2_t2v",
+ pipeline_cls_name="PAN2ModularPipeline",
+ )
+
# SD3.5
register_configs(
sampling_param_cls=None,
@@ -1404,6 +1422,8 @@ def _register_presets() -> None:
ALL_PRESETS as OPEN_H3_REGENERATE_2K_PRESETS, )
from fastvideo.pipelines.basic.mmaudio.presets import (
ALL_PRESETS as MMAUDIO_PRESETS, )
+ from fastvideo.pipelines.basic.pan2.presets import (
+ ALL_PRESETS as PAN2_PRESETS, )
from fastvideo.pipelines.basic.sd35.presets import (
ALL_PRESETS as SD35_PRESETS, )
from fastvideo.pipelines.basic.stable_audio.presets import (
@@ -1438,6 +1458,7 @@ def _register_presets() -> None:
PIKA_H3_FLEX_PRESETS,
OPEN_H3_REGENERATE_2K_PRESETS,
MMAUDIO_PRESETS,
+ PAN2_PRESETS,
SD35_PRESETS,
STABLE_AUDIO_PRESETS,
WAN_PRESETS,
diff --git a/fastvideo/tests/contract/test_env_policy.py b/fastvideo/tests/contract/test_env_policy.py
index e11c39bdc..d155182b8 100644
--- a/fastvideo/tests/contract/test_env_policy.py
+++ b/fastvideo/tests/contract/test_env_policy.py
@@ -90,6 +90,7 @@ EXTERNAL_ALLOWLIST = {
"CUDACXX": "CUDA compiler.",
"CUBLAS_WORKSPACE_CONFIG": "cuBLAS workspace setting.",
"TORCHDYNAMO_DISABLE": "PyTorch Dynamo switch.",
+ "TRITON_INTERPRET": "Triton interpreter switch; tests run kernels through it on CPU.",
"PYTORCH_MPS_*": "PyTorch MPS memory settings.",
"USER": "Login name of the current user.",
"PYTHONPATH": "Python module search path.",
diff --git a/fastvideo/tests/entrypoints/test_serving_family_baselines.py b/fastvideo/tests/entrypoints/test_serving_family_baselines.py
index bb7cf8b24..25bb903ba 100644
--- a/fastvideo/tests/entrypoints/test_serving_family_baselines.py
+++ b/fastvideo/tests/entrypoints/test_serving_family_baselines.py
@@ -25,6 +25,7 @@ from fastvideo.entrypoints.openai.request_adapter import (
prepare_reference_media,
)
from fastvideo.pipelines.basic.hunyuan15.presets import HUNYUAN15_I2V_480P_DISTILLED
+from fastvideo.pipelines.basic.pan2.presets import PAN2_T2V
ROOT = Path(__file__).resolve().parents[3]
@@ -37,6 +38,7 @@ BASELINES = {
"openai_kandinsky6_ti2va.yaml": "t2v",
"openai_kandinsky6_ti2va_distilled.yaml": "t2v",
"openai_ltx2_distilled.yaml": "t2v",
+ "openai_pan2.yaml": "t2v",
"openai_flux1_dev.yaml": "t2i",
"openai_flux2_dev.yaml": "t2i",
"openai_flux2_klein.yaml": "t2i",
@@ -112,6 +114,33 @@ def test_hunyuan_i2v_480p_keeps_the_registered_checkpoint_semantics(tmp_path: Pa
assert adapted.output.output_path.endswith(".mp4")
+def test_pan2_serving_matches_its_preset_and_runs_image_to_video_from_a_reference(tmp_path: Path) -> None:
+ config = _load("openai_pan2.yaml")
+ sampling = config.default_request.sampling
+ assert config.generator.model_path == "IFM/PAN2"
+ assert config.generator.engine.offload.dit is False
+ assert config.generator.engine.offload.dit_layerwise is False
+ for field in ("height", "width", "num_frames", "fps", "num_inference_steps", "guidance_scale"):
+ assert getattr(sampling, field) == PAN2_T2V.defaults[field], field
+ assert config.default_request.negative_prompt == PAN2_T2V.defaults["negative_prompt"]
+
+ image = tmp_path / "first-frame.png"
+ image.write_bytes(b"not-decoded-at-admission")
+ for reference in (None, {"image_url": str(image)}):
+ request = VideoGenerationRequest(prompt="a red fox", image_reference=reference)
+ adapted = build_generation_request(
+ "pan2_request",
+ request,
+ _args(config, "t2v"),
+ served_model_name=config.server.served_model_name or "",
+ output_dir=str(tmp_path),
+ default_request=config.default_request,
+ )
+ assert adapted.sampling.num_frames == 313
+ assert adapted.inputs.image_path == (str(image) if reference else None)
+ assert adapted.output.output_path.endswith(".mp4")
+
+
def test_stable_audio_seconds_override_produces_wav(tmp_path: Path) -> None:
config = _load("openai_stable_audio.yaml")
request = VideoGenerationRequest(prompt="rain on a metal roof", seconds=6)
diff --git a/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py b/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py
index 6344fafb8..04dd61018 100644
--- a/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py
+++ b/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-"""``_quantize_video_frames_to_uint8`` and its CUDA-OOM fallback.
+"""``_quantize_video_frames_to_uint8``, ``_video_frames`` and their CUDA-OOM fallback.
Found on a real GPU run of the K6 T2IVA-distilled example (30B params, ``dit_cpu_offload=False``): the worker process
keeps the transformer resident on the device after denoising and decode, so the client process's on-device quantize
@@ -75,3 +75,54 @@ def test_a_non_oom_runtime_error_is_not_swallowed(sample_video, monkeypatch):
monkeypatch.setattr(torch.Tensor, "clamp_", always_fails)
with pytest.raises(RuntimeError, match="unrelated failure"):
_quantize_video_frames_to_uint8(sample_video.clone())
+
+
+def _frames_by_grid(src: torch.Tensor) -> list:
+ """The frames as the generator built them before `_video_frames`: make_grid and a channels-last copy per frame."""
+ import torchvision
+
+ 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)
+ ]
+
+
+@pytest.mark.parametrize("shape", [(1, 3, 4, 5, 6), (1, 1, 4, 5, 6), (2, 3, 4, 5, 6)])
+def test_video_frames_match_the_per_frame_grid(shape):
+ import numpy as np
+
+ from fastvideo.entrypoints.video_generator import _video_frames
+
+ torch.manual_seed(1)
+ src = torch.rand(*shape) * 1.3 - 0.15
+ frames, expected = _video_frames(src.clone()), _frames_by_grid(src.clone())
+ assert len(frames) == len(expected) == shape[2]
+ for frame, reference in zip(frames, expected):
+ assert frame.dtype == np.uint8 and frame.shape == reference.shape
+ np.testing.assert_array_equal(frame, reference)
+
+
+def test_video_frames_fall_back_to_the_cpu_on_cuda_oom(monkeypatch, caplog):
+ import numpy as np
+
+ from fastvideo.entrypoints.video_generator import _video_frames
+
+ torch.manual_seed(2)
+ src = torch.rand(1, 3, 4, 5, 6)
+ expected = _frames_by_grid(src.clone())
+ calls = {"n": 0}
+ real_clamp_ = torch.Tensor.clamp_
+
+ def flaky_clamp_(self, *args, **kwargs):
+ calls["n"] += 1
+ if calls["n"] == 1:
+ raise torch.cuda.OutOfMemoryError("simulated")
+ return real_clamp_(self, *args, **kwargs)
+
+ monkeypatch.setattr(torch.Tensor, "clamp_", flaky_clamp_)
+ with caplog.at_level(logging.WARNING):
+ frames = _video_frames(src.clone())
+ assert calls["n"] == 2
+ for frame, reference in zip(frames, expected):
+ np.testing.assert_array_equal(frame, reference)
+ assert any("out of memory" in message.lower() for message in caplog.messages)
diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py
index 92a59aa1e..311552d95 100644
--- a/fastvideo/tests/entrypoints/test_video_generator.py
+++ b/fastvideo/tests/entrypoints/test_video_generator.py
@@ -980,6 +980,38 @@ def test_generate_single_video_latent_metadata_skips_cpu_materialization(tmp_pat
assert result["video_path"] is None
+@pytest.mark.parametrize(("preserve_space", "negative_prompt", "expected"), [
+ (False, " ", ""),
+ (True, " ", " "),
+ (True, " custom negative ", "custom negative"),
+])
+def test_generate_single_video_respects_single_space_negative_prompt_policy(
+ tmp_path, preserve_space, negative_prompt, expected):
+ output_batch = _single_video_output_batch(torch.zeros(1, 3, 2, 16, 16))
+ fastvideo_args = _single_video_args()
+ fastvideo_args.pipeline_config.preserve_single_space_negative_prompt = preserve_space
+ generator = _single_video_generator(output_batch, fastvideo_args)
+ captured = {}
+
+ def execute_forward(batch, args):
+ captured["batch"] = batch
+ return output_batch
+
+ generator.executor.execute_forward = execute_forward
+ sampling = _small_sampling_param()
+ sampling.negative_prompt = negative_prompt
+ generator._prepare_output_path = lambda *args: str(tmp_path / "unused.mp4")
+
+ generator._generate_single_video(
+ prompt="negative prompt",
+ sampling_param=sampling,
+ fastvideo_args=fastvideo_args,
+ output_path=str(tmp_path / "unused.mp4"),
+ )
+
+ assert captured["batch"].negative_prompt == expected
+
+
def test_generate_single_video_non_output_rank_suppresses_all_user_payloads(tmp_path):
audio = torch.ones(8)
audio_latents = torch.ones(2)
diff --git a/fastvideo/tests/stages/test_pan2_first_frame_and_attention.py b/fastvideo/tests/stages/test_pan2_first_frame_and_attention.py
new file mode 100644
index 000000000..973bf6ffd
--- /dev/null
+++ b/fastvideo/tests/stages/test_pan2_first_frame_and_attention.py
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2's first-frame loading (the stored pixels, as the reference reads them) and its FlashAttention-4 default."""
+
+import sys
+import types
+from types import SimpleNamespace
+
+import numpy as np
+import PIL.Image
+import pytest
+import torch
+
+from fastvideo import envs
+from fastvideo.pipelines.basic.pan2 import pan2_pipeline
+from fastvideo.pipelines.basic.pan2.stages import PAN2InputValidationStage
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+
+
+def _args():
+ vae = SimpleNamespace(arch_config=SimpleNamespace(spatial_compression_ratio=16, temporal_compression_ratio=4))
+ dit = SimpleNamespace(arch_config=SimpleNamespace(patch_size=(1, 2, 2)))
+ return SimpleNamespace(pipeline_config=SimpleNamespace(vae_config=vae, dit_config=dit))
+
+
+def test_an_exif_rotated_first_frame_keeps_its_stored_pixels(tmp_path, monkeypatch):
+ seen = []
+ monkeypatch.setattr(InputValidationStage, "forward",
+ lambda self, batch, args: seen.append((batch.image_path, batch.pil_image)) or batch)
+ stored = np.zeros((8, 16, 3), dtype=np.uint8)
+ stored[:, :8] = 255
+ path = tmp_path / "rotated.png"
+ exif = PIL.Image.Exif()
+ exif[0x0112] = 6 # Orientation: rotate 90 degrees clockwise to display upright.
+ PIL.Image.fromarray(stored).save(path, exif=exif)
+
+ batch = SimpleNamespace(image_path=str(path), pil_image=None, height=704, width=1248, num_frames=33)
+ out = PAN2InputValidationStage().forward(batch, _args())
+
+ # The base stage saw a loaded image and no path to load again; the path is restored for later stages.
+ assert seen[0][0] is None and seen[0][1] is out.pil_image
+ assert out.image_path == str(path)
+ assert out.pil_image.size == (16, 8) and (np.asarray(out.pil_image) == stored).all()
+
+
+@pytest.fixture
+def blackwell(monkeypatch):
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args: (10, 3))
+ monkeypatch.setitem(sys.modules, "fastvideo.attention.utils.flash_attn_cute",
+ types.ModuleType("fastvideo.attention.utils.flash_attn_cute"))
+ for name in pan2_pipeline._FA4_READERS:
+ monkeypatch.delitem(sys.modules, name, raising=False)
+ with envs.FASTVIDEO_FA4.override(None):
+ yield monkeypatch
+
+
+def test_pan2_turns_flash_attention_4_on_for_blackwell(blackwell):
+ assert pan2_pipeline.flash_attention_4_by_default() is True
+ assert not envs.FASTVIDEO_FA4.is_set()
+
+
+def test_flash_attention_4_is_set_only_while_the_pipeline_loads(blackwell):
+ seen = []
+ blackwell.setattr(pan2_pipeline.importlib, "import_module", lambda name: seen.append((name, envs.FASTVIDEO_FA4.get())))
+ blackwell.setattr(pan2_pipeline.ComposedPipelineBase, "__init__",
+ lambda self, *args, **kwargs: seen.append(("init", envs.FASTVIDEO_FA4.get())))
+ pan2_pipeline.PAN2Pipeline()
+ # The attention modules are imported, and the pipeline loads, with FA4 on; the variable is unset again after.
+ assert seen == [(name, True) for name in pan2_pipeline._FA4_READERS] + [("init", True)]
+ assert not envs.FASTVIDEO_FA4.is_set()
+
+
+@pytest.mark.parametrize("value", [False, True])
+def test_an_explicit_fastvideo_fa4_is_kept(blackwell, value):
+ envs.FASTVIDEO_FA4.set(value)
+ assert pan2_pipeline.flash_attention_4_by_default() is False
+ assert envs.FASTVIDEO_FA4.get() is value
+
+
+def test_hopper_keeps_the_default_backend(blackwell):
+ blackwell.setattr(torch.cuda, "get_device_capability", lambda *args: (9, 0))
+ assert pan2_pipeline.flash_attention_4_by_default() is False
+ assert not envs.FASTVIDEO_FA4.is_set()
+
+
+def test_without_fa4_installed_the_default_backend_stays(blackwell):
+ blackwell.setitem(sys.modules, "fastvideo.attention.utils.flash_attn_cute", None)
+ assert pan2_pipeline.flash_attention_4_by_default() is False
+ assert not envs.FASTVIDEO_FA4.is_set()
+
+
+def test_attention_set_up_before_the_pipeline_is_not_switched(blackwell):
+ blackwell.setitem(sys.modules, pan2_pipeline._FA4_READERS[0], types.ModuleType("already imported"))
+ assert pan2_pipeline.flash_attention_4_by_default() is False
+ assert not envs.FASTVIDEO_FA4.is_set()
diff --git a/fastvideo/tests/stages/test_pan2_first_latent_cfg.py b/fastvideo/tests/stages/test_pan2_first_latent_cfg.py
new file mode 100644
index 000000000..e74314091
--- /dev/null
+++ b/fastvideo/tests/stages/test_pan2_first_latent_cfg.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 image-to-video keeps the first latent frame out of classifier-free guidance; no weights or CUDA needed."""
+
+import pytest
+import torch
+
+from fastvideo.tests.stages._denoising_fixtures import (
+ NullProgressBar, TinyDenoiser, TinyScheduler, _patch_denoising_module, _tiny_args, _tiny_batch,
+)
+
+
+def _one_step(monkeypatch, env_overrides, *, image_to_video: bool, skip_first_latent: bool, guidance_rescale=0.0,
+ frame_values=(0.0, 0.0, 0.0)):
+ _patch_denoising_module(monkeypatch, env_overrides, None)
+ from fastvideo.pipelines.basic.pan2.stages import PAN2DenoisingStage
+
+ stage = PAN2DenoisingStage(TinyDenoiser(), TinyScheduler())
+ stage.progress_bar = lambda iterable=None, total=None: NullProgressBar()
+ batch = _tiny_batch()
+ batch.latents = torch.tensor(frame_values).reshape(1, 1, 3, 1, 1)
+ batch.raw_latent_shape = (1, 1, 3, 1, 1)
+ batch.timesteps = torch.tensor([0.0])
+ batch.num_inference_steps = 1
+ batch.pil_image = object() if image_to_video else None
+ batch.guidance_rescale = guidance_rescale
+ args = _tiny_args()
+ args.pipeline_config.i2v_skip_first_latent_cfg = skip_first_latent
+ args.pipeline_config.scheduler_step_in_fp32 = True
+ return stage.forward(batch, args).latents.flatten().tolist()
+
+
+def test_image_to_video_takes_the_positive_prediction_for_the_first_latent_frame(monkeypatch, env_overrides):
+ # With a zero latent and timestep the cond prediction is 0.5 and the uncond -0.5, so CFG 2 guides to 1.5.
+ latents = _one_step(monkeypatch, env_overrides, image_to_video=True, skip_first_latent=True)
+
+ assert latents == pytest.approx([-0.005, -0.015, -0.015])
+
+
+@pytest.mark.parametrize("image_to_video,skip_first_latent", [(True, False), (False, True)])
+def test_every_frame_is_guided_without_a_conditioning_image_or_with_the_skip_off(monkeypatch, env_overrides,
+ image_to_video, skip_first_latent):
+ latents = _one_step(monkeypatch, env_overrides, image_to_video=image_to_video, skip_first_latent=skip_first_latent)
+
+ assert latents == pytest.approx([-0.015, -0.015, -0.015])
+
+
+def test_first_latent_frame_keeps_the_positive_prediction_after_guidance_rescale(monkeypatch, env_overrides):
+ # Frames 0, 1, 2 predict 0.5, 0.7, 0.9 (cond) and 1.5, 1.7, 1.9 (CFG 2), so the rescale factor is 1 and only the
+ # first frame's replacement by 0.5 remains. Replacing before the rescale would scale it by std(cond)/std(mixed).
+ latents = _one_step(monkeypatch,
+ env_overrides,
+ image_to_video=True,
+ skip_first_latent=True,
+ guidance_rescale=1.0,
+ frame_values=(0.0, 1.0, 2.0))
+
+ assert latents == pytest.approx([-0.005, 0.983, 1.981])
diff --git a/fastvideo/tests/stages/test_pan2_guardrails.py b/fastvideo/tests/stages/test_pan2_guardrails.py
new file mode 100644
index 000000000..7a79d2ca8
--- /dev/null
+++ b/fastvideo/tests/stages/test_pan2_guardrails.py
@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 guardrails: the prompt and video checks against a stand-in `pan2_guardrail`; no weights or CUDA needed."""
+
+import importlib.machinery
+import sys
+import types
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
+from fastvideo.pipelines.basic.pan2 import guardrails as pan2_guardrails
+from fastvideo.pipelines.basic.pan2.guardrails import (
+ PAN2Guardrails,
+ PAN2GuardrailViolationError,
+ PAN2PromptGuardrailStage,
+ PAN2VideoGuardrailStage,
+)
+
+
+class FakeSafetyChecker:
+ """Blocks prompts containing "unsafe" and videos whose first pixel is white; records what it was given."""
+
+ def __init__(self):
+ self.text_guard = torch.nn.Linear(1, 1)
+ self.video_guard = torch.nn.Linear(1, 1)
+ self.prompts, self.videos, self.devices = [], [], []
+
+ @classmethod
+ def from_pretrained(cls):
+ return cls()
+
+ def to(self, device):
+ self.devices.append(str(device))
+ return self
+
+ def check_text_safety(self, prompt):
+ self.prompts.append(prompt)
+ return "unsafe" not in prompt
+
+ def check_video_safety(self, frames):
+ self.videos.append(frames)
+ return None if frames[0, 0, 0, 0] == 255 else frames
+
+
+@pytest.fixture
+def fake_package(monkeypatch):
+ module = types.ModuleType("pan2_guardrail")
+ module.__spec__ = importlib.machinery.ModuleSpec("pan2_guardrail", None)
+ module.PAN2SafetyChecker = FakeSafetyChecker
+ monkeypatch.setitem(sys.modules, "pan2_guardrail", module)
+ monkeypatch.setattr(pan2_guardrails, "get_local_torch_device", lambda: torch.device("cpu"))
+ return module
+
+
+def _args(output_type="pil"):
+ return SimpleNamespace(output_type=output_type)
+
+
+def test_guardrails_are_on_by_default():
+ config = PAN2PipelineConfig()
+ assert config.enable_guardrails is True
+ assert config.offload_guardrail_models is False
+
+
+def test_a_missing_package_says_how_to_install_or_disable(monkeypatch):
+ monkeypatch.setattr(pan2_guardrails.importlib.util, "find_spec", lambda name: None)
+ with pytest.raises(ImportError, match="pan2-guardrail.*enable_guardrails=False"):
+ PAN2Guardrails()
+
+
+def test_the_checker_loads_on_the_gpu_unless_offloaded(fake_package):
+ assert PAN2Guardrails().checker.devices == ["cpu"]
+ assert PAN2Guardrails(offload_to_cpu=True).checker.devices == []
+
+
+def test_a_safe_prompt_passes_and_an_unsafe_one_is_blocked(fake_package):
+ guardrails = PAN2Guardrails()
+ stage = PAN2PromptGuardrailStage(guardrails)
+
+ batch = SimpleNamespace(prompt="A red fox in the snow.")
+ assert stage.forward(batch, _args()) is batch
+ with pytest.raises(PAN2GuardrailViolationError, match="prompt was blocked"):
+ stage.forward(SimpleNamespace(prompt="an unsafe prompt"), _args())
+ assert guardrails.checker.prompts == ["A red fox in the snow.", "an unsafe prompt"]
+
+
+def test_latent_output_is_refused_before_any_compute(fake_package):
+ guardrails = PAN2Guardrails()
+ with pytest.raises(ValueError, match="latent output"):
+ PAN2PromptGuardrailStage(guardrails).forward(SimpleNamespace(prompt="A red fox."), _args("latent"))
+ assert guardrails.checker.prompts == []
+
+
+def test_each_video_is_checked_as_the_delivered_uint8_frames(fake_package):
+ guardrails = PAN2Guardrails()
+ stage = PAN2VideoGuardrailStage(guardrails)
+ # Two [3, T=2, H=4, W=5] videos in [0, 1]; values outside it saturate as in the generator.
+ output = torch.rand(2, 3, 2, 4, 5)
+ output[0, :, 0, 0, 0] = 0.5
+ output[1, :, 0, 0, 0] = -0.1
+
+ batch = SimpleNamespace(output=output)
+ assert stage.forward(batch, _args()) is batch
+ first, second = guardrails.checker.videos
+ assert first.shape == (2, 4, 5, 3) and first.dtype == "uint8"
+ assert first[0, 0, 0].tolist() == [127, 127, 127] and second[0, 0, 0].tolist() == [0, 0, 0]
+ assert (first == (output[0] * 255).clamp(0, 255).to(torch.uint8).permute(1, 2, 3, 0).numpy()).all()
+
+
+def test_a_blocked_video_raises(fake_package):
+ output = torch.zeros(1, 3, 2, 4, 5)
+ output[0, :, 0, 0, 0] = 1.0
+ with pytest.raises(PAN2GuardrailViolationError, match="video was blocked"):
+ PAN2VideoGuardrailStage(PAN2Guardrails()).forward(SimpleNamespace(output=output), _args())
+
+
+def test_offloaded_guards_visit_the_device_only_for_their_check(fake_package, monkeypatch):
+ moves = []
+ guardrails = PAN2Guardrails(offload_to_cpu=True)
+ for name in ("text_guard", "video_guard"):
+ guard = getattr(guardrails.checker, name)
+ monkeypatch.setattr(guard, "to", lambda device, name=name: moves.append((name, str(device))))
+
+ PAN2PromptGuardrailStage(guardrails).forward(SimpleNamespace(prompt="A red fox."), _args())
+ PAN2VideoGuardrailStage(guardrails).forward(SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2)), _args())
+ assert moves == [("text_guard", "cpu"), ("text_guard", "cpu"), ("video_guard", "cpu"), ("video_guard", "cpu")]
+
+
+class FakeGroup:
+ """A two-rank sequence-parallel group seen from ``rank``; a peer receives ``replies`` in order."""
+
+ def __init__(self, rank, replies=()):
+ self.rank_in_group, self.world_size = rank, 2
+ self.replies, self.broadcasts = list(replies), []
+
+ def broadcast_object(self, obj, src=0):
+ self.broadcasts.append(obj)
+ return obj if self.rank_in_group == src else self.replies.pop(0)
+
+
+def _sp_group(monkeypatch, group):
+ monkeypatch.setattr(pan2_guardrails, "model_parallel_is_initialized", lambda: True)
+ monkeypatch.setattr(pan2_guardrails, "get_sp_group", lambda: group)
+ return group
+
+
+def test_ranks_other_than_the_first_take_the_broadcast_verdict(fake_package, monkeypatch):
+ group = _sp_group(monkeypatch, FakeGroup(rank=1, replies=[(True, None), (True, False)]))
+ guardrails = PAN2Guardrails()
+ # This rank loads no checker and runs no check: it raises on rank 0's verdict.
+ assert guardrails.checker is None
+ with pytest.raises(PAN2GuardrailViolationError):
+ guardrails.check_prompt("A red fox.")
+ assert group.broadcasts == [None, None]
+
+
+def test_a_failed_load_on_the_first_rank_reaches_the_other_ranks_at_once(fake_package, monkeypatch):
+ def fail():
+ raise OSError("no guard weights")
+
+ monkeypatch.setattr(FakeSafetyChecker, "from_pretrained", staticmethod(fail))
+ first = _sp_group(monkeypatch, FakeGroup(rank=0))
+ with pytest.raises(OSError, match="no guard weights"):
+ PAN2Guardrails()
+ assert first.broadcasts == [(False, "OSError: no guard weights")]
+
+ _sp_group(monkeypatch, FakeGroup(rank=1, replies=first.broadcasts))
+ with pytest.raises(RuntimeError, match="failed to load the guardrails: OSError: no guard weights"):
+ PAN2Guardrails()
+
+
+def test_a_failed_check_on_the_first_rank_reaches_the_other_ranks_at_once(fake_package, monkeypatch):
+ first = _sp_group(monkeypatch, FakeGroup(rank=0))
+ guardrails = PAN2Guardrails()
+ monkeypatch.setattr(guardrails.checker, "check_text_safety", lambda prompt: 1 / 0)
+ with pytest.raises(ZeroDivisionError):
+ guardrails.check_prompt("A red fox.")
+ assert first.broadcasts[-1] == (False, "ZeroDivisionError: division by zero")
+
+ _sp_group(monkeypatch, FakeGroup(rank=1, replies=[(True, None), first.broadcasts[-1]]))
+ with pytest.raises(RuntimeError, match="failed to run the guardrail check"):
+ PAN2Guardrails().check_prompt("A red fox.")
+
+
+@pytest.mark.parametrize("config", [None, SimpleNamespace(enable_guardrails=True, offload_guardrail_models=True)])
+def test_the_pipeline_checks_the_prompt_first_and_the_video_last(monkeypatch, config):
+ from fastvideo.pipelines.basic.pan2 import pan2_pipeline
+
+ for name in ("PAN2InputValidationStage", "PAN2TextEncodingStage", "PAN2TimestepPreparationStage",
+ "PAN2ConditioningStage", "PAN2DenoisingStage", "DecodingStage", "LatentPreparationStage"):
+ monkeypatch.setattr(pan2_pipeline, name, lambda *args, **kwargs: None)
+ built = []
+ monkeypatch.setattr(pan2_pipeline, "PAN2Guardrails", lambda **kwargs: built.append(kwargs) or "guardrails")
+ pipeline = pan2_pipeline.PAN2Pipeline.__new__(pan2_pipeline.PAN2Pipeline)
+ stages = {}
+ pipeline.add_stage = lambda stage_name, stage: stages.__setitem__(stage_name, stage)
+ pipeline.get_module = lambda name: None
+
+ pipeline.create_pipeline_stages(SimpleNamespace(pipeline_config=config))
+
+ names = list(stages)
+ assert names[:2] == ["input_validation_stage", "prompt_guardrail_stage"]
+ assert names[-2:] == ["decoding_stage", "video_guardrail_stage"]
+ assert stages["prompt_guardrail_stage"].guardrails == stages["video_guardrail_stage"].guardrails == "guardrails"
+ assert built == [{"offload_to_cpu": config is not None}]
+
+
+def test_the_pipeline_has_no_guardrail_stages_when_they_are_off(monkeypatch):
+ from fastvideo.pipelines.basic.pan2 import pan2_pipeline
+
+ for name in ("PAN2InputValidationStage", "PAN2TextEncodingStage", "PAN2TimestepPreparationStage",
+ "PAN2ConditioningStage", "PAN2DenoisingStage", "DecodingStage", "LatentPreparationStage"):
+ monkeypatch.setattr(pan2_pipeline, name, lambda *args, **kwargs: None)
+ monkeypatch.setattr(pan2_pipeline, "PAN2Guardrails", lambda **kwargs: pytest.fail("guardrails were loaded"))
+ pipeline = pan2_pipeline.PAN2Pipeline.__new__(pan2_pipeline.PAN2Pipeline)
+ stages = {}
+ pipeline.add_stage = lambda stage_name, stage: stages.__setitem__(stage_name, stage)
+ pipeline.get_module = lambda name: None
+
+ pipeline.create_pipeline_stages(SimpleNamespace(pipeline_config=SimpleNamespace(enable_guardrails=False)))
+
+ assert not any("guardrail" in name for name in stages)
diff --git a/fastvideo/tests/stages/test_pan2_latent_preparation.py b/fastvideo/tests/stages/test_pan2_latent_preparation.py
new file mode 100644
index 000000000..de6954a09
--- /dev/null
+++ b/fastvideo/tests/stages/test_pan2_latent_preparation.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 draws its initial noise in fp32 whatever dtype the prompt embeddings arrive in; no weights or CUDA needed."""
+
+import types
+
+import pytest
+import torch
+
+from fastvideo.pipelines.basic.pan2 import pan2_pipeline
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages import LatentPreparationStage, latent_preparation
+
+
+@pytest.fixture()
+def cpu_device(monkeypatch):
+ monkeypatch.setattr(latent_preparation, "get_local_torch_device", lambda: torch.device("cpu"))
+
+
+def _args() -> types.SimpleNamespace:
+ arch = types.SimpleNamespace(temporal_compression_ratio=4, spatial_compression_ratio=16)
+ return types.SimpleNamespace(pipeline_config=types.SimpleNamespace(vae_config=types.SimpleNamespace(arch_config=arch)))
+
+
+def _noise(stage, embeds_dtype: torch.dtype) -> torch.Tensor:
+ batch = ForwardBatch(data_type="video",
+ prompt="a cat",
+ prompt_embeds=[torch.zeros(1, 4, 8, dtype=embeds_dtype)],
+ generator=torch.Generator("cpu").manual_seed(7),
+ height=32,
+ width=32,
+ num_frames=5)
+ return stage.forward(batch, _args()).latents
+
+
+def test_noise_is_fp32_and_independent_of_the_prompt_embedding_dtype(cpu_device):
+ transformer = types.SimpleNamespace(num_channels_latents=3)
+ stage = LatentPreparationStage(scheduler=types.SimpleNamespace(),
+ transformer=transformer,
+ latent_dtype=torch.float32)
+
+ from_bf16, from_fp32 = _noise(stage, torch.bfloat16), _noise(stage, torch.float32)
+
+ assert from_bf16.dtype == torch.float32
+ assert torch.equal(from_bf16, from_fp32)
+ assert _noise(LatentPreparationStage(types.SimpleNamespace(), transformer), torch.bfloat16).dtype == torch.bfloat16
+
+
+def test_pan2_pipeline_requests_fp32_latents(monkeypatch):
+ seen = {}
+
+ class _Stage:
+
+ def __init__(self, *args, **kwargs):
+ seen.setdefault("kwargs", []).append(kwargs)
+
+ for name in ("PAN2InputValidationStage", "PAN2TextEncodingStage", "PAN2TimestepPreparationStage",
+ "PAN2ConditioningStage", "PAN2DenoisingStage", "DecodingStage"):
+ monkeypatch.setattr(pan2_pipeline, name, _Stage)
+ monkeypatch.setattr(pan2_pipeline, "LatentPreparationStage", _Stage)
+ monkeypatch.setattr(pan2_pipeline, "PAN2Guardrails", lambda **kwargs: None)
+ pipeline = pan2_pipeline.PAN2Pipeline.__new__(pan2_pipeline.PAN2Pipeline)
+ stages = {}
+ pipeline.add_stage = lambda stage_name, stage: stages.__setitem__(stage_name, stage)
+ pipeline.get_module = lambda name: name
+
+ pipeline.create_pipeline_stages(None)
+
+ assert any(kwargs.get("latent_dtype") is torch.float32 for kwargs in seen["kwargs"])
+
+
+def test_pan2_pipeline_applies_its_performance_switches(monkeypatch):
+ for name in ("PAN2InputValidationStage", "PAN2TextEncodingStage", "PAN2TimestepPreparationStage",
+ "PAN2ConditioningStage", "PAN2DenoisingStage", "DecodingStage", "LatentPreparationStage"):
+ monkeypatch.setattr(pan2_pipeline, name, lambda *args, **kwargs: None)
+ monkeypatch.setattr(pan2_pipeline, "PAN2Guardrails", lambda **kwargs: None)
+
+ class _Transformer:
+ fused = None
+
+ def set_fused_kernels(self, enabled):
+ self.fused = enabled
+
+ modules = {"transformer": _Transformer(), "vae": types.SimpleNamespace(config=types.SimpleNamespace(decode_kernel=False))}
+ pipeline = pan2_pipeline.PAN2Pipeline.__new__(pan2_pipeline.PAN2Pipeline)
+ pipeline.add_stage = lambda stage_name, stage: None
+ pipeline.get_module = modules.get
+ config = types.SimpleNamespace(dit_fused_kernels=True, vae_decode_kernel=True, vae_decode_kernel_parallel=False)
+
+ pipeline.create_pipeline_stages(types.SimpleNamespace(pipeline_config=config))
+
+ assert modules["transformer"].fused is True
+ assert modules["vae"].config.decode_kernel is True and modules["vae"].config.decode_kernel_parallel is False
diff --git a/fastvideo/tests/stages/test_pan2_negative_prompt.py b/fastvideo/tests/stages/test_pan2_negative_prompt.py
new file mode 100644
index 000000000..96c3b59df
--- /dev/null
+++ b/fastvideo/tests/stages/test_pan2_negative_prompt.py
@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 falls back to the release's default negative prompt, rendered for the clip; no weights or CUDA needed."""
+
+import json
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
+from fastvideo.pipelines.basic.pan2.negative_prompt import (PAN2_DEFAULT_NEGATIVE_PROMPT,
+ render_pan2_default_negative_prompt)
+from fastvideo.pipelines.basic.pan2.stages import PAN2InputValidationStage, PAN2TextEncodingStage
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+
+def _time_ranges(node):
+ if isinstance(node, dict):
+ return [v for k, v in node.items() if k == "time_range_sec"] + sum(map(_time_ranges, node.values()), [])
+ if isinstance(node, list):
+ return sum(map(_time_ranges, node), [])
+ return []
+
+
+def test_the_default_spans_the_clip_and_is_otherwise_the_release_template():
+ rendered = json.loads(render_pan2_default_negative_prompt(481, 24))
+ assert _time_ranges(rendered) == ["0 to 20"] * 3
+ assert _time_ranges(PAN2_DEFAULT_NEGATIVE_PROMPT) == ["0 to 10"] * 3
+ assert render_pan2_default_negative_prompt(313, 24) == json.dumps(
+ json.loads(render_pan2_default_negative_prompt(313, 24)), ensure_ascii=False, separators=(",", ":"))
+ assert _time_ranges(json.loads(render_pan2_default_negative_prompt(313, 24))) == ["0 to 13"] * 3
+
+
+def test_single_space_preservation_is_a_pan2_policy_not_a_global_change():
+ assert PAN2PipelineConfig.preserve_single_space_negative_prompt is True
+ assert PipelineConfig.preserve_single_space_negative_prompt is False
+ assert "preserve_single_space_negative_prompt" not in {field.name for field in fields(PAN2PipelineConfig)}
+
+
+@pytest.fixture
+def encode(monkeypatch):
+ seen = []
+ monkeypatch.setattr(TextEncodingStage, "forward", lambda self, batch, args: seen.append(batch.negative_prompt))
+ stage = PAN2TextEncodingStage.__new__(PAN2TextEncodingStage)
+
+ def run(negative_prompt, cfg=True, num_frames=121, fps=24):
+ batch = SimpleNamespace(negative_prompt=negative_prompt, do_classifier_free_guidance=cfg,
+ num_frames=num_frames, fps=fps)
+ stage.forward(batch, None)
+ return seen[-1]
+
+ return run
+
+
+@pytest.mark.parametrize("negative_prompt", ["", None])
+def test_an_empty_negative_prompt_becomes_the_default(encode, negative_prompt):
+ assert encode(negative_prompt) == render_pan2_default_negative_prompt(121, 24)
+ assert encode(negative_prompt, num_frames=[241], fps=[24]) == render_pan2_default_negative_prompt(241, 24)
+ assert encode(negative_prompt, fps=None) == render_pan2_default_negative_prompt(121, 24)
+
+
+def test_a_given_negative_prompt_or_a_space_is_kept(encode):
+ assert encode("blurry") == "blurry"
+ assert encode(" ") == " "
+
+
+def test_without_guidance_the_negative_prompt_is_left_alone(encode):
+ assert encode("", cfg=False) == ""
+
+
+@pytest.mark.parametrize("negative_prompt", [None, "", " ", "blurry"])
+def test_validation_then_text_encoding_resolves_the_negative_prompt(monkeypatch, negative_prompt):
+ batch = ForwardBatch(
+ data_type="video",
+ prompt="a red fox",
+ negative_prompt=negative_prompt,
+ height=32,
+ width=32,
+ num_frames=5,
+ fps=24,
+ seed=0,
+ num_inference_steps=2,
+ guidance_scale=2.5,
+ )
+ args = SimpleNamespace(pipeline_config=SimpleNamespace(
+ text_encoder_configs=[SimpleNamespace()],
+ ti2v_task=False,
+ is_causal=False,
+ vae_config=SimpleNamespace(arch_config=SimpleNamespace(spatial_compression_ratio=8,
+ temporal_compression_ratio=4)),
+ dit_config=SimpleNamespace(arch_config=SimpleNamespace(patch_size=(1, 2, 2))),
+ ))
+ stage = PAN2TextEncodingStage([object()], [object()])
+ prompts = []
+
+ def encode_text(prompt, *args, **kwargs):
+ prompts.append(prompt)
+ return [torch.ones(1, 2, 3)], [torch.ones(1, 2, dtype=torch.long)]
+
+ monkeypatch.setattr(stage, "encode_text", encode_text)
+ assert PAN2InputValidationStage().forward(batch, args) is batch
+ assert stage.forward(batch, args) is batch
+
+ expected = negative_prompt or render_pan2_default_negative_prompt(5, 24)
+ assert batch.negative_prompt == expected
+ assert prompts == ["a red fox", expected]
+ assert len(batch.prompt_embeds) == len(batch.negative_prompt_embeds) == 1
diff --git a/fastvideo/tests/transformers/test_pan2_fusions.py b/fastvideo/tests/transformers/test_pan2_fusions.py
new file mode 100644
index 000000000..eaea99797
--- /dev/null
+++ b/fastvideo/tests/transformers/test_pan2_fusions.py
@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2's fused Q/K RMSNorm + RoPE and LayerNorm + modulation ops (``fastvideo.models.dits.pan2_fusions``).
+
+The PyTorch references must reproduce the eager arithmetic they replace, their backward formulas must be the true
+gradients (fp64 gradcheck), the Triton kernels must reproduce the references, and a PAN2 model with fused kernels must
+match the eager model forward and backward. On a GPU-less host the Triton kernels run under the Triton interpreter in
+a subprocess, in FP32: the interpreter truncates FP32 to BF16 instead of rounding to nearest even, so the BF16 kernel
+comparison needs a GPU.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig, PAN2VideoConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb, get_1d_rotary_pos_embed
+from fastvideo.models.dits import pan2 as pan2_module
+from fastvideo.models.dits import pan2_fusions
+from fastvideo.models.dits.pan2_fusions import layer_norm_modulate, qk_rmsnorm_rope
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+LATENT_CHANNELS = 4
+TINY_ARCH = dict(
+ patch_size=(1, 2, 2),
+ in_channels=2 * LATENT_CHANNELS + 1,
+ out_channels=LATENT_CHANNELS,
+ num_attention_heads=4,
+ attention_head_dim=16,
+ num_layers=2,
+ num_refiner_layers=1,
+ mlp_ratio=2.0,
+ context_mlp_ratio=2.0,
+ refiner_mlp_ratio=2.0,
+ text_embed_dim=16,
+ rope_axes_dim=(4, 6, 6),
+)
+
+
+@pytest.fixture()
+def cpu_attention(env_overrides):
+ env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+
+
+@pytest.fixture()
+def single_process_group(env_overrides):
+ """A 1x1 sequence/tensor-parallel group, which DistributedAttention needs even on one CPU process."""
+ from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+
+ env_overrides.enter_context(envs.override_external("MASTER_ADDR", os.environ.get("MASTER_ADDR") or "127.0.0.1"))
+ env_overrides.enter_context(envs.override_external("MASTER_PORT", os.environ.get("MASTER_PORT") or "29593"))
+ env_overrides.enter_context(envs.override_external("RANK", "0"))
+ env_overrides.enter_context(envs.override_external("WORLD_SIZE", "1"))
+ env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0"))
+ maybe_init_distributed_environment_and_model_parallel(1, 1)
+ try:
+ yield
+ finally:
+ cleanup_dist_env_and_memory()
+
+
+def _model() -> pan2_module.PAN2Transformer3DModel:
+ return pan2_module.PAN2Transformer3DModel(PAN2VideoConfig(arch_config=PAN2ArchConfig(**TINY_ARCH)), hf_config={})
+
+
+def _inputs(seed: int = 0) -> dict[str, torch.Tensor]:
+ generator = torch.Generator().manual_seed(seed)
+ return dict(
+ hidden_states=torch.randn(1, TINY_ARCH["in_channels"], 3, 4, 6, generator=generator),
+ encoder_hidden_states=torch.randn(1, 5, TINY_ARCH["text_embed_dim"], generator=generator),
+ timestep=torch.tensor([600.0]),
+ )
+
+
+def _randomize(model: torch.nn.Module, seed: int) -> None:
+ generator = torch.Generator().manual_seed(seed)
+ with torch.no_grad():
+ for parameter in model.parameters():
+ parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.1)
+
+
+def _qkv(batch: int, seq: int, heads: int, dim: int, dtype: torch.dtype, seed: int = 0):
+ generator = torch.Generator().manual_seed(seed)
+ qkv = torch.randn(batch, seq, 3 * heads * dim, generator=generator).to(dtype)
+ return qkv.unflatten(-1, (3, heads, -1)).unbind(2)
+
+
+def _rope_tables(seq: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+ positions = torch.arange(seq, dtype=torch.float32) * (16 / 24)
+ return get_1d_rotary_pos_embed(dim, positions, theta=256.0, use_real=False)
+
+
+def _eager_qk(x: torch.Tensor, weight: torch.Tensor, eps: float, rope: tuple[torch.Tensor, torch.Tensor] | None):
+ norm = RMSNorm(x.shape[-1], eps=eps)
+ with torch.no_grad():
+ norm.weight.copy_(weight)
+ y = norm.forward_native(x).to(x.dtype)
+ if rope is None:
+ return y
+ rows = y.flatten(0, 1)
+ cos, sin = (table.repeat(x.shape[0], 1) for table in rope)
+ return _apply_rotary_emb(rows, cos, sin, is_neox_style=False).unflatten(0, x.shape[:2])
+
+
+@pytest.mark.parametrize("with_rope", [False, True])
+def test_qk_reference_is_the_eager_chain_bit_for_bit(with_rope: bool) -> None:
+ q, k, _ = _qkv(2, 13, 3, 16, torch.bfloat16)
+ weights = torch.rand(16) + 0.5, torch.rand(16) + 0.5
+ rope = _rope_tables(13, 16) if with_rope else None
+ fused_q, fused_k = qk_rmsnorm_rope(q, k, *weights, 1e-6, *(rope or (None, None)))
+ assert torch.equal(fused_q, _eager_qk(q, weights[0], 1e-6, rope))
+ assert torch.equal(fused_k, _eager_qk(k, weights[1], 1e-6, rope))
+ assert fused_q.is_contiguous() and fused_q.dtype == torch.bfloat16
+
+
+def test_layer_norm_modulate_reference_is_the_autocast_arithmetic() -> None:
+ generator = torch.Generator().manual_seed(1)
+ x = (torch.randn(2, 11, 48, generator=generator) * 4 + 2).to(torch.bfloat16)
+ scale = (1 + torch.randn(2, 48, generator=generator)).to(torch.bfloat16)
+ shift = torch.randn(48, generator=generator).to(torch.bfloat16)
+ norm = torch.nn.LayerNorm(48, elementwise_affine=False, eps=1e-6)
+ # Under autocast the eager LayerNorm and modulation run in FP32 and the next linear rounds the result once.
+ expected = (norm(x.float()) * scale.float()[:, None] + shift.float()).to(torch.bfloat16)
+ actual = layer_norm_modulate(x, scale, shift, 1e-6)
+ assert actual.dtype == torch.bfloat16
+ torch.testing.assert_close(actual, expected, atol=0, rtol=2**-7)
+
+
+def test_fused_ops_have_true_gradients() -> None:
+ q, k, _ = _qkv(2, 5, 2, 8, torch.float64, seed=2)
+ q, k = q.detach().requires_grad_(), k.detach().requires_grad_()
+ weights = [(torch.rand(8, dtype=torch.float64) + 0.5).requires_grad_() for _ in range(2)]
+ cos, sin = (table.double() for table in _rope_tables(5, 8))
+ assert torch.autograd.gradcheck(lambda *args: qk_rmsnorm_rope(*args, 1e-6, cos, sin), (q, k, *weights))
+ assert torch.autograd.gradcheck(lambda *args: qk_rmsnorm_rope(*args, 1e-6), (q, k, *weights))
+
+ x = torch.randn(2, 4, 12, dtype=torch.float64, requires_grad=True)
+ scale = torch.randn(2, 12, dtype=torch.float64, requires_grad=True)
+ shift = torch.randn(12, dtype=torch.float64, requires_grad=True) # shared by the batch, as PAN2's text tables
+ assert torch.autograd.gradcheck(lambda *args: layer_norm_modulate(*args, 1e-6), (x, scale, shift))
+
+
+@pytest.mark.gpu
+def test_fused_ops_compile_with_their_backwards_on_gpu() -> None:
+ """Inside a fullgraph compiled region, as regional compile captures PAN2's blocks, both fused ops and their
+ backwards are opaque ops: tracing never launches a Triton kernel, and the values and gradients are eager's."""
+ if not torch.cuda.is_available() or not pan2_fusions.HAVE_TRITON:
+ pytest.skip("needs CUDA and Triton")
+ q, k, _ = (t.cuda() for t in _qkv(2, 21, 3, 16, torch.bfloat16, seed=6))
+ weights = [(torch.rand(16) + 0.5).cuda() for _ in range(2)]
+ cos, sin = (table.cuda() for table in _rope_tables(21, 16))
+ generator = torch.Generator().manual_seed(7)
+ x = torch.randn(2, 9, 40, generator=generator).cuda().to(torch.bfloat16)
+ scale = (1 + torch.randn(2, 40, generator=generator)).cuda()
+ shift = torch.randn(40, generator=generator).cuda()
+
+ def fn(q, k, q_weight, k_weight, x, scale, shift):
+ q_out, k_out = qk_rmsnorm_rope(q, k, q_weight, k_weight, 1e-6, cos, sin)
+ return (q_out.float().square().sum() + k_out.float().square().sum() +
+ layer_norm_modulate(x, scale, shift, 1e-6).float().square().sum())
+
+ def run(function):
+ leaves = [t.detach().clone().requires_grad_() for t in (q, k, *weights, x, scale, shift)]
+ loss = function(*leaves)
+ loss.backward()
+ return loss.detach(), [leaf.grad for leaf in leaves]
+
+ eager_loss, eager_grads = run(fn)
+ loss, grads = run(torch.compile(fn, fullgraph=True, dynamic=True))
+ torch.testing.assert_close(loss, eager_loss, atol=0, rtol=1e-5)
+ for actual, expected in zip(grads, eager_grads, strict=True):
+ torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+
+def check_triton_kernels(device: str, dtypes: tuple[torch.dtype, ...]) -> None:
+ """Every Triton kernel against its reference, forward and backward; raises on a mismatch."""
+ for dtype in dtypes:
+ tol = dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else dict(atol=0, rtol=2**-6)
+ # 21 rows: a partial row block. q and k are strided views of one projection, as in the model.
+ q, k, _ = (t.to(device) for t in _qkv(2, 21, 3, 16, dtype, seed=3))
+ weight = (torch.rand(16) + 0.5).to(device)
+ cos, sin = (table.to(device) for table in _rope_tables(21, 16))
+ grad = torch.randn(2, 21, 3, 16, generator=torch.Generator().manual_seed(4)).to(device=device, dtype=dtype)
+ for tables in ((cos, sin), (None, None)):
+ expected, expected_rstd = pan2_fusions._qk_forward_reference(q, weight, *tables, 1e-6)
+ actual, rstd = pan2_fusions._qk_forward_triton(q, weight, *tables, 1e-6)
+ torch.testing.assert_close(actual.float(), expected.float(), **tol)
+ torch.testing.assert_close(rstd, expected_rstd, atol=1e-6, rtol=1e-6)
+ expected_grads = pan2_fusions._qk_backward_reference(grad, q, weight, *tables, expected_rstd)
+ actual_grads = pan2_fusions._qk_backward_triton(grad, q, weight, *tables, rstd)
+ for a, e in zip(actual_grads, expected_grads, strict=True):
+ torch.testing.assert_close(a.float(), e.float(), atol=1e-4, rtol=1e-3 if dtype != torch.float32 else 1e-5)
+
+ generator = torch.Generator().manual_seed(5)
+ x = (torch.randn(2, 9, 40, generator=generator) * 3 + 1).to(device=device, dtype=dtype)
+ scale = (1 + torch.randn(2, 40, generator=generator)).to(device=device, dtype=dtype)
+ shift = torch.randn(40, generator=generator).to(device=device, dtype=dtype)
+ grad = torch.randn(2, 9, 40, generator=generator).to(device=device, dtype=dtype)
+ expected = pan2_fusions._ln_forward_reference(x, scale, shift, 1e-6)
+ actual = pan2_fusions._ln_forward_triton(x, scale, shift, 1e-6)
+ torch.testing.assert_close(actual[0].float(), expected[0].float(), **tol)
+ torch.testing.assert_close(actual[1:], expected[1:], atol=1e-5, rtol=1e-5)
+ expected_grads = pan2_fusions._ln_backward_reference(grad, x, scale, expected[1], expected[2])
+ actual_grads = pan2_fusions._ln_backward_triton(grad, x, scale, actual[1], actual[2])
+ for a, e in zip(actual_grads, expected_grads, strict=True):
+ torch.testing.assert_close(a.float(), e.float(), atol=1e-4, rtol=1e-3 if dtype != torch.float32 else 1e-5)
+
+
+def test_triton_kernels_match_their_references(env_overrides) -> None:
+ if not pan2_fusions.HAVE_TRITON:
+ pytest.skip("Triton is not installed")
+ if torch.cuda.is_available():
+ check_triton_kernels("cuda", (torch.float32, torch.bfloat16))
+ return
+ code = (f"import sys; sys.path.insert(0, {str(REPO_ROOT)!r}); "
+ "import fastvideo.models.dits.pan2_fusions as fusions; fusions._FORCE_TRITON = True; "
+ "import torch; from fastvideo.tests.transformers.test_pan2_fusions import check_triton_kernels; "
+ "check_triton_kernels('cpu', (torch.float32, ))")
+ # The child inherits this environment with the Triton interpreter switched on.
+ env_overrides.enter_context(envs.override_external("TRITON_INTERPRET", "1"))
+ process = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600)
+ assert process.returncode == 0, f"Triton interpreter check failed:\n{process.stdout}\n{process.stderr}"
+
+
+def _forward_backward(model: torch.nn.Module, inputs: dict[str, torch.Tensor]):
+ model.zero_grad(set_to_none=True)
+ with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")):
+ output = model(**inputs)
+ weight = torch.randn(output.shape, generator=torch.Generator().manual_seed(6), dtype=output.dtype)
+ (output * weight).sum().backward()
+ return output.detach(), {name: p.grad.detach().clone() for name, p in model.named_parameters() if p.grad is not None}
+
+
+def test_fused_model_matches_the_eager_model(cpu_attention, single_process_group) -> None:
+ model = _model()
+ _randomize(model, seed=7)
+ inputs = _inputs(seed=8)
+ eager_out, eager_grads = _forward_backward(model, inputs)
+ model.set_fused_kernels(True)
+ assert all(block.fused_kernels and block.attn.fused_kernels for block in model.transformer_blocks)
+ fused_out, fused_grads = _forward_backward(model, inputs)
+ torch.testing.assert_close(fused_out, eager_out, atol=1e-5, rtol=1e-5)
+ assert fused_grads.keys() == eager_grads.keys()
+ for name, grad in eager_grads.items():
+ torch.testing.assert_close(fused_grads[name], grad, atol=1e-5, rtol=1e-4, msg=name)
+
+
+@pytest.mark.gpu
+def test_fused_model_matches_the_eager_model_in_bf16_on_gpu(cpu_attention, single_process_group) -> None:
+ """The sampling numerics: BF16 activations under autocast, Triton kernels, on one GPU."""
+ if not torch.cuda.is_available():
+ pytest.skip("needs CUDA")
+ model = _model().cuda()
+ _randomize(model, seed=7)
+ inputs = {name: value.cuda() for name, value in _inputs(seed=8).items()}
+
+ def run():
+ model.zero_grad(set_to_none=True)
+ with torch.autocast("cuda", dtype=torch.bfloat16), set_forward_context(
+ current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")):
+ output = model(**inputs)
+ output.float().square().mean().backward()
+ return output.float(), {name: p.grad.float() for name, p in model.named_parameters() if p.grad is not None}
+
+ eager_out, eager_grads = run()
+ model.set_fused_kernels(True)
+ fused_out, fused_grads = run()
+ scale = eager_out.abs().mean()
+ assert float((fused_out - eager_out).abs().mean() / scale) < 1e-2
+ for name, grad in eager_grads.items():
+ error = float((fused_grads[name] - grad).norm() / grad.norm().clamp_min(1e-12))
+ assert error < 5e-2, f"{name}: relative gradient error {error:.3e}"
diff --git a/fastvideo/tests/vaes/test_wan_vae_decode_kernel.py b/fastvideo/tests/vaes/test_wan_vae_decode_kernel.py
new file mode 100644
index 000000000..a279f0163
--- /dev/null
+++ b/fastvideo/tests/vaes/test_wan_vae_decode_kernel.py
@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``AutoencoderKLWan.decode`` through fastvideo-kernel's Wan2.2 decoder (``WanVAEConfig.decode_kernel``).
+
+CPU tests: the switch is off unless asked for (PAN2 sampling asks), and a decode it cannot take runs the module's own
+decode without loading the kernel. GPU tests (sm_100 / sm_103, the PAN2 VAE's weights in
+``FASTVIDEO_TEST_PAN2_MODEL_DIR``): the kernel's decode against the module's fp32 decode, within three times the error
+of a bf16 decode (the kernel's own rule); another geometry decoding exactly as before; two ranks decoding on their own
+GPUs with no collective.
+"""
+
+from __future__ import annotations
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
+from fastvideo.models.wan import vae as wan_vae
+from fastvideo.models.wan.vae_config import WanVAEArchConfig, WanVAEConfig
+
+
+def _tiny_vae(decode_kernel: bool) -> wan_vae.AutoencoderKLWan:
+ torch.manual_seed(0)
+ arch = WanVAEArchConfig(base_dim=8, z_dim=4, dim_mult=(1, 1), num_res_blocks=1, temperal_downsample=(False, ),
+ latents_mean=(0.0, ) * 4, latents_std=(1.0, ) * 4)
+ config = WanVAEConfig(arch_config=arch)
+ config.load_encoder, config.load_decoder, config.decode_kernel = False, True, decode_kernel
+ return wan_vae.AutoencoderKLWan(config).eval()
+
+
+def test_the_kernel_is_opt_in_and_pan2_sampling_opts_in() -> None:
+ assert WanVAEConfig().decode_kernel is False and WanVAEConfig().decode_kernel_parallel is False
+ pan2 = PAN2PipelineConfig()
+ assert pan2.vae_decode_kernel is True and pan2.vae_decode_kernel_parallel is True and pan2.dit_fused_kernels
+
+
+def test_a_decode_the_kernel_cannot_take_runs_the_module(monkeypatch) -> None:
+ def no_kernel():
+ raise AssertionError("a CPU decode must not load the kernel")
+
+ monkeypatch.setattr(wan_vae, "_wan22_vae_decode", no_kernel)
+ z = torch.randn(1, 4, 2, 4, 4)
+ with torch.no_grad():
+ plain = _tiny_vae(False).decode(z)
+ routed = _tiny_vae(True).decode(z)
+ assert torch.equal(plain, routed)
+
+
+_ON_KERNEL_GPU = pytest.mark.skipif(
+ not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
+ reason="needs an sm_100 / sm_103 GPU")
+
+
+@pytest.fixture
+def pan2_vae_dir() -> str:
+ root = envs.FASTVIDEO_TEST_PAN2_MODEL_DIR.get()
+ if not root:
+ pytest.skip("set FASTVIDEO_TEST_PAN2_MODEL_DIR to a PAN2 root with vae/")
+ return os.path.join(root, "vae")
+
+
+def _pan2_vae(vae_dir: str) -> wan_vae.AutoencoderKLWan:
+ from fastvideo.models.loader.component_loader import VAELoader
+
+ config = WanVAEConfig()
+ config.load_encoder, config.load_decoder = False, True
+ loader_args = SimpleNamespace(model_paths={}, vae_cpu_offload=False,
+ pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32"))
+ return VAELoader().load(vae_dir, loader_args).to("cuda").eval()
+
+
+@_ON_KERNEL_GPU
+def test_the_kernel_decodes_pan2_latents_within_the_bf16_error(pan2_vae_dir) -> None:
+ vae = _pan2_vae(pan2_vae_dir)
+ z = torch.randn((1, 48, 3, 44, 78), generator=torch.Generator().manual_seed(3)).cuda()
+ with torch.no_grad():
+ vae.config.decode_kernel = False
+ reference = vae.decode(z)
+ with torch.autocast("cuda", dtype=torch.bfloat16):
+ bf16 = vae.decode(z)
+ vae.config.decode_kernel = True
+ kernel = vae.decode(z)
+ assert list(vae._kernel_decoders) == [(z.device, None)], "the kernel path was not taken"
+ assert kernel.shape == reference.shape == (1, 3, 9, 704, 1248) and kernel.dtype == torch.float32
+ assert float(kernel.abs().max()) <= 1.0
+ d, d_bf16 = kernel - reference, bf16.float() - reference
+ assert float(d.abs().max()) <= 3 * float(d_bf16.abs().max())
+ assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_bf16.pow(2).mean().sqrt())
+
+
+@_ON_KERNEL_GPU
+def test_another_geometry_decodes_as_before_on_the_gpu(pan2_vae_dir) -> None:
+ vae = _pan2_vae(pan2_vae_dir)
+ z = torch.randn((1, 48, 2, 32, 32), generator=torch.Generator().manual_seed(4)).cuda()
+ with torch.no_grad():
+ vae.config.decode_kernel = False
+ plain = vae.decode(z)
+ vae.config.decode_kernel = True
+ routed = vae.decode(z)
+ assert torch.equal(plain, routed) and not vae._kernel_decoders
+
+
+def _decode_alone(rank: int, store_path: str, vae_dir: str, z: torch.Tensor, out_path: str) -> None:
+ import torch.distributed as dist
+ from safetensors.torch import load_file
+
+ torch.cuda.set_device(rank)
+ dist.init_process_group("nccl", init_method=f"file://{store_path}", rank=rank, world_size=2)
+ try:
+ weights = load_file(os.path.join(vae_dir, "diffusion_pytorch_model.safetensors"))
+ decoder = wan_vae._wan22_vae_decode().WanVAEDecoder(weights, device=f"cuda:{rank}", local=True)
+ video = decoder.decode(z)
+ torch.save(video.cpu(), f"{out_path}.{rank}")
+ dist.barrier(device_ids=[rank])
+ finally:
+ dist.destroy_process_group()
+
+
+@_ON_KERNEL_GPU
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two GPUs")
+def test_two_ranks_decode_on_their_own_gpus(tmp_path, pan2_vae_dir) -> None:
+ import torch.multiprocessing as mp
+
+ # Build the extension here first, so the two ranks load it rather than both compiling into one directory.
+ wan_vae._wan22_vae_decode()._module()
+ z = torch.randn((1, 48, 2, 44, 78), generator=torch.Generator().manual_seed(5))
+ out = str(tmp_path / "video")
+ mp.start_processes(_decode_alone, args=(str(tmp_path / "store"), pan2_vae_dir, z, out), nprocs=2, join=True,
+ start_method="spawn")
+ first, second = torch.load(f"{out}.0"), torch.load(f"{out}.1")
+ assert first.shape == (1, 3, 5, 704, 1248) and torch.equal(first, second)
diff --git a/fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py b/fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py
new file mode 100644
index 000000000..4e2252f59
--- /dev/null
+++ b/fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``AutoencoderKLWan.decode`` through fastvideo-kernel's Wan2.2 decoder split over the sequence-parallel group
+(``WanVAEConfig.decode_kernel_parallel``), as PAN2 sampling decodes on several GPUs.
+
+Needs a multi-GPU torchrun launch (real NCCL collectives over the group), sm_100 / sm_103 GPUs and the PAN2 VAE's
+weights in ``FASTVIDEO_TEST_PAN2_MODEL_DIR``; skipped otherwise:
+
+ FASTVIDEO_TEST_PAN2_MODEL_DIR= torchrun --nproc-per-node=4 -m pytest \\
+ fastvideo/tests/vaes/test_wan_vae_decode_kernel_parallel_gpu.py -q
+
+Every rank must return the same video, the group's first rank's, within the kernel's own error rule against the
+module's fp32 decode (three times the error of a bf16 decode, as the one-GPU test); a video with fewer latent frames
+than GPUs decodes on each GPU alone, and a rank that cannot build the kernel makes the group decode alone without
+hanging the others. Over two GPUs the split decode is
+bit for bit the video one GPU decodes; over four and eight the bands of rows round differently, at the same error.
+"""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.models.wan.vae_config import WanVAEConfig
+
+_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1"))
+
+pytestmark = [
+ pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
+ reason="needs sm_100 / sm_103 GPUs"),
+ pytest.mark.skipif(_WORLD_SIZE < 2, reason="needs a torchrun launch with WORLD_SIZE > 1"),
+]
+
+
+@pytest.fixture(scope="module")
+def sp_group():
+ from fastvideo.distributed import get_sp_group, maybe_init_distributed_environment_and_model_parallel
+ maybe_init_distributed_environment_and_model_parallel(1, _WORLD_SIZE)
+ return get_sp_group()
+
+
+@pytest.fixture(scope="module")
+def pan2_vae_dir() -> str:
+ root = envs.FASTVIDEO_TEST_PAN2_MODEL_DIR.get()
+ if not root:
+ pytest.skip("set FASTVIDEO_TEST_PAN2_MODEL_DIR to a PAN2 root with vae/")
+ return os.path.join(root, "vae")
+
+
+def _load_pan2_vae(vae_dir: str):
+ from fastvideo.models.loader.component_loader import VAELoader
+
+ config = WanVAEConfig()
+ config.load_encoder, config.load_decoder, config.decode_kernel = False, True, True
+ loader_args = SimpleNamespace(model_paths={}, vae_cpu_offload=False,
+ pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32"))
+ return VAELoader().load(vae_dir, loader_args).to("cuda").eval()
+
+
+@pytest.fixture(scope="module")
+def pan2_vae(pan2_vae_dir):
+ return _load_pan2_vae(pan2_vae_dir)
+
+
+@pytest.mark.parametrize("latent_frames", (1, 2, 13))
+@torch.no_grad()
+def test_every_rank_returns_the_split_decode(sp_group, pan2_vae, latent_frames: int) -> None:
+ z = torch.randn((1, 48, latent_frames, 44, 78), generator=torch.Generator().manual_seed(latent_frames)).cuda()
+ pan2_vae.config.decode_kernel = False
+ reference = pan2_vae.decode(z)
+ with torch.autocast("cuda", dtype=torch.bfloat16):
+ bf16 = pan2_vae.decode(z)
+ pan2_vae.config.decode_kernel, pan2_vae.config.decode_kernel_parallel = True, False
+ alone = pan2_vae.decode(z)
+ pan2_vae.config.decode_kernel_parallel = True
+ split = pan2_vae.decode(z)
+ if latent_frames < sp_group.world_size: # fewer latent frames than GPUs: each GPU decodes alone
+ assert torch.equal(split, alone)
+ return
+ assert pan2_vae._kernel_decoders.get((z.device, tuple(sp_group.ranks))) is not None, "the split was not taken"
+ assert split.shape == reference.shape == (1, 3, 4 * latent_frames - 3, 704, 1248) and split.dtype == torch.float32
+ assert torch.equal(pan2_vae.decode(z), split), "a second split decode (the group's decoder reused) differs"
+
+ checksum = split.double().sum().reshape(1)
+ low, high = checksum.clone(), checksum.clone()
+ torch.distributed.all_reduce(low, op=torch.distributed.ReduceOp.MIN, group=sp_group.device_group)
+ torch.distributed.all_reduce(high, op=torch.distributed.ReduceOp.MAX, group=sp_group.device_group)
+ assert torch.equal(low, high), "the ranks hold different videos"
+
+ d, d_bf16 = split - reference, bf16.float() - reference
+ assert float(d.abs().max()) <= 3 * float(d_bf16.abs().max())
+ assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_bf16.pow(2).mean().sqrt())
+ if sp_group.world_size == 2:
+ assert torch.equal(split, alone)
+
+
+@torch.no_grad()
+def test_a_rank_without_the_kernel_makes_the_group_decode_alone(sp_group, pan2_vae_dir) -> None:
+ """Group rank 1 could not build the kernel (here: an earlier failure on that rank alone). Every rank still joins the
+ group's one agreement, so none waits on it forever, and they all decode alone: rank 1 with AutoencoderKLWan, the
+ others with the kernel on their own GPU."""
+ vae = _load_pan2_vae(pan2_vae_dir)
+ vae.config.decode_kernel_parallel = True
+ if sp_group.rank_in_group == 1:
+ vae._kernel_unavailable = "simulated: this rank could not build the kernel"
+ z = torch.randn((1, 48, 13, 44, 78), generator=torch.Generator().manual_seed(7)).cuda()
+ for _ in range(2): # the second decode reads the group's cached verdict without communicating
+ video = vae.decode(z)
+ assert video.shape == (1, 3, 49, 704, 1248) and video.dtype == torch.float32
+ assert vae._kernel_decoders[(z.device, tuple(sp_group.ranks))] is None
+ assert (vae._kernel_decoders[(z.device, None)] is None) == (sp_group.rank_in_group == 1)
diff --git a/tests/local_tests/encoders/test_pan2_qwen3_5_parity.py b/tests/local_tests/encoders/test_pan2_qwen3_5_parity.py
new file mode 100644
index 000000000..35a835d82
--- /dev/null
+++ b/tests/local_tests/encoders/test_pan2_qwen3_5_parity.py
@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3.5 text-encoder parity for PAN2 against Transformers' ``Qwen3_5TextModel``.
+
+Coverage scope: both. The checkpoint test loads the public tiny PAN2 text
+encoder through FastVideo's production ``TextEncoderLoader`` and feeds both
+models the exact PAN2 chat-templated token ids. The random-config tests use a
+larger hybrid stack (GQA, partial RoPE, more value than key heads in the Gated
+DeltaNet, sequences that cross several 64-token delta-rule chunks) and also
+cover bf16 and padded batches.
+
+Everything runs on CPU. Set ``PAN2_TINY_MODEL_DIR`` to a local copy of
+``wuqing157/tiny-pan2-modular-pipe``; otherwise the HF cache is used and the
+checkpoint test skips when it is absent.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+import socket
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextConfig
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.utils import set_default_torch_dtype
+
+PARITY_SCOPE = "both"
+TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe"
+PAN2_SYSTEM_MESSAGE = "You are a helpful assistant. Generate the video with the following instructions."
+PAN2_TEXT_ENCODER_LAYERS_TO_SKIP = 2
+PAN2_PROMPTS = (
+ "A red fox crosses fresh snow at sunrise.",
+ "A slow dolly shot glides through a rain-soaked neon market at night. Steam rises from food stalls while "
+ "vendors in bright raincoats shout over the hum of generators. Reflections of pink, teal and amber signs ripple "
+ "across puddles as a cyclist weaves between crowds, her umbrella catching droplets that sparkle in the light. "
+ "In the background, a tram rattles past an old clock tower, its windows glowing warm yellow. The camera rises to "
+ "reveal rooftops crowded with antennas and laundry lines, then tilts toward a thunderstorm rolling in over the "
+ "harbor, lightning briefly silhouetting cranes and container ships. Cinematic lighting, shallow depth of field, "
+ "35mm film grain, realistic motion blur.",
+)
+FP32_TOLERANCE = {"atol": 1e-5, "rtol": 1e-5}
+BF16_TOLERANCE = {"atol": 2e-2, "rtol": 2e-2}
+
+
+def _hf_qwen3_5() -> tuple[Any, Any]:
+ module = pytest.importorskip("transformers.models.qwen3_5", reason="transformers without Qwen3.5 support")
+ return module.Qwen3_5TextConfig, module.Qwen3_5TextModel
+
+
+def _tiny_model_dir() -> Path | None:
+ local = os.getenv("PAN2_TINY_MODEL_DIR")
+ if local:
+ return Path(local)
+ try:
+ from huggingface_hub import snapshot_download
+
+ return Path(snapshot_download(TINY_REPO_ID, local_files_only=True))
+ except Exception:
+ return None
+
+
+def _free_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("127.0.0.1", 0))
+ return sock.getsockname()[1]
+
+
+@pytest.fixture(scope="module")
+def cpu_single_rank():
+ """Single-rank CPU runtime; the loader would otherwise target ``mps`` on CPU-only Linux hosts."""
+ import fastvideo.distributed.parallel_state as parallel_state
+ import fastvideo.models.loader.component_loader as component_loader
+ from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+
+ with pytest.MonkeyPatch.context() as patch:
+ patch.setattr(parallel_state, "get_local_torch_device", lambda: torch.device("cpu"))
+ patch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu"))
+ for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+ ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+ patch.setenv(name, value)
+ maybe_init_distributed_environment_and_model_parallel(1, 1)
+ yield
+ cleanup_dist_env_and_memory()
+
+
+def _loader_args(config: Qwen3_5TextConfig, precision: str = "fp32") -> SimpleNamespace:
+ return SimpleNamespace(
+ pipeline_config=SimpleNamespace(text_encoder_configs=(config, ), text_encoder_precisions=(precision, )),
+ text_encoder_cpu_offload=False,
+ override_text_encoder_quant=None,
+ override_text_encoder_safetensors=None,
+ pin_cpu_memory=False,
+ disable_offload_on_unified_memory=lambda device_id, offload_flag=None: False,
+ )
+
+
+def _pan2_input_ids(tokenizer: Any, prompt: str) -> tuple[torch.Tensor, int]:
+ """Tokenize exactly like PAN2's ``encode_prompt`` and return ids plus the user-header crop."""
+ messages = [{"role": "system", "content": PAN2_SYSTEM_MESSAGE}, {"role": "user", "content": prompt or " "}]
+ input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+ return_dict=True).input_ids
+ header_ids = tokenizer("<|im_start|>user\n", 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,
+ )
+ assert crop_start is not None, "PAN2 chat template has no user header"
+ return torch.tensor([input_ids]), crop_start
+
+
+def _run_both(reference: torch.nn.Module, native: torch.nn.Module, **inputs: Any) -> tuple[Any, Any]:
+ with torch.inference_mode():
+ expected = reference(**inputs, output_hidden_states=True)
+ with set_forward_context(current_timestep=0, attn_metadata=None):
+ actual = native(**inputs, output_hidden_states=True)
+ return expected, actual
+
+
+def _assert_hidden_states_close(
+ expected: tuple[torch.Tensor, ...],
+ actual: tuple[torch.Tensor, ...] | None,
+ label: str,
+ tolerance: dict[str, float],
+ valid: torch.Tensor | None = None,
+) -> float:
+ assert actual is not None, f"{label}: native encoder returned no hidden_states"
+ assert len(actual) == len(expected), f"{label}: {len(actual)} hidden states, reference has {len(expected)}"
+ worst = 0.0
+ for index, (reference_state, native_state) in enumerate(zip(expected, actual, strict=True)):
+ assert native_state.shape == reference_state.shape, f"{label} hidden_states[{index}] shape"
+ assert torch.isfinite(native_state).all(), f"{label} hidden_states[{index}] has non-finite values"
+ if valid is not None:
+ reference_state = reference_state[valid]
+ native_state = native_state[valid]
+ drift = (native_state.float() - reference_state.float()).abs().max().item()
+ worst = max(worst, drift)
+ print(f"{label} hidden_states[{index}]: max_abs={drift:.3e}", flush=True)
+ assert_close(native_state, reference_state, **tolerance, msg=lambda m: f"{label} hidden_states[{index}]: {m}")
+ return worst
+
+
+def test_pan2_qwen3_5_arch_config_matches_transformers_rope_and_layer_types() -> None:
+ """Config.json fields resolve to the same RoPE settings and layer schedule as Transformers."""
+ hf_config_cls, _ = _hf_qwen3_5()
+ cases = (
+ {},
+ {
+ "rope_theta": 1_000_000.0
+ },
+ {
+ "partial_rotary_factor": 0.5,
+ "rope_parameters": {
+ "rope_type": "default",
+ "rope_theta": 5_000_000.0
+ }
+ },
+ {
+ "partial_rotary_factor": 0.5,
+ "rope_parameters": {
+ "rope_type": "default",
+ "rope_theta": 5_000_000.0,
+ "partial_rotary_factor": 0.25
+ }
+ },
+ {
+ "num_hidden_layers": 6,
+ "full_attention_interval": 3
+ },
+ {
+ "num_hidden_layers": 3,
+ "layer_types": ["mamba", "attention", "linear_attention"]
+ },
+ )
+ for case in cases:
+ reference = hf_config_cls.from_dict(dict(case))
+ config = Qwen3_5TextConfig()
+ config.update_model_arch(dict(case))
+ rope = config.arch_config.resolved_rope_parameters()
+ for key in ("rope_type", "rope_theta", "partial_rotary_factor"):
+ assert rope[key] == reference.rope_parameters[key], f"{case}: {key}"
+ assert config.arch_config.resolved_layer_types() == list(reference.layer_types), case
+
+
+def test_pan2_qwen3_5_tiny_checkpoint_loader_parity(cpu_single_rank) -> None:
+ """Production-loader encoder matches every HF hidden state on real PAN2 prompts (fp32 CPU)."""
+ _, hf_model_cls = _hf_qwen3_5()
+ model_dir = _tiny_model_dir()
+ if model_dir is None or not (model_dir / "text_encoder" / "config.json").is_file():
+ pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR")
+ if not (model_dir / "tokenizer").is_dir():
+ pytest.skip(f"{TINY_REPO_ID} tokenizer is not available locally")
+ from transformers import AutoTokenizer
+
+ from fastvideo.models.encoders.qwen3_5 import Qwen3_5TextModel
+ from fastvideo.models.loader.component_loader import TextEncoderLoader
+
+ text_encoder_dir = model_dir / "text_encoder"
+ tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer", local_files_only=True)
+ reference, loading_info = hf_model_cls.from_pretrained(
+ text_encoder_dir,
+ local_files_only=True,
+ dtype=torch.float32,
+ output_loading_info=True,
+ )
+ load_errors = {
+ name: loading_info.get(name)
+ for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") if loading_info.get(name)
+ }
+ assert not load_errors, f"Reference Qwen3.5 checkpoint did not load strictly: {load_errors}"
+ reference = reference.eval()
+
+ native = TextEncoderLoader().load(str(text_encoder_dir), _loader_args(Qwen3_5TextConfig()))
+ assert isinstance(native, Qwen3_5TextModel)
+ assert native.layer_types == list(reference.config.layer_types)
+ assert next(native.parameters()).dtype == torch.float32
+ assert_close(native.rotary_emb.inv_freq, reference.rotary_emb.inv_freq.float(), atol=0.0, rtol=0.0)
+
+ for prompt_index, prompt in enumerate(PAN2_PROMPTS):
+ input_ids, crop_start = _pan2_input_ids(tokenizer, prompt)
+ label = f"tiny prompt{prompt_index} (L={input_ids.shape[1]})"
+ expected, actual = _run_both(reference, native, input_ids=input_ids)
+ worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, label, FP32_TOLERANCE)
+ assert_close(actual.last_hidden_state, expected.last_hidden_state, **FP32_TOLERANCE)
+
+ skip = PAN2_TEXT_ENCODER_LAYERS_TO_SKIP + 1
+ reference_embeds = expected.hidden_states[-skip][:, crop_start:]
+ native_embeds = actual.hidden_states[-skip][:, crop_start:]
+ assert_close(native_embeds, reference_embeds, **FP32_TOLERANCE)
+ print(
+ f"{label} PAN2 prompt_embeds hidden_states[-{skip}][:, {crop_start}:] "
+ f"max_abs={(native_embeds - reference_embeds).abs().max().item():.3e} all_layers_max_abs={worst:.3e}",
+ flush=True,
+ )
+
+
+@pytest.fixture(scope="module")
+def random_pair(cpu_single_rank):
+ """A randomly initialized 8-layer hybrid reference and the native model loaded from its state dict."""
+ hf_config_cls, hf_model_cls = _hf_qwen3_5()
+ from fastvideo.models.encoders.qwen3_5 import Qwen3_5TextModel
+
+ torch.manual_seed(0)
+ hf_config = hf_config_cls(
+ vocab_size=1024,
+ hidden_size=128,
+ intermediate_size=256,
+ num_hidden_layers=8,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ head_dim=32,
+ linear_conv_kernel_dim=4,
+ linear_key_head_dim=16,
+ linear_value_head_dim=24,
+ linear_num_key_heads=2,
+ linear_num_value_heads=4,
+ layer_types=[
+ "linear_attention",
+ "linear_attention",
+ "full_attention",
+ "linear_attention",
+ "full_attention",
+ "linear_attention",
+ "linear_attention",
+ "full_attention",
+ ],
+ rope_parameters={
+ "rope_type": "default",
+ "rope_theta": 1_000_000.0,
+ "partial_rotary_factor": 0.25
+ },
+ initializer_range=0.08,
+ )
+ reference = hf_model_cls(hf_config).eval()
+ with torch.no_grad():
+ # Move the zero-centered norms, gates and decay parameters off their
+ # init values so every term contributes.
+ for name, parameter in reference.named_parameters():
+ if name.endswith("norm.weight") or name.endswith(("A_log", "dt_bias")):
+ parameter.add_(0.2 * torch.randn_like(parameter))
+
+ config = Qwen3_5TextConfig()
+ config.update_model_arch(hf_config.to_dict())
+ with set_default_torch_dtype(torch.float32):
+ native = Qwen3_5TextModel(config).eval()
+ loaded = native.load_weights(reference.state_dict().items())
+ expected_names = {name for name, _ in native.named_parameters()}
+ assert loaded == expected_names, f"unloaded native parameters: {sorted(expected_names - loaded)}"
+ assert set(reference.state_dict()) == expected_names
+ return reference, native
+
+
+def test_pan2_qwen3_5_random_config_fp32_parity(random_pair) -> None:
+ """Multi-chunk delta rule (L=301), GQA and partial RoPE match HF in fp32."""
+ reference, native = random_pair
+ generator = torch.Generator().manual_seed(1)
+ input_ids = torch.randint(0, reference.config.vocab_size, (2, 301), generator=generator)
+ expected, actual = _run_both(reference, native, input_ids=input_ids)
+ worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, "random fp32", FP32_TOLERANCE)
+ print(f"random fp32 all_layers_max_abs={worst:.3e}", flush=True)
+ # Each layer must move the residual stream, or the comparison is vacuous.
+ for index in range(1, len(expected.hidden_states) - 1):
+ change = (expected.hidden_states[index] - expected.hidden_states[index - 1]).abs().max().item()
+ assert change > 0.1, f"layer {index - 1} barely changes the hidden state ({change:.3e})"
+
+
+def test_pan2_qwen3_5_random_config_padding_parity(random_pair) -> None:
+ """Right- and left-padded rows match HF on every non-padded position."""
+ reference, native = random_pair
+ generator = torch.Generator().manual_seed(2)
+ input_ids = torch.randint(0, reference.config.vocab_size, (2, 200), generator=generator)
+ attention_mask = torch.ones_like(input_ids)
+ attention_mask[0, 150:] = 0
+ attention_mask[1, :37] = 0
+ expected, actual = _run_both(reference, native, input_ids=input_ids, attention_mask=attention_mask)
+ worst = _assert_hidden_states_close(
+ expected.hidden_states,
+ actual.hidden_states,
+ "random padded",
+ FP32_TOLERANCE,
+ valid=attention_mask.bool(),
+ )
+ print(f"random padded all_layers_max_abs={worst:.3e}", flush=True)
+
+
+def test_pan2_qwen3_5_random_config_bf16_parity(random_pair) -> None:
+ """Half-precision dtype handling (fp32 norms, gates and delta rule) follows HF."""
+ reference, native = (copy.deepcopy(model).to(torch.bfloat16) for model in random_pair)
+ generator = torch.Generator().manual_seed(3)
+ input_ids = torch.randint(0, reference.config.vocab_size, (1, 173), generator=generator)
+ expected, actual = _run_both(reference, native, input_ids=input_ids)
+ assert actual.hidden_states is not None and actual.hidden_states[-1].dtype == torch.bfloat16
+ worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, "random bf16", BF16_TOLERANCE)
+ print(f"random bf16 all_layers_max_abs={worst:.3e}", flush=True)
diff --git a/tests/local_tests/pan2/PORT_STATUS.md b/tests/local_tests/pan2/PORT_STATUS.md
new file mode 100644
index 000000000..5e979c6c1
--- /dev/null
+++ b/tests/local_tests/pan2/PORT_STATUS.md
@@ -0,0 +1,71 @@
+# PAN2 Port Status
+
+## Summary
+- model_family: pan2
+- workload_types: T2V, I2V (one pipeline; I2V when the request carries a first frame)
+- official_ref: vllm-project/vllm-omni#8602 at `e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74` (the diffusers integration is private)
+- official_ref_dir: none; rebuilt in `tests/local_tests/pan2/reference.py`
+- hf_weights_path: `IFM/PAN2` (not public); tests use `wuqing157/tiny-pan2-modular-pipe@3ff3130b`
+- local_weights_dir: HF cache, or `PAN2_TINY_MODEL_DIR`
+- source_layout: Modular Diffusers (`modular_model_index.json`), diffusers component subfolders
+- local_tests_readme: `tests/local_tests/pan2/README.md`
+
+## Current Phase
+- phase: 11 (post-parity review and handoff)
+- status: in_progress (GPU and real-weight validation pending, see I001 and I002)
+- owner: orchestrator
+- last_updated: 2026-10-20
+
+## Component Matrix
+| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
+|---|---|---|---|---|---|---|---|---|---|
+| transformer | DiT | port | vLLM-Omni `pan2_transformer.py` (diffusers `PAN2Transformer3DModel`) | `transformer/config.json` | `fastvideo/models/dits/pan2.py`, `fastvideo/configs/models/dits/pan2.py` | done | load-time key mapping (Q/K/V fused) | non_skip_pass | I002 |
+| text_encoder | encoder | port | transformers `Qwen3_5TextModel` | `text_encoder/config.json` | `fastvideo/models/encoders/qwen3_5.py`, `fastvideo/configs/models/encoders/qwen3_5.py` | done | passthrough | non_skip_pass | I002 |
+| tokenizer | tokenizer | reuse | `Qwen2Tokenizer` | `tokenizer/` | AutoTokenizer | n/a | passthrough | covered by pipeline parity | — |
+| vae | VAE | reuse | diffusers `AutoencoderKLWan` (Wan2.2, z_dim 48) | `vae/config.json` | `fastvideo/models/wan/vae.py` | done | passthrough | non_skip_pass | — |
+| scheduler | scheduler | reuse | diffusers `FlowMatchEulerDiscreteScheduler` | `scheduler/scheduler_config.json` (shift 7) | `fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py` | done | passthrough | non_skip_pass | — |
+
+## Conversion State
+- conversion_script: none; the transformer maps diffusers keys at load time (`PAN2ArchConfig.param_names_mapping`)
+- converted_weights_dir: n/a
+- source_layout: diffusers
+- strict_load_status: tiny checkpoint loads strictly (no missing or unexpected keys)
+- passthrough_components: text_encoder, tokenizer, vae, scheduler
+- retry_history: none
+
+## Parity Commands
+| Scope | Command | Last Result | Notes |
+|---|---|---|---|
+| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 13 passed (CPU) | SP=2/4 vs one rank over gloo, T2V/I2V, one/two pad tokens; max abs 7.7e-6 |
+| Text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | 5 passed (CPU) | exact |
+| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | 9 passed (CPU) | exact |
+| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | 9 passed (CPU) | exact |
+| Pipeline | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | 3 passed (CPU) | T2V and I2V latents exact |
+
+## Open Questions
+| ID | Question | Owner | Needed By Phase | Status | Resolution |
+|---|---|---|---|---|---|
+| Q001 | Which reference, with the diffusers integration private? | orchestrator | 2 | resolved | vLLM-Omni #8602, rebuilt from diffusers/transformers blocks in `reference.py` |
+| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Ported: `PAN2PromptGuardrailStage` and `PAN2VideoGuardrailStage` (`guardrails.py`), on by default (`enable_guardrails`). |
+| Q003 | Port Cache-DiT `quality` and CFG-parallel? | orchestrator | 7 | resolved | Not ported: vLLM-Omni serving features. Ulysses SP is supported through `DistributedAttention`. |
+
+## Issues And Blockers
+| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
+|---|---|---|---|---|---|---|---|---|
+| I001 | 9-10 | all | medium | No real-weight parity or SSIM regression | `IFM/PAN2` is not public | orchestrator | open | Add `fastvideo/tests/ssim/test_pan2_similarity.py` once weights are public. |
+| I002 | 9 | transformer, text_encoder | medium | No GPU parity run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | Local parity tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. |
+| I003 | 7 | pipeline | low | The shared denoiser casts latents to the model-output dtype in the scheduler step; PAN2 samples on fp32 latents | pipeline parity fails (8e-2) without the fix | orchestrator | resolved | `DenoisingStage.prepare_scheduler_model_output` hook; `PAN2DenoisingStage` returns `noise_pred.float()`. |
+
+## Escape Hatches
+| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
+|---|---|---|---|---|---|---|
+
+## Decisions
+| Date | Decision | Rationale | Impact |
+|---|---|---|---|
+| 2026-10-15 | Native Qwen3.5 encoder with a PyTorch chunked gated delta rule | No `fla`/`causal_conv1d` dependency; encoder runs once per prompt | Slower than fused kernels for long prompts |
+| 2026-10-15 | Fuse the per-stream Q/K/V projections at load time | One GEMM per stream instead of three | Mapping in `PAN2ArchConfig.param_names_mapping` |
+| 2026-10-15 | Text embeddings kept in fp32 after encoding | The shared latent stage draws noise in the embedding dtype; PAN2 uses fp32 noise and latents | Denoiser casts embeddings to the DiT dtype (`cast_prompt_embeds_to_dit_dtype`) |
+
+## Handoff Notes
+- The remaining work is I001 and I002. Both need hardware or weights, not code changes.
diff --git a/tests/local_tests/pan2/README.md b/tests/local_tests/pan2/README.md
new file mode 100644
index 000000000..2fac8e256
--- /dev/null
+++ b/tests/local_tests/pan2/README.md
@@ -0,0 +1,67 @@
+# PAN2 Local Tests
+
+Runbook for validating the PAN2 port (text-to-video and first-frame image-to-video). These tests run manually and are
+not part of the default CI suite.
+
+## Reference
+
+PAN2's diffusers integration (`PAN2Transformer3DModel`, `PAN2ModularPipeline`) is not public yet, and the full
+`IFM/PAN2` weights are not released. The reference is the vLLM-Omni port,
+[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `e0c4c731`), which reports
+parity with the diffusers implementation. `reference.py` rebuilds it from diffusers and transformers building blocks:
+the DiT from `AdaLayerNormZero`, `AdaLayerNormContinuous`, `FeedForward` (gelu-approximate and swiglu), `Timesteps`,
+`TimestepEmbedding` and the diffusers RoPE helpers, using the checkpoint's module names so the checkpoint loads
+strictly. The pipeline is `reference_generate`, a step-by-step copy of the vLLM-Omni `PAN2Pipeline.forward`
+using `Qwen3_5TextModel`, `AutoencoderKLWan` and `FlowMatchEulerDiscreteScheduler`.
+
+The checkpoint is the public random-weight [`wuqing157/tiny-pan2-modular-pipe`](https://huggingface.co/wuqing157/tiny-pan2-modular-pipe)
+(revision `3ff3130b`), the same one vLLM-Omni tests against. It has two DiT blocks, a four-layer Qwen3.5 text model
+(three Gated DeltaNet layers and one full-attention layer) and a small Wan2.2-layout VAE.
+
+## Setup
+
+```bash
+huggingface-cli download wuqing157/tiny-pan2-modular-pipe --revision 3ff3130b4527f61aaab7bde01465c9884dc84cdd
+# or point the tests at a local copy:
+export PAN2_TINY_MODEL_DIR=/path/to/tiny-pan2-modular-pipe
+```
+
+The tests need `diffusers` with `AutoencoderKLWan` and `transformers` with `Qwen3_5TextModel` (5.x). They run on CPU
+when no GPU is present. On a GPU-less host, `import fastvideo` also needs the `fastvideo_kernel` stub described in
+`fastvideo/tests/AGENTS.md`.
+
+## Tests
+
+| Area | Command | Result (CPU, fp32 unless noted) |
+|---|---|---|
+| DiT: tiny checkpoint, strict load + forward | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k tiny_checkpoint` | max abs 6.0e-7 |
+| DiT: random weights, batch 1 and 2 | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k random_weights` | max abs 0 and 1.9e-5 (outputs ~16) |
+| DiT: Ulysses SP=2/4 vs one rank, T2V/I2V with 15/18 tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sequence_parallel` | max abs 4.8e-6 / 7.7e-6 (gloo on CPU; NCCL on 2/4 GPUs) |
+| Qwen3.5 text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | all hidden states exact (fp32, bf16, padded batch) |
+| Wan2.2 VAE encode/decode + normalization | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | exact; normalization 2.4e-7 |
+| Flow-match scheduler, PAN2 sigmas + shift | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | exact |
+| Pipeline, T2V and I2V (bf16 DiT) | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | final latents and video exact |
+
+The pipeline test builds the FastVideo pipeline through the registry and component loaders and compares the final
+latents (relative L2 < 2e-2, cosine > 0.9995) and the decoded video (PSNR > 40 dB) with `reference_generate`. To
+confirm the threshold catches real deviations, three injected changes were each run against it: conditioning on
+`hidden_states[-2]`, letting the latents fall to bf16 between steps, and shifting the first-frame crop by one pixel.
+Each moved the latents by about 8e-2 relative L2.
+
+Upstream revision `e0c4c731` always masks video padding, fixing Ulysses=4 drift when padding was left unmasked.
+FastVideo already excludes this padding by passing the original video length to `DistributedAttention`, which trims
+it before RoPE and joint attention, then pads the output for the reverse all-to-all. The distributed test covers one
+pad token (15 tokens at SP=2/4), two pad tokens (18 at SP=4, matching 704x1248x313 modulo 4), and an unpadded
+control (18 at SP=2), with both zero T2V and first-frame I2V conditioning. No masking toggle is required.
+The pipeline math and checkpoint layout are unchanged between upstream revisions `59dd8c9a` and `e0c4c731`.
+
+On 2026-10-20, all 39 PAN2 local tests passed on CPU (Torch 2.12.0, Transformers 5.15.0, Diffusers 0.39.0).
+The broader API/contract suite had 499 passes and two failures also reproduced on clean `main` at `cdda1b753`:
+the schema inventory lacks MiniMax-H3's `audio_scheduler_shift`/`video_scheduler_shift`, and
+`test_load_kandinsky5.py` reads `FASTVIDEO_ATTENTION_BACKEND` outside the environment registry.
+Importing the PAN2 tests and running their fixtures preserves both an unset and a pre-existing attention backend.
+
+## Not covered
+
+- Real-weight parity and an SSIM regression: `IFM/PAN2` is not public.
+- GPU runs: everything above ran on CPU. The bf16 GPU path, FlashAttention, and SP over NCCL still need a GPU node.
diff --git a/tests/local_tests/pan2/reference.py b/tests/local_tests/pan2/reference.py
new file mode 100644
index 000000000..41a207580
--- /dev/null
+++ b/tests/local_tests/pan2/reference.py
@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 reference implementation for local parity tests.
+
+The diffusers ``PAN2Transformer3DModel`` and ``PAN2ModularPipeline`` are not public yet. The reference here is
+assembled from diffusers' and transformers' own building blocks (``AdaLayerNormZero``, ``AdaLayerNormContinuous``,
+``FeedForward`` with ``gelu-approximate``/``swiglu``, ``Timesteps``/``TimestepEmbedding``,
+``get_1d_rotary_pos_embed``/``apply_rotary_emb``, ``Qwen3_5TextModel``, ``AutoencoderKLWan``,
+``FlowMatchEulerDiscreteScheduler``) and follows the vLLM-Omni PAN2 port (vllm-project/vllm-omni#8602), which was
+validated against the diffusers implementation (upstream revision ``e0c4c731``). Module names are the checkpoint's,
+so checkpoints load strictly. Its unconditional SP padding mask is equivalent to unpadded reference attention.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm
+
+TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe"
+TINY_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
+SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions."
+TEXT_ENCODER_LAYERS_TO_SKIP = 2
+MAX_SEQUENCE_LENGTH = 1536
+
+
+def tiny_model_dir() -> Path | None:
+ """``PAN2_TINY_MODEL_DIR``, else the pinned tiny checkpoint in the local HF cache."""
+ local = os.getenv("PAN2_TINY_MODEL_DIR")
+ if local:
+ return Path(local)
+ try:
+ from huggingface_hub import snapshot_download
+
+ return Path(snapshot_download(TINY_REPO_ID, revision=TINY_REVISION, local_files_only=True))
+ except Exception:
+ return None
+
+
+def _sdpa(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+ out = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
+ return out.transpose(1, 2).flatten(2)
+
+
+class _RefRefinerAttention(nn.Module):
+
+ def __init__(self, dim: int, heads: int, dim_head: int, eps: float):
+ super().__init__()
+ self.heads = heads
+ self.to_q = nn.Linear(dim, heads * dim_head, bias=False)
+ self.to_k = nn.Linear(dim, heads * dim_head, bias=False)
+ self.to_v = nn.Linear(dim, heads * dim_head, bias=False)
+ self.norm_q = RMSNorm(dim_head, eps=eps)
+ self.norm_k = RMSNorm(dim_head, eps=eps)
+ self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim, bias=False)])
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
+ key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1)))
+ value = self.to_v(x).unflatten(-1, (self.heads, -1))
+ return self.to_out[0](_sdpa(query, key, value))
+
+
+class _RefRefinerBlock(nn.Module):
+
+ def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, eps: float):
+ super().__init__()
+ self.norm1 = RMSNorm(dim, eps=eps)
+ self.attn = _RefRefinerAttention(dim, heads, dim_head, eps)
+ self.norm2 = RMSNorm(dim, eps=eps)
+ self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="swiglu", bias=False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = x + self.attn(self.norm1(x))
+ return x + self.ff(self.norm2(x))
+
+
+class _RefJointAttention(nn.Module):
+
+ def __init__(self, dim: int, heads: int, dim_head: int, eps: float):
+ super().__init__()
+ self.heads = heads
+ self.to_q = nn.Linear(dim, heads * dim_head)
+ self.to_k = nn.Linear(dim, heads * dim_head)
+ self.to_v = nn.Linear(dim, heads * dim_head)
+ self.norm_q = RMSNorm(dim_head, eps=eps)
+ self.norm_k = RMSNorm(dim_head, eps=eps)
+ self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim)])
+ self.add_q_proj = nn.Linear(dim, heads * dim_head)
+ self.add_k_proj = nn.Linear(dim, heads * dim_head)
+ self.add_v_proj = nn.Linear(dim, heads * dim_head)
+ self.norm_added_q = RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = RMSNorm(dim_head, eps=eps)
+ self.to_add_out = nn.Linear(heads * dim_head, dim)
+
+ def forward(self, x: torch.Tensor, context: torch.Tensor,
+ rope: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+ query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
+ key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1)))
+ value = self.to_v(x).unflatten(-1, (self.heads, -1))
+ # Interleaved (GPT-J) rotary embedding on the video tokens only.
+ query = apply_rotary_emb(query, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
+ key = apply_rotary_emb(key, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
+
+ context_query = self.norm_added_q(self.add_q_proj(context).unflatten(-1, (self.heads, -1)))
+ context_key = self.norm_added_k(self.add_k_proj(context).unflatten(-1, (self.heads, -1)))
+ context_value = self.add_v_proj(context).unflatten(-1, (self.heads, -1))
+
+ out = _sdpa(torch.cat([query, context_query], dim=1), torch.cat([key, context_key], dim=1),
+ torch.cat([value, context_value], dim=1))
+ out, context_out = out.split([x.shape[1], context.shape[1]], dim=1)
+ return self.to_out[0](out), self.to_add_out(context_out)
+
+
+class _RefBlock(nn.Module):
+
+ def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, context_mlp_ratio: float):
+ super().__init__()
+ eps = 1e-6
+ self.norm1 = AdaLayerNormZero(dim)
+ self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim))
+ self.attn = _RefJointAttention(dim, heads, dim_head, eps)
+ self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="gelu-approximate")
+ self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.ff_context = FeedForward(dim, inner_dim=int(dim * context_mlp_ratio), activation_fn="gelu-approximate")
+
+ def forward(self, x, context, temb, rope):
+ norm_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(x, emb=temb)
+ c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
+ self.scale_shift_table_context.chunk(6, dim=-1))
+ norm_context = self.norm1_context(context) * (1 + c_scale_msa) + c_shift_msa
+ attn_out, context_attn_out = self.attn(norm_x, norm_context, rope)
+ x = x + attn_out * gate_msa.unsqueeze(1)
+ context = context + context_attn_out * c_gate_msa
+ norm_x = self.norm2(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+ norm_context = self.norm2_context(context) * (1 + c_scale_mlp) + c_shift_mlp
+ x = x + self.ff(norm_x) * gate_mlp.unsqueeze(1)
+ context = context + self.ff_context(norm_context) * c_gate_mlp
+ return x, context
+
+
+class ReferencePAN2Transformer(nn.Module):
+
+ def __init__(self, patch_size=(1, 2, 2), in_channels=97, out_channels=48, num_attention_heads=32,
+ attention_head_dim=128, num_layers=44, num_refiner_layers=4, mlp_ratio=8.0, context_mlp_ratio=4.0,
+ refiner_mlp_ratio=4.0, text_embed_dim=4096, rope_axes_dim=(16, 56, 56), rope_theta=256.0,
+ rope_temporal_scale=16 / 24, **_):
+ super().__init__()
+ self.patch_size = tuple(patch_size)
+ self.rope_axes_dim = tuple(rope_axes_dim)
+ self.rope_theta = rope_theta
+ self.rope_temporal_scale = rope_temporal_scale
+ inner_dim = num_attention_heads * attention_head_dim
+ self.x_embedder = nn.Conv3d(in_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size)
+ self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.time_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim)
+ self.context_embedder = nn.Sequential(RMSNorm(text_embed_dim, eps=1e-5), nn.Linear(text_embed_dim, inner_dim))
+ self.context_refiner = nn.ModuleList([
+ _RefRefinerBlock(inner_dim, num_attention_heads, attention_head_dim, refiner_mlp_ratio, eps=1e-5)
+ for _ in range(num_refiner_layers)
+ ])
+ self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim))
+ self.transformer_blocks = nn.ModuleList([
+ _RefBlock(inner_dim, num_attention_heads, attention_head_dim, mlp_ratio, context_mlp_ratio)
+ for _ in range(num_layers)
+ ])
+ self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(inner_dim, self.patch_size[0] * self.patch_size[1] * self.patch_size[2] * out_channels)
+
+ def _rope(self, num_frames: int, height: int, width: int, device) -> tuple[torch.Tensor, torch.Tensor]:
+ 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.rope_temporal_scale
+ cos, sin = zip(*[
+ get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), self.rope_theta, use_real=True,
+ repeat_interleave_real=True) for i, dim in enumerate(self.rope_axes_dim)
+ ])
+ return torch.cat(cos, dim=1).float(), torch.cat(sin, dim=1).float()
+
+ def forward(self, hidden_states, timestep, encoder_hidden_states):
+ batch_size, _, num_frames, height, width = hidden_states.shape
+ p_t, p_h, p_w = self.patch_size
+ t, h, w = num_frames // p_t, height // p_h, width // p_w
+ rope = self._rope(t, h, w, hidden_states.device)
+ x = self.x_embedder(hidden_states).flatten(2).transpose(1, 2)
+ temb = self.time_embedder(self.time_proj(timestep).to(x.dtype))
+ context = self.context_embedder(encoder_hidden_states)
+ for block in self.context_refiner:
+ context = block(context)
+ context = context + self.context_type_embedding
+ for block in self.transformer_blocks:
+ x, context = block(x, context, temb, rope)
+ x = self.proj_out(self.norm_out(x, temb))
+ x = x.reshape(batch_size, t, h, w, -1, p_t, p_h, p_w).permute(0, 4, 1, 5, 2, 6, 3, 7)
+ return x.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+
+def reference_encode_prompt(tokenizer, text_encoder, prompt: str, device) -> torch.Tensor:
+ 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
+ header_ids = tokenizer("<|im_start|>user\n", 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)
+ input_ids = torch.tensor([input_ids[:crop_start + MAX_SEQUENCE_LENGTH]], device=device)
+ hidden_states = text_encoder(input_ids=input_ids, output_hidden_states=True).hidden_states
+ return hidden_states[-(TEXT_ENCODER_LAYERS_TO_SKIP + 1)][:, crop_start:]
+
+
+def reference_condition_latents(vae, latents: torch.Tensor, image: PIL.Image.Image | None, height: int,
+ width: int) -> torch.Tensor:
+ batch_size, num_channels, num_frames, latent_height, latent_width = latents.shape
+ condition = latents.new_zeros(batch_size, num_channels + 1, num_frames, latent_height, latent_width)
+ if image is None:
+ return condition
+ from diffusers.video_processor import VideoProcessor
+
+ 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))
+ pixels = VideoProcessor(vae_scale_factor=16).preprocess(image, height=height, width=width)
+ pixels = pixels.unsqueeze(2).to(device=latents.device, dtype=vae.dtype)
+ image_latents = vae.encode(pixels).latent_dist.mode()
+ mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(image_latents)
+ std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(image_latents)
+ condition[:, :num_channels, :1] = (image_latents - mean) / std
+ condition[:, num_channels:, :1] = 1.0
+ return condition
+
+
+@torch.no_grad()
+def reference_generate(components: dict, *, prompt: str, negative_prompt: str, image: PIL.Image.Image | None,
+ height: int, width: int, num_frames: int, num_inference_steps: int, guidance_scale: float,
+ seed: int, dit_dtype: torch.dtype, device="cpu",
+ skip_first_latent_cfg: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
+ """Final latents and the decoded video in ``[0, 1]``, following the vLLM-Omni pipeline step by step.
+
+ ``skip_first_latent_cfg`` is the reference sampler's ``--i2v_skip_first_latent_cfg``: with an image, the first
+ latent frame takes the positive prediction instead of the guided one.
+ """
+ from diffusers.utils.torch_utils import randn_tensor
+
+ tokenizer, text_encoder = components["tokenizer"], components["text_encoder"]
+ transformer, vae, scheduler = components["transformer"], components["vae"], components["scheduler"]
+ do_cfg = guidance_scale > 1.0
+ prompt_embeds = reference_encode_prompt(tokenizer, text_encoder, prompt, device).to(dit_dtype)
+ negative_embeds = (reference_encode_prompt(tokenizer, text_encoder, negative_prompt, device).to(dit_dtype)
+ if do_cfg else None)
+
+ shape = (1, vae.config.z_dim, (num_frames - 1) // vae.config.scale_factor_temporal + 1,
+ height // vae.config.scale_factor_spatial, width // vae.config.scale_factor_spatial)
+ generator = [torch.Generator("cpu").manual_seed(seed)]
+ latents = randn_tensor(shape, generator=generator, device=torch.device(device), dtype=torch.float32)
+ condition = reference_condition_latents(vae, latents, image, height, width).to(dit_dtype)
+
+ scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device)
+ for t in scheduler.timesteps:
+ model_input = torch.cat([latents.to(dit_dtype), condition], dim=1)
+ timestep = t.expand(model_input.shape[0])
+ noise_pred = transformer(model_input, timestep, prompt_embeds)
+ if do_cfg:
+ noise_uncond = transformer(model_input, timestep, negative_embeds)
+ guided = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
+ if image is not None and skip_first_latent_cfg:
+ guided[:, :, :1] = noise_pred[:, :, :1]
+ noise_pred = guided
+ latents = scheduler.step(noise_pred.float(), t, latents, return_dict=False)[0]
+
+ mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents)
+ std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents)
+ video = vae.decode((latents * std + mean).to(vae.dtype), return_dict=False)[0]
+ return latents, (video.float() / 2 + 0.5).clamp(0, 1)
+
+
+def load_reference_components(model_dir: Path, dit_dtype: torch.dtype) -> dict:
+ import json
+
+ from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler
+ from safetensors.torch import load_file
+ from transformers import AutoTokenizer, Qwen3_5TextModel
+
+ transformer_dir = model_dir / "transformer"
+ transformer = ReferencePAN2Transformer(**json.loads((transformer_dir / "config.json").read_text()))
+ transformer.load_state_dict(load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors")), strict=True)
+ return {
+ "tokenizer": AutoTokenizer.from_pretrained(model_dir / "tokenizer"),
+ "text_encoder": Qwen3_5TextModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.float32).eval(),
+ "transformer": transformer.to(dit_dtype).eval(),
+ "vae": AutoencoderKLWan.from_pretrained(model_dir / "vae", torch_dtype=torch.float32).eval(),
+ "scheduler": FlowMatchEulerDiscreteScheduler.from_pretrained(model_dir / "scheduler"),
+ }
diff --git a/tests/local_tests/pipelines/test_pan2_pipeline_parity.py b/tests/local_tests/pipelines/test_pan2_pipeline_parity.py
new file mode 100644
index 000000000..e67ca7d6a
--- /dev/null
+++ b/tests/local_tests/pipelines/test_pan2_pipeline_parity.py
@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 end-to-end pipeline parity on the public tiny checkpoint.
+
+The FastVideo pipeline (built through the real registry and component loaders) is compared against
+``tests/local_tests/pan2/reference.py::reference_generate``, which follows the vLLM-Omni ``PAN2Pipeline.forward``
+step by step with transformers/diffusers components: prompt template and crop, fp32 noise, first-frame
+conditioning, the uniform sigma schedule with shift 7, sequential CFG, fp32 Euler steps and the VAE decode.
+
+Both sides run the DiT in bf16 (FastVideo's denoiser always feeds it bf16), so the comparison allows bf16 rounding.
+Runs on CPU when no GPU is present. Set ``PAN2_TINY_MODEL_DIR`` to a local copy of the tiny checkpoint.
+"""
+
+from __future__ import annotations
+
+import socket
+from contextlib import ExitStack
+
+import PIL.Image
+import pytest
+import torch
+
+from fastvideo import envs
+from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir
+
+HEIGHT, WIDTH, NUM_FRAMES = 64, 96, 9
+NUM_INFERENCE_STEPS = 3
+# A large scale makes the guided steps visible through the tiny random-weight VAE.
+GUIDANCE_SCALE = 30.0
+SEED = 42
+PROMPT = "a cat walks on the grass"
+NEGATIVE_PROMPT = "blurry"
+# Not the output aspect ratio, so the cover-resize and center crop are exercised.
+CONDITION_IMAGE = PIL.Image.radial_gradient("L").resize((120, 64)).convert("RGB")
+
+
+def _free_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("127.0.0.1", 0))
+ return sock.getsockname()[1]
+
+
+@pytest.fixture(scope="module")
+def pan2_runtime():
+ model_dir = tiny_model_dir()
+ if model_dir is None or not (model_dir / "modular_model_index.json").is_file():
+ pytest.skip("wuqing157/tiny-pan2-modular-pipe is not available locally; set PAN2_TINY_MODEL_DIR")
+ from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
+ from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+ import fastvideo.distributed.parallel_state as parallel_state
+ from fastvideo.fastvideo_args import FastVideoArgs
+ from fastvideo.pipelines import build_pipeline
+
+ with ExitStack() as stack:
+ stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+ monkeypatch = stack.enter_context(pytest.MonkeyPatch.context())
+ if not torch.cuda.is_available():
+ # Reach modules that imported the device helper by name on a GPU-less Linux host.
+ def _cpu_device() -> torch.device:
+ return torch.device("cpu")
+
+ monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__)
+ for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+ ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+ stack.enter_context(envs.override_external(name, value))
+ stack.callback(cleanup_dist_env_and_memory)
+ maybe_init_distributed_environment_and_model_parallel(1, 1)
+
+ # The parity checkpoint has random weights, so its videos are noise: compare without the guardrails.
+ pipeline_config = PAN2PipelineConfig(dit_precision="bf16",
+ vae_precision="fp32",
+ text_encoder_precisions=("fp32", ),
+ enable_guardrails=False)
+ args = FastVideoArgs(model_path=str(model_dir),
+ num_gpus=1,
+ pipeline_config=pipeline_config,
+ dit_cpu_offload=False,
+ dit_layerwise_offload=False,
+ text_encoder_cpu_offload=False,
+ vae_cpu_offload=False,
+ use_fsdp_inference=False)
+ pipeline = build_pipeline(args)
+ device = parallel_state.get_local_torch_device()
+ yield pipeline, args, model_dir, device
+
+
+def _fastvideo_generate(pipeline, args, image, guidance_scale: float = GUIDANCE_SCALE):
+ from fastvideo.pipelines import ForwardBatch
+
+ batch = ForwardBatch(data_type="video",
+ prompt=PROMPT,
+ negative_prompt=NEGATIVE_PROMPT,
+ seed=SEED,
+ height=HEIGHT,
+ width=WIDTH,
+ num_frames=NUM_FRAMES,
+ num_inference_steps=NUM_INFERENCE_STEPS,
+ guidance_scale=guidance_scale,
+ pil_image=image)
+ with torch.inference_mode():
+ batch = pipeline.forward(batch, args)
+ return batch.latents.float().cpu(), batch.output.float().cpu()
+
+
+def _psnr(a: torch.Tensor, b: torch.Tensor) -> float:
+ mse = (a - b).pow(2).mean().clamp_min(1e-12)
+ return float(10 * torch.log10(1.0 / mse))
+
+
+@pytest.mark.parametrize(
+ ("image", "skip_first_latent_cfg"),
+ [(None, True), (CONDITION_IMAGE, True), (CONDITION_IMAGE, False)],
+ ids=["t2v", "i2v", "i2v-guided-first-frame"],
+)
+def test_pan2_pipeline_matches_reference(pan2_runtime, monkeypatch, image, skip_first_latent_cfg):
+ pipeline, args, model_dir, device = pan2_runtime
+ monkeypatch.setattr(args.pipeline_config, "i2v_skip_first_latent_cfg", skip_first_latent_cfg)
+ latents, video = _fastvideo_generate(pipeline, args, image)
+
+ components = load_reference_components(model_dir, torch.bfloat16)
+ components = {
+ name: module.to(device) if isinstance(module, torch.nn.Module) else module
+ for name, module in components.items()
+ }
+ ref_latents, ref_video = reference_generate(components,
+ prompt=PROMPT,
+ negative_prompt=NEGATIVE_PROMPT,
+ image=image,
+ height=HEIGHT,
+ width=WIDTH,
+ num_frames=NUM_FRAMES,
+ num_inference_steps=NUM_INFERENCE_STEPS,
+ guidance_scale=GUIDANCE_SCALE,
+ seed=SEED,
+ dit_dtype=torch.bfloat16,
+ device=device,
+ skip_first_latent_cfg=skip_first_latent_cfg)
+ ref_latents, ref_video = ref_latents.float().cpu(), ref_video.float().cpu()
+
+ assert latents.dtype == torch.float32
+ assert latents.shape == ref_latents.shape == (1, 48, 3, HEIGHT // 16, WIDTH // 16)
+ assert video.shape == ref_video.shape == (1, 3, NUM_FRAMES, HEIGHT, WIDTH)
+ latent_rel = float((latents - ref_latents).norm() / ref_latents.norm())
+ cosine = float(torch.nn.functional.cosine_similarity(latents.flatten(), ref_latents.flatten(), dim=0))
+ psnr = _psnr(video, ref_video)
+ print(f"latents rel_l2={latent_rel:.3e} cosine={cosine:.6f}; video psnr={psnr:.1f} dB")
+ assert latent_rel < 2e-2
+ assert cosine > 0.9995
+ assert psnr > 40.0
+
+
+def test_pan2_image_conditioning_and_guidance_change_every_frame(pan2_runtime):
+ pipeline, args, _, _ = pan2_runtime
+ _, t2v = _fastvideo_generate(pipeline, args, None)
+ _, i2v = _fastvideo_generate(pipeline, args, CONDITION_IMAGE)
+ _, unguided = _fastvideo_generate(pipeline, args, None, guidance_scale=1.0)
+
+ def per_frame_max_diff(a, b):
+ return (a - b).abs().amax(dim=(0, 1, 3, 4))
+
+ assert per_frame_max_diff(t2v, i2v).max() > 1e-3
+ assert per_frame_max_diff(t2v, unguided).min() > 1e-3
diff --git a/tests/local_tests/schedulers/test_pan2_scheduler_parity.py b/tests/local_tests/schedulers/test_pan2_scheduler_parity.py
new file mode 100644
index 000000000..3fcdb7497
--- /dev/null
+++ b/tests/local_tests/schedulers/test_pan2_scheduler_parity.py
@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 flow-matching Euler schedule on FastVideo's FlowMatchEulerDiscreteScheduler.
+
+Coverage scope: both. PAN2 loads Diffusers' ``FlowMatchEulerDiscreteScheduler``
+from ``/scheduler`` (shift 7.0, no dynamic shifting), optionally
+overrides the shift with ``flow_shift``, calls
+``set_timesteps(sigmas=np.linspace(1.0, 0.0, N + 1)[:-1])`` and steps float32
+latents. The FastVideo side is loaded by the production ``SchedulerLoader`` and
+its schedule is built by the stock ``TimestepPreparationStage`` from
+``batch.sigmas``. Timesteps, sigmas, and every Euler update must be bit-exact.
+
+The default asset is the tiny PAN2 checkpoint (config only), so this runs on a
+CPU-only host::
+
+ pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py -v -s
+
+Set ``PAN2_MODEL_PATH`` to a local PAN2 root containing ``scheduler/`` to use
+another checkpoint.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from torch.testing import assert_close
+
+try:
+ from diffusers import FlowMatchEulerDiscreteScheduler as DiffusersFlowMatchEulerDiscreteScheduler
+except ImportError: # pragma: no cover - environment dependent
+ pytest.skip("diffusers with FlowMatchEulerDiscreteScheduler is required", allow_module_level=True)
+
+from fastvideo.models.loader.component_loader import SchedulerLoader
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+ FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+
+PARITY_SCOPE = "both"
+TINY_PAN2_REPO = "wuqing157/tiny-pan2-modular-pipe"
+TINY_PAN2_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
+MODEL_PATH_ENV = "PAN2_MODEL_PATH"
+LATENT_SHAPE = (1, 48, 3, 4, 6)
+
+
+def _pan2_root(subfolder: str) -> Path:
+ """Resolve a PAN2 root that contains ``subfolder``; skip when unavailable."""
+ override = os.getenv(MODEL_PATH_ENV)
+ if override:
+ root = Path(override)
+ else:
+ try:
+ from huggingface_hub import snapshot_download
+ except ImportError:
+ pytest.skip("huggingface_hub is required to resolve the tiny PAN2 checkpoint")
+ root = None
+ last_error: Exception | None = None
+ for local_files_only in (True, False):
+ try:
+ root = Path(
+ snapshot_download(
+ TINY_PAN2_REPO,
+ revision=TINY_PAN2_REVISION,
+ allow_patterns=[f"{subfolder}/*"],
+ local_files_only=local_files_only,
+ ))
+ break
+ except Exception as exc: # noqa: BLE001 - offline/missing cache must skip.
+ last_error = exc
+ if root is None:
+ pytest.skip(f"Tiny PAN2 checkpoint unavailable: {last_error}")
+ if not (root / subfolder).is_dir():
+ pytest.skip(f"PAN2 {subfolder}/ missing under {root}")
+ return root
+
+
+@pytest.fixture(scope="module")
+def pan2_root() -> Path:
+ root = _pan2_root("scheduler")
+ if not (root / "scheduler" / "scheduler_config.json").is_file():
+ pytest.skip(f"PAN2 scheduler config missing under {root}")
+ return root
+
+
+@pytest.fixture
+def cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Pin FastVideo's device helper to CPU; it reports ``mps`` on CPU-only Linux."""
+ monkeypatch.setattr("fastvideo.pipelines.stages.timestep_preparation.get_local_torch_device",
+ lambda: torch.device("cpu"))
+
+
+def _pan2_sigmas(num_steps: int) -> np.ndarray:
+ """PAN2's pre-shift grid: uniform from 1 down to, but excluding, 0."""
+ return np.linspace(1.0, 0.0, num_steps + 1)[:-1]
+
+
+def _reference(root: Path, flow_shift: float | None = None) -> DiffusersFlowMatchEulerDiscreteScheduler:
+ """Mirror the upstream PAN2 pipeline's scheduler construction."""
+ scheduler = DiffusersFlowMatchEulerDiscreteScheduler.from_pretrained(str(root),
+ subfolder="scheduler",
+ local_files_only=True)
+ if flow_shift is not None:
+ scheduler._shift = flow_shift
+ return scheduler
+
+
+def _fastvideo(root: Path, flow_shift: float | None = None) -> FastVideoFlowMatchEulerDiscreteScheduler:
+ """Load through the production SchedulerLoader, which applies ``flow_shift``."""
+ fastvideo_args = SimpleNamespace(pipeline_config=SimpleNamespace(flow_shift=flow_shift))
+ scheduler = SchedulerLoader().load(str(root / "scheduler"), fastvideo_args)
+ assert type(scheduler) is FastVideoFlowMatchEulerDiscreteScheduler
+ return scheduler
+
+
+def _prepare_fastvideo_timesteps(scheduler: FastVideoFlowMatchEulerDiscreteScheduler,
+ num_steps: int) -> ForwardBatch:
+ """Drive the stock TimestepPreparationStage with PAN2's custom sigmas."""
+ batch = ForwardBatch(data_type="video", num_inference_steps=num_steps, sigmas=_pan2_sigmas(num_steps).tolist())
+ stage = TimestepPreparationStage(scheduler)
+ assert stage.verify_input(batch, SimpleNamespace()).is_valid()
+ batch = stage.forward(batch, SimpleNamespace())
+ assert batch.timesteps is scheduler.timesteps
+ return batch
+
+
+def test_pan2_scheduler_loader_keeps_checkpoint_config(pan2_root: Path) -> None:
+ """Without ``flow_shift`` the loader keeps every released config value, including shift 7.0."""
+ raw = json.loads((pan2_root / "scheduler" / "scheduler_config.json").read_text())
+ assert raw["_class_name"] == "FlowMatchEulerDiscreteScheduler"
+ scheduler = _fastvideo(pan2_root)
+ for key, value in raw.items():
+ if not key.startswith("_"):
+ assert scheduler.config[key] == value, key
+ assert scheduler.shift == raw["shift"] == 7.0
+ assert not scheduler.config.use_dynamic_shifting
+
+
+@pytest.mark.parametrize("num_steps", [1, 4, 8, 30, 50])
+def test_pan2_scheduler_schedule_matches_diffusers(num_steps: int, pan2_root: Path, cpu_runtime: None) -> None:
+ """Timesteps and sigmas from ``batch.sigmas`` are bit-exact for short and long schedules."""
+ reference = _reference(pan2_root)
+ reference.set_timesteps(sigmas=_pan2_sigmas(num_steps), device="cpu")
+ fastvideo = _fastvideo(pan2_root)
+ batch = _prepare_fastvideo_timesteps(fastvideo, num_steps)
+
+ print(f"[pan2-scheduler] N={num_steps} timesteps[:4]={reference.timesteps[:4].tolist()} "
+ f"sigmas[-3:]={reference.sigmas[-3:].tolist()}")
+ assert len(batch.timesteps) == num_steps
+ assert_close(batch.timesteps, reference.timesteps, atol=0, rtol=0)
+ assert_close(fastvideo.sigmas, reference.sigmas, atol=0, rtol=0)
+ assert fastvideo.sigmas[0].item() == 1.0 and fastvideo.sigmas[-1].item() == 0.0
+
+
+@pytest.mark.parametrize("flow_shift", [None, 5.0], ids=["checkpoint_shift", "flow_shift_override"])
+def test_pan2_scheduler_steps_match_diffusers(flow_shift: float | None, pan2_root: Path, cpu_runtime: None) -> None:
+ """Five float32 Euler updates are bit-exact, with and without a shift override."""
+ num_steps = 5
+ reference = _reference(pan2_root, flow_shift)
+ reference.set_timesteps(sigmas=_pan2_sigmas(num_steps), device="cpu")
+ fastvideo = _fastvideo(pan2_root, flow_shift)
+ _prepare_fastvideo_timesteps(fastvideo, num_steps)
+ assert fastvideo.shift == reference.shift == (flow_shift or 7.0)
+ assert_close(fastvideo.timesteps, reference.timesteps, atol=0, rtol=0)
+
+ generator = torch.Generator().manual_seed(2026)
+ reference_latents = torch.randn(LATENT_SHAPE, generator=generator)
+ fastvideo_latents = reference_latents.clone()
+ max_diff = 0.0
+ for index in range(num_steps):
+ noise_pred = torch.randn(LATENT_SHAPE, generator=generator)
+ reference_latents = reference.step(noise_pred, reference.timesteps[index], reference_latents).prev_sample
+ fastvideo_latents = fastvideo.step(noise_pred, fastvideo.timesteps[index], fastvideo_latents).prev_sample
+ assert fastvideo_latents.dtype == torch.float32
+ max_diff = max(max_diff, (fastvideo_latents - reference_latents).abs().max().item())
+ assert_close(fastvideo_latents, reference_latents, atol=0, rtol=0)
+ print(f"[pan2-scheduler] flow_shift={flow_shift} steps={num_steps} max_abs={max_diff:.3e}")
+ assert fastvideo.step_index == reference.step_index == num_steps
+
+
+def test_pan2_scheduler_step_returns_model_output_dtype(pan2_root: Path, cpu_runtime: None) -> None:
+ """Both implementations cast ``prev_sample`` to the model output dtype.
+
+ PAN2 keeps float32 latents by stepping with ``noise_pred.float()``; a bf16
+ noise prediction would silently drop the latents to bf16 on either side.
+ """
+ reference = _reference(pan2_root)
+ reference.set_timesteps(sigmas=_pan2_sigmas(2), device="cpu")
+ fastvideo = _fastvideo(pan2_root)
+ _prepare_fastvideo_timesteps(fastvideo, 2)
+ latents = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(5))
+ noise_pred = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(6)).to(torch.bfloat16)
+ expected = reference.step(noise_pred, reference.timesteps[0], latents).prev_sample
+ actual = fastvideo.step(noise_pred, fastvideo.timesteps[0], latents).prev_sample
+ assert expected.dtype == actual.dtype == torch.bfloat16
+ assert_close(actual, expected, atol=0, rtol=0)
diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py
index f0a2e11e8..2f4e37822 100644
--- a/tests/local_tests/test_cookbook_serving.py
+++ b/tests/local_tests/test_cookbook_serving.py
@@ -264,3 +264,23 @@ def test_kandinsky6_cookbook_recipes_match_maintained_examples():
assert all("INPUT_VIDEO" in by_id[recipe_id]["command"] for recipe_id in ("kandinsky6-vsr", "kandinsky6-vsr-distilled"))
assert "VSR-distilled2steps" in by_id["kandinsky6-vsr-distilled"]["command"]
validate_cookbook()
+
+
+def test_pan2_recipe_serves_text_and_image_from_one_deployment():
+ recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"]
+ recipe = next(item for item in recipes if item["family"] == "pan2")
+ assert recipe["id"] == "pan2-ti2v"
+ assert recipe["model"] == "IFM/PAN2"
+ assert recipe["evidence"] == "Source-backed"
+ assert recipe["hardware"] == {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}
+ profile = cookbook_serving_profile(recipe)
+ assert profile["command"] == "fastvideo serve --config examples/serving/openai_pan2.yaml --server.host 127.0.0.1"
+ assert profile["model"] == "pan2"
+ assert profile["playground_url"] is None
+ assert profile["sampling"]["num_frames"] == 313
+ assert any("image_reference" in item for item in recipe["serving"]["limitations"])
+ manifest = yaml.safe_load((ROOT / "docs/cookbook/recipes/pan2.yaml").read_text())
+ deployment = manifest["deployments"]["cuda-rest"]
+ assert deployment["workload"] == "t2v"
+ assert deployment["defaults"] == recipe["serving"]["source"]
+ validate_cookbook()
diff --git a/tests/local_tests/transformers/test_pan2_transformer_parity.py b/tests/local_tests/transformers/test_pan2_transformer_parity.py
new file mode 100644
index 000000000..4df532f06
--- /dev/null
+++ b/tests/local_tests/transformers/test_pan2_transformer_parity.py
@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 transformer weight-schema and forward parity.
+
+The reference is ``tests/local_tests/pan2/reference.py``: diffusers building blocks arranged as the (not yet public)
+diffusers ``PAN2Transformer3DModel``, following the vLLM-Omni port. Its module names are the checkpoint's, so the
+public tiny checkpoint loads into it strictly.
+
+Set ``PAN2_TINY_MODEL_DIR`` to a local copy of ``wuqing157/tiny-pan2-modular-pipe``; otherwise the HF cache is
+used and the checkpoint test skips when it is absent.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import socket
+import subprocess
+import sys
+from contextlib import ExitStack
+from dataclasses import fields
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.testing import assert_close
+
+from fastvideo import envs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict
+from fastvideo.models.loader.utils import get_param_names_mapping
+from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+SP_WORLD_SIZES = (2, 4)
+
+
+def _free_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("127.0.0.1", 0))
+ return sock.getsockname()[1]
+
+
+def _small_kwargs() -> dict:
+ """Larger than the tiny checkpoint: several heads, asymmetric RoPE axes, more blocks, distinct MLP ratios."""
+ return {
+ "patch_size": (1, 2, 2),
+ "in_channels": 9,
+ "out_channels": 4,
+ "num_attention_heads": 4,
+ "attention_head_dim": 32,
+ "num_layers": 3,
+ "num_refiner_layers": 2,
+ "mlp_ratio": 3.0,
+ "context_mlp_ratio": 2.0,
+ "refiner_mlp_ratio": 1.5,
+ "text_embed_dim": 48,
+ "rope_axes_dim": (8, 12, 12),
+ "rope_theta": 256.0,
+ "rope_temporal_scale": 16 / 24,
+ }
+
+
+@pytest.fixture
+def single_rank_dist(monkeypatch):
+ from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+ import fastvideo.distributed.parallel_state as parallel_state
+
+ if not torch.cuda.is_available():
+ def _cpu_device() -> torch.device:
+ return torch.device("cpu")
+
+ monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__)
+ with ExitStack() as stack:
+ stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+ for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"),
+ ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
+ stack.enter_context(envs.override_external(name, value))
+ try:
+ maybe_init_distributed_environment_and_model_parallel(1, 1)
+ yield
+ finally:
+ cleanup_dist_env_and_memory()
+
+
+def _build_native(arch_kwargs: dict):
+ from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
+ from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
+
+ config = PAN2VideoConfig()
+ config.update_model_arch(arch_kwargs)
+ return PAN2Transformer3DModel(config=config, hf_config=arch_kwargs).float().eval()
+
+
+def _load_into_native(native, weights) -> None:
+ incompatible = load_model_from_full_model_state_dict(
+ native,
+ iter(weights),
+ device=torch.device("cpu"),
+ param_dtype=torch.float32,
+ strict=True,
+ param_names_mapping=get_param_names_mapping(native.param_names_mapping),
+ training_mode=False,
+ )
+ assert incompatible.missing_keys == []
+ assert incompatible.unexpected_keys == []
+
+
+def _randomize(module: nn.Module, seed: int) -> None:
+ generator = torch.Generator().manual_seed(seed)
+ with torch.no_grad():
+ for param in module.parameters():
+ param.copy_(torch.randn(param.shape, generator=generator) * 0.2)
+
+
+def _make_inputs(in_channels: int, text_dim: int, batch_size: int = 2, seed: int = 0) -> dict[str, torch.Tensor]:
+ generator = torch.Generator().manual_seed(seed)
+ return {
+ # 9 frames at 64x96 in latent space: 3 latent frames of 4x6.
+ "hidden_states": torch.randn(batch_size, in_channels, 3, 4, 6, generator=generator),
+ "timestep": torch.tensor([700.0, 125.5][:batch_size]),
+ "encoder_hidden_states": torch.randn(batch_size, 13, text_dim, generator=generator),
+ }
+
+
+def _compare(native, reference, inputs: dict[str, torch.Tensor], scope: str) -> None:
+ with torch.inference_mode():
+ expected = reference(**inputs)
+ with set_forward_context(current_timestep=0, attn_metadata=None):
+ actual = native(**inputs)
+ assert actual.shape == expected.shape
+ diff = (actual - expected).abs()
+ print(f"{scope}: shape={tuple(actual.shape)} ref_abs_mean={expected.abs().mean().item():.3e} "
+ f"diff_max={diff.max().item():.3e} diff_mean={diff.mean().item():.3e}")
+ assert torch.isfinite(expected).all() and expected.abs().mean() > 1e-3
+ assert_close(actual, expected, atol=2e-5, rtol=2e-5)
+
+
+def test_pan2_transformer_registry_resolves_native_class():
+ from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
+ from fastvideo.models.registry import ModelRegistry
+
+ model_cls, architecture = ModelRegistry.resolve_model_cls("PAN2Transformer3DModel")
+ assert model_cls is PAN2Transformer3DModel
+ assert architecture == "PAN2Transformer3DModel"
+
+
+@pytest.mark.parametrize(("module_name", "expected"), [
+ ("transformer_blocks.0", True),
+ ("transformer_blocks.43", True),
+ ("context_refiner.3", True),
+ ("transformer_blocks.0.attn.to_qkv", False),
+ ("transformer_blocks.0.ff.fc_in", False),
+])
+def test_pan2_fsdp_shard_conditions_match_only_blocks(module_name: str, expected: bool):
+ from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig
+
+ conditions = PAN2ArchConfig()._fsdp_shard_conditions
+ assert any(condition(module_name, None) for condition in conditions) is expected
+
+
+def test_pan2_default_arch_config_is_the_29b_model():
+ from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig
+
+ config = PAN2ArchConfig()
+ assert config.hidden_size == 4096
+ assert config.num_channels_latents == 48
+ assert config.in_channels == 2 * 48 + 1
+ assert sum(config.rope_axes_dim) == config.attention_head_dim
+
+
+def test_pan2_tiny_checkpoint_config_fields_are_arch_fields():
+ model_dir = tiny_model_dir()
+ if model_dir is None or not (model_dir / "transformer" / "config.json").is_file():
+ pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR")
+ from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig
+
+ checkpoint_config = json.loads((model_dir / "transformer" / "config.json").read_text())
+ assert checkpoint_config.pop("_class_name") == "PAN2Transformer3DModel"
+ checkpoint_config.pop("_diffusers_version", None)
+ assert set(checkpoint_config) <= {field.name for field in fields(PAN2ArchConfig)}
+
+
+def test_pan2_tiny_checkpoint_strict_load_and_forward_parity(single_rank_dist):
+ model_dir = tiny_model_dir()
+ if model_dir is None or not (model_dir / "transformer").is_dir():
+ pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR")
+ from safetensors.torch import load_file
+
+ transformer_dir = model_dir / "transformer"
+ arch_kwargs = json.loads((transformer_dir / "config.json").read_text())
+ state_dict = load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors"))
+
+ reference = ReferencePAN2Transformer(**arch_kwargs).float().eval()
+ reference.load_state_dict(state_dict, strict=True)
+ native = _build_native(arch_kwargs)
+ _load_into_native(native, state_dict.items())
+
+ inputs = _make_inputs(arch_kwargs["in_channels"], arch_kwargs["text_embed_dim"])
+ _compare(native, reference, inputs, "tiny checkpoint")
+
+
+def test_pan2_layerwise_offload_hooks_the_main_blocks(single_rank_dist):
+ native = _build_native(_small_kwargs())
+
+ first_module_list = next(child for child in native.children() if isinstance(child, torch.nn.ModuleList))
+
+ assert first_module_list is native.transformer_blocks
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+def test_pan2_random_weights_forward_parity(single_rank_dist, batch_size: int):
+ kwargs = _small_kwargs()
+ reference = ReferencePAN2Transformer(**kwargs).float().eval()
+ # Zero-initialized tables (text modulation, type embedding) must be exercised too.
+ _randomize(reference, seed=1234)
+ native = _build_native(kwargs)
+ _load_into_native(native, reference.state_dict().items())
+
+ inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"], batch_size=batch_size, seed=batch_size)
+ _compare(native, reference, inputs, f"random weights, batch {batch_size}")
+
+
+def _run_sp_worker(sp_size: int, output_path: Path) -> None:
+ from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig
+ from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+ from fastvideo.models.dits.pan2 import PAN2Transformer3DModel
+ import torch.distributed as dist
+
+ if torch.cuda.is_available():
+ device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+ torch.cuda.set_device(device)
+ torch.backends.cuda.matmul.allow_tf32 = False
+ torch.backends.cudnn.allow_tf32 = False
+ else:
+ # GPU-less host: gloo ranks on CPU. fastvideo resolves the local device to mps there; swapping the code
+ # object reaches every module that imported get_local_torch_device by name.
+ import fastvideo.distributed.parallel_state as parallel_state
+
+ def _cpu_device() -> torch.device:
+ return torch.device("cpu")
+
+ parallel_state.get_local_torch_device.__code__ = _cpu_device.__code__
+ device = torch.device("cpu")
+ try:
+ maybe_init_distributed_environment_and_model_parallel(1, sp_size)
+ kwargs = _small_kwargs()
+ reference = ReferencePAN2Transformer(**kwargs).float()
+ _randomize(reference, seed=1234)
+ config = PAN2VideoConfig()
+ config.update_model_arch(kwargs)
+ native = PAN2Transformer3DModel(config=config, hf_config=kwargs).float().eval()
+ _load_into_native(native, reference.state_dict().items())
+ native = native.to(device)
+ outputs = {}
+ # 15 tokens need one pad at SP=2/4; 18 need two at SP=4, as at 704x1248x313.
+ for latent_frames in (5, 6):
+ for image_conditioned in (False, True):
+ inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"])
+ video = torch.randn(2, kwargs["in_channels"], latent_frames, 2, 6,
+ generator=torch.Generator().manual_seed(5))
+ video[:, kwargs["out_channels"]:] = 0
+ if image_conditioned:
+ video[:, kwargs["out_channels"]:-1, :1] = torch.randn(
+ 2, kwargs["out_channels"], 1, 2, 6, generator=torch.Generator().manual_seed(6))
+ video[:, -1:, :1] = 1
+ inputs["hidden_states"] = video
+ inputs = {name: value.to(device) for name, value in inputs.items()}
+ with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+ output = native(**inputs)
+ outputs[f"{latent_frames}_{image_conditioned}"] = output.detach().cpu()
+ if int(os.environ.get("RANK", "0")) == 0:
+ torch.save(outputs, output_path)
+ dist.barrier()
+ finally:
+ cleanup_dist_env_and_memory()
+
+
+def _run_torchrun(nproc_per_node: int, output_path: Path) -> None:
+ command = [
+ sys.executable, "-m", "torch.distributed.run", "--nnodes", "1", "--nproc_per_node",
+ str(nproc_per_node), "--master_port",
+ str(_free_port()),
+ str(Path(__file__).resolve()), "--pan2-sp-worker", "--sp-size", str(nproc_per_node), "--output",
+ str(output_path)
+ ]
+ environment = os.environ.copy()
+ environment["PYTHONPATH"] = os.pathsep.join(filter(None, [str(REPO_ROOT), environment.get("PYTHONPATH")]))
+ environment["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA"
+ process = subprocess.run(command, capture_output=True, text=True, env=environment, timeout=180)
+ if process.returncode != 0:
+ raise RuntimeError(f"SP={nproc_per_node} worker failed with code {process.returncode}\n"
+ f"STDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}")
+
+
+@pytest.fixture(scope="module")
+def single_rank_sp_outputs(tmp_path_factory):
+ path = tmp_path_factory.mktemp("pan2-single") / "single.pt"
+ _run_torchrun(1, path)
+ return torch.load(path, map_location="cpu", weights_only=True)
+
+
+@pytest.mark.parametrize("sp_size", SP_WORLD_SIZES)
+def test_pan2_sequence_parallel_matches_single_rank(tmp_path: Path, single_rank_sp_outputs, sp_size: int):
+ """T2V and I2V joint attention excludes padding at Ulysses SP=2/4. Uses gloo on CPU or NCCL on 2/4 GPUs."""
+ if torch.cuda.is_available() and torch.cuda.device_count() < sp_size:
+ pytest.skip(f"PAN2 SP parity requires {sp_size} CUDA devices.")
+ sp_path = tmp_path / f"sp{sp_size}.pt"
+ _run_torchrun(sp_size, sp_path)
+ sp_outputs = torch.load(sp_path, map_location="cpu", weights_only=True)
+ assert sp_outputs.keys() == single_rank_sp_outputs.keys()
+ for case, single in single_rank_sp_outputs.items():
+ sp = sp_outputs[case]
+ diff = (sp - single).abs()
+ print(f"sp{sp_size}-vs-single {case} diff_max={diff.max().item():.3e}")
+ assert_close(sp, single, atol=1e-5, rtol=1e-5)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pan2-sp-worker", action="store_true")
+ parser.add_argument("--sp-size", type=int, choices=[1, *SP_WORLD_SIZES], required=True)
+ parser.add_argument("--output", type=str, required=True)
+ args = parser.parse_args()
+ if not args.pan2_sp_worker:
+ raise SystemExit("This module is intended to be run by pytest.")
+ _run_sp_worker(args.sp_size, Path(args.output))
diff --git a/tests/local_tests/vaes/test_pan2_vae_parity.py b/tests/local_tests/vaes/test_pan2_vae_parity.py
new file mode 100644
index 000000000..3f56210f3
--- /dev/null
+++ b/tests/local_tests/vaes/test_pan2_vae_parity.py
@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 reuse of FastVideo's Wan2.2 (48-channel) VAE.
+
+Coverage scope: both. PAN2 ships a Diffusers ``AutoencoderKLWan`` in the
+Wan2.2 layout (z_dim 48, residual blocks, 2x2 patchify, 16x spatial and 4x
+temporal compression). The FastVideo side loads it through ``VAELoader``, the
+path a pipeline uses, and is compared on CPU fp32 against
+``diffusers.AutoencoderKLWan.from_pretrained``:
+
+* posterior mode/logvar for a single first frame (PAN2's I2V conditioning),
+* posterior mode/logvar for a 9-frame clip (4x temporal downsampling + cache),
+* decode of 3 latent frames (first-chunk handling + temporal upsampling cache).
+
+The normalization tests use a copy of the checkpoint with non-trivial
+``latents_mean``/``latents_std`` (the tiny checkpoint ships 0/1, which would
+hide a missing or doubled normalization) and pin where it happens: the raw
+``encode``/``decode`` never normalize, ``PAN2ConditioningStage`` applies
+PAN2's ``(z - mean) / std`` once, and ``DecodingStage`` applies
+``z * std + mean`` once.
+
+The default asset is the random-weight tiny PAN2 checkpoint, so this runs on a
+CPU-only host::
+
+ pytest tests/local_tests/vaes/test_pan2_vae_parity.py -v -s
+
+Set ``PAN2_MODEL_PATH`` to a local PAN2 root containing ``vae/`` to use other
+weights.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file
+from torch.testing import assert_close
+
+try:
+ from diffusers import AutoencoderKLWan as DiffusersAutoencoderKLWan
+except ImportError: # pragma: no cover - environment dependent
+ pytest.skip("diffusers with AutoencoderKLWan is required", allow_module_level=True)
+
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.models.wan.vae import AutoencoderKLWan as FastVideoAutoencoderKLWan
+from fastvideo.models.wan.vae_config import WanVAEConfig
+
+PARITY_SCOPE = "both"
+TINY_PAN2_REPO = "wuqing157/tiny-pan2-modular-pipe"
+TINY_PAN2_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
+MODEL_PATH_ENV = "PAN2_MODEL_PATH"
+
+# H and W must be multiples of 32 (16x spatial compression times the 2x2
+# patchify inside the VAE).
+HEIGHT, WIDTH = 64, 96
+VIDEO_FRAMES = 9
+LATENT_FRAMES = (VIDEO_FRAMES - 1) // 4 + 1
+ATOL = RTOL = 1e-5
+
+# Structural fields of Wan-AI/Wan2.2-TI2V-5B-Diffusers vae/config.json. The
+# statistics do not affect parameter shapes, so placeholders keep this inline.
+WAN22_TI2V_VAE_CONFIG = {
+ "attn_scales": [],
+ "base_dim": 160,
+ "clip_output": False,
+ "decoder_base_dim": 256,
+ "dim_mult": [1, 2, 4, 4],
+ "dropout": 0.0,
+ "in_channels": 12,
+ "is_residual": True,
+ "latents_mean": [0.0] * 48,
+ "latents_std": [1.0] * 48,
+ "num_res_blocks": 2,
+ "out_channels": 12,
+ "patch_size": 2,
+ "scale_factor_spatial": 16,
+ "scale_factor_temporal": 4,
+ "temperal_downsample": [False, True, True],
+ "z_dim": 48,
+}
+
+
+def _pan2_root(subfolder: str) -> Path:
+ """Resolve a PAN2 root that contains ``subfolder``; skip when unavailable."""
+ override = os.getenv(MODEL_PATH_ENV)
+ if override:
+ root = Path(override)
+ else:
+ try:
+ from huggingface_hub import snapshot_download
+ except ImportError:
+ pytest.skip("huggingface_hub is required to resolve the tiny PAN2 checkpoint")
+ root = None
+ last_error: Exception | None = None
+ for local_files_only in (True, False):
+ try:
+ root = Path(
+ snapshot_download(
+ TINY_PAN2_REPO,
+ revision=TINY_PAN2_REVISION,
+ allow_patterns=[f"{subfolder}/*"],
+ local_files_only=local_files_only,
+ ))
+ break
+ except Exception as exc: # noqa: BLE001 - offline/missing cache must skip.
+ last_error = exc
+ if root is None:
+ pytest.skip(f"Tiny PAN2 checkpoint unavailable: {last_error}")
+ if not (root / subfolder).is_dir():
+ pytest.skip(f"PAN2 {subfolder}/ missing under {root}")
+ return root
+
+
+@pytest.fixture(scope="module")
+def vae_dir() -> Path:
+ path = _pan2_root("vae") / "vae"
+ for name in ("config.json", "diffusion_pytorch_model.safetensors"):
+ if not (path / name).is_file():
+ pytest.skip(f"PAN2 VAE asset missing: {path / name}")
+ return path
+
+
+@pytest.fixture
+def cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Pin FastVideo's device helper to CPU; it reports ``mps`` on CPU-only Linux."""
+ for module in (
+ "fastvideo.models.loader.component_loader",
+ "fastvideo.pipelines.stages.decoding",
+ "fastvideo.pipelines.basic.pan2.stages",
+ ):
+ monkeypatch.setattr(f"{module}.get_local_torch_device", lambda: torch.device("cpu"))
+
+
+def _pan2_vae_config() -> WanVAEConfig:
+ """The VAE config a PAN2 pipeline config would carry (arch filled by the loader)."""
+ config = WanVAEConfig()
+ config.load_encoder = True
+ config.load_decoder = True
+ return config
+
+
+def _load_fastvideo(vae_dir: Path) -> FastVideoAutoencoderKLWan:
+ """Load through the production VAELoader, exactly as a pipeline does."""
+ fastvideo_args = SimpleNamespace(
+ model_paths={},
+ pipeline_config=SimpleNamespace(vae_config=_pan2_vae_config(), vae_precision="fp32"),
+ vae_cpu_offload=False,
+ )
+ vae = VAELoader().load(str(vae_dir), fastvideo_args)
+ assert isinstance(vae, FastVideoAutoencoderKLWan)
+ assert fastvideo_args.model_paths["vae"] == str(vae_dir)
+ assert all(param.dtype == torch.float32 and param.device.type == "cpu" for param in vae.parameters())
+ return vae
+
+
+def _load_diffusers(vae_dir: Path) -> DiffusersAutoencoderKLWan:
+ model, info = DiffusersAutoencoderKLWan.from_pretrained(
+ str(vae_dir),
+ torch_dtype=torch.float32,
+ local_files_only=True,
+ output_loading_info=True,
+ )
+ problems = {key: value for key, value in info.items() if value}
+ assert not problems, f"Diffusers reference did not load strictly: {problems}"
+ return model.eval()
+
+
+def _pixels(num_frames: int, seed: int) -> torch.Tensor:
+ generator = torch.Generator().manual_seed(seed)
+ return torch.rand((1, 3, num_frames, HEIGHT, WIDTH), generator=generator) * 2.0 - 1.0
+
+
+def _latents(seed: int) -> torch.Tensor:
+ generator = torch.Generator().manual_seed(seed)
+ return torch.randn((1, 48, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16), generator=generator)
+
+
+def _report(name: str, actual: torch.Tensor, expected: torch.Tensor, clamp_bound: float | None = None) -> None:
+ assert actual.shape == expected.shape, f"{name}: {tuple(actual.shape)} != {tuple(expected.shape)}"
+ diff = (actual.float() - expected.float()).abs()
+ message = (f"[pan2-vae] {name}: shape={tuple(actual.shape)} max_abs={diff.max().item():.3e} "
+ f"mean_abs={diff.mean().item():.3e} ref_absmax={expected.abs().max().item():.3e}")
+ if clamp_bound is not None:
+ # Clamped pixels compare trivially; report how many there are.
+ message += f" clamped={(expected.abs() >= clamp_bound).float().mean().item():.2%}"
+ print(message)
+
+
+def _nontrivial_stats() -> tuple[list[float], list[float]]:
+ mean = [round(0.05 * (index - 24), 4) for index in range(48)]
+ std = [round(0.4 + 0.03 * index, 4) for index in range(48)]
+ return mean, std
+
+
+@pytest.fixture
+def stats_vae_dir(vae_dir: Path, tmp_path: Path) -> Path:
+ """Checkpoint copy with non-trivial latent statistics (weights symlinked)."""
+ config = json.loads((vae_dir / "config.json").read_text())
+ config["latents_mean"], config["latents_std"] = _nontrivial_stats()
+ target = tmp_path / "vae"
+ target.mkdir()
+ (target / "config.json").write_text(json.dumps(config))
+ (target / "diffusion_pytorch_model.safetensors").symlink_to(
+ (vae_dir / "diffusion_pytorch_model.safetensors").resolve())
+ return target
+
+
+def _diffusers_stats(model: DiffusersAutoencoderKLWan, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ mean = torch.tensor(model.config.latents_mean).view(1, -1, 1, 1, 1).to(like)
+ std = torch.tensor(model.config.latents_std).view(1, -1, 1, 1, 1).to(like)
+ return mean, std
+
+
+def test_pan2_vae_production_loader_matches_checkpoint(vae_dir: Path, cpu_runtime: None) -> None:
+ """The loader fills the tiny Wan2.2 arch from config.json and covers every weight.
+
+ ``VAELoader`` loads ``AutoencoderKLWan`` with ``strict=False``, so key
+ coverage is checked here instead of trusting the loader.
+ """
+ raw = json.loads((vae_dir / "config.json").read_text())
+ vae = _load_fastvideo(vae_dir)
+ arch = vae.config.arch_config
+ for key in ("base_dim", "decoder_base_dim", "z_dim", "num_res_blocks", "in_channels", "out_channels",
+ "is_residual", "patch_size", "scale_factor_spatial", "scale_factor_temporal", "dropout"):
+ assert getattr(arch, key) == raw[key], key
+ for key in ("dim_mult", "attn_scales", "temperal_downsample", "latents_mean", "latents_std"):
+ assert list(getattr(arch, key)) == list(raw[key]), key
+ assert vae.config.use_feature_cache and not vae.config.use_tiling
+ assert arch.spatial_compression_ratio == 16 and arch.temporal_compression_ratio == 4
+ assert tuple(vae.scaling_factor.shape) == tuple(vae.shift_factor.shape) == (1, 48, 1, 1, 1)
+ assert_close(vae.scaling_factor.flatten(), 1.0 / torch.tensor(raw["latents_std"]))
+ assert_close(vae.shift_factor.flatten(), torch.tensor(raw["latents_mean"]))
+
+ checkpoint = load_file(str(vae_dir / "diffusion_pytorch_model.safetensors"))
+ state = vae.state_dict()
+ assert set(state) == set(checkpoint), (sorted(set(checkpoint) - set(state)), sorted(set(state) - set(checkpoint)))
+ for key, tensor in checkpoint.items():
+ assert torch.equal(state[key], tensor.float()), key
+
+
+@pytest.mark.parametrize("variant", ["wan22_ti2v_5b", "tiny_pan2"])
+def test_pan2_vae_config_shapes_match_diffusers(variant: str, request: pytest.FixtureRequest) -> None:
+ """``WanVAEConfig`` expresses both the real Wan2.2 and the tiny PAN2 VAE configs.
+
+ Both models are built on the meta device, so the full-size config costs no
+ memory; parameter names and shapes must match one-to-one.
+ """
+ if variant == "tiny_pan2":
+ raw = json.loads((request.getfixturevalue("vae_dir") / "config.json").read_text())
+ config = {key: value for key, value in raw.items() if not key.startswith("_")}
+ else:
+ config = WAN22_TI2V_VAE_CONFIG
+ fastvideo_config = _pan2_vae_config()
+ fastvideo_config.update_model_arch(config)
+ with torch.device("meta"):
+ fastvideo = FastVideoAutoencoderKLWan(fastvideo_config)
+ reference = DiffusersAutoencoderKLWan.from_config(config)
+ fastvideo_shapes = {key: tuple(value.shape) for key, value in fastvideo.state_dict().items()}
+ reference_shapes = {key: tuple(value.shape) for key, value in reference.state_dict().items()}
+ assert fastvideo_shapes == reference_shapes
+
+
+def test_pan2_vae_wan22_topology_random_weight_parity() -> None:
+ """Numerically cover the Wan2.2 paths the tiny config skips.
+
+ The tiny checkpoint uses ``dim_mult`` [1, 1, 1, 1] and no separate decoder
+ width, so channel-changing shortcuts and resamplers never run. This narrow
+ random-weight model keeps the real topology (``dim_mult`` [1, 2, 4, 4], two
+ res blocks, ``decoder_base_dim`` > ``base_dim``).
+ """
+ config = dict(WAN22_TI2V_VAE_CONFIG, base_dim=8, decoder_base_dim=16)
+ torch.manual_seed(0)
+ reference = DiffusersAutoencoderKLWan.from_config(config).eval()
+ fastvideo_config = _pan2_vae_config()
+ fastvideo_config.update_model_arch(config)
+ fastvideo = FastVideoAutoencoderKLWan(fastvideo_config).eval()
+ fastvideo.load_state_dict(reference.state_dict(), strict=True)
+
+ video = _pixels(VIDEO_FRAMES, seed=11)
+ latents = _latents(seed=12)
+ with torch.inference_mode():
+ expected_mode = reference.encode(video).latent_dist.mode()
+ actual_mode = fastvideo.encode(video).mode()
+ expected_video = reference.decode(latents).sample
+ actual_video = fastvideo.decode(latents)
+ _report("wan22-topology encode mode", actual_mode, expected_mode)
+ _report("wan22-topology decode", actual_video, expected_video, clamp_bound=1.0)
+ assert_close(actual_mode, expected_mode, atol=ATOL, rtol=RTOL)
+ assert_close(actual_video, expected_video, atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize("num_frames", [1, VIDEO_FRAMES], ids=["first_frame", "video_9f"])
+def test_pan2_vae_encode_matches_diffusers(num_frames: int, vae_dir: Path, cpu_runtime: None) -> None:
+ """Posterior mode (PAN2 uses ``latent_dist.mode()``) and logvar match Diffusers.
+
+ FastVideo's ``encode`` returns the ``DiagonalGaussianDistribution`` itself,
+ so the port calls ``vae.encode(x).mode()`` instead of ``.latent_dist.mode()``.
+ """
+ reference = _load_diffusers(vae_dir)
+ fastvideo = _load_fastvideo(vae_dir)
+ pixels = _pixels(num_frames, seed=100 + num_frames)
+ with torch.inference_mode():
+ expected = reference.encode(pixels).latent_dist
+ actual = fastvideo.encode(pixels)
+ assert tuple(actual.mode().shape) == (1, 48, (num_frames - 1) // 4 + 1, HEIGHT // 16, WIDTH // 16)
+ _report(f"encode mode ({num_frames}f)", actual.mode(), expected.mode())
+ _report(f"encode logvar ({num_frames}f)", actual.logvar, expected.logvar)
+ assert_close(actual.mode(), expected.mode(), atol=ATOL, rtol=RTOL)
+ assert_close(actual.logvar, expected.logvar, atol=ATOL, rtol=RTOL)
+
+
+def test_pan2_vae_decode_matches_diffusers(vae_dir: Path, cpu_runtime: None) -> None:
+ """Decode three latent frames so first-chunk and temporal cache paths run."""
+ reference = _load_diffusers(vae_dir)
+ fastvideo = _load_fastvideo(vae_dir)
+ latents = _latents(seed=7)
+ with torch.inference_mode():
+ expected = reference.decode(latents).sample
+ actual = fastvideo.decode(latents)
+ assert tuple(actual.shape) == (1, 3, VIDEO_FRAMES, HEIGHT, WIDTH)
+ _report("decode", actual, expected, clamp_bound=1.0)
+ assert_close(actual, expected, atol=ATOL, rtol=RTOL)
+
+
+def test_pan2_conditioning_stage_matches_reference(stats_vae_dir: Path, cpu_runtime: None) -> None:
+ """``PAN2ConditioningStage`` builds the reference I2V condition: cover-resize and center crop the first frame,
+ encode it, apply ``(z - mean) / std`` once, and set the mask channel on the first latent frame only.
+
+ The raw ``encode`` is unnormalized, so the stage is the only place the statistics apply.
+ """
+ import PIL.Image
+
+ from fastvideo.pipelines.basic.pan2.stages import PAN2ConditioningStage
+ from tests.local_tests.pan2.reference import reference_condition_latents
+
+ reference = _load_diffusers(stats_vae_dir)
+ fastvideo = _load_fastvideo(stats_vae_dir)
+ # Not the output aspect ratio, so the resize and crop are exercised.
+ image = PIL.Image.radial_gradient("L").resize((120, 64)).convert("RGB")
+ latents = torch.zeros(1, 48, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16)
+ with torch.inference_mode():
+ pixels = PAN2ConditioningStage.preprocess_image(image, HEIGHT, WIDTH)
+ raw = reference.encode(pixels).latent_dist.mode()
+ assert_close(fastvideo.encode(pixels).mode(), raw, atol=ATOL, rtol=RTOL)
+ expected = reference_condition_latents(reference, latents, image, HEIGHT, WIDTH)
+
+ batch = SimpleNamespace(latents=latents, pil_image=image, height=HEIGHT, width=WIDTH, image_latent=None)
+ stage_args = SimpleNamespace(pipeline_config=SimpleNamespace(vae_precision="fp32"), vae_cpu_offload=False)
+ actual = PAN2ConditioningStage(fastvideo).forward(batch, stage_args).image_latent
+ _report("PAN2 I2V condition", actual, expected)
+ assert actual.shape == (1, 49, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16)
+ assert_close(actual, expected, atol=ATOL, rtol=RTOL)
+ assert torch.equal(actual[:, 48, 0], torch.ones_like(actual[:, 48, 0]))
+ assert not actual[:, :, 1:].any()
+ assert not torch.allclose(raw, expected[:, :48, :1], atol=1e-2), "statistics too close to identity"
+
+
+def test_pan2_decoding_stage_denormalizes_once(stats_vae_dir: Path, cpu_runtime: None) -> None:
+ """``DecodingStage`` performs PAN2's ``z * std + mean`` itself.
+
+ The PAN2 denoising loop must hand normalized latents to ``DecodingStage``.
+ Denormalizing beforehand (as the upstream pipeline does inline) would apply
+ the statistics twice, which the negative check below detects.
+ """
+ from fastvideo.pipelines.stages.decoding import DecodingStage
+
+ reference = _load_diffusers(stats_vae_dir)
+ fastvideo = _load_fastvideo(stats_vae_dir)
+ latents = _latents(seed=31)
+ # vae_tiling=True is the PipelineConfig default; the feature-cache decode ignores tiling, so it must not matter.
+ decode_args = SimpleNamespace(
+ pipeline_config=SimpleNamespace(vae_precision="fp32", vae_decode_precision="fp32", vae_tiling=True),
+ disable_autocast=False,
+ )
+ stage = DecodingStage(fastvideo)
+ with torch.inference_mode():
+ mean, std = _diffusers_stats(reference, latents)
+ expected = reference.decode(latents * std + mean).sample
+ expected = (expected * 0.5 + 0.5).clamp(0, 1)
+ actual = stage.decode(latents, decode_args)
+ doubled = stage.decode(latents * std + mean, decode_args)
+ _report("DecodingStage output", actual, expected)
+ assert_close(actual, expected, atol=ATOL, rtol=RTOL)
+ assert (doubled - expected).abs().max().item() > 1e-3, "double denormalization went undetected"