Merge remote-tracking branch 'origin/main' into push_pr_1761

# Conflicts: # docs/getting_started/installation/spark_performance.md # docs/inference/offloading.md # fastvideo/fastvideo_args.py
2026-08-31 13:19:10 -07:00
parent cc9bfd069d a28f2bab4b
commit 72c57b4ae8
91 changed files with 11269 additions and 349 deletions
@@ -51,9 +51,16 @@ python examples/inference/basic/mlx_fasth3.py \
  --output-path ./outputs/fasth3_int6.mp4
 ```

-Pass `--fast` for temporal RIFE fast mode. This MLX entrypoint currently
-supports T2VA only; FL2VA, Ref2VA, spatial fast mode, and two-pass refinement
-remain follow-up work. The complete setup and conversion commands are in the
+Pass `--fast` for temporal RIFE fast mode and `--fast-spatial` for spatial
+fast mode (reduced-canvas denoise + pixel-space upsample); the two compose.
+VSA is opt-in: convert with `--include-vsa` and pass `--vsa` (see the
+[Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/)).
+This MLX entrypoint currently supports T2VA only; FL2VA, Ref2VA, and
+two-pass refinement remain follow-up work. INT6/INT8/INT4 are
+weight-only; VSA attention activations stay BF16. Dense-only checkpoints keep
+working for dense inference.
+
+The complete setup and conversion commands are in the
 [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).

 For an example running DMD+VSA inference:
@@ -111,7 +118,7 @@ python examples/inference/basic/basic_fasth3.py \
  --warmup-seed 999
 ```

-`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median.
+`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. `--h3-sequential-load` / `--no-h3-sequential-load` override the auto split that releases Qwen3-VL before DiT/VAE load (on by default on GB10, off on discrete GPUs). The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median.

 One script covers each validated duration; regional compile is the fastest
 measured DiT route for all three:
@@ -105,6 +105,11 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser:
                        action=argparse.BooleanOptionalAction,
                        default=True,
                        help="round-robin VAE temporal chunks across sequence-parallel ranks")
+    parser.add_argument("--h3-sequential-load",
+                        action=argparse.BooleanOptionalAction,
+                        default=None,
+                        help="encode with Qwen3-VL, release it, then load DiT/VAEs. Default auto: on for "
+                        "unified-memory devices (GB10), off on discrete GPUs")
    parser.add_argument("--replicated-dit",
                        action=argparse.BooleanOptionalAction,
                        default=True,
@@ -235,6 +240,8 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
        "vae_parallel_decode": args.parallel_vae,
        "vae_parallel_decode_strategy": "gather",
    }
