diff --git a/examples/diffusers/Dockerfile b/examples/diffusers/Dockerfile index 0c195783d2b2..b557bfc6157a 100644 --- a/examples/diffusers/Dockerfile +++ b/examples/diffusers/Dockerfile @@ -53,11 +53,17 @@ ARG TORCH_CUDA_ARCH_LIST="10.0a" # Lower MAX_JOBS if the build OOMs (machines with <96GB RAM and many CPU cores). ARG MAX_JOBS=4 -# Persist arch list; CMAKE_ARGS pins fastvideo-kernel to sm_100 (Blackwell; -# CUDA 12.9 supports it) and disables the ThunderKittens kernel build (TK has -# crashed at worker startup on B200 in earlier revisions). +# Persist arch list; CMAKE_ARGS pins the fastvideo-kernel CUDA arch and disables +# the ThunderKittens kernel build (TK has crashed at worker startup on B200 in +# earlier revisions). CUDA_ARCH_CMAKE must match TORCH_CUDA_ARCH_LIST's target: +# - 100a (default) -> Blackwell B200/sm_100 (LTX-2.3 fleet) +# - 90a -> Hopper H100/sm_90 (FastWan-QAD fleet; build with +# --build-arg TORCH_CUDA_ARCH_LIST=9.0a --build-arg CUDA_ARCH_CMAKE=90a) +# It's a separate arg (not derived) because the cmake form drops the dot ("9.0a" +# -> "90a") and we keep that mapping explicit rather than sed-munging it. +ARG CUDA_ARCH_CMAKE=100a ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} -ENV CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=100a -DFASTVIDEO_KERNEL_BUILD_TK=OFF" +ENV CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCH_CMAKE} -DFASTVIDEO_KERNEL_BUILD_TK=OFF" # flash-attention FA4 from FastVideo's exact fork (XOR-op/flash-attention@ # fa4-compile -- the flash-attn-cute source in their pyproject). The cute @@ -118,5 +124,21 @@ ENV FASTVIDEO_X264_PRESET=ultrafast WORKDIR /opt/app COPY . /opt/app/ +# FastWan decodes with the TAEHV tiny autoencoder (the ~2.4s path) instead of +# the full Wan VAE (~5.8s). Bake the single-file taehv module + the Wan2.1 +# checkpoint into the fastwan family dir; fastwan.factory loads them. Gated so +# LTX images (which use the full VAE) don't carry it. Build the FastWan image +# with --build-arg INCLUDE_FASTWAN_TAEHV=1. +ARG INCLUDE_FASTWAN_TAEHV=0 +RUN if [ "$INCLUDE_FASTWAN_TAEHV" = "1" ]; then \ + . /opt/dynamo/venv/bin/activate && \ + uv pip install imageio imageio-ffmpeg && \ + curl -fsSL https://raw.githubusercontent.com/madebyollin/taehv/main/taehv.py \ + -o /opt/app/fastwan/taehv.py && \ + curl -fsSL https://github.com/madebyollin/taehv/raw/main/taew2_1.pth \ + -o /opt/app/fastwan/taew2_1.pth && \ + test -s /opt/app/fastwan/taew2_1.pth ; \ + fi + # Page-cache warming + LD_LIBRARY_PATH unset happen in entrypoint.sh. ENTRYPOINT ["/opt/app/entrypoint.sh"] diff --git a/examples/diffusers/fastwan/__init__.py b/examples/diffusers/fastwan/__init__.py new file mode 100644 index 000000000000..5cc165f93702 --- /dev/null +++ b/examples/diffusers/fastwan/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FastWan-QAD-FP8 video-pipeline integration. + +Model-specific glue (factory, shape menu, warmup) lives here. Generic +infrastructure (pool, backend, metrics, models, menu-hash) lives in the +sibling ``lib`` package; shared operational docs live in ``ltx23/`` +(RUNBOOK, CACHING, ARCHITECTURE) since the machinery is common. +""" diff --git a/examples/diffusers/fastwan/factory.py b/examples/diffusers/fastwan/factory.py new file mode 100644 index 000000000000..54bddd9b167c --- /dev/null +++ b/examples/diffusers/fastwan/factory.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FastWan-QAD-FP8 model factory (TAEHV tiny-decoder path). + +``load_model`` is the SINGLE source of truth for how the FastWan-QAD +``VideoGenerator`` is constructed, used by both the in-process path +(``lib.backend``) and the pool path (``lib.pool`` resolves it via +``--model-factory fastwan.factory:load_model``). Sharing it keeps the +torch.compile cache keys byte-identical between warmup-bake and serving. + +FastWan-QAD-FP8-1.3B is a Wan2.1-T2V-1.3B distillation. vs the LTX-2.3 factory: + * single-stage pipeline -- no refine/upsampler; + * FP8 e4m3 per-tensor linear quant, ALWAYS on (the QAD checkpoint's native + format -- baked scale_weight params), not gated on enable_optimizations; + * text-to-video only; + * 3 denoise steps, no CFG (per-request); + * **TAEHV tiny autoencoder for decode** instead of the full Wan VAE. This is + what makes it the fast model: ~2.4s end-to-end vs ~5.8s with the full VAE + (the DiT denoise is identical; only the decode differs). The generator is + built with ``output_type="latent"`` so the heavy VAE is bypassed, and the + latents are decoded by TAEHV (``taew2_1.pth``, baked into the image). This + mirrors FastVideo's own ``fp8_wan2_1_1_3b.py --taehv-checkpoint`` recipe. + +Heavy imports live inside the function so the dispatcher's argv-parse + dynamic +import path stays cheap for pool subprocesses. +""" + +import logging +import os +import sys +from typing import Any + +logger = logging.getLogger(__name__) + +# TAEHV (Wan2.1 tiny autoencoder) checkpoint, baked into the image (Dockerfile). +TAEHV_CKPT = os.environ.get("FASTWAN_TAEHV_CKPT", "/opt/app/fastwan/taew2_1.pth") + + +def _load_taehv(): + import torch + + repo_dir = os.path.dirname(TAEHV_CKPT) + if repo_dir and repo_dir not in sys.path: + sys.path.insert(0, repo_dir) + from taehv import TAEHV + + return TAEHV(checkpoint_path=TAEHV_CKPT).to("cuda", torch.float16) + + +def _decode_with_taehv(taehv_model, latents): + import torch + + with torch.no_grad(): + latents = latents.permute(0, 2, 1, 3, 4) + latents = latents.to( + device=next(taehv_model.parameters()).device, + dtype=next(taehv_model.parameters()).dtype, + ) + decoded = taehv_model.decode_video( + latents, parallel=False, show_progress_bar=False + ) + return [ + (f.clamp(0, 1) * 255).byte().cpu().permute(1, 2, 0).numpy() + for f in decoded[0] + ] + + +class _TaehvVideoGenerator: + """Adapter exposing the FastVideo ``generate_video(**kwargs)`` contract that + ``lib.backend``/``lib.pool`` call, but decoding via TAEHV. The underlying + generator is built with ``output_type="latent"`` (VAE bypassed), so + ``generate()`` returns latents in ``result.samples``; we TAEHV-decode them + and write the mp4 at the requested fps. width/height/num_frames ARE plumbed + into the latent request (via SamplingParam), so multi-shape serving renders + the requested dimensions -- e.g. portrait 480x832 vs landscape 832x480 -- + rather than the model's default shape.""" + + def __init__(self, gen: Any, taehv: Any) -> None: + self._gen = gen + self._taehv = taehv + + def generate_video( + self, + prompt: str, + output_path: str | None = None, + fps: int = 16, + num_inference_steps: int = 3, + guidance_scale: float = 1.0, + seed: int | None = None, + negative_prompt: str | None = None, + width: int | None = None, + height: int | None = None, + num_frames: int | None = None, + save_video: bool = True, + return_frames: bool = False, + **_ignored: Any, + ) -> Any: + import imageio + + sampling: dict[str, Any] = { + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + } + # Plumb the requested shape onto SamplingParam.{width,height,num_frames} + # so each pool subprocess renders (and compiles for) its actual shape. + # Without this the generator falls back to the model default (832x480) + # and portrait 480x832 silently comes out landscape. + if width is not None: + sampling["width"] = width + if height is not None: + sampling["height"] = height + if num_frames is not None: + sampling["num_frames"] = num_frames + if seed is not None: + sampling["seed"] = seed + request: dict[str, Any] = { + "prompt": prompt, + "sampling": sampling, + "output": {"save_video": False}, + } + if negative_prompt: + request["negative_prompt"] = negative_prompt + result = self._gen.generate(request=request) + frames = _decode_with_taehv(self._taehv, result.samples) + if output_path: + imageio.mimsave(output_path, frames, fps=fps, format="mp4") + return result + + +def load_model( + model_path: str, + num_gpus: int, + enable_optimizations: bool, +) -> Any: + """Build the FastWan-QAD-FP8 generator (FP8 + TAEHV decode). + + ``enable_optimizations`` is accepted for the shared pool/warmup factory + signature but intentionally unused (FP8 is always on; see module docstring). + """ + from fastvideo import VideoGenerator + from fastvideo.layers.quantization import get_quantization_config + + del enable_optimizations + quant_kwargs: dict[str, Any] = { + "transformer_quant": get_quantization_config("FP8")(granularity="tensor"), + } + logger.info( + "FastWan-QAD: FP8 e4m3 per-tensor linear quant (always on); " + "attention_backend=%s", + os.environ.get("FASTVIDEO_ATTENTION_BACKEND", ""), + ) + + gen = VideoGenerator.from_pretrained( + model_path, + num_gpus=num_gpus, + use_fsdp_inference=False, + dit_cpu_offload=False, + # output_type="latent" bypasses the full Wan VAE (TAEHV decodes instead), + # so the heavy VAE can be CPU-offloaded to save GPU VRAM. + vae_cpu_offload=True, + text_encoder_cpu_offload=False, + pin_cpu_memory=False, + enable_torch_compile=True, + output_type="latent", + **quant_kwargs, + ) + taehv = _load_taehv() + logger.info("FastWan-QAD: TAEHV tiny decoder loaded from %s", TAEHV_CKPT) + return _TaehvVideoGenerator(gen, taehv) diff --git a/examples/diffusers/fastwan/preflight_test.py b/examples/diffusers/fastwan/preflight_test.py new file mode 100755 index 000000000000..2e19d50c837a --- /dev/null +++ b/examples/diffusers/fastwan/preflight_test.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Standalone smoke-test of GenericVideoBackend.preflight() without spinning +up the Dynamo distributed runtime. Loads the model, runs the preflight +loop, exits. Logs per-shape timings and total wall time. + +Used to verify that a candidate ship image's preflight will: + - Find fastwan/shapes.json + - Successfully warm every shape's in-memory cache + - Complete in a reasonable wall-clock budget for pod startup +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import sys +import time + +# Put examples/diffusers/ on sys.path so the lib / fastwan packages +# resolve. This script lives at examples/diffusers/fastwan/preflight_test.py; +# parent dir = examples/diffusers/. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--model", + default="FastVideo/FastWan-QAD-FP8-1.3B", + help="HuggingFace model identifier", + ) + p.add_argument( + "--gpu-uuid", + required=True, + help="GPU UUID to pin to (sets CUDA_VISIBLE_DEVICES before torch import)", + ) + return p.parse_args() + + +async def _amain(model: str) -> int: + from fastwan.factory import load_model + from lib.backend import GenericVideoBackend + + backend_args = argparse.Namespace( + model=model, + served_model_name=None, + num_gpus=1, + enable_optimizations=False, + attention_backend="FLASH_ATTN", + ) + + backend = GenericVideoBackend( + args=backend_args, + model_factory_callable=load_model, + model_factory_dotted="fastwan.factory:load_model", + model_label="fastwan-qad-fp8", + ) + + # Match production: default the shapes-JSON path so preflight can + # find the menu without an explicit env var. + if "WARMUP_SHAPES_JSON_PATH" not in os.environ: + os.environ["WARMUP_SHAPES_JSON_PATH"] = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "shapes.json" + ) + + t_init = time.perf_counter() + print("[preflight-test] initialize_model() ...", flush=True) + await backend.initialize_model() + print( + "[preflight-test] initialize_model() done in %.1fs" + % (time.perf_counter() - t_init), + flush=True, + ) + + t_pre = time.perf_counter() + print("[preflight-test] preflight() ...", flush=True) + await backend.preflight() + print( + "[preflight-test] preflight() done in %.1fs" % (time.perf_counter() - t_pre), + flush=True, + ) + + print( + "[preflight-test] TOTAL boot-equivalent time: %.1fs" + % (time.perf_counter() - t_init), + flush=True, + ) + return 0 + + +def main() -> int: + args = _parse_args() + + # Pin GPU before any torch import (the factory imports torch). + os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_uuid + + # Match worker.py's logging setup so preflight's log lines look the same. + logging.basicConfig( + level=( + logging.DEBUG + if os.environ.get("FASTVIDEO_LOG_LEVEL") == "DEBUG" + else logging.INFO + ), + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + force=True, + ) + + return asyncio.run(_amain(args.model)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/diffusers/fastwan/shapes.json b/examples/diffusers/fastwan/shapes.json new file mode 100644 index 000000000000..a783708fe20f --- /dev/null +++ b/examples/diffusers/fastwan/shapes.json @@ -0,0 +1,20 @@ +{ + "_comment": "FastWan-QAD-FP8 v1 shape menu (deepinfra). TWO 480p/5s shapes: landscape 832x480 and portrait 480x832, both @ 81 frames, 3 denoise steps, guidance 1.0 (no CFG) -- the model's native QAD recipe (Wan2.1-T2V-1.3B distillation). Text-to-video ONLY (no i2v), so unlike ltx23 there is no 'warm_i2v' (the model has no image-conditioning graph to warm). Both shapes are kept boot-warm by the K=2 subprocess pool (VIDEO_POOL_MODE=1, VIDEO_POOL_MAX_SIZE=2) -- one resident shape-pinned process each with its own baked compile cache, so customers never hit cold-spawn churn switching orientation. Dims are multiples of 16 (Wan VAE spatial stride); 81 = 4*20+1 (Wan VAE temporal stride 4k+1). Adding/removing shapes changes IMAGE_SHAPE_HASH and requires a new image tag + a matching per-model menu in backend ttv.py _FASTWAN_SIZE_TABLE/FastWanQADFP8In (CI-checked). Compile-cache / boot-warm strategy: reuses lib/ + warmup.py (same machinery as ltx23/CACHING.md).", + "model": "/data/default", + "fps": 16, + "num_inference_steps": 3, + "guidance_scale": 1.0, + "seed": 42, + "shapes": [ + { + "width": 832, + "height": 480, + "num_frames": 81 + }, + { + "width": 480, + "height": 832, + "num_frames": 81 + } + ] +} diff --git a/examples/diffusers/fastwan/test_shapes.py b/examples/diffusers/fastwan/test_shapes.py new file mode 100644 index 000000000000..796fe7e4c5d4 --- /dev/null +++ b/examples/diffusers/fastwan/test_shapes.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Unit tests for fastwan/shapes.json + the menu-hash algorithm. + +The hash is the load-bearing identifier across the whole pipeline: + + - bake step embeds it in IMAGE_SHAPE_HASH (Dockerfile env) + - lib/menu.py asserts it matches at boot (refuses to start on mismatch) + - backend's i model-add admission check rejects images whose hash + doesn't match the vendored menu + +If the algorithm or the menu drifts, every layer downstream is wrong. +The canonical algorithm lives in examples/diffusers/lib/menu.py; the +backend's menu-hash admission check must agree with it. + +This module reimplements the canonical hash inline so it can run in CI +without FastVideo / torch. A second cross-check test imports the live +implementation from lib.menu and skips if heavy deps aren't available +-- that's the canary against the in-container algorithm drifting. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import sys + +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +SHAPES_JSON = os.path.join(HERE, "shapes.json") + +# Pinned hash for the current 2-shape FastWan-QAD-FP8 v1 menu +# (landscape 832x480 + portrait 480x832, both @ 81 frames). Updating +# shapes.json -- adding, removing, or renaming a shape -- changes this +# value, requires a new image bake, and requires updating this fixture. +EXPECTED_HASH = "8840f800" +EXPECTED_SHAPE_COUNT = 2 + +# Per-shape activation budget. A shape whose width*height*frames is too +# large makes the decoder's intermediate tensor exceed 2^31 elements, +# tripping PyTorch's int32 indexing limit in F.pad. Any shape we add must +# stay below that bound (the current 480p shapes are far under it). +INT32_LIMIT = 2**31 +ACTIVATION_BUDGET = INT32_LIMIT // 2 + + +def _canonical_menu_hash(shapes_json_path: str) -> tuple[str, int]: + """Reimplements ``lib.menu.compute_menu_hash`` inline so this test + file has no FastVideo / torch dependency. Must stay byte-identical to + the canonical algorithm in ``examples/diffusers/lib/menu.py::compute_menu_hash`` + (the same algorithm the bake step and the backend menu-hash admission + check rely on). + """ + with open(shapes_json_path, "r", encoding="utf-8") as f: + cfg = json.load(f) + shapes = sorted( + (int(s["width"]), int(s["height"]), int(s["num_frames"])) for s in cfg["shapes"] + ) + canonical = json.dumps(shapes, separators=(",", ":")) + return hashlib.sha256(canonical.encode()).hexdigest()[:8], len(shapes) + + +def _load_shapes() -> dict: + with open(SHAPES_JSON, "r", encoding="utf-8") as f: + return json.load(f) + + +def test_warmup_shapes_parses() -> None: + cfg = _load_shapes() + assert cfg["model"] == "/data/default" + assert "shapes" in cfg + assert isinstance(cfg["shapes"], list) + + +def test_warmup_shapes_required_keys() -> None: + cfg = _load_shapes() + for shape in cfg["shapes"]: + for key in ("width", "height", "num_frames"): + assert key in shape, f"shape {shape} missing {key}" + assert isinstance(shape[key], int) + assert shape[key] > 0 + + +def test_warmup_shapes_dimensions_multiple_of_16() -> None: + """Wan2.1 needs dims that are multiples of 16 (8x VAE spatial stride x + 2x2 DiT patch); non-multiples crash decode/patchify.""" + cfg = _load_shapes() + for shape in cfg["shapes"]: + assert shape["width"] % 16 == 0, f"width not /16: {shape}" + assert shape["height"] % 16 == 0, f"height not /16: {shape}" + + +def test_warmup_shapes_within_activation_budget() -> None: + """Reject any shape whose width*height*frames overflows the int32 limit + that broke 1080p@241f. Coarse proxy for VAE decode safety.""" + cfg = _load_shapes() + for shape in cfg["shapes"]: + product = shape["width"] * shape["height"] * shape["num_frames"] + assert product < ACTIVATION_BUDGET, ( + f"shape {shape} has product {product:,} >= " + f"budget {ACTIVATION_BUDGET:,} -- VAE decode will trip int32 indexing" + ) + + +def test_canonical_menu_hash_is_deterministic() -> None: + h1, n1 = _canonical_menu_hash(SHAPES_JSON) + h2, n2 = _canonical_menu_hash(SHAPES_JSON) + assert h1 == h2 + assert n1 == n2 == EXPECTED_SHAPE_COUNT + + +def test_canonical_menu_hash_matches_expected() -> None: + """The vendored menu's hash must match what the runtime image was + baked with. Update both this fixture and the image tag together.""" + actual, count = _canonical_menu_hash(SHAPES_JSON) + assert actual == EXPECTED_HASH, ( + f"Menu hash drift: warmup_shapes.json hashes to {actual} but " + f"EXPECTED_HASH={EXPECTED_HASH}. If the menu change is intentional, " + f"update EXPECTED_HASH and rebake the runtime image." + ) + assert count == EXPECTED_SHAPE_COUNT + + +def test_canonical_menu_hash_changes_when_menu_changes(tmp_path) -> None: + """A menu edit must produce a different hash; otherwise the boot + assertion can't catch drift.""" + cfg = _load_shapes() + cfg["shapes"].append({"width": 64, "height": 64, "num_frames": 17}) + new_path = tmp_path / "shapes_modified.json" + new_path.write_text(json.dumps(cfg)) + + original_hash, _ = _canonical_menu_hash(SHAPES_JSON) + modified_hash, modified_count = _canonical_menu_hash(str(new_path)) + assert modified_hash != original_hash + assert modified_count == EXPECTED_SHAPE_COUNT + 1 + + +def test_canonical_menu_hash_order_independent(tmp_path) -> None: + """Shape order in the JSON must not affect the hash; reviewers can + reorder for readability without breaking caches.""" + cfg = _load_shapes() + cfg["shapes"] = list(reversed(cfg["shapes"])) + reversed_path = tmp_path / "shapes_reversed.json" + reversed_path.write_text(json.dumps(cfg)) + h1, _ = _canonical_menu_hash(SHAPES_JSON) + h2, _ = _canonical_menu_hash(str(reversed_path)) + assert h1 == h2 + + +def test_lib_compute_menu_hash_matches_canonical() -> None: + """Cross-check: ``lib.menu.compute_menu_hash``'s live implementation + must agree with the inline algorithm. Skipped if the package isn't + importable in this env (e.g., CI runs without the package root on + sys.path).""" + # Put examples/diffusers/ on sys.path so ``lib.menu`` resolves. + sys.path.insert(0, os.path.dirname(HERE)) + # Pre-declare so static analyzers (CodeQL) don't flag the call site + # below as potentially-unbound. pytest.skip() raises so the call is + # actually unreachable on the ImportError path, but the analyzer + # doesn't model that. + compute_menu_hash = None + try: + from lib.menu import compute_menu_hash # type: ignore + except ImportError as exc: + pytest.skip(f"lib.menu not importable in this env: {exc}") + canonical_hash, _ = _canonical_menu_hash(SHAPES_JSON) + library_hash, _ = compute_menu_hash(SHAPES_JSON) + assert canonical_hash == library_hash + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/examples/diffusers/fastwan/warmup.py b/examples/diffusers/fastwan/warmup.py new file mode 100755 index 000000000000..bb9b905a2a39 --- /dev/null +++ b/examples/diffusers/fastwan/warmup.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Pre-compile FastWan-QAD-FP8 torch.compile / triton / inductor caches for every +production shape so the first production request is fast. + +Routes each shape through the same ``lib.pool.SubprocessPool`` code path +the serving worker uses, so the cache-building subprocess is byte-identical +in code path to the runtime serving subprocess. Per-shape +``TORCHINDUCTOR_CACHE_DIR`` / ``TRITON_CACHE_DIR`` are set by the pool itself +on the child env (/cache/per-shape//{torchinductor,triton}), +matching what production reads at serve time. + +CROSS-PROCESS CACHE PORTABILITY: the implicit on-disk per-shape cache (these +dirs) does NOT port to a fresh process -- a fresh worker recomputes a different +fxgraph key and RECOMPILES. The per-shape dirs here are only the in-process +compile scratch for THIS bake run. + +To actually make a fresh pod warm, use torch **Mega-Cache** +(save/load_cache_artifacts -- wired in fastvideo gpu_worker + lib/pool.py): it +DOES port across processes. The remaining residual is the torch.compile +FRONT-END (dynamo + AOTAutograd re-traces every process to produce the cache +key) -- not cacheable short of AOTInductor. The shared machinery and the full +investigation (instrumented on LTX-2.3, 2026-06-18) live in +**ltx23/CACHING.md** + ~/ltx23_cache_investigation_report.md. + +Usage: + python fastwan/warmup.py --shapes fastwan/shapes.json \\ + --output-dir /tmp/warmup --model /data/default +""" + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import json +import os +import sys +import time +from pathlib import Path + +# Put the parent directory (examples/diffusers/) on sys.path so we can +# import ``lib.pool`` regardless of where this script is invoked from. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +PROMPT = ( + "A close-up tracking shot of a golden retriever sprinting toward the camera " + "through a sunlit alpine meadow at golden hour, paws kicking up wildflowers, " + "ears flapping, tongue out in joyful pants, shallow depth of field, shot on " + "ARRI Alexa 65, 4K, photorealistic" +) + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _du_bytes(path: Path) -> int: + if not path.exists(): + return 0 + total = 0 + for root, _dirs, files in os.walk(path): + for f in files: + # File may have vanished between os.walk and stat; safe to skip. + with contextlib.suppress(OSError): + total += (Path(root) / f).stat().st_size + return total + + +def _human(n: int) -> str: + for unit in ("B", "K", "M", "G"): + if n < 1024 or unit == "G": + return f"{n:.1f}{unit}" + n /= 1024 + return f"{n:.1f}G" + + +def _read_shapes_config(args: argparse.Namespace) -> dict: + with open(args.shapes, "r", encoding="utf-8") as f: + cfg = json.load(f) + if not cfg.get("shapes"): + print("[warmup] shapes list is empty -- nothing to do", file=sys.stderr) + sys.exit(2) + return cfg + + +def _final_summary( + successes: list[str], + failures: list[tuple[str, str]], + total_shapes: int, + cache_before: int, + cache_after: int, + min_cache_growth_bytes: int, +) -> int: + cache_growth = cache_after - cache_before + print( + f"[warmup] done. success={len(successes)}/{total_shapes} " + f"failures={[t for t, _ in failures]} " + f"cache_before={_human(cache_before)} cache_after={_human(cache_after)} " + f"cache_growth={_human(cache_growth)}", + flush=True, + ) + if failures: + return 1 + if cache_growth < min_cache_growth_bytes: + print( + f"[warmup] ERROR cache grew by {_human(cache_growth)} " + f"(< --min-cache-growth-bytes {min_cache_growth_bytes}). " + f"Did compile actually run?", + file=sys.stderr, + flush=True, + ) + return 3 + return 0 + + +# ── driver ─────────────────────────────────────────────────────────────────── + + +def _run_driver(args: argparse.Namespace) -> int: + """ + Route each shape through the production ``lib.pool.SubprocessPool`` + code path (spawning ``worker.py --pool-worker``). A fresh K=1 pool is + built per shape, the request is routed once, and the pool is shut + down before the next shape so each shape gets a fresh pool subprocess + writing to its own ``/cache/per-shape//`` dir. + + Using the production pool subprocess as the cache-building subprocess + is the load-bearing property: torch.compile / inductor fx_graph_cache + keys are sensitive to invocation context (__main__ identity, argv, + sys.modules layout). The cache-building subprocess IS the cache-reading + subprocess by construction, so keys produced here match what the + runtime worker asks for at serve time. + """ + from lib.pool import SubprocessPool + + cfg = _read_shapes_config(args) + + model = cfg.get("model", "/data/default") + fps = int(cfg.get("fps", 24)) + num_inference_steps = int(cfg.get("num_inference_steps", 5)) + guidance_scale = float(cfg.get("guidance_scale", 1.0)) + seed = int(cfg.get("seed", 42)) + shapes = cfg["shapes"] + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + per_shape_root = Path("/cache/per-shape") + cache_before = _du_bytes(per_shape_root) + print( + f"[warmup] driver start: {len(shapes)} shape(s), " + f"model={model}, cache_before={_human(cache_before)} " + f"root={per_shape_root}", + flush=True, + ) + + successes: list[str] = [] + failures: list[tuple[str, str]] = [] + + async def _route_one_shape_with_timeout( + tag: str, request: dict, timeout: float + ) -> dict: + pool = SubprocessPool( + model_path=model, + num_gpus=1, + # Bake the cache so it MATCHES the serving path: fastwan.factory + # always loads FP8 (the QAD checkpoint's native format), so the same + # graph is compiled here and at serve time. Denoise steps (3) come + # from the shapes file. enable_optimizations is passed through but + # the fastwan factory ignores it (FP8 is unconditional). + enable_optimizations=True, + attention_backend="FLASH_ATTN", + model_factory_dotted="fastwan.factory:load_model", + model_label="fastwan-qad-fp8", + ) + try: + return await asyncio.wait_for(pool.route(tag, request), timeout=timeout) + finally: + await pool.shutdown() + + for idx, shape in enumerate(shapes, 1): + w = int(shape["width"]) + h = int(shape["height"]) + nf = int(shape["num_frames"]) + tag = f"{w}x{h}@{nf}f" + out_path = output_dir / f"shape_{tag}.mp4" + if out_path.exists(): + out_path.unlink() + + # The pool's _spawn sets TORCHINDUCTOR_CACHE_DIR / TRITON_CACHE_DIR + # on the child env using the shape_key, so we only need the dirs + # to exist (the pool subprocess writes into them). + shape_inductor = per_shape_root / tag / "torchinductor" + shape_triton = per_shape_root / tag / "triton" + shape_inductor.mkdir(parents=True, exist_ok=True) + shape_triton.mkdir(parents=True, exist_ok=True) + + request = { + "request_id": f"warmup_{tag}", + "prompt": PROMPT, + "width": w, + "height": h, + "num_frames": nf, + "fps": fps, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "seed": seed, + "negative_prompt": None, + "output_path": str(out_path), + } + + print( + f"[warmup] ({idx}/{len(shapes)}) launching {tag} -> {out_path}", + flush=True, + ) + print( + f"[warmup] ({idx}/{len(shapes)}) per-shape cache dir: " + f"{per_shape_root}/{tag}", + flush=True, + ) + + t0 = time.perf_counter() + try: + result = asyncio.run( + _route_one_shape_with_timeout(tag, request, args.per_shape_timeout) + ) + except asyncio.TimeoutError: + elapsed = time.perf_counter() - t0 + msg = f"timeout after {elapsed:.1f}s" + print(f"[warmup] FAIL {tag}: {msg}", flush=True) + failures.append((tag, msg)) + if args.fail_fast: + break + continue + except Exception as exc: + elapsed = time.perf_counter() - t0 + msg = f"{type(exc).__name__}: {exc}" + print(f"[warmup] FAIL {tag}: {msg} (after {elapsed:.1f}s)", flush=True) + failures.append((tag, msg)) + if args.fail_fast: + break + continue + + elapsed = time.perf_counter() - t0 + status = result.get("status") + if status == "DONE" and out_path.exists() and out_path.stat().st_size > 0: + size_mb = out_path.stat().st_size / 1_048_576 + print( + f"[warmup] OK {tag} in {elapsed:.1f}s ({size_mb:.1f}MB)", + flush=True, + ) + successes.append(tag) + else: + msg = ( + f"status={status} " + f"error={result.get('error', '?')} " + f"file_exists={out_path.exists()} " + f"size={out_path.stat().st_size if out_path.exists() else 0}" + ) + print(f"[warmup] FAIL {tag}: {msg} (after {elapsed:.1f}s)", flush=True) + failures.append((tag, msg)) + if args.fail_fast: + break + + cache_after = _du_bytes(per_shape_root) + return _final_summary( + successes, + failures, + len(shapes), + cache_before, + cache_after, + args.min_cache_growth_bytes, + ) + + +# ── argument parsing ───────────────────────────────────────────────────────── + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="FastWan-QAD-FP8 warmup / compile-cache populator. Routes each " + "shape through lib.pool.SubprocessPool so the cache-building " + "subprocess matches the runtime serving subprocess in code path. " + "Cache keys produced here match what the runtime asks for at " + "serve time.", + ) + + p.add_argument( + "--shapes", + default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "shapes.json"), + help="path to shapes JSON (default: fastwan/shapes.json next to this script)", + ) + p.add_argument( + "--output-dir", default="/tmp/warmup", help="where to save rendered MP4s" + ) + p.add_argument( + "--per-shape-timeout", + type=int, + default=1800, + help="per-shape subprocess timeout in seconds", + ) + p.add_argument( + "--fail-fast", + action="store_true", + help="stop on first shape failure (default: continue, report at end)", + ) + p.add_argument( + "--min-cache-growth-bytes", + type=int, + default=0, + help="fail if combined cache grew by fewer bytes than this", + ) + + # Override for the model path baked into shapes.json. The driver + # reads ``cfg.get("model", "/data/default")`` from shapes.json; this + # flag is currently a no-op (kept as a forward-compat hook). The + # production default is the local /data/default mount. + p.add_argument("--model", default="/data/default") + + return p.parse_args() + + +def main() -> int: + args = _parse_args() + return _run_driver(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/diffusers/fastwan/worker.py b/examples/diffusers/fastwan/worker.py new file mode 100644 index 000000000000..60a0e813dd60 --- /dev/null +++ b/examples/diffusers/fastwan/worker.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +FastWan-QAD-FP8 worker entry point for Dynamo (non-streaming). + +Registers a FastVideo VideoGenerator as a Dynamo backend endpoint +compatible with the ``/v1/videos`` frontend endpoint. The endpoint +generates a full video clip from the request parameters and returns +it as a single response containing the complete MP4 file +base64-encoded in ``data[0].b64_json``. + +Generation parameters (size, fps, num_frames, etc.) are taken from +the request body's ``nvext`` field, so the same worker instance can +serve requests with different resolutions and quality settings +without restarting. + +One request at a time (asyncio.Lock — VideoGenerator is not +re-entrant). + +Usage: + python worker.py [--model MODEL] [--num-gpus N] [--enable-optimizations] + [--attention-backend ATTENTION_BACKEND] + +This module is the FastWan-QAD-FP8-specific entry. The top-level +``examples/diffusers/worker.py`` shim is what the production +deployment / pool-subprocess invocations actually execute; the shim +dispatches pool-worker invocations to ``lib.pool`` directly (to skip +heavy imports here) and otherwise calls :func:`main_cli` from this +module. +""" + +import argparse +import asyncio +import logging +import os +import sys + +import uvloop +from lib.backend import GenericVideoBackend +from lib.dynamo_wiring import get_worker_namespace, register_model +from lib.menu import assert_shape_menu_hash_matches +from lib.metrics import VIDEO_REGISTRY + +from dynamo.common.utils.prometheus import register_engine_metrics_callback +from dynamo.runtime import DistributedRuntime + +from .factory import load_model + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "FastVideo/FastWan-QAD-FP8-1.3B" +# FLASH_ATTN: on H100, SageAttention2 measured equal to Flash for this +# 1.3B/480p/3-step workload (full SA2 vs Flash within ~3%), so we run Flash and +# skip the SageAttention source build. The FP4/SageAttention3 path is Blackwell +# sm_120-only and irrelevant here. TORCH_SDPA is the slow generic fallback. +DEFAULT_ATTENTION_BACKEND = "FLASH_ATTN" +DEFAULT_MODEL_LABEL = "fastwan-qad-fp8" + +# Where the FastWan shapes JSON lives in the production image. Production +# bake-time IMAGE_SHAPE_HASH is computed against this path; if the +# image build moves it, update both ends in lockstep. +DEFAULT_SHAPES_JSON_PATH = os.path.join(os.path.dirname(__file__), "shapes.json") +PRODUCTION_SHAPES_JSON_PATH = "/opt/app/fastwan/shapes.json" + + +def _attention_backend_choices() -> tuple[str, ...]: + # Lazy: FastVideo's enum is heavy. Only evaluate when the parent + # parses CLI args; pool subprocesses don't reach _parse_args. + from fastvideo.platforms.interface import AttentionBackendEnum + + return tuple( + backend_name + for backend_name in AttentionBackendEnum.__members__ + if backend_name != "NO_ATTENTION" + ) + + +def _parse_args() -> argparse.Namespace: + choices = _attention_backend_choices() + parser = argparse.ArgumentParser( + description="FastVideo Worker for Dynamo (non-streaming)" + ) + parser.add_argument( + "--model", + default=DEFAULT_MODEL, + help=f"HuggingFace model path (default: {DEFAULT_MODEL})", + ) + parser.add_argument( + "--served-model-name", + default=None, + help="Name advertised to the Dynamo discovery layer (default: same as --model)", + ) + parser.add_argument( + "--num-gpus", + type=int, + default=1, + dest="num_gpus", + help="Number of GPUs (default: 1)", + ) + parser.add_argument( + "--enable-optimizations", + action="store_true", + dest="enable_optimizations", + help="Enable FP4 quantization (if available) and torch.compile", + ) + parser.add_argument( + "--attention-backend", + choices=choices, + default=DEFAULT_ATTENTION_BACKEND, + dest="attention_backend", + help=( + "Attention backend to set via FASTVIDEO_ATTENTION_BACKEND " + f"(choices: {', '.join(choices)}; " + f"default: {DEFAULT_ATTENTION_BACKEND})" + ), + ) + return parser.parse_args() + + +async def _main(args: argparse.Namespace) -> None: + loop = asyncio.get_running_loop() + # Use Kubernetes discovery in-cluster and file discovery for local compose by default. + discovery_backend = os.environ.get("DYN_DISCOVERY_BACKEND") + if not discovery_backend: + discovery_backend = ( + "kubernetes" if os.environ.get("KUBERNETES_SERVICE_HOST") else "file" + ) + namespace_name = get_worker_namespace() + logger.info("Using discovery backend: %s", discovery_backend) + logger.info("Resolved worker namespace: %s", namespace_name) + # Pass enable_nats=False explicitly: the bundled ai-dynamo-runtime 1.0.0 + # is from before upstream commit af0ff07 ("remove enable_nats usage") + # and so its DistributedRuntime ctor still requires the 4th positional + # arg to gate the NATS client. Omitting it defaults to True, which + # makes the worker hard-fail at startup with "Failed to connect to + # NATS: Connection refused" because the cluster runs ZMQ-only + # (DYN_EVENT_PLANE=zmq, deepinfra has no NATS). + runtime = DistributedRuntime(loop, discovery_backend, "tcp", False) + + component_name = "backend" + endpoint_name = "generate" + endpoint = runtime.endpoint(f"{namespace_name}.{component_name}.{endpoint_name}") + logger.info( + "Serving endpoint %s/%s/%s", namespace_name, component_name, endpoint_name + ) + + model_label = args.served_model_name or DEFAULT_MODEL_LABEL + backend = GenericVideoBackend( + args=args, + model_factory_callable=load_model, + model_factory_dotted="fastwan.factory:load_model", + model_label=model_label, + ) + await backend.initialize_model() + # Eager warm-on-boot is the DEFAULT. The FastWan v1 menu is tiny (2 shapes), + # so warming every shape's pool subprocess at boot is cheap (~90s total). It + # runs BEFORE we register/serve below, so the first real request per shape is + # steady-state and traffic routing (register_model) is gated on it; existing + # warm pods cover traffic during a new pod's boot (min-instances >= 1). The + # startup-probe budget in backend kubernetes_utils_async.py is sized to cover + # this. Opt out with FASTWAN_EAGER_WARM=0 for debugging only. + if os.environ.get("FASTWAN_EAGER_WARM", "1") != "0": + await backend.preflight() + + # Wire video_pool_* series into the Dynamo runtime's /metrics scrape. + # Registers unconditionally: in legacy in-process mode (VIDEO_POOL_MODE=0) + # the series stay zero-valued, which is harmless and means dashboards + # don't break when toggling the env var. + register_engine_metrics_callback( + endpoint=endpoint, + registry=VIDEO_REGISTRY, + metric_prefix_filters=["video_"], + namespace_name=namespace_name, + component_name=component_name, + endpoint_name=endpoint_name, + model_name=backend.served_model_name, + ) + + # Pool metrics flow through dynamo's system-status server, which is + # disabled by default. If DYN_SYSTEM_PORT isn't set to a non-negative + # value, the /metrics endpoint never binds and our registered callback + # sits dormant -- operators would see zero pool visibility in + # production. Warn loudly so the misconfiguration is obvious in pod + # startup logs. + _system_port = os.environ.get("DYN_SYSTEM_PORT", "-1") + try: + _system_port_int = int(_system_port) + except ValueError: + _system_port_int = -1 + if _system_port_int < 0: + logger.warning( + "DYN_SYSTEM_PORT is not set (or is %r); pool metrics will NOT " + "be exposed. Set DYN_SYSTEM_PORT to a port (e.g. 9090) and " + "DYN_SYSTEM_HOST=0.0.0.0 to enable the /metrics endpoint. " + "See examples/diffusers/ltx23/RUNBOOK.md § Metrics.", + _system_port, + ) + else: + logger.info( + "Pool metrics enabled on DYN_SYSTEM_PORT=%d", + _system_port_int, + ) + + try: + await asyncio.gather( + endpoint.serve_endpoint(backend.create_video), # type: ignore[arg-type] + register_model(endpoint, backend.served_model_name, backend.model_name), + ) + finally: + if backend.pool is not None: + logger.info("shutting down subprocess pool") + await backend.pool.shutdown() + + +def main_cli() -> None: + """ + FastWan-QAD-FP8 worker entry. Called by the top-level ``worker.py`` shim + when the invocation is NOT a pool-worker subprocess (pool-worker + dispatch happens in the shim itself, before any of these heavy + imports run). + + Contract with the shim: ``--pool-worker`` is already handled. + Do not re-dispatch here. + """ + args = _parse_args() + logging.basicConfig( + level=( + logging.DEBUG + if os.environ.get("FASTVIDEO_LOG_LEVEL") == "DEBUG" + else logging.INFO + ), + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + force=True, + ) + # Default the shapes-JSON path for both the in-container production + # location and a dev-run-from-source location. The pool subprocess + # also inherits this via env, so they see the same shape file. + if "WARMUP_SHAPES_JSON_PATH" not in os.environ: + if os.path.isfile(PRODUCTION_SHAPES_JSON_PATH): + os.environ["WARMUP_SHAPES_JSON_PATH"] = PRODUCTION_SHAPES_JSON_PATH + else: + os.environ["WARMUP_SHAPES_JSON_PATH"] = DEFAULT_SHAPES_JSON_PATH + # Refuse to start if image cache and shape menu disagree. Runs before + # any model load or networking so we fail fast with a clear message. + assert_shape_menu_hash_matches(os.environ["WARMUP_SHAPES_JSON_PATH"]) + uvloop.install() + asyncio.run(_main(args)) + + +if __name__ == "__main__": + # If invoked directly (e.g. `python3 -m fastwan.worker`), guard the + # entry: pool-worker dispatch lives in the top-level shim. Direct + # invocation with --pool-worker would skip that dispatch path. + if "--pool-worker" in sys.argv: + raise SystemExit( + "fastwan.worker invoked directly with --pool-worker; " + "use the top-level worker.py shim, which dispatches " + "to lib.pool._pool_worker_dispatch_if_requested first." + ) + main_cli()