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

- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output. - WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32 accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it); decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the video, with the ranks agreeing on whether every one of them could build the kernel. - PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions, set_fused_kernels). - PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels). - VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame. - The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-28 00:47:23 +00:00
46 changed files with 6834 additions and 23 deletions
@@ -259,6 +259,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_H3_VSA_FP4_K_SMOOTH`                            | bool                                                               | `0`                                                   | attention   | Subtract the per-head key mean before FP4 key quantization, matching the emulator and QAD training numerics. Exact for the softmax; it only changes quantization error.                                                                                                                                                                                                                                                                     |
 | `FASTVIDEO_H3_GATE_CHUNK_ROWS`                             | int                                                                | `0`                                                   | performance | Rows per chunk when the VSA FP4 gate branch expands the pooled compression output back to token rows. 0 keeps the single full-size index_select; a positive value caps the transient at chunk_rows\*heads\*head_dim and is bit-identical.                                                                                                                                                                                                   |
 | `FASTVIDEO_H3_VSA_FINE_ATTENTION`                          | one of , emulator-mix8, emulator-mixmx, emulator-qk, emulator-bf16 | `""`                                                  | attention   | Replace tile-128 fine attention, preserving VSA selection and compression, with the Triton emulator of the sparse FP4 kernel: emulator-mix8 = NVFP4 Q/K video, per-token FP8 Q/K prefix; emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over FASTVIDEO_H3_VSA_FP4. Empty keeps the native route.                 |
+| `FASTVIDEO_WAN22_VAE_DECODE_DIR`                           | path                                                               | `""`                                                  | performance | Directory of the Wan2.2 VAE decoder kernels (fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds on first use, when its config sets decode_kernel. Empty takes the copy in this checkout.                                                                                                                                                                                                                            |
 | `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE`                        | bool                                                               | `0`                                                   | performance | VSA-H3 64/128-token tiles: scatter rows straight into the heads-first layout the sparse kernels read, removing their per-call transpose copy of query, key and value. Bit-identical; no-grad, uncompiled CUDA only.                                                                                                                                                                                                                         |
 | `FASTVIDEO_H3_VSA_TILE_FIRST`                              | bool                                                               | `0`                                                   | attention   | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters.                                                                                                                                                                                                                                                                                                                                 |
 | `FASTVIDEO_H3_VSA_SM89_KERNEL`                             | one of original, bf16, int8                                        | `original`                                            | attention   | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV).                                                                                                                                                                                                                                                                                                                                              |
@@ -328,6 +329,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER`                     | int                                                                | `0`                                                   | test        | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`.                                                                                                                                                                                                                                                                                                                                   |
 | `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT`                     | str                                                                | unset                                                 | test        | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`.                                                                                                                                                                                                                                                                                                                                            |
 | `FASTVIDEO_TEST_SD35_MODEL_DIR`                            | str                                                                | `stabilityai/stable-diffusion-3.5-medium`             | test        | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`.                                                                                                                                                                                                                                                                                                                                                           |
+| `FASTVIDEO_TEST_PAN2_MODEL_DIR`                            | path                                                               | unset                                                 | test        | PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU tests; unset skips them.                                                                                                                                                                                                                                                                                                                                              |
 | `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR`                       | str                                                                | unset                                                 | test        | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`.                                                                                                                                                                                                                                                                                                                                             |
 | `FASTVIDEO_TEST_WAN_ANIMATE_MODEL_DIR`                     | str                                                                | unset                                                 | test        | Local Wan2.2-Animate-14B checkpoint for the Wan-Animate weight tests. Deprecated names: `WAN_ANIMATE_MODEL_PATH`.                                                                                                                                                                                                                                                                                                                           |

@@ -19,7 +19,7 @@ from the `pan2_t2v` and `pan2_i2v` presets:
 | Resolution | 704x1248 (height and width must be multiples of 32) |
 | Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) |
 | Steps | 50, uniform sigmas from 1 to 0 with the checkpoint's flow shift (9 for `IFM/PAN2-TEST`, the reference sampler's value) |
-| Guidance | 3.0, classifier-free guidance. Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) |
+| Guidance | 2.5, classifier-free guidance, the reference launcher's default (3 before 2026-10-23). Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) |
 | Negative prompt | None unless the request carries one. The reference sampler always sends the release's structured negative caption (`negative_prompt.json` of the PAN2 repository) with its three `time_range_sec` values set to the clip length, `0 to 20` for 481 frames at 24 fps |

 PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
