[None][feat] Add FastWan2.2 TI2V-5B DMD pipeline (3-step text-to-video)

docs/source/models/visual-generation.md

@@ -39,6 +39,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
 | `nvidia/Cosmos3-Nano` | Text-to-Image, Text-to-Video, Image-to-Video |
 | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video |
+| `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | Text-to-Video |
 
 Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically.
 
@@ -51,6 +52,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 | **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
 | **Wan 2.1 VSA** [^2] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes |
 | **Wan 2.2** | Yes | Yes | Yes [^3] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| **FastWan 2.2** | Yes | Yes | No | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
 | **LTX-2** | Yes | Yes | Yes [^4] | Yes | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image** [^5] | Yes | Yes | No | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
 | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No |
@@ -65,6 +67,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 
 [^5]: Qwen-Image ships a native BF16 implementation with per-module numerical parity against `diffusers.QwenImagePipeline` (cosine similarity >= 0.999 on the full 20B transformer) and supports `trtllm-serve` / `/v1/images/generations`. VisualGen supports FP8 blockwise and NVFP4 dynamic quantization from BF16 checkpoints, as well as direct loading of statically quantized FP8 and NVFP4 ModelOpt checkpoints.
 
+
 ## Quick Start
 
 Here is a simple example to generate a video with Wan 2.1:

tensorrt_llm/_torch/visual_gen/models/__init__.py

@@ -37,7 +37,7 @@
 from .flux import Flux2Pipeline, FluxPipeline
 from .ltx2 import LTX2Pipeline  # noqa: F401
 from .qwen_image import QwenImagePipeline
-from .wan import WanImageToVideoPipeline, WanPipeline
+from .wan import WanDMDPipeline, WanImageToVideoPipeline, WanPipeline
 
 __all__ = [
     "AutoPipeline",
@@ -46,6 +46,7 @@
     "Flux2Pipeline",
     "QwenImagePipeline",
     "WanPipeline",
+    "WanDMDPipeline",
     "WanImageToVideoPipeline",
     "Cosmos3OmniMoTPipeline",
     "register_pipeline",

tensorrt_llm/_torch/visual_gen/models/wan/__init__.py

@@ -1,11 +1,13 @@
 from .parallel_vae import ParallelVAE_Wan
+from .pipeline_fastwan import WanDMDPipeline
 from .pipeline_wan import WanPipeline
 from .pipeline_wan_i2v import WanImageToVideoPipeline
 from .transformer_wan import WanTransformer3DModel
 from .wan_vae import WanVAE, WanVAEConfig
 
 __all__ = [
     "WanPipeline",
+    "WanDMDPipeline",
     "WanImageToVideoPipeline",
     "WanTransformer3DModel",
     "ParallelVAE_Wan",

tensorrt_llm/_torch/visual_gen/models/wan/defaults.py

@@ -131,6 +131,18 @@ def get_wan_default_params(
     return params
 
 
+def get_fastwan_default_params() -> dict:
+    """Default generation params for FastWan 2.2 TI2V-5B.
+
+    Same as the base Wan 2.2 TI2V-5B defaults, but distilled: a fixed 3-step
+    DMD schedule and CFG-free (guidance_scale 1.0).
+    """
+    params = dict(_WAN22_5B_PARAMS)
+    params["num_inference_steps"] = 3
+    params["guidance_scale"] = 1.0
+    return params
+
+
 def get_wan_extra_param_specs(is_wan22_14b: bool) -> Dict[str, ExtraParamSchema]:
     """Return extra_param_specs for a Wan model.
 

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py

@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""FastWan: distilled (DMD) Wan 2.2 TI2V-5B for 3-step text-to-video.
+
+FastWan shares Wan 2.2 TI2V-5B's architecture exactly — only the weights
+(distilled) and the sampling recipe differ. So this pipeline subclasses
+``WanPipeline`` and overrides only the denoising loop.
+"""
+
+import time
+
+import PIL.Image
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput
+from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline
+from tensorrt_llm._utils import nvtx_range
+from tensorrt_llm.logger import logger
+
+from .defaults import get_fastwan_default_params
+from .pipeline_wan import WanPipeline
+
+
+@register_pipeline(
+    "WanDMDPipeline",
+    hf_ids=[
+        "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+    ],
+    doc="FastWan 2.2 distilled (DMD) — 3-step Wan 2.2 TI2V-5B text-to-video.",
+)
+class WanDMDPipeline(WanPipeline):
+    """Wan 2.2 TI2V-5B with the DMD 3-step sampling loop.
+
+    DMD inference (Distribution Matching Distillation): at each step the model
+    predicts the clean video ``x0``, then we re-noise it with FRESH noise to the
+    next (lower) noise level. The last step keeps ``x0``. CFG-free — one
+    transformer forward per step. Faithful to FastVideo's ``DmdDenoisingStage``.
+    """
+
+    # Fixed DMD schedule for FastWan2.2-TI2V-5B (see model card / FastVideo config).
+    DMD_TIMESTEPS = (1000, 757, 522)
+    NUM_TRAIN_TIMESTEPS = 1000
+
+    @property
+    def default_generation_params(self):
+        """Return FastWan-specific default generation parameters (steps, guidance, frame rate)."""
+        return get_fastwan_default_params()
+
+    @nvtx_range("WanDMDPipeline.forward")
+    @torch.no_grad()
+    def forward(
+        self,
+        prompt: str | list[str],
+        seed: int,
+        negative_prompt: str | None = None,
+        height: int = 704,
+        width: int = 1280,
+        num_frames: int = 121,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | None = None,
+        guidance_scale_2: float | None = None,
+        boundary_ratio: float | None = None,
+        max_sequence_length: int = 512,
+        image: PIL.Image.Image | torch.Tensor | str | None = None,
+    ) -> PipelineOutput:
+        """Generate a video from a text prompt using the 3-step DMD sampling loop.
+
+        CFG-free: one transformer forward pass per step, no negative prompt support.
+        Unsupported parameters (num_inference_steps, guidance_scale, guidance_scale_2,
+        boundary_ratio) are accepted for base-class compatibility but logged as warnings
+        and ignored.
+        """
+        if image is not None:
+            raise NotImplementedError(
+                "FastWan currently supports text-to-video only (no image conditioning)."
+            )
+        if num_inference_steps not in (None, len(self.DMD_TIMESTEPS)):
+            logger.warning(
+                f"FastWan always uses {len(self.DMD_TIMESTEPS)} DMD steps; "
+                f"ignoring num_inference_steps={num_inference_steps}."
+            )
+        if guidance_scale not in (None, 1.0):
+            logger.warning(f"FastWan is CFG-free; ignoring guidance_scale={guidance_scale}.")
+        if guidance_scale_2 is not None:
+            logger.warning(
+                f"FastWan does not support guidance_scale_2; ignoring guidance_scale_2={guidance_scale_2}."
+            )
+        if boundary_ratio is not None:
+            logger.warning(
+                f"FastWan does not support boundary_ratio; ignoring boundary_ratio={boundary_ratio}."
+            )
+
+        pipeline_start = time.time()
+        timer = CudaPhaseTimer()
+        timer.mark_pre_start()
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        batch_size = len(prompt)
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        self.validate_resolution(height, width, num_frames)
+
+        if negative_prompt:
+            raise ValueError("FastWan is CFG-free and does not support negative prompts.")
+        prompt_embeds, _ = self._encode_prompt(prompt, "", max_sequence_length)
+
+        if self._fixed_latent is not None:
+            latents = self._fixed_latent.to(device=self.device, dtype=self.dtype)
+        else:
+            latents = self._prepare_latents(batch_size, height, width, num_frames, generator)
+
+        timer.mark_denoise_start()
+        latents = self._denoise(latents, prompt_embeds, generator)
+        timer.mark_post_start()
+
+        logger.info("Decoding video...")
+        decode_start = time.time()
+        video = self.decode_latents(latents, self._decode_latents)
+        if self.rank == 0:
+            logger.info(f"Video decoded in {time.time() - decode_start:.2f}s")
+            logger.info(f"Total pipeline time: {time.time() - pipeline_start:.2f}s")
+
+        timer.mark_end()
+        frame_rate = self.default_generation_params.get("frame_rate", 24.0)
+        return timer.fill(PipelineOutput(video=video, frame_rate=frame_rate))
+
+    @nvtx_range("WanDMDPipeline._denoise", color="blue")
+    def _denoise(
+        self, latents: torch.Tensor, prompt_embeds: torch.Tensor, generator: torch.Generator
+    ) -> torch.Tensor:
+        """3-step DMD loop: predict x0, then re-noise with fresh noise.
+
+        The per-step sigma reduces to ``sigma = t / NUM_TRAIN_TIMESTEPS`` because
+        FastVideo's scheduler keeps ``timesteps = sigmas * num_train_timesteps``
+        (co-scaled), so the lookup returns ``matched_timestep / 1000`` and the
+        flow-shift cancels (verified numerically for shift in {1, 5, 8, 12}).
+        """
+        timesteps = self.DMD_TIMESTEPS
+        num_steps = len(timesteps)
+        _, ph, pw = self.transformer.config.patch_size
+
+        # The DMD ops are linear (mul/sub/add), so they run in the transformer's
+        # native (bf16) dtype to match the FastVideo reference; no fp32 needed.
+        latents = latents.to(self.dtype)
+        # Patch grid dimensions are fixed for the entire denoising loop.
+        nf = latents.shape[2]
+        nh = latents.shape[3] // ph
+        nw = latents.shape[4] // pw
+        start = time.time()
+        for i, t in enumerate(timesteps):
+            t_tensor = torch.full((latents.shape[0], nf * nh * nw), float(t), device=latents.device)
+
+            pred_noise = self.transformer(
+                hidden_states=latents,
+                timestep=t_tensor / self.NUM_TRAIN_TIMESTEPS,
+                encoder_hidden_states=prompt_embeds,
+            )
+
+            sigma = t / self.NUM_TRAIN_TIMESTEPS
+            pred_video = latents - sigma * pred_noise  # clean-video estimate (x0)
+
+            if i < num_steps - 1:
+                sigma_next = timesteps[i + 1] / self.NUM_TRAIN_TIMESTEPS
+                noise = randn_tensor(
+                    latents.shape,
+                    generator=generator,
+                    device=latents.device,
+                    dtype=self.dtype,
+                )
+                latents = (1.0 - sigma_next) * pred_video + sigma_next * noise
+            else:
+                latents = pred_video
+
+            if self.rank == 0:
+                logger.info(f"DMD step {i + 1}/{num_steps} (t={t}, sigma={sigma:.3f})")
+
+        if self.rank == 0:
+            logger.info(f"DMD denoising done in {time.time() - start:.2f}s")
+        return latents

tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fastwan_lpips_golden_video.json

@@ -0,0 +1,16 @@
+{
+  "video": "fastwan_lpips_golden_video.mp4",
+  "model": "FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+  "source": "FastVideo VideoGenerator (trusted reference)",
+  "prompt": "A red fox walking through snow",
+  "negative_prompt": "",
+  "height": 704,
+  "width": 1280,
+  "num_frames": 121,
+  "num_inference_steps": 3,
+  "guidance_scale": 1.0,
+  "seed": 42,
+  "frame_rate": 24.0,
+  "lpips_net": "alex",
+  "lpips_threshold": 0.05
+}

tests/integration/defs/examples/visual_gen/test_visual_gen.py

@@ -120,6 +120,17 @@
 COSMOS3_LPIPS_FRAME_RATE = 24.0
 COSMOS3_LPIPS_THRESHOLD = 0.05
 
+# FastWan 2.2 TI2V-5B (DMD distilled) LPIPS configuration
+FASTWAN_LPIPS_PROMPT = "A red fox walking through snow"
+FASTWAN_LPIPS_NEGATIVE_PROMPT = ""
+FASTWAN_LPIPS_HEIGHT = 704
+FASTWAN_LPIPS_WIDTH = 1280
+FASTWAN_LPIPS_NUM_FRAMES = 121
+FASTWAN_LPIPS_NUM_INFERENCE_STEPS = 3
+FASTWAN_LPIPS_GUIDANCE_SCALE = 1.0
+FASTWAN_LPIPS_SEED = 42
+FASTWAN_LPIPS_FRAME_RATE = 24.0
+FASTWAN_LPIPS_THRESHOLD = 0.05
 # LTX-2 configuration
 LTX2_MODEL_CHECKPOINT_PATH = "LTX-2/ltx-2-19b-dev.safetensors"
 LTX2_TEXT_ENCODER_SUBPATH = "gemma-3-12b-it"
@@ -1134,6 +1145,42 @@ def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path):
     _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_fastwan_lpips_against_golden(tmp_path):
+    generated_path = tmp_path / "fastwan_generated.mp4"
+    golden_path = _golden_media_path(
+        tmp_path, "fastwan_lpips_golden_video.mp4", "FastWan LPIPS golden video"
+    )
+    _generate_wan_lpips_video(
+        _lpips_model_path("FastWan2.2-TI2V-5B-FullAttn-Diffusers"),
+        generated_path,
+        FASTWAN_LPIPS_PROMPT,
+        FASTWAN_LPIPS_NEGATIVE_PROMPT,
+        FASTWAN_LPIPS_HEIGHT,
+        FASTWAN_LPIPS_WIDTH,
+        FASTWAN_LPIPS_NUM_FRAMES,
+        FASTWAN_LPIPS_NUM_INFERENCE_STEPS,
+        FASTWAN_LPIPS_GUIDANCE_SCALE,
+        FASTWAN_LPIPS_SEED,
+        FASTWAN_LPIPS_FRAME_RATE,
+    )
+    score = _run_lpips_eval(
+        tmp_path,
+        "fastwan",
+        "video",
+        FASTWAN_LPIPS_PROMPT,
+        golden_path,
+        generated_path,
+    )
+    _assert_lpips_below_threshold(score, FASTWAN_LPIPS_THRESHOLD)
+
+
+@pytest.fixture(scope="session")
+def wan_trtllm_video_path(_visual_gen_deps, llm_venv, llm_root):
+    """Generate input video via models/wan_t2v.py and return path to trtllm_output.mp4."""
+    return _generate_wan_video(llm_venv, llm_root, WAN_T2V_MODEL_SUBPATH, "wan")
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 def test_cosmos3_nano_t2v_lpips_against_golden(_visual_gen_deps, tmp_path):
     generated_path = tmp_path / "cosmos3_nano_t2v_generated.mp4"

tests/integration/test_lists/test-db/l0_b200.yml

@@ -251,6 +251,8 @@ l0_b200:
   - unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py
   - unittest/_torch/visual_gen/test_wan_vsa_pipeline.py
   - unittest/_torch/visual_gen/test_wan_vae.py
+  - unittest/_torch/visual_gen/test_fastwan_dmd_math.py
+  - unittest/_torch/visual_gen/test_fastwan_pipeline.py
   - unittest/_torch/visual_gen/test_wan21_i2v_teacache.py
   - unittest/_torch/visual_gen/test_wan21_t2v_teacache.py
   - unittest/_torch/visual_gen/test_wan21_t2v_teacache_user_coefficients.py
@@ -370,6 +372,7 @@ l0_b200:
   - examples/visual_gen/test_visual_gen.py::test_qwenimage_cuda_graph_lpips_against_golden TIMEOUT (10)
   - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden TIMEOUT (10)
   - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15)
+  - examples/visual_gen/test_visual_gen.py::test_fastwan_lpips_against_golden TIMEOUT (10)
   - visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark
   - visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos]
   # ---- moved to post-merge (MoE CI optimization) ----