[misc]: add CPU regression tests for the on-device uint8 frame path

Two tests certifying #1362's frame semantics without a GPU:

- frames_match_legacy_cpu_loop: the quantize-then-grid path reproduces
  the legacy make_grid -> *255 -> uint8 per-frame loop bit-exactly for
  in-range fp32 pixels (batch>1 nrow=6 grid layout, odd frame count,
  uint8 HWC contract). Verified to pass against main's legacy loop too,
  so it pins both sides of the equivalence.
- frames_clamp_out_of_range_pixels: out-of-[0,1] VAE output saturates
  at 0/255 instead of wrapping mod 256; fails on the pre-#1362 loop.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

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

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

README.md

@@ -9,6 +9,7 @@
 **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**
 
 ## NEWS
+- `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac—follow the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/) and read the [Blog](https://haoailab.com/blogs/fastmetal/).
 - `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. See the [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), [Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/), and [blog](https://haoailab.com/blogs/fastwan-qad/).
 - `2026/03/17`: Release demo: Into the Dreamverse: Vibe Directing in FastVideo, check out the [Blog](https://haoailab.com/blogs/dreamverse/).
 - `2026/03/13`: Release demo: Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU, check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
@@ -62,6 +63,11 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo
 Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the
 [MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
 
+> **On an Apple Silicon Mac?** FastVideo runs FastWan text-to-video natively
+> through an MLX runtime — a 5-second 480p clip generated locally, no cloud,
+> no discrete GPU. Install with `uv pip install -e '.[mlx]'` and follow the
+> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
+
 Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
 
 > **On an NVIDIA DGX Spark (GB10 / ARM64 + CUDA 13)?** There's no prebuilt ARM wheel for the FastVideo CUDA kernel, so it's an editable from-source install (`UV_TORCH_BACKEND=cu130 uv pip install -e .`, which compiles that kernel for you) rather than `UV_TORCH_BACKEND=cu130 uv pip install fastvideo`. A compatible prebuilt ARM64 FlashAttention wheel is available separately. Follow the [DGX Spark install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/spark/).

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without

docs/design/apple_silicon_fast_mode.md

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

docs/getting_started/installation/mps.md

@@ -49,10 +49,12 @@ brew install ffmpeg
 
 ### Installation
 
+FastWan's native Apple Silicon runtime requires the `mlx` extra.
+
 #### With uv (recommended)
 
 ```bash
-uv pip install fastvideo
+uv pip install "fastvideo[mlx]"
 ```
 
 #### With Conda environment (alternative)
@@ -60,7 +62,7 @@ uv pip install fastvideo
 `uv` works inside an active conda env too, so prefer `uv pip` for the actual install:
 
 ```bash
-uv pip install fastvideo
+uv pip install "fastvideo[mlx]"
 ```
 
 ### Installation from Source
@@ -76,13 +78,13 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 Basic installation:
 
 ```bash
-uv pip install -e .
+uv pip install -e ".[mlx]"
 ```
 
 Alternative with Conda environment:
 
 ```bash
-uv pip install -e .
+uv pip install -e ".[mlx]"
 ```
 
 ## Development Environment Setup

docs/inference/support_matrix.md

@@ -178,6 +178,17 @@ optimizations: absence means **untested**, not incompatible.
 | Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |
 
+## Apple Silicon native runtime
+
+| Release path | Model | Mode | Validated hardware | Status |
+| --- | --- | --- | --- | --- |
+| MLX FastWan T2V | FastWan-QAD-INT8-1.3B `[release model ID pending]` | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB unified-memory class, MLX 0.31.2 | Release candidate; requires release-owner visual sign-off |
+
+This is a text-to-video-only source-install release. It is validated on the
+hardware listed above; MLX allocator caps are not evidence of support for a
+physical 16 GB Mac. See [Apple Silicon FastWan](../getting_started/installation/mps.md)
+for the supported command and release gates.
+
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 
 ***Lucy Edit Dev uses a non-commercial model license. FastVideo support is

examples/inference/basic/mlx_rife_smoke_test.py

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

examples/inference/basic/mlx_wan22_generate.py

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

examples/inference/basic/mlx_wan_decode_benchmark.py

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

examples/inference/basic/mlx_wan_quant_benchmark.py

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

examples/inference/basic/mlx_wan_video_quality.py

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

fastvideo-kernel/README.md

@@ -17,7 +17,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
 | Kernels | Where | Used when |
 |---|---|---|
 | Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
-| FA4 CuTe-DSL block-sparse (VSA-256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
+| FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
 | VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |
 
 ## What gets built where, and when
@@ -64,28 +64,49 @@ cd fastvideo-kernel
 ./build.sh --rocm
 ```
 
-### Optional: FA4 CuTe block-sparse backend (VSA-256 fastpath)
+### Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath)
 
-The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes to the
+The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the
 FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is
 an **optional** dependency: it is imported lazily, and `video_sparse_attn`
 transparently falls back to the Triton backend when it is absent (so the package is
 fully usable without it).
 
-The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`,
-`flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by
+The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`
+and the public/private forward-backward bridges in `flash_attn.cute.interface`) are provided upstream by
 [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
-commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
+commit `14c377950125c70b7a9dabf9c561fca53715ac7d`, the revision FastVideo pins as
 the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
-have an incompatible `_flash_attn_fwd` signature.
+have incompatible block-sparse forward/backward interfaces.
+
+Install it under its distribution name so its own runtime stack resolves with it.
+Do **not** pre-install `nvidia-cutlass-dsl` by hand: this revision pins
+`nvidia-cutlass-dsl==4.6.0.dev0` exactly, and a hand-installed 4.5.x floor either
+gets silently upgraded or, if something else holds it back, leaves the CuTe
+kernels broken.
 
 ```bash
-pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
-pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
+pip install torchvision
+pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@14c377950125c70b7a9dabf9c561fca53715ac7d#subdirectory=flash_attn/cute"
 ```
 
-The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against
-`tests/test_vsa256_forward*.py`.
+That resolves `nvidia-cutlass-dsl` to 4.6.0.dev0 and `quack-kernels` to 0.5.3, a
+combination this revision works with. A mismatched CuTe DSL only surfaces when the
+kernel JIT-compiles, so the error points at CuTe internals rather than at the
+install:
+
+| Error on first VSA-128/256 CuTe call | Cause |
+|---|---|
+| `TypeError: fmax() missing 1 required positional argument: 'b'` | `nvidia-cutlass-dsl` 4.5.x |
+| `AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'` | `quack-kernels` older than 0.5.1 |
+| `ImportError: cannot import name 'alloc_reserved_mbarrier'` | `quack-kernels` 0.6.2 or newer |
+
+An environment whose `flash_attn.cute` came from a prebuilt flash-attn wheel rather
+than from this pin hits the first row; that is what the overlay step in
+`docker/Dockerfile` works around.
+
+The CuTe kernels JIT-compile on first use. Forward and backward are verified on
+Blackwell (sm_100) against `tests/test_vsa128_*.py` and `tests/test_vsa256_*.py`.
 
 ## Usage
 
@@ -142,6 +163,14 @@ After building/installing `fastvideo-kernel`, run:
 ```bash
 cd fastvideo-kernel
 python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
+
+# VSA-256 FA4 CuTe forward/backward on Blackwell
+python benchmarks/bench_vsa.py --block_size 256 --use_cute \
+  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 20
+
+# VSA-128 FA4 CuTe forward/backward on Blackwell
+python benchmarks/bench_vsa.py --block_size 128 --use_cute \
+  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 40
 ```
 
 ### TurboDiffusion Kernels

fastvideo-kernel/benchmarks/bench_vsa.py

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

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py

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

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py

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

fastvideo-kernel/python/fastvideo_kernel/ops.py

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

fastvideo-kernel/tests/test_vsa128_backward.py

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

fastvideo-kernel/tests/test_vsa256_backward.py

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

fastvideo-kernel/tests/test_vsa256_triton.py

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

fastvideo-kernel/tests/test_vsa_varlen.py

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

fastvideo/attention/backends/video_sparse_attn_h3.py

@@ -319,8 +319,13 @@ class MiniMaxH3VSAImpl(AttentionImpl):
             v_pooled = _pool_tiles(value, attn_metadata.variable_block_sizes)
             out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)  # [B, H, n_tiles, D]
             out_c = out_c.permute(0, 2, 1, 3).to(out.dtype)  # [B, n_tiles, H, D]
-            batch, _, heads, dim = out.shape
+            batch, seq_len, heads, dim = out.shape
             n_tiles = attn_metadata.variable_block_sizes.numel()
-            out.view(batch, n_tiles, _TILE_ELEMS, heads,
-                     dim).addcmul_(out_c.unsqueeze(2), gate_compress.view(batch, n_tiles, _TILE_ELEMS, heads, dim))
+            # Out-of-place: on the CuTe backend ``out`` is the tensor FA4's
+            # autograd node saved for its backward, so an in-place add here
+            # bumps its version counter and backward dies with "one of the
+            # variables needed for gradient computation has been modified".
+            out_tiled = out.view(batch, n_tiles, _TILE_ELEMS, heads, dim)
+            gate_tiled = gate_compress.view(batch, n_tiles, _TILE_ELEMS, heads, dim)
+            out = (out_tiled + out_c.unsqueeze(2) * gate_tiled).view(batch, seq_len, heads, dim)
         return out

fastvideo/benchmarks/__init__.py

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

fastvideo/benchmarks/eval_metalfx_rife.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Benchmark generate-fewer-frames plus MLX RIFE interpolation.
+
+This script keeps the video diffusion path delegated to
+``fastvideo.benchmarks.mlx_fastwan_bench._generate_cell``. It only orchestrates
+two frame-count cells and the postprocess interpolation step.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import time
+from pathlib import Path
+from types import SimpleNamespace
+
+import imageio.v2 as imageio
+import numpy as np
+
+from examples.inference.basic.mlx_wan_prompt_to_video import encode_prompt, make_rotary_embeddings
+from fastvideo.benchmarks.mlx_fastwan_bench import _generate_cell, _ms_ssim
+from fastvideo.mlx_runtime.rife_interp import RIFEBackendError, interpolate, load_model
+from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits
+
+FOX_PROMPT = "A fox runs through a misty pine forest, leaves kicking up behind it."
+DEFAULT_MODEL_ROOT = Path("/Users/aryank/models/qad_int8_v2")
+
+
+def _parse_timesteps(raw: str) -> list[int]:
+    timesteps = [int(part.strip()) for part in raw.split(",") if part.strip()]
+    if not timesteps:
+        raise SystemExit("No DMD timesteps parsed from --dmd-denoising-steps")
+    return timesteps
+
+
+def _read_video(path: Path) -> list[np.ndarray]:
+    if not path.is_file():
+        raise FileNotFoundError(f"Video does not exist: {path}")
+    frames = [np.asarray(frame[:, :, :3], dtype=np.uint8) for frame in imageio.mimread(path)]
+    if not frames:
+        raise RuntimeError(f"No frames decoded from {path}")
+    return frames
+
+
+def _write_video(path: Path, frames: list[np.ndarray], fps: int) -> None:
+    if not frames:
+        raise ValueError("Cannot write an empty frame list")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with imageio.get_writer(str(path), fps=fps, macro_block_size=1, codec="libx264", quality=8) as writer:
+        for frame in frames:
+            writer.append_data(np.asarray(frame, dtype=np.uint8))
+
+
+def _copy_video(src: Path, dst: Path) -> Path:
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(src, dst)
+    return dst
+
+
+def _make_generation_inputs(args, num_frames: int):
+    import mlx.core as mx
+    import torch
+
+    config_path = args.model_root / "transformer" / "config.json"
+    checkpoint_path = args.model_root / "transformer" / "diffusion_pytorch_model.safetensors"
+    if not config_path.is_file():
+        raise SystemExit(f"Missing DiT config: {config_path}")
+    if not checkpoint_path.is_file():
+        raise SystemExit(f"Missing DiT checkpoint: {checkpoint_path}")
+
+    config = json.loads(config_path.read_text())
+    latent_frames = (num_frames - 1) // 4 + 1
+    latent_height = args.height // 8
+    latent_width = args.width // 8
+    freqs_cis = make_rotary_embeddings(
+        config,
+        latent_frames=latent_frames,
+        latent_height=latent_height,
+        latent_width=latent_width,
+    )
+
+    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    latents_seed = torch.randn(
+        (1, int(config["in_channels"]), latent_frames, latent_height, latent_width),
+        generator=generator,
+        dtype=torch.float32,
+    ).numpy()
+    timesteps = _parse_timesteps(args.dmd_denoising_steps)
+    renoise_by_step = [
+        torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy()
+        for _ in range(max(0,
+                           len(timesteps) - 1))
+    ]
+    prompt_embeds = encode_prompt(
+        model_root=args.model_root,
+        prompt=args.prompt,
+        max_sequence_length=args.max_sequence_length,
+        device_arg=args.torch_device,
+        dtype_arg=args.torch_dtype,
+    )
+    return {
+        "checkpoint_path": checkpoint_path,
+        "config_path": config_path,
+        "encoder_hidden_states": mx.array(prompt_embeds.numpy()),
+        "freqs_cis": freqs_cis,
+        "timesteps": timesteps,
+        "latents_seed": latents_seed,
+        "renoise_by_step": renoise_by_step,
+        "latent_frames": latent_frames,
+    }
+
+
+def _generate(args, num_frames: int):
+    cell_args = SimpleNamespace(
+        model_root=args.model_root,
+        height=args.height,
+        width=args.width,
+        num_frames=num_frames,
+        fps=args.fps,
+        flow_shift=args.flow_shift,
+        torch_device=args.torch_device,
+        torch_dtype=args.torch_dtype,
+        taehv_source_path=args.taehv_source_path,
+        taehv_checkpoint_path=args.taehv_checkpoint_path,
+        taehv_parallel=args.taehv_parallel,
+        mlx_checkpoint_cache=args.mlx_checkpoint_cache,
+        mlx_memory_limit_gib=args.mlx_memory_limit_gib,
+        mlx_cache_limit_gib=args.mlx_cache_limit_gib,
+        mlx_disable_cache=args.mlx_disable_cache,
+        mlx_wired_limit_gib=args.mlx_wired_limit_gib,
+        torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio,
+        torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio,
+        benchmark_preset="metalfx-rife",
+        current_prompt_id=f"fox-{num_frames}f",
+        current_prompt=args.prompt,
+        output_dir=args.output_dir,
+    )
+    inputs = _make_generation_inputs(args, num_frames)
+    print(f"=== generate: {num_frames} frames ({inputs['latent_frames']} latent frames) ===", flush=True)
+    return _generate_cell(
+        args=cell_args,
+        mode=args.mode,
+        decoder=args.decoder,
+        checkpoint_path=inputs["checkpoint_path"],
+        config_path=inputs["config_path"],
+        encoder_hidden_states=inputs["encoder_hidden_states"],
+        freqs_cis=inputs["freqs_cis"],
+        timesteps=inputs["timesteps"],
+        latents_seed=inputs["latents_seed"],
+        renoise_by_step=inputs["renoise_by_step"],
+    )
+
+
+def _relative_to_output(path: Path, output_dir: Path) -> str:
+    try:
+        return str(path.relative_to(output_dir))
+    except ValueError:
+        return str(path)
+
+
+def _write_report(args, result: dict) -> Path:
+    report_path = args.output_dir / "metalfx_rife_report.md"
+    rows = result["speed_rows"]
+    table_lines = [
+        "| path | denoise_s | decode_s | gen_total_s | rife_s | net_s | speedup_vs_81 |",
+        "| --- | ---: | ---: | ---: | ---: | ---: | ---: |",
+    ]
+    for row in rows:
+        table_lines.append(
+            "| {path} | {denoise_s:.3f} | {decode_s:.3f} | {gen_total_s:.3f} | {rife_s:.3f} | {net_s:.3f} | {speedup_vs_81:.3f}x |"
+            .format(**row))
+    text = f"""# MetalFX/RIFE Generate-Fewer-Frames Benchmark Run
+
+Prompt: {args.prompt}
+
+Resolution: {args.height}x{args.width}, fps={args.fps}, mode={args.mode}, decoder={args.decoder}
+
+| metric | value |
+| --- | ---: |
+| reconstruction_ms_ssim | {result['reconstruction_ms_ssim']:.6f} |
+| reference_frames | {result['reference_frames']} |
+| reduced_frames | {result['reduced_frames']} |
+| interpolated_frames | {result['interpolated_frames']} |
+
+{chr(10).join(table_lines)}
+
+Videos:
+
+- reference: `{result['videos']['reference']}`
+- reference drop-41 RIFE reconstruction: `{result['videos']['drop41_rife81']}`
+- generated 41 RIFE to 81: `{result['videos']['generated41_rife81']}`
+- generated 41 direct: `{result['videos']['generated41']}`
+
+Raw metrics are in `{_relative_to_output(args.output_dir / 'metrics.json', args.output_dir)}`.
+"""
+    report_path.write_text(text)
+    return report_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Evaluate 41-frame generation plus MLX RIFE interpolation vs 81-frame generation.")
+    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
+    parser.add_argument("--prompt", default=FOX_PROMPT)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=832)
+    parser.add_argument("--reference-frames", type=int, default=81)
+    parser.add_argument("--reduced-frames", type=int, default=41)
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--mode", default="int8", choices=("fp16", "bf16", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"))
+    parser.add_argument("--decoder", default="taehv", choices=("taehv", "wan-vae"))
+    parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
+    parser.add_argument("--flow-shift", type=float, default=8.0)
+    parser.add_argument("--max-sequence-length", type=int, default=512)
+    parser.add_argument("--torch-device", default="auto")
+    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
+    parser.add_argument("--output-dir", type=Path, default=Path("bench/metalfx_rife"))
+    parser.add_argument("--mlx-checkpoint-cache", type=Path, default=None)
+    parser.add_argument("--compile", action="store_true", help="Enable FASTVIDEO_MLX_COMPILE=1 for DiT denoise.")
+    parser.add_argument("--rife-scale", type=float, default=1.0)
+    parser.add_argument("--taehv-source-path", type=Path, default=None)
+    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
+    parser.add_argument("--taehv-parallel", action="store_true")
+    add_memory_limit_args(parser)
+    args = parser.parse_args()
+
+    if args.reference_frames != (args.reduced_frames - 1) * 2 + 1:
+        raise SystemExit(
+            "--reference-frames must equal (--reduced-frames - 1) * 2 + 1 for the default every-other-frame test")
+    if args.compile:
+        import os
+
+        os.environ["FASTVIDEO_MLX_COMPILE"] = "1"
+
+    args.model_root = args.model_root.expanduser().resolve()
+    args.output_dir = args.output_dir.expanduser().resolve()
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    if args.mlx_checkpoint_cache is None:
+        args.mlx_checkpoint_cache = args.output_dir / "mlx_checkpoint_cache"
+    else:
+        args.mlx_checkpoint_cache = args.mlx_checkpoint_cache.expanduser().resolve()
+
+    import mlx.core as mx
+    import torch
+
+    runtime_limits = apply_memory_limits(
+        mlx_memory_limit_gib=args.mlx_memory_limit_gib,
+        mlx_cache_limit_gib=args.mlx_cache_limit_gib,
+        mlx_disable_cache=args.mlx_disable_cache,
+        mlx_wired_limit_gib=args.mlx_wired_limit_gib,
+        torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio,
+        torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio,
+        mx_module=mx,
+    ).as_metrics()
+    mx.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    reference_cell = _generate(args, args.reference_frames)
+    reduced_cell = _generate(args, args.reduced_frames)
+
+    reference_video = _copy_video(reference_cell.video_path, args.output_dir / "fox_reference_81.mp4")
+    generated41_video = _copy_video(reduced_cell.video_path, args.output_dir / "fox_generated_41.mp4")
+
+    reference_frames = _read_video(reference_video)
+    dropped_reference_frames = reference_frames[::2]
+    if len(dropped_reference_frames) != args.reduced_frames:
+        raise RuntimeError(f"Expected {args.reduced_frames} dropped frames, got {len(dropped_reference_frames)}")
+
+    try:
+        rife_model = load_model("4.25")
+    except RIFEBackendError:
+        raise
+    except Exception as exc:  # noqa: BLE001 - preserve exact backend failure.
+        raise RIFEBackendError(f"Unexpected RIFE load failure: {exc}") from exc
+
+    print("=== RIFE reconstruction: drop every other reference frame, 41 -> 81 ===", flush=True)
+    start = time.perf_counter()
+    reconstructed_frames = interpolate(dropped_reference_frames, factor=2, model=rife_model, scale=args.rife_scale)
+    recon_rife_s = time.perf_counter() - start
+    reconstructed_video = args.output_dir / "fox_reference_drop41_rife81.mp4"
+    _write_video(reconstructed_video, reconstructed_frames, args.fps)
+    reconstruction_ms_ssim = _ms_ssim(reference_video, reconstructed_video, required=True)
+    if reconstruction_ms_ssim is None:
+        raise RuntimeError("MS-SSIM returned None for reconstruction comparison")
+
+    print("=== RIFE actual reduced path: generated 41 -> 81 ===", flush=True)
+    generated41_frames = _read_video(generated41_video)
+    start = time.perf_counter()
+    generated41_rife_frames = interpolate(generated41_frames, factor=2, model=rife_model, scale=args.rife_scale)
+    generated41_rife_s = time.perf_counter() - start
+    generated41_rife_video = args.output_dir / "fox_generated41_rife81.mp4"
+    _write_video(generated41_rife_video, generated41_rife_frames, args.fps)
+
+    direct_denoise_s = float(reference_cell.metrics["denoise_s"])
+    reduced_denoise_s = float(reduced_cell.metrics["denoise_s"])
+    direct_gen_total_s = float(reference_cell.metrics["denoise_s"]) + float(reference_cell.metrics["decode_s"])
+    reduced_gen_total_s = float(reduced_cell.metrics["denoise_s"]) + float(reduced_cell.metrics["decode_s"])
+    net_denoise_rife_s = reduced_denoise_s + generated41_rife_s
+    net_gen_rife_s = reduced_gen_total_s + generated41_rife_s
+
+    result = {
+        "prompt":
+        args.prompt,
+        "model_root":
+        str(args.model_root),
+        "rife_impl":
+        "rife-mlx vendored at fastvideo/third_party/rife_mlx, weights mlx-community/RIFE-4.25",
+        "runtime_limits":
+        runtime_limits,
+        "reference_frames":
+        len(reference_frames),
+        "reduced_frames":
+        len(generated41_frames),
+        "interpolated_frames":
+        len(generated41_rife_frames),
+        "reconstruction_ms_ssim":
+        reconstruction_ms_ssim,
+        "reconstruction_rife_s":
+        recon_rife_s,
+        "generated41_rife_s":
+        generated41_rife_s,
+        "reference_metrics":
+        reference_cell.metrics,
+        "reduced_metrics":
+        reduced_cell.metrics,
+        "speed_rows": [
+            {
+                "path": "generate_81",
+                "denoise_s": direct_denoise_s,
+                "decode_s": float(reference_cell.metrics["decode_s"]),
+                "gen_total_s": direct_gen_total_s,
+                "rife_s": 0.0,
+                "net_s": direct_gen_total_s,
+                "speedup_vs_81": 1.0,
+            },
+            {
+                "path": "generate_41_plus_rife81_denoise_only",
+                "denoise_s": reduced_denoise_s,
+                "decode_s": 0.0,
+                "gen_total_s": reduced_denoise_s,
+                "rife_s": generated41_rife_s,
+                "net_s": net_denoise_rife_s,
+                "speedup_vs_81": direct_denoise_s / net_denoise_rife_s,
+            },
+            {
+                "path": "generate_41_plus_rife81_decode_included",
+                "denoise_s": reduced_denoise_s,
+                "decode_s": float(reduced_cell.metrics["decode_s"]),
+                "gen_total_s": reduced_gen_total_s,
+                "rife_s": generated41_rife_s,
+                "net_s": net_gen_rife_s,
+                "speedup_vs_81": direct_gen_total_s / net_gen_rife_s,
+            },
+        ],
+        "videos": {
+            "reference": _relative_to_output(reference_video, args.output_dir),
+            "drop41_rife81": _relative_to_output(reconstructed_video, args.output_dir),
+            "generated41": _relative_to_output(generated41_video, args.output_dir),
+            "generated41_rife81": _relative_to_output(generated41_rife_video, args.output_dir),
+        },
+    }
+    metrics_path = args.output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(result, indent=2))
+    report_path = _write_report(args, result)
+
+    print(json.dumps(result["speed_rows"], indent=2))
+    print(f"reconstruction_ms_ssim={reconstruction_ms_ssim:.6f}")
+    print(f"wrote {metrics_path}")
+    print(f"wrote {report_path}")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/benchmarks/mlx_fastwan_bench.py

@@ -0,0 +1,826 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Prove-out benchmark for the MLX FastWan runtime (Apple Silicon).
+
+Sweeps ``{dtype/quant} x {decoder}``, generates a clip per cell, and records the
+latency breakdown, peak unified memory, and MS-SSIM (optionally LPIPS) against a
+reference video. It emits a JSON blob and a markdown table -- the artifact that
+turns "int8 + TAEHV looks good" into defensible numbers, and (via
+``--assert-min-ssim``) a regression gate for the ``mx.compile`` work.
+
+Design notes:
+- Generation reuses the hybrid POC helpers in
+  ``examples/inference/basic/mlx_wan_prompt_to_video.py`` (torch-MPS UMT5 encode
+  and Wan-VAE/TAEHV decode) plus the on-device MLX DMD sampler
+  (``fastvideo/mlx_runtime/sampling.py``); the denoise loop never leaves the
+  device.
+- Quality reuses the tested MS-SSIM primitive
+  ``fastvideo/tests/utils.py::compute_video_ssim_torchvision``.
+- Reference: by default each cell is scored against the highest-fidelity cell
+  in the sweep (``fp16`` + ``wan-vae``), which needs no CUDA box and answers
+  "how much does int8/int4/TAEHV degrade vs the best local config". Pass
+  ``--reference PATH`` to score against an external clip instead (e.g. the
+  torch-MPS or CUDA FastVideo output of the same model) for a "vs. the original
+  model" column.
+
+Run on an Apple Silicon Mac (needs ``mlx`` + a torch build with MPS):
+
+    python fastvideo/benchmarks/mlx_fastwan_bench.py \
+        --modes fp16,bf16,int8,int4 --decoders taehv,wan-vae
+"""
+
+from __future__ import annotations
+
+import argparse
+import html
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+from examples.inference.basic.mlx_wan_prompt_to_video import (
+    DEFAULT_MODEL_ROOT,
+    decode_latents_to_video,
+    encode_prompt,
+    make_rotary_embeddings,
+)
+from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits, cleanup_mlx
+
+# The highest-fidelity cell; used as the default SSIM reference when no external
+# reference video is supplied.
+REFERENCE_MODE = "fp16"
+REFERENCE_DECODER = "wan-vae"
+
+ALLOWED_MODES = ("fp16", "bf16", "int8", "int4", "mxfp8", "mxfp4", "nvfp4")
+ALLOWED_DECODERS = ("taehv", "wan-vae")
+
+
+@dataclass(frozen=True)
+class PromptCase:
+    id: str
+    prompt: str
+
+
+@dataclass(frozen=True)
+class BenchmarkPreset:
+    height: int
+    width: int
+    num_frames: int
+    modes: str
+    decoders: str
+    mlx_memory_limit_gib: float | None = None
+    mlx_disable_cache: bool = False
+    torch_mps_high_watermark_ratio: float | None = None
+    torch_mps_low_watermark_ratio: float | None = None
+
+
+PROMPT_SETS = {
+    "motion7": (
+        PromptCase("beach-sunset", "A slow cinematic sunset over ocean waves at a quiet beach."),
+        PromptCase("fox-forest", "A fox runs through a misty pine forest, leaves kicking up behind it."),
+        PromptCase("raccoon-sunflowers", "A raccoon walks through a sunflower field as petals move in the wind."),
+        PromptCase("surfing-cat", "A cat wearing sunglasses surfs across a bright blue ocean wave."),
+        PromptCase("burning-clock", "A vintage table clock burns on a wooden desk, flames flickering realistically."),
+        PromptCase("forest-walk", "Video game style footage of a man walking through a dense forest path."),
+        PromptCase("sea-dock-yachts", "Several yachts are parked at a sea dock while water ripples around them."),
+    ),
+}
+
+BENCHMARK_PRESETS = {
+    "default":
+    BenchmarkPreset(
+        height=480,
+        width=832,
+        num_frames=81,
+        modes="fp16,bf16,int8,int4",
+        decoders="taehv,wan-vae",
+    ),
+    "mac-16gb":
+    BenchmarkPreset(
+        height=448,
+        width=832,
+        num_frames=61,
+        modes="int8",
+        decoders="taehv",
+        mlx_memory_limit_gib=16.0,
+        mlx_disable_cache=True,
+        torch_mps_high_watermark_ratio=0.57,
+        torch_mps_low_watermark_ratio=0.0,
+    ),
+    "mac-32gb":
+    BenchmarkPreset(
+        height=480,
+        width=832,
+        num_frames=81,
+        modes="int8,fp16",
+        decoders="taehv",
+    ),
+    "mac-64gb":
+    BenchmarkPreset(
+        height=480,
+        width=832,
+        num_frames=81,
+        modes="int8,fp16",
+        decoders="taehv,wan-vae",
+    ),
+}
+
+
+@dataclass
+class Cell:
+    prompt_id: str
+    prompt: str
+    mode: str
+    decoder: str
+    video_path: Path
+    latents: np.ndarray
+    metrics: dict[str, float | int | str | bool | None] = field(default_factory=dict)
+
+
+def _mode_to_dtype_quant(mode: str) -> tuple[str, str | None]:
+    """Map a sweep mode to (MLX compute dtype, quantization spec).
+
+    Quantized modes keep fp16 activations and quantize only the DiT linear
+    weights (matching ``mlx_dit_from_diffusers_safetensors``).
+    """
+    if mode == "bf16":
+        return "bf16", None
+    if mode == "fp16":
+        return "fp16", None
+    # int8/int4/mxfp*/nvfp4 -> fp16 activations + quantized weights.
+    return "fp16", mode
+
+
+def _mx_dtype(mx, base: str):
+    return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[base]
+
+
+def _parse_list(raw: str, allowed: tuple[str, ...], label: str) -> list[str]:
+    items = [x.strip() for x in raw.split(",") if x.strip()]
+    unknown = sorted(set(items) - set(allowed))
+    if unknown:
+        raise ValueError(f"Unsupported {label}: {unknown} (allowed: {list(allowed)})")
+    return items
+
+
+def _safe_slug(value: str, *, fallback: str) -> str:
+    slug = "".join(ch.lower() if ch.isalnum() else "-" for ch in value.strip())
+    slug = "-".join(part for part in slug.split("-") if part)
+    return slug[:64] or fallback
+
+
+def _load_prompt_cases(prompt: str, prompt_file: Path | None, prompt_set: str = "single") -> list[PromptCase]:
+    """Load one prompt, a built-in prompt set, or a text/jsonl prompt file."""
+    if prompt_file is None:
+        if prompt_set == "single":
+            return [PromptCase(id="prompt-001", prompt=prompt)]
+        if prompt_set not in PROMPT_SETS:
+            raise ValueError(f"Unsupported prompt set: {prompt_set} (allowed: {sorted(PROMPT_SETS) + ['single']})")
+        return list(PROMPT_SETS[prompt_set])
+    cases: list[PromptCase] = []
+    for line_index, raw_line in enumerate(prompt_file.read_text().splitlines(), start=1):
+        line = raw_line.strip()
+        if not line or line.startswith("#"):
+            continue
+        prompt_id = f"prompt-{len(cases) + 1:03d}"
+        prompt_text = line
+        if prompt_file.suffix.lower() == ".jsonl":
+            item = json.loads(line)
+            prompt_text = str(item.get("prompt") or item.get("text") or item.get("caption") or "").strip()
+            if not prompt_text:
+                raise ValueError(f"{prompt_file}:{line_index} has no prompt/text/caption field")
+            prompt_id = str(item.get("id") or item.get("name") or prompt_id)
+        cases.append(PromptCase(id=_safe_slug(prompt_id, fallback=f"prompt-{len(cases) + 1:03d}"), prompt=prompt_text))
+    if not cases:
+        raise ValueError(f"No prompts found in {prompt_file}")
+    return cases
+
+
+def denoise_dmd_on_device(
+    *,
+    mx,
+    dit,
+    latents,
+    encoder_hidden_states,
+    freqs_cis,
+    timesteps: list[int],
+    renoise_by_step: list[np.ndarray],
+    schedule,
+    dmd_step,
+    mx_dtype,
+) -> tuple[np.ndarray, list[float]]:
+    """Run the FastWan DMD loop entirely on the MLX device.
+
+    Mirrors the loop in ``mlx_wan_prompt_to_video.py`` (fp32 affine math, MLX RNG
+    re-noise) so the benchmark measures exactly the shipped path.
+
+    Returns the final latents plus per-step wall times. The first step carries
+    one-time costs (mx.compile tracing, kernel warm-up), so first-vs-steady
+    step timing is how the benchmark separates cold-start from steady-state
+    denoise throughput.
+
+    All host-side tensors (timesteps, re-noise draws) are uploaded before the
+    loop starts, so the per-step body performs no bulk host->device transfers
+    and step timings measure device work rather than staging copies.
+    """
+    timesteps_mx = [mx.array([float(timestep)]).astype(mx.float32) for timestep in timesteps]
+    renoise_mx = [mx.array(renoise).astype(mx.float32) for renoise in renoise_by_step]
+    if timesteps_mx or renoise_mx:
+        mx.eval(*timesteps_mx, *renoise_mx)
+
+    step_times: list[float] = []
+    for step_index, timestep in enumerate(timesteps):
+        step_start = time.perf_counter()
+        noise_input_latent = latents
+        noise_pred = dit(latents.astype(mx_dtype), encoder_hidden_states, timesteps_mx[step_index], freqs_cis)
+
+        noise_input_f32 = noise_input_latent.astype(mx.float32)
+        pred_noise_f32 = noise_pred.astype(mx.float32)
+        if step_index < len(timesteps) - 1:
+            next_ts: float | None = float(timesteps[step_index + 1])
+            renoise = renoise_mx[step_index]
+        else:
+            next_ts, renoise = None, None
+        latents = dmd_step(
+            latents=noise_input_f32,
+            noise_input_latent=noise_input_f32,
+            pred_noise=pred_noise_f32,
+            schedule=schedule,
+            timestep=float(timestep),
+            next_timestep=next_ts,
+            noise=renoise,
+        ).astype(mx_dtype)
+        mx.eval(latents)
+        step_times.append(time.perf_counter() - step_start)
+    return np.array(latents.astype(mx.float32)), step_times
+
+
+def _peak_memory_bytes(mx) -> int:
+    try:
+        return int(mx.get_peak_memory())
+    except Exception:  # noqa: BLE001 - best-effort telemetry only.
+        return 0
+
+
+def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]:
+    diff = candidate.astype(np.float32) - baseline.astype(np.float32)
+    mse = float(np.mean(np.square(diff)))
+    signal = float(np.mean(np.square(baseline.astype(np.float32))))
+    return {
+        "latent_mse_vs_ref_mode": mse,
+        "latent_snr_db_vs_ref_mode": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"),
+    }
+
+
+def _ms_ssim(reference_video: Path, candidate_video: Path, *, required: bool = False) -> float | None:
+    """Mean MS-SSIM between two mp4s, via the repo's tested helper."""
+    if not reference_video.exists() or not candidate_video.exists():
+        return None
+    try:
+        from fastvideo.tests.utils import compute_video_ssim_torchvision
+    except ImportError as exc:
+        message = ("MS-SSIM is unavailable because `pytorch-msssim` is not installed. "
+                   "Install FastVideo with the test extra, e.g. `uv pip install -e '.[mlx,test]'`, "
+                   "or run without an SSIM assertion.")
+        if required:
+            raise RuntimeError(message) from exc
+        print(f"{message} Skipping MS-SSIM.")
+        return None
+
+    try:
+        ssim_values = compute_video_ssim_torchvision(str(reference_video), str(candidate_video), use_ms_ssim=True)
+    except ImportError as exc:
+        message = ("MS-SSIM is unavailable because `pytorch-msssim` is not installed. "
+                   "Install FastVideo with the test extra, e.g. `uv pip install -e '.[mlx,test]'`, "
+                   "or run without an SSIM assertion.")
+        if required:
+            raise RuntimeError(message) from exc
+        print(f"{message} Skipping MS-SSIM.")
+        return None
+    return float(ssim_values[0])
+
+
+def _markdown_table(rows: list[dict]) -> str:
+    columns = [
+        ("prompt_id", "prompt"),
+        ("mode", "mode"),
+        ("decoder", "decoder"),
+        ("status", "status"),
+        ("denoise_s", "denoise s"),
+        ("decode_s", "decode s"),
+        ("total_s", "total s"),
+        ("peak_gib", "peak GiB"),
+        ("ms_ssim_vs_ref", "MS-SSIM"),
+        ("lpips_vs_ref", "LPIPS"),
+    ]
+    header = "| " + " | ".join(label for _, label in columns) + " |"
+    sep = "| " + " | ".join("---" for _ in columns) + " |"
+    lines = [header, sep]
+    for row in rows:
+        cells = []
+        for key, _ in columns:
+            value = row.get(key)
+            if isinstance(value, float):
+                cells.append(f"{value:.3f}")
+            elif value is None:
+                cells.append("-")
+            else:
+                cells.append(str(value))
+        lines.append("| " + " | ".join(cells) + " |")
+    return "\n".join(lines)
+
+
+def _format_metric(value) -> str:
+    if isinstance(value, float):
+        return f"{value:.3f}"
+    if value is None:
+        return "-"
+    return str(value)
+
+
+def _html_grid(rows: list[dict]) -> str:
+    groups: dict[str, list[dict]] = {}
+    for row in rows:
+        groups.setdefault(str(row.get("prompt_id", "prompt")), []).append(row)
+
+    sections = []
+    for prompt_id, group_rows in groups.items():
+        prompt = next((str(row.get("prompt", "")) for row in group_rows if row.get("prompt")), "")
+        cards = []
+        for row in group_rows:
+            title = f"{row.get('mode', '-')} / {row.get('decoder', '-')}"
+            status = row.get("status", "-")
+            video_path = row.get("video_path")
+            if video_path:
+                media = f'<video src="{html.escape(str(video_path))}" muted loop controls playsinline></video>'
+            else:
+                media = f'<div class="missing">No video<br>{html.escape(str(row.get("error", "")))}</div>'
+            metrics = (f"status={status} · total={_format_metric(row.get('total_s'))}s · "
+                       f"denoise={_format_metric(row.get('denoise_s'))}s · "
+                       f"decode={_format_metric(row.get('decode_s'))}s · "
+                       f"peak={_format_metric(row.get('peak_gib'))}GiB")
+            cards.append("<article>"
+                         f"<h3>{html.escape(title)}</h3>"
+                         f"{media}"
+                         f"<p>{html.escape(metrics)}</p>"
+                         "</article>")
+        sections.append("<section>"
+                        f"<h2>{html.escape(prompt_id)}</h2>"
+                        f"<p class=\"prompt\">{html.escape(prompt)}</p>"
+                        f"<div class=\"grid\">{''.join(cards)}</div>"
+                        "</section>")
+
+    return """<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <title>FastVideo MLX benchmark grid</title>
+  <style>
+    body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; margin: 24px; background: #111; color: #eee; }
+    button { margin-right: 8px; padding: 8px 12px; border-radius: 8px; border: 1px solid #555; background: #222; color: #eee; }
+    section { margin-top: 28px; }
+    .prompt { color: #bbb; max-width: 900px; }
+    .grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(280px, 1fr)); gap: 16px; }
+    article { background: #1b1b1b; border: 1px solid #333; border-radius: 12px; padding: 12px; }
+    h1, h2, h3 { margin: 0 0 10px; }
+    video { width: 100%; border-radius: 8px; background: #000; }
+    article p { color: #bbb; font-size: 13px; line-height: 1.4; }
+    .missing { min-height: 160px; display: grid; place-items: center; text-align: center; color: #f5b5b5; background: #2a1515; border-radius: 8px; padding: 12px; }
+  </style>
+</head>
+<body>
+  <h1>FastVideo MLX benchmark grid</h1>
+  <p>Use the controls below to start/stop every clip together for side-by-side inspection.</p>
+  <button onclick="for (const v of document.querySelectorAll('video')) { v.currentTime = 0; v.play(); }">Restart + play all</button>
+  <button onclick="for (const v of document.querySelectorAll('video')) v.pause();">Pause all</button>
+""" + "\n".join(sections) + """
+</body>
+</html>
+"""
+
+
+def _write_html_grid(rows: list[dict], output_dir: Path) -> Path:
+    html_path = output_dir / "index.html"
+    html_path.write_text(_html_grid(rows))
+    return html_path
+
+
+def _generate_cell(
+    *,
+    args,
+    mode: str,
+    decoder: str,
+    checkpoint_path: Path,
+    config_path: Path,
+    encoder_hidden_states,
+    freqs_cis,
+    timesteps: list[int],
+    latents_seed: np.ndarray,
+    renoise_by_step: list[np.ndarray],
+) -> Cell:
+    import mlx.core as mx
+
+    from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
+    from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+    base_dtype, quantization = _mode_to_dtype_quant(mode)
+    mx_dtype = _mx_dtype(mx, base_dtype)
+
+    mx.clear_cache()
+    mx.reset_peak_memory()
+    load_start = time.perf_counter()
+    load_source = "diffusers"
+    if args.mlx_checkpoint_cache is not None:
+        from fastvideo.mlx_runtime.checkpoint import (
+            load_mlx_dit_checkpoint,
+            save_mlx_dit_checkpoint,
+        )
+
+        mode_ckpt_dir = args.mlx_checkpoint_cache / mode
+        if (mode_ckpt_dir / "mlx_dit.json").exists():
+            dit = load_mlx_dit_checkpoint(mode_ckpt_dir)
+            load_source = "mlx_checkpoint"
+        else:
+            dit = mlx_dit_from_diffusers_safetensors(
+                checkpoint_path,
+                config_path,
+                dtype=base_dtype,
+                quantization=quantization,
+            )
+            save_mlx_dit_checkpoint(dit, mode_ckpt_dir)
+            load_source = "diffusers_then_saved"
+    else:
+        dit = mlx_dit_from_diffusers_safetensors(
+            checkpoint_path,
+            config_path,
+            dtype=base_dtype,
+            quantization=quantization,
+        )
+    load_s = time.perf_counter() - load_start
+    load_peak = _peak_memory_bytes(mx)
+
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
+    schedule = MLXDMDSchedule.from_torch_scheduler(scheduler)
+
+    latents = mx.array(latents_seed).astype(mx_dtype)
+    mx.reset_peak_memory()
+    denoise_start = time.perf_counter()
+    latents_np, step_times = denoise_dmd_on_device(
+        mx=mx,
+        dit=dit,
+        latents=latents,
+        encoder_hidden_states=encoder_hidden_states.astype(mx_dtype),
+        freqs_cis=freqs_cis,
+        timesteps=timesteps,
+        renoise_by_step=renoise_by_step,
+        schedule=schedule,
+        dmd_step=dmd_step,
+        mx_dtype=mx_dtype,
+    )
+    denoise_s = time.perf_counter() - denoise_start
+    denoise_peak = _peak_memory_bytes(mx)
+    del dit, latents
+    cleanup_mlx(mx)
+
+    video_path = (args.output_dir / f"{args.current_prompt_id}" /
+                  f"video_{mode}_{decoder}_{args.height}x{args.width}x{args.num_frames}.mp4")
+    decode_start = time.perf_counter()
+    decode_latents_to_video(
+        model_root=args.model_root,
+        latents_np=latents_np,
+        output_path=video_path,
+        fps=args.fps,
+        device_arg=args.torch_device,
+        dtype_arg=args.torch_dtype,
+        backend=decoder,
+        taehv_source_path=args.taehv_source_path,
+        taehv_checkpoint_path=args.taehv_checkpoint_path,
+        taehv_parallel=args.taehv_parallel,
+    )
+    decode_s = time.perf_counter() - decode_start
+
+    metrics: dict[str, float | int | str | bool | None] = {
+        "prompt_id": args.current_prompt_id,
+        "prompt": args.current_prompt,
+        "benchmark_preset": args.benchmark_preset,
+        "mode": mode,
+        "decoder": decoder,
+        "status": "ok",
+        "video_path": str(video_path.relative_to(args.output_dir)),
+        "load_s": load_s,
+        "load_source": load_source,
+        "denoise_s": denoise_s,
+        # The first step carries one-time costs (mx.compile tracing, kernel
+        # warm-up); steady-state throughput is the median of the rest.
+        "denoise_first_step_s": step_times[0] if step_times else None,
+        "denoise_steady_step_s": (float(np.median(step_times[1:])) if len(step_times) > 1 else None),
+        "decode_s": decode_s,
+        "total_s": load_s + denoise_s + decode_s,
+        "load_peak_gib": load_peak / (1024**3),
+        "peak_gib": max(load_peak, denoise_peak) / (1024**3),
+        "quantization": quantization or "none",
+        "compute_dtype": base_dtype,
+        "compile": os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1",
+        "fast_norm": os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1",
+        "mlx_memory_limit_gib": args.mlx_memory_limit_gib,
+        "mlx_cache_limit_gib": args.mlx_cache_limit_gib,
+        "mlx_disable_cache": args.mlx_disable_cache,
+        "mlx_wired_limit_gib": args.mlx_wired_limit_gib,
+        "torch_mps_high_watermark_ratio": args.torch_mps_high_watermark_ratio,
+        "torch_mps_low_watermark_ratio": args.torch_mps_low_watermark_ratio,
+    }
+    return Cell(
+        prompt_id=args.current_prompt_id,
+        prompt=args.current_prompt,
+        mode=mode,
+        decoder=decoder,
+        video_path=video_path,
+        latents=latents_np,
+        metrics=metrics,
+    )
+
+
+def main() -> None:
+    preset_parser = argparse.ArgumentParser(add_help=False)
+    preset_parser.add_argument("--benchmark-preset", choices=tuple(BENCHMARK_PRESETS), default="default")
+    preset_args, _ = preset_parser.parse_known_args()
+    preset = BENCHMARK_PRESETS[preset_args.benchmark_preset]
+
+    parser = argparse.ArgumentParser(description="MLX FastWan prove-out benchmark (latency + quality).")
+    parser.add_argument("--benchmark-preset",
+                        choices=tuple(BENCHMARK_PRESETS),
+                        default=preset_args.benchmark_preset,
+                        help="Memory-tier benchmark defaults. Explicit CLI flags override preset values.")
+    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
+    parser.add_argument("--prompt", default="A paper boat sails through a shallow stream in a mossy forest.")
+    parser.add_argument(
+        "--prompt-file",
+        type=Path,
+        default=None,
+        help=
+        "Optional prompt set. Plain text uses one prompt per non-empty line; .jsonl accepts prompt/text/caption plus optional id/name.",
+    )
+    parser.add_argument(
+        "--prompt-set",
+        choices=("single", *PROMPT_SETS.keys()),
+        default="single",
+        help="Built-in standard prompt set. Ignored when --prompt-file is supplied.",
+    )
+    parser.add_argument("--height", type=int, default=preset.height)
+    parser.add_argument("--width", type=int, default=preset.width)
+    parser.add_argument("--num-frames", type=int, default=preset.num_frames)
+    parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
+    parser.add_argument("--flow-shift", type=float, default=8.0)
+    parser.add_argument("--max-sequence-length", type=int, default=512)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--modes", default=preset.modes)
+    parser.add_argument("--decoders", default=preset.decoders)
+    parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_fastwan_bench"))
+    parser.add_argument("--torch-device", default="auto")
+    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
+    parser.add_argument(
+        "--reference",
+        type=Path,
+        default=None,
+        help="External reference mp4 to score every cell against. Defaults to the fp16+wan-vae cell.",
+    )
+    parser.add_argument("--assert-min-ssim",
+                        type=float,
+                        default=None,
+                        help="Fail if any cell's MS-SSIM vs the reference falls below this value.")
+    parser.add_argument("--compile",
+                        action="store_true",
+                        help="Enable mx.compile on the DiT forward (sets FASTVIDEO_MLX_COMPILE=1).")
+    parser.add_argument("--lpips", action="store_true", help="Also compute LPIPS (needs the `lpips` package).")
+    parser.add_argument("--taehv-source-path", type=Path, default=None)
+    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
+    parser.add_argument("--taehv-parallel", action="store_true")
+    parser.add_argument(
+        "--mlx-checkpoint-cache",
+        type=Path,
+        default=None,
+        help="Directory of per-mode pre-quantized MLX checkpoints. The first cell of a mode "
+        "converts from Diffusers weights and saves here (load_source=diffusers_then_saved); "
+        "later cells and later runs reload without requantizing (load_source=mlx_checkpoint), "
+        "which is also how the checkpoint load-time win is measured.",
+    )
+    add_memory_limit_args(
+        parser,
+        mlx_memory_limit_gib=preset.mlx_memory_limit_gib,
+        mlx_disable_cache=preset.mlx_disable_cache,
+        torch_mps_high_watermark_ratio=preset.torch_mps_high_watermark_ratio,
+        torch_mps_low_watermark_ratio=preset.torch_mps_low_watermark_ratio,
+    )
+    args = parser.parse_args()
+
+    if args.compile:
+        os.environ["FASTVIDEO_MLX_COMPILE"] = "1"
+
+    import mlx.core as mx
+
+    runtime_limits = apply_memory_limits(
+        mlx_memory_limit_gib=args.mlx_memory_limit_gib,
+        mlx_cache_limit_gib=args.mlx_cache_limit_gib,
+        mlx_disable_cache=args.mlx_disable_cache,
+        mlx_wired_limit_gib=args.mlx_wired_limit_gib,
+        torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio,
+        torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio,
+        mx_module=mx,
+    ).as_metrics()
+    import torch
+
+    mx.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    modes = _parse_list(args.modes, ALLOWED_MODES, "modes")
+    decoders = _parse_list(args.decoders, ALLOWED_DECODERS, "decoders")
+
+    config_path = args.model_root / "transformer/config.json"
+    checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors"
+    config = json.loads(config_path.read_text())
+    latent_frames = (args.num_frames - 1) // 4 + 1
+    latent_height = args.height // 8
+    latent_width = args.width // 8
+    freqs_cis = make_rotary_embeddings(
+        config,
+        latent_frames=latent_frames,
+        latent_height=latent_height,
+        latent_width=latent_width,
+    )
+    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    latents_seed = torch.randn(
+        (1, int(config["in_channels"]), latent_frames, latent_height, latent_width),
+        generator=generator,
+        dtype=torch.float32,
+    ).numpy()
+
+    timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()]
+    # Keep DMD stochasticity identical across benchmark cells. Without this,
+    # FP16/INT8/decoder comparisons can accidentally measure different re-noise
+    # samples instead of only quantization or decoder differences.
+    renoise_by_step = [
+        torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy()
+        for _ in range(max(0,
+                           len(timesteps) - 1))
+    ]
+
+    from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError
+
+    prompt_cases = _load_prompt_cases(args.prompt, args.prompt_file, args.prompt_set)
+    cells: list[Cell] = []
+    unsupported_rows: list[dict] = []
+    for prompt_case in prompt_cases:
+        args.current_prompt_id = prompt_case.id
+        args.current_prompt = prompt_case.prompt
+        prompt_embeds = encode_prompt(
+            model_root=args.model_root,
+            prompt=prompt_case.prompt,
+            max_sequence_length=args.max_sequence_length,
+            device_arg=args.torch_device,
+            dtype_arg=args.torch_dtype,
+        )
+        encoder_hidden_states = mx.array(prompt_embeds.numpy())
+
+        for mode in modes:
+            for decoder in decoders:
+                print(f"=== cell: prompt={prompt_case.id} mode={mode} decoder={decoder} ===")
+                try:
+                    cells.append(
+                        _generate_cell(
+                            args=args,
+                            mode=mode,
+                            decoder=decoder,
+                            checkpoint_path=checkpoint_path,
+                            config_path=config_path,
+                            encoder_hidden_states=encoder_hidden_states,
+                            freqs_cis=freqs_cis,
+                            timesteps=timesteps,
+                            latents_seed=latents_seed,
+                            renoise_by_step=renoise_by_step,
+                        ))
+                except UnsupportedMLXQuantizationError as exc:
+                    # Record the cell as unsupported and keep sweeping: a partial
+                    # report on this MLX build beats crashing the whole run.
+                    print(f"skipping cell (unsupported by this MLX build): {exc}")
+                    unsupported_rows.append({
+                        "prompt_id": prompt_case.id,
+                        "prompt": prompt_case.prompt,
+                        "mode": mode,
+                        "decoder": decoder,
+                        "status": "unsupported_by_mlx",
+                        "error": str(exc),
+                    })
+
+    if not cells:
+        metrics_path = args.output_dir / "metrics.json"
+        metrics_path.write_text(json.dumps(unsupported_rows, indent=2))
+        raise SystemExit(f"No benchmark cell could run: every requested mode is unsupported by this MLX build. "
+                         f"Wrote {metrics_path}.")
+
+    # Resolve one internal reference per prompt. A single external reference, if
+    # supplied, is used for every prompt and only video metrics are computed.
+    reference_by_prompt: dict[str, tuple[Path, np.ndarray | None]] = {}
+    if args.reference is not None:
+        for prompt_case in prompt_cases:
+            reference_by_prompt[prompt_case.id] = (args.reference, None)
+    else:
+        for prompt_case in prompt_cases:
+            prompt_cells = [c for c in cells if c.prompt_id == prompt_case.id]
+            if not prompt_cells:
+                continue
+            ref_cell = next(
+                (c for c in prompt_cells if c.mode == REFERENCE_MODE and c.decoder == REFERENCE_DECODER),
+                prompt_cells[0],
+            )
+            reference_by_prompt[prompt_case.id] = (ref_cell.video_path, ref_cell.latents)
+            print(f"Using internal reference cell for {prompt_case.id}: "
+                  f"mode={ref_cell.mode} decoder={ref_cell.decoder}")
+
+    lpips_fn = _load_lpips() if args.lpips else None
+
+    rows: list[dict] = []
+    failures: list[str] = []
+    for cell in cells:
+        reference_video, reference_latents = reference_by_prompt[cell.prompt_id]
+        ms_ssim = _ms_ssim(Path(reference_video), cell.video_path, required=args.assert_min_ssim is not None)
+        cell.metrics["ms_ssim_vs_ref"] = ms_ssim
+        cell.metrics.update(runtime_limits)
+        if reference_latents is not None:
+            cell.metrics.update(_latent_delta_metrics(cell.latents, reference_latents))
+        cell.metrics["lpips_vs_ref"] = (_lpips_between(lpips_fn, Path(reference_video), cell.video_path)
+                                        if lpips_fn else None)
+        if args.assert_min_ssim is not None and ms_ssim is not None and ms_ssim < args.assert_min_ssim:
+            failures.append(
+                f"{cell.prompt_id}/{cell.mode}/{cell.decoder}: MS-SSIM {ms_ssim:.4f} < {args.assert_min_ssim}")
+        rows.append(dict(cell.metrics))
+        print(json.dumps(cell.metrics, indent=2))
+    rows.extend(unsupported_rows)
+
+    metrics_path = args.output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(rows, indent=2))
+    table_path = args.output_dir / "metrics.md"
+    table = _markdown_table(rows)
+    table_path.write_text(table + "\n")
+    html_path = _write_html_grid(rows, args.output_dir)
+    print("\n" + table)
+    print(f"\nWrote {metrics_path}, {table_path}, and {html_path}")
+
+    if failures:
+        raise SystemExit("SSIM regression gate failed:\n  " + "\n  ".join(failures))
+
+
+def _load_lpips() -> object | None:
+    """Return an LPIPS model, or ``None`` if the optional dep is unavailable."""
+    try:
+        import lpips  # noqa: PLC0415 - optional dependency.
+    except ImportError:
+        print("LPIPS requested but the `lpips` package is not installed; skipping (install `.[eval]`).")
+        return None
+    return lpips.LPIPS(net="alex")
+
+
+def _lpips_between(lpips_fn, reference_video: Path, candidate_video: Path) -> float | None:
+    if lpips_fn is None or not reference_video.exists() or not candidate_video.exists():
+        return None
+    import torch
+
+    ref = _read_video_frames(reference_video)
+    cand = _read_video_frames(candidate_video)
+    if ref is None or cand is None or ref.shape != cand.shape:
+        return None
+    # LPIPS expects NCHW in [-1, 1].
+    ref_t = torch.from_numpy(ref).permute(0, 3, 1, 2).float() / 127.5 - 1.0
+    cand_t = torch.from_numpy(cand).permute(0, 3, 1, 2).float() / 127.5 - 1.0
+    with torch.no_grad():
+        scores = lpips_fn(ref_t, cand_t)
+    return float(scores.mean().item())
+
+
+def _read_video_frames(path: Path) -> np.ndarray | None:
+    try:
+        import cv2
+    except ImportError:
+        return None
+    cap = cv2.VideoCapture(str(path))
+    frames = []
+    try:
+        while True:
+            ok, frame_bgr = cap.read()
+            if not ok:
+                break
+            frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
+    finally:
+        cap.release()
+    if not frames:
+        return None
+    return np.stack(frames, axis=0)
+
+
+if __name__ == "__main__":
+    main()

fastvideo/entrypoints/video_generator.py

@@ -808,14 +808,19 @@ class VideoGenerator:
         latent_batch_size = _infer_latent_batch_size(batch)
         is_latent_output = fastvideo_args.output_type == "latent"
         needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output)
-        needs_samples_buffer = batch.return_frames or needs_frame_output
-        # When ``output_type == "latent"`` the forward output has latent
-        # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``)
-        # rather than the pre-allocation's pixel shape. Skip the pinned
-        # ~50 MB buffer entirely. Also skip it for metadata-only calls;
-        # neither the result nor save path will consume the decoded tensor.
+        # A populated ``samples`` has exactly one consumer — the result
+        # dict (``"samples": samples if batch.return_frames else None``).
+        # Post-decode frame building reads ``output_batch.output``
+        # directly (the GPU ``vid_u8`` path), not ``samples``. So when
+        # ``return_frames=False`` the pinned fp32 alloc + D->H copy are
+        # dead weight — the CLI generate flow (``save_video=True``,
+        # ``return_frames=False``) hits this on every call.
+        # ``output_type == "latent"`` keeps its existing branch (shape
+        # mismatch falls through to ``.cpu()`` below) for callers that
+        # *do* ask for the latent samples via ``return_frames=True``.
         # ``skip_pixel_prealloc`` also gates the slow-path warning.
-        skip_pixel_prealloc = is_latent_output or not needs_samples_buffer
+        needs_samples_out = batch.return_frames
+        skip_pixel_prealloc = is_latent_output or not needs_samples_out
         if skip_pixel_prealloc:
             samples = torch.empty(0, device='cpu')
         else:
@@ -835,9 +840,11 @@ class VideoGenerator:
                                "This usually means the executor/pipeline failed earlier.")
 
         audio_only = bool(output_batch.extra.get("audio_only"))
-        if not needs_samples_buffer or (audio_only and not batch.return_frames):
-            # Metadata-only/audio-only request: keep the empty placeholder and
-            # avoid the decoded tensor D->H copy.
+        if not needs_samples_out:
+            # Nothing downstream reads ``samples`` (the result dict
+            # returns None when ``return_frames=False``); keep the empty
+            # placeholder allocated above and skip the fp32 D->H copy
+            # entirely.
             pass
         elif audio_only:
             # Audio-only return-frames requests expose the small placeholder
@@ -869,8 +876,13 @@ class VideoGenerator:
         # `GenerationResult.size` describes the produced media, not only the
         # base-stage request. Refiner pipelines can change the final pixel
         # dimensions, so derive this result metadata from the decoded output.
+        # Read the geometry from `output_batch.output` (a shape-only access,
+        # no D->H copy): when `return_frames=False` the `samples` mirror
+        # stays an empty placeholder and no longer carries the decoded
+        # shape. Metadata-only calls keep the request fallback and never
+        # inspect the (possibly dropped) worker output.
         output_size = _resolve_output_size(
-            samples,
+            output_batch.output if needs_frame_output else samples,
             (target_height, target_width, batch.num_frames),
             pixel_output=not is_latent_output and not audio_only,
         )
@@ -882,13 +894,26 @@ class VideoGenerator:
         elif not needs_frame_output:
             frames = None
         else:
-            videos = rearrange(samples, "b c t h w -> t b c h w")
-            frames = []
-            for x in videos:
-                x = torchvision.utils.make_grid(x, nrow=6)
-                x = x.permute(1, 2, 0).squeeze(-1)
-                x = (x * 255).to(torch.uint8)
-                frames.append(x.contiguous().cpu().numpy())
+            # Quantize on the source device (typically CUDA) BEFORE the
+            # device->host copy. `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. The old path
+            # paid a full fp32 video D->H copy (which scales with
+            # resolution x frames x batch) and then a single-threaded
+            # per-frame CPU *255/cast loop. Casting to uint8 on-device
+            # makes the transfer 4x smaller, ships it in a single copy,
+            # and moves the elementwise work onto the GPU. clamp_() also
+            # fixes a latent overflow bug: VAE output slightly outside
+            # [0, 1] wrapped mod 256 in the old unclamped cast.
+            # (Equivalence is SSIM-gated, not bit-exact: float->uint8
+            # differs <=1 LSB CPU vs GPU.)
+            src = output_batch.output
+            vid_u8 = (src * 255).clamp_(0, 255).to(torch.uint8)
+            vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w").cpu()
+            frames = [
+                torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8
+            ]
         postprocess_time = time.perf_counter() - postprocess_start
         logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
         if logging_info is not None:

fastvideo/mlx_runtime/__init__.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Experimental Apple MLX runtime helpers.
+
+This package is intentionally small for now. It exists to grow the Apple-native
+FastWan path in measurable steps: shape planning, primitive benchmarks, then
+Wan block parity, then full DiT/runtime support.
+"""
+
+from fastvideo.mlx_runtime.fastwan import (
+    FastWanShape,
+    MLXQuantizationSpec,
+    MLXWanDiT,
+    MLXWanTransformerBlock,
+    UnsupportedMLXQuantizationError,
+    ensure_quantization_supported,
+    fastwan_shape,
+    fastwan_shape_from_config,
+    mlx_dit_from_diffusers_safetensors,
+    mlx_block_weights_from_torch,
+    mlx_block_weights_from_diffusers_safetensors,
+    quantization_support_error,
+)
+from fastvideo.mlx_runtime.checkpoint import (
+    load_mlx_dit_checkpoint,
+    save_mlx_dit_checkpoint,
+)
+from fastvideo.mlx_runtime.memory import (
+    AppliedMemoryLimits,
+    add_memory_limit_args,
+    apply_memory_limits,
+    gib_to_bytes,
+)
+from fastvideo.mlx_runtime.refine import (
+    DEFAULT_REFINE_SIGMA,
+    RefinePlan,
+    TwoPassResult,
+    default_refine_timesteps,
+    plan_refine_resolutions,
+    prepare_refine_latents,
+    refine_sigma_from_schedule,
+    run_dmd_loop,
+    run_two_pass_dmd,
+    upsample_latents_spatial,
+)
+from fastvideo.mlx_runtime.frame_upsample import (
+    DEFAULT_PIXEL_UPSAMPLE_MODE,
+    PIXEL_UPSAMPLE_MODES,
+    unsharp,
+    upsample_frame,
+    upsample_frames,
+)
+from fastvideo.mlx_runtime.fast_spatial import (
+    DEFAULT_FAST_SPATIAL_SHARPEN,
+    FastSpatialPlan,
+    apply_fast_spatial_upsample,
+    plan_fast_spatial,
+    resolve_spatial_mode,
+)
+from fastvideo.mlx_runtime.prompt_enhance import (
+    DEFAULT_ENHANCE_SYSTEM_PROMPT,
+    DEFAULT_MLX_LM_MODEL,
+    EnhanceResult,
+    enhance_prompt,
+    enhance_prompt_template,
+    enhance_result_as_metrics,
+    load_or_enhance_prompt,
+)
+
+__all__ = [
+    "AppliedMemoryLimits",
+    "DEFAULT_ENHANCE_SYSTEM_PROMPT",
+    "DEFAULT_MLX_LM_MODEL",
+    "DEFAULT_FAST_SPATIAL_SHARPEN",
+    "DEFAULT_PIXEL_UPSAMPLE_MODE",
+    "DEFAULT_REFINE_SIGMA",
+    "EnhanceResult",
+    "FastSpatialPlan",
+    "FastWanShape",
+    "MLXQuantizationSpec",
+    "MLXWanDiT",
+    "MLXWanTransformerBlock",
+    "RefinePlan",
+    "TwoPassResult",
+    "UnsupportedMLXQuantizationError",
+    "add_memory_limit_args",
+    "apply_fast_spatial_upsample",
+    "apply_memory_limits",
+    "enhance_prompt",
+    "enhance_prompt_template",
+    "enhance_result_as_metrics",
+    "ensure_quantization_supported",
+    "fastwan_shape",
+    "fastwan_shape_from_config",
+    "gib_to_bytes",
+    "load_mlx_dit_checkpoint",
+    "load_or_enhance_prompt",
+    "mlx_dit_from_diffusers_safetensors",
+    "mlx_block_weights_from_diffusers_safetensors",
+    "mlx_block_weights_from_torch",
+    "PIXEL_UPSAMPLE_MODES",
+    "default_refine_timesteps",
+    "plan_fast_spatial",
+    "plan_refine_resolutions",
+    "prepare_refine_latents",
+    "quantization_support_error",
+    "refine_sigma_from_schedule",
+    "resolve_spatial_mode",
+    "run_dmd_loop",
+    "run_two_pass_dmd",
+    "save_mlx_dit_checkpoint",
+    "unsharp",
+    "upsample_frame",
+    "upsample_frames",
+    "upsample_latents_spatial",
+]

fastvideo/mlx_runtime/checkpoint.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pre-quantized MLX checkpoint save/load for the FastWan DiT.
+
+Loading the Diffusers fp32/fp16 checkpoint and quantizing at startup costs
+both download size and load time on every run. This module persists an already
+cast (and optionally already quantized) ``MLXWanDiT`` so 16 GB users download
+and load roughly half the bytes and skip requantization entirely:
+
+    dit = mlx_dit_from_diffusers_safetensors(ckpt, cfg, quantization="int8")
+    save_mlx_dit_checkpoint(dit, "FastWan2.1-T2V-1.3B-mlx-int8")
+    ...
+    dit = load_mlx_dit_checkpoint("FastWan2.1-T2V-1.3B-mlx-int8")
+
+Format (one directory):
+
+- ``mlx_dit.safetensors`` — every array, saved with ``mx.save_safetensors``.
+  Plain weights keep their key; a quantized weight ``K`` is stored as the
+  packed ``K`` plus ``K.scales`` (and ``K.biases`` for affine modes).
+- ``mlx_dit.json`` — format version, the model config, the quantization spec,
+  and which keys are quantized, so the loader can rebuild ``QuantizedMatrix``
+  objects without guessing.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any
+
+from fastvideo.logger import init_logger
+from fastvideo.mlx_runtime.fastwan import (
+    MLXQuantizationSpec,
+    MLXWanDiT,
+    MLXWanTransformerBlock,
+    QuantizedMatrix,
+    ensure_quantization_supported,
+)
+
+logger = init_logger(__name__)
+
+FORMAT_VERSION = 1
+WEIGHTS_FILENAME = "mlx_dit.safetensors"
+MANIFEST_FILENAME = "mlx_dit.json"
+
+_BLOCK_PREFIX = "blocks"
+
+_DTYPE_TO_NAME = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"}
+
+
+def _dtype_name(dtype) -> str:
+    """Return the manifest name for a supported MLX data type.
+
+    Parameters:
+        dtype: The MLX data type to convert.
+
+    Returns:
+        str: The manifest name corresponding to the data type.
+
+    Raises:
+        ValueError: If the data type is not supported for checkpointing.
+    """
+    import mlx.core as mx
+
+    for raw, name in _DTYPE_TO_NAME.items():
+        if dtype == getattr(mx, raw):
+            return name
+    raise ValueError(f"Unsupported MLX dtype for checkpointing: {dtype}")
+
+
+def _name_to_dtype(name: str):
+    """Convert a manifest dtype name to its corresponding MLX dtype.
+
+    Parameters:
+        name (str): Manifest name, such as ``"fp16"``, ``"bf16"``, or ``"fp32"``.
+
+    Returns:
+        The corresponding MLX dtype.
+    """
+    import mlx.core as mx
+
+    return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[name]
+
+
+def _flatten_weights(dit: MLXWanDiT) -> dict[str, Any]:
+    """Combine model and transformer-block weights into a single flattened mapping.
+
+    Parameters:
+        dit (MLXWanDiT): Model whose weights should be flattened.
+
+    Returns:
+        dict[str, Any]: Mapping containing top-level weights and indexed transformer-block weights.
+    """
+    flat: dict[str, Any] = dict(dit.weights)
+    for index, block in enumerate(dit.blocks):
+        for name, value in block.weights.items():
+            flat[f"{_BLOCK_PREFIX}.{index}.{name}"] = value
+    return flat
+
+
+def save_mlx_dit_checkpoint(dit: MLXWanDiT, checkpoint_dir: str | Path) -> Path:
+    """Save a plain or quantized MLX Wan DiT checkpoint to a directory.
+
+    Parameters:
+        dit (MLXWanDiT): Model whose weights and configuration will be saved.
+        checkpoint_dir (str | Path): Destination directory for the checkpoint.
+
+    Returns:
+        Path: Path to the checkpoint directory.
+    """
+    import mlx.core as mx
+
+    checkpoint_dir = Path(checkpoint_dir)
+    arrays: dict[str, Any] = {}
+    quantized: dict[str, dict[str, Any]] = {}
+    spec: MLXQuantizationSpec | None = None
+    for key, value in _flatten_weights(dit).items():
+        if isinstance(value, QuantizedMatrix):
+            if spec is not None and value.spec != spec:
+                raise ValueError(f"Mixed quantization specs in one checkpoint ({spec} vs {value.spec} at '{key}') "
+                                 "are not supported.")
+            spec = value.spec
+            arrays[key] = value.weight
+            arrays[f"{key}.scales"] = value.scales
+            if value.biases is not None:
+                arrays[f"{key}.biases"] = value.biases
+            quantized[key] = {
+                "dequantized_dtype": _dtype_name(value.dequantized_dtype),
+                "has_biases": value.biases is not None,
+            }
+        else:
+            arrays[key] = value
+
+    manifest = {
+        "format_version": FORMAT_VERSION,
+        "config": dit.config,
+        "num_blocks": len(dit.blocks),
+        "quantization": None if spec is None else {
+            "mode": spec.mode,
+            "bits": spec.bits,
+            "group_size": spec.group_size,
+        },
+        "quantized_keys": quantized,
+    }
+
+    manifest_json = json.dumps(manifest, indent=2)
+    checkpoint_dir.parent.mkdir(parents=True, exist_ok=True)
+    staging_dir = Path(tempfile.mkdtemp(dir=checkpoint_dir.parent, prefix=f".{checkpoint_dir.name}.staging-"))
+    backup_root: Path | None = None
+    try:
+        staged_weights = staging_dir / WEIGHTS_FILENAME
+        staged_manifest = staging_dir / MANIFEST_FILENAME
+        mx.save_safetensors(str(staged_weights), arrays)
+        staged_manifest.write_text(manifest_json)
+        if checkpoint_dir.exists():
+            backup_root = Path(tempfile.mkdtemp(dir=checkpoint_dir.parent, prefix=f".{checkpoint_dir.name}.backup-"))
+            try:
+                checkpoint_dir.replace(backup_root / checkpoint_dir.name)
+            except Exception:
+                shutil.rmtree(backup_root, ignore_errors=True)
+                raise
+        try:
+            staging_dir.replace(checkpoint_dir)
+        except Exception:
+            if backup_root is not None:
+                (backup_root / checkpoint_dir.name).replace(checkpoint_dir)
+                shutil.rmtree(backup_root, ignore_errors=True)
+            raise
+        if backup_root is not None:
+            shutil.rmtree(backup_root, ignore_errors=True)
+    finally:
+        shutil.rmtree(staging_dir, ignore_errors=True)
+    logger.info("Saved MLX DiT checkpoint (%d arrays, quantization=%s) to %s", len(arrays),
+                spec.label if spec else "none", checkpoint_dir)
+    return checkpoint_dir
+
+
+def load_mlx_dit_checkpoint(checkpoint_dir: str | Path, *, compile: bool = False) -> MLXWanDiT:
+    """
+    Reconstruct an MLXWanDiT model from a versioned checkpoint.
+
+    Parameters:
+        checkpoint_dir (str | Path): Directory containing the checkpoint manifest and weights.
+        compile (bool): Whether to configure the reconstructed model for compilation.
+
+    Returns:
+        MLXWanDiT: The reconstructed model.
+
+    Raises:
+        FileNotFoundError: If the checkpoint manifest or weights file is missing.
+        ValueError: If the checkpoint format is unsupported or block weights are incomplete.
+    """
+    import mlx.core as mx
+
+    checkpoint_dir = Path(checkpoint_dir)
+    manifest_path = checkpoint_dir / MANIFEST_FILENAME
+    weights_path = checkpoint_dir / WEIGHTS_FILENAME
+    if not manifest_path.exists() or not weights_path.exists():
+        raise FileNotFoundError(f"Not an MLX DiT checkpoint directory: {checkpoint_dir} "
+                                f"(expected {MANIFEST_FILENAME} and {WEIGHTS_FILENAME}).")
+
+    manifest = json.loads(manifest_path.read_text())
+    version = manifest.get("format_version")
+    if version != FORMAT_VERSION:
+        raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} has format_version={version}; "
+                         f"this FastVideo build reads version {FORMAT_VERSION}. Re-export the checkpoint.")
+
+    spec = None
+    if manifest["quantization"] is not None:
+        spec = MLXQuantizationSpec(**manifest["quantization"])
+        # The packed layout of mx.quantize output is mode-specific, so a build
+        # that cannot run the mode cannot use these arrays at all.
+        ensure_quantization_supported(spec)
+
+    arrays = mx.load(str(weights_path))
+    quantized_keys: dict[str, dict[str, Any]] = manifest["quantized_keys"]
+
+    def rebuild(key: str):
+        """
+        Reconstructs a weight array or quantized matrix from checkpoint data.
+
+        Parameters:
+            key (str): The weight key to rebuild.
+
+        Returns:
+            The stored array for an unquantized weight or a reconstructed quantized matrix.
+        """
+        if key not in quantized_keys:
+            return arrays[key]
+        info = quantized_keys[key]
+        assert spec is not None, f"Quantized key '{key}' in a checkpoint without a quantization spec"
+        return QuantizedMatrix(
+            weight=arrays[key],
+            scales=arrays[f"{key}.scales"],
+            biases=arrays[f"{key}.biases"] if info["has_biases"] else None,
+            spec=spec,
+            dequantized_dtype=_name_to_dtype(info["dequantized_dtype"]),
+        )
+
+    config = manifest["config"]
+    block_keys: dict[int, list[str]] = {}
+    top_level_keys: list[str] = []
+    for key in arrays:
+        if key.endswith(".scales") or key.endswith(".biases"):
+            continue
+        if key.startswith(f"{_BLOCK_PREFIX}."):
+            index_str, _, _ = key[len(_BLOCK_PREFIX) + 1:].partition(".")
+            block_keys.setdefault(int(index_str), []).append(key)
+        else:
+            top_level_keys.append(key)
+
+    weights = {key: rebuild(key) for key in top_level_keys}
+
+    num_blocks = int(manifest["num_blocks"])
+    if sorted(block_keys) != list(range(num_blocks)):
+        raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} is missing block weights: "
+                         f"manifest says {num_blocks} blocks, found indices {sorted(block_keys)}.")
+
+    inner_dim = int(config["num_attention_heads"]) * int(config["attention_head_dim"])
+    blocks = []
+    for index in range(num_blocks):
+        prefix = f"{_BLOCK_PREFIX}.{index}."
+        block_weights = {key[len(prefix):]: rebuild(key) for key in block_keys[index]}
+        blocks.append(
+            MLXWanTransformerBlock(
+                block_weights,
+                dim=inner_dim,
+                ffn_dim=int(config["ffn_dim"]),
+                num_heads=int(config["num_attention_heads"]),
+                eps=float(config["eps"]),
+            ))
+    return MLXWanDiT(weights, blocks, config, compile=compile)

fastvideo/mlx_runtime/fast_spatial.py

@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Spatial fast mode for the MLX Wan runtime (RIFE's spatial twin).
+
+RIFE ``--fast`` cuts *frames* (temporal). This module cuts *pixels*
+(spatial): denoise at ``target // scale``, decode at that size, then
+resample the decoded frames up to the target. No second denoise pass —
+that is ``--refine`` (quality). The two compose:
+
+* ``--fast-spatial`` alone → speed (≈ scale² fewer tokens)
+* ``--refine`` alone → quality two-pass (H3 / LTX-2)
+* ``--fast`` + ``--refine`` → fewer frames at base res, full-res refine
+* ``--fast`` + ``--fast-spatial`` → fewer frames *and* fewer pixels
+
+The upsample runs in **pixel** space, after the VAE decode. It used to run
+in latent space (bilinear over the latent H/W plane, sharing the refine
+hand-off primitive) and that is what made spatial fast mode incoherent: an
+interpolated Wan latent is off the decoder's manifold, so decode returned
+the right silhouette under a smeared veil. ``--refine`` can get away with
+the latent-space upsample because a second DMD pass re-denoises the result;
+spatial fast mode hands the latent straight to the decoder, so it cannot.
+See :mod:`fastvideo.mlx_runtime.frame_upsample` for the full rationale.
+
+MetalFX is intentionally not used: it needs game-engine motion vectors
+and depth that diffusion output lacks.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+import numpy as np
+
+from fastvideo.logger import init_logger
+from fastvideo.mlx_runtime.frame_upsample import (
+    DEFAULT_PIXEL_UPSAMPLE_MODE,
+    PIXEL_UPSAMPLE_MODES,
+    upsample_frames,
+)
+from fastvideo.mlx_runtime.refine import RefinePlan, plan_refine_resolutions
+
+logger = init_logger(__name__)
+
+# Resampling from a smaller decode loses high-frequency detail the same way
+# RIFE's flow warp does, so spatial fast mode borrows ``--fast``'s remedy: a
+# light unsharp pass. 0.4 recovers perceived crispness on Wan2.1 output at 2x
+# without the halos that show up by ~0.8.
+DEFAULT_FAST_SPATIAL_SHARPEN = 0.4
+
+
+@dataclass(frozen=True)
+class FastSpatialPlan:
+    """Resolved geometry for a spatial-fast (upsample-only) run."""
+
+    plan: RefinePlan
+    upsample_mode: str
+    sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN
+
+    @property
+    def enabled(self) -> bool:
+        """
+        Determine whether spatial scaling is enabled.
+
+        Returns:
+            `true` if the spatial scale is greater than one, `false` otherwise.
+        """
+        return self.plan.spatial_scale > 1
+
+    @property
+    def scale(self) -> int:
+        """Provides the configured spatial scaling factor.
+
+        Returns:
+            int: The spatial scaling factor.
+        """
+        return self.plan.spatial_scale
+
+    @property
+    def target_height(self) -> int:
+        """
+        Return the target output height for the spatial plan.
+
+        Returns:
+            int: Target output height in pixels.
+        """
+        return self.plan.target_height
+
+    @property
+    def target_width(self) -> int:
+        """Return the target image width in pixels.
+
+        Returns:
+            int: The target image width.
+        """
+        return self.plan.target_width
+
+    @property
+    def stage1_height(self) -> int:
+        """
+        Provide the stage-one latent height used for reduced-resolution processing.
+
+        Returns:
+            int: The stage-one latent height.
+        """
+        return self.plan.stage1_height
+
+    @property
+    def stage1_width(self) -> int:
+        """Get the stage-one latent width.
+
+        Returns:
+            int: The stage-one latent width.
+        """
+        return self.plan.stage1_width
+
+
+def plan_fast_spatial(
+    *,
+    height: int,
+    width: int,
+    num_frames: int,
+    spatial_scale: int = 2,
+    vae_spatial_compression: int = 8,
+    vae_temporal_compression: int = 4,
+    patch_size: tuple[int, int, int] = (1, 2, 2),
+    upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE,
+    sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN,
+    enabled: bool = True,
+) -> FastSpatialPlan:
+    """
+    Build a plan for reduced-resolution denoising followed by pixel-space upsampling.
+
+    Parameters:
+        upsample_mode (str): Pixel interpolation kernel, one of
+            :data:`~fastvideo.mlx_runtime.frame_upsample.PIXEL_UPSAMPLE_MODES`.
+        sharpen (float): Unsharp strength applied after the resize.
+
+    Returns:
+        FastSpatialPlan: The validated spatial-fast processing plan.
+
+    Raises:
+        ValueError: If the upsample mode is unsupported or ``sharpen`` is negative.
+    """
+    if upsample_mode not in PIXEL_UPSAMPLE_MODES:
+        raise ValueError(f"Unsupported upsample mode: {upsample_mode!r} "
+                         f"(expected one of {', '.join(PIXEL_UPSAMPLE_MODES)})")
+    if sharpen < 0.0:
+        raise ValueError(f"sharpen must be >= 0, got {sharpen}")
+    plan = plan_refine_resolutions(
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        spatial_scale=spatial_scale,
+        vae_spatial_compression=vae_spatial_compression,
+        vae_temporal_compression=vae_temporal_compression,
+        patch_size=patch_size,
+        enabled=enabled,
+        mode_label="fast-spatial",
+    )
+    if plan.spatial_scale > 1:
+        logger.info(
+            "[MLX fast-spatial] denoise+decode %dx%d → upsample %dx to %dx%d (%s, sharpen=%.2f)",
+            plan.stage1_width,
+            plan.stage1_height,
+            plan.spatial_scale,
+            plan.target_width,
+            plan.target_height,
+            upsample_mode,
+            sharpen,
+        )
+    return FastSpatialPlan(plan=plan, upsample_mode=upsample_mode, sharpen=sharpen)
+
+
+def apply_fast_spatial_upsample(
+    frames: Iterable[np.ndarray],
+    spatial: FastSpatialPlan,
+) -> list[np.ndarray]:
+    """Resample decoded stage-1 frames up to the target resolution.
+
+    This runs on decoded RGB frames, *not* on latents: see the module
+    docstring for why the latent-space version produced a blurred veil.
+
+    Parameters:
+        frames (Iterable[np.ndarray]): Decoded HxWx3 uint8 RGB frames, produced
+            by decoding at the stage-one resolution.
+        spatial (FastSpatialPlan): Plan defining the target size, interpolation
+            kernel, and unsharp strength.
+
+    Returns:
+        list[np.ndarray]: Frames at the target resolution. When spatial scaling
+            is disabled the frames are returned unchanged, as a list.
+    """
+    if not spatial.enabled:
+        return list(frames)
+    return upsample_frames(
+        frames,
+        width=spatial.target_width,
+        height=spatial.target_height,
+        mode=spatial.upsample_mode,
+        sharpen=spatial.sharpen,
+    )
+
+
+def resolve_spatial_mode(
+    *,
+    refine: bool,
+    fast_spatial: bool,
+) -> str:
+    """Select the active spatial processing mode, with refinement taking precedence.
+
+    Returns:
+        str: ``"refine"`` when refinement is enabled, ``"fast_spatial"`` when
+            spatial-fast processing is enabled, or ``"off"`` otherwise.
+    """
+    if refine:
+        return "refine"
+    if fast_spatial:
+        return "fast_spatial"
+    return "off"
+
+
+__all__ = [
+    "DEFAULT_FAST_SPATIAL_SHARPEN",
+    "FastSpatialPlan",
+    "apply_fast_spatial_upsample",
+    "plan_fast_spatial",
+    "resolve_spatial_mode",
+]

fastvideo/mlx_runtime/fastwan.py

@@ -0,0 +1,980 @@
+# SPDX-License-Identifier: Apache-2.0
+# mypy: disable-error-code=no-untyped-call
+"""FastWan-oriented helpers for the experimental MLX runtime path."""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import statistics
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from collections.abc import Callable
+
+from fastvideo.logger import init_logger
+
+if TYPE_CHECKING:
+    import mlx.core as mx
+    import torch
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class FastWanShape:
+    height: int
+    width: int
+    num_frames: int
+    latent_frames: int
+    latent_height: int
+    latent_width: int
+    patch_frames: int
+    patch_height: int
+    patch_width: int
+    tokens: int
+    hidden_size: int
+    num_heads: int
+    head_dim: int
+
+
+class UnsupportedMLXQuantizationError(ValueError):
+    """A quantization mode the installed MLX build cannot execute.
+
+    Raised by :func:`ensure_quantization_supported` before any model weights
+    are loaded, so callers (CLI flags, benchmark sweeps) can fail fast with an
+    actionable message -- or skip the mode -- instead of crashing deep inside
+    ``mx.quantize`` mid-load.
+    """
+
+
+@dataclass(frozen=True)
+class MLXQuantizationSpec:
+    """MLX quantized-matmul configuration for DiT linear weights."""
+
+    mode: str
+    bits: int | None = None
+    group_size: int | None = None
+
+    @classmethod
+    def from_name(cls, name: str | None) -> MLXQuantizationSpec | None:
+        if name is None or name in {"", "none", "fp16", "fp32"}:
+            return None
+        if name == "int8":
+            return cls(mode="affine", bits=8, group_size=64)
+        if name == "int4":
+            return cls(mode="affine", bits=4, group_size=64)
+        if name == "mxfp8":
+            return cls(mode="mxfp8")
+        if name == "mxfp4":
+            return cls(mode="mxfp4")
+        if name == "nvfp4":
+            return cls(mode="nvfp4")
+        raise ValueError(f"Unsupported MLX quantization mode: {name}")
+
+    @property
+    def label(self) -> str:
+        if self.mode == "affine":
+            return f"int{self.bits}"
+        return self.mode
+
+
+@dataclass(frozen=True)
+class QuantizedMatrix:
+    weight: mx.array
+    scales: mx.array
+    biases: mx.array | None
+    spec: MLXQuantizationSpec
+    dequantized_dtype: mx.Dtype
+
+
+def fastwan_shape(
+        *,
+        height: int,
+        width: int,
+        num_frames: int,
+        vae_temporal_compression: int = 4,
+        vae_spatial_compression: int = 8,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        num_heads: int = 12,
+        head_dim: int = 128,
+) -> FastWanShape:
+    """Return the approximate DiT token shape for Wan/FastWan T2V inference."""
+    latent_frames = (num_frames - 1) // vae_temporal_compression + 1
+    latent_height = height // vae_spatial_compression
+    latent_width = width // vae_spatial_compression
+    patch_frames = latent_frames // patch_size[0]
+    patch_height = latent_height // patch_size[1]
+    patch_width = latent_width // patch_size[2]
+    tokens = patch_frames * patch_height * patch_width
+    return FastWanShape(
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        latent_frames=latent_frames,
+        latent_height=latent_height,
+        latent_width=latent_width,
+        patch_frames=patch_frames,
+        patch_height=patch_height,
+        patch_width=patch_width,
+        tokens=tokens,
+        hidden_size=num_heads * head_dim,
+        num_heads=num_heads,
+        head_dim=head_dim,
+    )
+
+
+def fastwan_shape_from_config(
+    config_path: str | Path,
+    *,
+    height: int,
+    width: int,
+    num_frames: int,
+) -> FastWanShape:
+    config = json.loads(Path(config_path).read_text())
+    return fastwan_shape(
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        patch_size=tuple(config["patch_size"]),
+        num_heads=int(config["num_attention_heads"]),
+        head_dim=int(config["attention_head_dim"]),
+    )
+
+
+def replace_tokens(shape: FastWanShape, tokens: int) -> FastWanShape:
+    return FastWanShape(**{**shape.__dict__, "tokens": tokens})
+
+
+def median_ms(samples: list[float]) -> float:
+    return statistics.median(samples) * 1000.0
+
+
+def benchmark_mlx_attention(shape: FastWanShape, warmup: int, iters: int) -> float:
+    import mlx.core as mx
+
+    q = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
+    k = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
+    v = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
+    scale = shape.head_dim**-0.5
+
+    for _ in range(warmup):
+        y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+        mx.eval(y)
+
+    samples = []
+    for _ in range(iters):
+        start = time.perf_counter()
+        y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+        mx.eval(y)
+        samples.append(time.perf_counter() - start)
+    return median_ms(samples)
+
+
+def benchmark_mlx_linear(shape: FastWanShape, warmup: int, iters: int) -> float:
+    import mlx.core as mx
+
+    x = mx.random.normal((shape.tokens, shape.hidden_size), dtype=mx.float16)
+    w = mx.random.normal((shape.hidden_size, shape.hidden_size), dtype=mx.float16)
+    b = mx.zeros((shape.hidden_size, ), dtype=mx.float16)
+
+    for _ in range(warmup):
+        y = x @ w + b
+        mx.eval(y)
+
+    samples = []
+    for _ in range(iters):
+        start = time.perf_counter()
+        y = x @ w + b
+        mx.eval(y)
+        samples.append(time.perf_counter() - start)
+    return median_ms(samples)
+
+
+def benchmark_torch_mps_attention(shape: FastWanShape, warmup: int, iters: int) -> float | None:
+    try:
+        import torch
+        import torch.nn.functional as F
+    except ImportError:
+        return None
+
+    if not torch.backends.mps.is_available():
+        return None
+
+    device = torch.device("mps")
+    q = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
+    k = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
+    v = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
+
+    for _ in range(warmup):
+        y = F.scaled_dot_product_attention(q, k, v)
+        torch.mps.synchronize()
+        _ = y
+
+    samples = []
+    for _ in range(iters):
+        start = time.perf_counter()
+        y = F.scaled_dot_product_attention(q, k, v)
+        torch.mps.synchronize()
+        _ = y
+        samples.append(time.perf_counter() - start)
+    return median_ms(samples)
+
+
+def torch_to_mx(tensor) -> mx.array:
+    import mlx.core as mx
+
+    return mx.array(tensor.detach().cpu().float().numpy())
+
+
+def weight_dtype(weight):
+    if isinstance(weight, QuantizedMatrix):
+        return weight.dequantized_dtype
+    return weight.dtype
+
+
+_QUANT_SUPPORT_CACHE: dict[tuple[str, int | None, int | None], str | None] = {}
+
+
+def quantization_support_error(spec: MLXQuantizationSpec) -> str | None:
+    """Probe whether the installed MLX build supports ``spec``.
+
+    Runs a tiny ``mx.quantize`` + ``mx.quantized_matmul`` with exactly the
+    arguments :func:`quantize_matrix` / :func:`linear` use, so the result
+    reflects the real runtime path. The affine (int8/int4) modes are stable
+    across MLX releases, but the ``mxfp8``/``mxfp4``/``nvfp4`` mode strings
+    require newer MLX builds and raise otherwise. Returns ``None`` when the
+    mode works, else the underlying error message. Cached per spec.
+    """
+    key = (spec.mode, spec.bits, spec.group_size)
+    if key not in _QUANT_SUPPORT_CACHE:
+        import mlx.core as mx
+
+        try:
+            probe_dim = max(spec.group_size or 0, 64)
+            weight = mx.zeros((probe_dim, probe_dim), dtype=mx.float16)
+            quantized = quantize_matrix(weight, spec)
+            y = linear(mx.zeros((1, probe_dim), dtype=mx.float16), quantized)
+            mx.eval(y)
+            _QUANT_SUPPORT_CACHE[key] = None
+        except Exception as exc:  # noqa: BLE001 - MLX raises varied error types per backend/version.
+            _QUANT_SUPPORT_CACHE[key] = f"{type(exc).__name__}: {exc}"
+    return _QUANT_SUPPORT_CACHE[key]
+
+
+def ensure_quantization_supported(spec: MLXQuantizationSpec | None) -> None:
+    """Raise :class:`UnsupportedMLXQuantizationError` if ``spec`` cannot run here."""
+    if spec is None:
+        return
+    error = quantization_support_error(spec)
+    if error is None:
+        return
+    import mlx.core as mx
+
+    mlx_version = getattr(mx, "__version__", "unknown")
+    raise UnsupportedMLXQuantizationError(f"MLX quantization mode '{spec.label}' is not supported by the installed mlx "
+                                          f"({mlx_version}): {error}. Upgrade mlx or pick a supported mode "
+                                          f"(int8 is currently the most reliable quality/memory target).")
+
+
+def quantize_matrix(weight, spec: MLXQuantizationSpec | None):
+    if spec is None:
+        return weight
+    import mlx.core as mx
+
+    if len(weight.shape) < 2:
+        return weight
+    q = mx.quantize(weight, group_size=spec.group_size, bits=spec.bits, mode=spec.mode)
+    biases = q[2] if len(q) == 3 else None
+    eval_args = [q[0], q[1]]
+    if biases is not None:
+        eval_args.append(biases)
+    mx.eval(*eval_args)
+    return QuantizedMatrix(
+        weight=q[0],
+        scales=q[1],
+        biases=biases,
+        spec=spec,
+        dequantized_dtype=weight.dtype,
+    )
+
+
+def linear(x, weight, bias=None):
+    import mlx.core as mx
+
+    if isinstance(weight, QuantizedMatrix):
+        y = mx.quantized_matmul(
+            x,
+            weight.weight,
+            weight.scales,
+            weight.biases,
+            transpose=True,
+            group_size=weight.spec.group_size,
+            bits=weight.spec.bits,
+            mode=weight.spec.mode,
+        ).astype(x.dtype)
+    else:
+        y = x @ weight.T
+    if bias is not None:
+        y = y + bias
+    return y
+
+
+def _use_fast_norm() -> bool:
+    """Opt-in to MLX's fused ``mx.fast`` normalization kernels.
+
+    Off by default so the numerically-explicit reference path stays the
+    baseline. Set ``FASTVIDEO_MLX_FAST_NORM=1`` to route LayerNorm/RMSNorm
+    through single fused Metal kernels (fewer intermediates, less memory
+    traffic) and benchmark the speedup.
+    """
+    import os
+
+    return os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1"
+
+
+def layer_norm(x, weight=None, bias=None, eps: float = 1e-6):
+    import mlx.core as mx
+
+    if _use_fast_norm():
+        # Compute in fp32 (matching the reference below) so downstream dtype
+        # and precision are identical across call sites.
+        w = weight.astype(mx.float32) if weight is not None else None
+        b = bias.astype(mx.float32) if bias is not None else None
+        return mx.fast.layer_norm(x.astype(mx.float32), w, b, eps)
+
+    x_float = x.astype(mx.float32)
+    mean = mx.mean(x_float, axis=-1, keepdims=True)
+    var = mx.mean(mx.square(x_float - mean), axis=-1, keepdims=True)
+    y = (x_float - mean) * mx.rsqrt(var + eps)
+    if weight is not None:
+        y = y * weight
+    if bias is not None:
+        y = y + bias
+    return y
+
+
+def rms_norm(x, weight, eps: float = 1e-6):
+    import mlx.core as mx
+
+    if _use_fast_norm():
+        return mx.fast.rms_norm(x, weight, eps)
+
+    orig_dtype = x.dtype
+    x_float = x.astype(mx.float32)
+    variance = mx.mean(mx.square(x_float), axis=-1, keepdims=True)
+    y = x_float * mx.rsqrt(variance + eps)
+    return y.astype(orig_dtype) * weight
+
+
+def apply_rotary_emb(x, cos, sin, *, is_neox_style: bool = False):
+    """Apply FastVideo's rotary convention to MLX tensors.
+
+    Args:
+        x: [batch, seq, heads, head_dim]
+        cos/sin: [seq, head_dim] for Wan's full-dimension rotate-pair style,
+          or [seq, head_dim // 2] for traditional RoPE.
+    """
+    import mlx.core as mx
+
+    head_size = x.shape[-1]
+    rope_dim = cos.shape[-1]
+    cos = cos[None, :, None, :]
+    sin = sin[None, :, None, :]
+    x_float = x.astype(mx.float32)
+
+    if rope_dim == head_size:
+        x_pairs = x_float.reshape(*x.shape[:-1], -1, 2)
+        x_real = x_pairs[..., 0]
+        x_imag = x_pairs[..., 1]
+        x_rotated = mx.stack([-x_imag, x_real], axis=-1).reshape(*x.shape)
+        return (x_float * cos + x_rotated * sin).astype(x.dtype)
+
+    if is_neox_style:
+        x1, x2 = mx.split(x_float, 2, axis=-1)
+        o1 = x1 * cos - x2 * sin
+        o2 = x2 * cos + x1 * sin
+        return mx.concatenate([o1, o2], axis=-1).astype(x.dtype)
+
+    x1 = x_float[..., ::2]
+    x2 = x_float[..., 1::2]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    return mx.stack([o1, o2], axis=-1).reshape(*x.shape).astype(x.dtype)
+
+
+_WINDOWED_ATTENTION_WARNED = False
+
+
+def _warn_windowed_attention_once(window: int) -> None:
+    """Warn that FASTVIDEO_MLX_WINDOW degrades output on a dense-trained DiT.
+
+    Sliding-window self-attention is fast (6.6x at a +-3-frame window on 1.3B)
+    but these checkpoints were trained with dense attention, and restricting it
+    at inference produces heavy colour-block noise: structural agreement with
+    the dense baseline drops to 0.25 at +-3 frames and 0.03 at +-5. Sparsity of
+    this kind is a training-time method. Kept as a research knob, but it should
+    never be on by accident.
+    """
+    global _WINDOWED_ATTENTION_WARNED
+    if _WINDOWED_ATTENTION_WARNED:
+        return
+    _WINDOWED_ATTENTION_WARNED = True
+    logger.warning(
+        "FASTVIDEO_MLX_WINDOW=%d enables sliding-window self-attention. These "
+        "checkpoints are trained dense; expect severely degraded output. This is "
+        "a research knob, not a speed setting — use --fast-spatial for real "
+        "denoise savings.",
+        window,
+    )
+
+
+def gelu_tanh(x):
+    """tanh-approximate GELU, as used by Wan's FFN.
+
+    ``mlx.nn.gelu_approx`` is the same tanh approximation behind a fused
+    kernel. On the 1.3B FFN shape (32760x8960) it is bit-identical to the
+    expanded expression below and 3.3x faster — 28.9ms -> 8.7ms per layer,
+    which is 0.6s per denoise step across 30 layers.
+    """
+    import mlx.nn as nn
+
+    return nn.gelu_approx(x)
+
+
+def silu(x):
+    import mlx.core as mx
+
+    return x * mx.sigmoid(x)
+
+
+def timestep_embedding(t, dim: int, max_period: int = 10000):
+    import mlx.core as mx
+
+    half = dim // 2
+    freqs = mx.exp(-math.log(max_period) * mx.arange(0, half, dtype=mx.float32) / half)
+    args = t[:, None].astype(mx.float32) * freqs[None]
+    embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
+    if dim % 2:
+        embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1)
+    return embedding
+
+
+def scale_residual(residual, x, gate):
+    return residual + x * gate
+
+
+def scale_residual_layer_norm_scale_shift(residual, x, gate, shift, scale, weight=None, bias=None, eps: float = 1e-6):
+    if isinstance(gate, int):
+        assert gate == 1
+        residual_output = residual + x
+    else:
+        residual_output = residual + x * gate
+    normalized = layer_norm(residual_output, weight=weight, bias=bias, eps=eps)
+    modulated = normalized * (1.0 + scale) + shift
+    return modulated, residual_output
+
+
+class MLXWanT2VCrossAttention:
+
+    def __init__(self, weights: dict[str, mx.array], *, dim: int, num_heads: int, eps: float = 1e-6) -> None:
+        self.weights = weights
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.eps = eps
+
+    def __call__(self, x, context):
+        import mlx.core as mx
+
+        batch = x.shape[0]
+        q = linear(x, self.weights["attn2.to_q.weight"], self.weights.get("attn2.to_q.bias"))
+        q = rms_norm(q, self.weights["attn2.norm_q.weight"], eps=self.eps).reshape(batch, -1, self.num_heads,
+                                                                                   self.head_dim)
+
+        if context.shape[1] == 0:
+            attended = mx.zeros_like(q)
+        else:
+            k = linear(context, self.weights["attn2.to_k.weight"], self.weights.get("attn2.to_k.bias"))
+            k = rms_norm(k, self.weights["attn2.norm_k.weight"],
+                         eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim)
+            v = linear(context, self.weights["attn2.to_v.weight"],
+                       self.weights.get("attn2.to_v.bias")).reshape(batch, -1, self.num_heads, self.head_dim)
+            attended = mx.fast.scaled_dot_product_attention(
+                q.transpose(0, 2, 1, 3),
+                k.transpose(0, 2, 1, 3),
+                v.transpose(0, 2, 1, 3),
+                scale=self.head_dim**-0.5,
+            ).transpose(0, 2, 1, 3)
+
+        attended = attended.reshape(batch, -1, self.dim)
+        return linear(attended, self.weights["attn2.to_out.weight"], self.weights.get("attn2.to_out.bias"))
+
+
+class MLXWanTransformerBlock:
+    """Dense T2V Wan transformer block for the experimental MLX runtime.
+
+    This mirrors the non-VSA PyTorch block for single-process dense attention.
+    Rotary embeddings and sequence-parallel paths are intentionally left out of
+    this first parity target.
+    """
+
+    def __init__(self, weights: dict[str, mx.array], *, dim: int, ffn_dim: int, num_heads: int, eps: float = 1e-6):
+        self.weights = weights
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.eps = eps
+        self.attn2 = MLXWanT2VCrossAttention(weights, dim=dim, num_heads=num_heads, eps=eps)
+
+    def __call__(self, hidden_states, encoder_hidden_states, temb, freqs_cis=None):
+        import mlx.core as mx
+
+        orig_dtype = hidden_states.dtype
+        e = self.weights["scale_shift_table"] + temb.astype(mx.float32)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = mx.split(e, 6, axis=1)
+
+        norm_hidden_states = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
+        norm_hidden_states = (norm_hidden_states * (1.0 + scale_msa) + shift_msa).astype(orig_dtype)
+
+        query = linear(norm_hidden_states, self.weights["to_q.weight"], self.weights.get("to_q.bias"))
+        key = linear(norm_hidden_states, self.weights["to_k.weight"], self.weights.get("to_k.bias"))
+        value = linear(norm_hidden_states, self.weights["to_v.weight"], self.weights.get("to_v.bias"))
+
+        query = rms_norm(query, self.weights["norm_q.weight"],
+                         eps=self.eps).reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim)
+        key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(hidden_states.shape[0], -1,
+                                                                                 self.num_heads, self.head_dim)
+        value = value.reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim)
+
+        if freqs_cis is not None:
+            cos, sin = freqs_cis
+            query = apply_rotary_emb(query, cos, sin, is_neox_style=False)
+            key = apply_rotary_emb(key, cos, sin, is_neox_style=False)
+
+        # Self-attention only. FASTVIDEO_MLX_WINDOW=0/unset → full SDPA (byte-identical
+        # to the historical path). When >0, use chunked symmetric sliding-window
+        # attention (see windowed_attention.py). Cross-attn (attn2) stays full.
+        # Optional FASTVIDEO_MLX_WINDOW_SINK (default 0) adds global sink tokens.
+        q_bh = query.transpose(0, 2, 1, 3)  # (B, H, S, D)
+        k_bh = key.transpose(0, 2, 1, 3)
+        v_bh = value.transpose(0, 2, 1, 3)
+        scale = self.head_dim**-0.5
+        window = int(os.environ.get("FASTVIDEO_MLX_WINDOW", "0") or "0")
+        if window > 0:
+            from fastvideo.mlx_runtime.windowed_attention import windowed_attention
+
+            _warn_windowed_attention_once(window)
+            sink = int(os.environ.get("FASTVIDEO_MLX_WINDOW_SINK", "0") or "0")
+            attn_output = windowed_attention(q_bh, k_bh, v_bh, window=window, sink=sink, scale=scale)
+        else:
+            attn_output = mx.fast.scaled_dot_product_attention(q_bh, k_bh, v_bh, scale=scale)
+        attn_output = attn_output.transpose(0, 2, 1, 3)
+        attn_output = attn_output.reshape(hidden_states.shape[0], -1, self.dim)
+        attn_output = linear(attn_output, self.weights["to_out.weight"], self.weights.get("to_out.bias"))
+
+        norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift(
+            hidden_states,
+            attn_output,
+            gate_msa,
+            0.0,
+            0.0,
+            weight=self.weights["self_attn_residual_norm.norm.weight"],
+            bias=self.weights["self_attn_residual_norm.norm.bias"],
+            eps=self.eps,
+        )
+        norm_hidden_states = norm_hidden_states.astype(orig_dtype)
+        hidden_states = hidden_states.astype(orig_dtype)
+
+        attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
+        norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift(
+            hidden_states,
+            attn_output,
+            1,
+            c_shift_msa,
+            c_scale_msa,
+            eps=self.eps,
+        )
+        norm_hidden_states = norm_hidden_states.astype(orig_dtype)
+        hidden_states = hidden_states.astype(orig_dtype)
+
+        ff_output = linear(norm_hidden_states, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias"))
+        ff_output = gelu_tanh(ff_output)
+        ff_output = linear(ff_output, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias"))
+        hidden_states = scale_residual(hidden_states, ff_output, c_gate_msa)
+        return hidden_states.astype(orig_dtype)
+
+
+def mlx_block_weights_from_torch(torch_block) -> dict[str, mx.array]:
+    return {name: torch_to_mx(value) for name, value in torch_block.state_dict().items()}
+
+
+class MLXWanDiT:
+    """Experimental FP16 Wan/FastWan DiT forward path in MLX."""
+
+    def __init__(
+        self,
+        weights: dict[str, mx.array],
+        blocks: list[MLXWanTransformerBlock],
+        config: dict,
+        *,
+        compile: bool = False,
+    ) -> None:
+        import os
+
+        self.weights = weights
+        self.blocks = blocks
+        self.config = config
+        self.num_heads = int(config["num_attention_heads"])
+        self.head_dim = int(config["attention_head_dim"])
+        self.hidden_size = self.num_heads * self.head_dim
+        self.ffn_dim = int(config["ffn_dim"])
+        self.in_channels = int(config["in_channels"])
+        self.out_channels = int(config["out_channels"])
+        self.patch_size = tuple(config["patch_size"])
+        self.freq_dim = int(config["freq_dim"])
+        # Opt-in graph fusion. With fixed weights and static shapes, the whole
+        # denoise-step forward is a pure function of (latents, timestep) -- a
+        # good mx.compile target. Off by default so the eager path stays the
+        # baseline; enable via constructor or FASTVIDEO_MLX_COMPILE=1 and verify
+        # with the benchmark's SSIM ~= 1.0 check.
+        self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1"
+        self._compiled_forward: Callable[..., Any] | None = None
+        self._compiled_signature: tuple | None = None
+
+    def patch_embed(self, hidden_states):
+        batch, channels, frames, height, width = hidden_states.shape
+        pt, ph, pw = self.patch_size
+        patch_dim = channels * pt * ph * pw
+        x = hidden_states.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw)
+        x = x.transpose(0, 2, 4, 6, 1, 3, 5, 7).reshape(batch, -1, patch_dim)
+        return linear(x, self.weights["patch_embedding.weight"], self.weights.get("patch_embedding.bias"))
+
+    def condition(self, timestep, encoder_hidden_states):
+        t_freq = timestep_embedding(timestep, self.freq_dim).astype(
+            weight_dtype(self.weights["condition_embedder.time_embedder.linear_1.weight"]))
+        temb = linear(
+            t_freq,
+            self.weights["condition_embedder.time_embedder.linear_1.weight"],
+            self.weights["condition_embedder.time_embedder.linear_1.bias"],
+        )
+        temb = silu(temb)
+        temb = linear(
+            temb,
+            self.weights["condition_embedder.time_embedder.linear_2.weight"],
+            self.weights["condition_embedder.time_embedder.linear_2.bias"],
+        )
+        timestep_proj = silu(temb)
+        timestep_proj = linear(
+            timestep_proj,
+            self.weights["condition_embedder.time_proj.weight"],
+            self.weights["condition_embedder.time_proj.bias"],
+        ).reshape(timestep.shape[0], 6, self.hidden_size)
+
+        encoder_hidden_states = linear(
+            encoder_hidden_states,
+            self.weights["condition_embedder.text_embedder.linear_1.weight"],
+            self.weights["condition_embedder.text_embedder.linear_1.bias"],
+        )
+        encoder_hidden_states = gelu_tanh(encoder_hidden_states)
+        encoder_hidden_states = linear(
+            encoder_hidden_states,
+            self.weights["condition_embedder.text_embedder.linear_2.weight"],
+            self.weights["condition_embedder.text_embedder.linear_2.bias"],
+        )
+        return temb, timestep_proj, encoder_hidden_states
+
+    def output(self, hidden_states, temb, *, batch: int, frames: int, height: int, width: int):
+        pt, ph, pw = self.patch_size
+        post_patch_frames = frames // pt
+        post_patch_height = height // ph
+        post_patch_width = width // pw
+        shift, scale = mx_split_two(self.weights["scale_shift_table"] + temb[:, None, :], axis=1)
+        hidden_states = layer_norm(hidden_states, eps=float(self.config["eps"])) * (1.0 + scale) + shift
+        hidden_states = hidden_states.astype(weight_dtype(self.weights["proj_out.weight"]))
+        hidden_states = linear(hidden_states, self.weights["proj_out.weight"], self.weights["proj_out.bias"])
+        hidden_states = hidden_states.reshape(
+            batch,
+            post_patch_frames,
+            post_patch_height,
+            post_patch_width,
+            pt,
+            ph,
+            pw,
+            self.out_channels,
+        )
+        hidden_states = hidden_states.transpose(0, 7, 1, 4, 2, 5, 3, 6)
+        return hidden_states.reshape(batch, self.out_channels, frames, height, width)
+
+    def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin):
+        """Pure forward used both eagerly and as the mx.compile target.
+
+        ``cos``/``sin`` are passed as separate array args (rather than a tuple)
+        so the function traces cleanly under mx.compile.
+        """
+        batch, _, frames, height, width = hidden_states.shape
+        freqs_cis = (cos, sin) if cos is not None else None
+        hidden_states = self.patch_embed(hidden_states)
+        temb, timestep_proj, encoder_hidden_states = self.condition(timestep, encoder_hidden_states)
+        for block in self.blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis=freqs_cis)
+        return self.output(hidden_states, temb, batch=batch, frames=frames, height=height, width=width)
+
+    def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis):
+        cos, sin = freqs_cis if freqs_cis is not None else (None, None)
+        if self._enable_compile and cos is not None:
+            import mlx.core as mx
+
+            # mx.compile keeps one traced graph per input signature, and each
+            # graph pins its own materialization of the quantized weights. The
+            # two-pass modes (--refine) call the DiT at a second resolution, so
+            # keeping both graphs alive doubles resident DiT memory: 14B refine
+            # peaked at 34.7 GiB instead of 20.8 GiB. Retire the previous graph
+            # when the signature changes; the retrace costs far less than a
+            # second copy of the weights.
+            signature = (hidden_states.shape, encoder_hidden_states.shape, timestep.shape)
+            if self._compiled_forward is not None and signature != self._compiled_signature:
+                self._compiled_forward = None
+                self._compiled_signature = None
+                mx.clear_cache()
+            if self._compiled_forward is None:
+                self._compiled_forward = mx.compile(self._forward)
+                self._compiled_signature = signature
+            compiled_forward = self._compiled_forward
+            try:
+                return compiled_forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
+            except Exception as exc:  # noqa: BLE001 - some quant graphs may not trace; fall back to eager.
+                logger.warning("mx.compile forward failed (%s); falling back to eager execution.", exc)
+                self._enable_compile = False
+                self._compiled_forward = None
+                self._compiled_signature = None
+        return self._forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
+
+
+def mx_split_two(x, *, axis: int):
+    import mlx.core as mx
+
+    left, right = mx.split(x, 2, axis=axis)
+    return left, right
+
+
+def _load_safetensor_value(handle, name: str):
+    return handle.get_tensor(name)
+
+
+def _load_mx_array_from_safetensor(handle, name: str, dtype):
+    """Load a safetensors value and cast before creating the MLX array.
+
+    The FastWan Diffusers checkpoint is fp32. Creating an MLX array first and
+    then casting it to fp16 briefly materializes a large fp32 MLX allocation.
+    Casting the CPU tensor before crossing into MLX keeps the transient GPU-side
+    footprint lower.
+    """
+    import mlx.core as mx
+    import torch
+
+    tensor = handle.get_tensor(name)
+    if dtype == mx.float16:
+        tensor = tensor.to(torch.float16)
+    elif dtype == mx.float32:
+        tensor = tensor.to(torch.float32)
+    elif dtype == mx.bfloat16:
+        # NumPy has no bfloat16, so bridge through fp32 and cast on-device below.
+        tensor = tensor.to(torch.float32)
+    array = mx.array(tensor.numpy())
+    del tensor
+    if dtype is not None and array.dtype != dtype:
+        array = array.astype(dtype)
+    mx.eval(array)
+    return array
+
+
+def _eval_loaded_weight(value) -> None:
+    import mlx.core as mx
+
+    if isinstance(value, QuantizedMatrix):
+        eval_args = [value.weight, value.scales]
+        if value.biases is not None:
+            eval_args.append(value.biases)
+        mx.eval(*eval_args)
+    else:
+        mx.eval(value)
+
+
+# Diffusers-to-FastVideo key mapping for WanTransformerBlock weights.
+# Shared by both MLX and torch block loaders to keep mappings synchronized.
+_WAN_BLOCK_KEY_MAP = {
+    "scale_shift_table": "scale_shift_table",
+    "attn1.to_q.weight": "to_q.weight",
+    "attn1.to_q.bias": "to_q.bias",
+    "attn1.to_k.weight": "to_k.weight",
+    "attn1.to_k.bias": "to_k.bias",
+    "attn1.to_v.weight": "to_v.weight",
+    "attn1.to_v.bias": "to_v.bias",
+    "attn1.to_out.0.weight": "to_out.weight",
+    "attn1.to_out.0.bias": "to_out.bias",
+    "attn1.norm_q.weight": "norm_q.weight",
+    "attn1.norm_k.weight": "norm_k.weight",
+    "attn2.to_q.weight": "attn2.to_q.weight",
+    "attn2.to_q.bias": "attn2.to_q.bias",
+    "attn2.to_k.weight": "attn2.to_k.weight",
+    "attn2.to_k.bias": "attn2.to_k.bias",
+    "attn2.to_v.weight": "attn2.to_v.weight",
+    "attn2.to_v.bias": "attn2.to_v.bias",
+    "attn2.to_out.0.weight": "attn2.to_out.weight",
+    "attn2.to_out.0.bias": "attn2.to_out.bias",
+    "attn2.norm_q.weight": "attn2.norm_q.weight",
+    "attn2.norm_k.weight": "attn2.norm_k.weight",
+    "ffn.net.0.proj.weight": "ffn.fc_in.weight",
+    "ffn.net.0.proj.bias": "ffn.fc_in.bias",
+    "ffn.net.2.weight": "ffn.fc_out.weight",
+    "ffn.net.2.bias": "ffn.fc_out.bias",
+    "norm2.weight": "self_attn_residual_norm.norm.weight",
+    "norm2.bias": "self_attn_residual_norm.norm.bias",
+}
+
+
+def mlx_block_weights_from_diffusers_safetensors(
+    checkpoint_path: str | Path,
+    *,
+    block_index: int = 0,
+    quantization: str | MLXQuantizationSpec | None = None,
+    dtype=None,
+) -> dict[str, mx.array]:
+    """Load one Diffusers-format Wan block into the MLX dense-block key layout."""
+    from safetensors import safe_open
+
+    prefix = f"blocks.{block_index}."
+    key_map = _WAN_BLOCK_KEY_MAP
+
+    spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
+                                                           or isinstance(quantization, str)) else quantization
+    ensure_quantization_supported(spec)
+    matrix_targets = {target for target in key_map.values() if target.endswith(".weight") and "norm" not in target}
+    weights = {}
+    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
+        available = set(handle.keys())
+        for source_name, target_name in key_map.items():
+            full = prefix + source_name
+            if full not in available:
+                # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
+                # The block forward already fetches biases via ``.get(...)``.
+                if source_name.endswith(".bias"):
+                    continue
+                raise KeyError(f"missing required block weight: {full}")
+            array = _load_mx_array_from_safetensor(handle, full, dtype)
+            loaded = quantize_matrix(array, spec) if target_name in matrix_targets else array
+            _eval_loaded_weight(loaded)
+            weights[target_name] = loaded
+            del array
+    return weights
+
+
+def mlx_dit_from_diffusers_safetensors(
+    checkpoint_path: str | Path,
+    config_path: str | Path,
+    *,
+    dtype: str = "fp16",
+    num_blocks: int | None = None,
+    quantization: str | MLXQuantizationSpec | None = None,
+    compile: bool = False,
+) -> MLXWanDiT:
+    import mlx.core as mx
+    from safetensors import safe_open
+
+    config = json.loads(Path(config_path).read_text())
+    total_blocks = int(config["num_layers"])
+    if num_blocks is None:
+        num_blocks = total_blocks
+    cast_dtype = {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[dtype]
+    spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
+                                                           or isinstance(quantization, str)) else quantization
+    ensure_quantization_supported(spec)
+
+    top_level_names = [
+        "patch_embedding.weight",
+        "patch_embedding.bias",
+        "condition_embedder.time_embedder.linear_1.weight",
+        "condition_embedder.time_embedder.linear_1.bias",
+        "condition_embedder.time_embedder.linear_2.weight",
+        "condition_embedder.time_embedder.linear_2.bias",
+        "condition_embedder.time_proj.weight",
+        "condition_embedder.time_proj.bias",
+        "condition_embedder.text_embedder.linear_1.weight",
+        "condition_embedder.text_embedder.linear_1.bias",
+        "condition_embedder.text_embedder.linear_2.weight",
+        "condition_embedder.text_embedder.linear_2.bias",
+        "scale_shift_table",
+        "proj_out.weight",
+        "proj_out.bias",
+    ]
+    weights = {}
+    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
+        available = set(handle.keys())
+        for name in top_level_names:
+            if name not in available:
+                if name.endswith(".bias"):
+                    continue
+                raise KeyError(f"missing required weight: {name}")
+            array = _load_mx_array_from_safetensor(handle, name, cast_dtype)
+            if name == "patch_embedding.weight":
+                array = array.reshape(int(config["num_attention_heads"]) * int(config["attention_head_dim"]), -1)
+            if name.endswith(".weight") and name not in {"scale_shift_table"}:
+                loaded = quantize_matrix(array, spec)
+            else:
+                loaded = array
+            _eval_loaded_weight(loaded)
+            weights[name] = loaded
+            del array
+
+    blocks = []
+    for block_index in range(num_blocks):
+        block_weights = mlx_block_weights_from_diffusers_safetensors(
+            checkpoint_path,
+            block_index=block_index,
+            quantization=spec,
+            dtype=cast_dtype,
+        )
+        block_weights = {
+            name: (value if isinstance(value, QuantizedMatrix) else value.astype(cast_dtype))
+            for name, value in block_weights.items()
+        }
+        for value in block_weights.values():
+            _eval_loaded_weight(value)
+        blocks.append(
+            MLXWanTransformerBlock(
+                block_weights,
+                dim=int(config["num_attention_heads"]) * int(config["attention_head_dim"]),
+                ffn_dim=int(config["ffn_dim"]),
+                num_heads=int(config["num_attention_heads"]),
+                eps=float(config["eps"]),
+            ))
+    return MLXWanDiT(weights, blocks, config, compile=compile)
+
+
+def torch_block_state_from_diffusers_safetensors(
+    checkpoint_path: str | Path,
+    *,
+    block_index: int = 0,
+) -> dict[str, torch.Tensor]:
+    """Load one Diffusers-format Wan block into FastVideo's dense block keys."""
+    from safetensors import safe_open
+
+    prefix = f"blocks.{block_index}."
+    key_map = _WAN_BLOCK_KEY_MAP
+
+    state = {}
+    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
+        available = set(handle.keys())
+        for source_name, target_name in key_map.items():
+            full = prefix + source_name
+            if full not in available:
+                # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
+                if source_name.endswith(".bias"):
+                    continue
+                raise KeyError(f"missing required block weight: {full}")
+            state[target_name] = handle.get_tensor(full).float()
+    return state

fastvideo/mlx_runtime/frame_upsample.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pixel-space spatial resampling for decoded MLX Wan frames.
+
+Spatial fast mode denoises on a smaller latent grid and has to get back to
+the requested output size. That resize belongs *here* — after the VAE
+decode — and not in latent space.
+
+A Wan latent cell is a learned code for an 8x8 (Wan2.1) or 16x16 (Wan2.2)
+pixel block, not a low-pass sample of the image. Linearly blending two
+adjacent codes does not produce the code of the blended blocks; it produces
+a vector the decoder was never trained on. The decoder answers with smeared,
+ringing texture laid over otherwise-correct structure — the silhouette
+survives, the detail turns to haze. Measured on Wan2.1-1.3B at 480x832, a
+2x bilinear latent upsample destroys 62% of the latent's high-frequency
+energy while leaving its overall magnitude intact, which is exactly the
+signature of that veil.
+
+Resampling decoded RGB frames has none of that problem: an image *is* a
+sampled 2-D signal, so Lanczos/cubic interpolation is the operation it was
+defined for. The result is soft — it carries stage-1's real detail budget
+and no more — but it is clean and coherent.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+
+# Pixel-space interpolation kernels, best-quality first. ``lanczos`` is the
+# default: it holds edges better than cubic at 2x with no visible ringing on
+# decoder output, which is already band-limited.
+PIXEL_UPSAMPLE_MODES = ("lanczos", "cubic", "bilinear", "nearest")
+
+DEFAULT_PIXEL_UPSAMPLE_MODE = "lanczos"
+
+
+def _interpolation_flag(mode: str) -> int:
+    """
+    Map a pixel upsample mode name onto its OpenCV interpolation flag.
+
+    Parameters:
+        mode (str): One of :data:`PIXEL_UPSAMPLE_MODES`.
+
+    Returns:
+        int: The matching ``cv2.INTER_*`` flag.
+
+    Raises:
+        ValueError: If the mode is not a supported pixel upsample mode.
+    """
+    import cv2
+
+    flags = {
+        "lanczos": cv2.INTER_LANCZOS4,
+        "cubic": cv2.INTER_CUBIC,
+        "bilinear": cv2.INTER_LINEAR,
+        "nearest": cv2.INTER_NEAREST,
+    }
+    try:
+        return flags[mode]
+    except KeyError:
+        raise ValueError(f"Unsupported pixel upsample mode: {mode!r} "
+                         f"(expected one of {', '.join(PIXEL_UPSAMPLE_MODES)})") from None
+
+
+def unsharp(frame: np.ndarray, amount: float) -> np.ndarray:
+    """Light unsharp mask, used to counter resampling / optical-flow softening.
+
+    Parameters:
+        frame (np.ndarray): HxWx3 uint8 RGB frame.
+        amount (float): Strength; ``0`` returns the frame unchanged.
+
+    Returns:
+        np.ndarray: A new frame; the input is never modified in place.
+    """
+    if amount <= 0.0:
+        return frame
+    import cv2
+
+    blur = cv2.GaussianBlur(frame, (0, 0), 1.0)
+    return cv2.addWeighted(frame, 1.0 + amount, blur, -amount, 0)
+
+
+def upsample_frame(
+    frame: np.ndarray,
+    *,
+    width: int,
+    height: int,
+    mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE,
+    sharpen: float = 0.0,
+) -> np.ndarray:
+    """
+    Resample one decoded RGB frame to the target pixel size.
+
+    Parameters:
+        frame (np.ndarray): HxWx3 uint8 RGB frame.
+        width (int): Target width in pixels.
+        height (int): Target height in pixels.
+        mode (str): Interpolation kernel, one of :data:`PIXEL_UPSAMPLE_MODES`.
+        sharpen (float): Unsharp strength applied after the resize.
+
+    Returns:
+        np.ndarray: A new frame at ``height x width``; already-correct sizes
+            are still passed through ``sharpen``.
+
+    Raises:
+        ValueError: If the frame is not HxWx3, or the target size is not positive.
+    """
+    import cv2
+
+    array = np.asarray(frame)
+    if array.ndim != 3 or array.shape[2] != 3:
+        raise ValueError(f"frame must have shape HxWx3, got {array.shape}")
+    if width <= 0 or height <= 0:
+        raise ValueError(f"target size must be positive, got {width}x{height}")
+    if array.dtype != np.uint8:
+        array = np.clip(array, 0, 255).astype(np.uint8)
+
+    if (array.shape[0], array.shape[1]) != (height, width):
+        array = cv2.resize(array, (width, height), interpolation=_interpolation_flag(mode))
+    return unsharp(array, sharpen)
+
+
+def upsample_frames(
+    frames: Iterable[np.ndarray],
+    *,
+    width: int,
+    height: int,
+    mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE,
+    sharpen: float = 0.0,
+) -> list[np.ndarray]:
+    """
+    Resample every decoded frame to the target pixel size.
+
+    Parameters:
+        frames (Iterable[np.ndarray]): Decoded HxWx3 uint8 RGB frames.
+        width (int): Target width in pixels.
+        height (int): Target height in pixels.
+        mode (str): Interpolation kernel, one of :data:`PIXEL_UPSAMPLE_MODES`.
+        sharpen (float): Unsharp strength applied after each resize.
+
+    Returns:
+        list[np.ndarray]: New frames at the target size, in input order.
+    """
+    return [upsample_frame(frame, width=width, height=height, mode=mode, sharpen=sharpen) for frame in frames]
+
+
+__all__ = [
+    "DEFAULT_PIXEL_UPSAMPLE_MODE",
+    "PIXEL_UPSAMPLE_MODES",
+    "unsharp",
+    "upsample_frame",
+    "upsample_frames",
+]

fastvideo/mlx_runtime/memory.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Memory-tier helpers for Apple Silicon MLX/MPS experiments.
+
+macOS does not expose a perfect "pretend this machine only has 16 GB unified
+memory" switch. MLX can cap the allocator used by the Apple-native DiT path,
+and PyTorch MPS exposes process-level watermark environment variables for the
+hybrid prompt/decode stages. Applying both gives benchmark and generation
+entrypoints a practical, explicit way to exercise memory-tier presets.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+GIB = 1024**3
+
+
+@dataclass(frozen=True)
+class AppliedMemoryLimits:
+    """Memory limits applied for one Apple Silicon benchmark/generation process."""
+
+    mlx_memory_limit_gib: float | None = None
+    mlx_cache_limit_gib: float | None = None
+    mlx_disable_cache: bool = False
+    mlx_wired_limit_gib: float | None = None
+    torch_mps_high_watermark_ratio: float | None = None
+    torch_mps_low_watermark_ratio: float | None = None
+    applied_bytes: dict[str, int] = field(default_factory=dict)
+    previous_bytes: dict[str, int] = field(default_factory=dict)
+    errors: dict[str, str] = field(default_factory=dict)
+
+    def as_metrics(self) -> dict[str, int | float | str | bool | None]:
+        """Flatten the configured memory limits, applied values, previous values, and errors into a metrics dictionary.
+
+        Returns:
+            dict[str, int | float | str | bool | None]: Metrics keyed by limit names and their corresponding values.
+        """
+        metrics: dict[str, int | float | str | bool | None] = {
+            "mlx_memory_limit_gib": self.mlx_memory_limit_gib,
+            "mlx_cache_limit_gib": self.mlx_cache_limit_gib,
+            "mlx_disable_cache": self.mlx_disable_cache,
+            "mlx_wired_limit_gib": self.mlx_wired_limit_gib,
+            "torch_mps_high_watermark_ratio": self.torch_mps_high_watermark_ratio,
+            "torch_mps_low_watermark_ratio": self.torch_mps_low_watermark_ratio,
+        }
+        for name, value in self.applied_bytes.items():
+            metrics[f"{name}_bytes"] = value
+        for name, value in self.previous_bytes.items():
+            metrics[f"previous_{name}_bytes"] = value
+        for name, error in self.errors.items():
+            metrics[f"{name}_error"] = error
+        return metrics
+
+
+def gib_to_bytes(value: float | None) -> int | None:
+    """
+    Convert a positive memory limit from GiB to bytes.
+
+    Parameters:
+        value (float | None): Memory limit in GiB, or `None` when unset.
+
+    Returns:
+        int | None: The memory limit in bytes, or `None` when no limit is provided.
+
+    Raises:
+        ValueError: If `value` is zero or negative.
+    """
+    if value is None:
+        return None
+    if value <= 0:
+        raise ValueError(f"Memory limit must be positive GiB, got {value}")
+    return int(value * GIB)
+
+
+def cleanup_mlx(mx_module: Any | None = None) -> None:
+    """Collect unreachable MLX objects, then release their allocator cache."""
+    if mx_module is None:
+        import mlx.core as mx
+
+        mx_module = mx
+    gc.collect()
+    mx_module.clear_cache()
+
+
+def cleanup_torch_mps(torch_module: Any | None = None) -> None:
+    """Collect unreachable Torch objects, then release the MPS allocator cache."""
+    if torch_module is None:
+        import torch
+
+        torch_module = torch
+    gc.collect()
+    if torch_module.backends.mps.is_available():
+        torch_module.mps.empty_cache()
+
+
+def _set_mps_env(name: str, value: float | None) -> float | None:
+    """Set a PyTorch MPS watermark environment variable.
+
+    Parameters:
+        name (str): Name of the environment variable to set.
+        value (float | None): Watermark ratio, or `None` to leave the variable unchanged.
+
+    Returns:
+        float | None: The configured watermark ratio, or `None` when no value is provided.
+
+    Raises:
+        ValueError: If `value` is negative.
+    """
+    if value is None:
+        return None
+    if value < 0:
+        raise ValueError(f"{name} must be non-negative, got {value}")
+    os.environ[name] = str(value)
+    return value
+
+
+def apply_memory_limits(
+    *,
+    mlx_memory_limit_gib: float | None = None,
+    mlx_cache_limit_gib: float | None = None,
+    mlx_disable_cache: bool = False,
+    mlx_wired_limit_gib: float | None = None,
+    torch_mps_high_watermark_ratio: float | None = None,
+    torch_mps_low_watermark_ratio: float | None = None,
+    mx_module: Any | None = None,
+) -> AppliedMemoryLimits:
+    """Apply optional MLX allocator limits and PyTorch MPS watermarks.
+
+    PyTorch reads MPS watermark variables when the MPS backend initializes, so
+    call this before importing PyTorch. Specifying only a high watermark sets the
+    low watermark to ``0.0``. MLX limit-setting failures are recorded in the
+    result and do not prevent other limits from being applied.
+
+    Parameters:
+        mlx_memory_limit_gib (float | None): Maximum MLX memory in GiB.
+        mlx_cache_limit_gib (float | None): Maximum MLX cache size in GiB.
+        mlx_disable_cache (bool): Whether to disable the MLX cache.
+        mlx_wired_limit_gib (float | None): Maximum MLX wired memory in GiB.
+        torch_mps_high_watermark_ratio (float | None): PyTorch MPS high watermark
+            ratio.
+        torch_mps_low_watermark_ratio (float | None): PyTorch MPS low watermark
+            ratio.
+
+    Returns:
+        AppliedMemoryLimits: Configured values, applied and previous MLX byte
+            limits, MPS watermark values, and per-limit errors.
+    """
+    if torch_mps_high_watermark_ratio is not None and torch_mps_low_watermark_ratio is None:
+        torch_mps_low_watermark_ratio = 0.0
+
+    high = _set_mps_env("PYTORCH_MPS_HIGH_WATERMARK_RATIO", torch_mps_high_watermark_ratio)
+    low = _set_mps_env("PYTORCH_MPS_LOW_WATERMARK_RATIO", torch_mps_low_watermark_ratio)
+
+    memory_bytes = gib_to_bytes(mlx_memory_limit_gib)
+    cache_bytes = 0 if mlx_disable_cache else gib_to_bytes(mlx_cache_limit_gib)
+    wired_bytes = gib_to_bytes(mlx_wired_limit_gib)
+
+    applied: dict[str, int] = {}
+    previous: dict[str, int] = {}
+    errors: dict[str, str] = {}
+    if memory_bytes is not None or cache_bytes is not None or wired_bytes is not None:
+        if mx_module is None:
+            import mlx.core as mx
+
+            mx_module = mx
+
+        # Apply each limit independently; record failures without stopping.
+        limits = [
+            ("mlx_memory_limit", memory_bytes, mx_module.set_memory_limit),
+            ("mlx_cache_limit", cache_bytes, mx_module.set_cache_limit),
+            ("mlx_wired_limit", wired_bytes, mx_module.set_wired_limit),
+        ]
+        for name, value, setter in limits:
+            if value is not None:
+                try:
+                    previous[name] = int(setter(value))
+                    applied[name] = value
+                except Exception as exc:  # noqa: BLE001 - macOS/system-limit dependent.
+                    errors[name] = f"{type(exc).__name__}: {exc}"
+
+    return AppliedMemoryLimits(
+        mlx_memory_limit_gib=mlx_memory_limit_gib,
+        mlx_cache_limit_gib=mlx_cache_limit_gib,
+        mlx_disable_cache=mlx_disable_cache,
+        mlx_wired_limit_gib=mlx_wired_limit_gib,
+        torch_mps_high_watermark_ratio=high,
+        torch_mps_low_watermark_ratio=low,
+        applied_bytes=applied,
+        previous_bytes=previous,
+        errors=errors,
+    )
+
+
+def add_memory_limit_args(
+    parser: argparse.ArgumentParser,
+    *,
+    mlx_memory_limit_gib: float | None = None,
+    mlx_cache_limit_gib: float | None = None,
+    mlx_disable_cache: bool = False,
+    mlx_wired_limit_gib: float | None = None,
+    torch_mps_high_watermark_ratio: float | None = None,
+    torch_mps_low_watermark_ratio: float | None = None,
+) -> None:
+    """
+    Add configurable Apple Silicon memory-limit options to an argument parser.
+
+    Parameters:
+        parser (argparse.ArgumentParser): Parser to which the options are added.
+        mlx_memory_limit_gib (float | None): Default MLX memory limit in GiB.
+        mlx_cache_limit_gib (float | None): Default MLX cache limit in GiB.
+        mlx_disable_cache (bool): Whether the cache limit defaults to zero.
+        mlx_wired_limit_gib (float | None): Default MLX wired-memory limit in GiB.
+        torch_mps_high_watermark_ratio (float | None): Default PyTorch MPS high-watermark ratio.
+        torch_mps_low_watermark_ratio (float | None): Default PyTorch MPS low-watermark ratio.
+    """
+    parser.add_argument("--mlx-memory-limit-gib",
+                        type=float,
+                        default=mlx_memory_limit_gib,
+                        help="Set MLX memory limit in GiB for memory-tier testing (DiT path).")
+    parser.add_argument("--mlx-cache-limit-gib",
+                        type=float,
+                        default=mlx_cache_limit_gib,
+                        help="Set MLX cache limit in GiB. Use --mlx-disable-cache to force 0.")
+    parser.add_argument("--mlx-disable-cache",
+                        action="store_true",
+                        default=mlx_disable_cache,
+                        help="Set MLX cache limit to 0 for stricter memory-tier tests.")
+    parser.add_argument("--mlx-wired-limit-gib",
+                        type=float,
+                        default=mlx_wired_limit_gib,
+                        help="Set MLX wired-memory limit in GiB where supported by macOS/MLX.")
+    parser.add_argument("--torch-mps-high-watermark-ratio",
+                        type=float,
+                        default=torch_mps_high_watermark_ratio,
+                        help="Set PYTORCH_MPS_HIGH_WATERMARK_RATIO before importing torch.")
+    parser.add_argument("--torch-mps-low-watermark-ratio",
+                        type=float,
+                        default=torch_mps_low_watermark_ratio,
+                        help="Set PYTORCH_MPS_LOW_WATERMARK_RATIO before importing torch.")

fastvideo/mlx_runtime/prompt_cache.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Best-effort prompt-embedding cache shared by the MLX entrypoints."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+import logging
+import tempfile
+from pathlib import Path
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def fingerprint_digest(fingerprint: dict[str, object]) -> str:
+    payload = json.dumps(fingerprint, sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def text_encoder_fingerprint(model_root: Path) -> dict[str, object]:
+    """Return a cheap identity for the tokenizer and text-encoder files."""
+    root = model_root.resolve()
+    components = [path for name in ("tokenizer", "text_encoder") if (path := root / name).is_dir()]
+    scan_roots = components or [root]
+    files: list[list[object]] = []
+    complete = True
+    try:
+        for scan_root in scan_roots:
+            for path in sorted(scan_root.rglob("*")):
+                try:
+                    if not path.is_file():
+                        continue
+                    stat = path.stat()
+                    files.append([
+                        path.relative_to(root).as_posix(),
+                        stat.st_size,
+                        stat.st_mtime_ns,
+                        stat.st_ctime_ns,
+                    ])
+                except OSError:
+                    complete = False
+    except OSError:
+        complete = False
+    # ponytail: metadata avoids hashing multi-GB weights; use a model manifest
+    # if supported workflows ever preserve size, mtime, and ctime while mutating.
+    return {"root": str(root), "files": files, "complete": complete}
+
+
+def prompt_cache_meta_path(cache_path: Path) -> Path:
+    return cache_path.with_suffix(cache_path.suffix + ".json")
+
+
+def _fingerprint_is_complete(fingerprint: dict[str, object]) -> bool:
+    text_encoder = fingerprint.get("text_encoder")
+    return not isinstance(text_encoder, dict) or text_encoder.get("complete") is not False
+
+
+def load_prompt_cache(
+    cache_path: Path | None,
+    fingerprint: dict[str, object],
+) -> np.ndarray | None:
+    """Load a matching cache entry, treating every cache failure as a miss."""
+    if cache_path is None or not _fingerprint_is_complete(fingerprint):
+        return None
+    try:
+        metadata = json.loads(prompt_cache_meta_path(cache_path).read_text())
+        if not isinstance(metadata, dict):
+            return None
+        if metadata.get("fingerprint_sha256") != fingerprint_digest(fingerprint):
+            return None
+        payload = cache_path.read_bytes()
+        if metadata.get("data_sha256") != hashlib.sha256(payload).hexdigest():
+            return None
+        array = np.load(io.BytesIO(payload), allow_pickle=False)
+        return array if isinstance(array, np.ndarray) else None
+    except (EOFError, OSError, UnicodeError, ValueError) as exc:
+        logger.info("Prompt cache read skipped for %s: %s", cache_path, exc)
+        return None
+
+
+def _atomic_write(path: Path, payload: bytes) -> None:
+    temp_path: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+                mode="wb",
+                dir=path.parent,
+                prefix=f".{path.name}.",
+                suffix=".tmp",
+                delete=False,
+        ) as handle:
+            temp_path = Path(handle.name)
+            handle.write(payload)
+        temp_path.replace(path)
+    finally:
+        if temp_path is not None:
+            temp_path.unlink(missing_ok=True)
+
+
+def save_prompt_cache(
+    cache_path: Path | None,
+    embeds: np.ndarray,
+    fingerprint: dict[str, object],
+) -> bool:
+    """Atomically publish an integrity-bound cache entry when possible."""
+    if cache_path is None or not _fingerprint_is_complete(fingerprint):
+        return False
+    try:
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        buffer = io.BytesIO()
+        np.save(buffer, np.asarray(embeds), allow_pickle=False)
+        payload = buffer.getvalue()
+        metadata = (json.dumps(
+            {
+                "fingerprint_sha256": fingerprint_digest(fingerprint),
+                "data_sha256": hashlib.sha256(payload).hexdigest(),
+                "fingerprint": fingerprint,
+            },
+            indent=2) + "\n").encode("utf-8")
+        # Publish data first. Until metadata follows, old metadata's data digest
+        # makes the torn pair a harmless miss rather than a stale cache hit.
+        _atomic_write(cache_path, payload)
+        _atomic_write(prompt_cache_meta_path(cache_path), metadata)
+        return True
+    except (OSError, TypeError, ValueError) as exc:
+        logger.info("Prompt cache write skipped for %s: %s", cache_path, exc)
+        return False

fastvideo/mlx_runtime/prompt_enhance.py

@@ -0,0 +1,454 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Local prompt enrichment for the MLX Wan runtime (H3 Context-IR-style).
+
+Wan's training captions are long and cinematic; short user prompts leave
+quality on the table. This module expands a raw prompt into Wan-style
+shot language **on device** — no remote API, no training.
+
+Backends (first match wins):
+
+1. **mlx-lm** — optional local LLM (``--enhance-prompt-model``).
+2. **template** — deterministic cinematic expansion (always available).
+
+System-prompt contract matches the streaming server's enhancer defaults
+in ``fastvideo/entrypoints/streaming/prompt/enhancer.py`` so remote and
+local paths stay interchangeable.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Keep in lockstep with streaming PromptEnhancer defaults (enhance op).
+DEFAULT_ENHANCE_SYSTEM_PROMPT = ("You are a prompt enhancer for cinematic video generation. Given "
+                                 "a user prompt, produce an enhanced prompt that is more vivid, "
+                                 "specific, and concrete. Keep the subject intact; add lighting, "
+                                 "camera, and motion detail. Reply with just the enhanced prompt.")
+
+# Small default that fits 16 GB Macs alongside the 1.3B DiT when the user
+# opts into mlx-lm. Override with --enhance-prompt-model.
+DEFAULT_MLX_LM_MODEL = "mlx-community/Qwen2.5-0.5B-Instruct-4bit"
+
+_CAMERA_CUES = (
+    "cinematic",
+    "camera",
+    "lens",
+    "shot",
+    "bokeh",
+    "dolly",
+    "tracking",
+    "close-up",
+    "wide shot",
+    "handheld",
+    "steadicam",
+)
+_LIGHT_CUES = (
+    "light",
+    "lighting",
+    "sun",
+    "golden hour",
+    "neon",
+    "rim light",
+    "softbox",
+    "overcast",
+    "moonlight",
+    "volumetric",
+)
+_MOTION_CUES = (
+    "moving",
+    "motion",
+    "walk",
+    "run",
+    "flies",
+    "flying",
+    "drifts",
+    "sails",
+    "flows",
+    "pan",
+    "tilt",
+    "zoom",
+)
+
+
+@dataclass(frozen=True)
+class EnhanceResult:
+    """Outcome of a prompt enrichment call."""
+
+    original: str
+    enhanced: str
+    backend: str
+    elapsed_s: float
+    model: str | None = None
+
+    @property
+    def changed(self) -> bool:
+        """Indicates whether the enhanced prompt differs from the original after trimming surrounding whitespace.
+
+        Returns:
+            bool: `True` if the prompts differ, `False` otherwise.
+        """
+        return self.enhanced.strip() != self.original.strip()
+
+
+def _normalize_user_prompt(prompt: str) -> str:
+    """
+    Normalize a user prompt for enhancement.
+
+    Parameters:
+        prompt (str): User-provided prompt text.
+
+    Returns:
+        str: The prompt with leading and trailing whitespace removed and internal whitespace collapsed.
+
+    Raises:
+        ValueError: If the prompt is empty after whitespace normalization.
+    """
+    text = " ".join(prompt.strip().split())
+    if not text:
+        raise ValueError("prompt must be non-empty")
+    return text
+
+
+def _already_rich(prompt: str) -> bool:
+    """
+    Determine whether a prompt already contains substantial camera and lighting detail.
+
+    Returns:
+        bool: `true` if the prompt is at least 160 characters long and includes camera and lighting cues, `false` otherwise.
+    """
+    lower = prompt.lower()
+    has_camera = any(c in lower for c in _CAMERA_CUES)
+    has_light = any(c in lower for c in _LIGHT_CUES)
+    return len(prompt) >= 160 and has_camera and has_light
+
+
+def enhance_prompt_template(prompt: str) -> str:
+    """
+    Expand a prompt with cinematic camera, lighting, motion, and visual-quality details.
+
+    Rich prompts are preserved, while thinner prompts receive deterministic enhancements
+    without changing their subject.
+
+    Returns:
+        str: The original or expanded prompt with normalized whitespace and punctuation.
+    """
+    text = _normalize_user_prompt(prompt)
+    if _already_rich(text):
+        return text
+
+    lower = text.lower()
+    parts = [text.rstrip(".")]
+
+    if not any(c in lower for c in _CAMERA_CUES):
+        parts.append("shot on a 35mm anamorphic lens, gentle handheld micro-movement, "
+                     "shallow depth of field")
+    if not any(c in lower for c in _LIGHT_CUES):
+        parts.append("natural cinematic lighting with soft volumetric haze and subtle "
+                     "rim light separating subject from background")
+    if not any(c in lower for c in _MOTION_CUES):
+        parts.append("smooth continuous motion with grounded physics")
+
+    parts.append("highly detailed, coherent temporal continuity, film grain, "
+                 "color graded like a contemporary drama")
+    enhanced = ", ".join(parts)
+    # Single trailing period; collapse duplicate whitespace.
+    enhanced = re.sub(r"\s+", " ", enhanced).strip()
+    if not enhanced.endswith("."):
+        enhanced += "."
+    return enhanced
+
+
+def enhance_prompt_mlx_lm(
+    prompt: str,
+    *,
+    model: str = DEFAULT_MLX_LM_MODEL,
+    system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT,
+    max_tokens: int = 128,
+    temp: float = 0.6,
+) -> str:
+    """
+    Enhance a user prompt with a locally hosted mlx-lm instruction model.
+
+    Parameters:
+        prompt (str): The prompt to enhance.
+        model (str): The mlx-lm model identifier or path.
+        system_prompt (str): Instructions that guide prompt enhancement.
+        max_tokens (int): Maximum number of tokens to generate.
+        temp (float): Sampling temperature for generation.
+
+    Returns:
+        str: The enhanced prompt.
+
+    Raises:
+        RuntimeError: If mlx-lm is unavailable or produces an empty result.
+    """
+    try:
+        from mlx_lm import generate, load
+    except ImportError as exc:  # pragma: no cover - optional dep
+        raise RuntimeError("mlx-lm is not installed. `uv pip install mlx-lm` or use "
+                           "--enhance-prompt-backend template.") from exc
+
+    text = _normalize_user_prompt(prompt)
+    logger.info("[MLX enhance] loading %s", model)
+    mlx_model, tokenizer = load(model)
+
+    messages = [
+        {
+            "role": "system",
+            "content": system_prompt
+        },
+        {
+            "role": "user",
+            "content": text
+        },
+    ]
+    if hasattr(tokenizer, "apply_chat_template"):
+        chat = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+    else:  # pragma: no cover - ancient tokenizers
+        chat = f"{system_prompt}\n\nUser: {text}\nAssistant:"
+
+    raw = generate(
+        mlx_model,
+        tokenizer,
+        prompt=chat,
+        max_tokens=max_tokens,
+        temp=temp,
+        verbose=False,
+    )
+    enhanced = _clean_llm_output(raw, original=text)
+    if not enhanced:
+        raise RuntimeError("mlx-lm returned an empty enhance result")
+    return enhanced
+
+
+def _clean_llm_output(raw: str, *, original: str) -> str:
+    """
+    Clean generated prompt text and fall back to the original when the result is too short.
+
+    Parameters:
+        raw (str): Raw text produced by the language model.
+        original (str): Original prompt used as the fallback value.
+
+    Returns:
+        str: Cleaned first paragraph of the generated text, or the original prompt when the generated text is too short.
+    """
+    text = raw.strip()
+    # Drop common prefatory phrases.
+    for prefix in (
+            "enhanced prompt:",
+            "here's the enhanced prompt:",
+            "here is the enhanced prompt:",
+            "sure:",
+            "sure,",
+    ):
+        if text.lower().startswith(prefix):
+            text = text[len(prefix):].strip()
+    # Keep first non-empty paragraph only.
+    para = text.split("\n\n")[0].strip()
+    para = " ".join(para.split())
+    if len(para) < max(12, len(original) // 4):
+        return original
+    return para
+
+
+def enhance_prompt(
+    prompt: str,
+    *,
+    backend: str = "auto",
+    model: str | None = None,
+    system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT,
+    max_tokens: int = 128,
+) -> EnhanceResult:
+    """Enhance a prompt using the selected backend, falling back to a deterministic template when configured for automatic selection.
+
+    Parameters:
+        prompt (str): The prompt to enhance.
+        backend (str): The enhancement backend: ``"auto"``, ``"mlx-lm"``, or ``"template"``.
+        model (str | None): The MLX language model to use.
+        system_prompt (str): Instructions provided to the MLX language model.
+        max_tokens (int): Maximum number of tokens generated by the MLX language model.
+
+    Returns:
+        EnhanceResult: The original and enhanced prompts, selected backend, timing information, and model metadata.
+
+    Raises:
+        ValueError: If the prompt is empty or the backend is unsupported.
+        Exception: If the explicitly selected ``"mlx-lm"`` backend fails.
+    """
+    text = _normalize_user_prompt(prompt)
+    backend_norm = (backend or "auto").lower()
+    if backend_norm not in {"auto", "mlx-lm", "template"}:
+        raise ValueError(f"Unknown enhance backend: {backend}")
+
+    start = time.perf_counter()
+    used_model: str | None = None
+
+    if backend_norm in {"auto", "mlx-lm"}:
+        try:
+            used_model = model or DEFAULT_MLX_LM_MODEL
+            enhanced = enhance_prompt_mlx_lm(
+                text,
+                model=used_model,
+                system_prompt=system_prompt,
+                max_tokens=max_tokens,
+            )
+            return EnhanceResult(
+                original=text,
+                enhanced=enhanced,
+                backend="mlx-lm",
+                elapsed_s=time.perf_counter() - start,
+                model=used_model,
+            )
+        except Exception as exc:
+            if backend_norm == "mlx-lm":
+                raise
+            logger.info(
+                "[MLX enhance] mlx-lm unavailable (%s); using template backend",
+                exc,
+            )
+
+    enhanced = enhance_prompt_template(text)
+    return EnhanceResult(
+        original=text,
+        enhanced=enhanced,
+        backend="template",
+        elapsed_s=time.perf_counter() - start,
+        model=None,
+    )
+
+
+def enhance_cache_path(
+    prompt: str,
+    *,
+    backend: str,
+    model: str | None,
+    cache_dir: Path | None = None,
+) -> Path:
+    """Content-addressed cache file for an enhanced prompt string."""
+    root = cache_dir or (Path.home() / ".cache" / "fastvideo" / "enhanced_prompts")
+    key = "\0".join([prompt, backend, model or ""])
+    digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:24]
+    return root / f"{digest}.json"
+
+
+def load_or_enhance_prompt(
+    prompt: str,
+    *,
+    backend: str = "auto",
+    model: str | None = None,
+    system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT,
+    max_tokens: int = 128,
+    cache: bool = True,
+    cache_dir: Path | None = None,
+) -> EnhanceResult:
+    """
+    Enhance a prompt, reusing a cached result when available.
+
+    Parameters:
+        prompt (str): The prompt to enhance.
+        backend (str): Enhancement backend to use.
+        model (str | None): Optional model identifier.
+        system_prompt (str): System prompt for model-based enhancement.
+        max_tokens (int): Maximum number of tokens generated by the model.
+        cache (bool): Whether to read and write the on-disk cache.
+        cache_dir (Path | None): Optional directory for cached results.
+
+    Returns:
+        EnhanceResult: The enhanced prompt and backend metadata. Cached results are marked with the ``"cache"`` backend.
+    """
+    text = _normalize_user_prompt(prompt)
+    path = enhance_cache_path(text, backend=backend, model=model, cache_dir=cache_dir)
+    if cache and path.is_file():
+        try:
+            payload = json.loads(path.read_text())
+            return EnhanceResult(
+                original=str(payload.get("original", text)),
+                enhanced=str(payload["enhanced"]),
+                # Mark cache hits explicitly so metrics/logs can distinguish
+                # a free replay from a fresh template/mlx-lm call.
+                backend="cache",
+                elapsed_s=0.0,
+                model=payload.get("model"),
+            )
+        except (OSError, KeyError, json.JSONDecodeError):
+            pass
+
+    result = enhance_prompt(
+        text,
+        backend=backend,
+        model=model,
+        system_prompt=system_prompt,
+        max_tokens=max_tokens,
+    )
+    if cache:
+        try:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_text(
+                json.dumps(
+                    {
+                        "original": result.original,
+                        "enhanced": result.enhanced,
+                        "backend": result.backend,
+                        "model": result.model,
+                    },
+                    indent=2,
+                ))
+        except OSError as exc:  # pragma: no cover - cache is best-effort
+            logger.info("[MLX enhance] cache write skipped: %s", exc)
+    return result
+
+
+def enhance_result_as_metrics(result: EnhanceResult | None) -> dict[str, Any]:
+    """
+    Convert prompt enhancement results into metrics fields.
+
+    Parameters:
+        result (EnhanceResult | None): The enhancement result, or `None` when no enhancement was performed.
+
+    Returns:
+        dict[str, Any]: A metrics mapping containing enhancement status, backend metadata, timing, and original and enhanced prompts.
+    """
+    if result is None:
+        return {
+            "enhance_prompt": False,
+            "enhance_backend": None,
+            "enhance_model": None,
+            "enhance_elapsed_s": None,
+            "prompt_original": None,
+            "prompt_enhanced": None,
+        }
+    return {
+        "enhance_prompt": True,
+        "enhance_backend": result.backend,
+        "enhance_model": result.model,
+        "enhance_elapsed_s": result.elapsed_s,
+        "prompt_original": result.original,
+        "prompt_enhanced": result.enhanced,
+    }
+
+
+__all__ = [
+    "DEFAULT_ENHANCE_SYSTEM_PROMPT",
+    "DEFAULT_MLX_LM_MODEL",
+    "EnhanceResult",
+    "enhance_cache_path",
+    "enhance_prompt",
+    "enhance_prompt_mlx_lm",
+    "enhance_prompt_template",
+    "enhance_result_as_metrics",
+    "load_or_enhance_prompt",
+]

fastvideo/mlx_runtime/quant_backends.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MLX block-scaled quantization backends (affine INT8, MXFP8/4, NVFP4).
+
+Isolated experiment module: probes which ``mx.quantize`` modes the installed
+MLX build supports and exposes a thin wrapper around native quantized matmul.
+Depends only on ``mlx.core`` and the standard library — do not import the rest
+of FastVideo from here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Final
+from collections.abc import Mapping
+
+import mlx.core as mx
+
+# Probe matrix side length: must be divisible by every mode's group size
+# (affine g64, mxfp* g32, nvfp4 g16).
+_PROBE_DIM: Final[int] = 64
+
+
+class QuantBackend(str, Enum):
+    """Named MLX quantization backends evaluated for M5 Neural Accelerators."""
+
+    AFFINE_INT8_G64 = "affine_int8_g64"
+    MXFP8 = "mxfp8"
+    MXFP4 = "mxfp4"
+    NVFP4 = "nvfp4"
+
+
+BACKENDS: Final[tuple[str, ...]] = tuple(b.value for b in QuantBackend)
+
+# Backend name -> kwargs for mx.quantize / mx.quantized_matmul.
+# Affine baseline matches FastVideo DiT load path (INT8, group size 64).
+# MX/NV block-scaled modes use MLX defaults (see mx.quantize docs).
+_BACKEND_KWARGS: Final[Mapping[str, Mapping[str, object]]] = {
+    QuantBackend.AFFINE_INT8_G64.value: {
+        "mode": "affine",
+        "bits": 8,
+        "group_size": 64,
+    },
+    QuantBackend.MXFP8.value: {
+        "mode": "mxfp8",
+        "bits": None,
+        "group_size": None,
+    },
+    QuantBackend.MXFP4.value: {
+        "mode": "mxfp4",
+        "bits": None,
+        "group_size": None,
+    },
+    QuantBackend.NVFP4.value: {
+        "mode": "nvfp4",
+        "bits": None,
+        "group_size": None,
+    },
+}
+
+_SUPPORT_CACHE: dict[str, bool] = {}
+_SUPPORT_ERROR_CACHE: dict[str, str | None] = {}
+_BYTES_CACHE: dict[str, float] = {}
+
+
+@dataclass(frozen=True)
+class QuantizedWeight:
+    """Packed quantized weight plus scales/biases for one backend."""
+
+    weight: mx.array
+    scales: mx.array
+    biases: mx.array | None
+    backend: str
+    mode: str
+    bits: int | None
+    group_size: int | None
+    # Original (rows, cols) of the fp weight, used for bytes-per-element.
+    orig_shape: tuple[int, int]
+
+
+def _normalize_backend(backend: str) -> str:
+    """Normalize a quantization backend name and validate that it is supported.
+
+    Parameters:
+        backend (str): Backend name to normalize.
+
+    Returns:
+        str: The lowercase backend name without surrounding whitespace.
+
+    Raises:
+        ValueError: If the backend name is unknown.
+    """
+    name = backend.strip().lower()
+    if name not in _BACKEND_KWARGS:
+        known = ", ".join(BACKENDS)
+        raise ValueError(f"Unknown quant backend {backend!r}. Expected one of: {known}")
+    return name
+
+
+def _kwargs_for(backend: str) -> dict[str, object]:
+    """Return the MLX quantization arguments configured for a backend.
+
+    Parameters:
+        backend (str): Backend name to resolve.
+
+    Returns:
+        dict[str, object]: Quantization arguments for the normalized backend.
+    """
+    return dict(_BACKEND_KWARGS[_normalize_backend(backend)])
+
+
+def support_error(backend: str) -> str | None:
+    """
+    Check whether a quantization backend is supported by the current MLX runtime.
+
+    Parameters:
+        backend (str): Quantization backend name.
+
+    Returns:
+        str | None: An error description when the backend is unsupported, or `None` when supported.
+    """
+    name = _normalize_backend(backend)
+    if name in _SUPPORT_ERROR_CACHE:
+        return _SUPPORT_ERROR_CACHE[name]
+
+    kwargs = _kwargs_for(name)
+    try:
+        w = mx.zeros((_PROBE_DIM, _PROBE_DIM), dtype=mx.float16)
+        quantized = mx.quantize(
+            w,
+            group_size=kwargs["group_size"],  # type: ignore[arg-type]
+            bits=kwargs["bits"],  # type: ignore[arg-type]
+            mode=str(kwargs["mode"]),
+        )
+        w_q = quantized[0]
+        scales = quantized[1]
+        biases = quantized[2] if len(quantized) == 3 else None
+        x = mx.zeros((1, _PROBE_DIM), dtype=mx.float16)
+        y = mx.quantized_matmul(
+            x,
+            w_q,
+            scales,
+            biases,
+            transpose=True,
+            group_size=kwargs["group_size"],  # type: ignore[arg-type]
+            bits=kwargs["bits"],  # type: ignore[arg-type]
+            mode=str(kwargs["mode"]),
+        )
+        mx.eval(y)
+        _SUPPORT_ERROR_CACHE[name] = None
+        _SUPPORT_CACHE[name] = True
+    except Exception as exc:  # noqa: BLE001 - MLX raises varied types per mode/version.
+        msg = f"{type(exc).__name__}: {exc}"
+        _SUPPORT_ERROR_CACHE[name] = msg
+        _SUPPORT_CACHE[name] = False
+    return _SUPPORT_ERROR_CACHE[name]
+
+
+def is_supported(backend: str) -> bool:
+    """Return True if the installed MLX build can quantize/matmul with ``backend``."""
+    name = _normalize_backend(backend)
+    if name not in _SUPPORT_CACHE:
+        support_error(name)
+    return _SUPPORT_CACHE[name]
+
+
+def quantize_weight(w: mx.array, backend: str) -> QuantizedWeight:
+    """
+    Quantize a two-dimensional weight matrix using the specified native MLX backend.
+
+    Parameters:
+        w (mx.array): The two-dimensional weight matrix to quantize.
+        backend (str): The quantization backend to use.
+
+    Returns:
+        QuantizedWeight: The quantized weights and associated quantization metadata.
+
+    Raises:
+        ValueError: If the backend is unknown, the weight is not two-dimensional,
+            or its last dimension is not divisible by the backend's group size.
+        RuntimeError: If the backend is unsupported by the installed MLX build.
+    """
+    name = _normalize_backend(backend)
+    err = support_error(name)
+    if err is not None:
+        mlx_version = getattr(mx, "__version__", "unknown")
+        raise RuntimeError(f"Quant backend {name!r} is not supported by installed mlx "
+                           f"({mlx_version}): {err}")
+
+    if w.ndim != 2:
+        raise ValueError(f"quantize_weight expects a 2D weight, got shape {tuple(w.shape)}")
+
+    rows, cols = int(w.shape[0]), int(w.shape[1])
+    kwargs = _kwargs_for(name)
+    group_size = kwargs["group_size"]
+    # When group_size is None, MLX applies the mode default; only check when set.
+    if isinstance(group_size, int) and cols % group_size != 0:
+        raise ValueError(f"Weight last dim {cols} must be divisible by group_size={group_size} "
+                         f"for backend {name!r}")
+
+    quantized = mx.quantize(
+        w,
+        group_size=kwargs["group_size"],  # type: ignore[arg-type]
+        bits=kwargs["bits"],  # type: ignore[arg-type]
+        mode=str(kwargs["mode"]),
+    )
+    w_q = quantized[0]
+    scales = quantized[1]
+    biases = quantized[2] if len(quantized) == 3 else None
+    eval_args = [w_q, scales] if biases is None else [w_q, scales, biases]
+    mx.eval(*eval_args)
+
+    return QuantizedWeight(
+        weight=w_q,
+        scales=scales,
+        biases=biases,
+        backend=name,
+        mode=str(kwargs["mode"]),
+        bits=kwargs["bits"] if isinstance(kwargs["bits"], int) else None,
+        group_size=group_size if isinstance(group_size, int) else None,
+        orig_shape=(rows, cols),
+    )
+
+
+def quantized_matmul(x: mx.array, qw: QuantizedWeight) -> mx.array:
+    """Compute ``x @ w.T`` in the quantized domain via ``mx.quantized_matmul``."""
+    return mx.quantized_matmul(
+        x,
+        qw.weight,
+        qw.scales,
+        qw.biases,
+        transpose=True,
+        group_size=qw.group_size,
+        bits=qw.bits,
+        mode=qw.mode,
+    )
+
+
+def _artifact_nbytes(qw: QuantizedWeight) -> int:
+    """
+    Calculate the total storage size of a quantized weight artifact in bytes.
+
+    Parameters:
+        qw (QuantizedWeight): Quantized weight artifact whose packed weights, scales, and optional biases are measured.
+
+    Returns:
+        int: Total number of bytes used by the artifact's stored arrays.
+    """
+    total = int(qw.weight.nbytes) + int(qw.scales.nbytes)
+    if qw.biases is not None:
+        total += int(qw.biases.nbytes)
+    return total
+
+
+def bytes_per_weight(backend: str) -> float:
+    """
+    Measure the effective storage cost of a quantized weight.
+
+    Parameters:
+        backend (str): Quantization backend to measure.
+
+    Returns:
+        float: Stored bytes per original weight element, including packed weights,
+            scales, and optional biases.
+
+    Raises:
+        RuntimeError: If the backend is unsupported.
+    """
+    name = _normalize_backend(backend)
+    if name in _BYTES_CACHE:
+        return _BYTES_CACHE[name]
+
+    err = support_error(name)
+    if err is not None:
+        mlx_version = getattr(mx, "__version__", "unknown")
+        raise RuntimeError(f"Cannot measure bytes_per_weight for unsupported backend {name!r} "
+                           f"(mlx {mlx_version}): {err}")
+
+    probe = mx.zeros((_PROBE_DIM, _PROBE_DIM), dtype=mx.float16)
+    qw = quantize_weight(probe, name)
+    n_elem = qw.orig_shape[0] * qw.orig_shape[1]
+    value = _artifact_nbytes(qw) / float(n_elem)
+    _BYTES_CACHE[name] = value
+    return value

fastvideo/mlx_runtime/refine.py

@@ -0,0 +1,690 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Two-pass spatial refine for the MLX Wan runtime (H3 / LTX-2 pattern).
+
+Biggest quality lever on Apple Silicon without a new model or training:
+generate at base resolution, then run a second denoising pass with the
+*same* DiT at a higher resolution.
+
+This is the MLX-side port of the CUDA refine template in
+``fastvideo/pipelines/basic/ltx2/stages/ltx2_refine.py`` and the H3
+"base + regenerate" pattern documented in
+``docs/design/mac_qad_two_product_strategy.md``:
+
+1. :func:`plan_refine_resolutions` — split the request into stage-1
+   (base) and stage-2 (target) pixel sizes, validating VAE / patch
+   alignment the way :class:`LTX2RefineInitStage` does.
+2. :func:`upsample_latents_spatial` — 2× (or N×) spatial upsample of
+   clean latents. Wan has no learned latent upsampler on Mac, so this
+   is bilinear over the H×W plane (temporal axis untouched) — same
+   role as LTX-2's ``upsample_video`` hand-off, without the learned
+   residual.
+3. :func:`prepare_refine_latents` — upsample + re-noise the clean
+   stage-1 latents to the stage-2 sigma so the second denoise has
+   something to refine (mirrors :class:`LTX2UpsampleStage` +
+   ``apply_ltx2_gaussian_noiser``).
+4. :func:`run_two_pass_dmd` — orchestrate stage-1 denoise → refine
+   hand-off → stage-2 denoise with the same model / prompt embeds.
+
+No LoRA swap, no dedicated SR weights, no new training — pure pipeline
+work reusable by Wan2.1-14B and Wan2.2-5B on Apple Silicon.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+from collections.abc import Callable, Sequence
+
+import numpy as np
+
+from fastvideo.logger import init_logger
+from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, add_noise, dmd_step
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    import mlx.core as mx
+
+logger = init_logger(__name__)
+
+# Default stage-2 noise level when the caller does not supply a schedule.
+# Matches the first entry of LTX-2's STAGE_2_DISTILLED_SIGMA_VALUES in spirit
+# (start the refine denoise from a high-noise level) without hard-wiring the
+# LTX-2 distilled grid onto Wan's flow-match schedule.
+DEFAULT_REFINE_SIGMA = 0.909375
+
+
+@dataclass(frozen=True)
+class RefinePlan:
+    """Resolved stage-1 / stage-2 geometry for a two-pass refine run."""
+
+    target_height: int
+    target_width: int
+    stage1_height: int
+    stage1_width: int
+    spatial_scale: int
+    vae_spatial_compression: int
+    vae_temporal_compression: int
+    num_frames: int
+
+    @property
+    def stage1_latent_height(self) -> int:
+        """Return the stage-1 latent height after VAE spatial compression."""
+        return self.stage1_height // self.vae_spatial_compression
+
+    @property
+    def stage1_latent_width(self) -> int:
+        """Return the stage-one latent width after VAE spatial compression."""
+        return self.stage1_width // self.vae_spatial_compression
+
+    @property
+    def stage2_latent_height(self) -> int:
+        """Calculate the target-resolution latent height.
+
+        Returns:
+            int: The target height divided by the VAE spatial compression factor.
+        """
+        return self.target_height // self.vae_spatial_compression
+
+    @property
+    def stage2_latent_width(self) -> int:
+        """Return the target image width in latent-space units."""
+        return self.target_width // self.vae_spatial_compression
+
+    @property
+    def latent_frames(self) -> int:
+        """Calculate the number of latent frames after VAE temporal compression.
+
+        Returns:
+            int: The compressed latent frame count.
+        """
+        return (self.num_frames - 1) // self.vae_temporal_compression + 1
+
+
+def plan_refine_resolutions(
+        *,
+        height: int,
+        width: int,
+        num_frames: int,
+        spatial_scale: int = 2,
+        vae_spatial_compression: int = 8,
+        vae_temporal_compression: int = 4,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        enabled: bool = True,
+        mode_label: str = "Refine",
+) -> RefinePlan:
+    """
+    Validate the requested dimensions and create the stage-1 and target-resolution refinement plan.
+
+    Parameters:
+        height (int): Target image height in pixels.
+        width (int): Target image width in pixels.
+        num_frames (int): Number of frames in the input sequence.
+        spatial_scale (int): Factor used to reduce spatial dimensions for stage 1.
+        vae_spatial_compression (int): Spatial compression factor of the VAE.
+        vae_temporal_compression (int): Temporal compression factor of the VAE.
+        patch_size (tuple[int, int, int]): Temporal and spatial patch dimensions used to validate latent-grid alignment.
+        enabled (bool): Whether to use two-pass refinement.
+        mode_label (str): Name of the calling mode, used to prefix validation
+            errors so ``--fast-spatial`` failures do not read as refine failures.
+
+    Returns:
+        RefinePlan: The validated stage-1 and target-resolution plan.
+    """
+    if height <= 0 or width <= 0:
+        raise ValueError(f"height/width must be positive, got {height}x{width}")
+    if spatial_scale < 1:
+        raise ValueError(f"spatial_scale must be >= 1, got {spatial_scale}")
+    if num_frames <= 0:
+        raise ValueError(f"num_frames must be positive, got {num_frames}")
+    if vae_spatial_compression < 1 or vae_temporal_compression < 1:
+        raise ValueError("VAE compression factors must be positive")
+    if height % vae_spatial_compression != 0 or width % vae_spatial_compression != 0:
+        raise ValueError(f"height/width must be divisible by vae_spatial_compression={vae_spatial_compression} "
+                         f"(got {height}x{width}).")
+    if (num_frames - 1) % vae_temporal_compression != 0:
+        raise ValueError(f"num_frames must be 1 modulo vae_temporal_compression={vae_temporal_compression} "
+                         f"(got {num_frames}).")
+
+    if not enabled or spatial_scale == 1:
+        plan = RefinePlan(
+            target_height=height,
+            target_width=width,
+            stage1_height=height,
+            stage1_width=width,
+            spatial_scale=1,
+            vae_spatial_compression=vae_spatial_compression,
+            vae_temporal_compression=vae_temporal_compression,
+            num_frames=num_frames,
+        )
+        _validate_plan(plan, patch_size=patch_size, mode_label=mode_label)
+        return plan
+
+    if height % spatial_scale != 0 or width % spatial_scale != 0:
+        raise ValueError(f"{mode_label} requires height/width divisible by spatial_scale={spatial_scale} "
+                         f"(got {height}x{width}).")
+
+    stage1_height = height // spatial_scale
+    stage1_width = width // spatial_scale
+    # Stage-1 must land on a VAE-aligned grid so the first denoise produces
+    # valid latents; the LTX-2 init stage enforces the same constraint.
+    if (stage1_height % vae_spatial_compression != 0 or stage1_width % vae_spatial_compression != 0):
+        raise ValueError(f"{mode_label} requires height/width divisible by "
+                         f"{spatial_scale * vae_spatial_compression} "
+                         f"(got {height}x{width}, vae_spatial={vae_spatial_compression}).")
+
+    plan = RefinePlan(
+        target_height=height,
+        target_width=width,
+        stage1_height=stage1_height,
+        stage1_width=stage1_width,
+        spatial_scale=spatial_scale,
+        vae_spatial_compression=vae_spatial_compression,
+        vae_temporal_compression=vae_temporal_compression,
+        num_frames=num_frames,
+    )
+    _validate_plan(plan, patch_size=patch_size, mode_label=mode_label)
+    logger.info(
+        "[MLX refine] enabled: stage1=%dx%d stage2=%dx%d scale=%dx",
+        stage1_width,
+        stage1_height,
+        width,
+        height,
+        spatial_scale,
+    )
+    return plan
+
+
+def _validate_plan(plan: RefinePlan, *, patch_size: tuple[int, int, int], mode_label: str = "Refine") -> None:
+    """
+    Validate that both refinement stages have latent dimensions aligned to the patch grid.
+
+    Parameters:
+        patch_size (tuple[int, int, int]): Temporal, height, and width patch dimensions.
+
+    Raises:
+        ValueError: If a stage's spatial latent dimensions or the temporal latent
+            dimension is not divisible by the corresponding patch dimension.
+    """
+    pt, ph, pw = patch_size
+    for label, lh, lw in (
+        ("stage1", plan.stage1_latent_height, plan.stage1_latent_width),
+        ("stage2", plan.stage2_latent_height, plan.stage2_latent_width),
+    ):
+        if lh % ph != 0 or lw % pw != 0:
+            raise ValueError(f"{mode_label} {label} latent grid {lh}x{lw} is not divisible by "
+                             f"patch spatial size {ph}x{pw}.")
+        if plan.latent_frames % pt != 0:
+            raise ValueError(f"{mode_label} latent_frames={plan.latent_frames} is not divisible by "
+                             f"patch temporal size {pt}.")
+
+
+def upsample_latents_spatial(
+    latents: Any,
+    *,
+    scale: int = 2,
+    mode: str = "bilinear",
+) -> Any:
+    """
+    Upsample the spatial dimensions of 5-D latent arrays while preserving the batch, channel, and temporal dimensions.
+
+    Parameters:
+        latents (Any): Latents with shape ``(B, C, T, H, W)``.
+        scale (int): Integer factor for enlarging the spatial dimensions.
+        mode (str): Interpolation mode, either ``"nearest"`` or ``"bilinear"``.
+
+    Returns:
+        Any: Latents with shape ``(B, C, T, H * scale, W * scale)``.
+    """
+    if scale < 1:
+        raise ValueError(f"scale must be >= 1, got {scale}")
+    if scale == 1:
+        return latents
+
+    # Accept both mx.array and np.ndarray so unit tests can run without MLX.
+    is_mlx = hasattr(latents, "dtype") and type(latents).__module__.startswith("mlx")
+    if is_mlx:
+        return _upsample_latents_mlx(latents, scale=scale, mode=mode)
+    return _upsample_latents_numpy(np.asarray(latents), scale=scale, mode=mode)
+
+
+def _upsample_latents_numpy(
+    latents: np.ndarray,
+    *,
+    scale: int,
+    mode: str,
+) -> np.ndarray:
+    """Upsample 5-D latent arrays spatially using nearest-neighbor or bilinear interpolation.
+
+    Parameters:
+        latents (np.ndarray): Latents with shape ``(B, C, T, H, W)``.
+        scale (int): Spatial upsampling factor.
+        mode (str): Interpolation mode, either ``"nearest"`` or ``"bilinear"``.
+
+    Returns:
+        np.ndarray: Spatially upsampled latents with preserved batch, channel, and temporal dimensions.
+
+    Raises:
+        ValueError: If the latents are not five-dimensional or the interpolation mode is unsupported.
+    """
+    if latents.ndim != 5:
+        raise ValueError(f"Expected 5-D latents (B,C,T,H,W), got shape {latents.shape}")
+    b, c, t, h, w = latents.shape
+    if mode == "nearest":
+        # (B,C,T,H,1,W,1) -> broadcast to (B,C,T,H,scale,W,scale) -> merge.
+        out = np.repeat(np.repeat(latents, scale, axis=3), scale, axis=4)
+        return out
+
+    if mode != "bilinear":
+        raise ValueError(f"Unsupported upsample mode: {mode}")
+
+    # Bilinear over the spatial plane. Flatten (B,C,T) into a batch of 2-D
+    # maps so a single vectorized gather covers every frame/channel.
+    src = latents.reshape(b * c * t, h, w).astype(np.float32, copy=False)
+    out_h, out_w = h * scale, w * scale
+    # Map output pixel centers onto the input grid (align_corners=False).
+    ys = (np.arange(out_h, dtype=np.float32) + 0.5) * (h / out_h) - 0.5
+    xs = (np.arange(out_w, dtype=np.float32) + 0.5) * (w / out_w) - 0.5
+    ys = np.clip(ys, 0.0, h - 1.0)
+    xs = np.clip(xs, 0.0, w - 1.0)
+    y0 = np.floor(ys).astype(np.int64)
+    x0 = np.floor(xs).astype(np.int64)
+    y1 = np.minimum(y0 + 1, h - 1)
+    x1 = np.minimum(x0 + 1, w - 1)
+    wy = (ys - y0.astype(np.float32))[:, None]
+    wx = (xs - x0.astype(np.float32))[None, :]
+    # Gather the four corners: shape (N, out_h, out_w).
+    Ia = src[:, y0[:, None], x0[None, :]]
+    Ib = src[:, y0[:, None], x1[None, :]]
+    Ic = src[:, y1[:, None], x0[None, :]]
+    Id = src[:, y1[:, None], x1[None, :]]
+    wa = (1.0 - wy) * (1.0 - wx)
+    wb = (1.0 - wy) * wx
+    wc = wy * (1.0 - wx)
+    wd = wy * wx
+    out = wa * Ia + wb * Ib + wc * Ic + wd * Id
+    return out.reshape(b, c, t, out_h, out_w).astype(latents.dtype, copy=False)
+
+
+def _upsample_latents_mlx(latents: mx.array, *, scale: int, mode: str) -> mx.array:
+    """
+    Upsample MLX latent tensors along their spatial dimensions.
+
+    Parameters:
+        latents (mx.array): A latent tensor with shape `(B, C, T, H, W)`.
+        scale (int): The integer spatial upsampling factor.
+        mode (str): The interpolation mode, such as `"nearest"` or `"bilinear"`.
+
+    Returns:
+        mx.array: The spatially upsampled latent tensor with its original data type.
+    """
+    import mlx.core as mx
+
+    # Route through NumPy for the interpolation math. Latent tensors at Mac
+    # resolutions are small (e.g. 1×16×21×30×52 ≈ 1 MB) so the host hop is
+    # cheaper than carrying a bespoke Metal bilinear kernel, and it keeps
+    # the CPU-only unit tests and the MLX path on one implementation.
+    np_latents = np.array(latents.astype(mx.float32))
+    up = _upsample_latents_numpy(np_latents, scale=scale, mode=mode)
+    return mx.array(up).astype(latents.dtype)
+
+
+def prepare_refine_latents(
+    clean_latents: Any,
+    *,
+    scale: int = 2,
+    sigma: float = DEFAULT_REFINE_SIGMA,
+    noise: Any | None = None,
+    add_noise_flag: bool = True,
+    upsample_mode: str = "bilinear",
+    seed: int | None = None,
+) -> Any:
+    """
+    Upsample clean latents spatially and optionally mix them with Gaussian noise.
+
+    Parameters:
+        clean_latents: The stage-1 latent tensor.
+        sigma: Noise mixing factor between 0 and 1.
+        noise: Optional noise tensor to mix with the upsampled latents.
+        add_noise_flag: Whether to apply noise mixing.
+        upsample_mode: Spatial interpolation mode.
+        seed: Optional seed for generated noise.
+
+    Returns:
+        The upsampled latents, optionally mixed with noise.
+
+    Raises:
+        ValueError: If sigma is outside the range from 0 to 1.
+    """
+    if sigma < 0.0 or sigma > 1.0:
+        raise ValueError(f"sigma must be in [0, 1], got {sigma}")
+
+    upsampled = upsample_latents_spatial(clean_latents, scale=scale, mode=upsample_mode)
+    if not add_noise_flag or sigma == 0.0:
+        return upsampled
+
+    is_mlx = hasattr(upsampled, "dtype") and type(upsampled).__module__.startswith("mlx")
+    if noise is None:
+        noise = _draw_noise_like(upsampled, seed=seed, is_mlx=is_mlx)
+    return add_noise(upsampled, noise, float(sigma))
+
+
+def refine_sigma_from_schedule(
+    schedule: MLXDMDSchedule,
+    timesteps: Sequence[float | int],
+) -> float:
+    """Derive the refinement noise level from the first refinement timestep.
+
+    Parameters:
+        schedule (MLXDMDSchedule): Schedule used to map timesteps to noise levels.
+        timesteps (Sequence[float | int]): Refinement timesteps, whose first value determines the sigma.
+
+    Returns:
+        float: Sigma corresponding to the first refinement timestep.
+
+    Raises:
+        ValueError: If `timesteps` is empty.
+    """
+    if not timesteps:
+        raise ValueError("timesteps must be non-empty to derive a refine sigma")
+    return float(schedule.sigma_for(float(timesteps[0])))
+
+
+def default_refine_timesteps(
+    schedule: MLXDMDSchedule,
+    timesteps: Sequence[float | int],
+) -> list[float]:
+    """Derive stage-2 timesteps from the stage-1 DMD grid.
+
+    The stage-2 pass must start *below* full noise, otherwise the hand-off
+    ``(1 - sigma) * upsampled + sigma * noise`` weights stage 1 at zero and
+    the refine pass silently becomes a plain full-resolution generation at
+    twice the cost. FastWan's stage-1 grid opens at ``t=1000`` (``sigma``
+    exactly 1.0), so reusing it verbatim — which is what happens when
+    ``--refine-dmd-denoising-steps`` is left unset — discards stage 1.
+
+    Dropping the leading full-noise entries keeps the pass on timesteps the
+    distilled student was actually trained on (no off-grid ``t`` the DiT has
+    never seen) while letting the stage-1 structure through.
+
+    Parameters:
+        schedule (MLXDMDSchedule): Schedule used to map timesteps to noise levels.
+        timesteps (Sequence[float | int]): The stage-1 DMD timestep grid.
+
+    Returns:
+        list[float]: The stage-1 grid with leading full-noise timesteps removed.
+
+    Raises:
+        ValueError: If every timestep in the grid is at full noise, leaving no
+            usable refine step.
+    """
+    steps = [float(step) for step in timesteps]
+    first = 0
+    while first < len(steps) and schedule.sigma_for(steps[first]) >= 1.0:
+        first += 1
+    if first == len(steps):
+        raise ValueError(f"No usable refine timesteps in {steps}: every entry is at sigma >= 1 "
+                         "(full noise), which would discard the stage-1 result. Pass "
+                         "explicit stage-2 timesteps below the full-noise step.")
+    return steps[first:]
+
+
+def run_dmd_loop(
+    *,
+    dit: Any,
+    latents: Any,
+    encoder_hidden_states: Any,
+    freqs_cis: tuple[Any, Any],
+    timesteps: Sequence[float | int],
+    schedule: MLXDMDSchedule,
+    mx_dtype: Any,
+    seed: int | None = None,
+    step_callback: Callable[[int, int], None] | None = None,
+    label: str = "denoise",
+) -> Any:
+    """
+    Denoise latents over the supplied timesteps using the DMD schedule.
+
+    Parameters:
+        timesteps (Sequence[float | int]): Denoising timesteps in execution order.
+        seed (int | None): Seed for reproducible intermediate noise generation.
+        step_callback (Callable[[int, int], None] | None): Callback receiving the
+            completed step number and total step count.
+        label (str): Label used for progress output when no callback is provided.
+
+    Returns:
+        Any: The denoised latents.
+    """
+    import mlx.core as mx
+
+    renoise_rng = np.random.default_rng(seed) if seed is not None else None
+    latents_out = latents
+    n_steps = len(timesteps)
+    for step_index, timestep in enumerate(timesteps):
+        noise_input = latents_out
+        ts_val = float(timestep)
+        timestep_mx = mx.array([ts_val]).astype(mx.float32)
+        noise_pred = dit(
+            latents_out.astype(mx_dtype),
+            encoder_hidden_states,
+            timestep_mx,
+            freqs_cis,
+        )
+        noise_input_f32 = noise_input.astype(mx.float32)
+        pred_noise_f32 = noise_pred.astype(mx.float32)
+        if step_index < n_steps - 1:
+            next_ts: float | None = float(timesteps[step_index + 1])
+            if renoise_rng is not None:
+                renoise = mx.array(renoise_rng.standard_normal(tuple(noise_input_f32.shape)).astype(np.float32))
+            else:
+                renoise = mx.random.normal(noise_input_f32.shape).astype(mx.float32)
+        else:
+            next_ts, renoise = None, None
+        latents_out = dmd_step(
+            latents=noise_input_f32,
+            noise_input_latent=noise_input_f32,
+            pred_noise=pred_noise_f32,
+            schedule=schedule,
+            timestep=ts_val,
+            next_timestep=next_ts,
+            noise=renoise,
+        ).astype(mx_dtype)
+        mx.eval(latents_out)
+        if step_callback is not None:
+            step_callback(step_index + 1, n_steps)
+        else:
+            print(f"{label} step {step_index + 1}/{n_steps} complete")
+    return latents_out
+
+
+@dataclass(frozen=True)
+class TwoPassResult:
+    """Outputs of :func:`run_two_pass_dmd`."""
+
+    latents: Any
+    stage1_latents: Any
+    plan: RefinePlan
+    refine_sigma: float
+
+
+def run_two_pass_dmd(
+    *,
+    dit: Any,
+    encoder_hidden_states: Any,
+    noise_latents_stage1: Any,
+    freqs_cis_stage1: tuple[Any, Any],
+    freqs_cis_stage2: tuple[Any, Any] | None,
+    plan: RefinePlan,
+    schedule: MLXDMDSchedule,
+    timesteps: Sequence[float | int],
+    refine_timesteps: Sequence[float | int] | None = None,
+    mx_dtype: Any,
+    seed: int = 0,
+    add_noise_flag: bool = True,
+    upsample_mode: str = "bilinear",
+    refine_sigma: float | None = None,
+    step_callback: Callable[[str, int, int], None] | None = None,
+) -> TwoPassResult:
+    """
+    Run base denoising and, when enabled, spatial refinement denoising.
+
+    Parameters:
+        dit: DiT callable used for both denoising passes.
+        encoder_hidden_states: Prompt embeddings shared across both passes.
+        noise_latents_stage1: Initial stage-1 noise latents.
+        freqs_cis_stage1: RoPE tables for the stage-1 resolution.
+        freqs_cis_stage2: RoPE tables for the stage-2 resolution, required when refinement is enabled.
+        plan: Refinement geometry and configuration.
+        schedule: Flow-matching schedule used by both passes.
+        timesteps: Stage-1 denoising timesteps.
+        refine_timesteps: Stage-2 denoising timesteps. Uses `timesteps` when omitted.
+        mx_dtype: MLX dtype used for DiT inputs and outputs.
+        seed: Base seed for reproducible noise generation.
+        add_noise_flag: Whether to add noise to the upsampled stage-1 latents.
+        upsample_mode: Spatial upsampling mode, either `"bilinear"` or `"nearest"`.
+        refine_sigma: Stage-2 starting noise level. Derived from the first refinement timestep when omitted.
+        step_callback: Optional callback receiving the phase name, step index, and total step count.
+
+    Returns:
+        TwoPassResult containing the final latents, stage-1 latents, refinement plan, and applied refinement sigma.
+
+    Raises:
+        ValueError: If refinement is enabled without stage-2 RoPE tables, without refinement timesteps, or if upsampled latents do not match the planned stage-2 dimensions.
+    """
+    stage1_cb = None
+    stage2_cb = None
+    if step_callback is not None:
+        stage1_cb = lambda i, n: step_callback("stage1", i, n)  # noqa: E731
+        stage2_cb = lambda i, n: step_callback("stage2", i, n)  # noqa: E731
+
+    stage1_latents = run_dmd_loop(
+        dit=dit,
+        latents=noise_latents_stage1,
+        encoder_hidden_states=encoder_hidden_states,
+        freqs_cis=freqs_cis_stage1,
+        timesteps=timesteps,
+        schedule=schedule,
+        mx_dtype=mx_dtype,
+        seed=seed,
+        step_callback=stage1_cb,
+        label="stage1 denoise",
+    )
+
+    if plan.spatial_scale == 1:
+        return TwoPassResult(
+            latents=stage1_latents,
+            stage1_latents=stage1_latents,
+            plan=plan,
+            refine_sigma=0.0,
+        )
+
+    if freqs_cis_stage2 is None:
+        raise ValueError("freqs_cis_stage2 is required when refine spatial_scale > 1")
+
+    if refine_timesteps is not None:
+        stage2_timesteps = [float(step) for step in refine_timesteps]
+        if not stage2_timesteps:
+            raise ValueError("refine_timesteps must be non-empty when refine is enabled")
+    else:
+        # Not `list(timesteps)`: the stage-1 grid opens at full noise, which
+        # would weight the stage-1 result at zero. See default_refine_timesteps.
+        stage2_timesteps = default_refine_timesteps(schedule, timesteps)
+    grid_sigma = refine_sigma_from_schedule(schedule, stage2_timesteps)
+    sigma = float(refine_sigma) if refine_sigma is not None else grid_sigma
+    if refine_sigma is not None and abs(sigma - grid_sigma) > 1e-6:
+        # The loop tells the DiT `stage2_timesteps[0]`, which implies grid_sigma.
+        # Overriding the hand-off noise level breaks that correspondence, so the
+        # model is denoising from a level it was not told about. Useful for
+        # exploring schedules that bottom out too high, but say so out loud.
+        logger.warning(
+            "[MLX refine] refine_sigma=%.4f overrides the schedule's %.4f for timestep %g; "
+            "the DiT is told a timestep that no longer matches the noise it receives.",
+            sigma,
+            grid_sigma,
+            stage2_timesteps[0],
+        )
+
+    # A hand-off at sigma >= 1 is `0 * upsampled + 1 * noise`: stage 1 is
+    # thrown away and refine degrades to a plain full-res run at 2x the cost.
+    # Fail loudly rather than silently burning the first pass.
+    if add_noise_flag and sigma >= 1.0:
+        raise ValueError(f"Refine hand-off sigma={sigma:.4f} (from stage-2 timestep "
+                         f"{stage2_timesteps[0]:g}) discards the stage-1 result entirely: "
+                         "the upsampled latents are weighted (1 - sigma) = 0. Start the "
+                         "stage-2 grid below the full-noise timestep, or pass "
+                         "add_noise_flag=False to hand off the clean upsample.")
+
+    stage2_input = prepare_refine_latents(
+        stage1_latents,
+        scale=plan.spatial_scale,
+        sigma=sigma,
+        add_noise_flag=add_noise_flag,
+        upsample_mode=upsample_mode,
+        seed=seed + 1,
+    )
+
+    # Shape guard: upsampled latents must match the stage-2 RoPE grid.
+    expected_h = plan.stage2_latent_height
+    expected_w = plan.stage2_latent_width
+    got_h, got_w = int(stage2_input.shape[-2]), int(stage2_input.shape[-1])
+    if got_h != expected_h or got_w != expected_w:
+        raise ValueError(f"Refine upsample produced {got_h}x{got_w} latents, expected "
+                         f"{expected_h}x{expected_w} for target "
+                         f"{plan.target_height}x{plan.target_width}.")
+
+    logger.info(
+        "[MLX refine] stage2 start: latent=%dx%d sigma=%.4f steps=%d",
+        expected_w,
+        expected_h,
+        sigma,
+        len(stage2_timesteps),
+    )
+
+    stage2_latents = run_dmd_loop(
+        dit=dit,
+        latents=stage2_input,
+        encoder_hidden_states=encoder_hidden_states,
+        freqs_cis=freqs_cis_stage2,
+        timesteps=stage2_timesteps,
+        schedule=schedule,
+        mx_dtype=mx_dtype,
+        seed=seed + 2,
+        step_callback=stage2_cb,
+        label="stage2 refine",
+    )
+    return TwoPassResult(
+        latents=stage2_latents,
+        stage1_latents=stage1_latents,
+        plan=plan,
+        refine_sigma=sigma,
+    )
+
+
+__all__ = [
+    "DEFAULT_REFINE_SIGMA",
+    "RefinePlan",
+    "TwoPassResult",
+    "default_refine_timesteps",
+    "plan_refine_resolutions",
+    "prepare_refine_latents",
+    "refine_sigma_from_schedule",
+    "run_dmd_loop",
+    "run_two_pass_dmd",
+    "upsample_latents_spatial",
+]
+
+
+def _draw_noise_like(like: Any, *, seed: int | None, is_mlx: bool) -> Any:
+    """Generate Gaussian noise with the shape and array type of the input."""
+    shape = tuple(int(s) for s in like.shape)
+    if seed is not None:
+        rng = np.random.default_rng(seed)
+        noise_np = rng.standard_normal(shape).astype(np.float32)
+        if is_mlx:
+            import mlx.core as mx
+
+            return mx.array(noise_np).astype(mx.float32)
+        return noise_np.astype(np.asarray(like).dtype, copy=False)
+    if is_mlx:
+        import mlx.core as mx
+
+        return mx.random.normal(shape).astype(mx.float32)
+    return np.random.standard_normal(shape).astype(np.asarray(like).dtype, copy=False)

fastvideo/mlx_runtime/rife_interp.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Small MLX RIFE wrapper for frame interpolation experiments.
+
+The backend is the Apple-Silicon-native ``rife-mlx`` package, using the
+``mlx-community/RIFE-4.25`` weights. Frames are HWC RGB ``uint8`` arrays.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from functools import lru_cache
+
+import numpy as np
+from huggingface_hub.utils import LocalEntryNotFoundError
+
+
+class RIFEBackendError(RuntimeError):
+    """Raised when the MLX RIFE backend cannot be loaded or run."""
+
+
+class RIFEWeightsUnavailableError(RIFEBackendError):
+    """Raised when uncached RIFE weights cannot be downloaded."""
+
+
+def aligned_keyframe_count(target_frames: int, factor: int, temporal_compression: int = 4) -> int:
+    """Return the smallest VAE-aligned keyframe count that RIFE can expand to the target."""
+    if target_frames < 1:
+        raise ValueError(f"target_frames must be >= 1, got {target_frames}")
+    if factor < 1:
+        raise ValueError(f"factor must be >= 1, got {factor}")
+    if temporal_compression < 1:
+        raise ValueError(f"temporal_compression must be >= 1, got {temporal_compression}")
+    required_intervals = (target_frames - 1 + factor - 1) // factor
+    aligned_intervals = ((required_intervals + temporal_compression - 1) // temporal_compression * temporal_compression)
+    return aligned_intervals + 1
+
+
+def _require_hwc_rgb(frame: np.ndarray, index: int) -> np.ndarray:
+    array = np.asarray(frame)
+    if array.ndim != 3 or array.shape[2] != 3:
+        raise ValueError(f"frame {index} must have shape HxWx3, got {array.shape}")
+    if array.dtype != np.uint8:
+        array = np.clip(array, 0, 255).astype(np.uint8)
+    return np.ascontiguousarray(array)
+
+
+@lru_cache(maxsize=2)
+def load_model(version: str = "4.25", weights_dir: str | None = None):
+    """Load the MLX-native RIFE model.
+
+    ``weights_dir`` is passed through to ``build_model`` in the vendored ``rife_mlx``.
+    When it is ``None``, the package downloads/uses the Hugging Face
+    ``mlx-community/RIFE-4.25`` snapshot.
+    """
+    try:
+        from fastvideo.third_party.rife_mlx.utils.weights import build_model
+    except ImportError:
+        # Fall back to a separately installed upstream package, for anyone who
+        # already has one in the environment.
+        try:
+            from rife_mlx.utils.weights import build_model
+        except ImportError as exc:
+            raise RIFEBackendError("MLX RIFE backend is unavailable. It ships vendored under "
+                                   "fastvideo/third_party/rife_mlx, so this usually means MLX "
+                                   "itself is missing: install with `uv pip install -e '.[mlx]'`.") from exc
+
+    try:
+        return build_model(version, weights_dir=weights_dir)
+    except LocalEntryNotFoundError as exc:
+        raise RIFEWeightsUnavailableError(f"MLX RIFE {version} weights are unavailable: {exc}") from exc
+    except Exception as exc:  # noqa: BLE001 - preserve exact backend failure.
+        raise RIFEBackendError(f"Failed to load MLX RIFE {version}: {exc}") from exc
+
+
+def interpolate_pair(
+    frame_a: np.ndarray,
+    frame_b: np.ndarray,
+    timestep: float = 0.5,
+    *,
+    model=None,
+    scale: float = 1.0,
+) -> np.ndarray:
+    """Interpolate one RGB frame between two input RGB frames."""
+    if not 0.0 < timestep < 1.0:
+        raise ValueError(f"timestep must be inside (0, 1), got {timestep}")
+    img0 = _require_hwc_rgb(frame_a, 0)
+    img1 = _require_hwc_rgb(frame_b, 1)
+    if img0.shape != img1.shape:
+        raise ValueError(f"frame shapes must match, got {img0.shape} and {img1.shape}")
+
+    if model is None:
+        model = load_model()
+    try:
+        try:
+            from fastvideo.third_party.rife_mlx.pipeline_mlx import interpolate_pair as _interpolate_pair
+        except ImportError:
+            from rife_mlx.pipeline_mlx import interpolate_pair as _interpolate_pair
+
+        return _interpolate_pair(model, img0, img1, timestep=timestep, scale=scale)
+    except Exception as exc:  # noqa: BLE001 - preserve exact backend failure.
+        raise RIFEBackendError(f"MLX RIFE interpolation failed at timestep={timestep}: {exc}") from exc
+
+
+def interpolate(
+    frames: list[np.ndarray] | Iterable[np.ndarray],
+    factor: int = 2,
+    *,
+    model=None,
+    scale: float = 1.0,
+) -> list[np.ndarray]:
+    """Return an Nx interpolated frame list.
+
+    For ``len(frames)=41`` and ``factor=2``, the output length is 81:
+    ``(41 - 1) * 2 + 1``. Original keyframes are preserved in order and RIFE
+    fills ``factor - 1`` intermediate timesteps between each adjacent pair.
+    """
+    frame_list = [_require_hwc_rgb(frame, idx) for idx, frame in enumerate(frames)]
+    if factor < 1:
+        raise ValueError(f"factor must be >= 1, got {factor}")
+    if len(frame_list) < 2 or factor == 1:
+        return [frame.copy() for frame in frame_list]
+
+    first_shape = frame_list[0].shape
+    for idx, frame in enumerate(frame_list[1:], start=1):
+        if frame.shape != first_shape:
+            raise ValueError(f"all frames must have the same shape; frame 0={first_shape}, frame {idx}={frame.shape}")
+
+    if model is None:
+        model = load_model()
+
+    out: list[np.ndarray] = []
+    for left, right in zip(frame_list[:-1], frame_list[1:], strict=True):
+        out.append(left)
+        for step in range(1, factor):
+            out.append(interpolate_pair(left, right, step / factor, model=model, scale=scale))
+    out.append(frame_list[-1])
+    return out

fastvideo/mlx_runtime/sampling.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+"""On-device (MLX) DMD sampling for the FastWan runtime.
+
+The hybrid proof-of-concept ran the FastWan DiT in MLX but bounced every
+denoising step back through torch/NumPy to run the DMD scheduler math
+(``MLX -> np.array -> torch (CPU) -> np.array -> MLX``). That host round-trip
+forces a full device sync per step and defeats MLX's lazy graph execution.
+
+This module mirrors the exact DMD arithmetic from
+``fastvideo/models/utils.py::pred_noise_to_pred_video`` and
+``FlowMatchEulerDiscreteScheduler.add_noise`` while keeping every large tensor
+on the MLX device. The schedule lookup (``argmin`` over the ~1000-entry
+training schedule) is done once on the host in NumPy: it is tiny, it is the
+same value torch would compute, and it sidesteps the reduction-index quirk that
+affects ``argmin`` on the Metal/MPS backends (see the CPU fallbacks in
+``fastvideo/models/utils.py`` and ``scheduling_flow_match_euler_discrete.py``).
+
+Because the DMD loop applies a single scalar timestep per step, ``sigma`` is a
+scalar and the update is a plain elementwise affine combination — no
+permute/flatten reshaping is required.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    import mlx.core as mx
+
+
+@dataclass(frozen=True)
+class MLXDMDSchedule:
+    """Host-side copy of a flow-match scheduler's ``(sigmas, timesteps)``.
+
+    Holds the full training schedule so a DMD timestep (e.g. one of
+    ``1000, 757, 522``) can be mapped to its flow-match ``sigma`` with the same
+    nearest-timestep lookup the torch path uses.
+    """
+
+    sigmas: np.ndarray
+    timesteps: np.ndarray
+
+    @classmethod
+    def from_torch_scheduler(cls, scheduler: Any) -> MLXDMDSchedule:
+        """Snapshot ``scheduler.sigmas`` / ``scheduler.timesteps`` to NumPy.
+
+        Matches ``pred_noise_to_pred_video`` / ``add_noise``, which index the
+        scheduler's *full* training schedule (not the per-inference subset).
+        """
+        sigmas = scheduler.sigmas.detach().to("cpu").double().numpy()
+        timesteps = scheduler.timesteps.detach().to("cpu").double().numpy()
+        return cls(sigmas=np.asarray(sigmas), timesteps=np.asarray(timesteps))
+
+    def sigma_for(self, timestep: float) -> float:
+        """
+        Find the sigma associated with the scheduled timestep nearest to the given timestep.
+
+        Parameters:
+            timestep (float): Timestep for which to find the nearest scheduled sigma.
+
+        Returns:
+            float: Sigma associated with the nearest scheduled timestep.
+        """
+        idx = int(np.argmin(np.abs(self.timesteps - float(timestep))))
+        return float(self.sigmas[idx])
+
+
+def pred_noise_to_pred_video(
+    pred_noise: mx.array,
+    noise_input_latent: mx.array,
+    sigma: float,
+) -> mx.array:
+    """
+    Compute the clean latent prediction from a flow-matching noise prediction.
+
+    Parameters:
+        pred_noise (mx.array): Predicted noise.
+        noise_input_latent (mx.array): Noised latent input.
+        sigma (float): Noise level used for the prediction.
+
+    Returns:
+        mx.array: Predicted clean latent.
+    """
+    return noise_input_latent - sigma * pred_noise
+
+
+def add_noise(
+    clean_latent: mx.array,
+    noise: mx.array,
+    sigma: float,
+) -> mx.array:
+    """Flow-match forward noising, mirroring the scheduler's ``add_noise``.
+
+    ``sample = (1 - sigma) * clean_latent + sigma * noise``.
+    """
+    return (1.0 - sigma) * clean_latent + sigma * noise
+
+
+def dmd_step(
+    *,
+    latents: mx.array,
+    noise_input_latent: mx.array,
+    pred_noise: mx.array,
+    schedule: MLXDMDSchedule,
+    timestep: float,
+    next_timestep: float | None,
+    noise: mx.array | None = None,
+) -> mx.array:
+    """
+    Compute one DMD sampling update, optionally re-noising the clean latent prediction.
+
+    Args:
+        latents: Retained for call-site compatibility and not used in the update.
+        noise_input_latent: Noisy latent used to compute the clean prediction.
+        pred_noise: Predicted noise or velocity.
+        schedule: Flow-matching schedule used to map timesteps to sigmas.
+        timestep: Current sampling timestep.
+        next_timestep: Timestep for the next update, or `None` for the final step.
+        noise: Fresh noise used for re-noising intermediate steps.
+
+    Returns:
+        The re-noised latent for the next step or the clean latent prediction on
+        the final step.
+
+    Raises:
+        ValueError: If `next_timestep` is provided without `noise`.
+    """
+    del latents  # symmetry with the torch loop; not needed for the math.
+    sigma = schedule.sigma_for(timestep)
+    pred_video = pred_noise_to_pred_video(pred_noise, noise_input_latent, sigma)
+    if next_timestep is None:
+        return pred_video
+    if noise is None:
+        raise ValueError("dmd_step requires `noise` when `next_timestep` is set (re-noise step).")
+    sigma_next = schedule.sigma_for(next_timestep)
+    return add_noise(pred_video, noise, sigma_next)

fastvideo/mlx_runtime/taehv_decode.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional TAEHV decode helpers for Apple Silicon FastWan experiments.
+
+The TAEHV module itself is vendored at ``fastvideo/third_party/taehv`` (MIT,
+madebyollin/taehv), so no source code is downloaded or executed at runtime.
+Only the ``taew2_1.pth`` checkpoint is fetched on demand, and its sha256 is
+verified before use.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+
+TAEW2_1_CHECKPOINT_URL = "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth"
+# sha256 of the upstream taew2_1.pth this module was validated against
+# (fetched 2026-07-02). If upstream publishes a new checkpoint, revalidate the
+# decode path and update this pin.
+TAEW2_1_CHECKPOINT_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e"
+# Wan2.2 5B (z_dim=48) — see madebyollin/taehv taew2_2.pth; prefer
+# ``fastvideo.mlx_runtime.wan_vae.ensure_taehv_checkpoint(z_dim=48)`` for new code.
+TAEW2_2_CHECKPOINT_URL = ("https://raw.githubusercontent.com/madebyollin/taehv/"
+                          "563f40bdc820ed86bcad72ea515ee48f06bd22ec/taew2_2.pth")
+
+
+def _default_cache_dir() -> Path:
+    """Return the default directory used to cache TAEHV checkpoints.
+
+    Returns:
+        Path: The TAEHV checkpoint cache directory under the user's home directory.
+    """
+    return Path.home() / ".cache" / "fastvideo" / "taehv"
+
+
+def _sha256(path: Path) -> str:
+    """
+    Compute the SHA-256 digest of a file.
+
+    Parameters:
+        path (Path): The file whose contents are hashed.
+
+    Returns:
+        str: The file's SHA-256 digest in hexadecimal form.
+    """
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _verify_checkpoint(path: Path) -> None:
+    """Verify that a TAEW2.1 checkpoint matches the expected SHA-256 digest.
+
+    Parameters:
+        path (Path): Path to the checkpoint file.
+
+    Raises:
+        RuntimeError: If the checkpoint digest does not match the expected value.
+    """
+    actual = _sha256(path)
+    if actual != TAEW2_1_CHECKPOINT_SHA256:
+        raise RuntimeError(f"TAEHV checkpoint at {path} failed sha256 verification "
+                           f"(expected {TAEW2_1_CHECKPOINT_SHA256}, got {actual}). "
+                           "Delete the file to re-download it, or pass --taehv-checkpoint-path "
+                           "pointing at a checkpoint you trust.")
+
+
+def ensure_taew2_1_checkpoint(checkpoint_path: Path | None = None) -> Path:
+    """
+    Ensure the TAEW2.1 checkpoint is available locally.
+
+    A caller-provided path is treated as trusted and is only checked for existence.
+    The module-managed cached checkpoint is verified against the pinned SHA-256 digest
+    after downloading or before reuse.
+
+    Parameters:
+        checkpoint_path (Path | None): Optional path to a caller-provided checkpoint.
+
+    Returns:
+        Path: The available checkpoint path.
+
+    Raises:
+        FileNotFoundError: If a caller-provided checkpoint does not exist.
+        RuntimeError: If a module-managed checkpoint fails verification.
+    """
+    if checkpoint_path is not None:
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"TAEHV checkpoint not found: {checkpoint_path}")
+        return checkpoint_path
+
+    checkpoint_path = _default_cache_dir() / "taew2_1.pth"
+    if not checkpoint_path.exists():
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        print(f"Downloading {TAEW2_1_CHECKPOINT_URL} -> {checkpoint_path}")
+        import socket
+        import tempfile
+        # Download to a temporary file, verify, then atomically rename.
+        with tempfile.NamedTemporaryFile(
+                mode="wb",
+                dir=checkpoint_path.parent,
+                prefix=".tmp_taew2_1_",
+                suffix=".pth",
+                delete=False,
+        ) as tmp_file:
+            tmp_path = Path(tmp_file.name)
+            try:
+                old_timeout = socket.getdefaulttimeout()
+                socket.setdefaulttimeout(300)
+                try:
+                    urllib.request.urlretrieve(
+                        TAEW2_1_CHECKPOINT_URL,
+                        tmp_path,  # noqa: S310 - pinned public artifact, hash-verified below.
+                    )
+                finally:
+                    socket.setdefaulttimeout(old_timeout)
+                _verify_checkpoint(tmp_path)
+                tmp_path.replace(checkpoint_path)
+            except Exception:
+                tmp_path.unlink(missing_ok=True)
+                raise
+    else:
+        _verify_checkpoint(checkpoint_path)
+    return checkpoint_path
+
+
+def _load_taehv_class(source_path: Path | None):
+    """Load the TAEHV class from the vendored implementation or a local source override.
+
+    Parameters:
+        source_path (Path | None): Path to a local Python file defining `TAEHV`; `None` selects the vendored implementation.
+
+    Returns:
+        The loaded `TAEHV` class.
+
+    Raises:
+        RuntimeError: If the specified source cannot be loaded.
+    """
+    if source_path is None:
+        from fastvideo.third_party.taehv import TAEHV
+
+        return TAEHV
+    # Explicit local override for experimenting with a modified TAEHV; this is
+    # a user-supplied file on disk, never something this module downloads.
+    spec = importlib.util.spec_from_file_location("fastvideo_external_taehv", source_path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Could not load TAEHV source from {source_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.TAEHV
+
+
+def decode_latents_to_video_taehv(
+    *,
+    latents_np: np.ndarray,
+    output_path: Path,
+    fps: int,
+    device,
+    dtype,
+    parallel: bool,
+    source_path: Path | None = None,
+    checkpoint_path: Path | None = None,
+) -> None:
+    """Decode Wan/FastWan diffusion latents with TAEW2.1 and export MP4.
+
+    TAEHV's Wan wrapper expects the diffusion latents directly, without applying
+    the standard Wan VAE's `latents_mean` / `latents_std` shift.
+    """
+    import torch
+    from diffusers.utils import export_to_video
+
+    checkpoint_path = ensure_taew2_1_checkpoint(checkpoint_path)
+    TAEHV = _load_taehv_class(source_path)
+    taehv = TAEHV(str(checkpoint_path)).to(device=device, dtype=dtype)
+    taehv.eval()
+
+    latents = torch.from_numpy(latents_np).to(device=device, dtype=dtype)
+    with torch.no_grad():
+        video_ntchw = taehv.decode_video(
+            latents.transpose(1, 2),
+            parallel=parallel,
+            show_progress_bar=False,
+        )
+    video = video_ntchw.transpose(1, 2)
+    video_np = video[0].permute(1, 2, 3, 0).float().cpu().numpy()
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    export_to_video(video_np, str(output_path), fps=fps)

fastvideo/mlx_runtime/wan22.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan2.2-TI2V-5B dense MLX runtime — Track D.
+
+The Wan2.2 TI2V-5B (FullAttn) differs from the ported Wan2.1-T2V only in:
+
+- **Scale** (24 heads x 128, hidden 3072, ffn 14336) — pure config, block math
+  identical, so the dense loader ``mlx_dit_from_diffusers_safetensors`` loads the
+  weights unchanged and we re-wrap the blocks here.
+- **Per-token timestep conditioning** (``expand_timesteps=True``): the timestep is
+  ``[batch, seq_len]`` (a level per patch token — how TI2V keeps the conditioning
+  image frame at t=0 while the video frames are noised). ``timestep_proj`` becomes
+  ``[batch, seq_len, 6, dim]`` and the block/output modulation is per-token
+  (``[B, L, dim]``), a direct broadcast — this module implements exactly that.
+
+I2V rides on the same forward: encode the image, replace the first latent frame,
+and set that frame's timestep to 0 (handled by the caller / sampler). See
+``docs/design/ti2v_5b_port_guide.md``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from collections.abc import Callable
+
+from fastvideo.logger import init_logger
+from fastvideo.mlx_runtime.fastwan import (
+    MLXWanT2VCrossAttention,
+    gelu_tanh,
+    layer_norm,
+    linear,
+    mlx_dit_from_diffusers_safetensors,
+    rms_norm,
+    silu,
+    timestep_embedding,
+    weight_dtype,
+)
+
+if TYPE_CHECKING:
+    import mlx.core as mx
+
+logger = init_logger(__name__)
+
+
+class MLXWan22TransformerBlock:
+    """Dense Wan block with per-token (``[B, L, dim]``) timestep modulation."""
+
+    def __init__(
+        self,
+        weights: dict[str, mx.array],
+        *,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+    ):
+        self.weights = weights
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.eps = eps
+        self.attn2 = MLXWanT2VCrossAttention(weights, dim=dim, num_heads=num_heads, eps=eps)
+
+    def __call__(self, hidden_states, encoder_hidden_states, timestep_proj, cos, sin) -> mx.array:
+        import mlx.core as mx
+
+        orig_dtype = hidden_states.dtype
+        batch = hidden_states.shape[0]
+
+        # timestep_proj: [B, L, 6, dim] -> six per-token [B, L, dim] modulations.
+        e = self.weights["scale_shift_table"][None] + timestep_proj.astype(mx.float32)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = [
+            part.squeeze(2) for part in mx.split(e, 6, axis=2)
+        ]
+
+        # 1. Self-attention (dense, bidirectional) with per-token modulation.
+        norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
+        norm_hidden = (norm_hidden * (1.0 + scale_msa) + shift_msa).astype(orig_dtype)
+
+        query = linear(norm_hidden, self.weights["to_q.weight"], self.weights.get("to_q.bias"))
+        key = linear(norm_hidden, self.weights["to_k.weight"], self.weights.get("to_k.bias"))
+        value = linear(norm_hidden, self.weights["to_v.weight"], self.weights.get("to_v.bias"))
+
+        query = rms_norm(query, self.weights["norm_q.weight"],
+                         eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim)
+        key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(batch, -1, self.num_heads,
+                                                                                 self.head_dim)
+        value = value.reshape(batch, -1, self.num_heads, self.head_dim)
+
+        from fastvideo.mlx_runtime.fastwan import apply_rotary_emb
+        query = apply_rotary_emb(query, cos, sin, is_neox_style=False)
+        key = apply_rotary_emb(key, cos, sin, is_neox_style=False)
+
+        attn = mx.fast.scaled_dot_product_attention(
+            query.transpose(0, 2, 1, 3),
+            key.transpose(0, 2, 1, 3),
+            value.transpose(0, 2, 1, 3),
+            scale=self.head_dim**-0.5,
+        ).transpose(0, 2, 1, 3)
+        attn = attn.reshape(batch, -1, self.dim)
+        attn = linear(attn, self.weights["to_out.weight"], self.weights.get("to_out.bias"))
+
+        hidden_states = hidden_states + (attn * gate_msa).astype(orig_dtype)
+        norm_hidden = layer_norm(hidden_states.astype(mx.float32),
+                                 weight=self.weights["self_attn_residual_norm.norm.weight"],
+                                 bias=self.weights["self_attn_residual_norm.norm.bias"],
+                                 eps=self.eps).astype(orig_dtype)
+
+        # 2. Cross-attention, then per-token shift/scale modulation.
+        cross = self.attn2(norm_hidden, encoder_hidden_states)
+        hidden_states = hidden_states + cross
+        norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
+        norm_hidden = (norm_hidden * (1.0 + c_scale_msa) + c_shift_msa).astype(orig_dtype)
+
+        # 3. Feed-forward with per-token gate.
+        ff = linear(norm_hidden, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias"))
+        ff = gelu_tanh(ff)
+        ff = linear(ff, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias"))
+        hidden_states = hidden_states + (ff * c_gate_msa).astype(orig_dtype)
+        return hidden_states.astype(orig_dtype)
+
+
+class MLXWan22DiT:
+    """Wan2.2-TI2V-5B dense DiT with per-token timestep conditioning."""
+
+    def __init__(
+        self,
+        weights: dict[str, mx.array],
+        blocks: list[MLXWan22TransformerBlock],
+        config: dict,
+        *,
+        compile: bool = False,
+    ) -> None:
+        import os
+
+        self.weights = weights
+        self.blocks = blocks
+        self.config = config
+        self.num_heads = int(config["num_attention_heads"])
+        self.head_dim = int(config["attention_head_dim"])
+        self.hidden_size = self.num_heads * self.head_dim
+        self.freq_dim = int(config["freq_dim"])
+        self.patch_size = tuple(config["patch_size"])
+        self.out_channels = int(config["out_channels"])
+        self.eps = float(config.get("eps", 1e-6))
+        self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1"
+        self._compiled_forward: Callable[..., Any] | None = None
+        self._compiled_signature: tuple | None = None
+
+    def _patch_embed(self, hidden_states) -> mx.array:
+        batch, channels, frames, height, width = hidden_states.shape
+        pt, ph, pw = self.patch_size
+        patch_dim = channels * pt * ph * pw
+        x = hidden_states.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw)
+        x = x.transpose(0, 2, 4, 6, 1, 3, 5, 7).reshape(batch, -1, patch_dim)
+        return linear(x, self.weights["patch_embedding.weight"], self.weights.get("patch_embedding.bias"))
+
+    def _condition(self, timestep, encoder_hidden_states) -> tuple:
+        """Per-token conditioning. ``timestep`` is ``[B, L]`` (one level per token)."""
+        batch, seq = timestep.shape
+        t_freq = timestep_embedding(timestep.reshape(-1), self.freq_dim).astype(
+            weight_dtype(self.weights["condition_embedder.time_embedder.linear_1.weight"]))
+        temb = linear(t_freq, self.weights["condition_embedder.time_embedder.linear_1.weight"],
+                      self.weights["condition_embedder.time_embedder.linear_1.bias"])
+        temb = silu(temb)
+        temb = linear(temb, self.weights["condition_embedder.time_embedder.linear_2.weight"],
+                      self.weights["condition_embedder.time_embedder.linear_2.bias"])
+        timestep_proj = linear(silu(temb), self.weights["condition_embedder.time_proj.weight"],
+                               self.weights["condition_embedder.time_proj.bias"])
+        timestep_proj = timestep_proj.reshape(batch, seq, 6, self.hidden_size)
+
+        ehs = linear(encoder_hidden_states, self.weights["condition_embedder.text_embedder.linear_1.weight"],
+                     self.weights["condition_embedder.text_embedder.linear_1.bias"])
+        ehs = gelu_tanh(ehs)
+        ehs = linear(ehs, self.weights["condition_embedder.text_embedder.linear_2.weight"],
+                     self.weights["condition_embedder.text_embedder.linear_2.bias"])
+        temb_out = temb.reshape(batch, seq, self.hidden_size)
+        return temb_out, timestep_proj, ehs
+
+    def _output(self, hidden_states, temb_out, *, batch, frames, height, width) -> mx.array:
+        import mlx.core as mx
+
+        pt, ph, pw = self.patch_size
+        post_pt, post_ph, post_pw = frames // pt, height // ph, width // pw
+        # Per-token output modulation: scale_shift_table[1,2,dim] + temb[B,L,1,dim].
+        e = self.weights["scale_shift_table"][None] + temb_out[:, :, None, :].astype(mx.float32)
+        shift, scale = [part.squeeze(2) for part in mx.split(e, 2, axis=2)]
+        norm = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
+        norm = (norm * (1.0 + scale) + shift).astype(weight_dtype(self.weights["proj_out.weight"]))
+        out = linear(norm, self.weights["proj_out.weight"], self.weights["proj_out.bias"])
+        out = out.reshape(batch, post_pt, post_ph, post_pw, pt, ph, pw, self.out_channels)
+        out = out.transpose(0, 7, 1, 4, 2, 5, 3, 6)
+        return out.reshape(batch, self.out_channels, frames, height, width)
+
+    def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin) -> mx.array:
+        batch, _, frames, height, width = hidden_states.shape
+        hidden = self._patch_embed(hidden_states)
+        temb_out, timestep_proj, ehs = self._condition(timestep, encoder_hidden_states)
+        for block in self.blocks:
+            hidden = block(hidden, ehs, timestep_proj, cos, sin)
+        return self._output(hidden, temb_out, batch=batch, frames=frames, height=height, width=width)
+
+    def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> mx.array:
+        cos, sin = freqs_cis
+        if self._enable_compile and cos is not None:
+            import mlx.core as mx
+
+            # One traced graph per input signature, and each pins its own copy of
+            # the quantized weights. --refine denoises at two resolutions, so
+            # keeping both alive doubles resident DiT memory. Retire the previous
+            # graph when the signature changes.
+            signature = (hidden_states.shape, encoder_hidden_states.shape, timestep.shape)
+            if self._compiled_forward is not None and signature != self._compiled_signature:
+                self._compiled_forward = None
+                self._compiled_signature = None
+                mx.clear_cache()
+            if self._compiled_forward is None:
+                self._compiled_forward = mx.compile(self._forward)
+                self._compiled_signature = signature
+            compiled_forward = self._compiled_forward
+            try:
+                return compiled_forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
+            except Exception as exc:  # noqa: BLE001 - some graphs may not trace; fall back to eager.
+                logger.warning(
+                    "Wan2.2 mx.compile forward failed (%s); falling back to eager execution.",
+                    exc,
+                )
+                self._enable_compile = False
+                self._compiled_forward = None
+                self._compiled_signature = None
+        return self._forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
+
+
+def mlx_wan22_dit_from_diffusers_safetensors(
+    checkpoint_path: str | Path,
+    config_path: str | Path,
+    *,
+    dtype: str = "fp16",
+    num_blocks: int | None = None,
+    quantization=None,
+    compile: bool = False,
+) -> MLXWan22DiT:
+    """Load Wan2.2-TI2V-5B (FullAttn) into ``MLXWan22DiT`` via the dense loader."""
+    dense = mlx_dit_from_diffusers_safetensors(checkpoint_path,
+                                               config_path,
+                                               dtype=dtype,
+                                               num_blocks=num_blocks,
+                                               quantization=quantization)
+    inner_dim = int(dense.config["num_attention_heads"]) * int(dense.config["attention_head_dim"])
+    blocks = [
+        MLXWan22TransformerBlock(block.weights,
+                                 dim=inner_dim,
+                                 ffn_dim=int(dense.config["ffn_dim"]),
+                                 num_heads=int(dense.config["num_attention_heads"]),
+                                 eps=float(dense.config.get("eps", 1e-6))) for block in dense.blocks
+    ]
+    return MLXWan22DiT(dense.weights, blocks, dense.config, compile=compile)
+
+
+def mlx_wan22_dit_from_mlx_checkpoint(
+    checkpoint_dir: str | Path,
+    *,
+    compile: bool = False,
+) -> MLXWan22DiT:
+    """Rewrap a persisted MLX DiT checkpoint with Wan2.2 conditioning.
+
+    The generic checkpoint loader intentionally rebuilds ``MLXWanDiT`` because
+    it is also used by the Wan2.1 runtime.  Wan2.2 TI2V has the same weight
+    layout but needs per-token timestep modulation, so callers must rewrap the
+    loaded weights and blocks as :class:`MLXWan22DiT` before sampling.
+    """
+    from fastvideo.mlx_runtime.checkpoint import load_mlx_dit_checkpoint
+
+    dense = load_mlx_dit_checkpoint(checkpoint_dir)
+    inner_dim = int(dense.config["num_attention_heads"]) * int(dense.config["attention_head_dim"])
+    blocks = [
+        MLXWan22TransformerBlock(
+            block.weights,
+            dim=inner_dim,
+            ffn_dim=int(dense.config["ffn_dim"]),
+            num_heads=int(dense.config["num_attention_heads"]),
+            eps=float(dense.config.get("eps", 1e-6)),
+        ) for block in dense.blocks
+    ]
+    return MLXWan22DiT(dense.weights, blocks, dense.config, compile=compile)

fastvideo/mlx_runtime/wan22_sample.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dense DMD sampling for MLXWan22DiT (Wan2.2 per-token timestep).
+
+Matches the FastVideo pipeline's warped DMD schedule (``warp_denoising_step=True``,
+``dmd_denoising_steps=[1000,757,522]``, ``flow_shift=5.0`` for TI2V-5B) rather
+than treating raw step indices as continuous timesteps (a bug in early demos).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from collections.abc import Sequence
+
+import numpy as np
+
+from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step, pred_noise_to_pred_video
+
+if TYPE_CHECKING:
+    import mlx.core as mx
+
+    from fastvideo.mlx_runtime.wan22 import MLXWan22DiT
+
+
+def build_wan22_dmd_schedule(
+    dmd_denoising_steps: Sequence[int] | None = None,
+    *,
+    flow_shift: float = 5.0,
+    warp_denoising_step: bool = True,
+) -> tuple[MLXDMDSchedule, list[float]]:
+    """
+    Build the flow-matching schedule and continuous timesteps used for Wan2.2 DMD sampling.
+
+    Parameters:
+        dmd_denoising_steps (Sequence[int] | None): Denoising step values to use; defaults to 1000, 757, and 522.
+        flow_shift (float): Flow-matching shift applied when constructing the schedule.
+        warp_denoising_step (bool): Whether to convert denoising steps to scheduler-warped continuous timesteps.
+
+    Returns:
+        tuple[MLXDMDSchedule, list[float]]: The DMD schedule and corresponding continuous timesteps.
+    """
+    import torch
+
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+    steps = list(dmd_denoising_steps or [1000, 757, 522])
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift)
+    schedule = MLXDMDSchedule.from_torch_scheduler(scheduler)
+    step_idx = torch.tensor(steps, dtype=torch.long)
+    if warp_denoising_step:
+        warped = torch.cat((scheduler.timesteps.cpu(), torch.tensor([0.0], dtype=torch.float32)))
+        timesteps = [float(t) for t in warped[1000 - step_idx]]
+    else:
+        timesteps = [float(s) for s in steps]
+    return schedule, timesteps
+
+
+def sample_wan22_dmd(
+    model: MLXWan22DiT,
+    encoder_hidden_states: mx.array,
+    noise_latents: mx.array,
+    freqs_cis: tuple,
+    *,
+    dmd_denoising_steps: Sequence[int] | None = None,
+    flow_shift: float = 5.0,
+    warp_denoising_step: bool = True,
+    seed: int = 0,
+) -> mx.array:
+    """
+    Generate clean video latents from noisy latents using iterative DMD denoising.
+
+    Parameters:
+        noise_latents (mx.array): Initial noisy video latents.
+        freqs_cis (tuple): Rotary positional frequency tensors used by the model.
+        dmd_denoising_steps (Sequence[int] | None): DMD denoising steps, or the default schedule when omitted.
+        flow_shift (float): Flow-matching schedule shift.
+        warp_denoising_step (bool): Whether to warp the denoising timesteps.
+        seed (int): Seed for reproducible intermediate re-noising.
+
+    Returns:
+        mx.array: Denoised video latents.
+    """
+    import mlx.core as mx
+
+    schedule, timesteps = build_wan22_dmd_schedule(dmd_denoising_steps,
+                                                   flow_shift=flow_shift,
+                                                   warp_denoising_step=warp_denoising_step)
+    # NumPy RNG so re-noise is bit-reproducible across MLX / torch A/B dumps.
+    renoise_rng = np.random.default_rng(seed)
+    latents = noise_latents
+    batch, _c, frames, height, width = latents.shape
+    pt, ph, pw = model.patch_size
+    tokens = (frames // pt) * (height // ph) * (width // pw)
+    last = len(timesteps) - 1
+    for i, t in enumerate(timesteps):
+        ts = mx.full((batch, tokens), float(t), dtype=mx.float32)
+        pred = model(latents.astype(mx.float16), encoder_hidden_states, ts, freqs_cis)
+        ni = latents.astype(mx.float32)
+        pn = pred.astype(mx.float32)
+        if i < last:
+            renoise = mx.array(renoise_rng.standard_normal(tuple(latents.shape)).astype(np.float32))
+            latents = dmd_step(
+                latents=ni,
+                noise_input_latent=ni,
+                pred_noise=pn,
+                schedule=schedule,
+                timestep=float(t),
+                next_timestep=float(timesteps[i + 1]),
+                noise=renoise,
+            ).astype(latents.dtype)
+        else:
+            latents = pred_noise_to_pred_video(pn, ni, schedule.sigma_for(float(t))).astype(latents.dtype)
+        mx.eval(latents)
+    return latents

fastvideo/mlx_runtime/wan_vae.py

@@ -0,0 +1,557 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan VAE decode helpers for Apple Silicon MLX inference.
+
+Two decode backends:
+
+1. **TAEHV (primary / fast)** — Tiny AutoEncoder (madebyollin/taehv). Fully
+   MLX-native Conv2d path. ``taew2_1.pth`` for Wan2.1 (z_dim=16),
+   ``taew2_2.pth`` for Wan2.2 5B (z_dim=48, patch_size=2). Expected decode
+   wall-clock ~seconds vs ~minutes for the full 3D VAE on MPS.
+
+2. **Full AutoencoderKLWan (reference / quality)** — denormalize with
+   ``latents_mean`` / ``latents_std`` then torch decode (MPS preferred). Used
+   for parity gates and when TAEHV is unavailable. A pure-MLX 3D-conv port of
+   the residual Wan2.2 decoder is left as follow-up (causal feat-cache +
+   residual up blocks are large); TAEHV covers the product latency path.
+
+Diffusion latents from the DiT are **not** mean/std-normalized for TAEHV
+(matching ``taehv_decode.py``); full VAE decode **does** denormalize first
+(matching ``mlx_wan_prompt_to_video.decode_latents_to_video``).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+import numpy as np
+
+from fastvideo.mlx_runtime.memory import cleanup_mlx, cleanup_torch_mps
+
+GIB = 1024**3
+
+TAEW2_1_URL = "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth"
+TAEW2_2_URL = ("https://raw.githubusercontent.com/madebyollin/taehv/"
+               "563f40bdc820ed86bcad72ea515ee48f06bd22ec/taew2_2.pth")
+# Validated 2026-07-02 / 2026-07-09 against upstream madebyollin/taehv.
+TAEW2_1_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e"
+TAEW2_2_SHA256 = "d053e216ca50e2bb837bbcd79b85f0366bea00e5938025572382a773b74c559a"
+
+DecodeBackend = Literal["taehv", "taehv-torch", "wan-vae"]
+
+
+def _cache_dir() -> Path:
+    """Return the local directory used to cache TAEHV files."""
+    return Path.home() / ".cache" / "fastvideo" / "taehv"
+
+
+def _sha256(path: Path) -> str:
+    """Compute the SHA-256 digest of a file.
+
+    Parameters:
+        path (Path): Path to the file to hash.
+
+    Returns:
+        str: Lowercase hexadecimal SHA-256 digest.
+    """
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _verify_checkpoint(path: Path, expected_digest: str) -> None:
+    """
+    Verify a checkpoint's SHA-256 digest against a required expected digest.
+
+    Parameters:
+        path (Path): Path to the checkpoint file.
+        expected_digest (str): Expected lowercase SHA-256 digest.
+
+    Raises:
+        RuntimeError: If verification is enabled and the checkpoint digest does not match.
+    """
+    if len(expected_digest) != 64 or any(char not in "0123456789abcdef" for char in expected_digest):
+        raise ValueError("A valid lowercase SHA-256 digest is required for bundled TAEHV checkpoints")
+    actual = _sha256(path)
+    if actual != expected_digest:
+        raise RuntimeError(f"TAEHV checkpoint at {path} failed sha256 verification "
+                           f"(expected {expected_digest}, got {actual}). "
+                           "Delete the file to re-download it.")
+
+
+def ensure_taehv_checkpoint(*, z_dim: int, checkpoint_path: Path | None = None) -> Path:
+    """
+    Return a validated TAEHV checkpoint for the specified latent channel count.
+
+    Parameters:
+        z_dim (int): Number of latent channels, supported values are 16 and 48.
+        checkpoint_path (Path | None): Optional existing checkpoint path to validate and use.
+
+    Returns:
+        Path: Path to the validated TAEHV checkpoint.
+
+    Raises:
+        FileNotFoundError: If the supplied checkpoint path does not exist.
+        ValueError: If no checkpoint is mapped to the specified latent channel count.
+    """
+    if checkpoint_path is not None:
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"TAEHV checkpoint not found: {checkpoint_path}")
+        return checkpoint_path
+    if z_dim == 16:
+        name, url, expect = "taew2_1.pth", TAEW2_1_URL, TAEW2_1_SHA256
+    elif z_dim == 48:
+        name, url, expect = "taew2_2.pth", TAEW2_2_URL, TAEW2_2_SHA256
+    else:
+        raise ValueError(f"No TAEHV checkpoint mapped for z_dim={z_dim} (supported: 16, 48)")
+    path = _cache_dir() / name
+    if not path.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        print(f"Downloading {url} -> {path}")
+        import socket
+        import tempfile
+        # Download to a temporary file, verify, then atomically rename.
+        with tempfile.NamedTemporaryFile(
+                mode="wb",
+                dir=path.parent,
+                prefix=f".tmp_{name}_",
+                suffix=".pth",
+                delete=False,
+        ) as tmp_file:
+            tmp_path = Path(tmp_file.name)
+            try:
+                old_timeout = socket.getdefaulttimeout()
+                socket.setdefaulttimeout(300)
+                try:
+                    urllib.request.urlretrieve(url, tmp_path)  # noqa: S310 - public pinned artifact.
+                finally:
+                    socket.setdefaulttimeout(old_timeout)
+                _verify_checkpoint(tmp_path, expect)
+                tmp_path.replace(path)
+            except Exception:
+                tmp_path.unlink(missing_ok=True)
+                raise
+    else:
+        _verify_checkpoint(path, expect)
+    return path
+
+
+@dataclass(frozen=True)
+class WanVAEConfigView:
+    """Minimal config fields needed for denormalize + spatial scale."""
+
+    z_dim: int
+    latents_mean: tuple[float, ...]
+    latents_std: tuple[float, ...]
+    scale_factor_spatial: int = 8
+    scale_factor_temporal: int = 4
+    patch_size: int | None = None
+    vae_dir: Path | None = None
+
+    @classmethod
+    def from_vae_dir(cls, vae_dir: Path) -> WanVAEConfigView:
+        """
+        Load Wan VAE configuration values from a directory.
+
+        Parameters:
+            vae_dir (Path): Directory containing the VAE ``config.json`` file.
+
+        Returns:
+            WanVAEConfigView: Configuration loaded from the VAE directory.
+        """
+        cfg = json.loads((vae_dir / "config.json").read_text())
+        return cls(
+            z_dim=int(cfg["z_dim"]),
+            latents_mean=tuple(float(x) for x in cfg["latents_mean"]),
+            latents_std=tuple(float(x) for x in cfg["latents_std"]),
+            scale_factor_spatial=int(cfg.get("scale_factor_spatial", 8)),
+            scale_factor_temporal=int(cfg.get("scale_factor_temporal", 4)),
+            patch_size=cfg.get("patch_size"),
+            vae_dir=vae_dir,
+        )
+
+
+def denormalize_latents_np(latents: np.ndarray, config: WanVAEConfigView) -> np.ndarray:
+    """
+    Denormalize Wan VAE latent values using the configured means and standard deviations.
+
+    Parameters:
+        latents (np.ndarray): Latent values in normalized form.
+        config (WanVAEConfigView): Wan VAE latent statistics.
+
+    Returns:
+        np.ndarray: Denormalized latent values as float32.
+    """
+    mean = np.asarray(config.latents_mean, dtype=np.float32).reshape(1, -1, 1, 1, 1)
+    std = np.asarray(config.latents_std, dtype=np.float32).reshape(1, -1, 1, 1, 1)
+    return latents.astype(np.float32) * std + mean
+
+
+# ---------------------------------------------------------------------------
+# MLX TAEHV decoder (Conv2d stack — primary fully-MLX product path)
+# ---------------------------------------------------------------------------
+
+
+def _mlx_conv2d(x: Any, weight: Any, bias: Any, *, stride: int = 1) -> Any:
+    import mlx.core as mx
+
+    # x: NCHW, weight: OIHW
+    y = mx.conv2d(x.transpose(0, 2, 3, 1), weight.transpose(0, 2, 3, 1), stride=stride, padding=1)
+    y = y.transpose(0, 3, 1, 2)
+    if bias is not None:
+        y = y + bias.reshape(1, -1, 1, 1)
+    return y
+
+
+def _mlx_conv2d_1x1(x: Any, weight: Any, bias: Any = None) -> Any:
+    """
+    Applies a 1×1 convolution to an MLX tensor in channel-first layout.
+
+    Parameters:
+        x (Any): Input tensor with shape [batch, channels, height, width].
+        weight (Any): Convolution weights.
+        bias (Any, optional): Optional output-channel bias.
+
+    Returns:
+        Any: The convolved tensor with shape [batch, output_channels, height, width].
+    """
+    import mlx.core as mx
+
+    y = mx.conv2d(x.transpose(0, 2, 3, 1), weight.transpose(0, 2, 3, 1), stride=1, padding=0)
+    y = y.transpose(0, 3, 1, 2)
+    if bias is not None:
+        y = y + bias.reshape(1, -1, 1, 1)
+    return y
+
+
+def _load_torch_state(path: Path) -> dict[str, np.ndarray]:
+    """Load a PyTorch state dictionary as NumPy arrays.
+
+    Parameters:
+        path (Path): Path to the PyTorch checkpoint.
+
+    Returns:
+        dict[str, np.ndarray]: State dictionary with tensors converted to NumPy arrays.
+    """
+    import torch
+
+    sd = torch.load(path, map_location="cpu", weights_only=True)
+    return {k: v.detach().float().cpu().numpy() for k, v in sd.items()}
+
+
+class MLXTAEHVDecoder:
+    """Minimal MLX port of TAEHV ``decoder`` (parallel-over-time MemBlocks)."""
+
+    def __init__(self, checkpoint_path: Path, *, z_dim: int) -> None:
+        """Initialize the TAEHV decoder from a checkpoint for the specified latent dimensionality.
+
+        Parameters:
+            checkpoint_path (Path): Path to the TAEHV checkpoint.
+            z_dim (int): Number of latent channels, determining the decoder patch size.
+        """
+        import mlx.core as mx
+
+        self.checkpoint_path = Path(checkpoint_path)
+        self.latent_channels = z_dim
+        # Derive patch_size from z_dim: 48 channels → patch_size=2, 16 → patch_size=1
+        self.patch_size = 2 if z_dim == 48 else 1
+        self.image_channels = 3
+        self.frames_to_trim = 3  # TGrow strides (1,2,2) → 2**2 - 1 for w2.1/w2.2 defaults
+
+        sd = _load_torch_state(self.checkpoint_path)
+        # Patch TGrow kernels like upstream TAEHV.patch_tgrow_layers.
+        self.weights = {k: mx.array(v) for k, v in sd.items()}
+        self._n_f = [256, 128, 64, 64]
+
+    def decode_ntchw(self, latents_ntchw: Any) -> Any:
+        """
+        Decode latent video batches into clipped RGB frames.
+
+        Parameters:
+            latents_ntchw (Any): Latents with shape ``[N, T, C, H, W]`` and the
+                decoder's configured latent channel count.
+
+        Returns:
+            Any: Decoded frames with shape ``[N, T_out, 3, H_out, W_out]`` and values
+                clipped to the range ``[0, 1]``.
+        """
+        import mlx.core as mx
+
+        x = latents_ntchw
+        n, t, c, h, w = x.shape
+        if c != self.latent_channels:
+            raise ValueError(f"expected C={self.latent_channels}, got {c}")
+        x = x.reshape(n * t, c, h, w)
+        x = self._run_decoder_parallel(x, n=n)
+        # Pixel-shuffle if patch_size > 1: (NT, 3*p*p, H, W) -> (NT, 3, H*p, W*p)
+        if self.patch_size > 1:
+            p = self.patch_size
+            nt, c_out, hh, ww = x.shape
+            x = x.reshape(nt, self.image_channels, p, p, hh, ww)
+            x = x.transpose(0, 1, 4, 2, 5, 3).reshape(nt, self.image_channels, hh * p, ww * p)
+        _, c_out, hh, ww = x.shape
+        t_out = x.shape[0] // n
+        x = x.reshape(n, t_out, c_out, hh, ww)
+        if self.frames_to_trim > 0 and t_out > self.frames_to_trim:
+            x = x[:, self.frames_to_trim:]
+        return mx.clip(x, 0.0, 1.0)
+
+    def _run_decoder_parallel(self, x: Any, *, n: int) -> Any:
+        """
+        Apply the TAEHV decoder stack to flattened batch and temporal frames while preserving temporal memory.
+        """
+        import mlx.core as mx
+
+        w = self.weights
+
+        def memblock(base: int, xx: Any, past: Any) -> Any:
+            """
+            Apply a temporal memory block to the current and past feature tensors.
+
+            Parameters:
+                base (int): Decoder block index used to select the block weights.
+                xx (Any): Current feature tensor.
+                past (Any): Past feature tensor concatenated with the current features.
+
+            Returns:
+                Any: Activated feature tensor produced by the memory block.
+            """
+            cat = mx.concatenate([xx, past], axis=1)
+            h = _mlx_conv2d(cat, w[f"decoder.{base}.conv.0.weight"], w.get(f"decoder.{base}.conv.0.bias"))
+            h = mx.maximum(h, 0.0)
+            h = _mlx_conv2d(h, w[f"decoder.{base}.conv.2.weight"], w.get(f"decoder.{base}.conv.2.bias"))
+            h = mx.maximum(h, 0.0)
+            h = _mlx_conv2d(h, w[f"decoder.{base}.conv.4.weight"], w.get(f"decoder.{base}.conv.4.bias"))
+            skip_key = f"decoder.{base}.skip.weight"
+            skip = _mlx_conv2d_1x1(xx, w[skip_key], None) if skip_key in w else xx
+            return mx.maximum(h + skip, 0.0)
+
+        def upsample2(xx: Any) -> Any:
+            """
+            Upsample a four-dimensional tensor by a factor of two along its spatial dimensions.
+
+            Parameters:
+                xx (Any): Tensor with shape `(N, C, H, W)`.
+
+            Returns:
+                Any: Tensor with shape `(N, C, 2H, 2W)` containing replicated spatial values.
+            """
+            nt, c, h, ww = xx.shape
+            xx = xx.reshape(nt, c, h, 1, ww, 1)
+            xx = mx.broadcast_to(xx, (nt, c, h, 2, ww, 2))
+            return xx.reshape(nt, c, h * 2, ww * 2)
+
+        def tgrow(base: int, xx: Any, stride: int) -> Any:
+            wt = w[f"decoder.{base}.conv.weight"]
+            out_ch = int(xx.shape[1]) * stride
+            if int(wt.shape[0]) > out_ch:
+                wt = wt[-out_ch:]
+            y = _mlx_conv2d_1x1(xx, wt, None)
+            if stride == 1:
+                return y
+            # TGrow.forward: (NT, C*stride, H, W) -> (NT*stride, C, H, W)
+            nt, c, h, ww = y.shape
+            c_in = c // stride
+            y = y.reshape(nt, stride, c_in, h, ww).transpose(0, 1, 2, 3, 4)
+            return y.reshape(nt * stride, c_in, h, ww)
+
+        def mem_past(xx: Any) -> Any:
+            """
+            Build a temporal memory tensor containing a zero frame followed by the preceding frame at each time step.
+
+            Parameters:
+                xx (Any): Flattened batch and temporal tensor with shape ``(batch * time, channels, height, width)``.
+
+            Returns:
+                Any: Tensor with the same shape as ``xx`` containing the preceding frame for each temporal position.
+            """
+            nt, c, h, ww = xx.shape
+            t_cur = nt // n
+            x_ = xx.reshape(n, t_cur, c, h, ww)
+            # pad one zero frame at t=0, align past[t] = x[t-1]
+            past = mx.concatenate([mx.zeros_like(x_[:, :1]), x_[:, :-1]], axis=1)
+            return past.reshape(nt, c, h, ww)
+
+        # 0 Clamp, 1 conv, 2 ReLU
+        x = mx.tanh(x / 3.0) * 3.0
+        x = _mlx_conv2d(x, w["decoder.1.weight"], w.get("decoder.1.bias"))
+        x = mx.maximum(x, 0.0)
+        for mem_idx in (3, 4, 5):
+            x = memblock(mem_idx, x, mem_past(x))
+        x = upsample2(x)
+        x = tgrow(7, x, 1)
+        x = _mlx_conv2d(x, w["decoder.8.weight"], w.get("decoder.8.bias"))
+        for mem_idx in (9, 10, 11):
+            x = memblock(mem_idx, x, mem_past(x))
+        x = upsample2(x)
+        x = tgrow(13, x, 2)
+        x = _mlx_conv2d(x, w["decoder.14.weight"], w.get("decoder.14.bias"))
+        for mem_idx in (15, 16, 17):
+            x = memblock(mem_idx, x, mem_past(x))
+        x = upsample2(x)
+        x = tgrow(19, x, 2)
+        x = _mlx_conv2d(x, w["decoder.20.weight"], w.get("decoder.20.bias"))
+        x = mx.maximum(x, 0.0)
+        x = _mlx_conv2d(x, w["decoder.22.weight"], w.get("decoder.22.bias"))
+        return x
+
+
+def decode_latents_taehv_mlx(
+    latents_np: np.ndarray,
+    *,
+    z_dim: int | None = None,
+    checkpoint_path: Path | None = None,
+) -> np.ndarray:
+    """
+    Decode latent representations with the MLX TAEHV decoder.
+
+    Parameters:
+        latents_np (np.ndarray): Latents arranged as [B, C, T, H, W].
+        z_dim (int | None): Latent channel dimension used to select the decoder checkpoint.
+        checkpoint_path (Path | None): Optional path to a TAEHV checkpoint.
+
+    Returns:
+        np.ndarray: Decoded pixels arranged as [B, T, H, W, 3] with values in [0, 1].
+
+    Raises:
+        ValueError: If `latents_np` does not have five dimensions.
+    """
+    import mlx.core as mx
+
+    if latents_np.ndim != 5:
+        raise ValueError(f"expected [B,C,T,H,W], got {latents_np.shape}")
+    c = latents_np.shape[1]
+    z = z_dim if z_dim is not None else c
+    ckpt = ensure_taehv_checkpoint(z_dim=z, checkpoint_path=checkpoint_path)
+    dec = MLXTAEHVDecoder(ckpt, z_dim=z)
+    # NTCHW
+    x = mx.array(latents_np.transpose(0, 2, 1, 3, 4).astype(np.float32))
+    out = dec.decode_ntchw(x)  # N T C H W
+    mx.eval(out)
+    arr = np.array(out)
+    # B T H W C
+    return arr.transpose(0, 1, 3, 4, 2)
+
+
+def decode_latents_wan_vae_torch(
+    latents_np: np.ndarray,
+    *,
+    vae_dir: Path,
+    device: str = "auto",
+    dtype_name: str = "fp16",
+) -> np.ndarray:
+    """Full AutoencoderKLWan decode on torch (MPS/CPU) with mean/std denormalize.
+
+    Returns pixels ``[B, T, H, W, 3]`` float in [0, 1].
+    """
+    import torch
+    from diffusers import AutoencoderKLWan
+    from diffusers.video_processor import VideoProcessor
+
+    if device == "auto":
+        device = "mps" if torch.backends.mps.is_available() else "cpu"
+    dtype = torch.float16 if dtype_name == "fp16" and device == "mps" else torch.float32
+    config = WanVAEConfigView.from_vae_dir(vae_dir)
+    vae = AutoencoderKLWan.from_pretrained(vae_dir, torch_dtype=dtype, local_files_only=True).to(device)
+    vae.eval()
+    latents = torch.from_numpy(latents_np.astype(np.float32)).to(device=device, dtype=dtype)
+    mean = torch.tensor(config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+    inv_std = (1.0 / torch.tensor(config.latents_std, device=device, dtype=dtype)).view(1, -1, 1, 1, 1)
+    latents = latents / inv_std + mean  # matches prompt_to_video path
+    with torch.no_grad():
+        video = vae.decode(latents, return_dict=False)[0]
+    video = VideoProcessor(vae_scale_factor=config.scale_factor_spatial).postprocess_video(video, output_type="np")
+    return video  # [B, T, H, W, 3]
+
+
+def decode_latents_to_video(
+    latents_np: np.ndarray,
+    output_path: Path,
+    *,
+    fps: int = 16,
+    backend: DecodeBackend = "taehv",
+    vae_dir: Path | None = None,
+    z_dim: int | None = None,
+    taehv_checkpoint: Path | None = None,
+    torch_device: str = "auto",
+) -> dict[str, Any]:
+    """Decode latent video frames and export them as an MP4 file.
+
+    Parameters:
+        latents_np (np.ndarray): Latent video representation to decode.
+        output_path (Path): Destination path for the MP4 file.
+        fps (int): Output video frame rate.
+        backend (DecodeBackend): Decoder backend to use.
+        vae_dir (Path | None): Directory containing the full Wan VAE when using
+            the ``wan-vae`` backend.
+        z_dim (int | None): Latent channel count for TAEHV decoding.
+        taehv_checkpoint (Path | None): Optional TAEHV checkpoint path.
+        torch_device (str): PyTorch device selection for PyTorch-based decoding.
+
+    Returns:
+        dict[str, Any]: Decode time in seconds, backend name, output path, frame
+        count, and video resolution.
+
+    Raises:
+        ValueError: If the full VAE backend lacks ``vae_dir`` or the backend is
+        unknown.
+    """
+    import time
+
+    from diffusers.utils import export_to_video
+
+    t0 = time.perf_counter()
+    if backend in ("taehv", "taehv-torch"):
+        c = latents_np.shape[1] if z_dim is None else z_dim
+        if backend == "taehv":
+            video = decode_latents_taehv_mlx(latents_np, z_dim=c, checkpoint_path=taehv_checkpoint)
+        else:
+            # torch TAEHV (regression / parity reference)
+            import torch
+            from fastvideo.third_party.taehv import TAEHV
+
+            ckpt = ensure_taehv_checkpoint(z_dim=c, checkpoint_path=taehv_checkpoint)
+            if torch_device == "auto":
+                torch_device = "mps" if torch.backends.mps.is_available() else "cpu"
+            dtype = torch.float16 if torch_device == "mps" else torch.float32
+            model = TAEHV(str(ckpt)).to(device=torch_device, dtype=dtype).eval()
+            lat = torch.from_numpy(latents_np).to(device=torch_device, dtype=dtype)
+            with torch.no_grad():
+                out = model.decode_video(lat.transpose(1, 2), parallel=True, show_progress_bar=False)
+            video = out[0].permute(0, 2, 3, 1).float().cpu().numpy()[None, ...]
+            # out is NTCHW -> need BTHWC; decode_video returns NTCHW for batch
+            if video.ndim == 4:
+                video = video[None]
+    elif backend == "wan-vae":
+        if vae_dir is None:
+            raise ValueError("vae_dir required for wan-vae backend")
+        video = decode_latents_wan_vae_torch(latents_np, vae_dir=vae_dir, device=torch_device)
+    else:
+        raise ValueError(f"unknown backend {backend}")
+
+    decode_s = time.perf_counter() - t0
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    # export_to_video expects list/array of frames HxWxC
+    frames = video[0]
+    frames = np.clip(frames, 0.0, 1.0)
+    export_to_video(frames, str(output_path), fps=fps)
+    if backend == "taehv":
+        cleanup_mlx()
+    else:
+        if backend == "taehv-torch":
+            del model, lat, out
+        cleanup_torch_mps()
+    return {
+        "decode_s": decode_s,
+        "backend": backend,
+        "output_path": str(output_path),
+        "num_frames": int(frames.shape[0]),
+        "resolution": f"{frames.shape[2]}x{frames.shape[1]}" if frames.ndim == 4 else None,
+    }

fastvideo/mlx_runtime/windowed_attention.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Chunked non-causal sliding-window self-attention for MLX scaling studies.
+
+This module is intentionally standalone (``mlx.core`` + stdlib only) so it can
+be micro-benchmarked without pulling in the DiT / FastVideo stack.
+
+Window policy
+-------------
+**Symmetric** sliding window (non-causal).  For query index ``i`` the allowed
+key indices are:
+
+    sinks:   ``j in [0, sink)``  (always visible to every query, if ``sink > 0``)
+    local:   ``j in [max(0, i - half), min(S, i + half + 1))``
+             where ``half = window // 2``
+
+so each query sees roughly ``window + 1`` local keys (plus any sinks outside
+that range).  This is appropriate for a dense, bidirectional DiT denoise pass.
+
+Implementation note (FLOPs)
+---------------------------
+A full-size additive attention mask still materialises an ``O(S^2)`` score
+matrix inside SDPA and does **not** reduce work.  Instead we tile the sequence
+into query blocks and run ``mx.fast.scaled_dot_product_attention`` only against
+the union of keys that block needs (local slice ± sinks).  That makes
+per-block work ``O(chunk * (window + sink) * D)`` and total work
+``O(S * (window + sink) * D)``.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+def _default_scale(head_dim: int, scale: float | None) -> float:
+    """
+    Determine the attention scaling factor from an explicit value or head dimension.
+
+    Parameters:
+        head_dim (int): The attention head dimension used to derive the default scale.
+        scale (Optional[float]): An explicit scaling factor.
+
+    Returns:
+        float: The explicit scale converted to a float, or the reciprocal square root of `head_dim`.
+
+    Raises:
+        ValueError: If `scale` is not provided and `head_dim` is not positive.
+    """
+    if scale is not None:
+        return float(scale)
+    if head_dim <= 0:
+        raise ValueError(f"head_dim must be positive, got {head_dim}")
+    return head_dim**-0.5
+
+
+def _validate_qkv(q: mx.array, k: mx.array, v: mx.array) -> tuple[int, int, int, int]:
+    """
+    Validate compatible rank-4 query, key, and value tensors.
+
+    Parameters:
+        q (mx.array): Query tensor shaped `(B, H, S, D)`.
+        k (mx.array): Key tensor with the same shape as `q`.
+        v (mx.array): Value tensor with the same shape as `q`.
+
+    Returns:
+        tuple[int, int, int, int]: Batch size, head count, sequence length, and head dimension.
+
+    Raises:
+        ValueError: If the tensors are not rank 4, do not have identical shapes, or have an empty sequence or head dimension.
+    """
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError(f"q/k/v must be rank-4 (B, H, S, D); got shapes "
+                         f"{q.shape}, {k.shape}, {v.shape}")
+    b, h, s, d = q.shape
+    if k.shape != (b, h, s, d) or v.shape != (b, h, s, d):
+        raise ValueError(f"q/k/v shapes must match exactly; got q={q.shape}, k={k.shape}, v={v.shape}")
+    if s == 0:
+        raise ValueError("sequence length S must be > 0")
+    if d == 0:
+        raise ValueError("head dim D must be > 0")
+    return b, h, s, d
+
+
+def full_attention(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    scale: float | None = None,
+) -> mx.array:
+    """
+    Compute dense scaled dot-product attention over the full sequence.
+
+    Parameters:
+        scale (float, optional): Attention scaling factor. If omitted, uses the
+            inverse square root of the head dimension.
+
+    Returns:
+        mx.array: Attention output with shape ``(B, H, S, D)``.
+    """
+    _, _, _, d = _validate_qkv(q, k, v)
+    sc = _default_scale(d, scale)
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=sc)
+
+
+def _concat_kv_slices(
+    k: mx.array,
+    v: mx.array,
+    ranges: list[tuple[int, int]],
+) -> tuple[mx.array, mx.array]:
+    """Concatenate non-overlapping ``[start, end)`` key/value slices along seq."""
+    if not ranges:
+        raise ValueError("ranges must be non-empty")
+    if len(ranges) == 1:
+        s0, e0 = ranges[0]
+        return k[:, :, s0:e0, :], v[:, :, s0:e0, :]
+    k_parts = [k[:, :, s:e, :] for s, e in ranges]
+    v_parts = [v[:, :, s:e, :] for s, e in ranges]
+    return mx.concatenate(k_parts, axis=2), mx.concatenate(v_parts, axis=2)
+
+
+def _key_ranges_for_block(
+    qs: int,
+    qe: int,
+    seq_len: int,
+    half: int,
+    sink: int,
+) -> list[tuple[int, int]]:
+    """Return ordered, non-overlapping key ranges for a query block ``[qs, qe)``.
+
+    Symmetric local window over every query in the block, plus global sinks
+    ``[0, sink)``.  Overlap is merged into a single contiguous range when
+    possible so we avoid double-counting sink tokens.
+    """
+    local_start = max(0, qs - half)
+    # Last query index is ``qe - 1``; its right edge is ``qe - 1 + half + 1 = qe + half``.
+    local_end = min(seq_len, qe + half)
+    if local_start >= local_end:
+        # Degenerate (should not happen for valid qs < qe); fall back to sinks only.
+        if sink > 0:
+            return [(0, min(sink, seq_len))]
+        raise ValueError(f"empty local key range for query block [{qs}, {qe})")
+
+    if sink <= 0:
+        return [(local_start, local_end)]
+
+    sink_end = min(sink, seq_len)
+    if local_start <= sink_end:
+        # Sinks abut or overlap the local window — one contiguous slice from 0.
+        return [(0, max(local_end, sink_end))]
+    # Gap between sinks and local window: two slices, concat at SDPA time.
+    return [(0, sink_end), (local_start, local_end)]
+
+
+def windowed_attention(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    window: int,
+    sink: int = 0,
+    scale: float | None = None,
+    *,
+    chunk_size: int | None = None,
+) -> mx.array:
+    """
+    Apply symmetric sliding-window self-attention with optional global sink positions.
+
+    Parameters:
+        q (mx.array): Query tensor shaped `(B, H, S, D)`.
+        k (mx.array): Key tensor shaped `(B, H, S, D)`.
+        v (mx.array): Value tensor shaped `(B, H, S, D)`.
+        window (int): Symmetric attention window width in tokens; must be at least 1.
+        sink (int): Number of leading key positions available to every query; must
+            be between 0 and the sequence length.
+        scale (Optional[float]): Softmax scale. Defaults to `1 / sqrt(D)`.
+        chunk_size (Optional[int]): Query block length used for chunked processing.
+            Defaults to the smaller of `window` and 512.
+
+    Returns:
+        mx.array: Attention output with the same shape as `q`.
+
+    Raises:
+        ValueError: If the inputs or attention parameters are invalid.
+        RuntimeError: If a query block has no available keys.
+    """
+    _, _, seq_len, d = _validate_qkv(q, k, v)
+
+    if window < 1:
+        raise ValueError(f"window must be >= 1, got {window}")
+    if sink < 0:
+        raise ValueError(f"sink must be >= 0, got {sink}")
+    if sink > seq_len:
+        raise ValueError(f"sink ({sink}) cannot exceed sequence length ({seq_len})")
+
+    sc = _default_scale(d, scale)
+    half = window // 2
+
+    # When the requested window is at least the sequence length, every query can
+    # see every key under a symmetric policy — fall back to one dense SDPA.
+    # (Sinks are redundant once the full key set is used.)
+    if window >= seq_len:
+        return mx.fast.scaled_dot_product_attention(q, k, v, scale=sc)
+
+    chunk = min(window, 512) if chunk_size is None else int(chunk_size)
+    if chunk < 1:
+        raise ValueError(f"chunk_size must be >= 1, got {chunk}")
+    chunk = min(chunk, seq_len)
+
+    outputs: list[mx.array] = []
+    for qs in range(0, seq_len, chunk):
+        qe = min(seq_len, qs + chunk)
+        q_block = q[:, :, qs:qe, :]
+        ranges = _key_ranges_for_block(qs, qe, seq_len, half, sink)
+        k_block, v_block = _concat_kv_slices(k, v, ranges)
+        if k_block.shape[2] == 0:
+            raise RuntimeError(f"empty key set for query block [{qs}, {qe}) with window={window}, sink={sink}")
+        query_positions = mx.arange(qs, qe)[:, None]
+        key_positions = mx.array([position for start, end in ranges for position in range(start, end)])[None, :]
+        mask = mx.abs(query_positions - key_positions) <= half
+        if sink > 0:
+            mask = mask | (key_positions < sink)
+        out_block = mx.fast.scaled_dot_product_attention(q_block, k_block, v_block, scale=sc, mask=mask)
+        outputs.append(out_block)
+
+    return mx.concatenate(outputs, axis=2)

fastvideo/platforms/mps.py

@@ -48,7 +48,10 @@ class MpsPlatform(Platform):
     @classmethod
     def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int,
                              dtype: torch.dtype) -> str:
-        # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible
+        if selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
+            raise NotImplementedError("VIDEO_SPARSE_ATTN is not supported on MPS. Unset "
+                                      "FASTVIDEO_ATTENTION_BACKEND or set it to TORCH_SDPA.")
+        # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible.
         logger.info("Using Torch SDPA backend for MPS.")
         return "fastvideo.attention.backends.sdpa.SDPABackend"
 

fastvideo/tests/attention/test_vsa_h3_backward.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GPU backward checks for the VSA-H3 backend.
+
+The CuTe backend returns FA4's own output tensor, which FA4's autograd node
+saved for its backward. Composing the compression branch onto it in place
+therefore poisons the graph, and the failure only appears once the VSA-256
+CuTe path has a backward at all. These tests pin the composition.
+"""
+
+import pytest
+import torch
+
+from fastvideo.attention.backends.video_sparse_attn_h3 import (MiniMaxH3VSAImpl, MiniMaxH3VSAMetadataBuilder)
+
+_SPEC = dict(raw_latent_shape=(16, 16, 24), patch_size=(1, 2, 2), prefix_segments=(64, 32, 16))
+_HEADS = 2
+_DIM = 128
+
+
+def _build_meta(device, sparsity=0.5):
+    return MiniMaxH3VSAMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=_SPEC["raw_latent_shape"],
+        patch_size=_SPEC["patch_size"],
+        VSA_sparsity=sparsity,
+        prefix_segments=_SPEC["prefix_segments"],
+        device=device,
+    )
+
+
+def _select_backend(monkeypatch, backend):
+    if backend == "cute":
+        pytest.importorskip(
+            "flash_attn.cute.block_sparsity",
+            reason="optional FA4 CuTe build (flash_attn.cute) not installed",
+        )
+        monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
+        monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
+        monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
+    else:
+        monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
+        monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False)
+
+
+def _forward_backward(impl, meta, gate_compress, device):
+    seq = meta.total_seq_length
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(1, seq, _HEADS, _DIM, device=device, dtype=torch.bfloat16, requires_grad=True)
+               for _ in range(3))
+    tq, tk, tv = (impl.tile(t, meta).clone() for t in (q, k, v))
+
+    gate = None
+    if gate_compress:
+        gate = torch.randn(1, tq.shape[1], _HEADS, _DIM, device=device, dtype=torch.bfloat16) * 0.1
+
+    out = impl.forward(tq, tk, tv, gate, meta)
+    out = impl.postprocess_output(out, meta)
+    out.float().pow(2).sum().backward()
+    return out, (q, k, v)
+
+
+@pytest.mark.parametrize("backend", ["triton", "cute"])
+@pytest.mark.parametrize("gate_compress", [False, True])
+def test_h3_vsa_backward_runs(monkeypatch, backend: str, gate_compress: bool) -> None:
+    """Regression: with the CuTe backend and a non-zero gate this used to die
+    with "one of the variables needed for gradient computation has been
+    modified by an inplace operation ... output 0 of FlashAttnFuncBackward".
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    _select_backend(monkeypatch, backend)
+
+    device = torch.device("cuda")
+    meta = _build_meta(device)
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+
+    out, leaves = _forward_backward(impl, meta, gate_compress, device)
+
+    assert torch.isfinite(out).all().item()
+    for name, leaf in zip(("q", "k", "v"), leaves):
+        assert leaf.grad is not None, f"{name} received no gradient"
+        assert torch.isfinite(leaf.grad).all().item(), f"{name}.grad has non-finite values"
+        assert leaf.grad.abs().sum().item() > 0, f"{name}.grad is all zero"
+
+
+@pytest.mark.parametrize("gate_compress", [False, True])
+def test_h3_vsa_backward_cute_matches_triton(monkeypatch, gate_compress: bool) -> None:
+    """CuTe and Triton take different routes to the same math; their gradients
+    should agree to bf16 tolerance."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+
+    device = torch.device("cuda")
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+
+    grads = {}
+    for backend in ("triton", "cute"):
+        with monkeypatch.context() as m:
+            _select_backend(m, backend)
+            meta = _build_meta(device)
+            _, leaves = _forward_backward(impl, meta, gate_compress, device)
+            grads[backend] = [leaf.grad.detach().float() for leaf in leaves]
+
+    for name, ref, got in zip(("dq", "dk", "dv"), grads["triton"], grads["cute"]):
+        diff = (ref - got).abs()
+        avg_abs = diff.mean().item()
+        max_rel = (diff.max() / (ref.abs().mean() + 1e-6)).item()
+        print(f"[h3-vsa gate={gate_compress}] {name}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+        assert avg_abs < 1e-2, f"{name}: avg_abs {avg_abs:.3e}"
+        assert max_rel < 0.5, f"{name}: max_rel {max_rel:.3e}"

fastvideo/tests/entrypoints/test_video_generator.py

@@ -2,8 +2,10 @@ import os
 from types import SimpleNamespace
 import warnings
 
+import numpy as np
 import pytest
 import torch
+from einops import rearrange
 
 import fastvideo.entrypoints.video_generator as video_generator_module
 from fastvideo.api import (
@@ -271,6 +273,70 @@ def test_generate_single_video_return_frames_still_materializes_output(tmp_path)
     assert result["video_path"] is None
 
 
+def test_generate_single_video_frames_match_legacy_cpu_loop(tmp_path):
+    """The on-device quantize path (#1362) must reproduce the legacy
+    per-frame CPU loop (make_grid -> permute -> *255 -> uint8) bit-exactly
+    for in-range fp32 pixels: same uint8 dtype, same HWC grid layout with
+    nrow=6 (batch>1), odd frame count. CPU-only: on CUDA the float->uint8
+    cast may differ by <=1 LSB, but on CPU both orderings run identical
+    fp32 ops, so exact equality is required."""
+    torch.manual_seed(0)
+    output = torch.rand((2, 3, 3, 16, 16), dtype=torch.float32)
+    output_batch = _single_video_output_batch(output)
+    fastvideo_args = _single_video_args()
+    generator = _single_video_generator(output_batch, fastvideo_args)
+    sampling_param = _small_sampling_param(save_video=False, return_frames=True)
+    sampling_param.num_frames = 3
+    sampling_param.num_videos_per_prompt = 2
+
+    result = generator._generate_single_video(
+        prompt="grid parity",
+        sampling_param=sampling_param,
+        fastvideo_args=fastvideo_args,
+        output_path=str(tmp_path / "unused.mp4"),
+    )
+
+    legacy_frames = []
+    for x in rearrange(output, "b c t h w -> t b c h w"):
+        grid = video_generator_module.torchvision.utils.make_grid(x, nrow=6)
+        grid = grid.permute(1, 2, 0).squeeze(-1)
+        legacy_frames.append((grid * 255).to(torch.uint8).contiguous().cpu().numpy())
+
+    torch.testing.assert_close(result["samples"], output)
+    assert len(result["frames"]) == 3
+    for got, want in zip(result["frames"], legacy_frames, strict=True):
+        assert got.dtype == np.uint8
+        assert got.shape == want.shape
+        np.testing.assert_array_equal(got, want)
+
+
+def test_generate_single_video_frames_clamp_out_of_range_pixels(tmp_path):
+    """VAE output slightly outside [0, 1] must saturate at 0/255 in the
+    uint8 frames. The pre-#1362 unclamped cast wrapped mod 256 (e.g.
+    1.5 -> 126). CPU-only."""
+    output = torch.full((1, 3, 2, 16, 16), 1.5, dtype=torch.float32)
+    output[:, :, 1] = -0.5
+    output_batch = _single_video_output_batch(output)
+    fastvideo_args = _single_video_args()
+    generator = _single_video_generator(output_batch, fastvideo_args)
+
+    result = generator._generate_single_video(
+        prompt="clamp",
+        sampling_param=_small_sampling_param(save_video=False, return_frames=True),
+        fastvideo_args=fastvideo_args,
+        output_path=str(tmp_path / "unused.mp4"),
+    )
+
+    frames = result["frames"]
+    assert len(frames) == 2
+    # make_grid passes a single image through without grid padding, so
+    # every pixel comes from the (clamped) output tensor.
+    assert frames[0].dtype == np.uint8
+    assert frames[0].shape == (16, 16, 3)
+    assert (frames[0] == 255).all()
+    assert (frames[1] == 0).all()
+
+
 def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_path):
     output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5
     output_batch = _single_video_output_batch(output)
@@ -305,6 +371,37 @@ def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_p
     }
 
 
+def test_generate_single_video_save_only_reports_refined_output_size(monkeypatch, tmp_path):
+    """`GenerationResult.size` must describe the decoded media even when the
+    fp32 `samples` mirror is skipped (`return_frames=False`, the CLI save
+    flow). Refiner pipelines can change the final pixel geometry, so the size
+    has to come from `output_batch.output`, not the base request. CPU-only."""
+    # Refiner-style output: request asks for 2 frames of 16x16, pipeline
+    # produces 5 frames of 32x48.
+    output = torch.full((1, 3, 5, 32, 48), 0.5, dtype=torch.float32)
+    output_batch = _single_video_output_batch(output)
+    fastvideo_args = _single_video_args()
+    generator = _single_video_generator(output_batch, fastvideo_args)
+    saved = {}
+
+    def fake_mimsave(path, frames, *, fps, format):
+        saved["frame_count"] = len(frames)
+
+    monkeypatch.setattr(video_generator_module.imageio, "mimsave", fake_mimsave)
+
+    result = generator._generate_single_video(
+        prompt="refined save",
+        sampling_param=_small_sampling_param(save_video=True, return_frames=False),
+        fastvideo_args=fastvideo_args,
+        output_path=str(tmp_path / "refined.mp4"),
+    )
+
+    assert result["samples"] is None
+    assert result["frames"] is None
+    assert result["size"] == (32, 48, 5)
+    assert saved["frame_count"] == 5
+
+
 def test_generate_single_video_audio_only_metadata_returns_audio_without_frames(tmp_path):
     audio = torch.zeros((16, ), dtype=torch.float32)
     output_batch = _single_video_output_batch(

fastvideo/tests/mlx/test_dmd_sampling.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only correctness contract for the on-device MLX DMD sampler.
+
+The sampler's arithmetic uses only ``+``/``-``/``*``, so NumPy arrays are a
+valid stand-in for MLX arrays here: these tests validate the *math* against the
+torch reference formulas without requiring an MLX/Metal device.
+"""
+
+import numpy as np
+
+from fastvideo.mlx_runtime.sampling import (
+    MLXDMDSchedule,
+    add_noise,
+    dmd_step,
+    pred_noise_to_pred_video,
+)
+
+
+def _reference_schedule() -> MLXDMDSchedule:
+    # A small monotonic flow-match-like schedule: timesteps 0..999, sigmas in
+    # (0, 1]. Exact values are irrelevant; the lookup must pick the nearest.
+    """Create the deterministic schedule used by the reference tests.
+
+    Returns:
+        MLXDMDSchedule: A schedule with timesteps from 0 through 999 and linearly decreasing sigma values from 1.0 to 0.001.
+    """
+    timesteps = np.arange(1000, dtype=np.float64)
+    sigmas = np.linspace(1.0, 1e-3, 1000, dtype=np.float64)
+    return MLXDMDSchedule(sigmas=sigmas, timesteps=timesteps)
+
+
+def test_sigma_for_picks_nearest_timestep() -> None:
+    schedule = _reference_schedule()
+    # Exact hits.
+    assert schedule.sigma_for(0) == schedule.sigmas[0]
+    assert schedule.sigma_for(999) == schedule.sigmas[999]
+    # Nearest-neighbour rounding for a value between grid points.
+    assert schedule.sigma_for(522.4) == schedule.sigmas[522]
+    assert schedule.sigma_for(521.6) == schedule.sigmas[522]
+
+
+def test_pred_noise_to_pred_video_matches_reference_formula() -> None:
+    rng = np.random.default_rng(0)
+    pred_noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    noise_input = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    sigma = 0.37
+
+    got = pred_noise_to_pred_video(pred_noise, noise_input, sigma)
+    expected = noise_input - sigma * pred_noise
+    np.testing.assert_allclose(got, expected, rtol=0, atol=0)
+
+
+def test_add_noise_matches_reference_formula() -> None:
+    rng = np.random.default_rng(1)
+    clean = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    sigma = 0.62
+
+    got = add_noise(clean, noise, sigma)
+    expected = (1.0 - sigma) * clean + sigma * noise
+    np.testing.assert_allclose(got, expected, rtol=0, atol=0)
+
+
+def test_dmd_step_final_returns_clean_prediction() -> None:
+    schedule = _reference_schedule()
+    rng = np.random.default_rng(2)
+    pred_noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    noise_input = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+
+    out = dmd_step(
+        latents=noise_input,
+        noise_input_latent=noise_input,
+        pred_noise=pred_noise,
+        schedule=schedule,
+        timestep=522,
+        next_timestep=None,
+    )
+    sigma = schedule.sigma_for(522)
+    np.testing.assert_allclose(out, noise_input - sigma * pred_noise, rtol=0, atol=0)
+
+
+def test_dmd_step_intermediate_renoises_to_next_level() -> None:
+    """
+    Verifies that an intermediate DMD step re-noises the denoised prediction at the next schedule level.
+    """
+    schedule = _reference_schedule()
+    rng = np.random.default_rng(3)
+    pred_noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    noise_input = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+    renoise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32)
+
+    out = dmd_step(
+        latents=noise_input,
+        noise_input_latent=noise_input,
+        pred_noise=pred_noise,
+        schedule=schedule,
+        timestep=1000,
+        next_timestep=757,
+        noise=renoise,
+    )
+    sigma = schedule.sigma_for(1000)
+    sigma_next = schedule.sigma_for(757)
+    pred_video = noise_input - sigma * pred_noise
+    expected = (1.0 - sigma_next) * pred_video + sigma_next * renoise
+    np.testing.assert_allclose(out, expected, rtol=0, atol=0)
+
+
+def test_dmd_step_requires_noise_for_intermediate_step() -> None:
+    """Verify that an intermediate DMD step requires a noise tensor for re-noising."""
+    schedule = _reference_schedule()
+    zeros = np.zeros((1, 16, 3, 8, 8), dtype=np.float32)
+    try:
+        dmd_step(
+            latents=zeros,
+            noise_input_latent=zeros,
+            pred_noise=zeros,
+            schedule=schedule,
+            timestep=1000,
+            next_timestep=757,
+            noise=None,
+        )
+    except ValueError:
+        return
+    raise AssertionError("dmd_step must raise when re-noise is requested without noise")

fastvideo/tests/mlx/test_frame_upsample.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only contracts for pixel-space frame resampling."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from fastvideo.mlx_runtime.frame_upsample import (
+    PIXEL_UPSAMPLE_MODES,
+    unsharp,
+    upsample_frame,
+    upsample_frames,
+)
+from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
+
+
+def _frame(height: int = 24, width: int = 40, seed: int = 0) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    return rng.integers(0, 256, (height, width, 3), dtype=np.uint8)
+
+
+@pytest.mark.parametrize("mode", PIXEL_UPSAMPLE_MODES)
+def test_upsample_frame_hits_target_size_in_every_mode(mode: str) -> None:
+    out = upsample_frame(_frame(), width=80, height=48, mode=mode)
+    assert out.shape == (48, 80, 3)
+    assert out.dtype == np.uint8
+
+
+def test_upsample_frame_rejects_unknown_mode() -> None:
+    with pytest.raises(ValueError, match="Unsupported pixel upsample mode"):
+        upsample_frame(_frame(), width=80, height=48, mode="mitchell")
+
+
+def test_upsample_frame_rejects_bad_shape() -> None:
+    with pytest.raises(ValueError, match="HxWx3"):
+        upsample_frame(np.zeros((24, 40), dtype=np.uint8), width=80, height=48)
+
+
+def test_upsample_frame_rejects_nonpositive_size() -> None:
+    with pytest.raises(ValueError, match="must be positive"):
+        upsample_frame(_frame(), width=0, height=48)
+
+
+def test_upsample_frame_at_target_size_is_a_passthrough() -> None:
+    frame = _frame()
+    np.testing.assert_array_equal(upsample_frame(frame, width=40, height=24), frame)
+
+
+def test_nearest_upsample_is_exact_block_replication() -> None:
+    """A 2x nearest resize must duplicate pixels, not blend them."""
+    frame = _frame(height=4, width=6)
+    out = upsample_frame(frame, width=12, height=8, mode="nearest")
+    np.testing.assert_array_equal(out, np.repeat(np.repeat(frame, 2, axis=0), 2, axis=1))
+
+
+def test_unsharp_zero_amount_is_identity() -> None:
+    frame = _frame()
+    assert unsharp(frame, 0.0) is frame
+
+
+def test_unsharp_steepens_a_soft_edge() -> None:
+    # A blurred step edge: unsharp masking should raise its peak gradient.
+    import cv2
+
+    frame = np.zeros((24, 40, 3), dtype=np.uint8)
+    frame[:, 20:] = 200
+    frame = cv2.GaussianBlur(frame, (0, 0), 2.0)
+
+    def peak_gradient(image: np.ndarray) -> float:
+        return float(np.abs(np.diff(image[:, :, 0].astype(np.int32), axis=1)).max())
+
+    assert peak_gradient(unsharp(frame, 0.8)) > peak_gradient(frame)
+
+
+def test_upsample_frames_preserves_order_and_leaves_input_alone() -> None:
+    frames = [_frame(seed=index) for index in range(3)]
+    originals = [frame.copy() for frame in frames]
+    out = upsample_frames(frames, width=80, height=48)
+
+    assert len(out) == 3
+    assert all(frame.shape == (48, 80, 3) for frame in out)
+    # Distinct inputs stay distinct and in order.
+    assert not np.array_equal(out[0], out[1])
+    for frame, original in zip(frames, originals, strict=True):
+        np.testing.assert_array_equal(frame, original)
+
+
+@pytest.mark.parametrize(
+    ("target_frames", "factor", "keyframes"),
+    [(1, 999, 1), (81, 2, 41), (81, 3, 29), (82, 2, 45), (121, 4, 33)],
+)
+def test_aligned_keyframes_expand_to_exact_target(target_frames: int, factor: int, keyframes: int) -> None:
+    assert aligned_keyframe_count(target_frames, factor) == keyframes
+    interpolated_frames = (keyframes - 1) * factor + 1
+    assert interpolated_frames >= target_frames
+    assert keyframes % 4 == 1
+    if keyframes > 1:
+        assert (keyframes - 5) * factor + 1 < target_frames

fastvideo/tests/mlx/test_import_laziness.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Import hygiene for the lightweight Apple Silicon MLX test path."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+_HEAVY = "fastvideo.entrypoints.video_generator"
+
+
+def _modules_after(import_line: str) -> set[str]:
+    """Return sys.modules after running one import in a fresh interpreter."""
+    code = f"import sys; {import_line}; print('\\n'.join(sorted(sys.modules)))"
+    result = subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True)
+    return set(result.stdout.split())
+
+
+def test_mlx_runtime_import_adds_no_heavy_entrypoints() -> None:
+    """Importing the MLX runtime must not drag in the generator stack itself.
+
+    ``fastvideo/__init__.py`` imports ``VideoGenerator`` eagerly, so every
+    ``fastvideo.*`` import loads it transitively. That is a property of the
+    top-level package rather than of this runtime, so compare against a plain
+    ``import fastvideo`` baseline and assert the MLX runtime adds nothing heavy
+    on top of it.
+    """
+    baseline = _modules_after("import fastvideo")
+    with_runtime = _modules_after("import fastvideo.mlx_runtime.memory")
+    assert _HEAVY not in with_runtime - baseline

fastvideo/tests/mlx/test_memory_limits.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+
+from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits, gib_to_bytes
+
+
+class _FakeMLX:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, int]] = []
+
+    def set_memory_limit(self, value: int) -> int:
+        """
+        Record a requested memory limit and provide the prior limit.
+
+        Parameters:
+            value (int): The memory limit to record.
+
+        Returns:
+            int: The previously configured memory limit.
+        """
+        self.calls.append(("memory", value))
+        return 111
+
+    def set_cache_limit(self, value: int) -> int:
+        """Record and apply a cache memory limit.
+
+        Parameters:
+            value (int): Cache memory limit in bytes.
+
+        Returns:
+            int: The previously configured cache memory limit.
+        """
+        self.calls.append(("cache", value))
+        return 222
+
+    def set_wired_limit(self, value: int) -> int:
+        self.calls.append(("wired", value))
+        return 333
+
+
+def test_gib_to_bytes_rejects_non_positive_values() -> None:
+    assert gib_to_bytes(None) is None
+    assert gib_to_bytes(1.5) == int(1.5 * 1024**3)
+    with pytest.raises(ValueError, match="positive"):
+        gib_to_bytes(0)
+
+
+def test_apply_memory_limits_sets_mlx_limits_and_metrics(monkeypatch) -> None:
+    monkeypatch.delenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", raising=False)
+    monkeypatch.delenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", raising=False)
+    fake_mlx = _FakeMLX()
+
+    applied = apply_memory_limits(
+        mlx_memory_limit_gib=16,
+        mlx_disable_cache=True,
+        mlx_wired_limit_gib=12,
+        torch_mps_high_watermark_ratio=0.57,
+        mx_module=fake_mlx,
+    )
+
+    assert os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] == "0.57"
+    assert os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] == "0.0"
+    assert fake_mlx.calls == [
+        ("memory", 16 * 1024**3),
+        ("cache", 0),
+        ("wired", 12 * 1024**3),
+    ]
+    metrics = applied.as_metrics()
+    assert metrics["mlx_memory_limit_bytes"] == 16 * 1024**3
+    assert metrics["mlx_cache_limit_bytes"] == 0
+    assert metrics["previous_mlx_memory_limit_bytes"] == 111
+    assert metrics["previous_mlx_cache_limit_bytes"] == 222
+    assert metrics["previous_mlx_wired_limit_bytes"] == 333
+
+
+def test_apply_memory_limits_validates_watermarks() -> None:
+    with pytest.raises(ValueError, match="non-negative"):
+        apply_memory_limits(torch_mps_high_watermark_ratio=-1)
+
+
+def test_add_memory_limit_args_uses_defaults() -> None:
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    add_memory_limit_args(
+        parser,
+        mlx_memory_limit_gib=16,
+        mlx_disable_cache=True,
+        torch_mps_high_watermark_ratio=0.57,
+        torch_mps_low_watermark_ratio=0.0,
+    )
+
+    args = parser.parse_args([])
+    assert args.mlx_memory_limit_gib == 16
+    assert args.mlx_disable_cache is True
+    assert args.torch_mps_high_watermark_ratio == 0.57
+    assert args.torch_mps_low_watermark_ratio == 0.0

fastvideo/tests/mlx/test_mlx_checkpoint.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Round-trip tests for pre-quantized MLX DiT checkpoints.
+
+Saving an already-quantized DiT and reloading it must reproduce the original
+forward pass exactly (same packed weights, scales, and config), for both plain
+fp32 and int8-quantized models. Runs on any MLX backend (Metal or CPU).
+"""
+
+from __future__ import annotations
+
+import json
+
+import numpy as np
+import pytest
+
+pytest.importorskip("mlx.core", reason="MLX is required for checkpoint tests")
+
+from fastvideo.mlx_runtime.checkpoint import (  # noqa: E402
+    MANIFEST_FILENAME,
+    load_mlx_dit_checkpoint,
+    save_mlx_dit_checkpoint,
+)
+from fastvideo.mlx_runtime.fastwan import (  # noqa: E402
+    MLXQuantizationSpec,
+    QuantizedMatrix,
+)
+from fastvideo.tests.mlx.tiny_wan import (  # noqa: E402
+    build_hf_config,
+    build_inputs,
+    build_tiny_wan_config,
+    build_torch_model,
+    mlx_dit_from_torch_model,
+    mlx_output,
+    mlx_rotary_embeddings,
+)
+
+
+def _forward(dit) -> np.ndarray:
+    """Run the DiT model with standard inputs and return its output as a NumPy array."""
+    hidden_states, encoder_hidden_states, timestep = build_inputs()
+    return mlx_output(dit, hidden_states, encoder_hidden_states, timestep, mlx_rotary_embeddings(hidden_states))
+
+
+def test_plain_checkpoint_round_trip(distributed_setup, tmp_path) -> None:
+    dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
+    before = _forward(dit)
+
+    save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
+    loaded = load_mlx_dit_checkpoint(tmp_path / "ckpt")
+
+    # JSON canonicalizes tuples to lists (matching what the Diffusers
+    # config.json loader produces), so compare JSON-normalized configs.
+    assert json.loads(json.dumps(loaded.config)) == json.loads(json.dumps(dit.config))
+    assert len(loaded.blocks) == len(dit.blocks)
+    after = _forward(loaded)
+    np.testing.assert_array_equal(before, after)
+
+
+def test_int8_checkpoint_round_trip_skips_requantization(distributed_setup, tmp_path) -> None:
+    spec = MLXQuantizationSpec.from_name("int8")
+    dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()),
+                                   quantization=spec)
+    before = _forward(dit)
+
+    save_mlx_dit_checkpoint(dit, tmp_path / "ckpt-int8")
+    loaded = load_mlx_dit_checkpoint(tmp_path / "ckpt-int8")
+
+    # The loaded model must carry the same packed quantized weights -- no
+    # requantization pass -- and produce a bitwise-identical forward.
+    reloaded_q = loaded.weights["proj_out.weight"]
+    original_q = dit.weights["proj_out.weight"]
+    assert isinstance(reloaded_q, QuantizedMatrix)
+    assert reloaded_q.spec == spec
+    np.testing.assert_array_equal(np.array(reloaded_q.weight), np.array(original_q.weight))
+    np.testing.assert_array_equal(np.array(reloaded_q.scales), np.array(original_q.scales))
+
+    block_q = loaded.blocks[0].weights["to_q.weight"]
+    assert isinstance(block_q, QuantizedMatrix)
+
+    after = _forward(loaded)
+    np.testing.assert_array_equal(before, after)
+
+
+def test_missing_checkpoint_raises_clear_error(tmp_path) -> None:
+    with pytest.raises(FileNotFoundError, match="Not an MLX DiT checkpoint"):
+        load_mlx_dit_checkpoint(tmp_path / "does-not-exist")
+
+
+def test_existing_checkpoint_is_replaced(distributed_setup, tmp_path) -> None:
+    dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
+    checkpoint_dir = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
+    dit.config["checkpoint_revision"] = "second"
+
+    save_mlx_dit_checkpoint(dit, checkpoint_dir)
+
+    assert load_mlx_dit_checkpoint(checkpoint_dir).config["checkpoint_revision"] == "second"
+
+
+def test_failed_checkpoint_overwrite_preserves_existing(distributed_setup, tmp_path) -> None:
+    dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
+    checkpoint_dir = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
+    weights_before = (checkpoint_dir / "mlx_dit.safetensors").read_bytes()
+    manifest_before = (checkpoint_dir / MANIFEST_FILENAME).read_bytes()
+    dit.config["not_json_serializable"] = object()
+
+    with pytest.raises(TypeError, match="not JSON serializable"):
+        save_mlx_dit_checkpoint(dit, checkpoint_dir)
+
+    assert (checkpoint_dir / "mlx_dit.safetensors").read_bytes() == weights_before
+    assert (checkpoint_dir / MANIFEST_FILENAME).read_bytes() == manifest_before
+
+
+def test_future_format_version_is_rejected(distributed_setup, tmp_path) -> None:
+    dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
+    ckpt = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
+
+    manifest_path = ckpt / MANIFEST_FILENAME
+    manifest = json.loads(manifest_path.read_text())
+    manifest["format_version"] = 99
+    manifest_path.write_text(json.dumps(manifest))
+
+    with pytest.raises(ValueError, match="format_version=99"):
+        load_mlx_dit_checkpoint(ckpt)

fastvideo/tests/mlx/test_mlx_compile_parity.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Guard the ``mx.compile`` path of the MLX FastWan forward.
+
+`mx.compile` on the DiT forward was silently broken: a NumPy scalar multiplying
+a traced array (``np.sqrt(2/pi) * x`` in ``gelu_tanh``) dispatched through NumPy,
+which evals the traced array — illegal during compile. It raised "Attempting to
+eval an array during function transformations" (caught → eager fallback, no
+speedup) or segfaulted the process. These tests keep the compiled forward tracing
+cleanly and numerically identical to eager, so the ~1.4x compile speedup cannot
+silently regress. See ``fastvideo/mlx_runtime/fastwan.py`` and
+``docs/design/apple_silicon_benchmark_baseline.md``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+mx = pytest.importorskip("mlx.core", reason="MLX is required for the compile-parity tests")
+
+from fastvideo.mlx_runtime.fastwan import (  # noqa: E402
+    MLXWanTransformerBlock,
+    gelu_tanh,
+    timestep_embedding,
+)
+
+# One module-level generator so successive _rand calls of the same shape get
+# distinct values (re-seeding per call made every (dim, dim) weight identical).
+_RNG = np.random.default_rng(0)
+
+# Through MLX 0.31 the compiler emitted the same kernels as eager for these
+# graphs, so bit-identity was the contract. MLX 0.32 fuses more aggressively and
+# reassociates float ops, which moves a handful of elements by a few ULP (max
+# observed: 2.4e-07 absolute, 6.4e-06 relative, on 41 of 7680 elements). That is
+# reassociation, not a regression, so assert bit-identity where it is guaranteed
+# and a tight tolerance otherwise. Either way a silent fallback to eager or a
+# broken trace still fails the test.
+_MLX_VERSION = tuple(int(part) for part in mx.__version__.split(".")[:2])
+_COMPILE_IS_BITWISE = _MLX_VERSION < (0, 32)
+
+
+def _assert_compile_parity(eager: "mx.array", compiled: "mx.array") -> None:
+    """Compare an eager result against its compiled counterpart."""
+    if _COMPILE_IS_BITWISE:
+        np.testing.assert_array_equal(np.array(eager), np.array(compiled))
+    else:
+        np.testing.assert_allclose(np.array(eager), np.array(compiled), rtol=1e-5, atol=1e-6)
+
+
+def _rand(*shape: int, scale: float = 1.0) -> "mx.array":
+    """
+    Generate a float32 MLX array of normally distributed random values.
+
+    Parameters:
+        shape (int): Dimensions of the generated array.
+        scale (float): Factor applied to the sampled values.
+
+    Returns:
+        mx.array: The generated random array.
+    """
+    return mx.array((_RNG.standard_normal(shape) * scale).astype(np.float32))
+
+
+def test_gelu_tanh_compiles_and_matches_eager() -> None:
+    """The tanh-GELU (the op that broke compile) traces and matches eager."""
+    x = _rand(1, 120, 64)
+    eager = gelu_tanh(x)
+    mx.eval(eager)
+    compiled = mx.compile(gelu_tanh)(x)
+    mx.eval(compiled)
+    _assert_compile_parity(eager, compiled)
+
+
+def test_timestep_embedding_compiles_and_matches_eager() -> None:
+    """``timestep_embedding`` (math.log fix sibling of gelu_tanh) compiles cleanly.
+
+    Unlike gelu_tanh, the exp/sin/cos chain can reassociate under compile, so
+    we allow a tight tolerance rather than bit-identity.
+    """
+    t = mx.array(np.array([0.0, 250.0, 500.0, 1000.0], dtype=np.float32))
+    dim = 64
+    eager = timestep_embedding(t, dim)
+    mx.eval(eager)
+    compiled = mx.compile(lambda steps: timestep_embedding(steps, dim))(t)
+    mx.eval(compiled)
+    np.testing.assert_allclose(np.array(eager), np.array(compiled), rtol=1e-4, atol=1e-4)
+
+
+def _tiny_block_weights(dim: int, ffn: int) -> dict:
+    """
+    Create randomized weights for a small transformer block test fixture.
+
+    Parameters:
+        dim (int): Hidden dimension of the transformer block.
+        ffn (int): Intermediate dimension of the feed-forward network.
+
+    Returns:
+        dict: Randomized attention, normalization, modulation, and feed-forward weights and biases.
+    """
+    square = ["to_q", "to_k", "to_v", "to_out", "attn2.to_q", "attn2.to_k", "attn2.to_v", "attn2.to_out"]
+    weights = {f"{k}.weight": _rand(dim, dim, scale=0.05) for k in square}
+    weights.update({f"{k}.bias": _rand(dim, scale=0.05) for k in ["to_q", "to_k", "to_v", "to_out"]})
+    weights.update({
+        "scale_shift_table": _rand(1, 6, dim, scale=0.05),
+        "norm_q.weight": _rand(dim, scale=0.05),
+        "norm_k.weight": _rand(dim, scale=0.05),
+        "self_attn_residual_norm.norm.weight": _rand(dim, scale=0.05),
+        "self_attn_residual_norm.norm.bias": _rand(dim, scale=0.05),
+        "attn2.norm_q.weight": _rand(dim, scale=0.05),
+        "attn2.norm_k.weight": _rand(dim, scale=0.05),
+        "ffn.fc_in.weight": _rand(ffn, dim, scale=0.05),
+        "ffn.fc_in.bias": _rand(ffn, scale=0.05),
+        "ffn.fc_out.weight": _rand(dim, ffn, scale=0.05),
+        "ffn.fc_out.bias": _rand(dim, scale=0.05),
+    })
+    return weights
+
+
+def test_transformer_block_compiles_and_matches_eager() -> None:
+    """The full dense block (the mx.compile target's body) traces and matches."""
+    dim, num_heads, head_dim, ffn, seq, ctx = 64, 4, 16, 128, 120, 32
+    block = MLXWanTransformerBlock(_tiny_block_weights(dim, ffn), dim=dim, ffn_dim=ffn, num_heads=num_heads, eps=1e-6)
+    hidden = _rand(1, seq, dim, scale=0.05)
+    context = _rand(1, ctx, dim, scale=0.05)
+    temb = _rand(1, 6, dim, scale=0.05)
+    cos = _rand(seq, head_dim)
+    sin = _rand(seq, head_dim)
+
+    eager = block(hidden, context, temb, (cos, sin))
+    mx.eval(eager)
+    compiled = mx.compile(lambda h, e, t, c, s: block(h, e, t, (c, s)))(hidden, context, temb, cos, sin)
+    mx.eval(compiled)
+
+    _assert_compile_parity(eager, compiled)

fastvideo/tests/mlx/test_mlx_diffusers_compile_dispatch.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for preserving MLX compile flags on raw Diffusers loads."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+mx = pytest.importorskip("mlx.core", reason="MLX required for compile dispatch")
+
+
+def _write_tiny_diffusers_checkpoint(tmp_path: Path) -> tuple[Path, Path]:
+    import torch
+    from safetensors.torch import save_file
+
+    config = {
+        "num_layers": 0,
+        "num_attention_heads": 1,
+        "attention_head_dim": 2,
+        "ffn_dim": 4,
+        "in_channels": 16,
+        "out_channels": 16,
+        "patch_size": [1, 2, 2],
+        "freq_dim": 2,
+        "eps": 1e-6,
+    }
+    config_path = tmp_path / "config.json"
+    checkpoint_path = tmp_path / "diffusion_pytorch_model.safetensors"
+    config_path.write_text(json.dumps(config))
+
+    hidden = 2
+    weights = {
+        "patch_embedding.weight": torch.zeros(hidden, 64),
+        "condition_embedder.time_embedder.linear_1.weight": torch.zeros(hidden, 2),
+        "condition_embedder.time_embedder.linear_2.weight": torch.zeros(hidden, hidden),
+        "condition_embedder.time_proj.weight": torch.zeros(hidden * 6, hidden),
+        "condition_embedder.text_embedder.linear_1.weight": torch.zeros(hidden, 4096),
+        "condition_embedder.text_embedder.linear_2.weight": torch.zeros(hidden, hidden),
+        "scale_shift_table": torch.zeros(2, hidden),
+        "proj_out.weight": torch.zeros(64, hidden),
+    }
+    save_file(weights, checkpoint_path)
+    return checkpoint_path, config_path
+
+
+def test_diffusers_loader_preserves_compile_flag(tmp_path: Path) -> None:
+    from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
+
+    checkpoint_path, config_path = _write_tiny_diffusers_checkpoint(tmp_path)
+
+    compiled = mlx_dit_from_diffusers_safetensors(
+        checkpoint_path,
+        config_path,
+        compile=True,
+        num_blocks=0,
+        quantization=None,
+    )
+    eager = mlx_dit_from_diffusers_safetensors(
+        checkpoint_path,
+        config_path,
+        compile=False,
+        num_blocks=0,
+        quantization=None,
+    )
+
+    assert compiled._enable_compile is True
+    assert eager._enable_compile is False
+
+
+def test_prompt_to_video_cli_passes_compile_to_diffusers_loader() -> None:
+    root = Path(__file__).resolve().parents[3]
+    script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py"
+    source = script.read_text()
+
+    assert "compile=args.mlx_compile" in source

fastvideo/tests/mlx/test_mlx_dit_parity.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Full-DiT parity: the MLX Wan runtime vs the PyTorch reference model.
+
+This is the M1 "trustworthy baseline" gate for the Apple Silicon path: a tiny
+random-weight ``WanTransformer3DModel`` is run end to end (patch embed ->
+condition -> transformer blocks -> unpatchify) in PyTorch and in
+``fastvideo.mlx_runtime.fastwan.MLXWanDiT`` with identical weights, and the
+outputs must match within pinned fp32 tolerances.
+
+Runs anywhere MLX is installed: on Apple Silicon it exercises the Metal
+device, on Linux/CI it runs on MLX's CPU backend (``pip install 'mlx[cpu]'``)
+-- the graph is identical, so CPU parity is the CI-friendly golden variant.
+
+    pytest fastvideo/tests/mlx/test_mlx_dit_parity.py -v
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+pytest.importorskip("mlx.core", reason="MLX is required for DiT parity tests")
+
+from fastvideo.mlx_runtime.fastwan import MLXQuantizationSpec  # noqa: E402
+from fastvideo.tests.mlx.tiny_wan import (  # noqa: E402
+    build_hf_config,
+    build_inputs,
+    build_tiny_wan_config,
+    build_torch_model,
+    mlx_dit_from_torch_model,
+    mlx_output,
+    mlx_rotary_embeddings,
+    torch_reference_output,
+)
+
+# Pinned fp32 tolerances for the full forward (patch embed through unpatchify),
+# matching the single-block parity example. Measured max_abs_diff on the MLX
+# CPU backend is ~1.7e-6, so 2e-4 keeps ~100x headroom for accumulation-order
+# differences across MLX backends while still catching real math/layout bugs.
+FP32_ATOL = 2e-4
+FP32_RTOL = 2e-4
+
+# Quantized inference is lossy by design; gate it on signal-to-noise vs the
+# fp32 MLX output instead of elementwise closeness. int8 (group size 64) on
+# this tiny model measures ~43 dB; 20 dB leaves headroom without letting a
+# broken dequant path (which lands near 0 dB) slip through.
+INT8_MIN_SNR_DB = 20.0
+
+
+def test_full_dit_forward_matches_torch_reference(distributed_setup) -> None:
+    model = build_torch_model()
+    hidden_states, encoder_hidden_states, timestep = build_inputs()
+
+    torch_out = torch_reference_output(model, hidden_states, encoder_hidden_states, timestep)
+
+    dit = mlx_dit_from_torch_model(model, build_hf_config(build_tiny_wan_config()))
+    mlx_out = mlx_output(dit, hidden_states, encoder_hidden_states, timestep, mlx_rotary_embeddings(hidden_states))
+
+    assert mlx_out.shape == torch_out.shape
+    assert np.isfinite(mlx_out).all()
+    max_abs = float(np.abs(torch_out - mlx_out).max())
+    assert np.allclose(torch_out, mlx_out, atol=FP32_ATOL, rtol=FP32_RTOL), (
+        f"MLX full-DiT forward diverged from the torch reference: max_abs_diff={max_abs:.3e} "
+        f"(atol={FP32_ATOL}, rtol={FP32_RTOL})")
+
+
+def test_full_dit_forward_int8_stays_close_to_fp32(distributed_setup) -> None:
+    """
+    Ensure int8 quantization preserves the MLX DiT output within the required signal-to-noise ratio.
+
+    Parameters:
+        distributed_setup: Test fixture providing the distributed test environment.
+
+    Raises:
+        AssertionError: If the int8 output contains non-finite values or its signal-to-noise ratio is below the configured minimum.
+    """
+    model = build_torch_model()
+    hidden_states, encoder_hidden_states, timestep = build_inputs()
+    hf_config = build_hf_config(build_tiny_wan_config())
+    freqs_cis = mlx_rotary_embeddings(hidden_states)
+
+    fp32_out = mlx_output(
+        mlx_dit_from_torch_model(model, hf_config), hidden_states, encoder_hidden_states, timestep, freqs_cis)
+    int8_out = mlx_output(
+        mlx_dit_from_torch_model(model, hf_config, quantization=MLXQuantizationSpec.from_name("int8")),
+        hidden_states, encoder_hidden_states, timestep, freqs_cis)
+
+    assert np.isfinite(int8_out).all()
+    noise = float(np.mean(np.square(int8_out - fp32_out)))
+    signal = float(np.mean(np.square(fp32_out)))
+    snr_db = 10.0 * np.log10(signal / noise) if noise > 0 else float("inf")
+    assert snr_db >= INT8_MIN_SNR_DB, (
+        f"int8-quantized MLX DiT output is too far from fp32: SNR {snr_db:.1f} dB < {INT8_MIN_SNR_DB} dB")

fastvideo/tests/mlx/test_mlx_fast_spatial.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only contracts for spatial fast mode + flag composition."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from examples.inference.basic.mlx_wan22_generate import DEFAULT_HEIGHT, DEFAULT_NUM_FRAMES, DEFAULT_WIDTH
+from fastvideo.mlx_runtime.fast_spatial import (
+    DEFAULT_FAST_SPATIAL_SHARPEN,
+    apply_fast_spatial_upsample,
+    plan_fast_spatial,
+    resolve_spatial_mode,
+)
+from fastvideo.mlx_runtime.frame_upsample import upsample_frames
+
+
+def _frames(count: int, height: int, width: int, seed: int = 0) -> list[np.ndarray]:
+    rng = np.random.default_rng(seed)
+    return [rng.integers(0, 256, (height, width, 3), dtype=np.uint8) for _ in range(count)]
+
+
+def test_resolve_spatial_mode_priority() -> None:
+    assert resolve_spatial_mode(refine=False, fast_spatial=False) == "off"
+    assert resolve_spatial_mode(refine=False, fast_spatial=True) == "fast_spatial"
+    # Refine is the quality path and wins when both are requested.
+    assert resolve_spatial_mode(refine=True, fast_spatial=True) == "refine"
+    assert resolve_spatial_mode(refine=True, fast_spatial=False) == "refine"
+
+
+def test_plan_fast_spatial_matches_refine_geometry() -> None:
+    spatial = plan_fast_spatial(
+        height=480,
+        width=832,
+        num_frames=81,
+        spatial_scale=2,
+        vae_spatial_compression=8,
+    )
+    assert spatial.enabled
+    assert spatial.stage1_height == 240
+    assert spatial.stage1_width == 416
+    assert spatial.target_height == 480
+    assert spatial.target_width == 832
+    assert spatial.plan.stage1_latent_height == 30
+    assert spatial.plan.stage2_latent_width == 104
+    assert spatial.sharpen == DEFAULT_FAST_SPATIAL_SHARPEN
+
+
+def test_plan_fast_spatial_disabled() -> None:
+    spatial = plan_fast_spatial(height=480, width=832, num_frames=81, enabled=False)
+    assert not spatial.enabled
+    assert spatial.scale == 1
+
+
+def test_plan_fast_spatial_accepts_wan22_defaults() -> None:
+    assert (DEFAULT_HEIGHT, DEFAULT_WIDTH, DEFAULT_NUM_FRAMES) == (448, 832, 121)
+    spatial = plan_fast_spatial(
+        height=DEFAULT_HEIGHT,
+        width=DEFAULT_WIDTH,
+        num_frames=DEFAULT_NUM_FRAMES,
+        spatial_scale=2,
+        vae_spatial_compression=16,
+        vae_temporal_compression=4,
+        patch_size=(1, 2, 2),
+    )
+    assert (spatial.plan.stage1_latent_height, spatial.plan.stage1_latent_width) == (14, 26)
+    assert (spatial.plan.stage2_latent_height, spatial.plan.stage2_latent_width) == (28, 52)
+
+
+def test_apply_fast_spatial_upsample_resizes_decoded_frames() -> None:
+    """The upsample is pixel-space: it takes decoded frames, not latents.
+
+    Interpolating latents instead is what made spatial fast mode incoherent —
+    a blended Wan latent is off the decoder's manifold, so the decode came back
+    as the right silhouette under a smeared veil.
+    """
+    spatial = plan_fast_spatial(
+        height=480,
+        width=832,
+        num_frames=81,
+        spatial_scale=2,
+        vae_spatial_compression=8,
+        upsample_mode="lanczos",
+        sharpen=0.0,
+    )
+    frames = _frames(3, 240, 416)
+    out = apply_fast_spatial_upsample(frames, spatial)
+
+    assert len(out) == 3
+    assert all(frame.shape == (480, 832, 3) for frame in out)
+    assert all(frame.dtype == np.uint8 for frame in out)
+    expected = upsample_frames(frames, width=832, height=480, mode="lanczos", sharpen=0.0)
+    for got, want in zip(out, expected, strict=True):
+        np.testing.assert_array_equal(got, want)
+    # Inputs are not mutated.
+    assert all(frame.shape == (240, 416, 3) for frame in frames)
+
+
+def test_apply_fast_spatial_upsample_applies_sharpen() -> None:
+    frames = _frames(1, 240, 416)
+    soft = plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=2, sharpen=0.0)
+    crisp = plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=2, sharpen=0.8)
+    assert not np.array_equal(
+        apply_fast_spatial_upsample(frames, soft)[0],
+        apply_fast_spatial_upsample(frames, crisp)[0],
+    )
+
+
+def test_apply_fast_spatial_noop_when_disabled() -> None:
+    frames = _frames(2, 32, 32, seed=1)
+    spatial = plan_fast_spatial(height=32, width=32, num_frames=5, enabled=False, patch_size=(1, 1, 1))
+    out = apply_fast_spatial_upsample(frames, spatial)
+    for got, want in zip(out, frames, strict=True):
+        np.testing.assert_array_equal(got, want)
+
+
+def test_plan_fast_spatial_rejects_bad_mode() -> None:
+    with pytest.raises(ValueError, match="upsample mode"):
+        plan_fast_spatial(height=480, width=832, num_frames=81, upsample_mode="bicubic")
+
+
+def test_plan_fast_spatial_rejects_negative_sharpen() -> None:
+    with pytest.raises(ValueError, match="sharpen"):
+        plan_fast_spatial(height=480, width=832, num_frames=81, sharpen=-0.1)
+
+
+def test_geometry_errors_name_fast_spatial_not_refine() -> None:
+    """Both modes share the resolution splitter; the error must name the caller.
+
+    832 is not divisible by 3, so scale 3 is rejected at plan time.
+    """
+    with pytest.raises(ValueError, match=r"^fast-spatial requires"):
+        plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=3)

fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from fastvideo.benchmarks.mlx_fastwan_bench import (
+    ALLOWED_MODES,
+    BENCHMARK_PRESETS,
+    _html_grid,
+    _load_prompt_cases,
+    _mode_to_dtype_quant,
+    _parse_list,
+)
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [
+        ("fp16", ("fp16", None)),
+        ("bf16", ("bf16", None)),
+        ("int8", ("fp16", "int8")),
+        ("int4", ("fp16", "int4")),
+        ("mxfp8", ("fp16", "mxfp8")),
+        ("mxfp4", ("fp16", "mxfp4")),
+        ("nvfp4", ("fp16", "nvfp4")),
+    ],
+)
+def test_benchmark_modes_map_to_runtime_quantization(mode: str, expected: tuple[str, str | None]) -> None:
+    assert mode in ALLOWED_MODES
+    assert _mode_to_dtype_quant(mode) == expected
+
+
+def test_benchmark_rejects_unknown_modes() -> None:
+    with pytest.raises(ValueError, match="Unsupported modes"):
+        _parse_list("fp16,not_a_mode", ALLOWED_MODES, "modes")
+
+
+def test_load_prompt_cases_from_plain_text(tmp_path) -> None:
+    prompt_file = tmp_path / "prompts.txt"
+    prompt_file.write_text("\n# comment\nA fox runs through a forest.\nA raccoon walks in sunflowers.\n")
+
+    cases = _load_prompt_cases("unused", prompt_file)
+
+    assert [case.id for case in cases] == ["prompt-001", "prompt-002"]
+    assert [case.prompt for case in cases] == [
+        "A fox runs through a forest.",
+        "A raccoon walks in sunflowers.",
+    ]
+
+
+def test_load_prompt_cases_from_jsonl(tmp_path) -> None:
+    prompt_file = tmp_path / "prompts.jsonl"
+    prompt_file.write_text('{"id": "Fox Forest", "prompt": "A fox runs."}\n{"name": "clock", "caption": "A clock burns."}\n')
+
+    cases = _load_prompt_cases("unused", prompt_file)
+
+    assert [case.id for case in cases] == ["fox-forest", "clock"]
+    assert [case.prompt for case in cases] == ["A fox runs.", "A clock burns."]
+
+
+def test_load_builtin_prompt_set() -> None:
+    cases = _load_prompt_cases("unused", None, "motion7")
+    assert len(cases) == 7
+    assert cases[0].id == "beach-sunset"
+
+
+def test_benchmark_presets_include_memory_tiers() -> None:
+    assert BENCHMARK_PRESETS["mac-16gb"].modes == "int8"
+    assert BENCHMARK_PRESETS["mac-16gb"].decoders == "taehv"
+    assert BENCHMARK_PRESETS["mac-16gb"].mlx_memory_limit_gib == 16.0
+    assert BENCHMARK_PRESETS["mac-64gb"].decoders == "taehv,wan-vae"
+
+
+def test_html_grid_includes_video_and_sync_controls() -> None:
+    rendered = _html_grid([
+        {
+            "prompt_id": "fox",
+            "prompt": "A fox runs.",
+            "mode": "int8",
+            "decoder": "taehv",
+            "status": "ok",
+            "video_path": "fox/video_int8_taehv.mp4",
+            "total_s": 12.3,
+            "denoise_s": 10.0,
+            "decode_s": 1.0,
+            "peak_gib": 4.5,
+        }
+    ])
+
+    assert "Restart + play all" in rendered
+    assert "fox/video_int8_taehv.mp4" in rendered
+    assert "A fox runs." in rendered
+
+
+def test_denoise_dmd_on_device_runs_tiny_dit_and_reports_step_times(distributed_setup) -> None:
+    mx = pytest.importorskip("mlx.core", reason="MLX is required for the on-device denoise test")
+    import numpy as np
+    import torch
+
+    from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device
+    from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler,
+    )
+    from fastvideo.tests.mlx.tiny_wan import (
+        build_hf_config,
+        build_inputs,
+        build_tiny_wan_config,
+        build_torch_model,
+        mlx_dit_from_torch_model,
+        mlx_rotary_embeddings,
+    )
+
+    dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
+    hidden_states, encoder_hidden_states, _ = build_inputs()
+    schedule = MLXDMDSchedule.from_torch_scheduler(FlowMatchEulerDiscreteScheduler(shift=8.0))
+
+    timesteps = [1000, 757, 522]
+    generator = torch.Generator(device="cpu").manual_seed(7)
+    latents_seed = torch.randn(hidden_states.shape, generator=generator, dtype=torch.float32).numpy()
+    renoise_by_step = [
+        torch.randn(hidden_states.shape, generator=generator, dtype=torch.float32).numpy()
+        for _ in range(len(timesteps) - 1)
+    ]
+
+    latents_np, step_times = denoise_dmd_on_device(
+        mx=mx,
+        dit=dit,
+        latents=mx.array(latents_seed),
+        encoder_hidden_states=mx.array(encoder_hidden_states.numpy()),
+        freqs_cis=mlx_rotary_embeddings(hidden_states),
+        timesteps=timesteps,
+        renoise_by_step=renoise_by_step,
+        schedule=schedule,
+        dmd_step=dmd_step,
+        mx_dtype=mx.float32,
+    )
+
+    assert latents_np.shape == tuple(hidden_states.shape)
+    assert np.isfinite(latents_np).all()
+    assert len(step_times) == len(timesteps)
+    assert all(t > 0 for t in step_times)

fastvideo/tests/mlx/test_mlx_prompt_enhance.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only contracts for local prompt enrichment."""
+
+from __future__ import annotations
+
+import json
+
+from fastvideo.mlx_runtime.prompt_enhance import (
+    DEFAULT_ENHANCE_SYSTEM_PROMPT,
+    enhance_prompt,
+    enhance_prompt_template,
+    enhance_result_as_metrics,
+    load_or_enhance_prompt,
+)
+
+
+def test_template_adds_cinematic_cues_to_thin_prompt() -> None:
+    raw = "A red fox in the snow"
+    out = enhance_prompt_template(raw)
+    assert raw in out
+    lower = out.lower()
+    assert "lens" in lower or "camera" in lower or "cinematic" in lower
+    assert "light" in lower
+    assert out.endswith(".")
+
+
+def test_template_is_idempotent_on_rich_prompts() -> None:
+    rich = (
+        "A red fox trotting through a snowy pine forest at golden hour, "
+        "cinematic wide shot on a 35mm lens, soft volumetric lighting, "
+        "gentle camera dolly, highly detailed film grain and color graded."
+    )
+    assert enhance_prompt_template(rich) == " ".join(rich.split())
+
+
+def test_enhance_prompt_template_backend() -> None:
+    result = enhance_prompt("a paper boat on a stream", backend="template")
+    assert result.backend == "template"
+    assert result.changed
+    assert result.original == "a paper boat on a stream"
+    assert "paper boat" in result.enhanced.lower()
+
+
+def test_enhance_prompt_rejects_empty() -> None:
+    try:
+        enhance_prompt("   ", backend="template")
+    except ValueError as exc:
+        assert "non-empty" in str(exc)
+    else:  # pragma: no cover
+        raise AssertionError("expected ValueError")
+
+
+def test_load_or_enhance_prompt_caches(tmp_path) -> None:
+    result1 = load_or_enhance_prompt(
+        "a lantern in the fog",
+        backend="template",
+        cache=True,
+        cache_dir=tmp_path,
+    )
+    assert result1.backend == "template"
+    cached_files = list(tmp_path.glob("*.json"))
+    assert len(cached_files) == 1
+    payload = json.loads(cached_files[0].read_text())
+    assert payload["enhanced"] == result1.enhanced
+
+    result2 = load_or_enhance_prompt(
+        "a lantern in the fog",
+        backend="template",
+        cache=True,
+        cache_dir=tmp_path,
+    )
+    assert result2.enhanced == result1.enhanced
+    assert result2.backend == "cache"
+    assert result2.elapsed_s == 0.0
+
+
+def test_enhance_result_as_metrics_none() -> None:
+    metrics = enhance_result_as_metrics(None)
+    assert metrics["enhance_prompt"] is False
+    assert metrics["prompt_enhanced"] is None
+
+
+def test_default_system_prompt_matches_streaming_contract() -> None:
+    # Lockstep with fastvideo/entrypoints/streaming/prompt/enhancer.py
+    assert "cinematic video generation" in DEFAULT_ENHANCE_SYSTEM_PROMPT.lower()
+    assert "just the enhanced prompt" in DEFAULT_ENHANCE_SYSTEM_PROMPT.lower()

fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for the MLX prompt-to-video decode dispatch."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import numpy as np
+
+
+def _load_prompt_to_video_module():
+    root = Path(__file__).resolve().parents[3]
+    script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py"
+    spec = importlib.util.spec_from_file_location(
+        "mlx_wan_prompt_to_video_for_test", script
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_default_taehv_decode_uses_mlx_backend(monkeypatch, tmp_path):
+    module = _load_prompt_to_video_module()
+    import fastvideo.mlx_runtime.wan_vae as wan_vae
+
+    calls = []
+
+    def fake_decode_latents_to_video(latents_np, output_path, **kwargs):
+        calls.append((latents_np, output_path, kwargs))
+        output_path.write_bytes(b"fake mp4")
+        return {"backend": kwargs["backend"]}
+
+    monkeypatch.setattr(
+        wan_vae, "decode_latents_to_video", fake_decode_latents_to_video
+    )
+    latents = np.zeros((1, 16, 1, 2, 2), dtype=np.float32)
+    output_path = tmp_path / "out.mp4"
+
+    module.decode_latents_to_video(
+        model_root=tmp_path,
+        latents_np=latents,
+        output_path=output_path,
+        fps=16,
+        device_arg="cpu",
+        dtype_arg="fp32",
+        backend="taehv",
+        taehv_source_path=None,
+        taehv_checkpoint_path=None,
+        taehv_parallel=False,
+    )
+
+    assert output_path.read_bytes() == b"fake mp4"
+    assert len(calls) == 1
+    _, called_output_path, kwargs = calls[0]
+    assert called_output_path == output_path
+    assert kwargs["backend"] == "taehv"
+    assert kwargs["z_dim"] == 16
+
+
+def test_vendored_taehv_import_available():
+    from fastvideo.third_party.taehv import TAEHV
+
+    assert TAEHV.__name__ == "TAEHV"
+
+
+def test_default_model_resolution_downloads_only_needed_transformer(monkeypatch, tmp_path):
+    module = _load_prompt_to_video_module()
+    import huggingface_hub
+
+    calls = []
+
+    def fake_snapshot_download(model_id, allow_patterns):
+        calls.append((model_id, allow_patterns))
+        return str(tmp_path)
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+
+    assert module.resolve_model_root(None, include_transformer=True) == tmp_path
+    assert "transformer/*" in calls[-1][1]
+    assert module.resolve_model_root(None, include_transformer=False) == tmp_path
+    assert "transformer/*" not in calls[-1][1]
+    assert module.resolve_model_root(tmp_path) == tmp_path
+    assert len(calls) == 2

fastvideo/tests/mlx/test_mlx_refine.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only contracts for the MLX two-pass refine pipeline.
+
+Mirrors the LTX-2 refine stage math (init resolution split, spatial
+upsample, Gaussian re-noise) without requiring Metal / a real DiT. The
+orchestration helpers that need ``mx`` are exercised only when MLX is
+importable.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+
+from examples.inference.basic.mlx_wan22_generate import DEFAULT_HEIGHT, DEFAULT_NUM_FRAMES, DEFAULT_WIDTH
+from fastvideo.mlx_runtime.refine import (
+    DEFAULT_REFINE_SIGMA,
+    default_refine_timesteps,
+    plan_refine_resolutions,
+    prepare_refine_latents,
+    refine_sigma_from_schedule,
+    upsample_latents_spatial,
+)
+from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, add_noise
+
+
+def _fastwan_schedule() -> MLXDMDSchedule:
+    """FastWan's flow-match schedule: ``sigma == timestep / 1000``."""
+    timesteps = np.arange(1000, 0, -1, dtype=np.float64)
+    return MLXDMDSchedule(sigmas=timesteps / 1000.0, timesteps=timesteps)
+
+
+def test_plan_refine_resolutions_splits_even_target() -> None:
+    plan = plan_refine_resolutions(
+        height=480,
+        width=832,
+        num_frames=81,
+        spatial_scale=2,
+        vae_spatial_compression=8,
+        vae_temporal_compression=4,
+        patch_size=(1, 2, 2),
+    )
+    assert plan.stage1_height == 240
+    assert plan.stage1_width == 416
+    assert plan.target_height == 480
+    assert plan.target_width == 832
+    assert plan.stage1_latent_height == 30
+    assert plan.stage1_latent_width == 52
+    assert plan.stage2_latent_height == 60
+    assert plan.stage2_latent_width == 104
+    assert plan.latent_frames == 21
+    assert plan.spatial_scale == 2
+
+
+def test_plan_refine_disabled_collapses_to_single_pass() -> None:
+    plan = plan_refine_resolutions(
+        height=480,
+        width=832,
+        num_frames=81,
+        spatial_scale=2,
+        enabled=False,
+    )
+    assert plan.spatial_scale == 1
+    assert plan.stage1_height == plan.target_height == 480
+    assert plan.stage1_width == plan.target_width == 832
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "num_frames", "message"),
+    [(481, 832, 81, "height/width"), (480, 833, 81, "height/width"), (480, 832, 82, "num_frames")],
+)
+def test_plan_refine_rejects_requests_the_vae_would_truncate(
+    height: int,
+    width: int,
+    num_frames: int,
+    message: str,
+) -> None:
+    with pytest.raises(ValueError, match=message):
+        plan_refine_resolutions(height=height, width=width, num_frames=num_frames, enabled=False)
+
+
+def test_plan_refine_rejects_odd_target() -> None:
+    with pytest.raises(ValueError, match="divisible by"):
+        plan_refine_resolutions(height=481, width=832, num_frames=81, spatial_scale=2)
+
+
+def test_plan_refine_rejects_vae_misaligned_stage1() -> None:
+    # 500/2 = 250, 250 % 8 != 0 → rejected (same guard as LTX2RefineInitStage).
+    with pytest.raises(ValueError, match="divisible by"):
+        plan_refine_resolutions(
+            height=500,
+            width=800,
+            num_frames=81,
+            spatial_scale=2,
+            vae_spatial_compression=8,
+        )
+
+
+def test_plan_refine_wan22_5b_720p_geometry() -> None:
+    # Wan2.2 TI2V-5B native 704x1280 with 16x spatial VAE compression.
+    plan = plan_refine_resolutions(
+        height=704,
+        width=1280,
+        num_frames=121,
+        spatial_scale=2,
+        vae_spatial_compression=16,
+        vae_temporal_compression=4,
+        patch_size=(1, 2, 2),
+    )
+    assert plan.stage1_height == 352
+    assert plan.stage1_width == 640
+    assert plan.stage1_latent_height == 22
+    assert plan.stage1_latent_width == 40
+    assert plan.stage2_latent_height == 44
+    assert plan.stage2_latent_width == 80
+    assert plan.latent_frames == 31
+
+
+def test_plan_refine_wan22_default_geometry() -> None:
+    plan = plan_refine_resolutions(
+        height=DEFAULT_HEIGHT,
+        width=DEFAULT_WIDTH,
+        num_frames=DEFAULT_NUM_FRAMES,
+        spatial_scale=2,
+        vae_spatial_compression=16,
+        vae_temporal_compression=4,
+        patch_size=(1, 2, 2),
+    )
+    assert (plan.stage1_latent_height, plan.stage1_latent_width) == (14, 26)
+    assert (plan.stage2_latent_height, plan.stage2_latent_width) == (28, 52)
+    assert plan.latent_frames == 31
+
+
+def test_upsample_nearest_repeats_spatial_samples() -> None:
+    rng = np.random.default_rng(0)
+    latents = rng.standard_normal((1, 4, 2, 3, 5)).astype(np.float32)
+    up = upsample_latents_spatial(latents, scale=2, mode="nearest")
+    assert up.shape == (1, 4, 2, 6, 10)
+    # Every 2x2 block equals the source sample broadcast across the block.
+    for y in range(3):
+        for x in range(5):
+            block = up[0, :, :, 2 * y:2 * y + 2, 2 * x:2 * x + 2]
+            src = latents[0, :, :, y, x][:, :, None, None]
+            np.testing.assert_array_equal(block, np.broadcast_to(src, block.shape))
+
+
+def test_upsample_bilinear_preserves_constants() -> None:
+    latents = np.full((1, 2, 1, 4, 4), 3.5, dtype=np.float32)
+    up = upsample_latents_spatial(latents, scale=2, mode="bilinear")
+    assert up.shape == (1, 2, 1, 8, 8)
+    np.testing.assert_allclose(up, 3.5, rtol=0, atol=1e-6)
+
+
+def test_upsample_scale_one_is_identity() -> None:
+    rng = np.random.default_rng(1)
+    latents = rng.standard_normal((1, 2, 1, 4, 4)).astype(np.float32)
+    out = upsample_latents_spatial(latents, scale=1)
+    np.testing.assert_array_equal(out, latents)
+
+
+def test_prepare_refine_latents_matches_add_noise_formula() -> None:
+    rng = np.random.default_rng(2)
+    clean = rng.standard_normal((1, 4, 2, 3, 5)).astype(np.float32)
+    noise = rng.standard_normal((1, 4, 2, 6, 10)).astype(np.float32)
+    sigma = 0.75
+
+    got = prepare_refine_latents(
+        clean,
+        scale=2,
+        sigma=sigma,
+        noise=noise,
+        add_noise_flag=True,
+        upsample_mode="nearest",
+    )
+    up = upsample_latents_spatial(clean, scale=2, mode="nearest")
+    expected = add_noise(up, noise, sigma)
+    np.testing.assert_allclose(got, expected, rtol=0, atol=0)
+
+
+def test_prepare_refine_latents_no_noise_returns_upsample() -> None:
+    rng = np.random.default_rng(3)
+    clean = rng.standard_normal((1, 2, 1, 4, 4)).astype(np.float32)
+    got = prepare_refine_latents(clean, scale=2, add_noise_flag=False, upsample_mode="nearest")
+    expected = upsample_latents_spatial(clean, scale=2, mode="nearest")
+    np.testing.assert_array_equal(got, expected)
+
+
+def test_prepare_refine_latents_seed_is_deterministic() -> None:
+    rng = np.random.default_rng(4)
+    clean = rng.standard_normal((1, 2, 1, 3, 3)).astype(np.float32)
+    a = prepare_refine_latents(clean, scale=2, sigma=0.5, seed=123, upsample_mode="nearest")
+    b = prepare_refine_latents(clean, scale=2, sigma=0.5, seed=123, upsample_mode="nearest")
+    np.testing.assert_array_equal(a, b)
+
+
+def test_refine_sigma_from_schedule_uses_first_timestep() -> None:
+    timesteps = np.arange(1000, dtype=np.float64)
+    sigmas = np.linspace(1.0, 1e-3, 1000, dtype=np.float64)
+    schedule = MLXDMDSchedule(sigmas=sigmas, timesteps=timesteps)
+    # First refine timestep drives the stage-2 start sigma.
+    assert refine_sigma_from_schedule(schedule, [999, 757, 522]) == pytest.approx(float(sigmas[999]))
+    assert refine_sigma_from_schedule(schedule, [522]) == pytest.approx(float(sigmas[522]))
+
+
+def test_default_refine_sigma_matches_ltx2_stage2_head() -> None:
+    # Documented contract: DEFAULT_REFINE_SIGMA tracks LTX-2's first
+    # STAGE_2_DISTILLED_SIGMA_VALUES entry so the fallback is intentional.
+    assert math.isclose(DEFAULT_REFINE_SIGMA, 0.909375, rel_tol=0, abs_tol=0)
+
+
+def test_full_noise_handoff_would_discard_stage1() -> None:
+    """Why the stage-2 grid may not open at the stage-1 timestep.
+
+    FastWan's DMD grid starts at ``t=1000``, i.e. ``sigma == 1``. The hand-off
+    is ``(1 - sigma) * upsampled + sigma * noise``, so starting stage 2 there
+    weights the stage-1 result at exactly zero: refine silently degrades to a
+    plain full-resolution run costing two passes.
+    """
+    schedule = _fastwan_schedule()
+    assert refine_sigma_from_schedule(schedule, [1000, 757, 522]) == pytest.approx(1.0)
+
+    clean = np.random.default_rng(0).standard_normal((1, 2, 1, 4, 4)).astype(np.float32)
+    handoff = prepare_refine_latents(clean, scale=2, sigma=1.0, seed=7, upsample_mode="nearest")
+    upsampled = upsample_latents_spatial(clean, scale=2, mode="nearest")
+    # Nothing of stage 1 survives.
+    assert not np.allclose(handoff, upsampled)
+
+
+def test_default_refine_timesteps_drops_the_full_noise_step() -> None:
+    schedule = _fastwan_schedule()
+    steps = default_refine_timesteps(schedule, [1000, 757, 522])
+    assert steps == [757.0, 522.0]
+    # The resulting hand-off keeps a real share of stage 1.
+    sigma = refine_sigma_from_schedule(schedule, steps)
+    assert 0.0 < sigma < 1.0
+
+
+def test_default_refine_timesteps_keeps_an_already_valid_grid() -> None:
+    schedule = _fastwan_schedule()
+    assert default_refine_timesteps(schedule, [757, 522]) == [757.0, 522.0]
+
+
+def test_default_refine_timesteps_rejects_an_all_full_noise_grid() -> None:
+    schedule = _fastwan_schedule()
+    with pytest.raises(ValueError, match="No usable refine timesteps"):
+        default_refine_timesteps(schedule, [1000])

fastvideo/tests/mlx/test_mlx_rife_interpolation.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for the Apple-Silicon MLX RIFE interpolation backend."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from huggingface_hub.utils import LocalEntryNotFoundError
+
+mx = pytest.importorskip("mlx.core", reason="MLX required for RIFE tests")
+
+
+@pytest.mark.skipif(
+    not bool(getattr(mx, "metal", None) and mx.metal.is_available()),
+    reason="RIFE MLX regression requires Apple Silicon Metal",
+)
+def test_rife_interpolation_preserves_keyframes_shape_and_count() -> None:
+    from fastvideo.mlx_runtime.rife_interp import RIFEWeightsUnavailableError, interpolate, load_model
+
+    frame0 = np.zeros((64, 96, 3), dtype=np.uint8)
+    frame1 = np.zeros((64, 96, 3), dtype=np.uint8)
+    frame1[:, :, 0] = 255
+    try:
+        # The backend is vendored, but its weights are fetched from Hugging Face
+        # on first use, so an offline machine should skip rather than fail.
+        model = load_model()
+    except RIFEWeightsUnavailableError as exc:
+        pytest.skip(f"RIFE weights unavailable: {exc}")
+
+    frames = interpolate([frame0, frame1], factor=2, model=model)
+
+    assert len(frames) == 3
+    assert frames[1].shape == frame0.shape
+    assert frames[1].dtype == np.uint8
+    np.testing.assert_array_equal(frames[0], frame0)
+    np.testing.assert_array_equal(frames[-1], frame1)
+
+
+def test_rife_download_unavailable_has_specific_error(monkeypatch) -> None:
+    from fastvideo.mlx_runtime.rife_interp import RIFEWeightsUnavailableError, load_model
+    from fastvideo.third_party.rife_mlx.utils import weights
+
+    def unavailable(*args, **kwargs):
+        raise LocalEntryNotFoundError("offline")
+
+    monkeypatch.setattr(weights, "build_model", unavailable)
+    load_model.cache_clear()
+    with pytest.raises(RIFEWeightsUnavailableError):
+        load_model()
+
+
+def test_rife_backend_regression_is_not_skip_eligible(monkeypatch) -> None:
+    from fastvideo.mlx_runtime.rife_interp import RIFEBackendError, RIFEWeightsUnavailableError, load_model
+    from fastvideo.third_party.rife_mlx.utils import weights
+
+    def broken(*args, **kwargs):
+        raise AssertionError("backend regression")
+
+    monkeypatch.setattr(weights, "build_model", broken)
+    load_model.cache_clear()
+    with pytest.raises(RIFEBackendError) as caught:
+        load_model()
+    assert not isinstance(caught.value, RIFEWeightsUnavailableError)

fastvideo/tests/mlx/test_mlx_taehv_parity.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+"""TAEHV decode parity tests for the MLX implementation.
+
+TAEHV MLX vs torch is bit-close (atol 1e-5) for z_dim=16 and 48.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+mx = pytest.importorskip("mlx.core", reason="MLX required for Wan VAE/TAEHV tests")
+
+from fastvideo.mlx_runtime.wan_vae import (  # noqa: E402
+    decode_latents_taehv_mlx,
+    ensure_taehv_checkpoint,
+)
+
+
+@pytest.mark.parametrize("z_dim", [16, 48])
+def test_taehv_mlx_matches_torch(z_dim: int) -> None:
+    import torch
+    from fastvideo.third_party.taehv import TAEHV
+
+    ckpt = ensure_taehv_checkpoint(z_dim=z_dim)
+    rng = np.random.default_rng(0)
+    h = w = 8
+    t = 5
+    lat = (rng.standard_normal((1, z_dim, t, h, w)) * 0.5).astype(np.float32)
+
+    model = TAEHV(str(ckpt)).eval()
+    with torch.no_grad():
+        out_t = model.decode_video(torch.from_numpy(lat).transpose(1, 2), parallel=True, show_progress_bar=False)
+    torch_np = out_t[0].permute(0, 2, 3, 1).float().numpy()
+    mlx_np = decode_latents_taehv_mlx(lat, z_dim=z_dim)[0]
+    # Validate complete temporal output shape before comparing values.
+    assert mlx_np.shape[0] == torch_np.shape[0], (
+        f"MLX output frame count {mlx_np.shape[0]} != torch {torch_np.shape[0]}"
+    )
+    np.testing.assert_allclose(mlx_np, torch_np, atol=1e-5, rtol=1e-5)

fastvideo/tests/mlx/test_mlx_wan22_memory_handoff.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression test for the Wan2.2 MLX/MPS memory handoff."""
+
+from __future__ import annotations
+
+import gc
+import importlib.util
+import json
+import sys
+import weakref
+from pathlib import Path
+
+import numpy as np
+
+
+def _load_wan22_module():
+    root = Path(__file__).resolve().parents[3]
+    script = root / "examples/inference/basic/mlx_wan22_generate.py"
+    spec = importlib.util.spec_from_file_location("mlx_wan22_generate_memory_test", script)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_wan22_releases_compiled_dit_before_torch_decode_and_mlx_rife(monkeypatch, tmp_path: Path) -> None:
+    import mlx.core as mx
+    import torch
+    from diffusers import utils as diffusers_utils
+
+    from examples.inference.basic import mlx_wan_prompt_to_video as shared
+    from fastvideo.mlx_runtime import wan22, wan22_sample, wan_vae
+
+    module = _load_wan22_module()
+    checkpoint = tmp_path / "mlx"
+    checkpoint.mkdir()
+    (checkpoint / "mlx_dit.json").write_text(
+        json.dumps({"config": {"patch_size": [1, 2, 2], "in_channels": 1}})
+    )
+
+    events: list[str] = []
+    dit_ref = None
+
+    class CyclicCompiledDiT:
+        def __init__(self) -> None:
+            self.compiled = self.forward
+
+        def forward(self) -> None:
+            pass
+
+    def fake_load(*args, **kwargs):
+        nonlocal dit_ref
+        assert kwargs["compile"] is True
+        dit = CyclicCompiledDiT()
+        dit_ref = weakref.ref(dit)
+        return dit
+
+    def fake_decode(*args, **kwargs):
+        events.append("decode")
+        assert dit_ref is not None and dit_ref() is None
+        assert events[-3:] == ["gc", "mlx-clear", "decode"]
+        return np.zeros((1, 1, 2, 2, 3), dtype=np.float32)
+
+    def fake_postprocess(**kwargs) -> None:
+        events.append("rife")
+        assert events[-4:] == ["decode", "gc", "mps-clear", "rife"]
+
+    real_collect = gc.collect
+
+    def tracked_collect() -> int:
+        events.append("gc")
+        return real_collect()
+
+    monkeypatch.setattr(shared, "encode_prompt", lambda **kwargs: torch.zeros((1, 1, 1)))
+    monkeypatch.setattr(shared, "make_rotary_embeddings", lambda *args, **kwargs: None)
+    monkeypatch.setattr(shared, "_postprocess_video", fake_postprocess)
+    monkeypatch.setattr(wan22, "mlx_wan22_dit_from_mlx_checkpoint", fake_load)
+    monkeypatch.setattr(wan22_sample, "sample_wan22_dmd", lambda *args, **kwargs: args[2])
+    monkeypatch.setattr(wan_vae, "decode_latents_wan_vae_torch", fake_decode)
+    monkeypatch.setattr(diffusers_utils, "export_to_video", lambda *args, **kwargs: None)
+    monkeypatch.setattr(gc, "collect", tracked_collect)
+    monkeypatch.setattr(mx, "clear_cache", lambda: events.append("mlx-clear"))
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+    monkeypatch.setattr(torch.mps, "empty_cache", lambda: events.append("mps-clear"))
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "mlx_wan22_generate.py",
+            "--text-encoder-root",
+            str(tmp_path / "text"),
+            "--no-prompt-cache",
+            "--mlx-checkpoint",
+            str(checkpoint),
+            "--height",
+            "32",
+            "--width",
+            "32",
+            "--num-frames",
+            "1",
+            "--dmd-denoising-steps",
+            "1",
+            "--compile",
+            "--fast",
+            "--decode-backend",
+            "wan-vae",
+            "--vae-root",
+            str(tmp_path / "vae"),
+            "--output-path",
+            str(tmp_path / "out.mp4"),
+        ],
+    )
+
+    module.main()
+
+    assert events == ["gc", "mlx-clear", "decode", "gc", "mps-clear", "rife"]

fastvideo/tests/mlx/test_mlx_wan22_parity.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Track D Rung 2: per-token-timestep parity for the Wan2.2-TI2V-5B MLX port.
+
+The 5B (FullAttn) differs from the ported 2.1 only in scale (config) and
+*per-token* timestep conditioning (``expand_timesteps``: timestep is ``[B, L]``,
+one noise level per patch token — how TI2V keeps the image frame at t=0 while
+video frames are noised). Passing a 2-D timestep drives the torch model's
+per-token path (``ts_seq_len``), so this compares ``MLXWan22DiT`` against the
+torch reference on a tiny random-weight config. This is the run-6 prereq gate;
+runs on CPU in CI. Real 5B-weight parity is a separate Metal-gated test.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+mx = pytest.importorskip("mlx.core", reason="MLX is required for the Wan2.2 parity test")
+
+from fastvideo.forward_context import set_forward_context  # noqa: E402
+from fastvideo.mlx_runtime.fastwan import mlx_block_weights_from_torch  # noqa: E402
+from fastvideo.mlx_runtime.wan22 import (  # noqa: E402
+    MLXWan22DiT,
+    MLXWan22TransformerBlock,
+)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch  # noqa: E402
+from fastvideo.tests.mlx.tiny_wan import (  # noqa: E402
+    TINY_ARCH,
+    TOP_LEVEL_KEY_MAP,
+    build_hf_config,
+    build_tiny_wan_config,
+    build_torch_model,
+    mlx_rotary_embeddings,
+)
+
+
+def _mlx_wan22_from_torch(model, hf_config, *, compile: bool = False) -> MLXWan22DiT:
+    state = {name: value.detach().float() for name, value in model.state_dict().items()}
+    inner_dim = int(hf_config["num_attention_heads"]) * int(hf_config["attention_head_dim"])
+    weights = {}
+    for mlx_name, torch_name in TOP_LEVEL_KEY_MAP.items():
+        tensor = state[torch_name]
+        if mlx_name == "patch_embedding.weight":
+            tensor = tensor.reshape(inner_dim, -1)
+        weights[mlx_name] = mx.array(tensor.numpy())
+    blocks = [
+        MLXWan22TransformerBlock(
+            mlx_block_weights_from_torch(tb), dim=inner_dim, ffn_dim=int(hf_config["ffn_dim"]),
+            num_heads=int(hf_config["num_attention_heads"]), eps=float(hf_config["eps"])) for tb in model.blocks
+    ]
+    return MLXWan22DiT(weights, blocks, dict(hf_config), compile=compile)
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_wan22_per_token_timestep_matches_torch() -> None:
+    torch_model = build_torch_model()
+    hf_config = build_hf_config(build_tiny_wan_config())
+    mlx_model = _mlx_wan22_from_torch(torch_model, hf_config)
+
+    # Tiny latent: 4 frames, 8x8 latent, patch (1,2,2) -> 4*4*4 = 64 tokens (L).
+    frames, height, width = 4, 8, 8
+    p_t, p_h, p_w = TINY_ARCH["patch_size"]
+    num_tokens = (frames // p_t) * (height // p_h) * (width // p_w)
+    tokens_per_frame = num_tokens // (frames // p_t)
+
+    gen = torch.Generator().manual_seed(11)
+    hidden = torch.randn(1, TINY_ARCH["in_channels"], frames, height, width, generator=gen, dtype=torch.float32)
+    text = torch.randn(1, 8, TINY_ARCH["text_dim"], generator=gen, dtype=torch.float32)
+
+    # Per-token timestep [B, L]: frame 0 clean (t=0, I2V-style), rest noised (t=500).
+    per_frame = [0] + [500] * (frames // p_t - 1)
+    timestep = torch.tensor([[per_frame[i // tokens_per_frame] for i in range(num_tokens)]], dtype=torch.long)
+
+    with torch.no_grad(), set_forward_context(
+            current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")):
+        ref = torch_model(hidden_states=hidden, encoder_hidden_states=text, timestep=timestep).detach().float().numpy()
+
+    freqs_cis = mlx_rotary_embeddings(hidden)
+    out = mlx_model(mx.array(hidden.numpy()), mx.array(text.numpy()), mx.array(timestep.float().numpy()), freqs_cis)
+    mx.eval(out)
+    mlx_out = np.array(out.astype(mx.float32))
+
+    np.testing.assert_allclose(mlx_out, ref, atol=2e-3, rtol=2e-3)
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_wan22_compile_path_matches_eager() -> None:
+    torch_model = build_torch_model()
+    hf_config = build_hf_config(build_tiny_wan_config())
+    eager_model = _mlx_wan22_from_torch(torch_model, hf_config)
+    compiled_model = _mlx_wan22_from_torch(torch_model, hf_config, compile=True)
+
+    frames, height, width = 4, 8, 8
+    p_t, p_h, p_w = TINY_ARCH["patch_size"]
+    num_tokens = (frames // p_t) * (height // p_h) * (width // p_w)
+    tokens_per_frame = num_tokens // (frames // p_t)
+
+    gen = torch.Generator().manual_seed(12)
+    hidden = torch.randn(
+        1,
+        TINY_ARCH["in_channels"],
+        frames,
+        height,
+        width,
+        generator=gen,
+        dtype=torch.float32,
+    )
+    text = torch.randn(
+        1,
+        8,
+        TINY_ARCH["text_dim"],
+        generator=gen,
+        dtype=torch.float32,
+    )
+    per_frame = [0] + [500] * (frames // p_t - 1)
+    timestep = torch.tensor(
+        [[per_frame[i // tokens_per_frame] for i in range(num_tokens)]],
+        dtype=torch.long,
+    )
+    freqs_cis = mlx_rotary_embeddings(hidden)
+
+    eager = eager_model(
+        mx.array(hidden.numpy()),
+        mx.array(text.numpy()),
+        mx.array(timestep.float().numpy()),
+        freqs_cis,
+    )
+    compiled = compiled_model(
+        mx.array(hidden.numpy()),
+        mx.array(text.numpy()),
+        mx.array(timestep.float().numpy()),
+        freqs_cis,
+    )
+    mx.eval(eager, compiled)
+
+    np.testing.assert_allclose(
+        np.array(compiled.astype(mx.float32)),
+        np.array(eager.astype(mx.float32)),
+        atol=1e-4,
+        rtol=1e-4,
+    )

fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py

@@ -0,0 +1,318 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan2.2 prompt-embedding cache fingerprint regression tests."""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from fastvideo.mlx_runtime.prompt_cache import load_prompt_cache, save_prompt_cache
+
+
+def _load_wan22_module():
+    root = Path(__file__).resolve().parents[3]
+    script = root / "examples/inference/basic/mlx_wan22_generate.py"
+    spec = importlib.util.spec_from_file_location(
+        "mlx_wan22_generate_for_test", script
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_wan21_module():
+    root = Path(__file__).resolve().parents[3]
+    script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py"
+    spec = importlib.util.spec_from_file_location(
+        "mlx_wan_prompt_to_video_for_cache_test", script
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_wan22_prompt_cache_rejects_mismatched_fingerprint(tmp_path: Path) -> None:
+    module = _load_wan22_module()
+    cache_path = tmp_path / "prompt.npy"
+    text_encoder_root = tmp_path / "encoder"
+    text_encoder_root.mkdir()
+
+    fingerprint = module._prompt_cache_fingerprint(
+        prompt="a fox",
+        prompt_used="a fox, cinematic",
+        enhance_prompt=True,
+        enhance_prompt_backend="template",
+        text_encoder_root=text_encoder_root,
+        max_sequence_length=512,
+        dtype="fp16",
+    )
+    embeds = np.ones((1, 512, 4), dtype=np.float32)
+
+    save_prompt_cache(cache_path, embeds, fingerprint)
+
+    assert load_prompt_cache(cache_path, fingerprint) is not None
+
+    changed = dict(fingerprint)
+    changed["prompt"] = "a cat"
+    assert load_prompt_cache(cache_path, changed) is None
+
+
+def test_wan22_prompt_cache_rejects_missing_metadata(tmp_path: Path) -> None:
+    module = _load_wan22_module()
+    cache_path = tmp_path / "prompt.npy"
+    np.save(cache_path, np.zeros((1, 512, 4), dtype=np.float32))
+
+    fingerprint = module._prompt_cache_fingerprint(
+        prompt="a fox",
+        prompt_used="a fox",
+        enhance_prompt=False,
+        enhance_prompt_backend="template",
+        text_encoder_root=tmp_path,
+        max_sequence_length=512,
+        dtype="fp16",
+    )
+
+    assert load_prompt_cache(cache_path, fingerprint) is None
+
+
+def test_prompt_cache_is_best_effort_and_rejects_torn_data(tmp_path: Path) -> None:
+    cache_path = tmp_path / "prompt.npy"
+    fingerprint = {"prompt": "a fox"}
+    save_prompt_cache(cache_path, np.ones((1, 1, 1)), fingerprint)
+
+    np.save(cache_path, np.full((1, 1, 1), 9.0))
+    assert load_prompt_cache(cache_path, fingerprint) is None
+
+    invalid_parent = tmp_path / "not-a-directory"
+    invalid_parent.write_text("occupied")
+    assert not save_prompt_cache(
+        invalid_parent / "prompt.npy",
+        np.ones((1, 1, 1)),
+        fingerprint,
+    )
+
+    incomplete_path = tmp_path / "incomplete.npy"
+    incomplete = {"text_encoder": {"complete": False}}
+    assert not save_prompt_cache(incomplete_path, np.ones((1, 1, 1)), incomplete)
+    assert not incomplete_path.exists()
+
+
+def test_wan21_explicit_cache_reencodes_for_a_changed_prompt(monkeypatch, tmp_path: Path) -> None:
+    import torch
+
+    module = _load_wan21_module()
+    calls = []
+
+    def fake_encode_prompt(**kwargs):
+        calls.append(kwargs["prompt"])
+        return torch.tensor([[float(len(calls))]])
+
+    monkeypatch.setattr(module, "encode_prompt", fake_encode_prompt)
+    cache_path = tmp_path / "prompt.npy"
+    kwargs = {
+        "model_root": tmp_path / "model",
+        "max_sequence_length": 512,
+        "device_arg": "cpu",
+        "dtype_arg": "fp16",
+        "encode_mode": "inline",
+        "cache_path": cache_path,
+    }
+    first = module.get_prompt_embeds(prompt="first", **kwargs)
+    second = module.get_prompt_embeds(prompt="second", **kwargs)
+    second_cached = module.get_prompt_embeds(prompt="second", **kwargs)
+
+    assert calls == ["first", "second"]
+    assert not torch.equal(first, second)
+    assert torch.equal(second, second_cached)
+
+
+def test_wan22_missing_paths_download_only_selected_assets(monkeypatch, tmp_path: Path) -> None:
+    module = _load_wan22_module()
+    import huggingface_hub
+
+    calls = []
+
+    def fake_snapshot_download(model_id, allow_patterns):
+        calls.append((model_id, allow_patterns))
+        return str(tmp_path / ("wan21" if "2.1" in model_id else "wan22"))
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+    text_root, dit_checkpoint, dit_config, vae_root = module._resolve_model_paths(
+        text_encoder_root=None,
+        dit_checkpoint=None,
+        dit_config=None,
+        vae_root=None,
+        mlx_checkpoint=None,
+        decode_backend="taehv",
+    )
+
+    assert text_root == tmp_path / "wan21"
+    assert dit_checkpoint == tmp_path / "wan22/transformer/diffusion_pytorch_model.safetensors"
+    assert dit_config == tmp_path / "wan22/transformer/config.json"
+    assert vae_root is None
+    assert calls == [
+        (module.FASTWAN21_MODEL_ID, ["tokenizer/*", "text_encoder/*"]),
+        (
+            module.FASTWAN22_MODEL_ID,
+            ["transformer/diffusion_pytorch_model.safetensors", "transformer/config.json"],
+        ),
+    ]
+
+
+def test_wan22_explicit_mlx_paths_do_not_download(monkeypatch, tmp_path: Path) -> None:
+    module = _load_wan22_module()
+    import huggingface_hub
+
+    def unexpected_download(*args, **kwargs):
+        pytest.fail("explicit MLX and text paths must not download model assets")
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", unexpected_download)
+    text_root, dit_checkpoint, dit_config, vae_root = module._resolve_model_paths(
+        text_encoder_root=tmp_path / "text",
+        dit_checkpoint=None,
+        dit_config=None,
+        vae_root=None,
+        mlx_checkpoint=tmp_path / "mlx",
+        decode_backend="taehv",
+    )
+
+    assert text_root == tmp_path / "text"
+    assert dit_checkpoint is None
+    assert dit_config is None
+    assert vae_root is None
+
+
+def test_wan22_prompt_cache_has_a_default_path(tmp_path: Path) -> None:
+    """The cache must work without an explicit --prompt-embeds-cache.
+
+    It previously only engaged when handed a path, so every 5B run paid a full
+    UMT5 encode (~45s on an M4 Max) even for a repeat prompt, while the Wan2.1
+    entrypoint cached by default.
+    """
+    module = _load_wan22_module()
+    fingerprint = module._prompt_cache_fingerprint(
+        prompt="a fox",
+        prompt_used="a fox",
+        enhance_prompt=False,
+        enhance_prompt_backend="none",
+        text_encoder_root=tmp_path,
+        max_sequence_length=512,
+        dtype="fp16",
+    )
+    path = module._default_prompt_cache_path(fingerprint)
+    assert path.suffix == ".npy"
+    assert path.parent.name == "prompt_embeds"
+    assert path.name.startswith("wan22_")
+
+
+def test_wan22_default_cache_path_tracks_the_fingerprint(tmp_path: Path) -> None:
+    """Two prompts must not collide, and the same prompt must be stable."""
+    module = _load_wan22_module()
+
+    def fp(prompt: str):
+        return module._prompt_cache_fingerprint(
+            prompt=prompt,
+            prompt_used=prompt,
+            enhance_prompt=False,
+            enhance_prompt_backend="none",
+            text_encoder_root=tmp_path,
+            max_sequence_length=512,
+            dtype="fp16",
+        )
+
+    assert module._default_prompt_cache_path(fp("a fox")) == module._default_prompt_cache_path(fp("a fox"))
+    assert module._default_prompt_cache_path(fp("a fox")) != module._default_prompt_cache_path(fp("a cat"))
+
+
+def test_wan22_default_cache_path_tracks_encoder_files(tmp_path: Path) -> None:
+    module = _load_wan22_module()
+    weights = tmp_path / "text_encoder/weights.safetensors"
+    weights.parent.mkdir()
+    weights.write_bytes(b"one")
+
+    def path():
+        fingerprint = module._prompt_cache_fingerprint(
+            prompt="a fox",
+            prompt_used="a fox",
+            enhance_prompt=False,
+            enhance_prompt_backend="none",
+            text_encoder_root=tmp_path,
+            max_sequence_length=512,
+            dtype="fp16",
+        )
+        return module._default_prompt_cache_path(fingerprint)
+
+    before = path()
+    weights.write_bytes(b"two")
+    stat = weights.stat()
+    os.utime(weights, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1))
+    assert path() != before
+
+
+@pytest.mark.parametrize(
+    ("cache_args", "expected"),
+    [
+        ([], "default"),
+        (["--no-prompt-cache"], None),
+        (["--prompt-embeds-cache", "explicit.npy"], "explicit"),
+        (["--no-prompt-cache", "--prompt-embeds-cache", "explicit.npy"], "explicit"),
+    ],
+)
+def test_wan22_main_resolves_prompt_cache(
+    monkeypatch,
+    tmp_path: Path,
+    cache_args: list[str],
+    expected: str | None,
+) -> None:
+    module = _load_wan22_module()
+    checkpoint = tmp_path / "mlx"
+    checkpoint.mkdir()
+    (checkpoint / "mlx_dit.json").write_text(
+        json.dumps({"config": {"patch_size": [1, 2, 2], "in_channels": 1}})
+    )
+    explicit = tmp_path / "explicit.npy"
+
+    class CacheProbe(Exception):
+        pass
+
+    def probe(cache_path, fingerprint):
+        if expected == "default":
+            assert cache_path == module._default_prompt_cache_path(fingerprint)
+        elif expected == "explicit":
+            assert cache_path == explicit
+        else:
+            assert cache_path is None
+        raise CacheProbe
+
+    resolved_args = [str(explicit) if arg == "explicit.npy" else arg for arg in cache_args]
+    monkeypatch.setenv("HOME", str(tmp_path))
+    monkeypatch.setattr(module, "load_prompt_cache", probe)
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "mlx_wan22_generate.py",
+            "--text-encoder-root",
+            str(tmp_path / "text"),
+            "--mlx-checkpoint",
+            str(checkpoint),
+            "--height",
+            "32",
+            "--width",
+            "32",
+            "--num-frames",
+            "1",
+            *resolved_args,
+        ],
+    )
+
+    with pytest.raises(CacheProbe):
+        module.main()

fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Track D Rung 3: real-weight T2V parity for Wan2.2-TI2V-5B MLX vs torch.
+
+Loads the released ``FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`` transformer
+into both ``MLXWan22DiT`` and torch ``WanTransformer3DModel``, runs one forward
+with a 2-D per-token timestep (frame 0 at t=0 I2V-style, remaining tokens at
+t≈900), and asserts allclose. Gated on Metal + local weights so Linux
+``mlx[cpu]`` CI stays green.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+# The torch reference side brings up FastVideo's single-process distributed
+# environment, which rendezvouses over env://. Match the other torch-backed
+# tests in this repo and provide a local rendezvous before that happens.
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29513")
+
+mx = pytest.importorskip("mlx.core", reason="MLX is required for the Wan2.2 real-weight parity test")
+
+_ROOT = Path(os.environ.get("FASTVIDEO_WAN22_5B_ROOT", str(Path.home() / "models" / "fastwan22_5b")))
+_CHECKPOINT = _ROOT / "transformer" / "diffusion_pytorch_model.safetensors"
+_CONFIG = _ROOT / "transformer" / "config.json"
+
+_HAS_METAL = bool(getattr(mx, "metal", None) and mx.metal.is_available())
+_HAS_WEIGHTS = _CHECKPOINT.exists() and _CONFIG.exists() and _CHECKPOINT.stat().st_size > 1_000_000_000
+
+# The real-weight test loads ~10 GB fp16 weights + activations. Gate on the
+# documented 36 GB requirement unless explicitly opted in.
+def _check_memory_requirement() -> bool:
+    """Determine whether the system meets the memory requirement for the test.
+
+    Returns:
+        bool: `true` if the system has at least 36 GiB of memory or low-memory
+        execution is explicitly enabled, `false` otherwise.
+    """
+    if os.environ.get("FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY") == "1":
+        return True
+    try:
+        # On macOS, check unified memory via sysctl.
+        import subprocess
+        result = subprocess.run(
+            ["sysctl", "-n", "hw.memsize"],
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=5,
+        )
+        if result.returncode == 0:
+            mem_bytes = int(result.stdout.strip())
+            mem_gib = mem_bytes / (1024**3)
+            return mem_gib >= 36.0
+    except Exception:
+        pass
+    # Conservative fallback: require opt-in if we can't detect memory.
+    return False
+
+
+_HAS_SUFFICIENT_MEMORY = _check_memory_requirement()
+
+pytestmark = [
+    pytest.mark.skipif(not _HAS_METAL, reason="Metal required for real 5B fp16 parity"),
+    pytest.mark.skipif(
+        not _HAS_WEIGHTS,
+        reason=f"Wan2.2-5B checkpoint not found/incomplete under {_ROOT} "
+        "(set FASTVIDEO_WAN22_5B_ROOT; expected ~10 GB safetensors)",
+    ),
+    pytest.mark.skipif(
+        not _HAS_SUFFICIENT_MEMORY,
+        reason="Real-weight 5B test requires >= 36 GB unified memory "
+        "(set FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY=1 to override)",
+    ),
+]
+
+
+def _load_torch_wan22_from_diffusers(checkpoint: Path, config_path: Path, *, dtype):
+    """
+    Instantiate a Wan2.2 transformer and load its Diffusers-format checkpoint weights.
+
+    Parameters:
+        checkpoint (Path): Path to the safetensors checkpoint.
+        config_path (Path): Path to the Diffusers configuration file.
+        dtype: Data type to use for the returned model.
+
+    Returns:
+        WanTransformer3DModel: An evaluation-mode model with the loaded weights.
+    """
+    import torch
+    from safetensors.torch import load_file
+
+    from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig
+    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
+    from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict
+
+    hf_config = json.loads(config_path.read_text())
+    # Diffusers-only keys that are not arch fields.
+    arch_kwargs = {
+        k: v
+        for k, v in hf_config.items()
+        if k not in {"_class_name", "_name_or_path", "_diffusers_version"}
+    }
+    cfg = WanVideoConfig(arch_config=WanVideoArchConfig(**{
+        "num_attention_heads": int(arch_kwargs["num_attention_heads"]),
+        "attention_head_dim": int(arch_kwargs["attention_head_dim"]),
+        "in_channels": int(arch_kwargs["in_channels"]),
+        "out_channels": int(arch_kwargs["out_channels"]),
+        "text_dim": int(arch_kwargs["text_dim"]),
+        "freq_dim": int(arch_kwargs["freq_dim"]),
+        "ffn_dim": int(arch_kwargs["ffn_dim"]),
+        "num_layers": int(arch_kwargs["num_layers"]),
+        "patch_size": tuple(arch_kwargs["patch_size"]),
+        "cross_attn_norm": bool(arch_kwargs.get("cross_attn_norm", True)),
+        "qk_norm": arch_kwargs.get("qk_norm", "rms_norm_across_heads"),
+        "eps": float(arch_kwargs.get("eps", 1e-6)),
+        "rope_max_seq_len": int(arch_kwargs.get("rope_max_seq_len", 1024)),
+        "added_kv_proj_dim": arch_kwargs.get("added_kv_proj_dim"),
+        "image_dim": arch_kwargs.get("image_dim"),
+        "pos_embed_seq_len": arch_kwargs.get("pos_embed_seq_len"),
+    }))
+    model = WanTransformer3DModel(config=cfg, hf_config=hf_config).eval()
+    raw = load_file(str(checkpoint), device="cpu")
+    mapping = get_param_names_mapping(model.param_names_mapping)
+    custom_sd, _ = hf_to_custom_state_dict(raw.items(), mapping)
+    missing, unexpected = model.load_state_dict(custom_sd, strict=False)
+    # Diffusers layout should map cleanly; unexpected keys are always an error.
+    assert not unexpected, f"unexpected keys: {unexpected[:8]}"
+    # Missing keys must be either empty buffers or explicitly approved. Reject
+    # missing parameters to prevent parity checks from running with uninitialized weights.
+    if missing:
+        # Allowlist for known empty buffers (if any exist for Wan2.2-5B).
+        # Update this list if legitimate empty buffers are identified.
+        allowed_missing = set()
+        unapproved = [k for k in missing if k not in allowed_missing]
+        assert not unapproved, (
+            f"Missing keys in Wan2.2-5B state dict load: {unapproved[:8]}. "
+            "Cannot proceed with parity check; uninitialized parameters would "
+            "invalidate results."
+        )
+    return model.to(dtype=dtype)
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None:
+    """
+    Verify that the MLX and PyTorch Wan2.2 implementations produce equivalent outputs for per-token timesteps using real checkpoint weights.
+    """
+    import torch
+
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.mlx_runtime.wan22 import mlx_wan22_dit_from_diffusers_safetensors
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+    from examples.inference.basic.mlx_wan_prompt_to_video import make_rotary_embeddings
+
+    config = json.loads(_CONFIG.read_text())
+    in_ch = int(config["in_channels"])
+    text_dim = int(config["text_dim"])
+    # Small latent so activations fit alongside the ~10 GB fp16 weights on 36 GB.
+    frames, height, width = 2, 16, 16
+    p_t, p_h, p_w = tuple(config["patch_size"])
+    tokens_per_frame = (height // p_h) * (width // p_w)
+    num_tokens = (frames // p_t) * tokens_per_frame
+
+    # Frame 0 clean (t=0), remaining tokens noised (t=900) — I2V-style expand_timesteps.
+    per_frame_levels = [0] + [900] * (frames // p_t - 1)
+    timestep_1d = [per_frame_levels[i // tokens_per_frame] for i in range(num_tokens)]
+    timestep = torch.tensor([timestep_1d], dtype=torch.long)
+
+    rng = np.random.default_rng(2026)
+    hidden_np = (rng.standard_normal((1, in_ch, frames, height, width)) * 0.5).astype(np.float32)
+    text_np = (rng.standard_normal((1, 32, text_dim)) * 0.1).astype(np.float32)
+
+    # --- torch reference (CPU fp16, same weight dtype as MLX deploy path) ---
+    torch_model = _load_torch_wan22_from_diffusers(_CHECKPOINT, _CONFIG, dtype=torch.float16)
+    hidden_t = torch.from_numpy(hidden_np).to(torch.float16)
+    text_t = torch.from_numpy(text_np).to(torch.float16)
+    with torch.no_grad(), set_forward_context(
+            current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")):
+        ref = torch_model(hidden_states=hidden_t, encoder_hidden_states=text_t, timestep=timestep)
+        ref_np = ref.detach().float().cpu().numpy()
+
+    # Free torch weights before loading MLX so ~10 GB fp16 + activations fit in 36 GB.
+    del torch_model
+    import gc
+    gc.collect()
+
+    # --- MLX (fp16 weights, fp16 compute) ---
+    mlx_model = mlx_wan22_dit_from_diffusers_safetensors(_CHECKPOINT, _CONFIG, dtype="fp16")
+    freqs_cis = make_rotary_embeddings(
+        config, latent_frames=frames, latent_height=height, latent_width=width)
+    out = mlx_model(
+        mx.array(hidden_np).astype(mx.float16),
+        mx.array(text_np).astype(mx.float16),
+        mx.array(np.array(timestep_1d, dtype=np.float32)[None, :]),
+        freqs_cis,
+    )
+    mx.eval(out)
+    mlx_np = np.array(out.astype(mx.float32))
+
+    assert mlx_np.shape == ref_np.shape
+    assert np.isfinite(mlx_np).all()
+    max_abs = float(np.abs(mlx_np - ref_np).max())
+    mean_abs = float(np.abs(mlx_np - ref_np).mean())
+    flat_a, flat_b = mlx_np.ravel(), ref_np.ravel()
+    cosine = float(np.dot(flat_a, flat_b) / (np.linalg.norm(flat_a) * np.linalg.norm(flat_b) + 1e-12))
+    print(f"Wan2.2-5B real-weight MLX-fp16 vs torch-fp16: "
+          f"max|Δ|={max_abs:.3e} mean|Δ|={mean_abs:.3e} cosine={cosine:.6f}")
+    # Full 30-layer fp16 Metal vs torch SDPA drifts more than the tiny-config
+    # fp32 gate (2e-3) / CUDA-tiny dump (5e-3). Measured on M4 Max: mean 6.1e-3,
+    # cosine 0.999948, both stable across MLX versions. The worst single element
+    # is not: it was 9e-2 when this budget was written and reaches 1.65e-1 under
+    # MLX 0.32's kernel fusion, so the aggregate statistics carry the assertion
+    # and max|Δ| stays as a loose guard against a genuine divergence.
+    assert mean_abs <= 1e-2, f"mean|Δ|={mean_abs} over budget"
+    assert cosine >= 0.9995, f"cosine={cosine} too low"
+    assert max_abs <= 0.25, f"max|Δ|={max_abs} over budget"

fastvideo/tests/mlx/test_mlx_wan_vae_denormalization.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan VAE latent denormalization regression tests."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from fastvideo.mlx_runtime.wan_vae import (
+    WanVAEConfigView,
+    denormalize_latents_np,
+)
+
+
+def test_denormalize_matches_prompt_to_video_formula() -> None:
+    """``latents / (1/std) + mean`` used by the hybrid script."""
+    cfg = WanVAEConfigView(
+        z_dim=2,
+        latents_mean=(0.1, -0.2),
+        latents_std=(2.0, 0.5),
+    )
+    lat = np.ones((1, 2, 1, 1, 1), dtype=np.float32)
+    out = denormalize_latents_np(lat, cfg)
+    # z * std + mean
+    np.testing.assert_allclose(out[0, 0, 0, 0, 0], 1.0 * 2.0 + 0.1)
+    np.testing.assert_allclose(out[0, 1, 0, 0, 0], 1.0 * 0.5 + (-0.2))
+
+
+def test_full_vae_config_and_denormalize_formula() -> None:
+    """Config load + denormalize formula vs the hybrid script's tensor math."""
+    import torch
+
+    snaps = list(
+        Path.home()
+        .joinpath(
+            ".cache/huggingface/hub/"
+            "models--FastVideo--FastWan2.2-TI2V-5B-FullAttn-Diffusers/"
+            "snapshots"
+        )
+        .glob("*/vae")
+    )
+    if not snaps:
+        pytest.skip("Wan2.2 VAE not in HF cache")
+    vae_dir = snaps[0]
+    cfg = WanVAEConfigView.from_vae_dir(vae_dir)
+    assert cfg.z_dim == 48
+    rng = np.random.default_rng(1)
+    lat = (rng.standard_normal((1, 48, 3, 8, 8)) * 0.3).astype(np.float32)
+    lat_dn = denormalize_latents_np(lat, cfg)
+    mean = torch.tensor(cfg.latents_mean, dtype=torch.float32).view(
+        1, -1, 1, 1, 1
+    )
+    inv_std = (1.0 / torch.tensor(cfg.latents_std, dtype=torch.float32)).view(
+        1, -1, 1, 1, 1
+    )
+    lat_t = torch.from_numpy(lat)
+    ref = (lat_t / inv_std + mean).numpy()
+    np.testing.assert_allclose(lat_dn, ref, atol=1e-5, rtol=1e-5)

fastvideo/tests/mlx/test_quant_backends.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Accuracy + capability probe for MLX block-scaled quant backends.
+
+For each backend in :mod:`fastvideo.mlx_runtime.quant_backends`, if the installed
+MLX build supports the mode, assert that quantized matmul stays within a
+per-format relative Frobenius error budget vs. the fp16 reference ``x @ w.T``.
+Unsupported modes are skipped with an explicit reason.
+
+Run with ``-s`` to print the summary table (backend | supported | bytes/weight |
+rel-error).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("mlx.core", reason="MLX is required for quant backend tests")
+
+import mlx.core as mx  # noqa: E402
+
+from fastvideo.mlx_runtime.quant_backends import (  # noqa: E402
+    BACKENDS,
+    bytes_per_weight,
+    is_supported,
+    quantize_weight,
+    quantized_matmul,
+    support_error,
+)
+
+# Fixed shapes: weight cols must be divisible by all group sizes (64/32/16).
+_WEIGHT_ROWS = 1024
+_WEIGHT_COLS = 1024
+_BATCH = 8
+_SEED = 0
+
+# Relative Frobenius error ||y - ref||_F / ||ref||_F tolerances.
+# Measured once on mlx 0.31.2 / Apple Silicon for N(0,1) fp16 weights:
+#   affine_int8_g64 ~ 5e-3, mxfp8 ~ 8e-2, mxfp4 ~ 1.2e-1, nvfp4 ~ 1.0e-1.
+# Budgets keep ~2x headroom for RNG / platform variance.
+_REL_ERROR_TOL: dict[str, float] = {
+    "affine_int8_g64": 2e-2,
+    "mxfp8": 1.5e-1,
+    "mxfp4": 2.5e-1,
+    "nvfp4": 2.5e-1,
+}
+
+
+def _frobenius_rel_error(y: mx.array, ref: mx.array) -> float:
+    """||y - ref||_F / ||ref||_F in float32 (avoids fp16 sum overflow)."""
+    y32 = y.astype(mx.float32)
+    ref32 = ref.astype(mx.float32)
+    diff = y32 - ref32
+    num = mx.sqrt(mx.sum(diff * diff))
+    den = mx.sqrt(mx.sum(ref32 * ref32)) + mx.array(1e-12, dtype=mx.float32)
+    mx.eval(num, den)
+    return float(num / den)
+
+
+def _make_inputs() -> tuple[mx.array, mx.array, mx.array]:
+    """
+    Generate deterministic random weights and inputs with their matrix multiplication reference.
+
+    Returns:
+        tuple[mx.array, mx.array, mx.array]: The weights, inputs, and fp16 reference output.
+    """
+    mx.random.seed(_SEED)
+    w = mx.random.normal((_WEIGHT_ROWS, _WEIGHT_COLS)).astype(mx.float16)
+    x = mx.random.normal((_BATCH, _WEIGHT_COLS)).astype(mx.float16)
+    ref = x @ w.T
+    mx.eval(w, x, ref)
+    return w, x, ref
+
+
+@pytest.mark.parametrize("backend", list(BACKENDS))
+def test_quant_backend_matmul_accuracy(backend: str) -> None:
+    if not is_supported(backend):
+        err = support_error(backend) or "unknown error"
+        pytest.skip(reason=f"{backend} not supported by installed MLX: {err}")
+
+    w, x, ref = _make_inputs()
+    qw = quantize_weight(w, backend)
+    y = quantized_matmul(x, qw)
+    mx.eval(y)
+
+    rel = _frobenius_rel_error(y, ref)
+    tol = _REL_ERROR_TOL[backend]
+    assert rel < tol, (
+        f"{backend}: relative Frobenius error {rel:.6e} exceeds tolerance {tol:.6e}"
+    )
+
+
+def test_quant_backends_summary_table(capsys: pytest.CaptureFixture[str]) -> None:
+    """
+    Print a support, storage, and relative-error summary for each quantization backend.
+
+    The table reports unsupported backends as unavailable and asserts that the affine
+    int8 baseline is supported.
+    """
+    w, x, ref = _make_inputs()
+    rows: list[tuple[str, str, str, str]] = []
+
+    for backend in BACKENDS:
+        supported = is_supported(backend)
+        if not supported:
+            rows.append((backend, "no", "n/a", "n/a"))
+            continue
+
+        bpw = bytes_per_weight(backend)
+        qw = quantize_weight(w, backend)
+        y = quantized_matmul(x, qw)
+        mx.eval(y)
+        rel = _frobenius_rel_error(y, ref)
+        rows.append((backend, "yes", f"{bpw:.6f}", f"{rel:.6e}"))
+
+    # Header + rows for human inspection under pytest -s.
+    col_w = max(len(r[0]) for r in rows)
+    header = f"{'backend':<{col_w}}  supported  bytes/weight  rel-error"
+    print("\n" + header)
+    print("-" * len(header))
+    for backend, supported, bpw, rel in rows:
+        print(f"{backend:<{col_w}}  {supported:<9}  {bpw:>12}  {rel}")
+
+    # At least the affine baseline must work on any supported MLX install.
+    assert is_supported("affine_int8_g64"), (
+        "affine_int8_g64 must be supported; cannot evaluate other backends"
+    )

fastvideo/tests/mlx/test_quant_capability.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MLX quantization-mode capability detection.
+
+The affine int8/int4 modes work on every MLX build the repo supports; the
+mxfp8/mxfp4/nvfp4 mode strings only exist on newer MLX builds. These tests pin
+the contract: supported modes probe clean, and unsupported ones surface as
+``UnsupportedMLXQuantizationError`` with an actionable message *before* any
+model weights are loaded.
+
+Runs on any MLX backend (Metal or ``mlx[cpu]``).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("mlx.core", reason="MLX is required for quantization capability tests")
+
+from fastvideo.mlx_runtime.fastwan import (  # noqa: E402
+    MLXQuantizationSpec,
+    UnsupportedMLXQuantizationError,
+    _QUANT_SUPPORT_CACHE,
+    ensure_quantization_supported,
+    quantization_support_error,
+)
+
+
+def test_affine_modes_are_supported_on_any_mlx_build() -> None:
+    for name in ("int8", "int4"):
+        spec = MLXQuantizationSpec.from_name(name)
+        assert quantization_support_error(spec) is None, f"{name} must be supported on every MLX build"
+        ensure_quantization_supported(spec)  # must not raise
+
+
+def test_none_spec_is_always_supported() -> None:
+    ensure_quantization_supported(None)  # must not raise
+
+
+def test_probe_result_is_cached_per_spec() -> None:
+    spec = MLXQuantizationSpec.from_name("int8")
+    quantization_support_error(spec)
+    key = (spec.mode, spec.bits, spec.group_size)
+    assert key in _QUANT_SUPPORT_CACHE
+
+
+def test_unsupported_mode_raises_with_actionable_message(monkeypatch) -> None:
+    spec = MLXQuantizationSpec.from_name("nvfp4")
+    key = (spec.mode, spec.bits, spec.group_size)
+    monkeypatch.setitem(_QUANT_SUPPORT_CACHE, key, "ValueError: [quantize] Unknown mode 'nvfp4'")
+
+    with pytest.raises(UnsupportedMLXQuantizationError) as excinfo:
+        ensure_quantization_supported(spec)
+    message = str(excinfo.value)
+    assert "nvfp4" in message
+    assert "int8" in message  # points the user at the reliable fallback
+
+
+def test_loader_rejects_unsupported_mode_before_reading_weights(monkeypatch, tmp_path) -> None:
+    """Verify that unsupported quantization modes are rejected before checkpoint weights are read."""
+    from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
+
+    spec_key = ("mxfp4", None, None)
+    monkeypatch.setitem(_QUANT_SUPPORT_CACHE, spec_key, "ValueError: [quantize] Unknown mode 'mxfp4'")
+
+    config_path = tmp_path / "config.json"
+    config_path.write_text('{"num_layers": 1, "num_attention_heads": 4, "attention_head_dim": 16, '
+                           '"ffn_dim": 128, "in_channels": 16, "out_channels": 16, '
+                           '"patch_size": [1, 2, 2], "freq_dim": 64, "eps": 1e-6}')
+    # The checkpoint path deliberately does not exist: the capability check
+    # must fire before the loader ever opens the safetensors file.
+    with pytest.raises(UnsupportedMLXQuantizationError):
+        mlx_dit_from_diffusers_safetensors(tmp_path / "missing.safetensors", config_path, quantization="mxfp4")

fastvideo/tests/mlx/test_taehv_decode.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Offline contracts for the vendored TAEHV decoder helper."""
+
+import re
+
+import pytest
+
+from fastvideo.mlx_runtime import wan_vae
+from fastvideo.mlx_runtime.taehv_decode import TAEW2_1_CHECKPOINT_SHA256, ensure_taew2_1_checkpoint
+
+
+def test_taehv_checkpoint_pin_is_a_sha256() -> None:
+    assert len(TAEW2_1_CHECKPOINT_SHA256) == 64
+    assert int(TAEW2_1_CHECKPOINT_SHA256, 16) >= 0
+
+
+def test_explicit_taehv_checkpoint_path_is_not_downloaded(tmp_path) -> None:
+    checkpoint = tmp_path / "custom-taehv.pth"
+    checkpoint.write_bytes(b"local test checkpoint")
+    assert ensure_taew2_1_checkpoint(checkpoint) == checkpoint
+
+
+def test_taew2_2_checkpoint_pin_is_immutable_and_sha256() -> None:
+    assert "/main/" not in wan_vae.TAEW2_2_URL
+    assert re.search(r"/[0-9a-f]{40}/taew2_2\.pth$", wan_vae.TAEW2_2_URL)
+    assert re.fullmatch(r"[0-9a-f]{64}", wan_vae.TAEW2_2_SHA256)
+
+
+def test_cached_taew2_2_checkpoint_is_verified(monkeypatch, tmp_path) -> None:
+    monkeypatch.setattr(wan_vae, "_cache_dir", lambda: tmp_path)
+    (tmp_path / "taew2_2.pth").write_bytes(b"tampered")
+
+    with pytest.raises(RuntimeError, match="failed sha256 verification"):
+        wan_vae.ensure_taehv_checkpoint(z_dim=48)

fastvideo/tests/mlx/test_wan22_sample.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for Wan2.2 warped DMD schedule construction."""
+
+from __future__ import annotations
+
+import torch
+
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule
+
+
+def test_warped_dmd_timesteps_differ_from_raw_indices() -> None:
+    schedule, warped = build_wan22_dmd_schedule([1000, 757, 522], flow_shift=5.0, warp_denoising_step=True)
+    _, raw = build_wan22_dmd_schedule([1000, 757, 522], flow_shift=5.0, warp_denoising_step=False)
+    assert raw == [1000.0, 757.0, 522.0]
+    # Warping maps step indices into the continuous flow-match schedule.
+    assert warped[0] == 1000.0
+    assert warped[1] != 757.0
+    assert warped[2] != 522.0
+    # Sigmas are monotone-ish decreasing along the schedule.
+    sigmas = [schedule.sigma_for(t) for t in warped]
+    assert sigmas[0] >= sigmas[1] >= sigmas[2]
+
+
+def test_warped_dmd_timesteps_match_training_scheduler() -> None:
+    steps = torch.tensor([1000, 757, 522], dtype=torch.long)
+    scheduler = FlowMatchEulerDiscreteScheduler(shift=5.0)
+    expected = torch.cat((scheduler.timesteps, torch.tensor([0.0])))[1000 - steps]
+
+    _, warped = build_wan22_dmd_schedule(steps.tolist(), flow_shift=5.0, warp_denoising_step=True)
+
+    torch.testing.assert_close(torch.tensor(warped), expected)

fastvideo/tests/mlx/test_windowed_attention.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Correctness checks for chunked non-causal windowed attention (MLX)."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+mx = pytest.importorskip("mlx.core", reason="MLX is required for windowed-attention tests")
+
+from fastvideo.mlx_runtime.windowed_attention import (  # noqa: E402
+    _key_ranges_for_block,
+    full_attention,
+    windowed_attention,
+)
+
+
+def _rand_qkv(b: int, h: int, s: int, d: int, seed: int = 0):
+    """
+    Generate deterministic random query, key, and value tensors.
+
+    Parameters:
+        b (int): Batch size.
+        h (int): Number of attention heads.
+        s (int): Sequence length.
+        d (int): Feature dimension.
+        seed (int): Random number generator seed.
+
+    Returns:
+        tuple: Query, key, and value tensors with shape `(b, h, s, d)`.
+    """
+    rng = np.random.default_rng(seed)
+    q = mx.array(rng.standard_normal((b, h, s, d)).astype(np.float32))
+    k = mx.array(rng.standard_normal((b, h, s, d)).astype(np.float32))
+    v = mx.array(rng.standard_normal((b, h, s, d)).astype(np.float32))
+    return q, k, v
+
+
+def _mean_cosine(a: "mx.array", b: "mx.array") -> float:
+    """
+    Compute the mean cosine similarity between corresponding feature vectors.
+
+    Parameters:
+        a: The first tensor.
+        b: The second tensor.
+
+    Returns:
+        The mean cosine similarity, with protection against zero-norm vectors.
+    """
+    a32 = np.array(a.astype(mx.float32)).reshape(-1, a.shape[-1])
+    b32 = np.array(b.astype(mx.float32)).reshape(-1, b.shape[-1])
+    num = np.sum(a32 * b32, axis=-1)
+    den = np.linalg.norm(a32, axis=-1) * np.linalg.norm(b32, axis=-1)
+    return float(np.mean(num / np.maximum(den, 1e-12)))
+
+
+def test_windowed_matches_full_when_window_ge_seq() -> None:
+    """``window >= S`` must recover full attention (cosine > 0.999)."""
+    b, h, s, d = 1, 2, 64, 16
+    q, k, v = _rand_qkv(b, h, s, d, seed=7)
+    scale = d**-0.5
+
+    ref = full_attention(q, k, v, scale=scale)
+    out = windowed_attention(q, k, v, window=s, sink=0, scale=scale)
+    mx.eval(ref, out)
+
+    assert out.shape == q.shape
+    cos = _mean_cosine(ref, out)
+    assert cos > 0.999, f"expected cosine > 0.999 for window>=S, got {cos}"
+
+
+def test_windowed_output_shape_equals_input() -> None:
+    b, h, s, d = 2, 3, 48, 8
+    q, k, v = _rand_qkv(b, h, s, d, seed=11)
+    out = windowed_attention(q, k, v, window=16, sink=4)
+    mx.eval(out)
+    assert out.shape == q.shape
+    assert out.shape == (b, h, s, d)
+
+
+def test_windowed_with_sinks_shape_and_finite() -> None:
+    b, h, s, d = 1, 2, 40, 8
+    q, k, v = _rand_qkv(b, h, s, d, seed=19)
+    out = windowed_attention(q, k, v, window=8, sink=4, chunk_size=5)
+    mx.eval(out)
+    assert out.shape == (b, h, s, d)
+    arr = np.array(out)
+    assert np.isfinite(arr).all()
+
+
+def test_windowed_result_does_not_depend_on_chunk_size() -> None:
+    q, k, v = _rand_qkv(1, 2, 40, 8, seed=23)
+    per_query = windowed_attention(q, k, v, window=8, sink=4, chunk_size=1)
+    chunked = windowed_attention(q, k, v, window=8, sink=4, chunk_size=5)
+    mx.eval(per_query, chunked)
+
+    np.testing.assert_allclose(np.array(chunked), np.array(per_query), rtol=1e-5, atol=1e-5)
+
+
+def test_sink_range_extends_beyond_local_window() -> None:
+    assert _key_ranges_for_block(qs=0, qe=1, seq_len=12, half=1, sink=5) == [(0, 5)]
+
+
+def test_windowed_attention_enforces_per_query_local_and_sink_mask() -> None:
+    q = mx.zeros((1, 1, 6, 1))
+    k = mx.zeros_like(q)
+    v = mx.arange(6).reshape(1, 1, 6, 1)
+
+    out = windowed_attention(q, k, v, window=2, sink=3, chunk_size=4)
+    mx.eval(out)
+
+    np.testing.assert_allclose(np.array(out).reshape(-1), [1.0, 1.0, 1.5, 2.0, 2.5, 2.4], atol=1e-6)
+
+
+def test_invalid_window_raises() -> None:
+    q, k, v = _rand_qkv(1, 1, 8, 4, seed=1)
+    with pytest.raises(ValueError, match="window"):
+        windowed_attention(q, k, v, window=0)

fastvideo/tests/mlx/tiny_wan.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tiny random-weight Wan DiT fixtures shared by the MLX runtime tests.
+
+Builds a miniature ``WanTransformer3DModel`` (all matmul dims are multiples of
+the int8 group size 64, so the same model exercises the quantized paths) plus
+the conversions into ``fastvideo.mlx_runtime.fastwan.MLXWanDiT``.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29513")
+
+from fastvideo.configs.models.dits.wanvideo import (  # noqa: E402
+    WanVideoArchConfig,
+    WanVideoConfig,
+)
+from fastvideo.forward_context import set_forward_context  # noqa: E402
+from fastvideo.layers.rotary_embedding import get_rotary_pos_embed  # noqa: E402
+from fastvideo.mlx_runtime.fastwan import (  # noqa: E402
+    MLXQuantizationSpec,
+    MLXWanDiT,
+    MLXWanTransformerBlock,
+    mlx_block_weights_from_torch,
+    quantize_matrix,
+)
+from fastvideo.models.dits.wanvideo import WanTransformer3DModel  # noqa: E402
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch  # noqa: E402
+
+SEED = 2026
+
+# All matmul dims must be multiples of the int8 group size (64) so the
+# quantized variants of the tests can reuse the same tiny model.
+TINY_ARCH = dict(
+    num_attention_heads=4,
+    attention_head_dim=16,
+    in_channels=16,
+    out_channels=16,
+    text_dim=64,
+    freq_dim=64,
+    ffn_dim=128,
+    num_layers=2,
+    patch_size=(1, 2, 2),
+    rope_max_seq_len=64,
+)
+
+
+def build_tiny_wan_config() -> WanVideoConfig:
+    """Build the deterministic miniature Wan video configuration.
+
+    Returns:
+        WanVideoConfig: Configuration created from the tiny test architecture.
+    """
+    return WanVideoConfig(arch_config=WanVideoArchConfig(**TINY_ARCH))
+
+
+def build_hf_config(config: WanVideoConfig) -> dict[str, object]:
+    """
+    Builds the configuration dictionary required by the Hugging Face-style Wan model representation.
+
+    Parameters:
+        config (WanVideoConfig): Wan video configuration to convert.
+
+    Returns:
+        dict[str, object]: Configuration values for the Hugging Face-style model.
+    """
+    return {
+        "num_attention_heads": config.num_attention_heads,
+        "attention_head_dim": config.attention_head_dim,
+        "in_channels": config.in_channels,
+        "out_channels": config.out_channels,
+        "text_dim": config.text_dim,
+        "freq_dim": config.freq_dim,
+        "ffn_dim": config.ffn_dim,
+        "num_layers": config.num_layers,
+        "patch_size": config.patch_size,
+        "text_len": config.text_len,
+        "rope_max_seq_len": config.rope_max_seq_len,
+        "eps": 1e-6,
+    }
+
+
+def initialize_model_parameters(model: torch.nn.Module) -> None:
+    # ReplicatedLinear parameters are allocated with torch.empty and need an
+    # explicit initialization in tests to avoid undefined values.
+    """Initialize model parameters deterministically for reproducible tests.
+
+    Parameters:
+        model (torch.nn.Module): Model whose parameters will be initialized.
+    """
+    torch.manual_seed(SEED + 3)
+    with torch.no_grad():
+        for name, param in model.named_parameters():
+            if param.ndim <= 1:
+                if name.endswith("weight") and "norm" in name:
+                    param.fill_(1.0)
+                else:
+                    param.normal_(mean=0.0, std=0.02)
+                continue
+            torch.nn.init.xavier_uniform_(param)
+
+
+def build_torch_model() -> WanTransformer3DModel:
+    """
+    Build a deterministically initialized CPU float32 Wan transformer model.
+
+    Returns:
+        WanTransformer3DModel: The initialized model in evaluation mode.
+    """
+    config = build_tiny_wan_config()
+    model = WanTransformer3DModel(config=config, hf_config=build_hf_config(config))
+    model = model.to(device="cpu", dtype=torch.float32)
+    initialize_model_parameters(model)
+    model.eval()
+    return model
+
+
+def build_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Create deterministic latent, text-encoder, and timestep inputs for the tiny Wan model.
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]: The latent hidden states, text-encoder hidden states, and timestep."""
+    generator = torch.Generator(device="cpu").manual_seed(SEED + 1)
+    hidden_states = torch.randn(1, TINY_ARCH["in_channels"], 4, 8, 8, generator=generator, dtype=torch.float32)
+    encoder_hidden_states = torch.randn(1, 8, TINY_ARCH["text_dim"], generator=generator, dtype=torch.float32)
+    timestep = torch.tensor([10], dtype=torch.long)
+    return hidden_states, encoder_hidden_states, timestep
+
+
+# Top-level weights: the MLX runtime uses the Diffusers key layout; the torch
+# model uses FastVideo's module names (see WanVideoConfig.param_names_mapping).
+TOP_LEVEL_KEY_MAP = {
+    "patch_embedding.weight": "patch_embedding.proj.weight",
+    "patch_embedding.bias": "patch_embedding.proj.bias",
+    "condition_embedder.time_embedder.linear_1.weight": "condition_embedder.time_embedder.mlp.fc_in.weight",
+    "condition_embedder.time_embedder.linear_1.bias": "condition_embedder.time_embedder.mlp.fc_in.bias",
+    "condition_embedder.time_embedder.linear_2.weight": "condition_embedder.time_embedder.mlp.fc_out.weight",
+    "condition_embedder.time_embedder.linear_2.bias": "condition_embedder.time_embedder.mlp.fc_out.bias",
+    "condition_embedder.time_proj.weight": "condition_embedder.time_modulation.linear.weight",
+    "condition_embedder.time_proj.bias": "condition_embedder.time_modulation.linear.bias",
+    "condition_embedder.text_embedder.linear_1.weight": "condition_embedder.text_embedder.fc_in.weight",
+    "condition_embedder.text_embedder.linear_1.bias": "condition_embedder.text_embedder.fc_in.bias",
+    "condition_embedder.text_embedder.linear_2.weight": "condition_embedder.text_embedder.fc_out.weight",
+    "condition_embedder.text_embedder.linear_2.bias": "condition_embedder.text_embedder.fc_out.bias",
+    "scale_shift_table": "scale_shift_table",
+    "proj_out.weight": "proj_out.weight",
+    "proj_out.bias": "proj_out.bias",
+}
+
+
+def mlx_dit_from_torch_model(
+    model: WanTransformer3DModel,
+    hf_config: dict[str, object],
+    *,
+    quantization: MLXQuantizationSpec | None = None,
+) -> MLXWanDiT:
+    """
+    Convert a PyTorch Wan DiT model to an MLX Wan DiT model.
+
+    Parameters:
+        model (WanTransformer3DModel): PyTorch model whose parameters are converted.
+        hf_config (dict[str, object]): Model configuration used to construct the MLX representation.
+        quantization (MLXQuantizationSpec | None): Optional quantization specification for eligible weight matrices.
+
+    Returns:
+        MLXWanDiT: The converted MLX model.
+    """
+    import mlx.core as mx
+
+    state = {name: value.detach().float() for name, value in model.state_dict().items()}
+    inner_dim = int(hf_config["num_attention_heads"]) * int(hf_config["attention_head_dim"])  # type: ignore[arg-type]
+
+    weights = {}
+    for mlx_name, torch_name in TOP_LEVEL_KEY_MAP.items():
+        tensor = state[torch_name]
+        if mlx_name == "patch_embedding.weight":
+            tensor = tensor.reshape(inner_dim, -1)
+        array = mx.array(tensor.numpy())
+        if quantization is not None and mlx_name.endswith(".weight") and mlx_name != "scale_shift_table":
+            weights[mlx_name] = quantize_matrix(array, quantization)
+        else:
+            weights[mlx_name] = array
+
+    blocks = []
+    for torch_block in model.blocks:
+        block_weights = mlx_block_weights_from_torch(torch_block)
+        if quantization is not None:
+            block_weights = {
+                name: (quantize_matrix(value, quantization)
+                       if name.endswith(".weight") and "norm" not in name and len(value.shape) >= 2 else value)
+                for name, value in block_weights.items()
+            }
+        blocks.append(
+            MLXWanTransformerBlock(
+                block_weights,
+                dim=inner_dim,
+                ffn_dim=int(hf_config["ffn_dim"]),  # type: ignore[arg-type]
+                num_heads=int(hf_config["num_attention_heads"]),  # type: ignore[arg-type]
+                eps=float(hf_config["eps"]),  # type: ignore[arg-type]
+            ))
+    return MLXWanDiT(weights, blocks, dict(hf_config))
+
+
+def mlx_rotary_embeddings(hidden_states: torch.Tensor):
+    """
+    Generate MLX rotary-position embedding tables matching the tiny Wan model.
+
+    Parameters:
+        hidden_states (torch.Tensor): Input latent tensor whose spatial and temporal dimensions determine the embedding grid.
+
+    Returns:
+        Tuple of MLX arrays containing the cosine and sine rotary-position embeddings.
+    """
+    import mlx.core as mx
+
+    _, _, frames, height, width = hidden_states.shape
+    p_t, p_h, p_w = TINY_ARCH["patch_size"]
+    head_dim = TINY_ARCH["attention_head_dim"]
+    hidden_size = TINY_ARCH["num_attention_heads"] * head_dim
+    rope_dim_list = [head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)]
+    freqs_cos, freqs_sin = get_rotary_pos_embed(
+        (frames // p_t, height // p_h, width // p_w),
+        hidden_size,
+        TINY_ARCH["num_attention_heads"],
+        rope_dim_list,
+        dtype=torch.float64,
+        rope_theta=10000,
+    )
+    return (
+        mx.array(freqs_cos.float().numpy()).astype(mx.float32),
+        mx.array(freqs_sin.float().numpy()).astype(mx.float32),
+    )
+
+
+def torch_reference_output(model, hidden_states, encoder_hidden_states, timestep) -> np.ndarray:
+    """
+    Run the PyTorch model on the provided inputs and return its output as a NumPy array.
+
+    Parameters:
+        model: The PyTorch Wan transformer model.
+        hidden_states: The latent hidden states provided to the model.
+        encoder_hidden_states: The encoder hidden states provided to the model.
+        timestep: The diffusion timestep provided to the model.
+
+    Returns:
+        The model output as a CPU float32 NumPy array.
+    """
+    with torch.no_grad(), set_forward_context(
+            current_timestep=0,
+            attn_metadata=None,
+            forward_batch=ForwardBatch(data_type="dummy"),
+    ):
+        output = model(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+        )
+    return output.detach().float().cpu().numpy()
+
+
+def mlx_output(dit, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> np.ndarray:
+    """
+    Run the MLX DiT model on the provided inputs.
+
+    Parameters:
+        dit: The MLX DiT model to evaluate.
+        hidden_states: Input latent states.
+        encoder_hidden_states: Text encoder states.
+        timestep: Diffusion timestep input.
+        freqs_cis: Rotary-position embedding frequencies.
+
+    Returns:
+        np.ndarray: The model output as a float32 NumPy array.
+    """
+    import mlx.core as mx
+
+    out = dit(
+        mx.array(hidden_states.numpy()),
+        mx.array(encoder_hidden_states.numpy()),
+        mx.array(timestep.float().numpy()),
+        freqs_cis,
+    )
+    mx.eval(out)
+    return np.array(out.astype(mx.float32))

fastvideo/tests/platforms/test_mps_vsa_error.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MPS must reject VSA before a model reaches an incompatible SDPA call."""
+
+import pytest
+import torch
+
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.platforms.mps import MpsPlatform
+
+
+def test_mps_rejects_video_sparse_attention_with_actionable_error() -> None:
+    with pytest.raises(NotImplementedError, match="TORCH_SDPA"):
+        MpsPlatform.get_attn_backend_cls(AttentionBackendEnum.VIDEO_SPARSE_ATTN, 64, torch.float16)
+
+
+def test_mps_resolves_sdpa_for_supported_backend() -> None:
+    assert MpsPlatform.get_attn_backend_cls(AttentionBackendEnum.TORCH_SDPA, 64, torch.float16).endswith(
+        "SDPABackend"
+    )

fastvideo/third_party/rife_mlx/LICENSE

@@ -0,0 +1,26 @@
+MIT License
+
+This directory vendors the inference path of rife-mlx
+(https://github.com/xocialize/rife-mlx), an Apple MLX port of Practical-RIFE
+(https://github.com/hzwer/Practical-RIFE).
+
+Copyright (c) 2025 xocialize (rife-mlx, MLX port)
+Copyright (c) hzwer (Practical-RIFE / RIFE, original work)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

fastvideo/third_party/rife_mlx/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: MIT
+"""Vendored rife-mlx (RIFE frame interpolation on Apple MLX).
+
+Upstream: https://github.com/xocialize/rife-mlx, itself an MLX port of
+Practical-RIFE (https://github.com/hzwer/Practical-RIFE, (c) hzwer). Both are
+MIT; see ``LICENSE`` for the text.
+
+Only the inference path is vendored, which is what ``--fast`` calls:
+``utils.weights.build_model`` and ``pipeline_mlx.interpolate_pair``. Upstream's
+video I/O and torch-to-MLX conversion helpers are omitted, so this package
+needs no ``av``, ``PIL``, or torch import to load. Weights are still fetched at
+runtime from the ``mlx-community`` Hugging Face org, exactly as upstream does.
+"""
+
+from fastvideo.third_party.rife_mlx.config import DEFAULT_VERSION, VERSIONS, VersionConfig
+
+__all__ = ["DEFAULT_VERSION", "VERSIONS", "VersionConfig"]

fastvideo/third_party/rife_mlx/config.py

@@ -0,0 +1,45 @@
+"""Version registry — pinned config for the RIFE checkpoint(s) we port.
+
+Source of truth: the bundled `IFNet_HDv3.py` + `flownet.pkl` that ship together
+per version (Google Drive), NOT the GitHub repo (which omits the version arch).
+CONFIRM gate #3: pin the exact arch from the downloaded package, never guess.
+
+This round pins **RIFE 4.25** (upstream's recommended default; "anime scenes
+significantly improved"). Arch PINNED in P1 from the bundled IFNet_HDv3.py +
+flownet.pkl shapes: 5 IFBlocks (the 4.25 "more flow blocks" note), channels
+[192,128,96,64,32], in_planes [15,28,28,28,28] (blocks 1-4 add the running
+4-ch flow inside IFBlock.forward), scale_list [16,8,4,2,1] (÷ the --scale knob).
+Activation is LeakyReLU(0.2). All F.interpolate use align_corners=False; warp
+grid_sample uses align_corners=True. 40 teacher/caltime keys are train-only.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class VersionConfig:
+    name: str               # internal key
+    hf_name: str            # published HF repo name (mlx-community/<hf_name>)
+    gdrive_id: str          # per-version package (model .py + flownet.pkl)
+    # arch params — PINNED from bundled v4.25 IFNet_HDv3.py + flownet.pkl (P1)
+    scale_list: tuple[float, ...] = (16, 8, 4, 2, 1)
+    block_channels: tuple[int, ...] = (192, 128, 96, 64, 32)
+    block_in_planes: tuple[int, ...] = (15, 28, 28, 28, 28)
+    pad_to: int = 64  # 5 blocks downsample; pad H,W to a multiple of 64
+    # warp / interpolate semantics (verified from source)
+    grid_align_corners: bool = True      # warp grid_sample
+    grid_padding_mode: str = "border"
+    interp_align_corners: bool = False   # F.interpolate in IFBlock
+
+
+VERSIONS: dict[str, VersionConfig] = {
+    "4.25": VersionConfig(
+        name="4.25",
+        hf_name="RIFE-4.25",
+        gdrive_id="1ZKjcbmt1hypiFprJPIKW0Tt0lr_2i7bg",
+    ),
+}
+
+DEFAULT_VERSION = "4.25"

fastvideo/third_party/rife_mlx/model/IFNet_HDv3.py

@@ -0,0 +1,142 @@
+"""IFNet (RIFE 4.25) — MLX port, isomorphic to the bundled train_log/IFNet_HDv3.py.
+
+NHWC throughout. Translated line-by-line from the pinned v4.25 source; only
+PyTorch ops -> MLX ops and NCHW -> NHWC change. Parity-locked in P3 (Gate B).
+
+Arch (pinned): 5 IFBlocks c=[192,128,96,64,32], in_planes=[15,28,28,28,28];
+conv = Conv2d + LeakyReLU(0.2); ResConv = conv*beta + x -> LeakyReLU; lastconv =
+ConvTranspose2d(c,52,4,2,1) + PixelShuffle(2) -> 13ch (flow4+mask1+feat8); Head
+encoder -> 4ch; F.interpolate align_corners=False; warp align_corners=True.
+
+MLX param keys (post-conversion) match torch except: conv0.{j}.0.* -> conv0.{j}.conv.*
+and lastconv.0.* -> lastconv.* (see utils/convert.py).
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..ops.interpolate import interpolate_bilinear
+from ..ops.spatial import pixel_shuffle_nhwc
+from .warplayer import warp
+
+_SLOPE = 0.2
+
+
+def _lrelu(x):
+    return nn.leaky_relu(x, _SLOPE)
+
+
+class Conv(nn.Module):
+    """conv = Conv2d + LeakyReLU(0.2). key: <p>.conv.weight"""
+
+    def __init__(self, i, o, k=3, s=1, p=1):
+        super().__init__()
+        self.conv = nn.Conv2d(i, o, kernel_size=k, stride=s, padding=p)
+
+    def __call__(self, x):
+        return _lrelu(self.conv(x))
+
+
+class ResConv(nn.Module):
+    """conv*beta + x -> LeakyReLU. keys: .conv.weight/bias, .beta"""
+
+    def __init__(self, c):
+        super().__init__()
+        self.conv = nn.Conv2d(c, c, kernel_size=3, stride=1, padding=1)
+        self.beta = mx.ones((1, 1, 1, c))  # NHWC per-channel
+
+    def __call__(self, x):
+        return _lrelu(self.conv(x) * self.beta + x)
+
+
+class Head(nn.Module):
+    """Feature encoder. keys: cnn0,cnn1,cnn2 (Conv2d), cnn3 (ConvTranspose2d)."""
+
+    def __init__(self):
+        super().__init__()
+        self.cnn0 = nn.Conv2d(3, 16, 3, stride=2, padding=1)
+        self.cnn1 = nn.Conv2d(16, 16, 3, stride=1, padding=1)
+        self.cnn2 = nn.Conv2d(16, 16, 3, stride=1, padding=1)
+        self.cnn3 = nn.ConvTranspose2d(16, 4, 4, stride=2, padding=1)
+
+    def __call__(self, x):
+        x = _lrelu(self.cnn0(x))
+        x = _lrelu(self.cnn1(x))
+        x = _lrelu(self.cnn2(x))
+        return self.cnn3(x)
+
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c=64):
+        super().__init__()
+        self.conv0 = [Conv(in_planes, c // 2, 3, 2, 1), Conv(c // 2, c, 3, 2, 1)]
+        self.convblock = [ResConv(c) for _ in range(8)]
+        self.lastconv = nn.ConvTranspose2d(c, 4 * 13, 4, stride=2, padding=1)
+
+    def __call__(self, x, flow=None, scale=1):
+        x = interpolate_bilinear(x, scale_factor=1.0 / scale, align_corners=False)
+        if flow is not None:
+            flow = interpolate_bilinear(flow, scale_factor=1.0 / scale,
+                                        align_corners=False) * (1.0 / scale)
+            x = mx.concatenate([x, flow], axis=-1)
+        feat = self.conv0[0](x)
+        feat = self.conv0[1](feat)
+        for blk in self.convblock:
+            feat = blk(feat)
+        tmp = pixel_shuffle_nhwc(self.lastconv(feat), 2)
+        tmp = interpolate_bilinear(tmp, scale_factor=scale, align_corners=False)
+        flow = tmp[..., :4] * scale
+        mask = tmp[..., 4:5]
+        feat = tmp[..., 5:]
+        return flow, mask, feat
+
+
+class IFNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block0 = IFBlock(7 + 8, c=192)
+        self.block1 = IFBlock(8 + 4 + 8 + 8, c=128)
+        self.block2 = IFBlock(8 + 4 + 8 + 8, c=96)
+        self.block3 = IFBlock(8 + 4 + 8 + 8, c=64)
+        self.block4 = IFBlock(8 + 4 + 8 + 8, c=32)
+        self.encode = Head()
+
+    def __call__(self, x, timestep=0.5, scale_list=(16, 8, 4, 2, 1)):
+        """x: [N,H,W,6] (img0|img1 on channel axis). Returns (flow_list, mask, merged)."""
+        channel = x.shape[-1] // 2
+        img0 = x[..., :channel]
+        img1 = x[..., channel:]
+        N, H, W, _ = img0.shape
+        timestep = mx.full((N, H, W, 1), float(timestep))
+
+        f0 = self.encode(img0[..., :3])
+        f1 = self.encode(img1[..., :3])
+        flow_list, merged, mask_list = [], [], []
+        warped_img0, warped_img1 = img0, img1
+        flow = None
+        mask = None
+        blocks = [self.block0, self.block1, self.block2, self.block3, self.block4]
+        for i in range(5):
+            if flow is None:
+                flow, mask, feat = blocks[i](
+                    mx.concatenate([img0[..., :3], img1[..., :3], f0, f1, timestep], axis=-1),
+                    None, scale=scale_list[i])
+            else:
+                wf0 = warp(f0, flow[..., :2])
+                wf1 = warp(f1, flow[..., 2:4])
+                fd, m0, feat = blocks[i](
+                    mx.concatenate([warped_img0[..., :3], warped_img1[..., :3],
+                                    wf0, wf1, timestep, mask, feat], axis=-1),
+                    flow, scale=scale_list[i])
+                mask = m0
+                flow = flow + fd
+            mask_list.append(mask)
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[..., :2])
+            warped_img1 = warp(img1, flow[..., 2:4])
+            merged.append((warped_img0, warped_img1))
+        mask = mx.sigmoid(mask)
+        merged_final = warped_img0 * mask + warped_img1 * (1 - mask)
+        return flow_list, mask, merged_final

fastvideo/third_party/rife_mlx/model/RIFE_HDv3.py

@@ -0,0 +1,37 @@
+"""Model wrapper — MLX port of train_log/RIFE_HDv3.py.
+
+Wraps IFNet and exposes inference(img0, img1, timestep, scale) -> middle frame,
+matching upstream: scale_list = [16,8,4,2,1] / scale; H,W padded to a multiple of
+pad_to (64; the 5-block downsample factor) and cropped back after.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+from .IFNet_HDv3 import IFNet
+
+
+class Model:
+    def __init__(self, scale_list=(16, 8, 4, 2, 1), pad_to: int = 64) -> None:
+        self.flownet = IFNet()
+        self.scale_list = tuple(scale_list)
+        self.pad_to = pad_to
+
+    def inference(self, img0: mx.array, img1: mx.array,
+                  timestep: float = 0.5, scale: float = 1.0) -> mx.array:
+        """img0,img1: [N,H,W,3] in [0,1]. Returns the frame at `timestep`."""
+        N, H, W, _ = img0.shape
+        # pad must scale with 1/scale so the coarsest-scale interpolation
+        # round-trips exactly (else the recovered flow size != padded size).
+        pad = max(self.pad_to, int(round(self.pad_to / scale)))
+        ph = ((H - 1) // pad + 1) * pad
+        pw = ((W - 1) // pad + 1) * pad
+        pad = [(0, 0), (0, ph - H), (0, pw - W), (0, 0)]
+        i0 = mx.pad(img0, pad)
+        i1 = mx.pad(img1, pad)
+
+        scale_list = tuple(s / scale for s in self.scale_list)
+        x = mx.concatenate([i0, i1], axis=-1)
+        _flow, _mask, merged = self.flownet(x, timestep, scale_list)
+        return merged[:, :H, :W, :]

fastvideo/third_party/rife_mlx/model/__init__.py

@@ -0,0 +1,5 @@
+"""IFNet / RIFE model modules — isomorphic to upstream train_log/."""
+
+from .warplayer import warp
+
+__all__ = ["warp"]

fastvideo/third_party/rife_mlx/model/warplayer.py

@@ -0,0 +1,37 @@
+"""warp — MLX port of model/warplayer.py (backward warp via grid_sample).
+
+Isomorphic to upstream: builds a normalized identity grid, adds the
+per-pixel flow (normalized by (size-1)/2), and bilinear grid_samples the
+input with border padding + align_corners=True.
+
+NHWC throughout. flow is [N,H,W,2] with channel 0 = x (horizontal), 1 = y.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+from ..ops.grid_sample import grid_sample_bilinear
+
+_grid_cache: dict[tuple, mx.array] = {}
+
+
+def _identity_grid(N: int, H: int, W: int) -> mx.array:
+    key = (N, H, W)
+    if key not in _grid_cache:
+        xs = mx.linspace(-1.0, 1.0, W).reshape(1, 1, W, 1)
+        ys = mx.linspace(-1.0, 1.0, H).reshape(1, 1, H, 1)  # placeholder shape
+        xs = mx.broadcast_to(xs, (N, H, W, 1))
+        ys = mx.broadcast_to(mx.linspace(-1.0, 1.0, H).reshape(1, H, 1, 1), (N, H, W, 1))
+        _grid_cache[key] = mx.concatenate([xs, ys], axis=-1)  # [N,H,W,2]
+    return _grid_cache[key]
+
+
+def warp(tenInput: mx.array, tenFlow: mx.array) -> mx.array:
+    """tenInput: [N,H,W,C]; tenFlow: [N,H,W,2] (x,y pixel-space flow)."""
+    N, H, W, _ = tenFlow.shape
+    grid = _identity_grid(N, H, W)
+    fx = tenFlow[..., 0:1] / ((W - 1.0) / 2.0)
+    fy = tenFlow[..., 1:2] / ((H - 1.0) / 2.0)
+    g = grid + mx.concatenate([fx, fy], axis=-1)
+    return grid_sample_bilinear(tenInput, g, align_corners=True, padding_mode="border")

fastvideo/third_party/rife_mlx/ops/__init__.py

@@ -0,0 +1,6 @@
+"""MLX ops hand-rolled for RIFE (no native equivalents)."""
+
+from .grid_sample import grid_sample_bilinear
+from .interpolate import interpolate_bilinear
+
+__all__ = ["grid_sample_bilinear", "interpolate_bilinear"]

fastvideo/third_party/rife_mlx/ops/grid_sample.py

@@ -0,0 +1,60 @@
+"""Bilinear grid_sample (NHWC) — MLX equivalent of torch.nn.functional.grid_sample.
+
+RIFE's `warp` calls grid_sample(mode='bilinear', padding_mode='border',
+align_corners=True). No native MLX op — hand-rolled here. This is the crux
+surface of the port; parity-locked vs torch in P2 (Gate A).
+
+Conventions (match torch):
+- grid[..., 0] = x (width, normalized [-1,1]); grid[..., 1] = y (height).
+- align_corners=True: -1 -> pixel 0, +1 -> pixel (size-1).
+- padding_mode='border': neighbor integer indices are clamped to the valid
+  range; bilinear weights come from the UNCLAMPED continuous coordinate.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+def grid_sample_bilinear(inp: mx.array, grid: mx.array,
+                         align_corners: bool = True,
+                         padding_mode: str = "border") -> mx.array:
+    """inp: [N,H,W,C]; grid: [N,gH,gW,2] in [-1,1]. Returns [N,gH,gW,C]."""
+    if padding_mode != "border":
+        raise NotImplementedError("only padding_mode='border' is needed for RIFE")
+    N, H, W, C = inp.shape
+    _, gH, gW, _ = grid.shape
+
+    gx = grid[..., 0]  # [N,gH,gW]
+    gy = grid[..., 1]
+    if align_corners:
+        ix = (gx + 1) * 0.5 * (W - 1)
+        iy = (gy + 1) * 0.5 * (H - 1)
+    else:
+        ix = ((gx + 1) * W - 1) * 0.5
+        iy = ((gy + 1) * H - 1) * 0.5
+
+    x0 = mx.floor(ix); y0 = mx.floor(iy)
+    x1 = x0 + 1; y1 = y0 + 1
+    wx1 = ix - x0; wx0 = 1.0 - wx1
+    wy1 = iy - y0; wy0 = 1.0 - wy1
+
+    def clampx(a): return mx.clip(a, 0, W - 1).astype(mx.int32)
+    def clampy(a): return mx.clip(a, 0, H - 1).astype(mx.int32)
+    x0c, x1c = clampx(x0), clampx(x1)
+    y0c, y1c = clampy(y0), clampy(y1)
+
+    inp_flat = inp.reshape(N, H * W, C)
+
+    def gather(yc, xc):  # yc,xc: [N,gH,gW] int -> [N,gH,gW,C]
+        idx = (yc * W + xc).reshape(N, gH * gW, 1)
+        idx = mx.broadcast_to(idx, (N, gH * gW, C))
+        g = mx.take_along_axis(inp_flat, idx, axis=1)
+        return g.reshape(N, gH, gW, C)
+
+    v00 = gather(y0c, x0c); v01 = gather(y0c, x1c)
+    v10 = gather(y1c, x0c); v11 = gather(y1c, x1c)
+
+    w00 = (wy0 * wx0)[..., None]; w01 = (wy0 * wx1)[..., None]
+    w10 = (wy1 * wx0)[..., None]; w11 = (wy1 * wx1)[..., None]
+    return v00 * w00 + v01 * w01 + v10 * w10 + v11 * w11

fastvideo/third_party/rife_mlx/ops/interpolate.py

@@ -0,0 +1,51 @@
+"""Bilinear resize (NHWC) — MLX equivalent of F.interpolate(mode='bilinear').
+
+RIFE uses bilinear resize for the coarse-to-fine scale pyramid inside IFNet and
+for the `--scale` knob (0.25..4.0). align_corners is pinned from the bundled
+IFNet_HDv3.py in P3 (RIFE 4.x typically uses align_corners=False). Parity-locked
+in P2 alongside grid_sample.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+def _sample_coords(out: int, in_: int, align_corners: bool) -> mx.array:
+    dst = mx.arange(out, dtype=mx.float32)
+    if align_corners:
+        scale = (in_ - 1) / (out - 1) if out > 1 else 0.0
+        return dst * scale
+    scale = in_ / out
+    return (dst + 0.5) * scale - 0.5
+
+
+def _bilinear_1d(x: mx.array, axis: int, out: int, align_corners: bool) -> mx.array:
+    """Resample one spatial axis with bilinear weights + border clamp."""
+    in_ = x.shape[axis]
+    if in_ == out:
+        return x
+    src = _sample_coords(out, in_, align_corners)       # [out]
+    i0 = mx.floor(src); i1 = i0 + 1
+    w1 = src - i0; w0 = 1.0 - w1
+    i0 = mx.clip(i0, 0, in_ - 1).astype(mx.int32)
+    i1 = mx.clip(i1, 0, in_ - 1).astype(mx.int32)
+    g0 = mx.take(x, i0, axis=axis)
+    g1 = mx.take(x, i1, axis=axis)
+    shape = [1] * x.ndim; shape[axis] = out
+    w0 = w0.reshape(shape); w1 = w1.reshape(shape)
+    return g0 * w0 + g1 * w1
+
+
+def interpolate_bilinear(x: mx.array, size: tuple[int, int] | None = None,
+                         scale_factor: float | None = None,
+                         align_corners: bool = False) -> mx.array:
+    """x: [N,H,W,C]. Resize H,W to `size` or by `scale_factor` (bilinear)."""
+    N, H, W, C = x.shape
+    if size is not None:
+        oH, oW = size
+    else:
+        oH, oW = int(round(H * scale_factor)), int(round(W * scale_factor))
+    x = _bilinear_1d(x, axis=1, out=oH, align_corners=align_corners)
+    x = _bilinear_1d(x, axis=2, out=oW, align_corners=align_corners)
+    return x

fastvideo/third_party/rife_mlx/ops/spatial.py

@@ -0,0 +1,18 @@
+"""NHWC pixel-shuffle — verified (C,r,r) ordering matching torch.nn.PixelShuffle.
+
+Reused from the realesrgan-mlx port (mlx-porting pitfall #7). RIFE's IFBlock
+lastconv ends in nn.PixelShuffle(2).
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+def pixel_shuffle_nhwc(x: mx.array, r: int) -> mx.array:
+    """[N,H,W,C*r^2] -> [N,H*r,W*r,C], channel split (C,r,r) (torch-matching)."""
+    n, h, w, c_in = x.shape
+    c = c_in // (r * r)
+    x = x.reshape((n, h, w, c, r, r))
+    x = mx.transpose(x, (0, 1, 4, 2, 5, 3))  # (N,H,r_i,W,r_j,C)
+    return x.reshape((n, h * r, w * r, c))

fastvideo/third_party/rife_mlx/pipeline_mlx.py

@@ -0,0 +1,80 @@
+"""High-level interpolation + CLI (image-pair + full video). P5.
+
+    rife-mlx -i in.mp4 -o out.mp4 --multi 2 --scale 1.0
+    rife-mlx --img0 a.png --img1 b.png -t 0.5 -o mid.png
+
+--multi N inserts N-1 frames between each pair (at t = i/N); output fps = N*input.
+--scale {0.25,0.5,1,2,4} sets the coarse-to-fine pyramid scale (0.5 for 4K).
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import mlx.core as mx
+import numpy as np
+
+from .config import DEFAULT_VERSION
+from .utils.weights import build_model
+
+
+def _to_nhwc(img_u8: np.ndarray) -> mx.array:
+    return mx.array(img_u8.astype(np.float32)[None] / 255.0)
+
+
+def _to_u8(arr: mx.array) -> np.ndarray:
+    a = np.clip(np.array(arr)[0], 0.0, 1.0)
+    return (a * 255.0).round().astype(np.uint8)
+
+
+def interpolate_pair(model, img0_u8, img1_u8, timestep=0.5, scale=1.0) -> np.ndarray:
+    out = model.inference(_to_nhwc(img0_u8), _to_nhwc(img1_u8), timestep, scale)
+    mx.eval(out)
+    return _to_u8(out)
+
+
+def interpolate_sequence(model, frames, multi=2, scale=1.0):
+    """frames: list[HWC uint8] -> list with multi-1 frames inserted per pair."""
+    out = []
+    for i in range(len(frames) - 1):
+        out.append(frames[i])
+        for j in range(1, multi):
+            out.append(interpolate_pair(model, frames[i], frames[i + 1], j / multi, scale))
+    out.append(frames[-1])
+    return out
+
+
+def cli_main() -> None:
+    p = argparse.ArgumentParser(description="Practical-RIFE 4.25 MLX frame interpolation")
+    p.add_argument("-i", "--input", help="input video")
+    p.add_argument("--img0"); p.add_argument("--img1")
+    p.add_argument("-o", "--output", required=True)
+    p.add_argument("-t", "--timestep", type=float, default=0.5, help="image-pair mode")
+    p.add_argument("-m", "--multi", type=int, default=2, help="video fps multiplier")
+    p.add_argument("-s", "--scale", type=float, default=1.0,
+                   choices=[0.25, 0.5, 1.0, 2.0, 4.0])
+    p.add_argument("-n", "--version", default=DEFAULT_VERSION)
+    p.add_argument("--weights_dir", default=None)
+    args = p.parse_args()
+
+    model = build_model(args.version, weights_dir=args.weights_dir)
+
+    if args.img0 and args.img1:
+        from PIL import Image
+        a = np.asarray(Image.open(args.img0).convert("RGB"))
+        b = np.asarray(Image.open(args.img1).convert("RGB"))
+        mid = interpolate_pair(model, a, b, args.timestep, args.scale)
+        Image.fromarray(mid).save(args.output)
+        print(f"{args.img0} + {args.img1} @t={args.timestep} -> {args.output} {mid.shape}")
+        return
+
+    from .video import read_frames, write_video
+    frames, fps = read_frames(args.input)
+    out_frames = interpolate_sequence(model, frames, args.multi, args.scale)
+    write_video(args.output, out_frames, fps * args.multi, audio_from=args.input)
+    print(f"{args.input} ({len(frames)}f @{fps:.2f}) -> {args.output} "
+          f"({len(out_frames)}f @{fps * args.multi:.2f}, audio passthrough)")
+
+
+if __name__ == "__main__":
+    cli_main()

fastvideo/third_party/rife_mlx/utils/__init__.py

@@ -0,0 +1 @@
+"""Weight loading + conversion utilities."""

fastvideo/third_party/rife_mlx/utils/weights.py

@@ -0,0 +1,44 @@
+"""Model build + weight load (HF auto-download / local override).
+
+Resolution: weights_dir arg | $RIFE_MLX_WEIGHTS_DIR | dist/<hf_name> | HF
+mlx-community/<hf_name>. Each dir: model.safetensors + config.json.
+"""
+
+from __future__ import annotations
+
+import os
+
+import mlx.core as mx
+
+from ..config import VERSIONS, VersionConfig
+from ..model.RIFE_HDv3 import Model
+
+WEIGHTS_DIR_ENV = "RIFE_MLX_WEIGHTS_DIR"
+HF_ORG = "mlx-community"
+
+
+def resolve_repo_id(version: str) -> str:
+    return f"{HF_ORG}/{VERSIONS[version].hf_name}"
+
+
+def _resolve_dir(cfg: VersionConfig, weights_dir: str | None) -> str:
+    if weights_dir:
+        return weights_dir
+    env = os.environ.get(WEIGHTS_DIR_ENV)
+    if env and os.path.isdir(env):
+        return env
+    local = os.path.join("dist", cfg.hf_name)
+    if os.path.isdir(local):
+        return local
+    from huggingface_hub import snapshot_download
+    return snapshot_download(repo_id=f"{HF_ORG}/{cfg.hf_name}")
+
+
+def build_model(version: str = "4.25", weights_dir: str | None = None) -> Model:
+    cfg = VERSIONS[version]
+    wdir = _resolve_dir(cfg, weights_dir)
+    weights = dict(mx.load(os.path.join(wdir, "model.safetensors")).items())
+    model = Model(scale_list=cfg.scale_list, pad_to=cfg.pad_to)
+    model.flownet.load_weights(list(weights.items()), strict=True)
+    mx.eval(model.flownet.parameters())
+    return model

fastvideo/third_party/taehv/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Ollin Boer Bohan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

fastvideo/third_party/taehv/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: MIT
+"""Vendored TAEHV (Tiny AutoEncoder for Hunyuan/Wan video latents).
+
+Upstream: https://github.com/madebyollin/taehv -- see ``taehv.py``'s header for
+provenance and ``LICENSE`` for the MIT license text.
+"""
+
+from fastvideo.third_party.taehv.taehv import TAEHV
+
+__all__ = ["TAEHV"]

fastvideo/third_party/taehv/taehv.py

@@ -0,0 +1,475 @@
+# Vendored from https://github.com/madebyollin/taehv (MIT License, (c) 2025
+# Ollin Boer Bohan; see LICENSE in this directory).
+#
+# Source file: https://raw.githubusercontent.com/madebyollin/taehv/main/taehv.py
+# Accessed: 2026-07-02
+# Upstream file sha256: 1e228d34d47e2f95b17f1aa190f551d89ce5ef4af0f82be459a46bea81ba3fd2
+#
+# Vendored unmodified below this header so the Apple Silicon TAEHV decode path
+# (fastvideo/mlx_runtime/taehv_decode.py) does not download and execute remote
+# code at runtime. To update, re-fetch the file, refresh the sha256 above, and
+# re-run fastvideo/tests/mlx/.
+#!/usr/bin/env python3
+"""
+Tiny AutoEncoder for Hunyuan Video
+(DNN for encoding / decoding videos to Hunyuan Video's latent space)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+from collections import namedtuple
+
+TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+class MemBlock(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out))
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.act = nn.ReLU(inplace=True)
+    def forward(self, x, past):
+        return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
+
+class TPool(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f*stride,n_f, 1, bias=False)
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        return self.conv(x.reshape(-1, self.stride * C, H, W))
+
+class TGrow(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False)
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        x = self.conv(x)
+        return x.reshape(-1, C, H, W)
+
+def apply_model_with_memblocks_parallel(model, x, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input,
+    with parallelization over the time axis and iteration over blocks.
+
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - show_progress_bar: if True, enables tqdm progressbar display
+
+    Returns NTCHW tensor of output data.
+    """
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    N, T, C, H, W = x.shape
+    x = x.reshape(N*T, C, H, W)
+
+    # parallel over input timesteps, iterate over blocks
+    for b in tqdm(model, disable=not show_progress_bar):
+        if isinstance(b, MemBlock):
+            NT, C, H, W = x.shape
+            T = NT // N
+            _x = x.reshape(N, T, C, H, W)
+            # pad with zeros along time axis (i.e. empty memory), slice
+            block_memory = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape)
+            x = b(x, block_memory)
+        else:
+            x = b(x)
+    NT, C, H, W = x.shape
+    T = NT // N
+    return x.view(N, T, C, H, W)
+
+def apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar=None):
+    """
+    Process the work queue (a graph traversal over blocks and timesteps)
+    until an output frame is produced or the queue is empty.
+    Mutates memory and work_queue in place.
+
+    Returns N1CHW output tensor, or None if the queue needs more input.
+    """
+    while work_queue:
+        xt, i = work_queue.pop(0)
+        if progress_bar is not None and i == 0:
+            progress_bar.update(1)
+        if i == len(model):
+            return xt.unsqueeze(1)
+        b = model[i]
+        if isinstance(b, MemBlock):
+            # mem blocks are simple since we're visiting the graph in causal order
+            if memory[i] is None:
+                xt_new = b(xt, xt * 0)
+            else:
+                xt_new = b(xt, memory[i])
+            memory[i] = xt
+            work_queue.insert(0, TWorkItem(xt_new, i+1))
+        elif isinstance(b, TPool):
+            # pool blocks accumulate inputs until they have enough to pool
+            if memory[i] is None:
+                memory[i] = []
+            memory[i].append(xt)
+            if len(memory[i]) > b.stride:
+                raise ValueError(f"TPool memory overflow: {len(memory[i])} items for stride {b.stride}")
+            elif len(memory[i]) == b.stride:
+                N, C, H, W = xt.shape
+                xt = b(torch.cat(memory[i], 1).view(N*b.stride, C, H, W))
+                memory[i] = []
+                work_queue.insert(0, TWorkItem(xt, i+1))
+        elif isinstance(b, TGrow):
+            xt = b(xt)
+            NT, C, H, W = xt.shape
+            for xt_next in reversed(xt.view(NT//b.stride, b.stride*C, H, W).chunk(b.stride, 1)):
+                work_queue.insert(0, TWorkItem(xt_next, i+1))
+        else:
+            xt = b(xt)
+            work_queue.insert(0, TWorkItem(xt, i+1))
+    return None
+
+def apply_model_with_memblocks_sequential(model, x, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input,
+    with iteration over timesteps as well as blocks.
+
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - show_progress_bar: if True, enables tqdm progressbar display
+
+    Returns NTCHW tensor of output data.
+    """
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    work_queue = [TWorkItem(xt, 0) for xt in x.unbind(1)]
+    memory = [None] * len(model)
+    progress_bar = tqdm(range(len(work_queue)), disable=not show_progress_bar)
+    out = []
+    while work_queue:
+        xt = apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar)
+        if xt is not None:
+            out.append(xt)
+    progress_bar.close()
+    return torch.cat(out, 1)
+
+def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input.
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - parallel: if True, parallelize over timesteps (fast but uses O(T) memory)
+        if False, each timestep will be processed sequentially (slow but uses O(1) memory)
+    - show_progress_bar: if True, enables tqdm progressbar display
+
+    Returns NTCHW tensor of output data.
+    """
+    if parallel:
+        return apply_model_with_memblocks_parallel(model, x, show_progress_bar)
+    else:
+        return apply_model_with_memblocks_sequential(model, x, show_progress_bar)
+
+class TAEHV(nn.Module):
+    def __init__(self, checkpoint_path="taehv.pth", encoder_time_downscale=(True, True, False), decoder_time_upscale=(False, True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16):
+        """Initialize pretrained TAEHV from the given checkpoint.
+
+        Arg:
+            checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1.
+            encoder_time_downscale: whether temporal downsampling is enabled for each block.
+            decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            patch_size: input/output pixelshuffle patch-size for this model.
+            latent_channels: number of latent channels (z dim) for this model.
+        """
+        super().__init__()
+        self.patch_size = patch_size
+        self.latent_channels = latent_channels
+        self.image_channels = 3
+        if len(decoder_time_upscale) == 2:
+            decoder_time_upscale = (False, *decoder_time_upscale)
+        self.is_cogvideox = checkpoint_path is not None and "taecvx" in checkpoint_path
+        if checkpoint_path is not None and "taew2_2" in checkpoint_path:
+            self.patch_size, self.latent_channels = 2, 48
+        if checkpoint_path is not None and "taehv1_5" in checkpoint_path:
+            self.patch_size, self.latent_channels = 2, 32
+        if checkpoint_path is not None and "taeltx" in checkpoint_path: # same for both 2 and 2.3
+            self.patch_size, self.latent_channels, encoder_time_downscale, decoder_time_upscale = 4, 128, (True, True, True), (True, True, True)
+        self.encoder = nn.Sequential(
+            conv(self.image_channels*self.patch_size**2, 64), nn.ReLU(inplace=True),
+            TPool(64, 2 if encoder_time_downscale[0] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 2 if encoder_time_downscale[1] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 2 if encoder_time_downscale[2] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            conv(64, self.latent_channels),
+        )
+        n_f = [256, 128, 64, 64]
+        self.decoder = nn.Sequential(
+            Clamp(), conv(self.latent_channels, n_f[0]), nn.ReLU(inplace=True),
+            MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1), conv(n_f[0], n_f[1], bias=False),
+            MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1), conv(n_f[1], n_f[2], bias=False),
+            MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1), conv(n_f[2], n_f[3], bias=False),
+            nn.ReLU(inplace=True), conv(n_f[3], self.image_channels*self.patch_size**2),
+        )
+        # computed properties
+        self.t_downscale = 2**sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool))
+        self.t_upscale = 2**sum(t.stride == 2 for t in self.decoder if isinstance(t, TGrow))
+        self.frames_to_trim = self.t_upscale - 1
+
+        if checkpoint_path is not None:
+            self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True)))
+
+    def patch_tgrow_layers(self, sd):
+        """Patch TGrow layers to use a smaller kernel if needed.
+
+        Args:
+            sd: state dict to patch
+        """
+        new_sd = self.state_dict()
+        for i, layer in enumerate(self.decoder):
+            if isinstance(layer, TGrow):
+                key = f"decoder.{i}.conv.weight"
+                if sd[key].shape[0] > new_sd[key].shape[0]:
+                    # take the last-timestep output channels
+                    sd[key] = sd[key][-new_sd[key].shape[0]:]
+        return sd
+
+    def preprocess_input_frames(self, x):
+        """Preprocess RGB input frames prior to the main encoder sequence."""
+        if self.patch_size > 1: x = F.pixel_unshuffle(x, self.patch_size)
+        return x
+
+    def encode_video(self, x, parallel=True, show_progress_bar=True):
+        """Encode a sequence of frames.
+
+        Args:
+            x: input NTCHW RGB (C=3) tensor with values in [0, 1].
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW latent tensor with ~Gaussian values.
+        """
+        x = self.preprocess_input_frames(x)
+        if x.shape[1] % self.t_downscale != 0:
+            # pad at end to multiple of self.t_downscale
+            n_pad = self.t_downscale - x.shape[1] % self.t_downscale
+            padding = x[:, -1:].repeat_interleave(n_pad, dim=1)
+            x = torch.cat([x, padding], 1)
+        return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
+
+    def postprocess_output_frames(self, x):
+        """Postprocess RGB frames after the main decoder sequence."""
+        if self.patch_size > 1: x = F.pixel_shuffle(x, self.patch_size)
+        return x.clamp_(0, 1)
+
+    def decode_video(self, x, parallel=True, show_progress_bar=True):
+        """Decode a sequence of frames.
+
+        Args:
+            x: input NTCHW latent (C=self.latent_channels) tensor with ~Gaussian values.
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW RGB tensor with ~[0, 1] values.
+        """
+        skip_trim = self.is_cogvideox and x.shape[1] % 2 == 0
+        x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
+        x = self.postprocess_output_frames(x)
+        if skip_trim:
+            # skip trimming for cogvideox to make frame counts match.
+            # this still doesn't have correct temporal alignment for certain frame counts
+            # (cogvideox seems to pad at the start?), but for multiple-of-4 it's fine.
+            return x
+        return x[:, self.frames_to_trim:]
+
+class StreamingTAEHV(nn.Module):
+    def __init__(self, taehv):
+        """Streaming wrapper around TAEHV for real-time use-cases (where not all inputs are available immediately).
+
+        Encode-decode (video-to-video) usage:
+            streaming = StreamingTAEHV(taehv)
+            for frame in video_frames:
+                latent = streaming.encode(frame_tensor)
+                decoded = streaming.decode(latent)  # feeds latent if not None, then returns next frame
+                if decoded is not None:
+                    display(decoded)
+            for frame in streaming.flush():
+                display(frame)
+
+        Decode-only (world model) usage:
+            streaming = StreamingTAEHV(taehv)
+            while running:
+                latent = world_model.step()       # latent represents t_upscale frames
+                frame = streaming.decode(latent)   # returns first frame immediately
+                while frame is not None:           # retrieve remaining frames from this latent
+                    display(frame)
+                    frame = streaming.decode()
+        """
+        super().__init__()
+        self.taehv = taehv
+        self.reset()
+
+    def reset(self):
+        """Reset all internal state. Call this to start encoding/decoding a new stream."""
+        self.encoder_work_queue, self.encoder_memory = [], [None] * len(self.taehv.encoder)
+        self.decoder_work_queue, self.decoder_memory = [], [None] * len(self.taehv.decoder)
+        self.n_frames_encoded, self.n_frames_decoded = 0, 0
+        self._last_encoder_input_frame = None
+
+    def encode(self, x=None):
+        """Feed an input frame (optional) and try to produce an encoder output.
+
+        The encoder accumulates t_downscale input frames before producing one latent,
+        so most calls will return None. Use flush_encoder() at end-of-stream to pad and
+        drain any remaining latents.
+
+        Args:
+            x: NTCHW RGB frame tensor with values in [0, 1], or None to just process pending work.
+        Returns: N1CHW latent tensor, or None if not enough input has been accumulated.
+        """
+        if x is not None:
+            assert x.ndim == 5 and x.shape[2] == self.taehv.image_channels, f"Expected NTCHW frames but got {x.shape=}"
+            self._last_encoder_input_frame = x[:, -1:]
+            x = self.taehv.preprocess_input_frames(x)
+            self.encoder_work_queue.extend(TWorkItem(xt, 0) for xt in x.unbind(1))
+            self.n_frames_encoded += x.shape[1]
+        xt = apply_model_with_memblocks_sequential_single_step(
+            self.taehv.encoder, self.encoder_memory, self.encoder_work_queue)
+        return xt
+
+    def decode(self, x=None):
+        """Feed a latent (optional) and try to produce a decoded frame.
+
+        Each latent produces t_upscale output frames due to temporal upscaling. The first
+        decode(latent) call returns the first of these frames; call decode() with no argument
+        to retrieve the rest, one at a time. Each call does the minimum decoder work needed to
+        produce one frame.
+
+        Startup frames (the first frames_to_trim raw decoder outputs, used for causal alignment
+        with the reference VAE) are consumed internally and never returned.
+
+        Args:
+            x: NTCHW latent tensor, or None to retrieve the next pending frame.
+        Returns: N1CHW decoded RGB frame tensor, or None if the queue needs more input.
+        """
+        if x is not None:
+            assert x.ndim == 5 and x.shape[2] == self.taehv.latent_channels, f"Expected NTCHW latents but got {x.shape=}"
+            self.decoder_work_queue.extend(TWorkItem(xt, 0) for xt in x.unbind(1))
+        while True:
+            xt = apply_model_with_memblocks_sequential_single_step(
+                self.taehv.decoder, self.decoder_memory, self.decoder_work_queue)
+            if xt is None:
+                return None
+            self.n_frames_decoded += 1
+            # skip startup frames (to match decode_video trim behavior)
+            if not self.taehv.is_cogvideox and self.n_frames_decoded <= self.taehv.frames_to_trim:
+                continue
+            return self.taehv.postprocess_output_frames(xt)
+
+    def flush_encoder(self):
+        """Pad (if needed) and drain all remaining latents from the encoder.
+
+        Returns list of N1CHW latent tensors.
+        """
+        latents = []
+        if self._last_encoder_input_frame is not None and self.n_frames_encoded % self.taehv.t_downscale != 0:
+            n_pad = self.taehv.t_downscale - self.n_frames_encoded % self.taehv.t_downscale
+            for _ in range(n_pad):
+                lat = self.encode(self._last_encoder_input_frame)
+                if lat is not None:
+                    latents.append(lat)
+        while (lat := self.encode()) is not None:
+            latents.append(lat)
+        return latents
+
+    def flush_decoder(self):
+        """Drain all remaining decoded frames from the decoder.
+
+        Returns list of N1CHW decoded RGB frame tensors.
+        """
+        frames = []
+        while (frame := self.decode()) is not None:
+            frames.append(frame)
+        return frames
+
+    def flush(self):
+        """Flush encoder (with padding) and decoder, returning all remaining decoded frames.
+
+        Returns list of N1CHW decoded RGB frame tensors.
+        """
+        frames = []
+        for latent in self.flush_encoder():
+            frame = self.decode(latent)
+            if frame is not None:
+                frames.append(frame)
+        frames.extend(self.flush_decoder())
+        return frames
+
+@torch.no_grad()
+def main():
+    """Run TAEHV roundtrip reconstruction on the given video paths."""
+    import os
+    import sys
+    import cv2 # no highly esteemed deed is commemorated here
+
+    class VideoTensorReader:
+        def __init__(self, video_file_path):
+            self.cap = cv2.VideoCapture(video_file_path)
+            assert self.cap.isOpened(), f"Could not load {video_file_path}"
+            self.fps = self.cap.get(cv2.CAP_PROP_FPS)
+        def __iter__(self):
+            return self
+        def __next__(self):
+            ret, frame = self.cap.read()
+            if not ret:
+                self.cap.release()
+                raise StopIteration  # End of video or error
+            return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW
+
+    class VideoTensorWriter:
+        def __init__(self, video_file_path, width_height, fps=30):
+            self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height)
+            assert self.writer.isOpened(), f"Could not create writer for {video_file_path}"
+        def write(self, frame_tensor):
+            assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??"
+            self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC
+        def __del__(self):
+            if hasattr(self, 'writer'): self.writer.release()
+
+    dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+    dtype = torch.float16
+    checkpoint_path = os.getenv("TAEHV_CHECKPOINT_PATH", "taehv.pth")
+    checkpoint_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
+    print(f"Using device \033[31m{dev}\033[0m, dtype \033[32m{dtype}\033[0m, checkpoint \033[34m{checkpoint_name}\033[0m ({checkpoint_path})")
+    taehv = TAEHV(checkpoint_path=checkpoint_path).to(dev, dtype)
+    for video_path in sys.argv[1:]:
+        print(f"Processing {video_path}...")
+        video_in = VideoTensorReader(video_path)
+        video = torch.stack(list(video_in), 0)[None]
+        vid_dev = video.to(dev, dtype).div_(255.0)
+        # convert to device tensor
+        if video.numel() < 100_000_000:
+            print(f"  {video_path} seems small enough, will process all frames in parallel")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        else:
+            print(f"  {video_path} seems large, will process each frame sequentially")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev, parallel=False)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc, parallel=False)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        video_out_path = video_path + f".reconstructed_by_{checkpoint_name}.mp4"
+        video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps)))
+        for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]:
+            video_out.write(frame)
+        print(f"  Saved to {video_out_path}")
+
+if __name__ == "__main__":
+    main()

mkdocs.yml

@@ -151,6 +151,7 @@ nav:
   - Home: index.md
   - Getting Started:
     - Installation: getting_started/installation.md
+    - Apple Silicon FastWan: getting_started/installation/mps.md
     - Quick Start: getting_started/quick_start.md
     - V1 API: getting_started/v1_api.md
   - Inference:

pyproject.toml

@@ -122,9 +122,11 @@ fastvideo-kernel = [
 # uv reads UV_TORCH_BACKEND / --torch-backend (not a pyproject setting). The
 # supported backends (cu126, cu130) both provide torch 2.12.0.
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
-# FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
-# comes from FastVideo's own custom_op wrappers.
-flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
+# FA4 cute. This revision pins nvidia-cutlass-dsl==4.6.0.dev0; holding cutlass-dsl
+# back at 4.5.x breaks the CuTe kernels at JIT time, not at import
+# (see fastvideo-kernel/README.md).
+# torch.compile support comes from FastVideo's own custom_op wrappers.
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "14c377950125c70b7a9dabf9c561fca53715ac7d", subdirectory = "flash_attn/cute" }
 
 [project.optional-dependencies]
 
@@ -201,6 +203,12 @@ rocm = [
     "amdsmi",
 ]
 
+# Apple Silicon FastWan release runtime. The range is intentionally pinned to
+# the MLX version used for the checked-in parity and checkpoint tests.
+mlx = [
+    "mlx>=0.31.2,<0.33; platform_system == 'Darwin' and platform_machine == 'arm64'",
+]
+
 [project.scripts]
 fastvideo = "fastvideo.entrypoints.cli.main:main"
 dreamverse-server = "dreamverse.server_entry:cli"
@@ -212,6 +220,8 @@ exclude = ["assets*", "docker*", "docs", "scripts*", "apps*", "dreamverse.tests*
 
 [tool.setuptools.package-data]
 dreamverse = ["prompts/*.md"]
+"fastvideo.third_party.rife_mlx" = ["LICENSE"]
+"fastvideo.third_party.taehv" = ["LICENSE"]
 
 [tool.wheel]
 exclude = ["assets*", "docker*", "docs", "scripts*", "apps*"]