@@ -67,6 +67,20 @@ The trimmed sequence can use either FlashAttention or Torch SDPA.
 PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py
 ```

+## Performance switches
+
+PAN2 sampling turns these on by default (`PAN2PipelineConfig`); pass any of them as `False` to
+`VideoGenerator.from_pretrained` to turn it off:
+
+| Switch | What it does |
+|---|---|
+| `vae_decode_kernel` | Decodes with fastvideo-kernel's Wan2.2 decoder on B200 / GB200 / B300 (704x1248, batch 1): bf16 with fp32 accumulation, 7x faster than the fp32 decode on a B300 at PSNR 61-63 dB and SSIM 0.998 against it, at a peak of about 20 GiB for a 5 s video. It builds on first use (minutes) into `WAN22_VAE_DECODE_BUILD_DIR`, against fastvideo-kernel's CUTLASS submodule (`git submodule update --init fastvideo-kernel/include/cutlass`, or set `CUTLASS_DIR`). Other GPUs and sizes keep the fp32 decode. |
+| `vae_decode_kernel_parallel` | With 2, 4 or 8 GPUs, splits that decode over them instead of decoding the whole video on each: a 20 s video takes 3.0 s on one B300, 1.55 s on two, 0.79 s on four and 0.47 s on eight. |
+| `dit_fused_kernels` | Runs the DiT's fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation. Measured with a distilled PAN2 model on eight B300s, they cut the denoising time by 21%. The rounding differs from the unfused path; record by record, fused and unfused videos were equally close to the 50-step reference. |
+
+On GPUs with room for every component (B200, B300), keep the text encoder resident as well:
+`text_encoder_cpu_offload=False`.
+
 ## Memory and offloading

 The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards
@@ -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

@@ -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)
@@ -0,0 +1,98 @@
+# Wan2.2 VAE decoder (sm_100a, sm_103a)
+
+The decoder of the Wan2.2 VAE (the VAE of Wan2.2-TI2V-5B: 48 latent channels, 16 x 16 spatial and 4x temporal
+compression) at 704 x 1248, on one, two or four B200 / GB200 GPUs of a node. It decodes latents [1, 48, T, 44, 78]
+into the video [1, 3, 4T - 3, 704, 1248] for any T >= 1, as diffusers' `AutoencoderKLWan.decode(z).sample` does;
+`reference.py` defines the result exactly, a plain-PyTorch decoder bit-exact with diffusers in fp32 and in bf16.
+
+## Use
+
+```python
+import sys
+sys.path.insert(0, "fastvideo-kernel/wan22_vae_decode")
+from wan22_vae_decode import WanVAEDecoder
+
+decoder = WanVAEDecoder(vae.state_dict(), device="cuda")   # diffusers' AutoencoderKLWan, or its state dict
+video = decoder.decode(latents)   # [1, 48, T, 44, 78], as the VAE's decode takes them -> [1, 3, 4T-3, 704, 1248] bf16
+```
+
+- **Over several GPUs**, every rank builds a decoder with the same process group (`WanVAEDecoder(sd, device, pg)`,
+  default: torch.distributed's default group) and calls `decode` with the full latents; rank 0 gets the video and the
+  other ranks `None`. On one GPU without torch.distributed, the decoder initializes a single-rank NCCL group.
+- The first call builds the extension from `solution/` (about a minute) into `$WAN22_VAE_DECODE_BUILD_DIR` (default
+  `~/.cache/fastvideo/wan22_vae_decode`, one directory per `LOCAL_RANK`), and later processes reuse the build.
+- `WanVAEDecoder(..., frames_per_pass=4)` sets the latent frames each GPU carries through one pass (see below).
+- `WanVAEDecoder` prepares the weights once (layout, phase and Winograd transforms of the weights). The raw entry
+  points are `prepare(weights, rank, world, pg)`, `run(z, weights, rank, world, pg, out)` (4 frames per pass) and
+  `run_with(..., out, frames_per_pass)` in `solution/binding.cpp`, destination-passing, on the current stream;
+  `weights` are the 110 tensors of `weight_names()` in bf16.
+
+## Performance and memory
+
+At PAN2's default 704 x 1248, 24 fps, on GB200 GPUs, against PAN2's decode today (diffusers' `AutoencoderKLWan` in bf16
+under autocast, one GPU, one latent frame at a time). Memory is the peak allocated and, after it, the peak reserved by
+PyTorch's caching allocator (what nvidia-smi shows), per GPU; the default is 4 latent frames per GPU per pass:
+
+| Decoder | 8 s (193 frames) | 12 s (289 frames) | 20 s (481 frames) |
+|---|---|---|---|
+| Release, 1 GPU | 5.88 s, 12.9 / 17.1 GiB | 8.84 s, 13.9 / 19.4 GiB | 14.76 s, 15.8 / 91.5 GiB |
+| **This decoder, 1 GPU** | **1.17 s** (x5.0), 20.7 / 36.3 GiB | **1.74 s** (x5.1), 21.2 / 36.8 GiB | **2.89 s** (x5.1), 22.2 / 37.8 GiB |
+| **This decoder, 2 GPUs** | **619 ms**, 18.7 / 33.0 GiB | **929 ms**, 19.2 / 33.5 GiB | **1.54 s**, 20.2 / 34.5 GiB |
+| **This decoder, 4 GPUs** | **328 ms**, 17.8 / 41.4 GiB | **475 ms**, 18.3 / 34.1 GiB | **787 ms** (x18.8), 19.2 / 42.9 GiB |
+
+`frames_per_pass` trades memory for speed and leaves the video unchanged. On one GPU:
+
+| Latent frames per pass | 8 s | 12 s | 20 s |
+|---|---|---|---|
+| 2 | 1.25 s, 14.4 / 24.4 GiB | 1.88 s, 14.9 / 24.9 GiB | 3.13 s, 15.9 / 25.8 GiB |
+| 4 (default) | 1.17 s, 20.7 / 36.3 GiB | 1.74 s, 21.2 / 36.8 GiB | 2.89 s, 22.2 / 37.8 GiB |
+| 8 | 1.10 s, 33.7 / 62.1 GiB | 1.64 s, 34.2 / 62.6 GiB | 2.72 s, 35.2 / 63.6 GiB |
+| 16 | 1.06 s, 59.8 / 121.0 GiB | 1.58 s, 60.2 / 121.7 GiB | 2.67 s, 61.2 / 122.5 GiB |
+
+- With `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, the default reserves 30.6-32.0 GiB on one GPU and 2 frames
+  per pass 19.6-21.1 GiB, at the same speed.
+- On the optimization run's nine workloads (`tools/bench.py`: 121, 241 and 481 frames at 1, 2 and 4 GPUs, SM clocks
+  1.6-1.95 GHz), the default's geomean is 758 ms, against 1043 ms for the run's first solution (x1.38) and 698 ms at
+  16 frames per pass.
+- The CUTLASS Winograd component products take 60 % of the GPU time, the transforms and epilogues 30 %
+  (`notes/optimization.md`).
+
+## Correctness
+
+Against the fp32 decode of real latents, with `d = video - reference` per clip: every element finite and in [-1, 1],
+max |d| at most three times what the release's bf16 decode errs by, and the RMS of d likewise. On this tree, every
+clip at worlds 1, 2 and 4 and the short clips `f1`, `f5`, `f17` at worlds 1 and 4 pass, using 0.30-0.66 of the max
+allowance and 0.35-0.42 of the RMS one (PSNR 56-60 dB against fp32; the release's bf16 decode has 58-61 dB).
+
+```bash
+cd fastvideo-kernel/wan22_vae_decode
+DATA=/path/to/data python tools/bench.py --solution solution            # every workload this job has the GPUs for
+DATA=/path/to/data python tools/bench.py --workload f5x1 f17x1 --no-time  # correctness of the short clips
+cd .. && WAN22_VAE_WEIGHTS=/path/to/vae python -m pytest tests/test_wan22_vae_decode.py   # the wrapper, one GPU
+```
+
+`tools/bench.py` reads `$DATA`: the latents (`latents_<source>.pt`), the weights in `weight_names()` order
+(`decoder_weights_bf16.pt`, `weight_names.json`), the fp32 decodes (`ref_<clip>.pt`) and the release decode's errors
+(`calibration.json`); `tools/spec.py` documents each. The tests need only the VAE's weights.
+
+## Limits
+
+- sm_100a and sm_103a (B200 / GB200, B300): `tcgen05` with tensor memory, TMA and 2-CTA clusters; the loader builds
+  for the GPU it runs on. CUDA 13, torch 2.12 tested; the tests pass on B300 too, built for sm_103a.
+- 704 x 1248 video (latents 44 x 78) and batch 1; any T >= 1.
+- 1, 2 or 4 GPUs of one node were checked; the code accepts up to 8 ranks.
+- bf16 weights and activations, fp32 accumulation and fp32 RMS-norm statistics.
+- In FastVideo, `WanVAEConfig.decode_kernel` routes `AutoencoderKLWan.decode` here where it applies (on one GPU per
+  rank, `WanVAEDecoder(..., local=True)`), and PAN2 sampling turns it on (`PAN2PipelineConfig.vae_decode_kernel`). On
+  a B300 with PAN2 latents it decodes 7x faster than the fp32 decode, at PSNR 61-63 dB and SSIM 0.998 against it.
+
+## Layout
+
+- `solution/`: the kernels and their binding; `solution.json` holds the build flags. The CUTLASS and CuTe headers
+  come from fastvideo-kernel's CUTLASS submodule (`include/cutlass`, or `$CUTLASS_DIR`); the decoder's structure follows
+  diffusers' (`WAN_LICENSE.txt`).
+- `wan22_vae_decode.py`: `WanVAEDecoder`, `decoder_weights` and `weight_names`.
+- `reference.py`: the definition of the result.
+- `tools/`: `bench.py` (check, audit and time), `loader.py` (the build), `spec.py` (workloads, weights, data),
+  `audit.py`, `bounds.py`.
+- `notes/optimization.md`: the design, what each merge bought, what did not pay, and where to go next.
@@ -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).
@@ -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)
@@ -0,0 +1,217 @@
+// Mid-block self-attention with our own kernels, replacing ATen's
+// scaled_dot_product_attention (fmha_cutlassF, 8.5 ms of f121x1's 954 ms).
+//
+// The reference computes, per latent frame, one head of softmax(q k^T/32) v over
+// the H*W=3432 positions with C=1024 channels and no mask. ATen's path cannot use
+// flash attention here (head dim 1024 > 256) and falls back to the xformers
+// cutlass kernel, which reads and writes far more per frame than the products
+// need: the decode's whole attention is 1.5 TFLOP (0.63 ms at 16-bit peak for
+// f121x1) against 8.5 ms measured, ~13x off.
+//
+// Three stages, each at its own bound:
+//   1. pack: q,k into contiguous [HW,C] per frame and v transposed into the
+//      [C,PAD] weight layout the convolution GEMM wants (~21 MB/frame).
+//   2. S = q k^T per frame (fp32 out) and O = p v per frame (bf16 out), both as
+//      1x1x1 implicit GEMMs through the vendored CUTLASS SM100 convolution
+//      collective (2-CTA tcgen05 MMA, TMA, fp32 accumulation) -- the same
+//      machinery conv.cu uses, so no new GEMM path is introduced.
+//   3. row softmax: S row -> exp(S/sqrt(C) - max)/sum in bf16, one CTA per row,
+//      the row held in shared memory (one read of S, one write of P).
+// P's rows are padded to a multiple of 64 columns and v's transpose is zero
+// there, so the GEMM's tail contributes nothing.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include "attn.h"
+
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/convnd_problem_shape.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/conv/dispatch_policy.hpp"
+#include "cutlass/conv/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/conv/device/conv_universal_adapter.hpp"
+#include "cutlass/conv/kernel/conv_universal.hpp"
+
+using namespace cute;
+using ElementBF16 = cutlass::bfloat16_t;
+using LayoutNDHWC = cutlass::layout::TensorNDHWC;
+using ArchTag = cutlass::arch::Sm100;
+using OperatorClass = cutlass::arch::OpClassTensorOp;
+constexpr auto Fprop = cutlass::conv::Operator::kFprop;
+
+// One 1x1x1 implicit GEMM, y = x (per position) times w, fp32 accumulation and
+// ElementD out. x: channels-last3d [1,C,D,1,1] (may be a strided view); w:
+// [K,C,1,1,1] with explicit strides; y: packed channels-last3d [1,K,D,1,1].
+template <class ElementD, int M, int N, int CL>
+static void gemm1x1(const at::Tensor& x, const at::Tensor& w, const at::Tensor& y) {
+  using Tile = Shape<Int<M>, Int<N>, Shape<_64>>;
+  using Cluster = Shape<Int<CL>, _1, _1>;
+  using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto,
+      float, float, ElementD, LayoutNDHWC, 8, ElementD, LayoutNDHWC, 8,
+      cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+  using Mainloop = typename cutlass::conv::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Fprop, ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, float,
+      Tile, Cluster,
+      cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
+      cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+  using Problem = cutlass::conv::ConvProblemShape<Fprop, 3>;
+  using Kernel = cutlass::conv::kernel::ConvUniversal<Problem, Mainloop, Epilogue>;
+  using Conv = cutlass::conv::device::ConvUniversalAdapter<Kernel>;
+  int C = x.size(1), D = x.size(2);
+  int K = w.size(0);
+  TORCH_CHECK(y.size(1) == K && y.size(2) == D, "attn gemm: output shape mismatch");
+  Problem problem(cutlass::conv::Mode::kCrossCorrelation,
+                  {1, D, 1, 1, C}, {x.stride(0), x.stride(2), x.stride(3), x.stride(4), x.stride(1)},
+                  {K, 1, 1, 1, C}, {w.stride(0), w.stride(2), w.stride(3), w.stride(4), w.stride(1)},
+                  {0, 0, 0}, {0, 0, 0}, {1, 1, 1}, {1, 1, 1}, 1);
+  typename Kernel::StrideC sc{};
+  typename Kernel::StrideD sd{};
+  cute::for_each(cute::make_seq<cute::rank<0>(sc)>{}, [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
+  cute::for_each(cute::make_seq<cute::rank<0>(sd)>{}, [&](auto i) { cute::get<0, i>(sd) = problem.stride_C[Problem::RankT - 2 - i]; });
+  typename Conv::Arguments args{problem, {(const ElementBF16*)x.data_ptr(), (const ElementBF16*)w.data_ptr()},
+                                {{1.f, 0.f}, (const ElementD*)y.data_ptr(), sc, (ElementD*)y.data_ptr(), sd}};
+  Conv conv;
+  TORCH_CHECK(conv.can_implement(args) == cutlass::Status::kSuccess,
+              "attn gemm: CUTLASS cannot implement D=", D, " N=", K, " C=", C);
+  size_t size = Conv::get_workspace_size(args);
+  auto workspace = at::empty({(int64_t)size}, x.options().dtype(at::kByte));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  TORCH_CHECK(conv.initialize(args, workspace.data_ptr(), stream) == cutlass::Status::kSuccess,
+              "attn gemm: initialization failed");
+  TORCH_CHECK(conv.run(stream) == cutlass::Status::kSuccess, "attn gemm: run failed");
+}
+
+// Channels-last-3d views with explicit strides: a size-1 dim built by unsqueeze
+// carries stride 1, which the im2col TMA descriptor rejects.
+static at::Tensor nhwc_in(const at::Tensor& m, int64_t C, int64_t D) {  // [D,C] -> [1,C,D,1,1]
+  return m.as_strided({1, C, D, 1, 1}, {C * D, 1, C, C, C});
+}
+static at::Tensor nhwc_w(const at::Tensor& m, int64_t K, int64_t C) {  // [K,C] -> [K,C,1,1,1]
+  return m.as_strided({K, C, 1, 1, 1}, {C, 1, C, C, C});
+}
+static at::Tensor nhwc_out(const at::Tensor& m, int64_t K, int64_t D) {  // [D,K] -> [1,K,D,1,1]
+  return m.as_strided({1, K, D, 1, 1}, {K * D, 1, K, K, K});
+}
+
+// v [T,1,HW,3C] at offset 2C -> vt [T,C,PAD] (transposed, [T,C,PAD] zeroed first).
+static __global__ void transpose_v_kernel(const __nv_bfloat16* __restrict__ qkv, int T, int HW, int C, int PAD,
+                                   __nv_bfloat16* __restrict__ vt) {
+  __shared__ __nv_bfloat16 tile[32][33];
+  int i0 = blockIdx.x * 32, c0 = blockIdx.y * 32, t = blockIdx.z;
+  {
+    int i = i0 + threadIdx.y, c = c0 + threadIdx.x;
+    if (i < HW && c < C) tile[threadIdx.y][threadIdx.x] = qkv[((size_t)t * HW + i) * (3 * C) + 2 * C + c];
+  }
+  __syncthreads();
+  {
+    int i = i0 + threadIdx.x, c = c0 + threadIdx.y;
+    if (i < HW && c < C) vt[((size_t)t * C + c) * PAD + i] = tile[threadIdx.x][threadIdx.y];
+  }
+}
+
+// One CTA per row: p = exp(s*scale - max)/sum in bf16, then zero the padded tail.
+static __global__ void row_softmax_kernel(const float* __restrict__ s, int cols, int PAD, float scale,
+                                   __nv_bfloat16* __restrict__ p) {
+  extern __shared__ float row[];
+  const int n4 = cols / 4;
+  const float4* src = (const float4*)(s + (size_t)blockIdx.x * cols);
+  __nv_bfloat16* dst = p + (size_t)blockIdx.x * PAD;
+  float4* r4 = (float4*)row;
+  for (int j = threadIdx.x; j < n4; j += blockDim.x) {
+    float4 v = src[j];
+    v.x *= scale;
+    v.y *= scale;
+    v.z *= scale;
+    v.w *= scale;
+    r4[j] = v;
+  }
+  __syncthreads();
+  float m = -INFINITY;
+  for (int j = threadIdx.x; j < n4; j += blockDim.x) {
+    float4 v = r4[j];
+    m = fmaxf(m, fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w)));
+  }
+  for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off));
+  __shared__ float red[32];
+  if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = m;
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    m = threadIdx.x < blockDim.x / 32 ? red[threadIdx.x] : -INFINITY;
+    for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off));
+    if (threadIdx.x == 0) red[0] = m;
+  }
+  __syncthreads();
+  m = red[0];
+  __syncthreads();
+  float sum = 0.f;
+  for (int j = threadIdx.x; j < n4; j += blockDim.x) {
+    float4 v = r4[j];
+    v.x = __expf(v.x - m);
+    v.y = __expf(v.y - m);
+    v.z = __expf(v.z - m);
+    v.w = __expf(v.w - m);
+    r4[j] = v;
+    sum += (v.x + v.y) + (v.z + v.w);
+  }
+  for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, off);
+  if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = sum;
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    sum = threadIdx.x < blockDim.x / 32 ? red[threadIdx.x] : 0.f;
+    for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor_sync(0xffffffff, sum, off);
+    if (threadIdx.x == 0) red[0] = sum;
+  }
+  __syncthreads();
+  const float inv = 1.f / red[0];
+  const int npair = cols / 2;
+  for (int j = threadIdx.x; j < npair; j += blockDim.x)
+    *(__nv_bfloat162*)(dst + 2 * j) = __floats2bfloat162_rn(row[2 * j] * inv, row[2 * j + 1] * inv);
+  for (int j = cols + threadIdx.x; j < PAD; j += blockDim.x) dst[j] = __float2bfloat16(0.f);
+}
+
+at::Tensor attn_forward(const at::Tensor& qkv) {
+  TORCH_CHECK(qkv.dim() == 4 && qkv.size(1) == 1 && qkv.size(3) % 3 == 0, "attn: expected [T,1,HW,3C]");
+  int T = qkv.size(0), HW = qkv.size(2), C = qkv.size(3) / 3;
+  TORCH_CHECK(C % 64 == 0, "attn: channels must be a multiple of 64");
+  int PAD = (HW + 63) / 64 * 64;
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto bf = qkv.options();
+  auto vt = at::zeros({T, C, PAD}, bf);
+  auto s = at::empty({T, HW, HW}, bf.dtype(at::kFloat));
+  auto p = at::empty({T, HW, PAD}, bf);
+  auto o = at::empty({T, 1, HW, C}, bf);
+  const __nv_bfloat16* qkv_p = (const __nv_bfloat16*)qkv.data_ptr();
+
+  transpose_v_kernel<<<dim3((HW + 31) / 32, (C + 31) / 32, T), dim3(32, 32), 0, stream>>>(
+      qkv_p, T, HW, C, PAD, (__nv_bfloat16*)vt.data_ptr());
+
+  // S = q k^T : M = positions, N = keys, K = channels. q and k are read in place:
+  // their channel stride is 1 and their position stride 3C, which TMA accepts.
+  for (int t = 0; t < T; ++t) {
+    auto qv = qkv.select(0, t).select(0, 0);  // [HW, 3C]
+    auto x = qv.narrow(1, 0, C).as_strided({1, C, HW, 1, 1}, {(int64_t)3 * C * HW, 1, 3 * C, 3 * C, 3 * C});
+    auto w = qv.narrow(1, C, C).as_strided({HW, C, 1, 1, 1}, {3 * C, 1, 3 * C, 3 * C, 3 * C});
+    auto y = nhwc_out(s.select(0, t), HW, HW);
+    gemm1x1<float, 256, 256, 2>(x, w, y);
+  }
+
+  float* s_p = (float*)s.data_ptr();
+  float scale = 1.f / sqrtf((float)C);
+  row_softmax_kernel<<<T * HW, 256, HW * sizeof(float), stream>>>(
+      s_p, HW, PAD, scale, (__nv_bfloat16*)p.data_ptr());
+
+  // O = p v : M = positions, N = channels, K = keys (padded).
+  for (int t = 0; t < T; ++t) {
+    auto x = nhwc_in(p.select(0, t), PAD, HW);
+    auto w = nhwc_w(vt.select(0, t), C, PAD);
+    auto y = nhwc_out(o.select(0, t).select(0, 0), C, HW);
+    gemm1x1<ElementBF16, 256, 256, 2>(x, w, y);
+  }
+  return o;
+}
@@ -0,0 +1,6 @@
+#pragma once
+#include <ATen/ATen.h>
+// One head of self-attention over the H*W positions of every latent frame of
+// one decoder chunk: qkv is [T,1,HW,3C] bf16 (q, k, v concatenated on the last
+// dim, as binding.cpp's attention() builds it); the result is [T,1,HW,C] bf16.
+at::Tensor attn_forward(const at::Tensor& qkv);
@@ -0,0 +1,459 @@
+// Wan2.2 VAE decoder seed. Derived from the reference architecture and
+// indexing in diffusers' AutoencoderKLWan (Wan Team / Hugging Face Team,
+// Apache-2.0; see WAN_LICENSE.txt). Native CUDA execution is implemented here.
+// All invocation state, including temporal history, belongs to Decoder below.
+#include <torch/extension.h>
+#include <ATen/ops/cudnn_convolution.h>
+#include <ATen/ops/scaled_dot_product_attention.h>
+#include <ATen/record_function.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include "kernels.h"
+#include "prefix.h"
+#include "comm.h"
+#include "resample.h"
+#include "attn.h"
+#include "head.h"
+#include <array>
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+
+at::Tensor custom_conv(const at::Tensor&,const at::Tensor&,int64_t,int64_t,int64_t);
+
+// The temporal upsampler, writing the frame-interleaved video directly (solution/timeconv.cu).
+at::Tensor time_upsample(const at::Tensor&,const at::Tensor&,const at::Tensor&,const at::Tensor&,int64_t,int64_t,bool);
+
+using PG = c10::intrusive_ptr<c10d::ProcessGroup>;
+using at::Tensor;
+using Weights = std::vector<Tensor>;
+
+Weights prepare(const Weights& weights,int64_t rank,int64_t world,const PG& pg) {
+  TORCH_CHECK(weights.size()==110,"Expected the 110 decoder tensors");
+  at::NoGradGuard guard;
+  Weights result;
+  result.reserve(weights.size());
+  for(const auto& w:weights) {
+    // A 1x1x1 tensor is "trivially channels-last", so contiguous() leaves it with unit strides on
+    // its size-1 spatial dims; CUTLASS's im2col TMA needs those strides to be 16-byte multiples.
+    // The same bytes with dense strides: as_strided keeps the NDHWC order but gives every spatial
+    // dim the channel-count stride.
+    if(w.dim()==5 && w.size(2)==1 && w.size(3)==1 && w.size(4)==1) {
+      auto wc=w.contiguous();int64_t C=wc.size(1);
+      result.push_back(wc.as_strided({wc.size(0),C,1,1,1},{C,1,C,C,C}));
+    }
+    else if(w.dim()==5) result.push_back(w.contiguous(at::MemoryFormat::ChannelsLast3d));
+    else if(w.dim()==4 && w.size(2)>0 && w.size(1)>1)
+      result.push_back(w.unsqueeze(2).contiguous(at::MemoryFormat::ChannelsLast3d));
+    else result.push_back(w.contiguous());
+  }
+  // The output head: conv_out's nine spatial taps become the output channels of a 3x1x1
+  // (temporal-only) convolution, so its 112 outputs are (tap i,j) x (channel c) and head.cu's
+  // epilogue sums the nine shifted copies. W'[(i*3+j)*12+c, ch, t] = W[c, ch, t, i, j]; the four
+  // channels past 108 are zero and are never read. N=112 rather than 108 keeps the epilogue's
+  // channel count a multiple of 16.
+  {
+    auto orig=result[108];
+    auto t=orig.permute({3,4,0,1,2}).contiguous();          // [i,j,c,ch,t] = [3,3,12,256,3]
+    auto filled=at::zeros({112,256,3},orig.options());      // [n,ch,t], n=(i*3+j)*12+c
+    filled.narrow(0,0,108).copy_(t.reshape({108,256,3}));
+    auto kt=filled.transpose(1,2).contiguous();             // [n,t,ch] contiguous
+    // The layout a [K,C,3,3,3] weight has under ChannelsLast3d: C innermost with stride 1, every
+    // other stride a multiple of 16 bytes. torch normalises a size-1 H/W away, hence as_strided.
+    result[108]=kt.transpose(1,2).as_strided({112,256,3,1,1},{768,1,256,256,256});
+  }
+  // Nearest-x2 + 3x3 convolution equals four 2x2 low-resolution phases.
+  // Form the sums in fp32 and store the transformed weights in bf16.
+  for(int spatial : {39,61,85}) {
+    auto orig=result[spatial].to(at::kFloat);
+    int C=orig.size(0);
+    for(int a=0;a<2;++a)for(int b=0;b<2;++b) {
+      auto transformed=at::zeros({C,C,1,2,2},orig.options());
+      for(int u=0;u<2;++u)for(int v=0;v<2;++v) {
+        auto sum=at::zeros({C,C},orig.options());
+        for(int r=0;r<3;++r)for(int c=0;c<3;++c) {
+          int dy=a+r-1,dx=b+c-1;
+          int iy=dy<0?-1:dy/2,ix=dx<0?-1:dx/2;
+          if(iy==u+a-1&&ix==v+b-1)sum.add_(orig.select(2,0).select(2,r).select(2,c));
+        }
+        transformed.select(2,0).select(2,u).select(2,v).copy_(sum);
+      }
+      result.push_back(transformed.to(at::kBFloat16).contiguous(at::MemoryFormat::ChannelsLast3d));
+    }
+  }
+  // F(4x4,3x3) Winograd weights for up blocks 1, 2, and 3.
+  // The 36 components each retain the causal 3-tap temporal convolution.
+  const float G[6][3]={{.25f,0,0},{-1.f/6,-1.f/6,-1.f/6},{-1.f/6,1.f/6,-1.f/6},{1.f/24,1.f/12,1.f/6},{1.f/24,-1.f/12,1.f/6},{0,0,1}};
+  for(int idx : {44,47,50,53,56,59,66,69,74,77,80,83,88,91,96,99,102,105}) {
+    auto orig=result[idx].to(at::kFloat);
+    for(int a=0;a<6;a++)for(int b=0;b<6;b++) {
+      auto sum=at::zeros({orig.size(0),orig.size(1),3},orig.options());
+      for(int r=0;r<3;r++)for(int c=0;c<3;c++)if(G[a][r]*G[b][c]!=0)
+        sum.add_(orig.select(3,r).select(3,c),G[a][r]*G[b][c]);
+      result.push_back(sum.unsqueeze(3).unsqueeze(4).to(at::kBFloat16).contiguous(at::MemoryFormat::ChannelsLast3d));
+    }
+  }
+  // The 1x1x1 shortcut of a channel-changing resnet is added to the second convolution's
+  // bias + output, out = round(acc2 + b2) + round(acc1 + b1). Fold b1 into b2 instead (one more
+  // bf16 rounding of the bias sum and none of the operand): the shortcut tensor then needs no
+  // separate bias pass, which was a full read+write of it.
+  for(auto pair : {std::pair<int,int>{72,70},std::pair<int,int>{94,92}}) {
+    auto sum=result[pair.second].to(at::kFloat)+result[pair.first].to(at::kFloat);
+    result[pair.second]=sum.to(at::kBFloat16).contiguous();
+  }
+  return result;
+}
+
+struct Decoder {
+  // This must not implicitly bind to a fused-epilogue weight index when
+  // independently extended conv() interfaces are merged.
+  struct TemporalTrim {
+    int frames;
+    explicit constexpr TemporalTrim(int n) : frames(n) {}
+  };
+
+  const Weights& w;
+  int rank,world;
+  PG pg;
+  std::array<Tensor,110> cache;
+  bool sharded=false;
+  // A convolution whose input was already normalized, activated and padded by the previous
+  // Winograd convolution's fused epilogue: (tensor, the conv index it is prepared for, and the
+  // sharding it was prepared under — a mismatch there would be silently wrong, not a crash).
+  Tensor prepared; int prepared_for=-1; bool prepared_sharded=false;
+  // The conv index whose y store was suppressed because this conv reserved it for `prepared`.
+  // It is the RUNTIME half of the y_dead contract: the producer states an intent and this records
+  // it, so a later edit that stops the consumer taking the `pre` edge aborts instead of handing a
+  // never-written tensor to norm_pad. Cleared when the edge is taken, and asserted closed at the
+  // end of every chunk beside prepared_for, because a conv index restarts at the same values.
+  int y_suppressed_for=-1;
+  Decoder(const Weights& weights,int r,int n,const PG& p):w(weights),rank(r),world(n),pg(p){}
+
+  Tensor conv(Tensor x,int idx,int gamma=-1,bool activate=true,bool add_bias=true,const Tensor& residual={},TemporalTrim trim=TemporalTrim{0},int fg=-1,int fnext=-1,bool y_dead=false) {
+    RECORD_FUNCTION("wan.conv."+std::to_string(idx), std::vector<c10::IValue>());
+    const int temporal_trim=trim.frames;
+    int kt=w[idx].size(2),kh=w[idx].size(3),kw=w[idx].size(4);
+    TORCH_CHECK(temporal_trim>=0 && temporal_trim<kt,"Invalid causal prefix trim");
+    int pt=kt-1-temporal_trim,ph=(sharded&&world>1&&kh==3)?1:0;
+    Tensor p=x;
+    bool pre=(prepared_for==idx);
+    // The consumer asserts the producer's y_dead contract AT RUNTIME, not in prose. If a conv had
+    // its y store suppressed for this idx and this conv does not take the fused `prepared` edge,
+    // it is about to read a tensor that was never written -- so abort instead of producing a wrong
+    // video. The existing metadata checks at `pre` cannot catch this, because they live inside the
+    // branch that is not being taken.
+    TORCH_CHECK(pre || y_suppressed_for!=idx,
+                "conv ",idx,": its y store was suppressed FOR THIS CONV, which then did not take "
+                "the prepared edge; the tensor it is about to read is unwritten");
+    // Prefix pruning never uses the fused, full-length prepared-input edge.
+    // A future caller must implement a trimmed prepared layout before mixing them.
+    TORCH_CHECK(temporal_trim==0 || (!pre && fg<0 && fnext<0),
+                "Causal prefix trim cannot use or produce a fused prepared input");
+    if(pre) {
+      TORCH_CHECK(prepared.defined(),"conv ",idx," is marked prepared but no tensor was stored");
+      TORCH_CHECK(prepared_sharded==sharded,"prepared tensor crosses a change of sharding");
+      // The graph edge checks itself. pt/ph here are the consumer's own, read from this conv's
+      // weight, and the producer computed the same two from w[fnext] with fnext==idx -- so these
+      // equalities are true by construction today and are a proof obligation on any later edit:
+      // a wrong fnext would otherwise hand a wrong-shaped tensor to the convolution, and a
+      // wrong-but-same-shaped one would pass every check the kernel makes.
+      TORCH_CHECK(prepared.size(1)==x.size(1) && prepared.size(2)==x.size(2)+pt &&
+                  prepared.size(3)==x.size(3)+2*ph && prepared.size(4)==x.size(4),
+                  "conv ",idx,": prepared ",prepared.sizes()," does not fit input ",x.sizes(),
+                  " (pt ",pt,", ph ",ph,", sharded ",(int)sharded,")");
+      p=prepared; prepared=Tensor(); prepared_for=-1; prepared_sharded=false;
+      y_suppressed_for=-1;   // the suppressed y was reserved for this edge and this edge took it
+    }
+    else if(pt||ph||gamma>=0) p=norm_pad(x,gamma>=0?w[gamma]:Tensor(),cache[idx],pt,ph,activate);
+    // Memory: the unpadded input is dead once the padded copy exists (p holds it when no copy was made).
+    x=Tensor();
+    // The causal cache carries this convolution's input history to the next chunk. `p` may be the
+    // tensor the previous fused epilogue produced and this conv has just consumed, so the copy is
+    // ordered after that kernel by the stream they share -- a side stream here would race it.
+    if(pt) cache[idx]=last_frames(p,pt,ph);
+    if(ph) {
+      RECORD_FUNCTION("wan.halo",std::vector<c10::IValue>());
+      auto recv=halo_exchange(halo_pack(p),pg,rank,world);
+      halo_unpack(p,recv,rank,world);
+    }
+    const int Ho=p.size(3)-2*ph,Wo=p.size(4);
+    std::vector<int64_t> padding={0,ph?0:kh/2,kw/2};
+    bool use_custom=(p.size(4)==78 && kt==3 && kh==3 && w[idx].size(0)==1024 && w[idx].size(1)==1024) || idx==61 || idx==85
+                    || idx==71 || idx==93;
+    // Only the expensive spatial 3x3 convolutions use Winograd. Low-resolution
+    // convolutions remain direct implicit GEMMs; 1x1 and temporal-only layers
+    // use cuDNN. Each transformed GEMM accumulates bf16 operands in fp32.
+    int wi=-1;
+    const int wids[18]={44,47,50,53,56,59,66,69,74,77,80,83,88,91,96,99,102,105};
+    for(int i=0;i<18;i++)if(idx==wids[i])wi=i;
+    Tensor y;
+    if(wi>=0) {
+      auto transformed=winograd_input_fast(p,ph);
+      // Memory: the padded input is dead once transformed (its history went to the cache above), and the transform
+      // (2.25x the input) once every component product is enqueued; the caching allocator reuses both in stream order.
+      p=Tensor();
+      std::vector<Tensor> products;
+      for(int a=0;a<36;a++) {
+        auto input=transformed.select(0,a).unsqueeze(0).permute({0,4,1,2,3});
+        products.push_back(custom_conv(input,w[122+36*wi+a],0,0,3));
+      }
+      transformed=Tensor();
+      // The SiLU this epilogue applies is the CONSUMER's norm activation, not this conv's
+      // `activate`: fusing is valid only when the consumer would have activated. Every fused
+      // consumer in the decoder (the two convs of a residual block) has activate=true, so the
+      // test below holds by construction; a non-activating consumer would have to pass its own.
+      bool fuse=(fg>=0&&fnext>=0&&activate);
+      // write_y is false exactly where the caller asserts nobody reads this launch's y AND the
+      // fused `prepared` edge exists to carry the activation instead. It is recorded so the
+      // consumer can ASSERT the edge rather than trust it; see y_suppressed_for above.
+      const bool write_y=!(y_dead && fuse);
+      int pt2=pt,ph2=ph;
+      if(fuse) { pt2=w[fnext].size(2)-1; ph2=(sharded&&world>1&&w[fnext].size(3)==3)?1:0; }
+      auto out=wino_epilogue(products,Ho,Wo,add_bias?w[idx+1]:Tensor(),residual,
+                             fuse?w[fg]:Tensor(),fuse?cache[fnext]:Tensor(),pt2,ph2,write_y);
+      y=out.y;
+      if(fuse) { prepared=out.n; prepared_for=fnext; prepared_sharded=sharded;
+                 if(!write_y) y_suppressed_for=fnext; }
+    } else y=use_custom?custom_conv(p,w[idx],padding[1],padding[2],3):at::cudnn_convolution(p,w[idx],padding,{1,1,1},{1,1,1},1,false,false,false);
+    if(add_bias && wi<0) finish(y,w[idx+1],residual);
+    return y;
+  }
+
+  Tensor resnet(Tensor x,int base,bool change=false,int nx=-1,int nxg=-1) {
+    RECORD_FUNCTION("wan.resnet."+std::to_string(base),std::vector<c10::IValue>());
+    Tensor skip=change?conv(x,base+6,-1,true,false):x;
+    // Memory: moved into the first convolution, the block's input is released there once its norm has read it
+    // (skip still holds it when it is the residual).
+    auto y=conv(std::move(x),base+1,base,true,true,Tensor(),TemporalTrim{0},base+3,base+4, /*y_dead=*/true);
+    return conv(y,base+4,base+3,true,true,skip,TemporalTrim{0},nxg,nx);
+  }
+
+  Tensor attention(const Tensor& x) {
+    RECORD_FUNCTION("wan.attention",std::vector<c10::IValue>());
+    int T=x.size(2),H=x.size(3),W=x.size(4),C=x.size(1);
+    Tensor qkv=conv(x,5,4,false);
+    auto p=qkv.permute({0,2,3,4,1}).view({T,1,H*W,3*C});
+    Tensor y=attn_forward(p);
+    y=y.view({1,T,H,W,C}).permute({0,4,1,2,3});
+    return conv(y,7,-1,true,true,x);
+  }
+
+  // The output head: RMS norm + SiLU + the causal padding, the 3x1x1 convolution to 112 channels
+  // (108 real) and the 3x3 stencil with the bias and clamp (head.cu). Replaces conv(x,108,107),
+  // whose 12 output channels waste 244 of a 256-wide tile's columns.
+  Tensor output_head(const Tensor& x) {
+    RECORD_FUNCTION("wan.head",std::vector<c10::IValue>());
+    const int pt=2;
+    int ph=world>1?1:0;
+    Tensor p=norm_pad(x,w[107],cache[108],pt,ph,true);
+    if(ph) {
+      RECORD_FUNCTION("wan.halo",std::vector<c10::IValue>());
+      // The third halo site. PR204 deleted pack_edges/unpack_edges, and the resolution that
+      // compiles-and-passes while silently dropping the vectorised pack is to re-add them:
+      // this is the same two pad rows through comm.cu's halo_pack/halo_unpack, which is what
+      // the other two sites (conv/resnet) already use. Census: post_edges == wait_edges == 3.
+      auto recv=halo_exchange(halo_pack(p),pg,rank,world);
+      halo_unpack(p,recv,rank,world);
+    }
+    cache[108]=last_frames(p,pt,ph);
+    Tensor g=conv_n128(p,w[108]);
+    int rows=x.size(3);
+    auto out=at::empty({1,12,g.size(2),rows,g.size(4)},x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+    // 352 is the pixel-plane height (44 latent rows x 8) and lo this rank's first pixel-plane row;
+    // both are the fixed 704x1248 geometry, and lo must stay 8*(44*rank/world) because run()
+    // passes lo*16 - the same row on the 704-tall video grid - to write_video.
+    head_stencil(g,w[109],out,352,ph,8*(44*rank/world));
+    return out;
+  }
+
+  Tensor up(Tensor input,int block,bool first) {
+    RECORD_FUNCTION("wan.up_block."+std::to_string(block),std::vector<c10::IValue>());
+    const int bases[4]={21,43,65,87};
+    int base=bases[block];
+    bool change=block>=2;
+    int b2=base+(change?8:6),b3=base+(change?14:12);
+    // Memory: only the resampler's DupUp3D shortcut reads the block's input again, and block 3 has none.
+    auto x=resnet(block==3?std::move(input):input,base,change,b2+1,b2);
+    x=resnet(x,b2,false,b3+1,b3);
+    x=resnet(x,b3,false);
+    if(block==3) return x;
+    int spatial=base+(change?20:18);
+    if(block<2) {
+      // The 3x1x1 causal time convolution writes the frame-interleaved tensor directly (a stride-2 D epilogue
+      // in solution/timeconv.cu), so `up_time` and `finish` are gone from this seam.  `activate=false` is
+      // INERT: with no gain the dispatcher takes norm.cu's copy-only mode and never reads the flag.
+      int start=first?1:0;
+      Tensor head=first?x.narrow(2,0,1):Tensor();
+      int nf=x.size(2)-start;
+      Tensor p;
+      if(nf>0) {
+        p=norm_pad(x.narrow(2,start,nf),Tensor(),cache[spatial+2],2,0,false);
+        cache[spatial+2]=last_frames(p,2,0);
+      }
+      x=time_upsample(p,head,w[spatial+2],w[spatial+3],x.size(2),nf,first);
+    }
+    if(world>1) {
+      x=norm_pad(x,Tensor(),Tensor(),0,1,false);
+      auto recv=halo_exchange(halo_pack(x),pg,rank,world);
+      halo_unpack(x,recv,rank,world);
+    }
+    // Four-phase spatial upsampler: the four 2x2 phase convolutions, then one bandwidth-bound pass that
+    // interleaves them with the bias and the DupUp3D shortcut (resample.cu).
+    std::vector<Tensor> pw;
+    for(int a=0;a<2;++a)for(int b=0;b<2;++b) pw.push_back(w[110+4*block+2*a+b]);
+    return resample_upsample(x,pw,w[spatial+1],input,world>1,block<2,first);
+  }
+
+  // Stateless exact prefix interval. Three causal convolutions require six
+  // preceding latent frames, but the support is pruned by two after each conv.
+  Tensor prefix_interval(const Tensor& z,int begin,int count) {
+    sharded=false;
+    const int end=begin+count;
+    const int lo6=std::max(0,begin-6),lo4=std::max(0,begin-4),lo2=std::max(0,begin-2);
+    auto raw=z.narrow(2,lo6,end-lo6).contiguous(at::MemoryFormat::ChannelsLast3d);
+    auto x=conv(raw,0);
+    // Reducing left temporal padding prunes the invalid outputs without ever
+    // computing them. Near global frame zero, the remaining padding is real.
+    x=conv(x,2,-1,true,true,Tensor(),TemporalTrim{lo4-lo6});
+    auto y=conv(x,10,9,true,true,Tensor(),TemporalTrim{lo2-lo4});
+    // Batch is one and T is the outer physical NDHWC dimension, so this
+    // frame slice remains dense for finish()'s linear residual indexing.
+    auto skip=x.narrow(2,begin-lo4,count);
+    TORCH_CHECK(skip.is_contiguous(at::MemoryFormat::ChannelsLast3d),"Prefix residual must be dense NDHWC");
+    y=conv(y,13,12,true,true,skip,TemporalTrim{begin-lo2});
+    return attention(y);
+  }
+
+  Tensor distributed_prefix(const Tensor& z,int start,int count) {
+    // No persistent prefix history: temporal ownership changes from batch to
+    // batch. The full latents on each rank provide exact predecessor context.
+    Decoder prefix(w,rank,world,pg);
+    Tensor x;
+    if(count>=2*world) {
+      const int lo=count*rank/world,hi=count*(rank+1)/world;
+      x=prefix.prefix_interval(z,start+lo,hi-lo);
+      x=prefix_temporal_to_rows(x,count,rank,world,pg);
+    } else {
+      // Tiny clips/tails do not amortize the transpose, and may have empty
+      // temporal owners. A stateless replicated prefix remains exact.
+      x=prefix.prefix_interval(z,start,count);
+      const int lo=44*rank/world,hi=44*(rank+1)/world;
+      x=x.narrow(3,lo,hi-lo).contiguous(at::MemoryFormat::ChannelsLast3d);
+    }
+    return x;
+  }
+
+  Tensor tail(Tensor x,bool first) {
+    // Every prepared edge is consumed within its residual/up-block chain.
+    // Tail history survives chunks, but no prepared current-chunk tensor may.
+    TORCH_CHECK(prepared_for<0 && !prepared.defined() && y_suppressed_for<0,
+                "Prepared input crossed a tail-chunk boundary");
+    sharded=true;
+    x=resnet(x,15);
+    for(int i=0;i<4;++i) x=up(std::move(x),i,first);
+    return output_head(x);
+  }
+
+  Tensor chunk(const Tensor& z,bool first) {
+    sharded=false;
+    // The invariant tail() checks must be CHECKED here too, not cleared. Clearing it first makes
+    // tail()'s check vacuous on the only path world 1 takes, and a leaked fused tensor is not
+    // merely retained: prepared_for is a conv INDEX and every chunk restarts at the same indices,
+    // so the next chunk's conv of that index CONSUMES it (`pre`) and the video is silently wrong
+    // at world 1 while world 4 errors out.
+    //
+    // The reset was dead on this tree, so the check costs one compare per chunk and nothing else,
+    // but the closure argument is NOT "every fused pair lives inside one resnet()" -- that is
+    // false. resnet(b)'s second conv is called with (nxg=b2, nx=b2+1) and produces for
+    // fnext=nx=b2+1, which is the FIRST conv of the next resnet; the same edge appears at every
+    // up block (47->50, 53->56 at block 1; 69->74, 77->80 at block 2). The chain closes at the
+    // END OF THE UP BLOCK instead: up() runs resnet(base, nxg=b2, nx=b2+1), then
+    // resnet(b2, nxg=b3, nx=b3+1), then resnet(b3) with the DEFAULT nx=-1, so the third resnet's
+    // last conv has fg=fnext=-1, produces nothing, and no fused tensor leaves up(). up() is the
+    // only Winograd caller and nothing after it is Winograd (output_head goes through conv_n128),
+    // so the triple is clear at every chunk entry.
+    // The predicate is tail()'s, EXACTLY, and that identity is the point rather than a
+    // coincidence: chunk() must accept everything tail() accepts, or the subset property this
+    // change exists to establish is not established. A third conjunct would also be unreachable --
+    // all three of prepared / prepared_for / prepared_sharded are written together at the single
+    // producer site (conv()'s `if(fuse)`) and cleared together at the single consumption site, so
+    // `prepared_sharded` can never differ from `prepared_for` and cannot fail on its own.
+    TORCH_CHECK(prepared_for<0 && !prepared.defined(),
+                "Prepared input crossed a chunk boundary");
+    auto x=conv(z,0);
+    x=conv(x,2);
+    x=resnet(x,9);
+    x=attention(x);
+    if(world>1) {
+      int lo=44*rank/world,hi=44*(rank+1)/world;
+      x=x.narrow(3,lo,hi-lo).contiguous(at::MemoryFormat::ChannelsLast3d);
+    }
+    return tail(x,first);
+  }
+};
+
+void decode_video(const Tensor& z,const Weights& weights,int64_t rank,int64_t world,const PG& pg,Tensor out,
+                  int64_t frames_per_pass) {
+  at::NoGradGuard no_grad;
+  c10::cuda::CUDAGuard device_guard(z.device());
+  TORCH_CHECK(z.scalar_type()==at::kBFloat16&&z.dim()==5&&z.size(0)==1&&z.size(1)==48&&z.size(2)>=1&&z.size(3)==44&&z.size(4)==78,"Expected bf16 [1,48,T,44,78]");
+  TORCH_CHECK(world>=1&&world<=8&&rank>=0&&rank<world,"Invalid rank/world");
+  TORCH_CHECK(weights.size()==770,"Call prepare() once before run()");
+  TORCH_CHECK(frames_per_pass>=1,"frames_per_pass must be at least 1, not ",frames_per_pass);
+  halo_stats()=HaloStats{};  // per call: the debug tail below prints THIS run's exchange, not the process's
+  Decoder decoder(weights,rank,world,pg);
+  // A pass carries frames_per_pass latent frames per GPU, frames_per_pass * world in all. The working memory of a GPU
+  // grows with it (3.5-4.2 GiB per latent frame at 704 x 1248), and so does the speed: a pass's GEMMs at 44 x 78 and
+  // 88 x 156 fill the GPU only with enough frames, and above world 1 every pass also pays a fixed cost (its prefix
+  // and its halo exchanges). The causal history crosses passes in `cache`, so the video does not depend on the size.
+  // At most 61 latent frames: the largest pass this tree was measured with (wino_epilogue.cu's indices are 64-bit).
+  static const int64_t kMaxChunk=61;
+  const int chunk_size=(int)std::min<int64_t>(frames_per_pass*world,kMaxChunk);
+  int T=z.size(2);
+  for(int start=0;start<T;start+=chunk_size) {
+    int count=std::min((int)chunk_size,T-start);
+    Tensor pixels;
+    if(world>1) pixels=decoder.tail(decoder.distributed_prefix(z,start,count),start==0);
+    else {
+      auto input=z.narrow(2,start,count).contiguous(at::MemoryFormat::ChannelsLast3d);
+      pixels=decoder.chunk(input,start==0);
+    }
+    int frame_offset=start==0?0:4*start-3;
+    if(world==1) write_video(pixels,out,frame_offset,0);
+    else if(rank==0) {
+      write_video(pixels,out,frame_offset,0);
+      for(int peer=1;peer<world;++peer) {
+        int lo=44*peer/world,hi=44*(peer+1)/world;
+        auto part=nhwc_empty(12,pixels.size(2),(hi-lo)*8,624,pixels);
+        std::vector<Tensor> buffers={part};
+        pg->recv(buffers,peer,0)->wait();
+        write_video(part,out,frame_offset,lo*16);
+      }
+    } else {
+      std::vector<Tensor> buffers={pixels};
+      pg->send(buffers,0,0)->wait();
+    }
+  }
+  // Debug only, and never set by CI: cached once per process so the measured span runs no getenv.
+  static const bool show_halo_stats=std::getenv("WAN_HALO_STATS")!=nullptr;
+  if(show_halo_stats) {
+    const HaloStats& s=halo_stats();
+    printf("[halo rank %d/%d] exchanges %lld  sent %.1f MB  received %.1f MB\n",(int)rank,(int)world,
+           (long long)s.exchanges,1e-6*s.bytes_sent,1e-6*s.bytes_received);
+    fflush(stdout);
+  }
+}
+
+// The default pass: 4 latent frames per GPU. On one GB200 at 704 x 1248 it peaks at 21-22 GiB allocated and takes
+// 1.09x the time of 16 frames per GPU, which peaks at 60-61 GiB (8 to 20 s of video).
+constexpr int64_t kFramesPerPass=4;
+
+void run(const Tensor& z,const Weights& weights,int64_t rank,int64_t world,const PG& pg,Tensor out) {
+  decode_video(z,weights,rank,world,pg,out,kFramesPerPass);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME,m) {
+  m.def("prepare",&prepare);
+  m.def("run",&run,py::call_guard<py::gil_scoped_release>());
+  m.def("run_with",&decode_video,py::call_guard<py::gil_scoped_release>());
+}
@@ -0,0 +1,213 @@
+// Halo edge transport: neighbour-only point to point and the seed's ring all-gather, over an
+// eight-channels-per-thread edge pack/scatter. Which arm runs is use_allgather's world rule below --
+// the ring only at world 2, neighbour-only at world 4 and at every world above it -- which WAN_HALO
+// can force either way for a measurement.
+// THIS POLICY IS STATED IN THREE PLACES AND THEY MUST AGREE: this file header, the block comment on
+// use_allgather, and the WAN_HALO contract in comm.h. An earlier revision of this change corrected one
+// of the three and was rightly blocked on the other two, so the cross-reference is the fix for the class
+// and not for the instance.
+//
+// Why the point to point ops go in one coalesced group: ProcessGroupNCCL launches an unbatched p2p op as a
+// group of its own on the group's stream, so a rank's p2p ops are stream-ordered against each other, and two
+// ranks each waiting to send into the other's not-yet-issued receive is a cycle. (NCCL documents that
+// serialization; w24's M1387 reaches the same deadlock from a direct all-to-all of unbatched p2p ops.)
+// ProcessGroup::startCoalescing/endCoalescing puts all of this rank's neighbour ops into one NCCL group, which
+// removes that ordering: the exchange is one hop, posted and closed with no host block.
+#include "comm.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <cstdlib>
+#include <cstring>
+
+namespace {
+
+cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+
+// One uint4 is eight bf16 channels; C is a multiple of 8 for every layer that sharded space (256/512/1024).
+__global__ void pack_edge8(const uint4* __restrict__ x, uint4* __restrict__ y, int64_t n8, int CV, int T, int H,
+                           int W) {
+  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t)gridDim.x * blockDim.x) {
+    int64_t row = i / (W * CV), rv = i % (W * CV);
+    int t = (int)(row % T), edge = (int)(row / T);
+    y[i] = x[((int64_t)t * H + (edge ? H - 2 : 1)) * W * CV + rv];
+  }
+}
+
+// Scatters the two received edges into the pad rows, eight channels per thread.
+__global__ void unpack_edge8(uint4* __restrict__ x, const uint4* __restrict__ e, int64_t n8, int CV, int T, int H,
+                             int W, int rank, int world) {
+  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t)gridDim.x * blockDim.x) {
+    int64_t row = i / (W * CV), rv = i % (W * CV);
+    int t = (int)(row % T), edge = (int)(row / T);
+    int peer = rank + (edge ? 1 : -1);  // pad row 0 comes from rank-1, pad row H-1 from rank+1
+    uint4 v = make_uint4(0u, 0u, 0u, 0u);
+    if (peer >= 0 && peer < world) v = e[((int64_t)peer * 2 * T + (int64_t)(1 - edge) * T + t) * W * CV + rv];
+    x[((int64_t)t * H + (edge ? H - 1 : 0)) * W * CV + rv] = v;
+  }
+}
+
+int grid_for(int64_t work) {
+  int64_t blocks = (work + 255) / 256;
+  return (int)(blocks > 2048 ? 2048 : (blocks < 1 ? 1 : blocks));
+}
+
+// Where each transport is used, when WAN_HALO does not force one. With P the bytes of one edge (half the packed
+// buffer): the ring sends (world-1)*2P and receives world*2P, of which this exchange reads 2P (one edge from each
+// neighbour); the neighbour path sends and receives one edge per neighbour -- 2P in and out for an interior rank,
+// P at the ends. So it is 6P -> 2P each way at world 4 and 2P/4P -> P/P at world 2, i.e. the received bytes fall
+// by 3x and 4x. So the policy is the ring only where it has been measured to cost nothing, and the
+// neighbour path everywhere else:
+//   world 2   the two arms measured the same, and the copy counts are equal there too. J817, three
+//             alternating rounds in one job at f121x2, minima 698.6 ms (ring) against 696.5 ms
+//             (neighbour), 0.3% apart, J570's 2.3% the other way sitting inside its own run-to-run
+//             spread; and J3057 ran none / ag / p2p at f481x2 in one job, 2305.6 / 2309.7 / 2293.2.
+//   world 4   neighbour. J2211's own `for A in ag p2p ag p2p` at f481x4: ag 2462.9 / 2244.5 against
+//             p2p 1010.6 / 1017.5, the p2p arm's A/A 0.68 %. The ring moves (world-1) copies each
+//             way, the neighbour path one.
+//   above 4   neighbour, and this rung is chosen on COPY COUNTS rather than on a measurement: the
+//             ring moves (world-1) copies each way against the neighbour path's 1, so the ratio grows
+//             with the world size, and no world above 4 has been measured on either arm here -- a
+//             cross-tray one has never run. The world-4 gap (2.32x) is the only scaled evidence and
+//             it is the smallest world where the two differ. Whoever first gets eight GPUs should
+//             treat this rung as the thing to measure, not as settled.
+// What that leaves: every world the harness runs is unchanged -- 1 by the early return, 2 true, 4 false
+// -- and only world 3 and world 5 and up change.
+// WAN_HALO forces either arm for a measurement.
+bool use_allgather(int world) {
+  if (world < 2) return true;  // one rank has no neighbour to talk to
+  const int forced = halo_transport();
+  if (forced == kAllGather) return true;
+  if (forced == kP2P) return false;
+  return world == 2;
+}
+
+}  // namespace
+
+HaloStats& halo_stats() {
+  static HaloStats stats;
+  return stats;
+}
+
+int halo_transport() {
+  static const int mode = [] {
+    const char* s = std::getenv("WAN_HALO");
+    if (!s) return (int)kPolicy;
+    if (std::strcmp(s, "ag") == 0 || std::strcmp(s, "allgather") == 0) return (int)kAllGather;
+    if (std::strcmp(s, "p2p") == 0 || std::strcmp(s, "neighbour") == 0) return (int)kP2P;
+    return (int)kPolicy;
+  }();
+  return mode;
+}
+
+at::Tensor halo_pack(const at::Tensor& padded) {
+  const int C = (int)padded.size(1), T = (int)padded.size(2), H = (int)padded.size(3), W = (int)padded.size(4);
+  TORCH_CHECK(C % 8 == 0, "halo_pack needs C a multiple of 8");
+  at::Tensor e = at::empty({2, T, W, C}, padded.options());
+  const int CV = C / 8;
+  int64_t n8 = (int64_t)2 * T * W * CV;
+  pack_edge8<<<grid_for(n8), 256, 0, stream()>>>((const uint4*)padded.data_ptr(), (uint4*)e.data_ptr(), n8, CV, T, H,
+                                                 W);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return e;
+}
+
+void halo_unpack(at::Tensor& padded, const at::Tensor& edges, int rank, int world) {
+  const int C = (int)padded.size(1), T = (int)padded.size(2), H = (int)padded.size(3), W = (int)padded.size(4);
+  TORCH_CHECK(C % 8 == 0, "halo_unpack needs C a multiple of 8");
+  const int CV = C / 8;
+  int64_t n8 = (int64_t)2 * T * W * CV;
+  unpack_edge8<<<grid_for(n8), 256, 0, stream()>>>((uint4*)padded.data_ptr(), (const uint4*)edges.data_ptr(), n8, CV,
+                                                   T, H, W, rank, world);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+HaloHandle post_edges(const PG& pg, const at::Tensor& send, at::Tensor& recv, int rank, int world,
+                      cudaStream_t comm) {
+  HaloHandle h;
+  h.started = true;
+  // Both buffers are flat here: the p2p slices are element offsets (narrowing the packed [2,T,W,C] on dim 0 would
+  // index a dimension of size 2, which is what round 1 of this PR did).
+  TORCH_CHECK(send.is_contiguous(), "halo edges must be contiguous");
+  TORCH_CHECK(recv.is_contiguous(), "the halo receive buffer must be contiguous");
+  h.send = send.reshape({-1});
+  h.recv = recv.reshape({-1});
+  HaloStats& s = halo_stats();
+  ++s.exchanges;
+  const int64_t bytes = send.numel() * (int64_t)send.element_size();
+  // The transport is ordered on `comm`: ProcessGroupNCCL makes its own stream wait for the current stream and
+  // the current stream wait for its end event, so this guard is what puts the collective after the pack.
+  c10::cuda::CUDAStreamGuard guard(c10::cuda::getStreamFromExternal(comm, c10::cuda::current_device()));
+  if (use_allgather(world)) {
+    // A ring all-gather: each rank's buffer goes round the ring, (world-1) copies each way.
+    at::Tensor in = send;
+    h.keep.push_back(in);
+    h.keep.push_back(recv);
+    h.works.push_back(pg->_allgather_base(recv, in));
+    s.bytes_sent += bytes * (world - 1);
+    s.bytes_received += bytes * (world - 1);
+    return h;
+  }
+  // The packed edges are [2,T,W,C]; the p2p slices below are element offsets, so they are taken flat (above).
+  h.keep.push_back(send);
+  h.keep.push_back(h.send);
+  h.keep.push_back(recv);
+  // Split on rank 0 = rows [0, H/world), so rank+1 is the band BELOW me and rank-1 the band above. Each rank
+  // sends its first real row up to rank-1 and its last real row down to rank+1 -- one edge per neighbour, not
+  // both to each -- and reads back only the two neighbour slots of `recv`; its other world-2 slots are never
+  // written and stay uninitialised, which halo_unpack is careful not to read (it touches slot rank-1's edge 1
+  // and slot rank+1's edge 0 and nothing else). On the all-gather path every slot is filled. Both leave the
+  // same two edges in the same places.
+  const int64_t half = send.numel() / 2;  // one edge: [T,W,C]; edge 0 first, edge 1 second
+  const int neighbours = (rank > 0 ? 1 : 0) + (rank + 1 < world ? 1 : 0);
+  s.bytes_sent += (int64_t)neighbours * half * send.element_size();
+  s.bytes_received += (int64_t)neighbours * half * send.element_size();
+  // Every tensor operation happens before the group opens: a throw between startCoalescing and endCoalescing
+  // would leave the NCCL group open and hang every rank instead of failing this one.
+  const bool hasLower = rank + 1 < world, hasUpper = rank > 0;
+  std::vector<at::Tensor> recvLower, sendUpper, recvUpper, sendLower;
+  if (hasLower) {
+    at::Tensor v = h.recv.narrow(0, (int64_t)(rank + 1) * 2 * half, half);  // below: its edge 0
+    h.keep.push_back(v);
+    recvLower.push_back(v);
+  }
+  if (hasUpper) {
+    at::Tensor v = h.send.narrow(0, 0, half);  // above: my first real row
+    h.keep.push_back(v);
+    sendUpper.push_back(v);
+  }
+  if (hasUpper) {
+    at::Tensor v = h.recv.narrow(0, (int64_t)(rank - 1) * 2 * half + half, half);  // above: its edge 1
+    h.keep.push_back(v);
+    recvUpper.push_back(v);
+  }
+  if (hasLower) {
+    at::Tensor v = h.send.narrow(0, half, half);  // below: my last real row
+    h.keep.push_back(v);
+    sendLower.push_back(v);
+  }
+  pg->startCoalescing(c10::DeviceType::CUDA);
+  if (hasLower) pg->recv(recvLower, rank + 1, 0);
+  if (hasUpper) pg->send(sendUpper, rank - 1, 0);
+  if (hasUpper) pg->recv(recvUpper, rank - 1, 0);
+  if (hasLower) pg->send(sendLower, rank + 1, 0);
+  // The ops inside a coalesced group carry no work of their own: the group's work is the one to wait on.
+  h.works.push_back(pg->endCoalescing(c10::DeviceType::CUDA));
+  return h;
+}
+
+void wait_edges(HaloHandle& h) {
+  if (!h.started) return;
+  // Each wait installs the collective's completion as an event wait on the current stream and does not block
+  // the host, so the caller's next kernel on that stream is ordered after the transfer.
+  for (auto& w : h.works) w->wait();
+  h.works.clear();
+  h.started = false;
+}
+
+at::Tensor halo_exchange(const at::Tensor& send, const PG& pg, int rank, int world) {
+  at::Tensor recv = at::empty({world * send.numel()}, send.options());
+  HaloHandle h = post_edges(pg, send, recv, rank, world, at::cuda::getCurrentCUDAStream());
+  wait_edges(h);
+  return recv;
+}
@@ -0,0 +1,83 @@
+// Halo edge transport for the row-band split: the two boundary rows a 3x3 convolution needs from its
+// neighbours, moved with the cheapest mechanism available. Representation-agnostic by design: it carries
+// whatever `halo_pack` packs (today the normalized NDHWC bf16 edge rows; the fusion work may change that),
+// and it leaves the packed layout and `halo_unpack` untouched, so the decoded video is bit-identical to the
+// all-gather it replaces.
+//
+// Two ways in:
+//   halo_exchange(...)                 one call: allocates the recv buffer, transports and closes. What
+//                                      Decoder::conv uses today.
+//   post_edges(...) / wait_edges(...)   the split seam for overlapping the transport with interior compute:
+//                                      post enqueues the exchange and returns a handle without a host block,
+//                                      wait closes it. `halo_pack` may run on one stream and the transport on
+//                                      another: post_edges takes the stream the transport is ordered on.
+#pragma once
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <c10/util/intrusive_ptr.h>
+#include <cuda_runtime.h>
+#include <vector>
+
+using PG = c10::intrusive_ptr<c10d::ProcessGroup>;
+
+// Which transport the halo exchange uses, from WAN_HALO once per process. With nothing set the world-size rule
+// in comm.cu decides; the variable is for forcing an arm in a measurement (both must be the same on every rank
+// of a world: one rank in a collective and one in a p2p group would never meet).
+//   unset: the rule -- the ring all-gather only at world 2, where both arms have been measured equal, and the
+//          neighbour path at world 4 and at every world above it. (World 1 has no exchange at all: ph is 0
+//          there, so nothing calls this.)
+//   "ag" or "allgather": the ring collective, whatever the world size.
+//   "p2p" or "neighbour": the neighbour path, whatever the world size. This is what the rule selects above
+//          world 4, and the neighbour path across trays has never been measured: nothing at that world size
+//          has, on either arm.
+enum HaloTransport { kPolicy = 0, kAllGather = 1, kP2P = 2 };
+int halo_transport();
+
+// The exchange sites and the bytes they moved. Debug instrumentation only: nothing on the run path reads them,
+// and with WAN_HALO_STATS unset they cost one increment per exchange.
+struct HaloStats {
+  int64_t exchanges = 0, bytes_sent = 0, bytes_received = 0;
+};
+HaloStats& halo_stats();
+
+// Gather the two edges a rank must send: the first and the last real row of the padded NDHWC tensor
+// [1,C,T,H,W] (H includes exactly one pad row per side), as [2,T,W,C] contiguous, edge-major: edge 0 is the
+// first real row (H index 1), edge 1 the last (H index H-2). Same layout as the seed's pack_edges, eight
+// channels per thread.
+at::Tensor halo_pack(const at::Tensor& padded);
+
+// Scatter the received edges into the pad rows: row 0 takes rank-1's edge 1, row H-1 takes rank+1's edge 0,
+// and a rank with no neighbour on that side gets zeros. `edges` is the transport's `recv` buffer, and rank+1
+// is the band *below* (higher rows), since the split gives rank 0 the rows [0, H/world).
+void halo_unpack(at::Tensor& padded, const at::Tensor& edges, int rank, int world);
+
+// The transport, split in two. `send` is halo_pack's [2,T,W,C]; `recv` is the caller's buffer of
+// world*2*T*W*C elements in the all-gather's layout (slot r is rank r's two edges), which is what
+// halo_unpack reads. The handle owns the works and keeps every send and recv view alive until it is waited.
+struct HaloHandle {
+  std::vector<c10::intrusive_ptr<c10d::Work>> works;
+  std::vector<at::Tensor> keep;
+  at::Tensor send, recv;  // flat views the p2p slices index
+  bool started = false;
+  HaloHandle() = default;
+  HaloHandle(HaloHandle&&) = default;
+  HaloHandle& operator=(HaloHandle&&) = default;
+  HaloHandle(const HaloHandle&) = delete;
+  HaloHandle& operator=(const HaloHandle&) = delete;
+};
+
+// Enqueue only: no host block, and nothing added to the compute stream. `comm` is the stream the transport is
+// ordered on; it must already carry whatever ordering the caller needs (a wait on the pack, say), and the
+// unpack must not run on a different stream before wait_edges has been called on it.
+HaloHandle post_edges(const PG& pg, const at::Tensor& send, at::Tensor& recv, int rank, int world,
+                      cudaStream_t comm);
+
+// Close it: this orders the CALLER'S CURRENT STREAM after the transfer and returns at once. WorkNCCL::wait is
+// an event wait on the stream that is current at the call, not a CPU block. So the next kernel enqueued on
+// that stream sees the edges -- and a kernel enqueued on any other stream does not.
+void wait_edges(HaloHandle& h);
+
+// Allocate the recv buffer, post, close and return the edges in halo_unpack's layout.
+at::Tensor halo_exchange(const at::Tensor& send, const PG& pg, int rank, int world);
@@ -0,0 +1,78 @@
+// CUTLASS 4.2 Blackwell implicit convolution: 2-CTA 256x256x64 tiles,
+// TMA im2col loads, tcgen05 MMA, fp32 accumulation, bf16 epilogue.
+// Every included CUTLASS/CuTe dependency is vendored beside this source.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/convnd_problem_shape.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/conv/dispatch_policy.hpp"
+#include "cutlass/conv/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/conv/device/conv_universal_adapter.hpp"
+#include "cutlass/conv/kernel/conv_universal.hpp"
+using namespace cute;
+using ElementBF16=cutlass::bfloat16_t;
+using LayoutNDHWC=cutlass::layout::TensorNDHWC;
+using ArchTag=cutlass::arch::Sm100;
+using OperatorClass=cutlass::arch::OpClassTensorOp;
+constexpr auto Fprop=cutlass::conv::Operator::kFprop;
+
+template<int M,int N,int CL>
+at::Tensor launch(const at::Tensor& x,const at::Tensor& w,int ph,int pw,int phu=-1,int pwu=-1) {
+  if(phu<0)phu=ph;if(pwu<0)pwu=pw;
+  using Tile=Shape<Int<M>,Int<N>,Shape<_64>>;
+  using Cluster=Shape<Int<CL>,_1,_1>;
+  using Epilogue=typename cutlass::epilogue::collective::CollectiveBuilder<
+    ArchTag,OperatorClass,Tile,Cluster,cutlass::epilogue::collective::EpilogueTileAuto,
+    float,float,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+  using Mainloop=typename cutlass::conv::collective::CollectiveBuilder<
+    ArchTag,OperatorClass,Fprop,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,float,Tile,Cluster,
+    cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
+    cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+  using Problem=cutlass::conv::ConvProblemShape<Fprop,3>;
+  using Kernel=cutlass::conv::kernel::ConvUniversal<Problem,Mainloop,Epilogue>;
+  using Conv=cutlass::conv::device::ConvUniversalAdapter<Kernel>;
+  int C=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4);
+  int K=w.size(0),KT=w.size(2),KH=w.size(3),KW=w.size(4);
+  auto y=at::empty({1,K,D-KT+1,H+ph+phu-KH+1,W+pw+pwu-KW+1},x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+  Problem problem(cutlass::conv::Mode::kCrossCorrelation,
+                  {1,D,H,W,C},{x.stride(0),x.stride(2),x.stride(3),x.stride(4),x.stride(1)},
+                  {K,KT,KH,KW,C},{w.stride(0),w.stride(2),w.stride(3),w.stride(4),w.stride(1)},
+                  {0,ph,pw},{0,phu,pwu},{1,1,1},{1,1,1},1);
+  typename Kernel::StrideC sc{};
+  typename Kernel::StrideD sd{};
+  cute::for_each(cute::make_seq<cute::rank<0>(sc)>{},[&](auto i){cute::get<0,i>(sc)=problem.stride_C[Problem::RankT-2-i];});
+  cute::for_each(cute::make_seq<cute::rank<0>(sd)>{},[&](auto i){cute::get<0,i>(sd)=problem.stride_C[Problem::RankT-2-i];});
+  typename Conv::Arguments args{problem,{(const ElementBF16*)x.data_ptr(),(const ElementBF16*)w.data_ptr()},
+    {{1.f,0.f},(const ElementBF16*)y.data_ptr(),sc,(ElementBF16*)y.data_ptr(),sd}};
+  Conv conv;
+  TORCH_CHECK(conv.can_implement(args)==cutlass::Status::kSuccess,"CUTLASS cannot implement shape");
+  size_t size=Conv::get_workspace_size(args);
+  auto workspace=at::empty({(int64_t)size},x.options().dtype(at::kByte));
+  auto stream=at::cuda::getCurrentCUDAStream();
+  auto status=conv.initialize(args,workspace.data_ptr(),stream);
+  TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS initialization failed ",int(status));
+  status=conv.run(stream);
+  TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS convolution failed ",int(status));
+  return y;
+}
+at::Tensor custom_conv(const at::Tensor& x,const at::Tensor& w,int64_t ph,int64_t pw,int64_t variant) {
+  return launch<256,256,2>(x,w,ph,pw);
+}
+
+// The output head's 3x1x1 convolution: N=128 because the head has 112 output channels, and a
+// 256-wide tile would spend 144 of its 256 columns on nothing. Measured at one clock (1951-2022
+// MHz, both arms in J1044 and J1343): N=128 is 0.83 PFLOP/s = 34% of peak, N=256 is 0.69 = 28%.
+// The time is what decides it and it is not close: 6.30 ms against 15.2 ms.
+at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w) {
+  return launch<256,128,2>(x,w,0,0);
+}
+
+at::Tensor phase_conv(const at::Tensor& x,const at::Tensor& w,int a,int b,bool halo) {
+  at::Tensor v=halo?x.narrow(3,a,x.size(3)-1):x;
+  return launch<256,256,2>(v,w,halo?0:1-a,1-b,halo?0:a,b);
+}
@@ -0,0 +1,82 @@
+// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to
+// [-1,1], from the 112-channel plane conv.cu's `conv_n128` (a 3x1x1 convolution) writes to the
+// 12-channel plane write_video consumes.
+//
+// Why the convolution writes 108 channels and this pass sums 9 of them. conv_out is 256x27 MACs
+// and its per-position INPUT WINDOW is 6912*2 B (the weight is 12x that and stays resident), i.e.
+// 12 FLOP/byte against the machine's ~375, so any form that reads the 27 taps per position is
+// operand-bound. Moving the nine spatial taps into the output channel dimension of a 3x1x1
+// (temporal-only) convolution makes the operand 768*2 B for 108 outputs instead of 13.8 KB for
+// 12, and the nine shifted copies are summed here.
+//
+// The value path, so that the rounding is on the record: the three temporal taps are accumulated
+// in the convolution's fp32 accumulator and rounded to bf16 in g; the nine spatial taps are then
+// summed here in fp32 from nine bf16 g values, the bias is added in fp32, and the result is
+// rounded once to bf16 on the store. Main's path is also two roundings (cuDNN's bf16 conv_out,
+// then finish_scalar's bf16 of x+bias); the count is the same and what differs is WHERE the first
+// one sits - after three taps instead of after twenty-seven - which is why max|d| moves and rms
+// does not. The reference accumulates all 27 taps in one fp32 accumulator.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_bf16.h>
+#include "head.h"
+
+// One thread per output position. Its 12 accumulators take the 12 channels of the same (i,j)
+// tap from nine neighbouring positions; the nine shifted reads of each 216-byte channel vector
+// are served from L1, so the pass moves each element of g about once.
+static __global__ void head_stencil_kernel(const __nv_bfloat16* __restrict__ g,
+                                    const __nv_bfloat16* __restrict__ bias,
+                                    __nv_bfloat16* __restrict__ out,
+                                    int D,int rows,int W,int H,int g_off,int lo) {
+  int idx=blockIdx.x*blockDim.x+threadIdx.x;
+  if(idx>=D*rows*W)return;
+  int w=idx%W, t=idx/W, r=t%rows, d=t/rows;
+  int gh0=lo+r;
+  const int gH=rows+2*g_off;
+  float acc[12];
+#pragma unroll
+  for(int c=0;c<12;++c)acc[c]=__bfloat162float(bias[c]);
+  for(int i=0;i<3;++i) {
+    int gh=gh0+i-1;
+    if(gh<0||gh>=H)continue;                       // conv3d's zero padding, in the plane
+    int gr=r+i-1+g_off;
+    if(gr<0||gr>=gH)continue;
+    for(int j=0;j<3;++j) {
+      int ww=w+j-1;
+      if(ww<0||ww>=W)continue;                     // conv3d's zero padding, across the row
+      const __nv_bfloat162* p=(const __nv_bfloat162*)(g+((((size_t)d*gH+gr)*W)+ww)*112+(i*3+j)*12);
+#pragma unroll
+      for(int k=0;k<6;++k) {
+        float2 f=__bfloat1622float2(p[k]);
+        acc[2*k]+=f.x;
+        acc[2*k+1]+=f.y;
+      }
+    }
+  }
+  __nv_bfloat162* o=(__nv_bfloat162*)(out+((((size_t)d*rows+r)*W)+w)*12);
+#pragma unroll
+  for(int k=0;k<6;++k)
+    o[k]=__floats2bfloat162_rn(fminf(1.f,fmaxf(-1.f,acc[2*k])),fminf(1.f,fmaxf(-1.f,acc[2*k+1])));
+}
+
+void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out,
+                  int64_t H_global,int64_t g_off,int64_t lo) {
+  int D=out.size(2),rows=out.size(3),W=out.size(4);
+  TORCH_CHECK(g.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&g.size(1)==112,
+              "head_stencil: g must be [1,112,D,rows+2*g_off,W] channels-last");
+  TORCH_CHECK(out.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&out.size(1)==12,
+              "head_stencil: out must be [1,12,D,rows,W] channels-last");
+  // The kernel strides g by gH = rows+2*g_off taken from `out` and reads D*rows*W positions of it,
+  // so a g of the wrong depth or height reads out of bounds rather than failing here.
+  TORCH_CHECK(g.size(2)==out.size(2)&&g.size(3)==out.size(3)+2*g_off&&g.size(4)==out.size(4),
+              "head_stencil: g must be [1,112,D,rows+2*g_off,W] of out's [1,12,D,rows,W]");
+  c10::cuda::CUDAGuard guard(g.device());
+  int64_t total=(int64_t)D*rows*W;
+  int threads=256;
+  int blocks=(int)((total+threads-1)/threads);
+  head_stencil_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(
+      (const __nv_bfloat16*)g.data_ptr(),(const __nv_bfloat16*)bias.data_ptr(),
+      (__nv_bfloat16*)out.data_ptr(),D,rows,W,(int)H_global,(int)g_off,(int)lo);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
@@ -0,0 +1,19 @@
+#pragma once
+#include <ATen/ATen.h>
+
+// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to
+// [-1,1], from the 112-channel plane the 3x1x1 convolution writes (conv.cu's `conv_n128`) to the
+// 12-channel plane write_video consumes.
+//   g:   [1,112,D,rows+2*g_off,W] channels-last-3d bf16, n = (i*3+j)*12 + c for the tap (i,j);
+//        channels 108..111 are the zero-weight lanes and are never read.
+//   out: [1,12,D,rows,W] channels-last-3d bf16, written whole.
+// H_global is the height of the plane this stencil filters (352 = 44 latent rows x 8; the pixel
+// plane is 704 tall and run()'s video row is twice this one). lo is the first row of this rank's
+// band IN THAT PLANE, 8*(44*rank/world), and it must stay that: run() passes lo*16 to write_video,
+// the same row on the 704-tall video grid, and the two derivations have to agree.
+void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out,
+                  int64_t H_global,int64_t g_off,int64_t lo);
+
+// conv.cu: the head's 3x1x1 convolution, at the tile its 112 output channels want. Declared here
+// so the two-sided declaration of a CUTLASS tile cannot drift from its definition in conv.cu.
+at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w);
@@ -0,0 +1,99 @@
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+using B = __nv_bfloat16;
+union Pack { uint4 v; B b[8]; __nv_bfloat162 h[4]; };
+static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+static const B* ptr(const Tensor& t) { return t.defined() ? (const B*)t.data_ptr<at::BFloat16>() : nullptr; }
+static B* ptr(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
+Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like) {
+  return at::empty({1,c,t,h,w},like.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+}
+Tensor last_frames(const Tensor& x,int n,int ph) {
+  int C=x.size(1),T=x.size(2),H=x.size(3)-2*ph,W=x.size(4);
+  Tensor y=nhwc_empty(C,n,H,W,x);
+  size_t row=(size_t)H*W*C*sizeof(B),src_pitch=(size_t)(H+2*ph)*W*C*sizeof(B);
+  const B* source=ptr(x)+((int64_t)(T-n)*(H+2*ph)+ph)*W*C;
+  C10_CUDA_CHECK(cudaMemcpy2DAsync(ptr(y),row,source,src_pitch,row,n,cudaMemcpyDeviceToDevice,stream()));
+  return y;
+}
+// The halo edge pack/scatter lives in comm.cu (halo_pack/halo_unpack, vectorised), with the transport.
+__global__ void finish_vector(B* x,const B* b,const B* s,int64_t n,int C) {
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n/8;i+=(int64_t)gridDim.x*blockDim.x) {
+    Pack v, bias, skip;
+    v.v=((uint4*)x)[i];bias.v=((const uint4*)b)[i%(C/8)];
+    if(s) skip.v=((const uint4*)s)[i];
+    #pragma unroll
+    for(int j=0;j<4;j++) { v.h[j]=__hadd2(v.h[j],bias.h[j]); if(s) v.h[j]=__hadd2(v.h[j],skip.h[j]); }
+    ((uint4*)x)[i]=v.v;
+  }
+}
+__global__ void finish_scalar(B* x,const B* b,const B* s,int64_t n,int C) {
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n;i+=(int64_t)gridDim.x*blockDim.x) {
+    B v=__float2bfloat16(__bfloat162float(x[i])+__bfloat162float(b[i%C]));
+    if(s) v=__float2bfloat16(__bfloat162float(v)+__bfloat162float(s[i])); x[i]=v;
+  }
+}
+void finish(Tensor& x,const Tensor& b,const Tensor& s) {
+  int C=x.size(1);int64_t n=x.numel();
+  if(C%8==0) finish_vector<<<std::min((int64_t)2048,(n+2047)/2048),256,0,stream()>>>(ptr(x),ptr(b),ptr(s),n,C);
+  else finish_scalar<<<std::min((int64_t)2048,(n+255)/256),256,0,stream()>>>(ptr(x),ptr(b),ptr(s),n,C);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+__global__ void write_kernel(const B* x,B* out,int64_t n,int T,int H,int W,int fullT,int fullH,int offsetT,int offsetH) {
+  int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;
+  if(i<n) {
+    int c=i%12;int64_t p=i/12;int w=p%W,h=(p/W)%H,t=p/((int64_t)W*H);
+    int color=c/4,dy=c%2,dx=(c%4)/2;
+    int64_t o=(((int64_t)color*fullT+t+offsetT)*fullH+offsetH+2*h+dy)*(2*W)+2*w+dx;
+    out[o]=__float2bfloat16(fminf(1.f,fmaxf(-1.f,__bfloat162float(x[i]))));
+  }
+}
+void write_video(const Tensor& x,Tensor& out,int frame_offset,int row_offset) {
+  write_kernel<<<(x.numel()+255)/256,256,0,stream()>>>(ptr(x),ptr(out),x.numel(),x.size(2),x.size(3),x.size(4),out.size(2),out.size(3),frame_offset,row_offset);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+// phases_kernel used to sit here: PR17 replaces it, together with finish_phases below, with one
+// bandwidth-bound pass in resample.cu that writes the interleaved tensor directly.
+// F(4x4,3x3), with fp32 spatial transforms and bf16 transformed operands.
+__device__ __forceinline__ void bt6(float v0,float v1,float v2,float v3,float v4,float v5,float* y) {
+  float a=v4-4*v2,b=v3-4*v1,c=v4-v2,d=2*(v3-v1);
+  y[0]=4*v0-5*v2+v4;y[1]=a+b;y[2]=a-b;y[3]=c+d;y[4]=c-d;y[5]=4*v1-5*v3+v5;
+}
+__global__ void winograd_input_kernel(const B* x,B* v,int C,int D,int H,int W,int PH,int TH,int TW) {
+  int HP=H+2*PH,CV=C/2;int64_t n=(int64_t)D*TH*TW*CV;
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n;i+=(int64_t)gridDim.x*blockDim.x) {
+    int c=i%CV,p=i/CV,tx=p%TW,ty=(p/TW)%TH,t=p/(TW*TH);
+    float d[2][6][6];
+    #pragma unroll
+    for(int a=0;a<6;a++) {
+      int y=4*ty-1+a+PH;
+      #pragma unroll
+      for(int b=0;b<6;b++) {int xx=4*tx-1+b;float2 f=(y>=0&&y<HP&&xx>=0&&xx<W)?__bfloat1622float2(((const __nv_bfloat162*)x)[(((int64_t)t*HP+y)*W+xx)*CV+c]):make_float2(0,0);d[0][a][b]=f.x;d[1][a][b]=f.y;}
+    }
+    #pragma unroll
+    for(int lane=0;lane<2;lane++) {
+      #pragma unroll
+      for(int b=0;b<6;b++) {float q[6];bt6(d[lane][0][b],d[lane][1][b],d[lane][2][b],d[lane][3][b],d[lane][4][b],d[lane][5][b],q);
+        #pragma unroll
+        for(int a=0;a<6;a++)d[lane][a][b]=q[a];}
+    }
+    #pragma unroll
+    for(int a=0;a<6;a++) {float q0[6],q1[6];bt6(d[0][a][0],d[0][a][1],d[0][a][2],d[0][a][3],d[0][a][4],d[0][a][5],q0);bt6(d[1][a][0],d[1][a][1],d[1][a][2],d[1][a][3],d[1][a][4],d[1][a][5],q1);
+      #pragma unroll
+      for(int b=0;b<6;b++)((__nv_bfloat162*)v)[(6*a+b)*n+i]=__floats2bfloat162_rn(q0[b],q1[b]);}
+  }
+}
+Tensor winograd_input(const Tensor& x,int ph) {
+  int C=x.size(1),D=x.size(2),H=x.size(3)-2*ph,W=x.size(4),TH=(H+3)/4,TW=(W+3)/4;
+  auto v=at::empty({36,D,TH,TW,C},x.options());int64_t n=(int64_t)D*TH*TW*C;
+  winograd_input_kernel<<<std::min((int64_t)2048,(n+255)/256),256,0,stream()>>>(ptr(x),ptr(v),C,D,H,W,ph,TH,TW);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();return v;
+}
+// (WinoPtrs, at6, winograd_output_kernel and winograd_output are deleted: with the fused inverse
+//  in wino_epilogue.cu nothing calls them. That file's a6v reads the same A^T out of shared memory
+//  instead of 36 separate global pointers, and it is the only copy of the transform left there.)
+
@@ -0,0 +1,21 @@
+#pragma once
+#include <ATen/ATen.h>
+using at::Tensor;
+Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like);
+Tensor norm_pad(const Tensor& x, const Tensor& gamma, const Tensor& history, int pt, int ph, bool silu);
+Tensor last_frames(const Tensor& padded, int n, int ph);
+void finish(Tensor& x, const Tensor& bias, const Tensor& residual);
+void write_video(const Tensor& x, Tensor& out, int frame_offset, int row_offset);
+
+Tensor winograd_input(const Tensor& x,int ph);
+Tensor winograd_input_fast(const Tensor& x,int ph);
+// winograd_output / winograd_output_kernel are deleted: the fused inverse in wino_epilogue.cu
+// replaced their only caller, and they had no live user left (PR172 round-1 review).
+
+// Fused inverse transform + bias/residual + RMS norm/SiLU.  y = A^T M A + bias + residual;
+// n (when gamma is given) = silu(rmsnorm(y)*sqrt(C)*gain) in the (D+pt, H+2ph, W, C) padded
+// layout the next convolution consumes, with its first pt frames taken from `hist`.
+struct WinoOut { Tensor y; Tensor n; };
+WinoOut wino_epilogue(const std::vector<Tensor>& m,int H,int W,const Tensor& bias,
+                      const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph,
+                      bool write_y=true);
@@ -0,0 +1,198 @@
+// norm.cu — RMS norm + SiLU + pad for the Wan2.2 VAE decoder, one warp per position.
+//
+// This replaces the seed's norm_vector<C> / pad_vector pair in kernels.cu with the same
+// call contract (norm_pad(x, gamma, hist, pt, ph, silu)) and the same layout (NDHWC,
+// [T+pt][H+2ph][W][C], history frames copied through, pad rows zeroed).
+//
+// Why it is faster (measured in probe/norm_probe2.cu and norm_probe3.cu, one GB200):
+//  - the seed is XU-pipe bound (w08's and w14's ncu: 78% SM busy, 16% DRAM, XU 75%):
+//    per element it pays cvt bf16->f32 for the sum, five cvts for the reference's three
+//    bf16 roundings, __expf and an IEEE divide. This does
+//      * bf16 -> f32 by an integer shift (ALU pipe, exact, no XU op),
+//      * one rounding for the whole normalize*sqrt(C)*gain product instead of three,
+//        which is also closer to the fp32 reference the decode is graded against,
+//      * SiLU as u + u*tanh(u) with u = 0.5*pre-activation (one tanh.approx.f32),
+//      * bf16 stores two elements at a time (cvt.rn.bf16x2.f32),
+//    i.e. ~1.5 XU ops per element instead of ~11.
+//  - rows are blocked so the per-position integer divisions of the seed's flat-pixel loop
+//    happen once per row, not once per position (at C=256 a position is only 8 elements).
+//  - P positions per warp in flight, the number that measured best per channel count.
+//
+// Numerics: one bf16 rounding for the whole normalize*gain product instead of three, and
+// SiLU in fp32 through tanh.approx.f32. That approximation's documented worst case is an
+// absolute error of about 2^-11 (5e-4) in tanh, which is harmless here because it enters
+// the result as u*2^-11 and the stored value is bf16, so the perturbation stays below the
+// bf16 quantum of u + u*tanh(u) over the range these pre-activations take. Measured, and
+// the place to look if that ever changes (bench.py, f121x1 vs the fp32 reference):
+//   seed   max 0.1709 (0.34x budget) rms 0.002242 (0.38x budget)
+//   this   max 0.1992 (0.40x budget) rms 0.002187 (0.37x budget)
+// i.e. the max moves up a little, the rms moves down, both far inside the budget.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+#include <cstdlib>
+#include <type_traits>
+
+using B = __nv_bfloat16;
+
+namespace {
+
+__device__ __forceinline__ void split2(uint32_t w, float& lo, float& hi) {
+  lo = __uint_as_float(w << 16);      // exact bf16 -> f32, ALU pipe
+  hi = __uint_as_float(w & 0xffff0000u);
+}
+__device__ __forceinline__ float tanh_approx(float x) {
+  float r; asm("tanh.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); return r;
+}
+__device__ __forceinline__ uint32_t pack2(float a, float b) {
+  uint32_t r; asm("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(r) : "f"(b), "f"(a)); return r;
+}
+
+// MODE 0: RMS norm * sqrt(C) * gain, then SiLU.  MODE 1: the same, no SiLU.  MODE 2: copy only.
+template <int C, int P, int MODE>
+__global__ void norm_pad_kernel(const B* __restrict__ x, const B* __restrict__ gain,
+                                const B* __restrict__ hist, B* __restrict__ y,
+                                int T, int H, int W, int PT, int PH) {
+  constexpr int CV = C / 8;      // uint4 per position
+  constexpr int NP = C / 256;    // uint4 per lane per position
+  const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5;
+  const int HP = H + 2 * PH;
+  const int64_t rows = (int64_t)(T + PT) * HP;
+  const float sqrtC = sqrtf((float)C);
+  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+    const int h = (int)(row % HP), t = (int)(row / HP);
+    const bool interior = (h >= PH && h < H + PH);
+    const bool current = (t >= PT);       // a history frame is copied through, never re-normalized
+    const int ph = h - PH;
+    const uint4* src = nullptr;
+    if (interior) {
+      if (current) src = (const uint4*)x + ((int64_t)(t - PT) * H + ph) * W * CV;
+      else if (hist) src = (const uint4*)hist + ((int64_t)t * H + ph) * W * CV;
+    }
+    for (int w0 = warp * P; w0 < W; w0 += warps * P) {
+      uint4 v[P][NP];
+      #pragma unroll
+      for (int q = 0; q < P; q++) {
+        const int w = w0 + q;
+        #pragma unroll
+        for (int j = 0; j < NP; j++) {
+          v[q][j] = make_uint4(0, 0, 0, 0);
+          if (src && w < W) v[q][j] = src[(int64_t)w * CV + lane + j * 32];
+        }
+      }
+      if (MODE < 2 && interior && current) {
+        float s[P][2];
+        #pragma unroll
+        for (int q = 0; q < P; q++) { s[q][0] = 0.f; s[q][1] = 0.f; }
+        #pragma unroll
+        for (int q = 0; q < P; q++) {
+          #pragma unroll
+          for (int j = 0; j < NP; j++) {
+            float f[8];
+            split2(v[q][j].x, f[0], f[1]); split2(v[q][j].y, f[2], f[3]);
+            split2(v[q][j].z, f[4], f[5]); split2(v[q][j].w, f[6], f[7]);
+            #pragma unroll
+            for (int k = 0; k < 8; k += 2) {
+              s[q][0] = fmaf(f[k], f[k], s[q][0]);
+              s[q][1] = fmaf(f[k + 1], f[k + 1], s[q][1]);
+            }
+          }
+        }
+        #pragma unroll
+        for (int q = 0; q < P; q++) {
+          float t0 = s[q][0] + s[q][1];
+          #pragma unroll
+          for (int d = 16; d; d >>= 1) t0 += __shfl_xor_sync(0xffffffff, t0, d);
+          // 0.5 * sqrt(C) / max(sqrt(sum), 1e-12): one multiply per element, u = 0.5*pre-activation
+          const float hs = 0.5f * sqrtC * rsqrtf(fmaxf(t0, 1e-24f));
+          #pragma unroll
+          for (int j = 0; j < NP; j++) {
+            const uint4 gn = ((const uint4*)gain)[lane + j * 32];
+            float g[8], xf[8], yf[8];
+            split2(gn.x, g[0], g[1]); split2(gn.y, g[2], g[3]);
+            split2(gn.z, g[4], g[5]); split2(gn.w, g[6], g[7]);
+            split2(v[q][j].x, xf[0], xf[1]); split2(v[q][j].y, xf[2], xf[3]);
+            split2(v[q][j].z, xf[4], xf[5]); split2(v[q][j].w, xf[6], xf[7]);
+            #pragma unroll
+            for (int k = 0; k < 8; k++) {
+              const float u = xf[k] * hs * g[k];
+              yf[k] = (MODE == 0) ? fmaf(u, tanh_approx(u), u) : (2.f * u);
+            }
+            v[q][j].x = pack2(yf[0], yf[1]); v[q][j].y = pack2(yf[2], yf[3]);
+            v[q][j].z = pack2(yf[4], yf[5]); v[q][j].w = pack2(yf[6], yf[7]);
+          }
+        }
+      }
+      #pragma unroll
+      for (int q = 0; q < P; q++) {
+        const int w = w0 + q;
+        if (w < W) {
+          uint4* dst = (uint4*)y + row * W * CV + (int64_t)w * CV + lane;
+          #pragma unroll
+          for (int j = 0; j < NP; j++) dst[j * 32] = v[q][j];
+        }
+      }
+    }
+  }
+}
+
+// Copy + pad for any C divisible by 8: the calls with no gain (post_quant_conv's input, the temporal convs).
+__global__ void copy_row_kernel(const uint4* __restrict__ x, const uint4* __restrict__ hist,
+                                uint4* __restrict__ y, int T, int H, int W, int PT, int PH, int CV) {
+  const int HP = H + 2 * PH;
+  const int64_t rows = (int64_t)(T + PT) * HP;
+  const int64_t nrow = (int64_t)W * CV;
+  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+    const int h = (int)(row % HP), t = (int)(row / HP);
+    const bool interior = (h >= PH && h < H + PH);
+    const uint4* src = nullptr;
+    if (interior) {
+      const int tt = t - PT, ph = h - PH;
+      if (tt >= 0) src = x + ((int64_t)tt * H + ph) * nrow;
+      else if (hist) src = hist + ((int64_t)t * H + ph) * nrow;
+    }
+    uint4* dst = y + row * nrow;
+    const uint4 zz = make_uint4(0, 0, 0, 0);
+    for (int64_t i = threadIdx.x; i < nrow; i += blockDim.x) dst[i] = src ? src[i] : zz;
+  }
+}
+
+}  // namespace
+
+at::Tensor norm_pad(const at::Tensor& x, const at::Tensor& gamma, const at::Tensor& hist,
+                    int pt, int ph, bool silu) {
+  const int C = (int)x.size(1), T = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);
+  auto y = at::empty({1, C, T + pt, H + ph * 2, W},
+                     x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+  const B* xp = (const B*)x.data_ptr<at::BFloat16>();
+  const B* gp = gamma.defined() ? (const B*)gamma.data_ptr<at::BFloat16>() : nullptr;
+  const B* hp = hist.defined() ? (const B*)hist.data_ptr<at::BFloat16>() : nullptr;
+  B* yp = (B*)y.data_ptr<at::BFloat16>();
+  const int64_t rows = (int64_t)(T + pt) * (H + ph * 2);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  const int blocks = (int)std::min<int64_t>(rows, 1024);
+  const int threads = 256;
+  auto go = [&](auto tag) {
+    constexpr int CC = decltype(tag)::value;
+    // P positions in flight per warp: measured best at 4 for C<=512, 2 at C=1024
+    constexpr int P = (CC == 1024) ? 2 : 4;
+    if (!gp) norm_pad_kernel<CC, P, 2><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
+    else if (silu) norm_pad_kernel<CC, P, 0><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
+    else norm_pad_kernel<CC, P, 1><<<blocks, threads, 0, stream>>>(xp, gp, hp, yp, T, H, W, pt, ph);
+  };
+  const bool fast = (C == 1024 || C == 512 || C == 256);
+  TORCH_CHECK(!gp || fast, "Unsupported RMS width ", C);
+  if (fast) {
+    if (C == 1024) go(std::integral_constant<int, 1024>{});
+    else if (C == 512) go(std::integral_constant<int, 512>{});
+    else go(std::integral_constant<int, 256>{});
+  } else {
+    TORCH_CHECK(C % 8 == 0, "norm_pad copy needs C divisible by 8, got ", C);
+    copy_row_kernel<<<blocks, threads, 0, stream>>>((const uint4*)xp, (const uint4*)hp, (uint4*)yp, T, H, W,
+                                                    pt, ph, C / 8);
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}
@@ -0,0 +1,56 @@
+// Prefix ownership transpose: no arithmetic, only bf16 NDHWC packets.
+#include "prefix.h"
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_runtime.h>
+
+namespace {
+// Blocks for destination bands are disjoint; each packet is [T,H_peer,W,C].
+__global__ void pack_prefix_rows(const uint4* x,uint4* send,int T,int H,int WC8,int world) {
+  const int peer=blockIdx.y;
+  const int lo=H*peer/world, hi=H*(peer+1)/world;
+  const int64_t n=(int64_t)T*(hi-lo)*WC8;
+  const int64_t offset=(int64_t)T*lo*WC8;
+  for(int64_t i=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;i<n;i+=(int64_t)gridDim.x*blockDim.x) {
+    const int64_t t=i/((int64_t)(hi-lo)*WC8);
+    const int64_t rowcol=i%((int64_t)(hi-lo)*WC8);
+    send[offset+i]=x[(t*H+lo)*WC8+rowcol];
+  }
+}
+}
+
+at::Tensor prefix_temporal_to_rows(const at::Tensor& local,int64_t total_frames,
+    int rank,int world,const c10::intrusive_ptr<c10d::ProcessGroup>& pg) {
+  const int C=local.size(1),T=local.size(2),H=local.size(3),W=local.size(4);
+  TORCH_CHECK(C%8==0 && local.is_contiguous(at::MemoryFormat::ChannelsLast3d),"Prefix needs aligned NDHWC");
+  TORCH_CHECK(T==total_frames*(rank+1)/world-total_frames*rank/world,"Prefix ownership mismatch");
+  auto send=at::empty({local.numel()},local.options());
+  const int band_height=H*(rank+1)/world-H*rank/world;
+  auto output=nhwc_empty(C,total_frames,band_height,W,local);
+  // ProcessGroup splits count bf16 elements, not bytes. Receive into the
+  // physical NDHWC order: the source ranks own chronological frame intervals.
+  auto physical=output.permute({0,2,3,4,1});
+  TORCH_CHECK(physical.is_contiguous(),"Prefix receive must be dense NDHWC");
+  auto recv=physical.view({-1});
+  if(T) {
+    dim3 grid(256,world);
+    pack_prefix_rows<<<grid,256,0,at::cuda::getCurrentCUDAStream()>>>(
+        (const uint4*)local.data_ptr(),(uint4*)send.data_ptr(),T,H,W*(C/8),world);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+  std::vector<int64_t> in_split(world),out_split(world);
+  int64_t sent=0,received=0;
+  for(int p=0;p<world;++p) {
+    in_split[p]=(int64_t)T*(H*(p+1)/world-H*p/world)*W*C;
+    out_split[p]=(total_frames*(p+1)/world-total_frames*p/world)*band_height*W*C;
+    sent+=in_split[p]; received+=out_split[p];
+  }
+  TORCH_CHECK(sent==send.numel() && received==recv.numel(),"Prefix packet size mismatch");
+  // WorkNCCL::wait orders the caller's stream after the collective; it need
+  // not block the host. ProcessGroup's tensor/allocator stream tracking keeps
+  // send/recv storage live through that dependency. Do not drop this wait
+  // when tail kernels consume output on the caller's current stream.
+  pg->alltoall_base(recv,send,out_split,in_split)->wait();
+  return output;
+}
@@ -0,0 +1,7 @@
+#pragma once
+#include <ATen/ATen.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+// Exchange chronological full-frame prefix slices for chronological row bands.
+// Source rank r owns [floor(T*r/world), floor(T*(r+1)/world)).
+at::Tensor prefix_temporal_to_rows(const at::Tensor& local,int64_t total_frames,
+    int rank,int world,const c10::intrusive_ptr<c10d::ProcessGroup>& pg);
@@ -0,0 +1,96 @@
+// Resampler stage: the four-phase spatial upsamplers of up_blocks 0,1,2 of the Wan2.2 VAE decoder.
+// See resample.h for why the phase convolutions stay direct and what this pass replaces.
+#include "resample.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+using B = __nv_bfloat16;
+// conv.cu: a 2x2 low-resolution convolution of one output parity, padded per phase; `halo` is set when
+// the input carries a one-row halo from the neighbour rank, so the input is narrowed instead of padded.
+at::Tensor phase_conv(const at::Tensor&, const at::Tensor&, int, int, bool);
+
+static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+static const B* pc(const Tensor& t) { return (const B*)t.data_ptr<at::BFloat16>(); }
+static B* pm(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
+
+union P8 { uint4 v; B b[8]; __nv_bfloat162 h[4]; };
+
+// One thread: eight channels of one 2x2 low-resolution tile, all four phases.  The grid carries the
+// coordinates (tx, t, ty) so there is no integer division on the hot path, and every access is a
+// 16-byte vector: four phase outputs and the DupUp3D source in, four interleaved positions out.
+__global__ void upsampler_kernel(const B* __restrict__ p0, const B* __restrict__ p1,
+                                 const B* __restrict__ p2, const B* __restrict__ p3,
+                                 B* __restrict__ y, const B* __restrict__ bias,
+                                 const B* __restrict__ skip, int C, int HL, int WL, int TC,
+                                 bool temporal, bool first) {
+  const int cg = threadIdx.x;                       // 8-channel group
+  const int tx = blockIdx.x, t = blockIdx.y, ty = blockIdx.z;
+  const int CP = C / 8, SP = TC / 8;
+  const P8 b = {((const uint4*)bias)[cg]};
+  const int st = temporal ? (first ? (t + 1) / 2 : t / 2) : t;
+  const int64_t wC = 2 * (int64_t)WL * C;            // one row of the interleaved output
+  const int64_t tC = 2 * (int64_t)HL * wC;           // one frame of it
+#pragma unroll
+  for (int i = 0; i < 2; ++i) {
+    const int yy = 2 * ty + i;
+    if (yy >= HL) continue;
+#pragma unroll
+    for (int j = 0; j < 2; ++j) {
+      const int xx = 2 * tx + j;
+      if (xx >= WL) continue;
+      const int64_t q0 = ((int64_t)t * HL + yy) * WL + xx;
+      P8 ph[4];
+      ph[0].v = ((const uint4*)p0)[q0 * CP + cg];
+      ph[1].v = ((const uint4*)p1)[q0 * CP + cg];
+      ph[2].v = ((const uint4*)p2)[q0 * CP + cg];
+      ph[3].v = ((const uint4*)p3)[q0 * CP + cg];
+      const int64_t s = ((int64_t)st * HL + yy) * WL + xx;
+      P8 s0, s1;
+      if (temporal) {
+        s0.v = ((const uint4*)skip)[s * SP + cg];
+        s1.v = s0.v;                                 // nearest x2 in time: both row parities alike
+      } else {
+        // up block 2: output channel J takes source channel 2*J + (row parity), so read both parities.
+        s0.v = ((const uint4*)skip)[s * SP + 2 * cg];
+        s1.v = ((const uint4*)skip)[s * SP + 2 * cg + 1];
+      }
+      B* out = y + (int64_t)t * tC + (int64_t)(2 * yy) * wC + (int64_t)(2 * xx) * C + 8 * cg;
+      if (!temporal) {
+        // channel J of the output takes source channel 2J + a: even channels from s0, odd from s1.
+        P8 e, o;
+#pragma unroll
+        for (int k = 0; k < 4; ++k) {
+          e.b[k] = s0.b[2 * k];
+          e.b[k + 4] = s1.b[2 * k];
+          o.b[k] = s0.b[2 * k + 1];
+          o.b[k + 4] = s1.b[2 * k + 1];
+        }
+        s0.v = e.v;
+        s1.v = o.v;
+      }
+      P8 r;
+#pragma unroll
+      for (int p = 0; p < 4; ++p) {
+        const P8& sv = (p & 2) ? s1 : s0;            // row parity of this phase
+#pragma unroll
+        for (int k = 0; k < 4; ++k) r.h[k] = __hadd2(__hadd2(ph[p].h[k], b.h[k]), sv.h[k]);
+        const int a = p >> 1, bcol = p & 1;
+        ((uint4*)(out + (int64_t)a * wC + (int64_t)bcol * C))[0] = r.v;
+      }
+    }
+  }
+}
+
+Tensor resample_upsample(const Tensor& x, const std::vector<Tensor>& pw, const Tensor& bias,
+                         const Tensor& skip, bool halo, bool temporal, bool first) {
+  TORCH_CHECK(pw.size() == 4, "resample_upsample needs the four phase kernels");
+  std::vector<Tensor> q(4);
+  for (int a = 0; a < 2; ++a) for (int b = 0; b < 2; ++b) q[2 * a + b] = phase_conv(x, pw[2 * a + b], a, b, halo);
+  const int C = q[0].size(1), D = q[0].size(2), HL = q[0].size(3), WL = q[0].size(4);
+  Tensor y = at::empty({1, C, D, 2 * HL, 2 * WL}, x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
+  dim3 block(C / 8, 1), grid((WL + 1) / 2, D, (HL + 1) / 2);
+  upsampler_kernel<<<grid, block, 0, stream()>>>(pc(q[0]), pc(q[1]), pc(q[2]), pc(q[3]), pm(y), pc(bias),
+                                                 pc(skip), C, HL, WL, (int)skip.size(1), temporal, first);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}
@@ -0,0 +1,51 @@
+#pragma once
+// Resampler stage of the Wan2.2 VAE decoder: the four-phase spatial upsamplers of up_blocks 0,1,2.
+//
+// A resampler is "nearest x2 in space, then a 3x3x3 convolution", which equals four 2x2 low-resolution
+// convolutions, one per output parity (their repeated taps already summed into per-phase kernels by
+// prepare()).  Those convolutions stay direct (4 taps, no transform: F(2x2,2x2) would cut the MACs 1.78x
+// but materialise 4 x 9/4 = 9 x the low-resolution tensor of transformed input and 2.25x the output of
+// products, which prices out as a wash - measured and posted).
+//
+// What is fixed here is the pass after them: the seed's interleave kernel walked the four phase tensors
+// and the DupUp3D source with four 64-bit division/modulo chains per 16 bytes and ran at 3.9 TB/s, half
+// of HBM.  This one walks tiles and channel groups with the grid carrying the coordinates, so the pass
+// is bandwidth bound: it reads the four phase outputs and the DupUp3D source once and writes the
+// interleaved high-resolution tensor once, with bias and the shortcut added in bf16 at the same
+// rounding points the seed used.
+#include <ATen/ATen.h>
+#include <vector>
+using at::Tensor;
+
+// One resampler over one latent chunk.  x is the low-resolution input, already padded by one row top
+// and bottom ([1,C,D,H+2,W] channels-last) with the halo installed at world > 1 (`halo`), and by the two
+// causal history frames in time; pw are the four per-phase 2x2 kernels ([C_out,C_in,1,2,2], phase =
+// 2*row_parity + col_parity), bias [C_out], and skip the DupUp3D source ([1,TC,Ds,HL,WL]).  `temporal`
+// is true for up_blocks 0 and 1 (the shortcut and the phase kernels are nearest-x2 in time as well),
+// false for up block 2, whose shortcut is the channel-pair-to-row rearrangement.  `first` marks the
+// chunk that holds latent frame 0.
+// Returns the interleaved high-resolution tensor [1,C_out,D,2*HL,2*WL] channels-last.
+//
+// The pass moves 100.5 GB of logical traffic for f121x1 (the four phase tensors read 42.5, the DupUp3D
+// source 15.5, the output written 42.5) in 14.8 ms = 6.8 TB/s.
+//
+// THAT 6.8 IS NOT A BOUND AND MUST NOT BE TRACED TO ONE.  It is solved from 100.5 / 14.8, so quoting it
+// against a ceiling is an identity - w06 M18769/M18811, w22 M18463.  The bounds on this board that could
+// apply are: the MIX ceiling for a read+write stage, 5.07-5.59 TB/s at this stage's own f_r (w08's
+// f_r-form, J2654/J2944); w14's per-thread store-loop write arm, 4.86-5.15 (J3082); and w12's 4-byte-
+// per-lane burst arm, 7.35-7.50 (J3029), which is BURST-FLATTERED at 2 KB against the tree's 128 B and
+// which w12 withdrew.  6.9 TB/s was a deflation of that withdrawn arm, so it inherits the arm's defect:
+// it has no f_r, it belongs to a limb with no read, and this is a read+write stage.
+//
+// So the honest statement is that the stage moves 100.5 GB at 6.8 TB/s and that NO CEILING FOR ITS OWN
+// STREAM HAS BEEN MEASURED.  The counters on this pass would supply one; that run has not been made.
+// Do not write "at its byte bound", and do not write "the model is conservative" (106.7 GB / 14.8 ms is
+// 7.21 TB/s, which is ABOVE the mixed ceiling - the same fact as "the model is 6 % low", stated with the
+// sign that does not flatter it).
+//
+// J2574 also counts 1.37x the sectors of its logical read.  That is sector waste on the access pattern
+// and not work the algorithm needs: this kernel reads each phase tensor and the shortcut at the
+// position's own address, with no neighbour re-read and no stencil anywhere in it.
+
+Tensor resample_upsample(const Tensor& x, const std::vector<Tensor>& pw, const Tensor& bias,
+                         const Tensor& skip, bool halo, bool temporal, bool first);
@@ -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":[]}
@@ -0,0 +1,122 @@
+// w11: a CUTLASS implicit convolution whose epilogue stores into a caller-given destination with arbitrary
+// (N,D,H,W) strides, and can accumulate on what the destination already holds (beta=1, source == destination).
+// One user: the temporal upsampler, which writes the frame-interleaved video straight out (stride 2 in D).
+// The resampler's phase tensors come from conv.cu's phase_conv, not from this file.
+// Nothing else changes: same tile (256x256x64), same 2-CTA cluster, same TMA im2col, same bf16 epilogue.
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include "cutlass/cutlass.h"
+#include "cute/tensor.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/conv/convolution.h"
+#include "cutlass/conv/convnd_problem_shape.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/conv/dispatch_policy.hpp"
+#include "cutlass/conv/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/conv/device/conv_universal_adapter.hpp"
+#include "cutlass/conv/kernel/conv_universal.hpp"
+at::Tensor nhwc_empty(int64_t c, int64_t t, int64_t h, int64_t w, const at::Tensor& like);
+using namespace cute;
+using B = __nv_bfloat16;
+using ElementBF16 = cutlass::bfloat16_t;
+using LayoutNDHWC = cutlass::layout::TensorNDHWC;
+using ArchTag = cutlass::arch::Sm100;
+using OperatorClass = cutlass::arch::OpClassTensorOp;
+constexpr auto Fprop = cutlass::conv::Operator::kFprop;
+
+// x: activation [1,C,D,H,W] channels-last; w: [K,C,KT,KH,KW] channels-last; y: the destination, with its own strides.
+template <int M, int N, int CL>
+static void launch_into(const at::Tensor& x, const at::Tensor& w, at::Tensor& y, int ph, int pw, float beta,
+                        int phu = -1, int pwu = -1) {
+  if (phu < 0) phu = ph;
+  if (pwu < 0) pwu = pw;
+  using Tile = Shape<Int<M>, Int<N>, Shape<_64>>;
+  using Cluster = Shape<Int<CL>, _1, _1>;
+  using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Tile, Cluster, cutlass::epilogue::collective::EpilogueTileAuto, float, float,
+      ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8,
+      cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
+  using Mainloop = typename cutlass::conv::collective::CollectiveBuilder<
+      ArchTag, OperatorClass, Fprop, ElementBF16, LayoutNDHWC, 8, ElementBF16, LayoutNDHWC, 8, float, Tile, Cluster,
+      cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
+      cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
+  using Problem = cutlass::conv::ConvProblemShape<Fprop, 3>;
+  using Kernel = cutlass::conv::kernel::ConvUniversal<Problem, Mainloop, Epilogue>;
+  using Conv = cutlass::conv::device::ConvUniversalAdapter<Kernel>;
+  int C = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4);
+  int K = w.size(0), KT = w.size(2), KH = w.size(3), KW = w.size(4);
+  int OD = y.size(2), OH = y.size(3), OW = y.size(4);
+  TORCH_CHECK(OD == D - KT + 1 && OH == H + ph + phu - KH + 1 && OW == W + pw + pwu - KW + 1,
+              "launch_into: destination shape does not match the convolution");
+  TORCH_CHECK(y.size(1) == K && y.stride(1) == 1, "launch_into: destination must be channel-innermost");
+  Problem problem(cutlass::conv::Mode::kCrossCorrelation,
+                  {1, D, H, W, C}, {x.stride(0), x.stride(2), x.stride(3), x.stride(4), x.stride(1)},
+                  {K, KT, KH, KW, C}, {w.stride(0), w.stride(2), w.stride(3), w.stride(4), w.stride(1)},
+                  {0, ph, pw}, {0, phu, pwu}, {1, 1, 1}, {1, 1, 1}, 1);
+  // The output's own strides: index 0 N, 1 D, 2 H, 3 W, 4 C (the layout the C strides already use).
+  problem.stride_C[0] = y.stride(0);
+  problem.stride_C[1] = y.stride(2);
+  problem.stride_C[2] = y.stride(3);
+  problem.stride_C[3] = y.stride(4);
+  problem.stride_C[4] = 1;
+  typename Kernel::StrideC sc{};
+  typename Kernel::StrideD sd{};
+  cute::for_each(cute::make_seq<cute::rank<0>(sc)>{},
+                 [&](auto i) { cute::get<0, i>(sc) = problem.stride_C[Problem::RankT - 2 - i]; });
+  cute::for_each(cute::make_seq<cute::rank<0>(sd)>{},
+                 [&](auto i) { cute::get<0, i>(sd) = problem.stride_C[Problem::RankT - 2 - i]; });
+  const ElementBF16* dst = (const ElementBF16*)y.data_ptr();
+  typename Conv::Arguments args{problem, {(const ElementBF16*)x.data_ptr(), (const ElementBF16*)w.data_ptr()},
+                                {{1.f, beta}, dst, sc, (ElementBF16*)y.data_ptr(), sd}};
+  Conv conv;
+  TORCH_CHECK(conv.can_implement(args) == cutlass::Status::kSuccess, "CUTLASS cannot implement shape");
+  size_t size = Conv::get_workspace_size(args);
+  auto workspace = at::empty({(int64_t)size}, x.options().dtype(at::kByte));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto status = conv.initialize(args, workspace.data_ptr(), stream);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS initialization failed ", int(status));
+  status = conv.run(stream);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "CUTLASS convolution failed ", int(status));
+}
+
+// The temporal upsampler: y[t] = w[:,:,2] x[t-2] + w[:,:,1] x[t-1] + w[:,:,0] x[t] + bias, with frame 0 passed
+// through unchanged when `first`, the two temporal halves of the 2C channels going to consecutive frames.
+// p is the activation already padded with its two causal history frames ([1,C,nf+2,H,W], channels-last).
+// INVARIANT: every element of y is written by head_bias_kernel, and the two parity arms below cover the
+// non-head frames disjointly, each exactly once -- which is what makes beta=1 onto an at::empty safe.
+static __global__ void head_bias_kernel(const B* head, const B* bias, B* y, int C, int T2, int64_t HW, int off) {
+  int64_t n = (int64_t)T2 * HW * (C / 8);
+  for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (int64_t)gridDim.x * blockDim.x) {
+    int cg = (int)(i % (C / 8)); int64_t r = i / (C / 8);
+    int64_t pos = r % HW; int t = (int)(r / HW);
+    uint4 v;
+    if (t < off) v = ((const uint4*)head)[pos * (C / 8) + cg];
+    else v = ((const uint4*)(bias + ((t - off) & 1) * C))[cg];
+    ((uint4*)y)[i] = v;
+  }
+}
+
+at::Tensor time_upsample(const at::Tensor& p, const at::Tensor& head, const at::Tensor& w, const at::Tensor& bias,
+                         int64_t T, int64_t nf, bool first) {
+  const at::Tensor& ref = p.defined() ? p : head;
+  int C = bias.size(0) / 2, H = ref.size(3), W = ref.size(4);
+  int T2 = (int)(2 * T - (first ? 1 : 0));
+  at::Tensor y = nhwc_empty(C, T2, H, W, ref);
+  int off = first ? 1 : 0;
+  head_bias_kernel<<<1024, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
+      first ? (const B*)head.data_ptr() : nullptr, (const B*)bias.data_ptr(), (B*)y.data_ptr(), C, T2,
+      (int64_t)H * W, off);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  if (nf > 0) {
+    int64_t sN = y.stride(0), sD = 2 * y.stride(2), sH = y.stride(3), sW = y.stride(4);
+    for (int par = 0; par < 2; par++) {
+      auto yv = y.as_strided({1, (int64_t)C, nf, (int64_t)H, (int64_t)W}, {sN, 1, sD, sH, sW},
+                             (int64_t)(off + par) * y.stride(2));
+      launch_into<256, 256, 2>(p, w.narrow(0, (int64_t)par * C, C), yv, 0, 0, 1.f);
+    }
+  }
+  return y;
+}
@@ -0,0 +1,400 @@
+// Fused F(4x4,3x3) inverse transform for the causal 3x3x3 Winograd convolutions.
+// One pass over the 36 component tensors M produces BOTH
+//   y = A^T M A + bias + residual         (raw conv output: the resnet's residual operand)
+//   n = silu(rmsnorm(y) * sqrt(C) * gain) (the next convolution's normalized, activated,
+//                                          temporal+spatial padded input)
+// so the normalized tensor is never written and read again by a separate norm pass.
+// The RMS statistic needs every output channel of a position, so one block holds all of C
+// for one 4x4 output tile: blockDim.x == C/2, one channel pair per thread.
+//
+// M layout: 36 separate allocations, each [D][TH][TW][C] bf16, C innermost, so a plane's
+// position stride is 1 in C and its plane offset is (((t*TH+ty)*TW+tx)*C + ch).
+// y is [D][H][W][C]; n is [D+PT][H+2PH][W][C].
+//
+// Two lessons from the guild, both measured elsewhere and both applied here:
+//  - w01's J576 (input transform): the head carries no divisions, no predicates and no 64-bit
+//    address chain. 3-D grid (tx,ty,t), a separate EXACT instantiation with the bounds test
+//    removed, and a 32-bit plane byte offset. EXACT is dispatched per call: it holds at
+//    88x156, 176x312 and 352x624, whose widths are multiples of 4 -- NOT at 44x78, where
+//    W=78 leaves the last tile two columns short and the checked path is taken.
+//  - w05's J627 (inverse transform, 5.5 TB/s at 64 registers): no thread loads 4 B of one plane.
+//    The block stages the 36 components of its tile into shared memory cooperatively -- one
+//    plane's C channels is C*2 contiguous bytes, so a warp pass is 16 B per lane and one fully
+//    coalesced 512 B request -- and the consumer reads smem[comp][2c] as a 4-byte pair, which is
+//    128 contiguous bytes per warp and therefore conflict-free. No transpose on either side.
+//
+// Per-element work follows w02's recipe (M137, PR15): bf16->f32 by an integer shift so the
+// conversions leave the XU pipe, one bf16 rounding for the whole normalize*sqrt(C)*gain product
+// instead of the reference's three, SiLU as u + u*tanh(u) with u = pre-activation/2 and
+// tanh.approx.f32, and 0.5*sqrt(C)*rsqrt(sum) as one per-position scalar.
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+using B=__nv_bfloat16;
+static cudaStream_t e_stream(){return at::cuda::getCurrentCUDAStream();}
+static const B* e_ptr(const Tensor& t){return t.defined()?(const B*)t.data_ptr<at::BFloat16>():nullptr;}
+static B* e_mut(Tensor& t){return t.defined()?(B*)t.data_ptr<at::BFloat16>():nullptr;}
+
+struct EpiPtrs{const B* p[36];};
+
+// A of F(4x4,3x3) applied along one axis: out[i] = sum_v A[i][v] v[v], two fp32 lanes per
+// instruction (a6v below is the only form this file uses). add/sub/fma.f32x2 need sm_100+; this
+// tree is built for sm_100a only (solution.json), so no scalar fallback is kept beside it -- a
+// scalar copy would be dead source here, which is why the scalar `a6` was removed rather than left.
+// f32x2: PTX 8.6 gives add/sub/mul/fma on a .b64 operand holding two f32 lanes (sm_100a).
+// This kernel already carries two independent fp32 lanes per thread -- the channel pair a thread
+// owns -- so `ch` IS the SIMD lane and nothing crosses a lane. Each instruction performs exactly
+// the arithmetic the scalar instruction performs, in the same association.
+struct V2{
+  float x,y;
+  __device__ V2(){} __device__ V2(float a,float b):x(a),y(b){}
+};
+__device__ __forceinline__ unsigned long long pk(V2 a){
+  return (unsigned long long)__float_as_uint(a.y)<<32|__float_as_uint(a.x);
+}
+__device__ __forceinline__ V2 upk(unsigned long long u){
+  V2 r; r.x=__uint_as_float((unsigned)u); r.y=__uint_as_float((unsigned)(u>>32)); return r;
+}
+__device__ __forceinline__ V2 vadd(V2 a,V2 b){
+  unsigned long long r; asm("add.rn.f32x2 %0,%1,%2;":"=l"(r):"l"(pk(a)),"l"(pk(b))); return upk(r);
+}
+__device__ __forceinline__ V2 vsub(V2 a,V2 b){
+  unsigned long long r; asm("sub.rn.f32x2 %0,%1,%2;":"=l"(r):"l"(pk(a)),"l"(pk(b))); return upk(r);
+}
+__device__ __forceinline__ V2 vfmas(V2 a,float c,V2 acc){
+  V2 k; k.x=k.y=c; unsigned long long r;
+  asm("fma.rn.f32x2 %0,%1,%2,%3;":"=l"(r):"l"(pk(a)),"l"(pk(k)),"l"(pk(acc))); return upk(r);
+}
+__device__ __forceinline__ V2 b2v(uint32_t w){
+  V2 r; r.x=__uint_as_float(w<<16); r.y=__uint_as_float(w&0xffff0000u); return r;
+}
+// The transform of A above, two lanes per instruction.
+__device__ __forceinline__ void a6v(V2 v0,V2 v1,V2 v2,V2 v3,V2 v4,V2 v5,V2* y){
+  V2 a=vadd(v1,v2), b=vsub(v1,v2), c=vadd(v3,v4), d=vsub(v3,v4);
+  y[0]=vadd(vadd(v0,a),c); y[1]=vfmas(d,2.f,b); y[2]=vfmas(c,4.f,a); y[3]=vadd(vfmas(d,8.f,b),v5);
+}
+// Exact bf16 pair -> two floats without the XU pipe: a bf16 is the top half of an fp32.
+__device__ __forceinline__ float2 b2f(uint32_t w){
+  return make_float2(__uint_as_float(w<<16),__uint_as_float(w&0xffff0000u));
+}
+__device__ __forceinline__ float tanh_ap(float x){float r;asm("tanh.approx.f32 %0, %1;":"=f"(r):"f"(x));return r;}
+// silu(x) with u = x/2: x*sigmoid(x) = u + u*tanh(u).
+__device__ __forceinline__ float silu2(float u){return __fmaf_rn(u,tanh_ap(u),u);}
+__device__ __forceinline__ uint32_t add2(uint32_t a,uint32_t b){
+  __nv_bfloat162 x=*(const __nv_bfloat162*)&a,y=*(const __nv_bfloat162*)&b;
+  __nv_bfloat162 r=__hadd2(x,y); return *(const uint32_t*)&r;
+}
+__device__ __forceinline__ uint32_t pack2(float x,float y){
+  __nv_bfloat162 r=__floats2bfloat162_rn(x,y); return *(const uint32_t*)&r;
+}
+__device__ __forceinline__ float sq8(const uint4& v){
+  float2 a=b2f(v.x),b=b2f(v.y),c=b2f(v.z),d=b2f(v.w);
+  return (__fmaf_rn(a.x,a.x,a.y*a.y)+__fmaf_rn(b.x,b.x,b.y*b.y))
+        +(__fmaf_rn(c.x,c.x,c.y*c.y)+__fmaf_rn(d.x,d.x,d.y*d.y));
+}
+
+// Stage the 36 components of one 4x4 tile into shared memory: 36*C/8 tasks of 16 bytes,
+// 9 per thread at any C. Each warp's 32 lanes take 32 consecutive 16-byte slots, i.e. 512
+// contiguous bytes of one plane.
+//
+// The copy is `cp.async.cg` (16 B, global->shared, L2-cached, NO register round trip) rather than
+// the `uint4` load-then-store it replaces.  Measured reason (J3871, wino_epi_kernel<1024,1>, ncu,
+// 3 passes): the binding FIXABLE stall is `mio_throttle 2.11` cycles per issued instruction --
+// second only to `not_selected 2.73`, which is not a defect (it means the SM already has eligible
+// warps to pass over) -- and `l1tex__throughput` is the highest pipe at 62.7-64.7 %.  The old form
+// is a global load AND a shared store per 16 B, so it spends two MIO slots and a writeback
+// register per slot; this form spends one slot and no register.  Per block at C=1024 that is
+// 9*512 = 4,608 pairs, i.e. 9,216 MIO instructions becoming 4,608.
+//
+// Bytes moved, alignment and the resulting V are IDENTICAL: same comp/off decomposition, same
+// 16-byte slot, same plane offset inside smem, so the transform reads what it read before and the
+// output is bit-for-bit the old output.  This is a scheduling change, not a numerical one.
+//
+// The commit/wait pair sits here rather than at the call site because the transform consumes the
+// staging immediately -- there is no independent work to overlap the transfer with inside one
+// tile.  Overlap across tiles needs a persistent loop and is a separate change, not this one.
+// The form above is measured, not preferred: `cp.async.cg` beats the `uint4` load+store it
+// replaced by 12.7 ms at f121x1 (J4745, one job, one clock, A/A 0.2 ms).
+template<int C>
+__device__ __forceinline__ void stage36(B* sm,const EpiPtrs& mp,unsigned src_tile){
+  constexpr int CV=C/2;
+  constexpr int PER=C/8;                 // 16-byte slots per plane
+  const unsigned smbase=(unsigned)__cvta_generic_to_shared(sm);
+  #pragma unroll
+  for(int i=0;i<9;i++){
+    const int task=threadIdx.x+i*CV;
+    const int comp=task/PER, off=task%PER;
+    const unsigned dst=smbase+(unsigned)((comp*C+off*8)*2);
+    const char* src=(const char*)mp.p[comp]+src_tile+off*16;
+    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"::"r"(dst),"l"(src));
+  }
+  asm volatile("cp.async.commit_group;\n"::);
+  asm volatile("cp.async.wait_group 0;\n"::);
+}
+
+template<int C,bool EXACT>
+__global__ __launch_bounds__(C/2) void wino_epi_kernel(
+    EpiPtrs mp,B* __restrict__ y,B* __restrict__ n,
+    const B* __restrict__ bias,const B* __restrict__ residual,const B* __restrict__ gain,
+    int H,int W,int TH,int TW,int PT,int PH,bool wy){
+  // `wy` is the CALLER's assertion that nobody reads this launch's y. It is only ever false where
+  // the very next conv has idx == fnext and therefore takes the fused `prepared` edge (see
+  // binding.cpp: resnet's first conv, whose consumer is its own second conv). The n tensor is
+  // written either way, so this can never be false without a n to carry the activation.
+  constexpr int CV=C/2;          // channel pairs per position
+  constexpr int WP=CV/32;        // warps per block
+  const int c=threadIdx.x;
+  const int tx=blockIdx.x, ty=blockIdx.y, t=blockIdx.z;
+  extern __shared__ __align__(16) B smem[];
+  B* sm=smem;                                              // 36 planes x C, this tile
+  // 32-bit plane byte offset: D*TH*TW*C*2 bytes, where D is the PIXEL frame count at this level
+  // (4*chunk-3 at up block 1, not the latent chunk size -- the earlier "34" was the latent count).
+  // chunk_size is NOW min(16*world, 61) -- this PR is what moved the rule off `8*world`, so chunk
+  // 32 at world 4 no longer exists and the worst case moves with it. Over the shipping configs the
+  // largest src_tile is world 2 / chunk 32 = 4.39e8 bytes and world 1 / chunk 16 = 4.29e8, both
+  // ~10x inside UINT32_MAX and both enforced at run time by the guard further down this file.
+  // The `8.8e8` that stood here was 4x too large and it described a configuration this PR retires.
+  // A chunk size that more than doubles again, or a wider clip, must widen this to unsigned long long.
+  const unsigned src_tile=(unsigned)((((t*TH+ty)*TW+tx)*C)*2u);
+
+  stage36<C>(sm,mp,src_tile);
+  __syncthreads();
+
+  V2 acc[4][4];
+  #pragma unroll
+  for(int i=0;i<4;i++)
+    #pragma unroll
+    for(int b=0;b<4;b++) acc[i][b]=V2(0.f,0.f);
+
+  const B* q0=sm+2*c;                     // this thread's channel pair, in each plane
+  #pragma unroll
+  for(int u=0;u<6;u++){
+    V2 f[6];
+    #pragma unroll
+    for(int v=0;v<6;v++) f[v]=b2v(*(const uint32_t*)(q0+(6*u+v)*C));
+    V2 q[4];
+    a6v(f[0],f[1],f[2],f[3],f[4],f[5],q);
+    #pragma unroll
+    for(int b=0;b<4;b++){
+      const V2 v=q[b];
+      if(u==0) acc[0][b]=vadd(acc[0][b],v);
+      else if(u==1){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vadd(acc[1][b],v);acc[2][b]=vadd(acc[2][b],v);acc[3][b]=vadd(acc[3][b],v);}
+      else if(u==2){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vsub(acc[1][b],v);acc[2][b]=vadd(acc[2][b],v);acc[3][b]=vsub(acc[3][b],v);}
+      else if(u==3){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vfmas(v,2.f,acc[1][b]);acc[2][b]=vfmas(v,4.f,acc[2][b]);acc[3][b]=vfmas(v,8.f,acc[3][b]);}
+      else if(u==4){acc[0][b]=vadd(acc[0][b],v);acc[1][b]=vfmas(v,-2.f,acc[1][b]);acc[2][b]=vfmas(v,4.f,acc[2][b]);acc[3][b]=vfmas(v,-8.f,acc[3][b]);}
+      else {acc[3][b]=vadd(acc[3][b],v);}
+    }
+  }
+
+  const int Hp=H+2*PH;
+  float2 biasv=make_float2(0.f,0.f),gainv=make_float2(0.f,0.f);
+  if(bias) biasv=b2f(*(const uint32_t*)&((const __nv_bfloat162*)bias)[c]);
+  if(gain) gainv=b2f(*(const uint32_t*)&((const __nv_bfloat162*)gain)[c]);
+  const long yb=(((long)t*H+4*ty)*W+4*tx)*CV+c;
+
+  const __nv_bfloat162 bv=__floats2bfloat162_rn(biasv.x,biasv.y);
+  if(!n){
+  // No prepared tensor: write the raw conv output in the natural per-thread layout.
+  #pragma unroll
+  for(int i=0;i<4;i++){
+    const int yy=4*ty+i;
+    #pragma unroll
+    for(int b=0;b<4;b++){
+      const int xx=4*tx+b;
+      const long off=(long)(i*W+b)*CV;
+      const bool ok=EXACT?true:(yy<H&&xx<W);
+      __nv_bfloat162 v=__floats2bfloat162_rn(acc[i][b].x,acc[i][b].y);
+      if(bias) v=__hadd2(v,bv);
+      if(residual&&ok) v=__hadd2(v,((const __nv_bfloat162*)residual)[yb+off]);
+      if(ok) ((__nv_bfloat162*)y)[yb+off]=v;
+    }
+  }
+  }
+
+  if(n){
+    // Remapped epilogue (w01's M655, measured 1.58x on the standalone inverse): the thread that
+    // computes the transform owns a channel pair, so it would write 4 bytes per output, which
+    // caps the pass at the 4-byte row of the stream table. Instead the 16 outputs go to shared
+    // memory with the bias already applied, one barrier, and the epilogue is re-laid so a warp
+    // owns one position and each lane an octet of channels: the residual read and the y write are
+    // one 512-byte transaction per warp instruction, and the sum of squares becomes a single-warp
+    // butterfly because a warp now holds every channel of its position.
+    __syncthreads();                       // the components are dead; their buffer is reused
+    B* so=smem;                            // so[p*C + ch], bias applied
+    #pragma unroll
+    for(int i=0;i<4;i++)
+      #pragma unroll
+      for(int b=0;b<4;b++){
+        __nv_bfloat162 v=__floats2bfloat162_rn(acc[i][b].x,acc[i][b].y);
+        if(bias) v=__hadd2(v,bv);
+        ((__nv_bfloat162*)so)[(i*4+b)*CV+c]=v;
+      }
+    __syncthreads();
+    constexpr int NPW=16/WP;               // positions per warp (4, 2 or 1)
+    constexpr int CW=C/32;                 // channels per lane (8, 16 or 32)
+    constexpr int NU=CW/8;                 // 16-byte chunks per lane (1, 2 or 4)
+    const int wp=c>>5, ln=c&31;
+    const float hs=0.5f*sqrtf((float)C);
+    uint4 gv[NU];
+    // Chunk map: lane ln takes chunk (k*32+ln), i.e. channel (k*32+ln)*8, so the 32 lanes of one
+    // warp instruction cover 32 consecutive chunks = 512 contiguous bytes. The previous map gave
+    // lane ln the block [ln*CW, ln*CW+CW), which puts neighbours CW*2 bytes apart: contiguous at
+    // C=256 (16 B), half a line at C=512 (32 B), a quarter at C=1024 (64 B) -- and at C=1024 that
+    // is a 4x sector over-fetch on the residual read and on both stores. w21's J1096 gives it in
+    // counters (L1 hit 5.32%, L2 hit 0.20%, DRAM 65.6%): no reuse, so every over-fetched sector
+    // is paid at DRAM. Measured at D=8, min of three rounds: C=256 1.000x (identical expression at
+    // NU=1), C=512 1.244x, C=1024 1.431x. It permutes which channel each k reads, so the 4-term
+    // sum of squares accumulates in a different order and s rounds differently in its last bits --
+    // f121x1 max|d| 0.1605 -> 0.147, rms unchanged to five digits.
+    #pragma unroll
+    for(int k=0;k<NU;k++)
+      gv[k]=gain?*(const uint4*)((const char*)gain+((long)(k*32+ln)*8)*2)
+               :make_uint4(0x3f803f80u,0x3f803f80u,0x3f803f80u,0x3f803f80u);
+    #pragma unroll
+    for(int j=0;j<NPW;j++){
+      const int p=wp*NPW+j, yy=4*ty+(p>>2), xx=4*tx+(p&3);
+      const bool ok=EXACT?true:(yy<H&&xx<W);
+      uint4 yv4[NU]; float sq=0.f;
+      #pragma unroll
+      for(int k=0;k<NU;k++){
+        const int off=(k*32+ln)*8;
+        uint4 v=*(const uint4*)((const char*)so+((long)p*C+off)*2);
+        if(residual&&ok){
+          // (long)t first: the whole offset must be 64-bit, not the last factor. `H` here is this
+          // RANK'S BAND height, not the full 352 -- chunk() narrows dim 3 of the latent grid (44
+          // rows) to [44r/W,44(r+1)/W) before resnet15 and every up block, so block 3 carries 88
+          // rows at world 4 and 44 at world 8. The band shrinks as the temporal chunk grows, which
+          // is why a rank's working set is near-constant -- BUT THE CENSUS THAT STOOD HERE WAS
+          // TRUE ONLY OF `chunk_size = 8*world`, AND THIS PR MOVES THE RULE TO `min(16*world, 61)`.
+          // Post-change THREE cells pass INT_MAX and THE LARGEST IS WORLD 2, so it is listed
+          // first: narrowing against world 4 because "chunk 61 looks biggest" is exactly the
+          // mislabelled-maximum this paragraph exists to prevent.
+          //   world 2 / chunk 32   D=125, rows 176, 624x256   = 3.51e9   <-- THE MAXIMUM
+          //   world 1 / chunk 16   D=61,  rows 352            = 3.43e9
+          //   world 4 / chunk 61   D=241, rows 88             = 3.39e9
+          // World 2's chunk 32 is THIS PR's OWN value at world 2 (min(16*2, 61) = 32), not a
+          // foreign or historical form. ALL THREE ARE PAST INT_MAX. This site is
+          // correct because it chains `(long)` here and again on `nbase` below, so the 64-bit chain
+          // is now LOAD-BEARING at those three cells rather than merely belt-and-braces. Do not
+          // narrow it back on the strength of an `inside INT_MAX` reading from the old rule.
+          const uint4 r=*(const uint4*)((const char*)residual
+                        +((long)((((long)t*H+yy)*W+xx)*C+off))*2);
+          v.x=add2(v.x,r.x); v.y=add2(v.y,r.y); v.z=add2(v.z,r.z); v.w=add2(v.w,r.w);
+        }
+        yv4[k]=v; sq+=sq8(v);
+      }
+      #pragma unroll
+      for(int d=16;d;d>>=1) sq+=__shfl_xor_sync(0xffffffffu,sq,d);
+      if(ok){
+        const float s=hs*rsqrtf(fmaxf(sq,1e-24f));
+        const long ybase=((((long)t*H+yy)*W+xx)*C)*2;
+        const long nbase=(((((long)(t+PT)*Hp+PH+yy)*W+xx)*C))*2;
+        #pragma unroll
+        for(int k=0;k<NU;k++){
+          const long cof=(long)((k*32+ln)*8)*2;
+          if(wy) *(uint4*)((char*)y+ybase+cof)=yv4[k];
+          uint4 u=yv4[k], g=gv[k], o;
+          float2 f0=b2f(u.x), f1=b2f(u.y), f2=b2f(u.z), f3=b2f(u.w);
+          float2 g0=b2f(g.x), g1=b2f(g.y), g2=b2f(g.z), g3=b2f(g.w);
+          o.x=pack2(silu2(f0.x*s*g0.x),silu2(f0.y*s*g0.y));
+          o.y=pack2(silu2(f1.x*s*g1.x),silu2(f1.y*s*g1.y));
+          o.z=pack2(silu2(f2.x*s*g2.x),silu2(f2.y*s*g2.y));
+          o.w=pack2(silu2(f3.x*s*g3.x),silu2(f3.y*s*g3.y));
+          *(uint4*)((char*)n+nbase+cof)=o;
+        }
+      }
+    }
+  }
+}
+
+// Zero the spatial padding rows of every frame of the prepared tensor.
+__global__ void epi_border(B* n,int C,int T,int Hp,int W){
+  const int CV=C/2; const long per=(long)2*W*CV;
+  const long total=(long)T*per;
+  for(long i=(long)blockIdx.x*blockDim.x+threadIdx.x;i<total;i+=(long)gridDim.x*blockDim.x){
+    const int t=(int)(i/per); const long r=i%per; const int edge=(int)(r/((long)W*CV));
+    const long r2=r%((long)W*CV); const int w=(int)(r2/CV),cc=(int)(r2%CV);
+    const long idx=(((long)t*Hp+(edge?Hp-1:0))*W+w)*CV+cc;
+    ((__nv_bfloat162*)n)[idx]=__floats2bfloat162_rn(0.f,0.f);
+  }
+}
+__global__ void epi_zero(B* n,long pairs){
+  const __nv_bfloat162 z=__floats2bfloat162_rn(0.f,0.f);
+  for(long i=(long)blockIdx.x*blockDim.x+threadIdx.x;i<pairs;i+=(long)gridDim.x*blockDim.x)
+    ((__nv_bfloat162*)n)[i]=z;
+}
+
+template<int C,bool EXACT>
+static void launch_epi(const dim3& grid,const dim3& block,size_t smem,
+                       const EpiPtrs& p,B* y,B* n,
+                       const B* b,const B* r,const B* g,
+                       int H,int W,int TH,int TW,int PT,int PH,bool wy,cudaStream_t s){
+  auto k=wino_epi_kernel<C,EXACT>;
+  if(smem>48*1024){
+    static bool done[2]={false,false};
+    if(!done[EXACT?1:0]){
+      C10_CUDA_CHECK(cudaFuncSetAttribute(k,cudaFuncAttributeMaxDynamicSharedMemorySize,36*C*2));
+      done[EXACT?1:0]=true;
+    }
+  }
+  k<<<grid,block,smem,s>>>(p,y,n,b,r,g,H,W,TH,TW,PT,PH,wy);
+}
+
+WinoOut wino_epilogue(const std::vector<Tensor>& m,int H,int W,const Tensor& bias,
+                      const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph,
+                      bool write_y){
+  // `write_y=false` is admissible ONLY when the caller's next conv has idx == the `fnext` this call
+  // prepares for, so that next conv takes the `prepared` edge and never reads the returned y. Out.y
+  // is still a distinct allocation: its unread contents are undefined, not zero.
+  TORCH_CHECK(m.size()==36&&m[0].dim()==5,"wino_epilogue: M list ",m.size()," rank ",m[0].dim());
+  const int C=m[0].size(1),D=m[0].size(2),TH=m[0].size(3),TW=m[0].size(4);
+  TORCH_CHECK(C==256||C==512||C==1024,"wino_epilogue: unsupported channel count");
+  auto s=e_stream();
+  EpiPtrs p; for(int i=0;i<36;i++) p.p[i]=(const B*)m[i].data_ptr<at::BFloat16>();
+  WinoOut out;
+  out.y=nhwc_empty(C,D,H,W,m[0]);
+  const int Hp=H+2*ph;
+  if(gamma.defined()){
+    out.n=nhwc_empty(C,D+pt,Hp,W,m[0]);
+    B* np=e_mut(out.n);
+    const long pairs=(long)(D+pt)*Hp*W*(C/2);
+    if(pt>0&&!hist.defined())
+      epi_zero<<<(int)std::min((long)2048,((long)pt*Hp*W*(C/2)+4095)/4096),256,0,s>>>(np,(long)pt*Hp*W*(C/2));
+    if(ph>0)
+      epi_border<<<(int)std::min((long)2048,(pairs+4095)/4096),256,0,s>>>(np,C,D+pt,Hp,W);
+    if(pt>0&&hist.defined()){
+      // `hist` is a `last_frames` output: a COMPACT [C,pt,H,W] channels-last-3d tensor, so its
+      // per-frame SOURCE pitch is H*W*C -- unpadded. Only the DESTINATION is padded: out.n is
+      // [C,D+pt,Hp,W] with pitch Hp*W*C, and these frames land in its interior rows [ph,ph+H).
+      // The two pitches are one of each and neither is a typo: passing Hp*W*C as the source pitch
+      // would read pt frames of garbage, and passing H*W*C as the destination pitch would write
+      // the pad rows away. If `last_frames` changes its shape or its pitch, this copy changes.
+      const size_t row=(size_t)H*W*C*sizeof(B);
+      C10_CUDA_CHECK(cudaMemcpy2DAsync(np+(long)ph*W*C,(size_t)Hp*W*C*sizeof(B),
+                                       e_ptr(hist),row,row,pt,cudaMemcpyDeviceToDevice,s));
+    }
+  }
+  // src_tile in the kernel is the only 32-bit index in this file, so assert here what it relies on
+  // rather than trusting the kernel's comment. (The y and n offsets in the kernel body are formed in
+  // 64-bit and are not bounded by this.) Checked against UINT32_MAX, not INT32_MAX: the index is
+  // unsigned, and a check against the signed bound would refuse shapes that are in fact safe.
+  { const unsigned long long off=(((unsigned long long)(D-1)*TH+(TH-1))*TW+(TW-1))*(unsigned long long)C*2ull;
+    TORCH_CHECK(off<=0xFFFFFFFFull,"wino_epilogue: 32-bit tile offset ",off,
+                " exceeds UINT32_MAX; widen src_tile in wino_epilogue.cu"); }
+  const bool exact=(TH*4==H&&TW*4==W);
+  const dim3 grid(TW,TH,D);
+  const dim3 block(C/2);
+  const size_t shmem=(size_t)36*C*2;
+  B* yp=e_mut(out.y); B* np=e_mut(out.n);
+  const B* bp=e_ptr(bias); const B* rp=e_ptr(residual); const B* gp=e_ptr(gamma);
+  if(C==256){ if(exact) launch_epi<256,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s);
+              else launch_epi<256,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); }
+  else if(C==512){ if(exact) launch_epi<512,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s);
+                   else launch_epi<512,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); }
+  else { if(exact) launch_epi<1024,true>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s);
+         else launch_epi<1024,false>(grid,block,shmem,p,yp,np,bp,rp,gp,H,W,TH,TW,pt,ph,write_y,s); }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return out;
+}
@@ -0,0 +1,247 @@
+// F(4x4,3x3) spatial input transform of the Winograd path of the Wan2.2 VAE decoder.
+//
+// ONE PASS: `winograd_input_fast` replaces `kernels.cu`'s `winograd_input` for the 18 Winograd
+// convolutions (1024/512/256 channels; any other channel count keeps the general kernel), and
+// its output V is BIT-IDENTICAL to it -- same patch origin, same block-uniform fast-path
+// condition, same B^T-then-B order with the in-place column update, same 6a+b plane index, same
+// single bf16 rounding at the store, same plane stride.  Nothing else in the tree moves but the
+// one call site in binding.cpp and one declaration in kernels.h.
+//
+// Why it is faster, from the profile of the kernel it replaces (w01's J549): 581.8 M
+// warp-instructions in 1.031 ms at 59.5% issue-active, of which 43.5% are ALU/integer against
+// 47.6% FMA and 4.9% LSU; DRAM 23.6%, FMA pipe 32.4%.  It sits at 82% of its own instruction
+// issue bound and the integer address math is the only large removable block:
+//   1. a 3-D grid, blockIdx = (tx, ty, t) and threadIdx = the channel pair, so nothing divides;
+//   2. one block-uniform fast path for the interior tiles, where the 36 loads carry no
+//      predicate at all -- only the first and last tile of a row and the first and last tile
+//      row can touch the time pad or the frame edge -- and the patch row pointer advances by
+//      the row stride, so the 36 addresses are 36 LSU instructions with folded offsets;
+//   3. the 36 plane stores walk one pointer by V's plane stride instead of doing an int64
+//      multiply per store.
+//
+// MORE BLOCKS PER TILE INSTEAD OF ONE: SPLIT = 4 AT C = 1024, 2 AT C = 512 AND 256.
+// A tile's C/2 channel pairs used to be covered by one block of C/2 threads, so `blockDim` and
+// `C` were the same variable and the grid held exactly one block per tile.  A block is now the
+// tile's width divided by SPLIT, and SPLIT blocks of the grid cover one tile.  At C = 1024 that
+// moves the level from one block of 16 warps (32.08 ms) to five blocks of 4 warps (22.42 ms,
+// 20 warps resident); at C = 512 and C = 256 the same move is worth 1.05 % and -0.4 % -- i.e.
+// nothing -- so those two stay at SPLIT = 2.  See the table below for why.
+//
+// `__launch_bounds__(CP, 512 / CP)` IS LOAD-BEARING AND IS HALF THE MEASUREMENT.  It pins the
+// register budget at 65536/512 = 128 at EVERY width -- exactly what main's
+// `__launch_bounds__(512, 1)` grants at its one wide instantiation.  Without it ptxas may spend
+// up to the architectural cap of 255 registers per thread (65536/CP is itself 512 at CP = 128,
+// so it is the 255 that binds there), floor(65536/(CP * REG)) falls, the resident warps follow,
+// and the same arm runs 2.25x SLOWER.  The counterfactual is measured at ONE instantiation:
+// J3705 ran `wino_in_kernel_st<256,2>` both ways -- `__launch_bounds__(256, 1)` at 59.11 ms
+// against `(256, 2)` at 26.26 ms, one integer and one body.  At the shipped CP = 128 the no-pin
+// case gives floor(65536/(128*255)) = 2 blocks of 4 warps against the pinned 5, i.e. 20 resident
+// warps down to 8 -- a 2.5x cut, not a halving, and the reason C = 1024 needs the pin most.
+// The declaration and the launch are one change; neither is the arm without the other.
+//
+// Measured (J3868, one build, one process, A B B B B A over SPLIT = 1 2 4 8 2 1, f121x1
+// --kernels, 1132 MHz on every leg, exit 0).  Both control legs are shown so the drift is visible:
+//                     SPLIT=1      SPLIT=2          SPLIT=4        SPLIT=8      2 -> 4
+//   C = 1024 (shipped)  32.08/32.08  26.26/26.53      22.42          22.27       1.1775
+//   C =  512 (2 kept)   64.06/63.92  54.60/55.44      54.45          55.45       1.0105
+//   C =  256 (2 kept)   60.47/60.66  59.56/60.05      60.04        (n/a)         0.9962
+//   CALL               864.7        847.6            844.2          845.2
+//   A/A 0.3 %, B/B 1.0-1.5 %; max |d| 0.1543 of 0.4984, rms 0.002093 of 0.005847, PSNR 59.6 dB,
+//   IDENTICAL to every printed digit on all six legs.
+//   The C = 256 SPLIT = 8 cell is `n/a`, not a decline: J3868's probe only instantiated the CP
+//   values it needed and SPLIT = 8 at C = 256 would ask for CP = 16, which that build did not
+//   contain.  The head's launcher is an exhaustive switch with no fall-back path of any kind.
+//
+// ONLY THE 2 -> 4 MOVE IS A C = 1024-ONLY GAIN, and that is why the other two keep SPLIT = 2:
+//   the shipped SPLIT = 2 already takes the 1 -> 2 step on BOTH of them -- 64.06 -> 54.60 at
+//   C = 512 is 1.1733x and 60.47 -> 59.56 at C = 256 is 1.0183x, the +15.5 % and +1.5 % below --
+//   so neither is a rung that gains nothing.  What is declined on them is the FOURTH block:
+//   C = 512's 1.0105 is inside the B/B drift (its own legs read 54.60 and 55.44, a 1.5 % spread);
+//   C = 256's 0.9962 REVERSES -- its legs read 59.56, 60.05 and 60.66, a 1.8 % spread, and 60.04
+//   sits inside it, so the option is declined rather than merely unpriced;
+//   SPLIT = 8 at C = 1024 is 22.27 against 22.42 -- 0.7 % at a 1.0 % drift, a tie, not a win.
+//
+// WHY C = 1024 AND NOT THE OTHERS.  With REG 96 the realized resident warps are
+// floor(65536/(96*blockDim)) * blockDim/32, i.e. 16 at blockDim 512 or 256, 20 at 128 or 64, and
+// 21 at 32.  Main's C = 1024 case used blockDim 512, which rounds to ONE block and never reaches
+// past 16 warps although 2048/96 = 21.3 is its ceiling; `__launch_bounds__` cannot fix that,
+// because 2 * 512 * 96 = 98,304 > 65,536.  Only a narrower block can, and narrowing it is SPLIT.
+// C = 512 and C = 256 already reach 20 warps at SPLIT = 2, so there is nothing there to buy.
+// The three levels' rates were 2.609 / 3.340 / 3.760 TB/s (J3195's counted bytes).  DERIVED, not
+// re-measured: at fixed byte counts rate_after/rate_before = t_before/t_after, so from J3868's own
+// legs the after-rates are 2.609*32.08/22.42 = 3.733, 3.340*64.06/55.44 = 3.859 and
+// 3.760*60.47/59.56 = 3.817 TB/s -- a 44.1 % spread closed to 3.4 %.  The largest move is the
+// slow level's own (+43 %), and the fast two barely move (+15.5 % and +1.5 %), so what happened is
+// that the slow level was walked up to where the other two already sat, not that a ceiling moved.
+// An earlier draft of this comment stated 3.952 for C = 256, which needs 60.47 -> 57.5 ms and no
+// leg of J3868 is below 59.56; it was a number carried over from a withdrawn row and it is wrong.
+//
+// `J3761`'s rebuild of the round-1 source: BUILD_RC=0, REG:96 STACK:0 SHARED:0 LOCAL:0.
+// `J3889`'s build of this source: the same, with `wino_in_kernel<128,4>` at REG:95 (one register
+// lower, which changes no block count: 65536/(128*95) = 5.39 -> 5, the same as at 96).
+// Correctness at worlds 1 and 4 on all three short clips: J3889 f1x1/f5x1/f17x1 ok at
+// x0.41/x0.41/x0.65 and J3908 f1x4/f5x4/f17x4 ok at x0.45/x0.45/x0.66, every leg inside main's
+// own margins.
+//
+// The share of the call is flat in clip length (17.6 % at f121x1, 17.9 % at f481x1) and this
+// function takes no rank, no world and no process group, so a world-1 cell is the geomean cell.
+//
+// Measured by w04 on main's own tree (kernel#1, this change alone, f121x1 at 1849 MHz): the
+// three instantiations total 114.4 ms against 221.8 ms of `winograd_input_kernel`, and the call
+// 969.9 -> 862.5 ms (-11.1%), with V bit-identical -- f1/f5/f17 keep the seed's max |d| and rms
+// to the digit -- and f481x1 ok at x0.38 of the max and x0.41 of the rms allowance.  w21 wrote
+// the kernel; w08 owns this file.
+#include "kernels.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+#include <algorithm>
+
+using B = __nv_bfloat16;
+static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
+
+__device__ __forceinline__ void wbt6(float v0, float v1, float v2, float v3, float v4, float v5, float* y) {
+  float a = v4 - 4 * v2, b = v3 - 4 * v1, c = v4 - v2, d = 2 * (v3 - v1);
+  y[0] = 4 * v0 - 5 * v2 + v4; y[1] = a + b; y[2] = a - b; y[3] = c + d; y[4] = c - d; y[5] = 4 * v1 - 5 * v3 + v5;
+}
+
+// One block per (part-tile, frame); threadIdx.x = the channel pair within this block's share.
+// CP = C/2/SPLIT is the block size.  The shipped dispatch is SPLIT = 4 at C = 1024 and SPLIT = 2
+// at C = 512 and 256, so blockDim is 128 / 128 / 64 at C = 1024 / 512 / 256 -- but SPLIT is NOT
+// uniform and no comment here may claim it is.  The invariant that actually guarantees coverage is
+//     CP * SPLIT == C / 2                                    (128*4 = 128*2 = 64*2 = C/2)
+// enforced by the `static_assert` in `winograd_input_fast` below, together with `grid.x == SPLIT*TW`
+// from the same `constexpr`.  `(blockIdx.x % SPLIT) * CP + threadIdx.x` then walks [0, SPLIT*CP) =
+// [0, C/2) exactly once -- every channel pair of a tile is written once and no pair twice -- and
+// `blockIdx.x / SPLIT` over a grid of SPLIT*TW hits each tile exactly SPLIT times, for any SPLIT:
+// the remainder cycles 0..SPLIT-1 within each group of SPLIT consecutive blockIdx, so TW need not be
+// a multiple of SPLIT.
+//
+// SPLIT = 1 with CP = C/2 gives back main's one-block-per-tile BODY, but not its attribute (this
+// one is `__launch_bounds__(CP, 512/CP)`, which equals main's `(CP, 1)` only at CP = 512) and not
+// its grid shape, so do not use it as an exact revert; revert by checking out the parent commit.
+//
+// The launched blockDim must equal CP.  It does at all three call sites below, from the same
+// template value, so the two cannot drift apart.
+template <int CP, int SPLIT>
+__global__ void __launch_bounds__(CP, 512 / CP)
+wino_in_kernel(const B* __restrict__ x, B* __restrict__ v, int C, int D, int H, int W, int PH, int TH, int TW) {
+  const int cp = (blockIdx.x % SPLIT) * CP + threadIdx.x;
+  const int tx = blockIdx.x / SPLIT;
+  const int ty = blockIdx.y, t = blockIdx.z;
+  const int HP = H + 2 * PH;
+  const int y0 = 4 * ty - 1 + PH, x0 = 4 * tx - 1;
+  const int64_t rowStep = (int64_t)W * C;
+  const B* base = x + (int64_t)t * HP * rowStep + (int64_t)y0 * rowStep + (int64_t)x0 * C + 2 * cp;
+  float d[2][6][6];
+  // Uniform across the block, so it is one branch and not 36 predicates.
+  if ((y0 >= 0) && (y0 + 5 < HP) && (x0 >= 0) && (x0 + 5 < W)) {
+    const B* prow = base;
+#pragma unroll
+    for (int a = 0; a < 6; a++) {
+#pragma unroll
+      for (int b = 0; b < 6; b++) {
+        float2 f = __bfloat1622float2(*(const __nv_bfloat162*)(prow + b * C));
+        d[0][a][b] = f.x; d[1][a][b] = f.y;
+      }
+      prow += rowStep;
+    }
+  } else {
+    const B* prow = base;
+#pragma unroll
+    for (int a = 0; a < 6; a++) {
+      const int y = y0 + a;
+      const bool ok = (y >= 0) && (y < HP);
+#pragma unroll
+      for (int b = 0; b < 6; b++) {
+        const int xx = x0 + b;
+        if (ok && xx >= 0 && xx < W) {
+          float2 f = __bfloat1622float2(*(const __nv_bfloat162*)(prow + b * C));
+          d[0][a][b] = f.x; d[1][a][b] = f.y;
+        } else {
+          d[0][a][b] = 0.f; d[1][a][b] = 0.f;
+        }
+      }
+      prow += rowStep;
+    }
+  }
+  // B^T d down the patch's 6 columns, then B along its 6 rows, one bf16 rounding at the end.
+#pragma unroll
+  for (int k = 0; k < 2; k++)
+#pragma unroll
+    for (int b = 0; b < 6; b++) {
+      float q[6];
+      wbt6(d[k][0][b], d[k][1][b], d[k][2][b], d[k][3][b], d[k][4][b], d[k][5][b], q);
+#pragma unroll
+      for (int a = 0; a < 6; a++) d[k][a][b] = q[a];
+    }
+  B* dst = v + (int64_t)(((t * TH + ty) * TW + tx) * C + 2 * cp);
+  const int64_t plane = (int64_t)D * TH * TW * C;
+#pragma unroll
+  for (int a = 0; a < 6; a++) {
+    float q[2][6];
+#pragma unroll
+    for (int k = 0; k < 2; k++)
+      wbt6(d[k][a][0], d[k][a][1], d[k][a][2], d[k][a][3], d[k][a][4], d[k][a][5], q[k]);
+#pragma unroll
+    for (int b = 0; b < 6; b++, dst += plane) *(__nv_bfloat162*)dst = __floats2bfloat162_rn(q[0][b], q[1][b]);
+  }
+}
+
+Tensor winograd_input_fast(const Tensor& x, int ph) {
+  const int C = x.size(1), D = x.size(2), H = x.size(3) - 2 * ph, W = x.size(4);
+  const int TH = (H + 3) / 4, TW = (W + 3) / 4;
+  // The 18 Winograd convolutions are the 1024-, 512- and 256-channel ones; anything else
+  // keeps the general kernel.
+  if (C != 1024 && C != 512 && C != 256) return winograd_input(x, ph);
+  auto v = at::empty({36, D, TH, TW, C}, x.options());
+  // SPLIT = 4 at C = 1024, 2 at C = 512 and 256.  CP * SPLIT = C/2 at every case, so a tile's
+  // channel pairs are covered exactly once; grid.x is SPLIT * TW, a block count and not a tile
+  // count.  blockDim is CP, taken from the same template value, so the two cannot drift.
+  //
+  // Why 4 only at C = 1024 (J3868: 2 -> 4 on the rung at each C, one build, one process):
+  //   C=1024  <256,2> 26.26/26.53 -> <128,4> 22.42   1.1775   warps 16 -> 20
+  //   C= 512  <128,2> 54.60/55.44 -> < 64,4> 54.45   1.0105   warps 20 -> 20   inside the 1.5 % drift
+  //   C= 256  < 64,2> 59.56/60.05 -> < 32,4> 60.04   0.9962   warps 20 -> 21   it REVERSES
+  // So the resolved gain is C = 1024 alone, and it is the only rung whose split buys warps:
+  // with REG 96 the realized resident warps are floor(65536/(96*blockDim))*blockDim/32, i.e.
+  // 16 at blockDim 512 or 256, 20 at 128 or 64, 21 at 32.  The step off 16 is the whole of it.
+  // SPLIT = 8 at C = 1024 measured 22.27 against 22.42 -- 0.7 % at a 1.0 % B/B drift, not
+  // resolvable, so it is not taken.  A tie is not a win.
+  //
+  // ONE `constexpr` PER CASE BINDS THE TEMPLATE ARGUMENT, THE BLOCK WIDTH, THE GRID MULTIPLIER AND
+  // THE COVERAGE INVARIANT.  Three literals in three places is a desync waiting to become a bug:
+  // raise SPLIT without raising grid.x and `tx = blockIdx.x / SPLIT` stays inside TW, the kernel
+  // silently computes one block's share of every tile and the rest of V is never written; raise
+  // grid.x without SPLIT and `tx` runs past TW, computing x0 = 4*tx-1 for tiles that do not exist,
+  // which is an out-of-bounds read AND write rather than a slowdown.  The `static_assert` makes the
+  // first a compile error and the single `constexpr` makes the second unrepresentable.
+  const B* xp = (const B*)x.data_ptr<at::BFloat16>();
+  B* vp = (B*)v.data_ptr<at::BFloat16>();
+  auto st = stream();
+  switch (C) {
+    case 1024: {
+      constexpr int SPLIT = 4, CP = 1024 / 2 / SPLIT;
+      static_assert(CP * SPLIT == 1024 / 2, "CP * SPLIT must equal C/2 or a tile is not covered");
+      const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D);
+      wino_in_kernel<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
+      break;
+    }
+    case 512: {
+      constexpr int SPLIT = 2, CP = 512 / 2 / SPLIT;
+      static_assert(CP * SPLIT == 512 / 2, "CP * SPLIT must equal C/2 or a tile is not covered");
+      const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D);
+      wino_in_kernel<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
+      break;
+    }
+    case 256: {
+      constexpr int SPLIT = 2, CP = 256 / 2 / SPLIT;
+      static_assert(CP * SPLIT == 256 / 2, "CP * SPLIT must equal C/2 or a tile is not covered");
+      const dim3 g((unsigned)(SPLIT * TW), (unsigned)TH, (unsigned)D);
+      wino_in_kernel<CP, SPLIT><<<g, CP, 0, st>>>(xp, vp, C, D, H, W, ph, TH, TW);
+      break;
+    }
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return v;
+}
@@ -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()
@@ -0,0 +1,771 @@
+"""Correctness, timing and library-kernel audit of a solution: the Wan2.2 VAE decode of latents into video, on one, two
+or four GB200 GPUs of one node. The task's CI runs this same file (see ci.py).
+
+The task (spec.py has the details)
+--------
+A solution decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T-3, 704, 1248] (bf16, in [-1, 1]) with
+the Wan2.2 VAE's decoder, as diffusers' AutoencoderKLWan decodes them (reference.py in the kernel repo). It is a
+directory (default solution/, see loader.py) compiled into one extension module that exports
+
+    run(z, weights, rank, world, pg, out)                    positional, destination passing
+    prepare(weights, rank, world, pg) -> list[Tensor]       optional
+
+called in one process per GPU: z the full latents on every rank; weights the decoder's 110 bf16 tensors on the rank's
+GPU, in the order of spec.WEIGHTS (or what prepare returned); pg torch.distributed's NCCL ProcessGroup over the ranks;
+out the video on rank 0, from torch.empty and never zeroed (every element must be written), and an empty tensor on the
+other ranks. When run has returned on every rank, rank 0's out holds the video. prepare, if exported, runs once per
+process before any timed call and is not timed; its result replaces the weights in every run.
+
+The workloads (spec.WORKLOADS) are clips decoded at a world size, named <clip>x<world>. A clip is the first T latent
+frames of one of three real 20 s clips, named by its video frames: f1, f5, f17 (T = 1, 2, 5) and f121, f241, f481
+(T = 31, 61, 121). The long clips at 1, 2 and 4 GPUs are the nine timed workloads (f121x1 ... f481x4), whose geomean
+is the objective; the short clips at 1 and 4 GPUs (f1x1 ... f17x4) are checked and never timed.
+
+Running it
+----------
+In a job with as many GPUs as the largest world size it runs, with the task's data at $DATA (default /data), from the
+repo root:
+
+    python tools/bench.py                             every workload this job has the GPUs for: check, audit, time
+    python tools/bench.py --workload f5x1 f17x1       a subset; the short clips take seconds
+    python tools/bench.py --workload f121x1 --no-time correctness only
+    python tools/bench.py --kernels                   also list the audited calls' kernels by GPU time
+    python tools/bench.py --solution other            another solution directory
+
+A run
+-----
+The solution is built once, by this process, exactly as it is deployed (loader.py; a development copy, tools/, takes a
+build its node already made of the same sources, flags and toolchain from the node's cache, unless --no-cache). Then,
+per world size in turn (1, 2, 4), torch.multiprocessing spawns one process per rank, rank r on GPU r of the job's
+CUDA_VISIBLE_DEVICES. Each initializes NCCL over 127.0.0.1 with its device (the default group, the solution's pg) and
+a gloo group for the harness's own barriers, imports the build (in its working directory, as deployed), moves the
+weights ($DATA/decoder_weights_bf16.pt, checked against spec.WEIGHTS) to its GPU and calls prepare, if the module
+exports it. The ranks then run that world size's workloads in spec.WORKLOADS' order. A rank that raises, dies or
+hangs (WORKLOAD_SEC per workload, SESSION_START_SEC before the first) fails the workload it was at: all the ranks'
+processes are stopped, and the workloads after it run in new ones. This process runs no solution code.
+
+Every call gets z at a fresh address (a copy of the clip's latents) and out at a fresh address, 256-byte aligned and
+filled with NaN on rank 0, so that an element run leaves unwritten fails the check. run must not write z, nor the
+weights (compared after each workload with what prepare left), nor start threads: each fails the workload.
+
+Correctness
+-----------
+The first call of a workload, its warm-up, is checked on rank 0 against the clip's fp32 decode ($DATA/ref_<clip>.pt):
+every element finite and within [-1, 1]; and with d = out - ref in fp32, max |d| <= min(2, 3 x bf16_max_abs) and
+rms(d) <= 3 x bf16_rms, the error of the release's bf16 decode of that clip ($DATA/calibration.json; spec.py,
+Correctness). The line shows max |d| and rms(d), each with its allowance and the share of it used, and the PSNR,
+10 log10(4 / mse). Every later call's output (the audited call's, the trials') must agree with the checked one:
+lose none of the elements it wrote (none left NaN or overflowed), and be within the same allowance of it everywhere
+else. A race, or a state carried from call to call, fails the workload.
+
+Timing (the timed workloads)
+------
+After the warm-up come the trials: 3 when the warm-up took under 1.5 s, else 2. A trial is one call timed by wall
+clock between barriers on all ranks: synchronize the device, barrier, start, run, synchronize, barrier, stop. The
+time of record is rank 0's; the latency of record the mean of the trials, and the fastest trial is shown too. While
+the trials run, this process reads the SM clock of the world's GPUs through NVML every CLOCK_PERIOD_S (a rank inside
+run may hold its GIL for the whole call, so the ranks do not read it themselves); the mean of rank 0's GPU is shown
+with the latency, and the range over the world's GPUs where they differ by 2% or more: the clocks are not locked
+alike on every GPU, and this says at what clock the times were taken. Which processes hold the world's GPUs is
+asked (nvidia-smi) before every trial and after the last: a workload timed while another process held one of them is
+marked SHARED, since its times do not compare with those of another job.
+
+Audit
+-----
+At world 1, one more call per timed clip runs under CUPTI, and each kernel is attributed to the library that launched
+it (see audit.py). Library kernels (cuBLAS/cuBLASLt, cuDNN, ATen) are allowed as placeholders: their count per call
+and their share of the call's GPU time show where the solution still depends on them. NCCL's kernels are
+communication, reported apart and not counted as placeholders. --kernels lists the audited calls' kernels.
+
+Bound
+-----
+Each timed workload's latency is shown with its bound (bounds.py with spec.bounds: the decode's FLOPs spread evenly
+over the world's GPUs at the 16-bit tensor-core peak, or the video and the weights at the memory bandwidth, whichever
+is slower) and the gap, latency / bound. The peak assumes the SM clock GUILD_SM_CLOCK_MHZ where the environment sets
+it (the guild's: 2062 MHz, the GB200's highest), else the GPU's highest.
+
+In CI, and from tools/
+----------------------
+ci.py runs this file in child processes: --save DIR has rank 0 save each workload's checked output there (node-local
+scratch) instead of checking it, so that CI judges it in a process that runs no solution code; --no-check skips the
+check (the timing runs), --no-audit the audit, and --json writes the report. The kernel repo holds copies of this
+file, spec.py, loader.py, audit.py and bounds.py in tools/, which run as they do here: they read the data at $DATA
+(default /data) and the solution from solution/ under the current directory.
+"""
+
+import argparse
+import collections
+import ctypes
+import functools
+import json
+import math
+import os
+import random
+import socket
+import statistics
+import subprocess
+import sys
+import threading
+import time
+import traceback
+from datetime import timedelta
+from multiprocessing import connection
+
+# The caching allocator stays as torch configures it by default: without expandable segments, whose memory
+# cudaIpcGetMemHandle cannot export, since a solution may share its buffers between ranks by CUDA IPC.
+import torch
+import torch.distributed as dist
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import audit  # noqa: E402
+import bounds  # noqa: E402
+import loader  # noqa: E402
+import spec  # noqa: E402
+
+# Trials per timed workload, fewer when one call takes long.
+TRIALS, LONG_TRIALS, LONG_CALL_S = 3, 2, 1.5
+# How often this process reads rank 0's SM clock while ranks run.
+CLOCK_PERIOD_S = 0.02
+NVML_CLOCK_SM = 1  # the nvmlClockType_t of the SM clock
+# Limits that stop a runaway: a session's start (NCCL, the module, the weights, prepare), and one workload (its calls,
+# its check or its save); a call of the release's bf16 decode takes 15 s at the longest clip.
+SESSION_START_SEC = 900
+WORKLOAD_SEC = 1800
+STOP_GRACE_SEC = 15
+# Elements of a video compared at once, in fp32 (1 GiB per operand).
+CHUNK_ELEMENTS = 1 << 28
+
+
+def gpu_holders():
+    """{GPU UUID (lower case, without "GPU-"): ids of the processes that hold it} now, as nvidia-smi gives them, or
+    None when it cannot say."""
+    try:
+        got = subprocess.run(["nvidia-smi", "--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader"],
+                             capture_output=True, text=True, timeout=30)
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    if got.returncode != 0:
+        return None
+    out = collections.defaultdict(set)
+    for line in got.stdout.splitlines():
+        gpu, _, pid = line.partition(",")
+        if pid.strip().isdigit():
+            out[gpu.strip().lower().removeprefix("gpu-")].add(int(pid))
+    return out
+
+
+def others(questions, uuids, ours):
+    """How many processes besides this run's held one of these GPUs at any of the questions (gpu_holders' answers), or
+    None when one could not be asked. This run's processes are known by their ids where nvidia-smi gives the ids this
+    process sees (one PID namespace, the usual case: then a rank that also holds another rank's GPU is still one of
+    ours); otherwise each GPU is taken to hold one of ours, the rank on it."""
+    if not questions or any(q is None for q in questions):
+        return None
+    total = 0
+    for uuid in uuids:
+        seen = set().union(*(q.get(uuid, set()) for q in questions))
+        total += len(seen - ours) if seen & ours else max(0, len(seen) - 1)
+    return total
+
+
+@functools.cache
+def sm_clock(uuid):
+    """A function that reads the SM clock of the GPU with this UUID now, in MHz (None when a read fails), or None where
+    NVML cannot be loaded or does not know the GPU. NVML in this process takes microseconds per read."""
+    try:
+        nvml = ctypes.CDLL("libnvidia-ml.so.1")
+        if nvml.nvmlInit_v2() != 0:
+            return None
+        handle, mhz = ctypes.c_void_p(), ctypes.c_uint()
+        if not uuid or nvml.nvmlDeviceGetHandleByUUID(f"GPU-{uuid}".encode(), ctypes.byref(handle)) != 0:
+            return None
+    except (OSError, AttributeError):
+        return None
+
+    def read():
+        ok = nvml.nvmlDeviceGetClockInfo(handle, NVML_CLOCK_SM, ctypes.byref(mhz)) == 0
+        return mhz.value if ok and mhz.value else None
+
+    return read if read() else None
+
+
+class Clock:
+    """The SM clocks of the world's GPUs (rank order), read every CLOCK_PERIOD_S by a thread of this process, as
+    (monotonic time, [MHz or None per GPU])."""
+
+    def __init__(self, uuids):
+        self.reads, self.samples, self.done = [sm_clock(u) for u in uuids], [], threading.Event()
+        self.thread = threading.Thread(target=self._run, daemon=True) if any(self.reads) else None
+        if self.thread:
+            self.thread.start()
+
+    def _run(self):
+        while not self.done.wait(CLOCK_PERIOD_S):
+            self.samples.append((time.monotonic(), [read() if read else None for read in self.reads]))
+
+    def stop(self):
+        self.done.set()
+        if self.thread:
+            self.thread.join()
+
+    def means(self, windows):
+        """Per GPU, the mean of its reads within these (start, stop) windows of time.monotonic(), or None without
+        any."""
+        inside = [mhz for t, mhz in list(self.samples) if any(a <= t <= b for a, b in windows)]
+        out = []
+        for i in range(len(self.reads)):
+            got = [m[i] for m in inside if m[i]]
+            out.append(statistics.mean(got) if got else None)
+        return out
+
+
+class Pool:
+    """A workload's buffers for one tensor: each call's at a fresh address, shifted by multiples of 256 bytes."""
+
+    ALIGN, MAX_SHIFT = 256, 2048
+
+    def __init__(self, shape, dtype, device, calls, seed):
+        rng = random.Random(seed)
+        self.offsets = [0]
+        for _ in range(calls - 1):
+            self.offsets.append(self.offsets[-1] + rng.randint(1, self.MAX_SHIFT // self.ALIGN))
+        self.shape, self.numel = tuple(shape), math.prod(shape)
+        self.block = max(1, self.ALIGN // torch.empty(0, dtype=dtype).element_size())
+        self.buffer = torch.empty(self.numel + self.offsets[-1] * self.block, dtype=dtype, device=device)
+        self.calls = 0
+
+    def next(self):
+        start = self.offsets[self.calls] * self.block
+        self.calls += 1
+        return self.buffer[start:start + self.numel].view(self.shape)
+
+
+def fingerprint(tensors):
+    """Each tensor's bytes summed as integers: whether run() wrote the weights."""
+    out = []
+    for t in tensors:
+        raw = t.detach().contiguous().view(-1).view(torch.uint8)
+        try:
+            raw = raw.view(torch.int32)
+        except RuntimeError:
+            pass
+        out.append(int(raw.sum(dtype=torch.int64)))
+    return out
+
+
+def judge(out, ref, allowed):
+    """A video against the reference, or a later call's video against the checked one's: {"passed", "nonfinite",
+    "outside" (elements beyond [-1, 1]), "max_abs", "rms", "psnr_db", "max_allowed", "rms_allowed"}, d = out - ref where
+    both are finite (rms and PSNR over those elements, None where there is none), in fp32 on out's GPU (the current one
+    for a video on the CPU), in chunks of frames, the squares summed in fp64."""
+    device = out.device if out.is_cuda else torch.device("cuda", torch.cuda.current_device())
+    step = max(1, CHUNK_ELEMENTS // max(1, out[:, :, :1].numel()))
+    nonfinite = outside = compared = unmatched = 0
+    max_abs = squares = 0.0
+    for f in range(0, out.shape[2], step):
+        x = out[:, :, f:f + step].to(device, torch.float32)
+        y = ref[:, :, f:f + step].to(device, torch.float32)
+        finite = torch.isfinite(x)
+        nonfinite += int((~finite).sum())
+        outside += int((x.abs() > 1).sum())
+        # Where the checked call itself left an element unfinite, a later call is not compared there (the check said).
+        both = finite & torch.isfinite(y)
+        compared += int(both.sum())
+        unmatched += int((~finite & torch.isfinite(y)).sum())
+        d = torch.where(both, x - y, 0.0)
+        max_abs = max(max_abs, float(d.abs().max()))
+        squares += float(d.square().sum(dtype=torch.float64))
+        del x, y, finite, both, d
+    mse = squares / compared if compared else None
+    rms = math.sqrt(mse) if compared else None
+    passed = (nonfinite == 0 and outside == 0 and compared == out.numel() and max_abs <= allowed["max_abs"]
+              and rms <= allowed["rms"])
+    return {"passed": passed, "nonfinite": nonfinite, "unmatched": unmatched, "outside": outside,
+            "max_abs": max_abs if compared else None, "rms": rms, "psnr_db": 10 * math.log10(4 / mse) if mse else None,
+            "max_allowed": allowed["max_abs"], "rms_allowed": allowed["rms"]}
+
+
+def verdict(c):
+    """A check (judge's result) in a few words."""
+    parts = []
+    if c["nonfinite"]:
+        parts.append(f"{c['nonfinite']} elements not finite (NaN: never written, or overflowed)")
+    if c["outside"]:
+        parts.append(f"{c['outside']} elements outside [-1, 1]")
+    if c["max_abs"] is None:
+        parts.append("no element to compare")
+    else:
+        parts.append(f"max |d| {c['max_abs']:.4g} of {c['max_allowed']:.4g} (x{c['max_abs'] / c['max_allowed']:.2f})")
+        parts.append(f"rms {c['rms']:.4g} of {c['rms_allowed']:.4g} (x{c['rms'] / c['rms_allowed']:.2f})")
+        parts.append(f"PSNR {c['psnr_db']:.1f} dB" if c["psnr_db"] is not None else "identical")
+    return f"{'ok  ' if c['passed'] else 'FAIL'} ({', '.join(parts)})"
+
+
+def run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg):
+    """One workload on one rank, in step with the other ranks (see the module docstring); rank 0's result."""
+    entry = spec.WORKLOADS[w]
+    clip, device = entry["clip"], torch.device("cuda", rank)
+    timed = entry["timed"] and cfg["time"]
+    audited = entry["timed"] and cfg["audit"] and world == 1
+    calls = 1 + audited + (TRIALS if timed else 0)
+    seed = list(spec.WORKLOADS).index(w)
+    z0 = spec.make_inputs(w, device)[0]
+    zs = Pool(z0.shape, z0.dtype, device, calls, seed)
+    outs = Pool(spec.output_shape(w), torch.bfloat16, device, calls, seed + 1) if rank == 0 else None
+    allowed = spec.allowance(clip)
+    res = {"clip": clip, "world": world, "timed": entry["timed"]}
+    problems = []
+    threads = threading.active_count()
+
+    def arguments():
+        z = zs.next()
+        z.copy_(z0)
+        out = outs.next().fill_(math.nan) if outs else torch.empty(0, dtype=torch.bfloat16, device=device)
+        return z, out
+
+    def unchanged(z):
+        if not torch.equal(z, z0) and "run() wrote its input z" not in problems:
+            problems.append("run() wrote its input z")
+
+    def call():
+        z, out = arguments()
+        torch.cuda.synchronize(device)
+        dist.barrier(group=harness)
+        start = time.monotonic()
+        sol.fn(z, weights, rank, world, pg, out)
+        torch.cuda.synchronize(device)
+        dist.barrier(group=harness)
+        stop = time.monotonic()
+        unchanged(z)
+        return out, start, stop
+
+    # The warm-up, which is the call checked.
+    out, start, stop = call()
+    warm = stop - start
+    first, worst = None, 0.0
+    if rank == 0:
+        if cfg["save"]:
+            path = os.path.join(cfg["save"], f"{w}.pt")
+            torch.save(out.cpu(), path)
+            res["saved"] = path
+        elif cfg["check"]:
+            res["check"] = judge(out, spec.load_reference(clip), allowed)
+        if calls > 1:
+            first = out.clone()
+
+    def agree(out, what):
+        nonlocal worst
+        # It agrees when it loses no element the checked call wrote, and is within the allowance of it elsewhere;
+        # the checked call's own faults are its check's to report.
+        c = judge(out, first, allowed)
+        worst = max(worst, c["max_abs"] or 0.0)
+        close = c["max_abs"] is None or (c["max_abs"] <= allowed["max_abs"] and c["rms"] <= allowed["rms"])
+        if c["unmatched"] or not close:
+            lost = f"{c['unmatched']} elements not finite where it was; " if c["unmatched"] else ""
+            problems.append(f"{what}'s output does not agree with the checked call's: {lost}"
+                            f"{verdict(c).removeprefix('FAIL').removeprefix('ok').strip()}")
+
+    # World 1 alone: one call under CUPTI, which attributes every kernel.
+    if audited:
+        z, out = arguments()
+        ops = audit.trace(sol.fn, z, weights, rank, world, pg, out)
+        unchanged(z)
+        agree(out, "the audited call")
+        res |= {"kernels": audit.kernels(ops), "launches": sum(1 for op in ops if op.kind == "kernel"),
+                "lib": sum(1 for op in ops if audit.placeholder(op.library)), "lib_share": audit.library_share(ops),
+                "comm": sum(1 for op in ops if op.library in audit.COMMUNICATION), "comm_share": audit.comm_share(ops)}
+        del ops
+
+    if timed:
+        n = [TRIALS if warm < LONG_CALL_S else LONG_TRIALS]
+        dist.broadcast_object_list(n, src=0, group=harness)
+        questions, trials, windows = [], [], []
+        for k in range(n[0]):
+            if rank == 0:
+                questions.append(gpu_holders())
+            out, start, stop = call()
+            trials.append(1e3 * (stop - start))
+            windows.append((start, stop))
+            if rank == 0:
+                agree(out, f"trial {k + 1}")
+        if rank == 0:
+            questions.append(gpu_holders())
+            res |= {"trials": trials, "ms": statistics.mean(trials), "min_ms": min(trials), "windows": windows}
+            held = others(questions, [uuid for _, uuid in ranks], {pid for pid, _ in ranks} | {os.getppid()})
+            if held:
+                res["shared"] = held
+    if first is not None:
+        res["agree_max_abs"] = worst
+    if fingerprint(weights) != marks:
+        problems.append("run() wrote its weights (what prepare returned is read-only to run: keep scratch elsewhere)")
+    if threading.active_count() > threads:
+        problems.append(f"the solution started threads ({threading.active_count() - threads} more)")
+
+    # What every rank found, so that rank 0 reports it.
+    found = [None] * world
+    dist.all_gather_object(found, problems, group=harness)
+    text = [p if world == 1 else f"rank {r}: {p}" for r, ps in enumerate(found) for p in ps]
+    if text and rank == 0:
+        res["error"] = "; ".join(text)
+    return res
+
+
+def worker(rank, world, port, cfg, conn):
+    """The process of one rank (see A run in the module docstring): it reports to the process that spawned it through
+    conn, ("ready", rank, uuid, prepare seconds), ("start", workload) and ("result", workload, result) on rank 0,
+    ("done", rank), or ("error", rank, workload or None, phase, traceback), after which it exits at once."""
+    phase, current = "start", None
+    try:
+        torch.cuda.set_device(rank)
+        device = torch.device("cuda", rank)
+        timeout = timedelta(seconds=WORKLOAD_SEC)
+        dist.init_process_group("nccl", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world,
+                                device_id=device, timeout=timeout)
+        harness = dist.new_group(backend="gloo", timeout=timeout)
+        pg = dist.group.WORLD
+        phase = "load"
+        sol = loader.import_built(cfg["workdir"], cfg["spec"])
+        phase = "weights"
+        weights = spec.load_weights(device)
+        phase = "prepare"
+        prepare_s = None
+        if sol.prepare is not None:
+            torch.cuda.synchronize(device)
+            start = time.monotonic()
+            got = sol.prepare(weights, rank, world, pg)
+            torch.cuda.synchronize(device)
+            prepare_s = time.monotonic() - start
+            if not isinstance(got, (list, tuple)) or not all(isinstance(t, torch.Tensor) for t in got):
+                raise TypeError(f"prepare returned {type(got).__name__}, not a list of tensors")
+            weights = list(got)
+        uuid = str(torch.cuda.get_device_properties(device).uuid).lower().removeprefix("gpu-")
+        ranks = [None] * world
+        dist.all_gather_object(ranks, (os.getpid(), uuid), group=harness)
+        conn.send(("ready", rank, uuid, prepare_s))
+        phase = "workload"
+        marks = fingerprint(weights)
+        for w in cfg["workloads"]:
+            current = w
+            if rank == 0:
+                conn.send(("start", w))
+            res = run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg)
+            if rank == 0:
+                conn.send(("result", w, res))
+            current = None
+            torch.cuda.empty_cache()
+        phase = "end"
+        dist.barrier(group=harness)
+        conn.send(("done", rank))
+        dist.destroy_process_group()
+    except BaseException:
+        try:
+            conn.send(("error", rank, current, phase, traceback.format_exc(limit=8)))
+        finally:
+            sys.stdout.flush()
+            sys.stderr.flush()
+            # At once: NCCL's teardown would wait for the other ranks, which may wait for this one.
+            os._exit(1)
+
+
+def free_port():
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def stop(procs, grace=0.0):
+    """Ends rank processes: those that do not exit within `grace` seconds get SIGTERM, then SIGKILL after
+    STOP_GRACE_SEC."""
+    deadline = time.monotonic() + grace
+    for p in procs:
+        p.join(max(0.0, deadline - time.monotonic()))
+    for p in procs:
+        if p.is_alive():
+            p.terminate()
+    deadline = time.monotonic() + STOP_GRACE_SEC
+    for p in procs:
+        p.join(max(0.1, deadline - time.monotonic()))
+        if p.is_alive():
+            p.kill()
+            p.join(10)
+
+
+def exit_text(code):
+    if code is None:
+        return "did not exit"
+    if code < 0:
+        import signal
+
+        try:
+            return f"was killed by {signal.Signals(-code).name}"
+        except ValueError:
+            return f"was killed by signal {-code}"
+    return f"exited with code {code}"
+
+
+def session(world, workloads, cfg, report, show):
+    """One set of rank processes running these workloads at one world size, in order; results go into report as they
+    arrive, and show(w) prints each. Returns the workloads it did not reach: after a failure at one workload (which gets
+    the error), those after it; [] when every one ran, or when the failure came before any (every one gets it)."""
+    ctx = torch.multiprocessing.get_context("spawn")
+    port = free_port()
+    pipes = [ctx.Pipe(duplex=False) for _ in range(world)]
+    procs = [ctx.Process(target=worker, args=(r, world, port, dict(cfg, workloads=workloads), pipes[r][1]),
+                         daemon=True) for r in range(world)]
+    for p in procs:
+        p.start()
+    for _, send in pipes:
+        send.close()
+    receivers = {pipes[r][0]: r for r in range(world)}
+    done, errors, finished, uuids = set(), [], [], {}
+    current, clock, failure = None, None, None
+    deadline = time.monotonic() + SESSION_START_SEC
+    started = time.monotonic()
+
+    def handle(msg):
+        nonlocal current, clock, deadline
+        kind = msg[0]
+        if kind == "ready":
+            _, rank, uuid, prepare_s = msg
+            uuids[rank] = uuid
+            if rank == 0:
+                report.setdefault("prepare_s", {})[str(world)] = prepare_s
+            if len(uuids) == world:
+                clock = Clock([uuids[r] for r in range(world)])
+                ready = f", prepare {report['prepare_s'][str(world)]:.1f} s" if prepare_s is not None else ""
+                print(f"world {world}: {world} rank{'s' if world > 1 else ''} up in {time.monotonic() - started:.0f} s"
+                      f"{ready}", flush=True)
+        elif kind == "start":
+            current = msg[1]
+            deadline = time.monotonic() + WORKLOAD_SEC
+        elif kind == "result":
+            _, w, res = msg
+            windows = res.pop("windows", None)
+            means = clock.means(windows) if clock and windows else []
+            if means and means[0]:
+                res["clock_mhz"] = means[0]
+            if world > 1 and means and all(means):
+                res["world_clock_mhz"] = [min(means), max(means)]
+            report["workloads"][w] = res
+            finished.append(w)
+            current = None
+            deadline = time.monotonic() + WORKLOAD_SEC
+            show(w)
+        elif kind == "error":
+            errors.append(msg[1:])
+        elif kind == "done":
+            done.add(msg[1])
+
+    try:
+        while len(done) < world:
+            # A rank that exited cleanly may still have its last messages in its pipe: it fails once that is read.
+            bad = [(r, p.exitcode) for r, p in enumerate(procs) if r not in done and p.exitcode is not None
+                   and (p.exitcode != 0 or pipes[r][0] not in receivers)]
+            if bad:
+                failure = "; ".join(f"rank {r} {exit_text(code)}" for r, code in bad)
+                break
+            left = deadline - time.monotonic()
+            if left <= 0:
+                failure = (f"timed out after {WORKLOAD_SEC} s" if current else
+                           f"the ranks did not start within {SESSION_START_SEC} s")
+                break
+            waits = list(receivers) + [p.sentinel for p in procs if p.exitcode is None]
+            for obj in connection.wait(waits, timeout=min(left, 5.0)):
+                if obj in receivers:
+                    try:
+                        handle(obj.recv())
+                    except (EOFError, OSError):
+                        del receivers[obj]
+        # What the other ranks still say, such as the error that made one exit.
+        drain = time.monotonic() + 2.0
+        while failure and receivers and time.monotonic() < drain:
+            for obj in connection.wait(list(receivers), timeout=max(0.0, drain - time.monotonic())):
+                try:
+                    handle(obj.recv())
+                except (EOFError, OSError):
+                    del receivers[obj]
+    finally:
+        stop(procs, 0.0 if failure else STOP_GRACE_SEC)
+        if clock:
+            clock.stop()
+    if not failure:
+        return []
+
+    # The first error a rank reported says what failed; else how a rank ended.
+    if errors:
+        rank, at, phase, text = errors[0]
+        if phase == "load":
+            report["load_error"] = text.strip().splitlines()[-1]
+        text = (f"rank {rank}: " if world > 1 else "") + text.strip()
+    else:
+        text = failure
+    rest = [w for w in workloads if w not in finished]
+    at = current or (errors and errors[0][1])
+    if at not in rest:
+        at = rest[0] if rest and finished else None
+    if not at or report.get("load_error"):
+        for w in rest:
+            report["workloads"][w] = {"error": text, "world": world, "clip": spec.WORKLOADS[w]["clip"],
+                                      "timed": spec.WORKLOADS[w]["timed"]}
+            show(w)
+        return []
+    report["workloads"][at] = {"error": text, "world": world, "clip": spec.WORKLOADS[at]["clip"],
+                               "timed": spec.WORKLOADS[at]["timed"]}
+    show(at)
+    return rest[rest.index(at) + 1:]
+
+
+def kernel_lines(res, limit=25):
+    """The audited call's kernels by GPU time, one line each."""
+    rows = res.get("kernels", [])
+    total = sum(r[3] for r in rows) or 1
+    lines = [f"    {'[' + lib + ']' if lib else '':<9} {ns / 1e6:10.2f} ms {100 * ns / total:5.1f}%  x{n:<6} "
+             f"{name if name.startswith('(') else audit.short_name(name, 110)}" for lib, name, n, ns in rows[:limit]]
+    if len(rows) > limit:
+        rest = rows[limit:]
+        lines.append(f"    {'':<9} {sum(r[3] for r in rest) / 1e6:10.2f} ms "
+                     f"{100 * sum(r[3] for r in rest) / total:5.1f}%  {len(rest)} more kinds")
+    return lines
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--solution", default="solution", help="the solution directory (default: solution)")
+    p.add_argument("--workload", nargs="*", help="workloads to run (default: every one this job has the GPUs for)")
+    p.add_argument("--no-time", action="store_true", help="correctness only: no audit, no trials")
+    p.add_argument("--no-check", action="store_true", help="do not check the outputs (CI's timing runs)")
+    p.add_argument("--no-audit", action="store_true", help="no call under CUPTI")
+    p.add_argument("--kernels", action="store_true", help="list the audited calls' kernels by GPU time")
+    p.add_argument("--build-dir", help="write and build the solution here; reused when it holds the same build")
+    p.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds")
+    p.add_argument("--save", help="CI: save each workload's checked output here instead of checking it")
+    p.add_argument("--json", help="write the results to this file")
+    a = p.parse_args()
+    unknown = [w for w in a.workload or [] if w not in spec.WORKLOADS]
+    if unknown:
+        p.error(f"unknown workloads {unknown}; known: {list(spec.WORKLOADS)}")
+    # Absolute paths, the data's too: the ranks run in the solution's working directory.
+    solution = os.path.abspath(a.solution)
+    if os.path.isfile(solution):
+        solution = os.path.dirname(solution)
+    a.build_dir, a.save, a.json = (x and os.path.abspath(x) for x in (a.build_dir, a.save, a.json))
+    os.environ["DATA"] = os.path.abspath(spec.data_dir())
+    report = {"workloads": {}}
+
+    def finish():
+        if a.json:
+            with open(a.json, "w") as f:
+                json.dump(report, f)
+
+    rules = audit.rule_table(getattr(spec, "RULES", None))
+    report["rules"] = audit.check(solution, rules)
+    for v in report["rules"]:
+        print(f"rule violated: {v}")
+    print(f"stack rules: {'ok' if not report['rules'] else 'FAIL'} ({len(rules)} rules)")
+
+    # The workloads: those named, or every one this job has the GPUs for.
+    gpus = torch.cuda.device_count()
+    chosen = a.workload or [w for w, e in spec.WORKLOADS.items() if e["world"] <= gpus]
+    skipped = [w for w in spec.WORKLOADS if w not in chosen and not a.workload]
+    if skipped:
+        print(f"skipped {', '.join(skipped)}: this job has {gpus} GPU{'s' if gpus != 1 else ''}")
+    order = [w for w in spec.WORKLOADS if w in chosen]
+    timed = [w for w in order if spec.WORKLOADS[w]["timed"] and not a.no_time]
+    limits, dev, peak = bounds.workload_bounds(spec, timed)
+    if limits:
+        print(bounds.describe(dev, peak))
+
+    try:
+        workdir, build_spec, seconds = loader.stage(solution, a.build_dir,
+                                                    cache=None if a.no_cache else loader.cache_dir())
+    except Exception as e:
+        report["load_error"] = str(e) if isinstance(e, loader.SolutionError) else traceback.format_exc(limit=6)
+        print(f"FAILED to build {solution}:\n{report['load_error']}", flush=True)
+        finish()
+        sys.exit(1)
+    report |= {"language": build_spec["languages"][0], "build_s": seconds}
+    built = f"built in {seconds:.1f} s" if seconds is not None else "a build reused"
+    print(f"{build_spec['entry_point']} ({', '.join(build_spec['languages'])}): {built}", flush=True)
+
+    width = max(len(w) for w in spec.WORKLOADS)
+
+    def show(w):
+        res = report["workloads"][w]
+        line = f"{w:<{width}}"
+        if "ms" in res:
+            at = f", SM clock {res['clock_mhz']:.0f} MHz" if "clock_mhz" in res else ""
+            low, high = res.get("world_clock_mhz", (None, None))
+            if low and "clock_mhz" in res and high - low >= 0.02 * res["clock_mhz"]:
+                at += f", the world's GPUs {low:.0f}-{high:.0f} MHz"
+            line += f" {res['ms']:9.1f} ms (fastest trial {res['min_ms']:.1f}{at})"
+            if w in limits:
+                res |= {"bound_ms": limits[w], "gap": res["ms"] / limits[w]}
+                line += f"  bound {limits[w]:.1f} ms, gap x{res['gap']:.2f}"
+        if "check" in res:
+            line += "  " + verdict(res["check"])
+        elif "saved" in res:
+            line += "  saved"
+        if not res.get("timed", True):
+            line += "  (correctness only)"
+        if "lib" in res:
+            line += f"  library kernels {res['lib']} of {res['launches']} ({100 * res['lib_share']:.0f}% of GPU time)"
+        if res.get("comm"):
+            line += f", NCCL {res['comm']} ({100 * res['comm_share']:.0f}%)"
+        if "shared" in res:
+            line += f"  SHARED: {res['shared']} other process(es) held a GPU while it was timed"
+        if "error" in res:
+            line += "  ERROR\n" + res["error"]
+        print(line, flush=True)
+        if a.kernels and "kernels" in res:
+            print("\n".join(kernel_lines(res)), flush=True)
+        finish()
+
+    cfg = {"workdir": workdir, "spec": build_spec, "time": not a.no_time, "check": not a.no_check,
+           "audit": not a.no_audit and (not a.no_time or a.kernels), "save": a.save}
+    for world in sorted({spec.WORKLOADS[w]["world"] for w in order}):
+        todo = [w for w in order if spec.WORKLOADS[w]["world"] == world]
+        if world > gpus or report.get("load_error"):
+            why = report.get("load_error") or f"needs {world} GPUs; this job has {gpus}"
+            for w in todo:
+                report["workloads"][w] = {"error": why, "world": world, "clip": spec.WORKLOADS[w]["clip"],
+                                          "timed": spec.WORKLOADS[w]["timed"]}
+                show(w)
+            continue
+        while todo:
+            todo = session(world, todo, cfg, report, show)
+
+    results = [report["workloads"][w] for w in order if w in report["workloads"]]
+    times = [report["workloads"].get(w, {}).get("ms") for w in timed]
+    if timed and all(times):
+        gaps = [report["workloads"][w]["gap"] for w in timed if "gap" in report["workloads"][w]]
+        gap = f", gap to the bound x{bounds.geomean(gaps):.2f}" if len(gaps) == len(timed) else ""
+        clocks = [report["workloads"][w].get("clock_mhz") for w in timed]
+        at = f", mean SM clock {statistics.mean(clocks):.0f} MHz" if all(clocks) else ""
+        print(f"geomean {bounds.geomean(times):.1f} ms over {len(timed)} timed workloads (mean of trials){gap}{at}")
+    shared = [w for w in order if "shared" in report["workloads"].get(w, {})]
+    if shared:
+        print(f"SHARED GPU: another process held one while {', '.join(shared)} "
+              f"{'was' if len(shared) == 1 else 'were'} timed, so those times do not compare with another job's")
+    checked = [r for r in results if "check" in r]
+    if checked:
+        ok = sum(1 for r in checked if r["check"]["passed"] and "error" not in r)
+        print(f"correctness: ok at {ok} of {len(checked)} workloads checked")
+    failed = [w for w in order if "error" in report["workloads"].get(w, {})]
+    if failed:
+        print(f"errors at {', '.join(failed)}")
+    audited = [r for r in results if "kernels" in r]
+    if audited:
+        libs = collections.Counter()
+        for r in audited:
+            for lib, _, n, _ in r["kernels"]:
+                if audit.placeholder(lib):
+                    libs[lib] += n
+        detail = ", ".join(f"{k} {v}" for k, v in libs.most_common())
+        share = statistics.mean(r["lib_share"] for r in audited)
+        print(f"library kernels: {sum(libs.values())} launches in {len(audited)} audited calls"
+              f"{' (' + detail + ')' if detail else ''}, {100 * share:.0f}% of GPU time on average")
+    finish()
+
+
+if __name__ == "__main__":
+    main()
@@ -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()
@@ -0,0 +1,542 @@
+"""Builds and loads a solution exactly as it is deployed, so that what runs here runs there.
+
+A solution is a directory holding exactly the files that are deployed, plus `solution.json`, the deployment's build
+spec, with its field names and meaning:
+
+    {"languages": ["cuda_cpp"], "entry_point": "binding.cpp::run", "destination_passing_style": true,
+     "compile_options": {"cflags": [...], "cuda_cflags": [...], "ld_flags": [...]}, "dependencies": []}
+
+- `languages`: always ["cuda_cpp"].
+- `entry_point`: "<file>::<function>", the function the deployment calls, in a C++/CUDA file.
+- `destination_passing_style`: required, and true for this task: run() writes its output into `out`, its last
+  argument, and returns nothing. The calling convention must never be left to inference.
+- `compile_options` (optional): extra flags; see the build below. `dependencies` (optional): a list of names, which
+  does not change the build. `binding` (optional): null or "torch".
+Any other field is refused: the deployment would ignore it, so a misspelled field would silently fall back to its
+default. Every other file in the directory is a source: non-empty UTF-8 text at a relative path, a plain file (no
+links; `__pycache__` and `.pyc` files are ignored).
+
+The module exports the entry, called positionally as the deployment calls it (argument names play no part):
+
+    run(z, weights, rank, world, pg, out)                    every call; writes out on rank 0
+    prepare(weights, rank, world, pg) -> list[Tensor]       optional, under that name: once per process, before run
+
+(spec.py says what each argument is.) Load one solution per process.
+
+stage(solution_dir) writes the sources into a fresh working directory and compiles every .cpp/.cc/.cxx/.c/.cu file at
+its top into one extension module, in a separate process, with torch.utils.cpp_extension.load, build_directory = the
+working directory, include paths [working directory, $CUTLASS_DIR/include, $CUTLASS_DIR/tools/util/include] ($CUTLASS_DIR
+defaults to fastvideo-kernel/include/cutlass, the CUTLASS submodule), and these
+flags on top of torch's own:
+  - no compile_options: cuda_cflags [], cflags [], ld_flags ["-lcuda"];
+  - compile_options given: its lists, where a missing one defaults to cflags [], cuda_cflags ["-O3",
+    "--use_fast_math"], ld_flags ["-lcuda"];
+  - unless a cuda_cflag names an architecture (-gencode, -arch), a -gencode flag for the GPU's own architecture with
+    its arch-specific features goes first (sm_100a on compute capability 10.0), as the deployment does for its GPU.
+Files in subdirectories are not compiled, but can be included. import_built(working directory) then makes it the
+current directory, as the deployment does, imports the module and takes run (and prepare) from it; load() does both.
+The module's name is the build's, not the solution's: PYBIND11_MODULE must be given TORCH_EXTENSION_NAME, as the
+deployment builds the module under a name of its own. bench.py stages a build once and its ranks import that build.
+
+Development builds are kept on each node, so that a job does not compile again what an earlier job on its node built:
+a content-addressed cache in a directory of the node that setup/node.sh makes for the agent user (CACHE_DIR, or
+$GUILD_BUILD_CACHE), keyed by everything that determines the build: the sources, solution.json, the flags, the target
+architecture, this loader, and the toolchain (the machine, Python, torch with its CUDA, nvcc, the host compiler, and
+the variables they read). A build found there is copied into the working directory in place of compiling it; one made
+is put there with its sources and module, without what the compiler left beside them. Only the commands of a
+development copy of the harness use it (tools/ in a repo: `loader.py check` and `build`, `bench.py`), unless they are
+given --no-cache. The copy in the task's checks/, which CI, the hooks and the task's commands run, never reads or
+writes it and builds fresh: the jobs of agents write that directory, so what it holds is theirs to forge. Past a
+quarter of its filesystem, a limit that only stops a runaway, the least recently used builds are removed.
+
+    python tools/loader.py check [solution] [--build-dir DIR]   build and load it, print what was loaded
+    python tools/loader.py build [solution] [DIR]               only write and build it in DIR (default build/): with
+                                                                --arch (e.g. sm_100a), on a machine without a GPU,
+                                                                such as a CPU builder, where the module is not loaded
+"""
+
+import argparse
+import atexit
+import functools
+import hashlib
+import importlib.util
+import json
+import os
+import platform
+import re
+import shutil
+import stat
+import subprocess
+import sys
+import tempfile
+import time
+
+SPEC_FILE = "solution.json"
+# The name the build gives the module; solutions never name it themselves (see the module docstring). It is this
+# kernel's own: another kernel's loader may build an extension of a generic name (solution_ext), and in one process the
+# second extension of a name loads as the first one.
+MODULE = "wan22_vae_decode_ext"
+# The optional function a solution exports beside its entry, called once per process before any call of the entry.
+PREPARE = "prepare"
+BINDINGS = ("torch",)
+CPP_ENTRY_SUFFIXES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh")
+# The fields solution.json may hold.
+SPEC_FIELDS = ("languages", "entry_point", "dependencies", "destination_passing_style", "binding", "compile_options")
+OPTION_FIELDS = ("cflags", "cuda_cflags", "ld_flags")
+# The languages a solution may declare, and the files compiled.
+LANGUAGES = ("cuda_cpp",)
+CPP_SOURCES = (".cu", ".cpp", ".cc", ".cxx", ".c")
+DEFAULT_OPTIONS = {"cflags": [], "cuda_cflags": ["-O3", "--use_fast_math"], "ld_flags": ["-lcuda"]}
+BUILD_TIMEOUT_SEC = 1800
+# The first argument of a PYBIND11_MODULE in a source.
+MODULE_NAME = re.compile(r"\bPYBIND11_MODULE\s*\(\s*([A-Za-z_]\w*)")
+HERE = os.path.dirname(os.path.abspath(__file__))
+# The node's cache of development builds (see the module docstring), and what each build there holds beside its sources.
+CACHE_DIR = "/scratch/kdev-builds"
+STAMP, SIZE = ".stamp", ".size"
+# The share of its filesystem past which the cache removes its least recently used builds: a limit that stops a runaway.
+CACHE_SHARE = 0.25
+# What a job that stopped while it stored a build left in the cache is removed after this long.
+STALE_SEC = 3600
+# The variables the compilers read, which change what they make.
+BUILD_VARIABLES = ("CC", "CXX", "CUDA_HOME", "CUDA_PATH", "CUTLASS_DIR", "NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS",
+                   "CPATH", "C_INCLUDE_PATH", "CPLUS_INCLUDE_PATH", "LIBRARY_PATH")
+
+
+class SolutionError(Exception):
+    """The solution is malformed, or fails to build or load."""
+
+
+def _nonempty_str(x):
+    return isinstance(x, str) and len(x) >= 1
+
+
+def _str_list(x):
+    return isinstance(x, list) and all(isinstance(s, str) for s in x)
+
+
+def spec_problems(spec):
+    """Problems of a build spec, as a list of strings (empty: valid); see the module docstring.
+
+    Stricter than the deployment only where it would coerce a value (a "true" string as a boolean) or fall back to a
+    default (an unknown field, a missing destination_passing_style)."""
+    if not isinstance(spec, dict):
+        return [f"{SPEC_FILE} must hold an object"]
+    out = []
+    unknown = sorted(set(spec) - set(SPEC_FIELDS))
+    if unknown:
+        out.append(f"unknown fields {unknown} (known: {list(SPEC_FIELDS)})")
+    langs = spec.get("languages")
+    if not isinstance(langs, list) or not langs or any(lang not in LANGUAGES for lang in langs):
+        out.append(f"languages: required, {list(LANGUAGES)}")
+    entry = spec.get("entry_point")
+    if not _nonempty_str(entry):
+        out.append("entry_point: required, a non-empty string")
+    elif entry.count("::") != 1:
+        out.append(f'entry_point: {entry!r} is not "<file_path>::<function_name>"')
+    elif os.path.splitext(entry.split("::")[0])[1] not in CPP_ENTRY_SUFFIXES:
+        out.append(f"entry_point: needs a C++/CUDA entry file, not {entry.split('::')[0]!r}")
+    deps = spec.get("dependencies", [])
+    if not isinstance(deps, list) or not all(_nonempty_str(d) for d in deps):
+        out.append("dependencies: a list of non-empty strings")
+    if "destination_passing_style" not in spec:
+        out.append("destination_passing_style: must be given, and true here (run writes out, its last argument)")
+    elif spec["destination_passing_style"] is not True:
+        out.append("destination_passing_style: true: this task's run() writes the video into out, its last argument, "
+                   "and returns nothing")
+    if spec.get("binding") not in (None,) + BINDINGS:
+        out.append(f"binding: null or one of {list(BINDINGS)}")
+    options = spec.get("compile_options")
+    if options is not None:
+        if not isinstance(options, dict):
+            out.append("compile_options: null or an object")
+        else:
+            for k in OPTION_FIELDS:
+                if k in options and not _str_list(options[k]):
+                    out.append(f"compile_options.{k}: a list of strings")
+    return out
+
+
+def read_spec(solution_dir):
+    """solution.json, checked against the build spec's rules (see the module docstring)."""
+    path = os.path.join(solution_dir, SPEC_FILE)
+    try:
+        with open(path) as f:
+            spec = json.load(f)
+    except (OSError, ValueError) as e:
+        raise SolutionError(f"cannot read {path}: {e}") from e
+    problems = spec_problems(spec)
+    if problems:
+        raise SolutionError(f"{SPEC_FILE}: " + "; ".join(problems))
+    return spec
+
+
+def sources(solution_dir):
+    """{relative path: text} of every file of the solution except solution.json."""
+    out = {}
+    for root, dirs, files in os.walk(solution_dir):
+        dirs[:] = sorted(d for d in dirs if d != "__pycache__")
+        for name in sorted(files):
+            path = os.path.join(root, name)
+            rel = os.path.relpath(path, solution_dir).replace(os.sep, "/")
+            if rel == SPEC_FILE or name.endswith(".pyc"):
+                continue
+            if os.path.islink(path):
+                raise SolutionError(f"{rel}: sources must be plain files, not links")
+            try:
+                with open(path, encoding="utf-8") as f:
+                    text = f.read()
+            except (OSError, UnicodeDecodeError) as e:
+                raise SolutionError(f"{rel}: sources must be UTF-8 text ({e})") from e
+            if not text:
+                raise SolutionError(f"{rel}: sources must not be empty")
+            out[rel] = text
+    if not out:
+        raise SolutionError(f"{solution_dir} holds no sources")
+    named = sorted({m for text in out.values() for m in MODULE_NAME.findall(text) if m != "TORCH_EXTENSION_NAME"})
+    if named:
+        raise SolutionError(f"PYBIND11_MODULE names its module {', '.join(named)}: give it TORCH_EXTENSION_NAME, the "
+                            f"name the build chooses (the deployment builds the module under a name of its own)")
+    return out
+
+
+def read(solution_dir):
+    """(spec, {relative path: text}) of a solution, checked as stage() checks it before building: solution.json against
+    the build spec's rules, the sources as the deployment takes them, the entry file among them."""
+    spec = read_spec(solution_dir)
+    files = sources(solution_dir)
+    entry = spec["entry_point"].split("::")[0]
+    if entry not in files:
+        raise SolutionError(f"entry file {entry} is not among the sources")
+    return spec, files
+
+
+def gpu_arch():
+    """The architecture solutions are built for, e.g. sm_100a: the GPU's own, with its arch-specific features. It asks
+    the GPU's properties only, which makes no CUDA context."""
+    import torch
+
+    if not torch.cuda.is_available():
+        raise SolutionError("there is no GPU here to take the architecture from: name it (--arch sm_100a, or the "
+                            "architecture of the GPUs the build is for)")
+    major, minor = torch.cuda.get_device_capability()
+    return f"sm_{major}{minor}" + ("a" if major >= 9 else "")
+
+
+def compile_flags(spec, arch):
+    """(cflags, cuda_cflags, ld_flags) added to torch's own flags (see the module docstring)."""
+    given = spec.get("compile_options")
+    options = {"cflags": [], "cuda_cflags": [], "ld_flags": ["-lcuda"]} if given is None else {
+        k: list(given.get(k, v)) for k, v in DEFAULT_OPTIONS.items()}
+    if not any("-gencode" in f or "-arch" in f for f in options["cuda_cflags"]):
+        options["cuda_cflags"].insert(0, f"-gencode=arch=compute_{arch[3:]},code={arch}")
+    return options["cflags"], options["cuda_cflags"], options["ld_flags"]
+
+
+def cache_dir():
+    """The node's cache of development builds, or None where a build is made fresh: in the copy of the harness in the
+    task's checks/ (CI, hooks and the task's commands run it, which must not load what a job could have put there), and
+    on a node without a cache that this user may write."""
+    checks = os.environ.get("CHECKS")
+    if checks and os.path.realpath(checks) == os.path.realpath(HERE):
+        return None
+    cache = os.environ.get("GUILD_BUILD_CACHE", CACHE_DIR)
+    return cache if os.path.isdir(cache) and os.access(cache, os.W_OK | os.X_OK) else None
+
+
+@functools.lru_cache(maxsize=None)
+def toolchain():
+    """What determines a build beside the solution, its flags and the architecture: this loader, the machine, Python,
+    torch with its CUDA, nvcc, the host compiler, and the variables the compilers read."""
+    import torch
+    from torch.utils import cpp_extension
+
+    def version(cmd):
+        try:
+            p = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
+            return (p.stdout + p.stderr).strip()
+        except (OSError, subprocess.TimeoutExpired) as e:
+            return f"unavailable: {e}"
+
+    with open(os.path.abspath(__file__), "rb") as f:
+        loader = hashlib.sha256(f.read()).hexdigest()
+    return {"loader": loader, "machine": platform.machine(), "python": sys.version,
+            "torch": [torch.__version__, torch.version.cuda, os.path.dirname(torch.__file__)],
+            "nvcc": version([os.path.join(cpp_extension.CUDA_HOME or "/usr/local/cuda", "bin", "nvcc"), "--version"]),
+            "cxx": version([os.environ.get("CXX", "c++"), "--version"]),
+            "variables": {k: os.environ[k] for k in BUILD_VARIABLES if k in os.environ}}
+
+
+def cache_key(spec, files, flags, arch, tools):
+    """The key of a build in the cache: a hash of everything that determines it (see the module docstring)."""
+    what = {"spec": spec, "files": sorted(files.items()), "flags": flags, "arch": arch, "toolchain": tools}
+    return hashlib.sha256(json.dumps(what, sort_keys=True).encode()).hexdigest()
+
+
+def disk_usage(path):
+    """The bytes a directory takes on disk."""
+    total = 0
+    for root, _, names in os.walk(path):
+        for name in names:
+            try:
+                st = os.lstat(os.path.join(root, name))
+            except OSError:
+                continue
+            total += getattr(st, "st_blocks", 0) * 512 or st.st_size
+    return total
+
+
+def fetch(cache, key, build_dir):
+    """Copies the build of a key from the cache into build_dir, in its place, and marks the build used; False when the
+    cache has none, or it went while it was copied."""
+    entry = os.path.join(cache, key)
+    if not os.path.exists(os.path.join(entry, MODULE + ".so")):
+        return False
+    shutil.rmtree(build_dir, ignore_errors=True)
+    try:
+        shutil.copytree(entry, build_dir, symlinks=True, ignore=lambda d, _: [SIZE] if d == entry else [])
+        os.utime(entry)
+    except OSError:
+        shutil.rmtree(build_dir, ignore_errors=True)
+        return False
+    return True
+
+
+def store(cache, key, build_dir, files):
+    """Puts the build in build_dir into the cache under its key: its sources, module and stamp, without what the
+    compiler left beside them. A build of the key that another job put there first stays. Then the cache is kept within
+    its limit. A cache that cannot take it (a full disk, say) is left as it is: the build stands without it."""
+    entry = os.path.join(cache, key)
+    if os.path.exists(entry):
+        return
+    try:
+        tmp = tempfile.mkdtemp(prefix=f".{key[:16]}-", dir=cache)
+    except OSError:
+        return
+    try:
+        for rel in [*files, MODULE + ".so", STAMP]:
+            dest = os.path.join(tmp, rel)
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+            shutil.copy2(os.path.join(build_dir, rel), dest)
+        with open(os.path.join(tmp, SIZE), "w") as f:
+            f.write(str(disk_usage(tmp)))
+        os.rename(tmp, entry)
+    except OSError:
+        shutil.rmtree(tmp, ignore_errors=True)
+        return
+    try:
+        st = os.statvfs(cache)
+        evict(cache, CACHE_SHARE * st.f_blocks * st.f_frsize)
+    except OSError:
+        pass
+
+
+def evict(cache, limit):
+    """Removes the least recently used builds of the cache until it holds at most `limit` bytes, and what jobs that
+    stopped while they stored a build left there (after STALE_SEC). Anything but a directory is left alone."""
+    now = time.time()
+    builds = []
+    for name in os.listdir(cache):
+        path = os.path.join(cache, name)
+        try:
+            st = os.lstat(path)
+        except OSError:
+            continue
+        if not stat.S_ISDIR(st.st_mode):
+            continue
+        if name.startswith("."):
+            if now - st.st_mtime > STALE_SEC:
+                shutil.rmtree(path, ignore_errors=True)
+            continue
+        try:
+            with open(os.path.join(path, SIZE)) as f:
+                size = int(f.read())
+        except (OSError, ValueError):
+            size = disk_usage(path)
+        builds.append((st.st_mtime, path, size))
+    total = sum(size for _, _, size in builds)
+    for _, path, size in sorted(builds):
+        if total <= limit:
+            break
+        shutil.rmtree(path, ignore_errors=True)
+        total -= size
+
+
+def stage(solution_dir, build_dir=None, arch=None, extra_cuda_cflags=(), cache=None):
+    """Writes the sources into build_dir (a fresh temporary directory if None) and builds them there.
+
+    Returns (workdir, spec, build seconds or None when an earlier build was reused). A build_dir that already holds
+    a build of the same sources and flags is reused, so CI builds each solution once for its rounds. With a cache
+    (cache_dir(): a development build), a build of the same key is copied from it, and one made is put there."""
+    spec, files = read(solution_dir)
+    arch = arch or gpu_arch()
+    flags = compile_flags(spec, arch)
+    # Flags a tool adds that do not change the code, such as -lineinfo for source-level profiles.
+    flags[1].extend(extra_cuda_cflags)
+    digest = hashlib.sha256(json.dumps([spec, flags, sorted(files.items())]).encode()).hexdigest()
+    if build_dir is None:
+        build_dir = tempfile.mkdtemp(prefix="solution-")
+        atexit.register(shutil.rmtree, build_dir, True)
+    build_dir = os.path.abspath(build_dir)
+    stamp = os.path.join(build_dir, STAMP)
+    try:
+        with open(stamp) as f:
+            if f.read() == digest and os.path.exists(os.path.join(build_dir, MODULE + ".so")):
+                return build_dir, spec, None
+    except OSError:
+        pass
+    key = cache and cache_key(spec, files, flags, arch, toolchain())
+    if key and fetch(cache, key, build_dir):
+        return build_dir, spec, None
+    shutil.rmtree(build_dir, ignore_errors=True)
+    os.makedirs(build_dir)
+    for rel, text in files.items():
+        dest = os.path.join(build_dir, rel)
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        with open(dest, "w", encoding="utf-8") as f:
+            f.write(text)
+    start = time.monotonic()
+    env = dict(os.environ, PYTORCH_ALLOC_CONF="expandable_segments:True")
+    cmd = [sys.executable, "-s", os.path.abspath(__file__), "_build", json.dumps(flags)]
+    try:
+        p = subprocess.run(cmd, cwd=build_dir, env=env, capture_output=True, text=True, timeout=BUILD_TIMEOUT_SEC)
+    except subprocess.TimeoutExpired as e:
+        raise SolutionError(f"the build timed out after {BUILD_TIMEOUT_SEC} s") from e
+    seconds = time.monotonic() - start
+    if p.returncode != 0 or not os.path.exists(os.path.join(build_dir, MODULE + ".so")):
+        log = (p.stdout + p.stderr).strip()
+        raise SolutionError(f"the build failed after {seconds:.0f} s:\n{_errors(log)}")
+    with open(stamp, "w") as f:
+        f.write(digest)
+    if key:
+        store(cache, key, build_dir, files)
+    return build_dir, spec, seconds
+
+
+def _errors(log, limit=6000):
+    """The compiler's error lines with some context, else the log's tail."""
+    lines = log.splitlines()
+    keep = [i for i, line in enumerate(lines) if "error" in line.lower()]
+    if not keep:
+        return log[-limit:]
+    picked = sorted({j for i in keep for j in range(max(0, i - 2), min(len(lines), i + 3))})
+    return "\n".join(lines[j] for j in picked)[-limit:]
+
+
+# fastvideo-kernel's CUTLASS submodule (git submodule update --init fastvideo-kernel/include/cutlass).
+_CUTLASS_SUBMODULE = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "include",
+                                  "cutlass")
+
+
+def _build(flags):
+    """The compile step, run in the working directory by stage() in its own process."""
+    import torch.utils.cpp_extension as ext
+
+    cflags, cuda_cflags, ld_flags = json.loads(flags)
+    here = os.getcwd()
+    srcs = sorted(os.path.join(here, p) for p in os.listdir(here)
+                  if p.endswith(CPP_SOURCES) and os.path.isfile(os.path.join(here, p)))
+    if not srcs:
+        raise RuntimeError("No CUDA/C++ source files found in working directory")
+    cutlass = os.environ.get("CUTLASS_DIR") or _CUTLASS_SUBMODULE
+    try:
+        ext.load(name=MODULE, sources=srcs, extra_cuda_cflags=cuda_cflags, extra_cflags=cflags, extra_ldflags=ld_flags,
+                 extra_include_paths=[here, f"{cutlass}/include", f"{cutlass}/tools/util/include"],
+                 build_directory=here, verbose=True)
+    except ImportError as e:
+        # Built, and not loadable here: a machine without a GPU has no driver to load it against, and whoever loads it
+        # says whether it loads where it runs. Where a GPU is, a module that does not load (an undefined symbol, a
+        # missing library) is a build that failed, and is neither stamped nor kept in the cache.
+        import torch
+
+        if torch.cuda.is_available():
+            raise
+        print(f"built; not loaded on this machine: {e}", flush=True)
+    built = [p for p in os.listdir(here) if p.startswith(MODULE) and p.endswith(".so") and p != MODULE + ".so"]
+    if built:
+        os.rename(built[0], MODULE + ".so")
+
+
+class Solution:
+    """A loaded solution: fn, the entry (run), called as the deployment calls it, fn(*inputs, *outputs); prepare, the
+    module's prepare function or None; spec (solution.json), entry, languages, language, module, workdir, and
+    build_seconds (None when an earlier build was reused, or when the build was staged elsewhere)."""
+
+    def __init__(self, spec, fn, prepare, module, workdir, build_seconds):
+        self.spec = spec
+        self.languages = list(spec["languages"])
+        self.language = self.languages[0]
+        self.entry = spec["entry_point"]
+        self.fn, self.prepare, self.module = fn, prepare, module
+        self.workdir, self.build_seconds = workdir, build_seconds
+
+
+def import_built(workdir, spec, chdir=True, build_seconds=None):
+    """Imports the module that stage() built in workdir and takes its entry and prepare from it. chdir: make the
+    working directory the current one, as in deployment."""
+    import torch  # noqa: F401  (the module needs libtorch loaded first)
+
+    if chdir:
+        os.chdir(workdir)
+    file, func = spec["entry_point"].split("::")
+    module_spec = importlib.util.spec_from_file_location(MODULE, os.path.join(workdir, MODULE + ".so"))
+    module = importlib.util.module_from_spec(module_spec)
+    try:
+        module_spec.loader.exec_module(module)
+    except Exception as e:
+        raise SolutionError(f"the built module does not load: {e}") from e
+    fn = getattr(module, func, None)
+    if not callable(fn):
+        raise SolutionError(f"{file} defines no function {func} in its PYBIND11_MODULE")
+    prepare = getattr(module, PREPARE, None)
+    if prepare is not None and not callable(prepare):
+        raise SolutionError(f"the module's {PREPARE} is not a function")
+    return Solution(spec, fn, prepare, module, workdir, build_seconds)
+
+
+def load(solution_dir, build_dir=None, chdir=True, extra_cuda_cflags=(), cache=None):
+    """Builds and loads a solution as it is deployed; see the module docstring.
+
+    build_dir: where to write and build it (a fresh temporary directory if None; reused if it holds the same build).
+    chdir: make the working directory the current one, as in deployment. extra_cuda_cflags: flags a tool adds to the
+    deployed ones, that must not change the code (e.g. -lineinfo). cache: the node's cache of development builds
+    (cache_dir()), or None to build fresh."""
+    workdir, spec, seconds = stage(solution_dir, build_dir, extra_cuda_cflags=extra_cuda_cflags, cache=cache)
+    return import_built(workdir, spec, chdir=chdir, build_seconds=seconds)
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    c = sub.add_parser("check", help="build and load a solution")
+    c.add_argument("solution", nargs="?", default="solution")
+    c.add_argument("--build-dir")
+    d = sub.add_parser("build", help="write and build a solution in a directory, without loading it")
+    d.add_argument("solution", nargs="?", default="solution")
+    d.add_argument("build_dir", nargs="?", default="build")
+    d.add_argument("--arch", help="the architecture to build for, e.g. sm_100a (default: this machine's GPU's)")
+    for q in (c, d):
+        q.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds")
+    b = sub.add_parser("_build")
+    b.add_argument("flags")
+    a = p.parse_args()
+    if a.cmd == "_build":
+        _build(a.flags)
+        return
+    cache = None if a.no_cache else cache_dir()
+    try:
+        if a.cmd == "build":
+            where, _, seconds = stage(os.path.abspath(a.solution), os.path.abspath(a.build_dir), a.arch, cache=cache)
+            module = os.path.join(where, MODULE + ".so")
+            print(f"built in {seconds:.1f} s: {module}" if seconds is not None else f"reused a build: {module}")
+            return
+        sol = load(os.path.abspath(a.solution), a.build_dir and os.path.abspath(a.build_dir), cache=cache)
+    except SolutionError as e:
+        sys.exit(f"FAILED: {e}")
+    built = f", built in {sol.build_seconds:.1f} s" if sol.build_seconds is not None else ", a build reused"
+    prepare = f", with {PREPARE}" if sol.prepare is not None else f", without {PREPARE}"
+    print(f"loaded {sol.entry} ({', '.join(sol.languages)}){prepare}{built}")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,428 @@
+"""The task: decoding Wan2.2 VAE latents into video on one, two or four GB200 GPUs of one node, as bench.py, ci.py and
+bounds.py read it. It defines the clips and workloads, the weights' order, the data, the error allowance and the work
+of one decode.
+
+What is computed
+----------------
+reference.py (in the kernel repo) defines the result: diffusers' AutoencoderKLWan._decode of the Wan2.2-TI2V-5B VAE,
+latents z [1, 48, T, 44, 78] to a video [1, 3, 4T-3, 704, 1248] in [-1, 1]: post_quant_conv, then the decoder one
+latent frame at a time with its causal caches (the first latent frame yields one frame, every later one four),
+unpatchify 2 x 2, clamp. The judge compares a solution with the fp32 decode of the same latents, cached as videos in
+$DATA (see Correctness).
+
+The interface (as loader.py builds it and bench.py calls it, positional, destination passing)
+-------------
+    run(z, weights, rank, world, pg, out)
+    prepare(weights, rank, world, pg) -> list[Tensor]       optional
+
+- z: [1, 48, T, 44, 78] bf16, contiguous, the full latents on every rank: what a pipeline passes to
+  vae.decode(latents), rounded to bf16.
+- weights: the 110 tensors of WEIGHTS below, bf16, on this rank's GPU; or what prepare returned.
+- rank, world: this process's rank and the number of ranks (1, 2 or 4), rank r on GPU r.
+- pg: torch.distributed's NCCL ProcessGroup over the ranks; a C++ entry takes it as
+  `const c10::intrusive_ptr<c10d::ProcessGroup>&` (from <torch/csrc/distributed/c10d/ProcessGroup.hpp>).
+- out: on rank 0, the video [1, 3, 4T-3, 704, 1248] bf16, contiguous, from torch.empty and not zeroed (every element
+  must be written); on the other ranks an empty tensor. When run has returned on every rank, rank 0's out holds the
+  video.
+- prepare, when the module exports it, is called once per process before any timed call, and is not timed. It sees
+  the weights, not the latents; its result replaces `weights` in every run.
+
+Numerics: bf16 weights and activations with fp32 accumulation, and fp32 statistics in the RMS norms; nothing narrower
+(no FP8, FP6 or FP4). This is a rule for the reviewers: no check reads it from the code.
+
+Workloads
+---------
+A workload is a clip decoded at a world size, named <clip>x<world>. A clip is the first T latent frames of one of
+three real 20 s clips encoded by the VAE (latents_<source>.pt in $DATA):
+
+    clip   source       latent frames T   video frames 4T-3   worlds    timed
+    f1     row_000168          1                   1           1, 4      no
+    f5     row_000168          2                   5           1, 4      no
+    f17    row_000168          5                  17           1, 4      no
+    f121   row_000168         31                 121           1, 2, 4   yes
+    f241   row_000129         61                 241           1, 2, 4   yes
+    f481   row_000090        121                 481           1, 2, 4   yes
+
+The nine timed workloads (f121x1 ... f481x4) make the objective, their geomean; the six short ones (f1x1 ... f17x4)
+are checked for correctness and never timed: a decoder sees every clip length, and the first latent frame is a
+special case.
+
+Correctness
+-----------
+On rank 0's out, against ref_<clip>.pt, the fp32 decode of the clip's latents (fp32 latents, TF32 off), stored as fp16:
+every element finite and within [-1, 1]; and with d = out - ref in fp32,
+
+    max |d|  <= min(MAX_ABS_CAP, ERROR_FACTOR x bf16_max_abs)
+    rms(d)   <= ERROR_FACTOR x bf16_rms
+
+bf16_max_abs and bf16_rms being the error of the release's bf16 decode (bf16 weights under bf16 autocast, the
+latents rounded to bf16) against that fp32 decode on the same clip, from calibration.json. PSNR, 10 log10(4 / mse), is
+shown beside them.
+
+The weights
+-----------
+run() gets the decoder's parameters and post_quant_conv's, 110 tensors, as a list in diffusers' registration order
+(that of AutoencoderKLWan.named_parameters(), post_quant_conv first, then the decoder from conv_in to conv_out;
+$DATA/weight_names.json lists the same), bf16, in PyTorch's layouts: a conv3d's weight is [out, in, kt, kh, kw], a
+conv2d's [out, in, kh, kw], a bias [out], an RMS norm's gain [C, 1, 1, 1] (the attention's [C, 1, 1]).
+$DATA/decoder_weights_bf16.pt is that list: the release's fp32 weights rounded to bf16 to nearest even, as diffusers'
+from_pretrained(torch_dtype=torch.bfloat16) casts them. The order is not the order of execution: the mid block
+registers its attention before its two residual blocks, though it runs resnets.0, the attention, then resnets.1; an
+upsampler registers its 3x3 conv (resample.1) before its time conv, though the time conv runs first. The attention's
+to_qkv gives q, k and v as its output channels 0-1023, 1024-2047 and 2048-3071. A time conv C -> 2C gives frame 2t of
+its output as channels 0..C-1 and frame 2t+1 as channels C..2C-1. conv_out's 12 channels are (colour k, column offset
+j, row offset i) of 2 x 2 pixel patches: video[k, f, 2y + i, 2x + j] = channel 4k + 2j + i at (f, y, x).
+
+      #  name                                                shape                  role
+      0  post_quant_conv.weight                              [48, 48, 1, 1, 1]      1x1x1 conv 48 -> 48, on the latents
+      1  post_quant_conv.bias                                [48]
+      2  decoder.conv_in.weight                              [1024, 48, 3, 3, 3]    causal 3x3x3 conv 48 -> 1024
+      3  decoder.conv_in.bias                                [1024]
+      4  decoder.mid_block.attentions.0.norm.gamma           [1024, 1, 1]           RMS norm gain, before to_qkv
+      5  decoder.mid_block.attentions.0.to_qkv.weight        [3072, 1024, 1, 1]     1x1 conv2d 1024 -> 3072 (q, k, v)
+      6  decoder.mid_block.attentions.0.to_qkv.bias          [3072]
+      7  decoder.mid_block.attentions.0.proj.weight          [1024, 1024, 1, 1]     1x1 conv2d 1024 -> 1024
+      8  decoder.mid_block.attentions.0.proj.bias            [1024]
+      9  decoder.mid_block.resnets.0.norm1.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv1
+     10  decoder.mid_block.resnets.0.conv1.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     11  decoder.mid_block.resnets.0.conv1.bias              [1024]
+     12  decoder.mid_block.resnets.0.norm2.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv2
+     13  decoder.mid_block.resnets.0.conv2.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     14  decoder.mid_block.resnets.0.conv2.bias              [1024]
+     15  decoder.mid_block.resnets.1.norm1.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv1
+     16  decoder.mid_block.resnets.1.conv1.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     17  decoder.mid_block.resnets.1.conv1.bias              [1024]
+     18  decoder.mid_block.resnets.1.norm2.gamma             [1024, 1, 1, 1]        RMS norm gain, before conv2
+     19  decoder.mid_block.resnets.1.conv2.weight            [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     20  decoder.mid_block.resnets.1.conv2.bias              [1024]
+     21  decoder.up_blocks.0.resnets.0.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     22  decoder.up_blocks.0.resnets.0.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     23  decoder.up_blocks.0.resnets.0.conv1.bias            [1024]
+     24  decoder.up_blocks.0.resnets.0.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     25  decoder.up_blocks.0.resnets.0.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     26  decoder.up_blocks.0.resnets.0.conv2.bias            [1024]
+     27  decoder.up_blocks.0.resnets.1.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     28  decoder.up_blocks.0.resnets.1.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     29  decoder.up_blocks.0.resnets.1.conv1.bias            [1024]
+     30  decoder.up_blocks.0.resnets.1.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     31  decoder.up_blocks.0.resnets.1.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     32  decoder.up_blocks.0.resnets.1.conv2.bias            [1024]
+     33  decoder.up_blocks.0.resnets.2.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     34  decoder.up_blocks.0.resnets.2.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     35  decoder.up_blocks.0.resnets.2.conv1.bias            [1024]
+     36  decoder.up_blocks.0.resnets.2.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     37  decoder.up_blocks.0.resnets.2.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     38  decoder.up_blocks.0.resnets.2.conv2.bias            [1024]
+     39  decoder.up_blocks.0.upsampler.resample.1.weight     [1024, 1024, 3, 3]     3x3 conv2d 1024 -> 1024, after x2
+     40  decoder.up_blocks.0.upsampler.resample.1.bias       [1024]
+     41  decoder.up_blocks.0.upsampler.time_conv.weight      [2048, 1024, 3, 1, 1]  causal 3x1x1 conv 1024 -> 2048
+     42  decoder.up_blocks.0.upsampler.time_conv.bias        [2048]
+     43  decoder.up_blocks.1.resnets.0.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     44  decoder.up_blocks.1.resnets.0.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     45  decoder.up_blocks.1.resnets.0.conv1.bias            [1024]
+     46  decoder.up_blocks.1.resnets.0.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     47  decoder.up_blocks.1.resnets.0.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     48  decoder.up_blocks.1.resnets.0.conv2.bias            [1024]
+     49  decoder.up_blocks.1.resnets.1.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     50  decoder.up_blocks.1.resnets.1.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     51  decoder.up_blocks.1.resnets.1.conv1.bias            [1024]
+     52  decoder.up_blocks.1.resnets.1.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     53  decoder.up_blocks.1.resnets.1.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     54  decoder.up_blocks.1.resnets.1.conv2.bias            [1024]
+     55  decoder.up_blocks.1.resnets.2.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     56  decoder.up_blocks.1.resnets.2.conv1.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     57  decoder.up_blocks.1.resnets.2.conv1.bias            [1024]
+     58  decoder.up_blocks.1.resnets.2.norm2.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv2
+     59  decoder.up_blocks.1.resnets.2.conv2.weight          [1024, 1024, 3, 3, 3]  causal 3x3x3 conv 1024 -> 1024
+     60  decoder.up_blocks.1.resnets.2.conv2.bias            [1024]
+     61  decoder.up_blocks.1.upsampler.resample.1.weight     [1024, 1024, 3, 3]     3x3 conv2d 1024 -> 1024, after x2
+     62  decoder.up_blocks.1.upsampler.resample.1.bias       [1024]
+     63  decoder.up_blocks.1.upsampler.time_conv.weight      [2048, 1024, 3, 1, 1]  causal 3x1x1 conv 1024 -> 2048
+     64  decoder.up_blocks.1.upsampler.time_conv.bias        [2048]
+     65  decoder.up_blocks.2.resnets.0.norm1.gamma           [1024, 1, 1, 1]        RMS norm gain, before conv1
+     66  decoder.up_blocks.2.resnets.0.conv1.weight          [512, 1024, 3, 3, 3]   causal 3x3x3 conv 1024 -> 512
+     67  decoder.up_blocks.2.resnets.0.conv1.bias            [512]
+     68  decoder.up_blocks.2.resnets.0.norm2.gamma           [512, 1, 1, 1]         RMS norm gain, before conv2
+     69  decoder.up_blocks.2.resnets.0.conv2.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     70  decoder.up_blocks.2.resnets.0.conv2.bias            [512]
+     71  decoder.up_blocks.2.resnets.0.conv_shortcut.weight  [512, 1024, 1, 1, 1]   1x1x1 conv 1024 -> 512, shortcut
+     72  decoder.up_blocks.2.resnets.0.conv_shortcut.bias    [512]
+     73  decoder.up_blocks.2.resnets.1.norm1.gamma           [512, 1, 1, 1]         RMS norm gain, before conv1
+     74  decoder.up_blocks.2.resnets.1.conv1.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     75  decoder.up_blocks.2.resnets.1.conv1.bias            [512]
+     76  decoder.up_blocks.2.resnets.1.norm2.gamma           [512, 1, 1, 1]         RMS norm gain, before conv2
+     77  decoder.up_blocks.2.resnets.1.conv2.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     78  decoder.up_blocks.2.resnets.1.conv2.bias            [512]
+     79  decoder.up_blocks.2.resnets.2.norm1.gamma           [512, 1, 1, 1]         RMS norm gain, before conv1
+     80  decoder.up_blocks.2.resnets.2.conv1.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     81  decoder.up_blocks.2.resnets.2.conv1.bias            [512]
+     82  decoder.up_blocks.2.resnets.2.norm2.gamma           [512, 1, 1, 1]         RMS norm gain, before conv2
+     83  decoder.up_blocks.2.resnets.2.conv2.weight          [512, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 512
+     84  decoder.up_blocks.2.resnets.2.conv2.bias            [512]
+     85  decoder.up_blocks.2.upsampler.resample.1.weight     [512, 512, 3, 3]       3x3 conv2d 512 -> 512, after x2
+     86  decoder.up_blocks.2.upsampler.resample.1.bias       [512]
+     87  decoder.up_blocks.3.resnets.0.norm1.gamma           [512, 1, 1, 1]         RMS norm gain, before conv1
+     88  decoder.up_blocks.3.resnets.0.conv1.weight          [256, 512, 3, 3, 3]    causal 3x3x3 conv 512 -> 256
+     89  decoder.up_blocks.3.resnets.0.conv1.bias            [256]
+     90  decoder.up_blocks.3.resnets.0.norm2.gamma           [256, 1, 1, 1]         RMS norm gain, before conv2
+     91  decoder.up_blocks.3.resnets.0.conv2.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+     92  decoder.up_blocks.3.resnets.0.conv2.bias            [256]
+     93  decoder.up_blocks.3.resnets.0.conv_shortcut.weight  [256, 512, 1, 1, 1]    1x1x1 conv 512 -> 256, shortcut
+     94  decoder.up_blocks.3.resnets.0.conv_shortcut.bias    [256]
+     95  decoder.up_blocks.3.resnets.1.norm1.gamma           [256, 1, 1, 1]         RMS norm gain, before conv1
+     96  decoder.up_blocks.3.resnets.1.conv1.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+     97  decoder.up_blocks.3.resnets.1.conv1.bias            [256]
+     98  decoder.up_blocks.3.resnets.1.norm2.gamma           [256, 1, 1, 1]         RMS norm gain, before conv2
+     99  decoder.up_blocks.3.resnets.1.conv2.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+    100  decoder.up_blocks.3.resnets.1.conv2.bias            [256]
+    101  decoder.up_blocks.3.resnets.2.norm1.gamma           [256, 1, 1, 1]         RMS norm gain, before conv1
+    102  decoder.up_blocks.3.resnets.2.conv1.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+    103  decoder.up_blocks.3.resnets.2.conv1.bias            [256]
+    104  decoder.up_blocks.3.resnets.2.norm2.gamma           [256, 1, 1, 1]         RMS norm gain, before conv2
+    105  decoder.up_blocks.3.resnets.2.conv2.weight          [256, 256, 3, 3, 3]    causal 3x3x3 conv 256 -> 256
+    106  decoder.up_blocks.3.resnets.2.conv2.bias            [256]
+    107  decoder.norm_out.gamma                              [256, 1, 1, 1]         RMS norm gain, before conv_out
+    108  decoder.conv_out.weight                             [12, 256, 3, 3, 3]     causal 3x3x3 conv 256 -> 12
+    109  decoder.conv_out.bias                               [12]
+
+Up blocks 0 and 1 run at 1024 channels, up block 2 from 1024 to 512, up block 3 from 512 to 256; a residual block's
+first conv changes the channel count, and its 1x1x1 conv_shortcut does so on the shortcut.
+
+Data
+----
+$DATA (default /data), read-only: latents_<source>.pt ({"latents": fp32 [1, 48, 121, 44, 78]}), ref_<clip>.pt
+({"video": fp16 [1, 3, 4T-3, 704, 1248], "source", "latent_frames"}), calibration.json (per clip: bf16_max_abs,
+bf16_mean_abs, bf16_rms, bf16_psnr_db, fp32_decode_s, bf16_decode_s), weight_names.json and decoder_weights_bf16.pt.
+
+The work of one decode (bounds)
+-------------------------------
+FLOPs, a multiply-add counting 2: every convolution at its full kernel over its output positions (2 x the weight's
+elements x frames x rows x columns, zero padding included, which adds a few percent at 44 x 78 and less than one
+beyond), plus the attention's two products per latent frame (2 x 2 x 3432^2 x 1024). The frames at each level, for T
+latent frames: T at 44 x 78 (post_quant_conv, conv_in, the mid block, up block 0's residual blocks); T - 1 for up
+block 0's time conv, which skips the first latent frame; 2T - 1 at 88 x 156 (up block 0's 3x3 conv, up block 1's
+residual blocks); 2T - 2 for up block 1's time conv; 4T - 3 at 176 x 312 (up block 1's 3x3 conv, up block 2's
+residual blocks) and at 352 x 624 (up block 2's 3x3 conv, up block 3, conv_out). The FLOPs are divided over the
+workload's GPUs. Bytes: the video written (on rank 0) and the weights read once.
+"""
+
+import json
+import math
+import os
+
+LATENT_CHANNELS = 48
+LATENT_HW = (44, 78)
+VIDEO_HW = (704, 1248)
+# How far a solution may err, against what the release's bf16 decode errs by on the same clip (see Correctness).
+ERROR_FACTOR = 3.0
+MAX_ABS_CAP = 2.0
+WEIGHTS_FILE = "decoder_weights_bf16.pt"
+NAMES_FILE = "weight_names.json"
+
+CLIPS = {
+    "f1": {"source": "row_000168", "latent_frames": 1},
+    "f5": {"source": "row_000168", "latent_frames": 2},
+    "f17": {"source": "row_000168", "latent_frames": 5},
+    "f121": {"source": "row_000168", "latent_frames": 31},
+    "f241": {"source": "row_000129", "latent_frames": 61},
+    "f481": {"source": "row_000090", "latent_frames": 121},
+}
+# The short clips are checked at one GPU and at four, never timed; the long ones are checked and timed at 1, 2 and 4.
+SHORT_CLIPS, LONG_CLIPS = ("f1", "f5", "f17"), ("f121", "f241", "f481")
+WORKLOADS = {
+    **{f"{c}x{n}": {"clip": c, "world": n, "timed": False} for c in SHORT_CLIPS for n in (1, 4)},
+    **{f"{c}x{n}": {"clip": c, "world": n, "timed": True} for c in LONG_CLIPS for n in (1, 2, 4)},
+}
+
+# decoder.up_blocks.{i}: (in channels, out channels, temporal upsampling, spatial upsampling)
+UP_BLOCKS = ((1024, 1024, True, True), (1024, 1024, True, True), (1024, 512, False, True), (512, 256, False, False))
+RESNETS_PER_UP_BLOCK = 3
+
+
+def _weights():
+    """(name, shape) of run()'s weights in diffusers' registration order (see The weights)."""
+    out = []
+
+    def conv(name, shape):
+        out.extend([(f"{name}.weight", tuple(shape)), (f"{name}.bias", (shape[0],))])
+
+    def resnet(name, cin, cout):
+        out.append((f"{name}.norm1.gamma", (cin, 1, 1, 1)))
+        conv(f"{name}.conv1", (cout, cin, 3, 3, 3))
+        out.append((f"{name}.norm2.gamma", (cout, 1, 1, 1)))
+        conv(f"{name}.conv2", (cout, cout, 3, 3, 3))
+        if cin != cout:
+            conv(f"{name}.conv_shortcut", (cout, cin, 1, 1, 1))
+
+    conv("post_quant_conv", (48, 48, 1, 1, 1))
+    conv("decoder.conv_in", (1024, 48, 3, 3, 3))
+    out.append(("decoder.mid_block.attentions.0.norm.gamma", (1024, 1, 1)))
+    conv("decoder.mid_block.attentions.0.to_qkv", (3072, 1024, 1, 1))
+    conv("decoder.mid_block.attentions.0.proj", (1024, 1024, 1, 1))
+    for j in range(2):
+        resnet(f"decoder.mid_block.resnets.{j}", 1024, 1024)
+    for i, (cin, cout, temporal, spatial) in enumerate(UP_BLOCKS):
+        for j in range(RESNETS_PER_UP_BLOCK):
+            resnet(f"decoder.up_blocks.{i}.resnets.{j}", cin if j == 0 else cout, cout)
+        if spatial:
+            conv(f"decoder.up_blocks.{i}.upsampler.resample.1", (cout, cout, 3, 3))
+        if temporal:
+            conv(f"decoder.up_blocks.{i}.upsampler.time_conv", (2 * cout, cout, 3, 1, 1))
+    out.append(("decoder.norm_out.gamma", (256, 1, 1, 1)))
+    conv("decoder.conv_out", (12, 256, 3, 3, 3))
+    return tuple(out)
+
+
+WEIGHTS = _weights()
+WEIGHT_NAMES = [name for name, _ in WEIGHTS]
+
+
+def data_dir():
+    """The task's data: $DATA, else /data, where the jobs mount it."""
+    return os.environ.get("DATA") or "/data"
+
+
+def frames(latent_frames):
+    """The video frames of T latent frames: one for the first, four for every later one."""
+    return 4 * latent_frames - 3
+
+
+def output_shape(workload):
+    """The shape of rank 0's out at a workload."""
+    return (1, 3, frames(CLIPS[WORKLOADS[workload]["clip"]]["latent_frames"]), *VIDEO_HW)
+
+
+def _load(path):
+    import torch
+
+    try:
+        return torch.load(path, map_location="cpu", weights_only=True, mmap=True)
+    except (OSError, RuntimeError) as e:
+        raise ValueError(f"cannot read {path}: {e}") from e
+
+
+def make_inputs(workload, device):
+    """[z]: the clip's latents, the first T latent frames of its source, rounded to bf16, contiguous on `device`."""
+    import torch
+
+    clip = CLIPS[WORKLOADS[workload]["clip"]]
+    path = os.path.join(data_dir(), f"latents_{clip['source']}.pt")
+    latents = _load(path)["latents"]
+    t = clip["latent_frames"]
+    if latents.dim() != 5 or tuple(latents.shape[:2]) != (1, LATENT_CHANNELS) or tuple(latents.shape[3:]) != LATENT_HW \
+            or latents.shape[2] < t:
+        raise ValueError(f"{path}: latents {tuple(latents.shape)}, expected [1, 48, >= {t}, 44, 78]")
+    return [latents[:, :, :t].to(device=device, dtype=torch.bfloat16).contiguous()]
+
+
+def load_weights(device):
+    """run()'s weights on `device`: decoder_weights_bf16.pt, checked against WEIGHTS (and weight_names.json where it
+    is), so that a solution indexes what this file documents."""
+    import torch
+
+    listed = os.path.join(data_dir(), NAMES_FILE)
+    if os.path.exists(listed):
+        with open(listed) as f:
+            names = [(e["name"], tuple(e["shape"])) for e in json.load(f)]
+        if names != list(WEIGHTS):
+            raise ValueError(f"{listed} does not list spec.WEIGHTS: the weights' order or shapes changed")
+    path = os.path.join(data_dir(), WEIGHTS_FILE)
+    weights = _load(path)
+    if not isinstance(weights, list) or len(weights) != len(WEIGHTS):
+        raise ValueError(f"{path}: expected a list of {len(WEIGHTS)} tensors")
+    for i, (t, (name, shape)) in enumerate(zip(weights, WEIGHTS)):
+        if not isinstance(t, torch.Tensor) or t.dtype != torch.bfloat16 or tuple(t.shape) != shape:
+            got = f"{t.dtype} {tuple(t.shape)}" if isinstance(t, torch.Tensor) else type(t).__name__
+            raise ValueError(f"{path}: tensor {i} is {got}, expected {name}, bf16 {shape}")
+    return [t.to(device).contiguous() for t in weights]
+
+
+def load_reference(clip):
+    """The clip's fp32 decode, as stored: a CPU fp16 tensor [1, 3, 4T-3, 704, 1248] (memory-mapped)."""
+    import torch
+
+    path = os.path.join(data_dir(), f"ref_{clip}.pt")
+    ref = _load(path)
+    t, source = CLIPS[clip]["latent_frames"], CLIPS[clip]["source"]
+    video = ref.get("video") if isinstance(ref, dict) else None
+    shape = (1, 3, frames(t), *VIDEO_HW)
+    if not isinstance(video, torch.Tensor) or tuple(video.shape) != shape or video.dtype != torch.float16:
+        raise ValueError(f"{path}: expected video, fp16 {list(shape)}")
+    if ref.get("source") != source or ref.get("latent_frames") != t:
+        raise ValueError(f"{path} decodes {ref.get('source')}'s first {ref.get('latent_frames')} latent frames, not "
+                         f"{source}'s first {t}")
+    return video
+
+
+def calibration(clip):
+    """The release's bf16 decode of the clip against its fp32 decode (calibration.json's entry)."""
+    path = os.path.join(data_dir(), "calibration.json")
+    try:
+        with open(path) as f:
+            entry = json.load(f)[clip]
+    except (OSError, ValueError, KeyError) as e:
+        raise ValueError(f"{path} has no entry for {clip}: {e}") from e
+    if entry.get("source") != CLIPS[clip]["source"] or entry.get("latent_frames") != CLIPS[clip]["latent_frames"]:
+        raise ValueError(f"{path}: {clip} is calibrated on another clip ({entry.get('source')}, "
+                         f"{entry.get('latent_frames')} latent frames)")
+    return entry
+
+
+def allowance(clip):
+    """{"max_abs", "rms"}: how far a solution's video may be from the fp32 decode of the clip (see Correctness)."""
+    entry = calibration(clip)
+    return {"max_abs": min(MAX_ABS_CAP, ERROR_FACTOR * entry["bf16_max_abs"]), "rms": ERROR_FACTOR * entry["bf16_rms"]}
+
+
+def _conv_level(name, t):
+    """(output frames, upsampling factor in space) of the conv whose weight is `name`, decoding t latent frames."""
+    if name.startswith(("post_quant_conv.", "decoder.conv_in.", "decoder.mid_block.", "decoder.up_blocks.0.resnets.")):
+        return t, 1
+    if name.startswith("decoder.up_blocks.0.upsampler.time_conv."):
+        return t - 1, 1
+    if name.startswith(("decoder.up_blocks.0.upsampler.resample.", "decoder.up_blocks.1.resnets.")):
+        return 2 * t - 1, 2
+    if name.startswith("decoder.up_blocks.1.upsampler.time_conv."):
+        return 2 * t - 2, 2
+    if name.startswith(("decoder.up_blocks.1.upsampler.resample.", "decoder.up_blocks.2.resnets.")):
+        return 4 * t - 3, 4
+    return 4 * t - 3, 8
+
+
+def decode_flops(t):
+    """FLOPs of one decode of t latent frames (see The work of one decode)."""
+    h, w = LATENT_HW
+    total = 0
+    for name, shape in WEIGHTS:
+        if name.endswith(".weight") and len(shape) >= 4:
+            count, scale = _conv_level(name, t)
+            total += 2 * math.prod(shape) * count * (h * scale) * (w * scale)
+    return total + 2 * 2 * (h * w) ** 2 * 1024 * t
+
+
+def bounds(workload):
+    """{"flops", "bytes", "gpus"} of one call at the workload, for bounds.py: the decode's FLOPs (spread over the
+    world's GPUs), and the bytes that must move: the video written and the weights read."""
+    entry = WORKLOADS[workload]
+    t = CLIPS[entry["clip"]]["latent_frames"]
+    video = math.prod(output_shape(workload)) * 2
+    weights = sum(math.prod(shape) for _, shape in WEIGHTS) * 2
+    return {"flops": decode_flops(t), "bytes": video + weights, "gpus": entry["world"]}
+
+
+# The stack rules beside audit.py's: run() reads nothing but its arguments.
+RULES = {
+    "file-io": {
+        "include": r"(fstream|fcntl\.h|sys/mman\.h|filesystem|experimental/filesystem)",
+        "text": r"\b(fopen|fopen64|freopen|fdopen|fread|fwrite|ifstream|ofstream|fstream|filebuf|mmap|mmap64|shm_open|"
+                r"torch::load|torch::save|torch::pickle_load|torch::pickle_save|torch::jit::load)\b|"
+                r"::\s*(open|openat|read|pread|write|pwrite)\s*\(",
+        "why": "run() and prepare() read nothing but their arguments and write nothing but out: no files (the "
+               "reference videos lie on the job's disk) and nothing that outlives a call"},
+}
+
+
+if __name__ == "__main__":
+    for w in WORKLOADS:
+        b = bounds(w)
+        print(f"{w:<8} {b['flops'] / 1e15:8.4f} PFLOP over {b['gpus']} GPU(s), {b['bytes'] / 1e9:6.2f} GB")
@@ -0,0 +1,143 @@
+"""The decoder of the Wan2.2 VAE (Wan2.2-TI2V-5B, 48 latent channels) on sm_100a (B200 / GB200), on one or several GPUs
+of a node, built from solution/ on first use.
+
+It decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T - 3, 704, 1248] (bf16, in [-1, 1]), as
+diffusers' AutoencoderKLWan.decode(z).sample computes it (reference.py defines the result), for any T >= 1. Other
+latent sizes and batches above one are refused: the kernels are specialised to 704 x 1248.
+
+    decoder = WanVAEDecoder(vae.state_dict(), device="cuda")    # diffusers' AutoencoderKLWan, or its state dict
+    video = decoder.decode(latents)                              # [1, 3, 4T - 3, 704, 1248] bf16
+
+Over several GPUs, every rank builds a decoder with the same process group and calls decode with the full latents;
+rank 0 gets the video and the other ranks None. Without a process group, decode uses torch.distributed's default
+group, and initializes a single-rank NCCL group when there is none (one GPU only). With local=True each process decodes
+on its own GPU alone, whatever torch.distributed holds: one rank of one, so the kernels never communicate (comm.cu
+returns before any exchange below two ranks) and no group has to be created collectively.
+
+frames_per_pass sets how many latent frames each GPU carries through one pass of the decoder, which trades memory for
+speed; the video is the same at any value. On one GB200 at 8 to 20 s of video: 2 frames peak at 14-16 GiB allocated
+(x0.85 the speed of 16 frames), 4 (the default) at 21-22 GiB (x0.92), 8 at 34-35 GiB (x0.97), 16 at 60-61 GiB. The
+allocator reserves more than it allocates; PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True keeps the difference small.
+"""
+
+import atexit
+import functools
+import importlib.util
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+LATENT_CHANNELS, LATENT_HW, VIDEO_HW = 48, (44, 78), (704, 1248)
+
+
+def _own(path, name):
+    """The module at path, imported under a name of this kernel's own. Other kernels built the same way
+    keep tools and references of the same file names; a bare import would hand whichever loaded first to both."""
+    name = f"wan22_vae_decode_{name}"
+    module = sys.modules.get(name)
+    if module is None:
+        module_spec = importlib.util.spec_from_file_location(name, path)
+        module = importlib.util.module_from_spec(module_spec)
+        sys.modules[name] = module
+        try:
+            module_spec.loader.exec_module(module)
+        except BaseException:
+            del sys.modules[name]
+            raise
+    return module
+
+
+@functools.lru_cache(maxsize=None)
+def _tools():
+    """tools/loader.py and tools/spec.py."""
+    return (_own(os.path.join(HERE, "tools", "loader.py"), "loader"),
+            _own(os.path.join(HERE, "tools", "spec.py"), "spec"))
+
+
+@functools.lru_cache(maxsize=None)
+def _module():
+    """The extension built from solution/, once per process. Builds are kept in WAN22_VAE_DECODE_BUILD_DIR (default
+    ~/.cache/fastvideo/wan22_vae_decode), one directory per local rank, and reused while the sources and flags match."""
+    loader, _ = _tools()
+    base = os.environ.get("WAN22_VAE_DECODE_BUILD_DIR", os.path.expanduser("~/.cache/fastvideo/wan22_vae_decode"))
+    build_dir = os.path.join(base, f"rank{os.environ.get('LOCAL_RANK', '0')}")
+    os.makedirs(base, exist_ok=True)
+    return loader.load(os.path.join(HERE, "solution"), build_dir=build_dir, chdir=False).module
+
+
+def weight_names() -> list[str]:
+    """The 110 parameters the decoder takes, under diffusers' AutoencoderKLWan names, in the order it takes them."""
+    return list(_tools()[1].WEIGHT_NAMES)
+
+
+def decoder_weights(state_dict, device=None) -> list[Tensor]:
+    """The decoder's weights from an AutoencoderKLWan state dict (or the module itself): its post_quant_conv and
+    decoder parameters, in weight_names() order, as contiguous bf16 tensors on `device`."""
+    if isinstance(state_dict, torch.nn.Module):
+        state_dict = state_dict.state_dict()
+    _, spec = _tools()
+    weights = []
+    for name, shape in spec.WEIGHTS:
+        if name not in state_dict:
+            raise KeyError(f"the state dict has no {name}: is it the Wan2.2-TI2V-5B VAE's?")
+        w = state_dict[name]
+        if tuple(w.shape) != shape:
+            raise ValueError(f"{name} is {tuple(w.shape)}, expected {shape}")
+        weights.append(w.detach().to(device=device, dtype=torch.bfloat16).contiguous())
+    return weights
+
+
+def _destroy_if_initialized():
+    # The group may be gone already (a caller destroyed it), and destroying it twice raises.
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def _default_group():
+    if not dist.is_initialized():
+        dist.init_process_group("nccl", store=dist.HashStore(), rank=0, world_size=1)
+        atexit.register(_destroy_if_initialized)
+    return dist.group.WORLD
+
+
+class WanVAEDecoder:
+    """The decoder with its weights prepared once (the layout and the Winograd transforms the kernels read)."""
+
+    def __init__(self, state_dict, device=None, pg=None, local=False, frames_per_pass=4):
+        if frames_per_pass < 1:
+            raise ValueError(f"frames_per_pass must be at least 1, not {frames_per_pass}")
+        self.frames_per_pass = int(frames_per_pass)
+        self.device = torch.device(device if device is not None else "cuda")
+        if self.device.index is None:
+            self.device = torch.device("cuda", torch.cuda.current_device())
+        with torch.cuda.device(self.device):
+            if local:
+                if pg is not None:
+                    raise ValueError("local=True decodes on this GPU alone; it takes no process group")
+                # At world 1 the group is only carried, never used: any initialized group serves.
+                self.pg = dist.group.WORLD if dist.is_initialized() else _default_group()
+                self.rank, self.world = 0, 1
+            else:
+                self.pg = pg if pg is not None else _default_group()
+                self.rank, self.world = dist.get_rank(self.pg), dist.get_world_size(self.pg)
+            if self.world > 8:
+                raise ValueError(f"the decoder splits over at most 8 GPUs, not {self.world}")
+            self.weights = _module().prepare(decoder_weights(state_dict, self.device), self.rank, self.world,
+                                             self.pg)
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor | None:
+        """The video of latents z [1, 48, T, 44, 78] on rank 0, None on the other ranks."""
+        if z.dim() != 5 or z.shape[0] != 1 or z.shape[1] != LATENT_CHANNELS or tuple(z.shape[3:]) != LATENT_HW:
+            raise ValueError(f"latents {tuple(z.shape)}: expected [1, 48, T, 44, 78]")
+        z = z.to(device=self.device, dtype=torch.bfloat16).contiguous()
+        frames = 4 * z.shape[2] - 3
+        with torch.cuda.device(self.device):
+            out = torch.empty((1, 3, frames, *VIDEO_HW) if self.rank == 0 else (0,), dtype=torch.bfloat16,
+                              device=self.device)
+            _module().run_with(z, self.weights, self.rank, self.world, self.pg, out, self.frames_per_pass)
+        return out if self.rank == 0 else None
@@ -24,6 +24,17 @@ class PAN2PipelineConfig(PipelineConfig):

    dit_precision: str = "bf16"
    vae_precision: str = "fp32"
+    # Decode the video with fastvideo-kernel's Wan2.2 decoder (B200 / GB200 / B300; a 704 x 1248 video, batch 1): bf16
+    # with fp32 accumulation, about 7x faster than the fp32 decode on a B300, at PSNR 61-63 dB and SSIM 0.998 against
+    # it. Other devices and sizes, and False, keep the fp32 decode.
+    vae_decode_kernel: bool = True
+    # With vae_decode_kernel, split the decode over the sequence-parallel GPUs (2, 4 or 8) rather than decoding the
+    # whole video on each of them. On B300s a 20 s video decodes in 3.0 s on one GPU, 1.55 s on 2 (the same video, bit
+    # for bit), 0.79 s on 4 and 0.47 s on 8 (rounded differently, at the same PSNR against the fp32 decode).
+    vae_decode_kernel_parallel: bool = True
+    # Run the DiT's fused Triton kernels (Q/K RMSNorm + RoPE, LayerNorm + modulation): 21% less denoising time on 8
+    # B300s, as far from the eager path's samples as bf16 rounding differences carry them.
+    dit_fused_kernels: bool = True
    # Unset: the checkpoint's scheduler config decides, which is shift 9 for IFM/PAN2-TEST, the reference sampler's
    # fixed value. A fixed default here would override the released shift with whatever a test checkpoint declares.
    flow_shift: float | None = None
@@ -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:
@@ -1041,10 +1066,7 @@ class VideoGenerator:
            # `samples` above is just the pinned-CPU mirror of `output_batch.output` (`samples.copy_(output)` or
            # `output.cpu()`) with no intervening preprocessing, so reading `output_batch.output` here is the same
            # data. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback.
-            vid_u8 = _quantize_video_frames_to_uint8(output_batch.output)
-            frames = [
-                torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8
-            ]
+            frames = _video_frames(output_batch.output)
        postprocess_time = time.perf_counter() - postprocess_start
        logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
        if logging_info is not None:
@@ -635,6 +635,12 @@ FASTVIDEO_H3_VSA_FINE_ATTENTION = EnvChoice(
    "emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized "
    "BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over "
    "FASTVIDEO_H3_VSA_FP4. Empty keeps the native route.")
+FASTVIDEO_WAN22_VAE_DECODE_DIR = EnvPath("",
+                                         category="performance",
+                                         doc="Directory of the Wan2.2 VAE decoder kernels "
+                                         "(fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds "
+                                         "on first use, when its config sets decode_kernel. Empty takes the copy in "
+                                         "this checkout.")
 FASTVIDEO_H3_VSA_HEADS_FIRST_TILE = EnvBool(False,
                                            category="performance",
                                            doc="VSA-H3 64/128-token tiles: scatter rows straight into the heads-first "
@@ -924,6 +930,10 @@ FASTVIDEO_TEST_SD35_MODEL_DIR = EnvStr("stabilityai/stable-diffusion-3.5-medium"
                                       category="test",
                                       doc="Model for the Stable Diffusion 3.5 SSIM test.",
                                       deprecated_names=("SD35_MODEL_DIR", ))
+FASTVIDEO_TEST_PAN2_MODEL_DIR = EnvPath(None,
+                                        category="test",
+                                        doc="PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU "
+                                        "tests; unset skips them.")
 FASTVIDEO_TEST_TAEH3_REFERENCE_DIR = EnvStr(None,
                                            category="test",
                                            doc="Upstream taehv checkout for the MLX TAEH3 parity test.",
@@ -16,16 +16,26 @@ 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_world_size
+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]``."""

@@ -102,6 +112,9 @@ class PAN2TokenRefinerBlock(nn.Module):
 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")
@@ -111,12 +124,18 @@ class PAN2AdaLayerNormZero(nn.Module):
                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)
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        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")
@@ -125,6 +144,8 @@ class PAN2AdaLayerNormOut(nn.Module):
    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, :]


@@ -132,9 +153,13 @@ 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.
+    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,
@@ -172,16 +197,25 @@ class PAN2Attention(nn.Module):
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
-        freqs_cis: tuple[torch.Tensor, 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)
-        query = self.norm_q(query).to(value.dtype)
-        key = self.norm_k(key).to(value.dtype)
-
        encoder_query, encoder_key, encoder_value = self._project(self.add_qkv_proj, encoder_hidden_states)
-        encoder_query = self.norm_added_q(encoder_query).to(encoder_value.dtype)
-        encoder_key = self.norm_added_k(encoder_key).to(encoder_value.dtype)
+        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,
@@ -198,6 +232,9 @@ class PAN2Attention(nn.Module):

 class PAN2TransformerBlock(nn.Module):

+    # Set by PAN2Transformer3DModel.set_fused_kernels.
+    fused_kernels = False
+
    def __init__(self,
                 dim: int,
                 num_heads: int,
@@ -234,6 +271,7 @@ class PAN2TransformerBlock(nn.Module):
        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)
        (
@@ -244,11 +282,28 @@ class PAN2TransformerBlock(nn.Module):
            context_scale_mlp,
            context_gate_mlp,
        ) = self.scale_shift_table_context.to(temb.dtype).chunk(6, dim=-1)
-        norm_encoder_hidden_states = (self.norm1_context(encoder_hidden_states) *
-                                      (1 + context_scale_msa) + context_shift_msa)
+        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)
+                                                     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

@@ -329,8 +384,34 @@ class PAN2Transformer3DModel(BaseDiT):
                                         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."""
@@ -366,6 +447,9 @@ class PAN2Transformer3DModel(BaseDiT):

        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)
@@ -379,10 +463,10 @@ class PAN2Transformer3DModel(BaseDiT):
        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)
