feat: fastwan22_5b family — FastWan2.2-TI2V-5B-FullAttn (recommended recipe)

examples/diffusers/Dockerfile

@@ -79,6 +79,25 @@ RUN export FLASH_ATTN_CUDA_ARCHS=$(echo "${TORCH_CUDA_ARCH_LIST}" | sed 's/ /;/g
  && uv pip install ./flash_attn/cute \
  && rm -rf ../flash-attention
 
+# FA4 (flash-attn-cute) source swap: the XOR-op fork installed above is stale
+# against cutlass-dsl 4.5+ (cute.core.ThrMma renamed away -> AttributeError at
+# import, which kills any worker that selects FLASH_ATTN). FastVideo hit the
+# same break and fixed it upstream (hao-ai-lab/FastVideo#1564) by pinning
+# flash-attn-4 from Dao-AILab/flash-attention at a CuTe-DSL-4.6-compatible
+# commit. Reproduce that exact pin (same as Dockerfile.dreamverse); do not
+# bump the ref without re-validating.
+ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+RUN . /opt/dynamo/venv/bin/activate \
+ && uv pip uninstall flash-attn-cute \
+ && git clone https://github.com/Dao-AILab/flash-attention.git /tmp/fa-upstream \
+ && cd /tmp/fa-upstream && git checkout ${FA4_CUTE_REF} \
+ && cd / \
+ && uv pip install --no-deps --no-build-isolation /tmp/fa-upstream/flash_attn/cute \
+ && rm -rf /tmp/fa-upstream
+# --no-deps: the pin declares nvidia-cutlass-dsl==4.6.0.dev0 (pre-release),
+# which uv refuses and which would needlessly downgrade the venv's 4.6.0
+# release; the environment already satisfies the real requirement.
+
 # Install Dynamo with /v1/videos support (our serving layer).
 RUN uv pip install ai-dynamo==1.1.1
 
@@ -98,6 +117,12 @@ COPY patches/ltx23_gpu_worker_megacache.patch /tmp/megacache.patch
 # default 32, env-overridable via FASTVIDEO_X264_THREADS. See PR description.
 COPY patches/x264-threads-cap.patch /tmp/x264-threads-cap.patch
 
+# The upstream FA4 pin returns 4-tuples from _flash_attn_fwd; fastvideo's
+# wrapper unpacks 2. Same [:2] fix FastVideo ships today (see
+# Dockerfile.dreamverse, where this patch has been applied since the LTX-2.3
+# FA4 work).
+COPY patches/flash-attn-cute-fa4-tuple-fix.patch /tmp/flash-attn-cute-fa4-tuple-fix.patch
+
 # Install FastVideo at the pinned SHA. No pyproject sed needed: the base is now
 # cu128, which is exactly what FastVideo's [tool.uv.sources] targets, so the
 # torch family resolves natively without any index gymnastics.
@@ -110,6 +135,7 @@ RUN echo "Building FastVideo @ ${FASTVIDEO_SHA} for TORCH_CUDA_ARCH_LIST=${TORCH
  && git submodule update --init --recursive \
  && git apply --verbose /tmp/megacache.patch \
  && git apply --verbose /tmp/x264-threads-cap.patch \
+ && git apply --verbose /tmp/flash-attn-cute-fa4-tuple-fix.patch \
  && cd / \
  && uv pip install --no-build-isolation /tmp/FastVideo \
  && rm -rf /tmp/FastVideo
@@ -129,6 +155,21 @@ COPY . /opt/app/
 # 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.
+# FastWan2.2-5B decodes with the taew2_2_super TAEHV tiny autoencoder (the
+# ~3s path on B200) instead of the full Wan2.2 VAE. Pinned to a taehv commit
+# because the HF mirror went private and main is a moving ref. Gated so
+# other-family images don't carry it; build with --build-arg
+# INCLUDE_FASTWAN22_5B_TAEHV=1.
+ARG INCLUDE_FASTWAN22_5B_TAEHV=0
+ARG TAEHV_REF=093b918971d59001a0bad6dfd6e0409b5e1752cf
+RUN if [ "$INCLUDE_FASTWAN22_5B_TAEHV" = "1" ]; then \
+      curl -fsSL "https://raw.githubusercontent.com/madebyollin/taehv/${TAEHV_REF}/taehv.py" \
+        -o /opt/app/fastwan22_5b/taehv.py && \
+      curl -fsSL "https://raw.githubusercontent.com/madebyollin/taehv/${TAEHV_REF}/taew2_2_super.pth" \
+        -o /opt/app/fastwan22_5b/taew2_2_super.pth && \
+      test -s /opt/app/fastwan22_5b/taew2_2_super.pth ; \
+    fi
+
 ARG INCLUDE_FASTWAN_TAEHV=0
 RUN if [ "$INCLUDE_FASTWAN_TAEHV" = "1" ]; then \
       . /opt/dynamo/venv/bin/activate && \

examples/diffusers/fastwan22_5b/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""FastWan2.2-TI2V-5B 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.
+"""

examples/diffusers/fastwan22_5b/factory.py

@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""FastWan2.2-TI2V-5B-FullAttn model factory (optimized: tiny decoder + compile).
+
+``load_model`` is the SINGLE source of truth for how this family's
+``VideoGenerator`` is constructed, used by both the in-process path
+(``lib.backend``) and the pool path (``lib.pool`` resolves it via
+``--model-factory fastwan22_5b.factory:load_model``).
+
+Config (Johan quality-approved 2026-08-03 via seed-locked split-screen A/B
+vs the full-VAE recipe -- "equally good to the eye"; ~3x faster/cheaper):
+  * bf16 -- the model's native precision; NO quantization;
+  * **TAEHV tiny autoencoder (taew2_2_super) for decode** instead of the full
+    Wan2.2 VAE: generator built with ``output_type="latent"`` and the heavy
+    VAE CPU-offloaded; latents decoded by the wrapper below. B200 measured
+    ~3.0s/clip vs 9.8s full-VAE;
+  * torch.compile ON (quality-neutral, ~19% on the bench);
+  * 3 denoise steps, no CFG (per-request, matching the DMD schedule);
+  * ``pipeline_config`` is PINNED, see below.
+"""
+
+import logging
+import os
+import sys
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Wan2.2-VAE tiny decoder (48-ch; taew2_1 does NOT fit this model family),
+# baked into the image (Dockerfile).
+TAEHV_CKPT = os.environ.get(
+    "FASTWAN22_5B_TAEHV_CKPT", "/opt/app/fastwan22_5b/taew2_2_super.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. width/height/
+    num_frames ARE plumbed into the latent request (via SamplingParam) so each
+    pool subprocess renders its actual shape -- without this, portrait
+    704x1280 silently comes out landscape (the QAD bug)."""
+
+    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 = 24,
+        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,
+        }
+        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 FastWan2.2-TI2V-5B generator (bf16 + compile + TAEHV decode).
+
+    ``enable_optimizations`` is accepted for the shared pool/warmup factory
+    signature but intentionally unused (the optimized recipe is always on).
+    """
+    from fastvideo import VideoGenerator
+    from fastvideo.configs.pipelines.wan import FastWan2_2_TI2V_5B_Config
+
+    del enable_optimizations
+
+    # The serving mount anonymizes the weights path to /data/default, which
+    # defeats fastvideo's path-based preset resolution -- and this model's
+    # model_index.json ``_class_name`` (WanDMDPipeline) then matches the
+    # FastWan2.1 detector, silently selecting the 480p Wan2.1 config.
+    # Pinning the config makes the selection explicit; the check makes any
+    # future rename/refactor of the config class loud instead of silent.
+    # (Empirically verified 2026-08-03: RESOLVED_CONFIG=FastWan2_2_TI2V_5B_Config
+    # flow_shift=5.0 dmd=[1000, 757, 522] at /data/default.)
+    pipeline_config = FastWan2_2_TI2V_5B_Config()
+    if not getattr(pipeline_config, "dmd_denoising_steps", None):
+        raise RuntimeError(
+            "FastWan2_2_TI2V_5B_Config has no dmd_denoising_steps; refusing "
+            "to serve a distill checkpoint with a non-distill config"
+        )
+    logger.info(
+        "fastwan22_5b: pinned pipeline_config=%s flow_shift=%s dmd_steps=%s "
+        "attention_backend=%s taehv_ckpt=%s",
+        type(pipeline_config).__name__,
+        getattr(pipeline_config, "flow_shift", None),
+        getattr(pipeline_config, "dmd_denoising_steps", None),
+        os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "<default>"),
+        TAEHV_CKPT,
+    )
+
+    gen = VideoGenerator.from_pretrained(
+        model_path,
+        num_gpus=num_gpus,
+        pipeline_config=pipeline_config,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        # output_type="latent" bypasses the full Wan2.2 VAE (TAEHV decodes
+        # instead), so the heavy VAE is 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",
+    )
+    taehv = _load_taehv()
+    logger.info("fastwan22_5b: TAEHV tiny decoder loaded from %s", TAEHV_CKPT)
+    return _TaehvVideoGenerator(gen, taehv)

examples/diffusers/fastwan22_5b/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 fastwan22_5b/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 / fastwan22_5b packages
+# resolve. This script lives at examples/diffusers/fastwan22_5b/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/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+        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 fastwan22_5b.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="fastwan22_5b.factory:load_model",
+        model_label="fastwan22-ti2v-5b",
+    )
+
+    # 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())