+    if args.h3_sequential_load is not None:
+        experimental["h3_sequential_load"] = args.h3_sequential_load
    if use_vsa:
        experimental.update({
            "VSA_sparsity": args.vsa_sparsity,
@@ -19,8 +19,22 @@ denoising fewer video frames, then uses MLX RIFE 4.25 to reconstruct the
 requested frame count. A 1280x720 request runs on H3's 1280x736 grid and is
 center-cropped after decode.

+``--fast-spatial`` is spatial fast mode, ``--fast``'s spatial twin. It
+denoises and decodes on the smallest 32px-aligned canvas covering
+height/width divided by ``--fast-spatial-scale``, then resamples the decoded
+frames up to the requested size in pixel space. The two modes compose.
+This trades fine detail for speed: the output carries the reduced canvas's
+detail budget and reads softer than a native-resolution render, so it stays
+off by default.
+
 This entrypoint currently supports text-to-video-with-audio only. It does not
-yet wire FL2VA, Ref2VA, spatial fast mode, or two-pass refinement.
+yet wire FL2VA, Ref2VA, or two-pass refinement.
+
+VSA is off by default; existing dense MLX checkpoints remain supported.
+H3 uses fused MLX RMSNorm, which can change BF16 rounding compared with the
+older explicit normalization path. Convert with ``--include-vsa`` and pass
+``--vsa`` to enable the sparse path. Attention activations stay BF16; INT6/INT8/INT4 apply only to
+linear weights, including the optional gate projection.
 """

 from __future__ import annotations
@@ -30,6 +44,13 @@ import json
 from pathlib import Path


+def _dense_layers(value: str) -> tuple[int, ...]:
+    layers = tuple(int(part.strip()) for part in value.split(",") if part.strip())
+    if any(layer < 0 for layer in layers):
+        raise argparse.ArgumentTypeError("dense layer indices must be non-negative")
+    return layers
+
+
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
@@ -58,9 +79,31 @@ def parse_args() -> argparse.Namespace:
                        help="temporal reduction target for --fast (default: 2)")
    parser.add_argument("--fast-sharpen", type=float, default=0.6,
                        help="unsharp strength after RIFE interpolation (0 disables)")
+    parser.add_argument(
+        "--fast-spatial",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="denoise and decode at height/width // fast-spatial-scale on H3's 32px grid, "
+        "then resample the decoded frames up to the requested size; composes with --fast. "
+        "Trades fine detail for speed",
+    )
+    parser.add_argument("--fast-spatial-scale", type=int, default=2,
+                        help="spatial reduction factor for --fast-spatial (default: 2)")
+    parser.add_argument("--fast-spatial-upsample-mode",
+                        choices=("lanczos", "cubic", "bilinear", "nearest"),
+                        default="lanczos",
+                        help="pixel interpolation kernel for the post-decode upsample")
+    parser.add_argument("--fast-spatial-sharpen", type=float, default=0.4,
+                        help="unsharp strength after the upsample (0 disables)")
    parser.add_argument("--rife-weights-dir", type=Path, default=None,
                        help="optional local mlx-community/RIFE-4.25 snapshot")
    parser.add_argument("--vae-dtype", choices=("fp32", "fp16", "bf16"), default="fp32")
+    parser.add_argument("--video-decode-backend", choices=("h3-vae", "taeh3"), default="h3-vae",
+                        help="full H3 VAE or approximate TAEH3 preview decoder; audio is unchanged")
+    parser.add_argument("--taeh3-checkpoint", type=Path, default=None,
+                        help="local TAEH3 safetensors; otherwise download hash-verified upstream weights")
+    parser.add_argument("--taeh3-chunk-size", type=int, default=5,
+                        help="latent frames per TAEH3 chunk; causal memory persists across chunks")
    parser.add_argument("--prompt-cache-dir", type=Path, default=None,
                        help="directory for reusable prompt embedding caches")
    parser.add_argument(
@@ -69,6 +112,58 @@ def parse_args() -> argparse.Namespace:
        default=True,
        help="decode with the reference 256px overlapping VAE tiles (disable only for diagnostics)",
    )
+    parser.add_argument(
+        "--vsa",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help=(
+            "enable MiniMax H3 VSA; requires a VSA-capable MLX checkpoint "
+            "from --include-vsa"
+        ),
+    )
+    parser.add_argument(
+        "--vsa-sparsity",
+        type=float,
+        default=0.9,
+        help="VSA sparsity in [0, 1); 0.9 is the trained FastH3 policy",
+    )
+    parser.add_argument(
+        "--vsa-tile-size",
+        type=int,
+        default=64,
+        choices=(64, 256),
+        help="VSA tile size in tokens",
+    )
+    parser.add_argument(
+        "--vsa-prefix-mode",
+        choices=("exempt", "compete"),
+        default="exempt",
+        help=(
+            "prefix-key policy: always keep (exempt) or FLOP-matched top-k "
+            "(compete)"
+        ),
+    )
+    parser.add_argument(
+        "--vsa-dense-first-n-steps",
+        type=int,
+        default=0,
+        help="run the first N denoise steps dense",
+    )
+    parser.add_argument(
+        "--vsa-dense-layers",
+        type=_dense_layers,
+        default=(),
+        help="comma-separated layer indices forced dense",
+    )
+    parser.add_argument(
+        "--vsa-impl",
+        choices=("auto", "reference", "simd"),
+        default="auto",
+        help=(
+            "sparse attention implementation; auto uses chunked gather+SDPA "
+            "(simd is opt-in)"
+        ),
+    )
    return parser.parse_args()


@@ -80,6 +175,9 @@ def main() -> None:
        model_root=args.model_root,
        mlx_dit_checkpoint=args.mlx_checkpoint,
        vae_dtype=args.vae_dtype,
+        video_decode_backend=args.video_decode_backend,
+        taeh3_checkpoint=args.taeh3_checkpoint,
+        taeh3_chunk_size=args.taeh3_chunk_size,
        prompt_cache_dir=args.prompt_cache_dir,
    )
    result = pipeline.generate(
@@ -95,11 +193,24 @@ def main() -> None:
        fast_factor=args.fast_factor,
        fast_sharpen=args.fast_sharpen,
        rife_weights_dir=args.rife_weights_dir,
+        fast_spatial=args.fast_spatial,
+        fast_spatial_scale=args.fast_spatial_scale,
+        fast_spatial_upsample_mode=args.fast_spatial_upsample_mode,
+        fast_spatial_sharpen=args.fast_spatial_sharpen,
+        vsa=args.vsa,
+        vsa_sparsity=args.vsa_sparsity,
+        vsa_tile_size=args.vsa_tile_size,
+        vsa_prefix_mode=args.vsa_prefix_mode,
+        vsa_dense_first_n_steps=args.vsa_dense_first_n_steps,
+        vsa_dense_layers=args.vsa_dense_layers,
+        vsa_impl=args.vsa_impl,
    )
    print(json.dumps({
        "video_path": result.video_path,
        "timings_s": {k: round(v, 2) for k, v in result.timings.items()},
        "peak_memory_gib": {k: round(v, 2) for k, v in result.peak_memory_gib.items()},
+        "vsa": result.vsa,
+        "video_decode_backend": result.video_decode_backend,
        "audio_samples": int(result.waveform.shape[-1]),
    }, indent=2))

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compare H3 decoders on saved normalized video rows without rerunning denoise.
+
+The input NPZ must contain a ``video`` array of packed diffusion rows. Geometry
+and the DiT manifest must match the generation that produced those rows.
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import platform
+import resource
+import subprocess
+import time
+from pathlib import Path
+
+import numpy as np
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--latents", type=Path, required=True)
+    parser.add_argument("--model-root", type=Path, required=True)
+    parser.add_argument("--mlx-checkpoint", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=832)
+    parser.add_argument("--num-frames", type=int, default=124)
+    parser.add_argument("--backends", nargs="+", choices=("h3-vae", "taeh3"), default=["h3-vae", "taeh3"])
+    parser.add_argument("--taeh3-checkpoint", type=Path)
+    parser.add_argument("--taeh3-chunk-size", type=int, default=5)
+    parser.add_argument("--vae-dtype", choices=("fp32", "fp16", "bf16"), default="fp32")
+    parser.add_argument("--repeats", type=int, default=1)
+    args = parser.parse_args()
+    if args.repeats < 1:
+        parser.error("--repeats must be positive")
+    if args.output_dir.exists():
+        parser.error("--output-dir must be a new directory to preserve previous results")
+
+    import mlx.core as mx
+    from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline, _cleanup_mlx
+    from fastvideo.mlx_runtime.minimax_h3_taeh3 import ensure_taeh3_checkpoint
+
+    with np.load(args.latents) as archive:
+        rows = archive["video"]
+    if not np.isfinite(rows).all():
+        raise ValueError("The saved video rows contain non-finite values")
+    args.output_dir.mkdir(parents=True)
+    checksum = hashlib.sha256()
+    with args.latents.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 20), b""):
+            checksum.update(chunk)
+    checkpoint = ensure_taeh3_checkpoint(args.taeh3_checkpoint) if "taeh3" in args.backends else None
+    report = {
+        "mlx": mx.__version__,
+        "platform": platform.platform(),
+        "device": mx.device_info(),
+        "latents": str(args.latents.resolve()),
+        "latents_sha256": checksum.hexdigest(),
+        "geometry": [args.height, args.width, args.num_frames],
+        "dtype": args.vae_dtype,
+        "chunk_size": args.taeh3_chunk_size,
+        "checkpoint_download_excluded": True,
+        "decoder_loading_included": True,
+        "trials": [],
+    }
+    for repeat in range(args.repeats):
+        # Reverse each paired trial's order to expose warmup/order effects.
+        order = args.backends if repeat % 2 == 0 else args.backends[::-1]
+        for backend in order:
+            pipeline = MiniMaxH3MLXPipeline(model_root=args.model_root,
+                                           mlx_dit_checkpoint=args.mlx_checkpoint,
+                                           video_decode_backend=backend,
+                                           taeh3_checkpoint=checkpoint if backend == "taeh3" else None,
+                                           taeh3_chunk_size=args.taeh3_chunk_size,
+                                           vae_dtype=args.vae_dtype)
+            _cleanup_mlx()
+            mx.reset_peak_memory()
+            before = subprocess.check_output(["sysctl", "-n", "vm.swapusage"], text=True).strip()
+            started = time.perf_counter()
+            frames = pipeline.decode_video(rows, height=args.height, width=args.width, num_frames=args.num_frames)
+            elapsed = time.perf_counter() - started
+            trial = {
+                "repeat": repeat,
+                "backend": backend,
+                "decode_s": elapsed,
+                "shape": list(frames.shape),
+                "peak_active_gib": mx.get_peak_memory() / 2**30,
+                "process_lifetime_rss_peak_gib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**30,
+                "swap_before": before,
+                "swap_after": subprocess.check_output(["sysctl", "-n", "vm.swapusage"], text=True).strip(),
+            }
+            if repeat == 0:
+                np.save(args.output_dir / f"{backend}_frames.npy", frames)
+            report["trials"].append(trial)
+            (args.output_dir / "report.json").write_text(json.dumps(report, indent=2) + "\n")
+            print(json.dumps(trial), flush=True)
+            del frames, pipeline
+            _cleanup_mlx()
+
+
+if __name__ == "__main__":
+    main()