+                    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)
+                                                             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)
@@ -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",
+]
@@ -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):
@@ -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
@@ -39,7 +39,17 @@ class PAN2Pipeline(ComposedPipelineBase):
                       stage=PAN2DenoisingStage(transformer=self.get_module("transformer"),
                                                scheduler=self.get_module("scheduler"),
                                                pipeline=self))
-        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
+        # The performance switches of PAN2PipelineConfig: the DiT's fused kernels and the VAE's decode kernel.
+        pipeline_config = getattr(fastvideo_args, "pipeline_config", None)
+        transformer = self.get_module("transformer")
+        if hasattr(transformer, "set_fused_kernels"):
+            transformer.set_fused_kernels(bool(getattr(pipeline_config, "dit_fused_kernels", False)))
+        vae = self.get_module("vae")
+        vae_config = getattr(vae, "config", None)
+        if hasattr(vae_config, "decode_kernel"):
+            vae_config.decode_kernel = bool(getattr(pipeline_config, "vae_decode_kernel", False))
+            vae_config.decode_kernel_parallel = bool(getattr(pipeline_config, "vae_decode_kernel_parallel", False))
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=vae, pipeline=self))


 class PAN2ModularPipeline(PAN2Pipeline):
@@ -10,7 +10,8 @@ _DENOISE_STAGE = PresetStageSpec(
    allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}),
 )

-# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 3. The preset sends no negative prompt (an empty one is encoded
+# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 2.5 (the release launcher's default since ifm-ai/PAN2
+# 86d5023, 2026-10-23; 3 before). The preset sends no negative prompt (an empty one is encoded
 # as a single space); the reference sampler sends the release's structured negative caption, so pass it as the
 # request's ``negative_prompt`` to follow that recipe.
 _SHARED_DEFAULTS = {
@@ -19,7 +20,7 @@ _SHARED_DEFAULTS = {
    "num_frames": 313,
    "fps": 24,
    "num_inference_steps": 50,
-    "guidance_scale": 3.0,
+    "guidance_scale": 2.5,
    "negative_prompt": "",
 }

@@ -33,6 +34,7 @@ PAN2_T2V = InferencePreset(
    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,
@@ -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)
@@ -65,3 +65,26 @@ def test_pan2_pipeline_requests_fp32_latents(monkeypatch):
    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)
+
+    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
@@ -0,0 +1,292 @@
+# 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])
+
+
+# ---- references against the eager arithmetic -----------------------------------------------------------------------
+
+
+@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)
+
+
+# ---- the Triton kernels against the references ----------------------------------------------------------------------
+
+
+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}"
+
+
+# ---- the fused model against the eager model -----------------------------------------------------------------------
+
+
+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}"
@@ -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)
@@ -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=<PAN2 root> 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)
				`@@ -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":[]}`