[TRTLLM-16673][feat] Add FastH3 support with VSA

docs/source/models/supported-models.md

@@ -289,6 +289,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | `Lightricks/LTX-2` | Text-to-Video (with Audio), Image-to-Video (with Audio) |
 | `Lightricks/LTX-2.3` | Text-to-Video (with Audio), Video-to-Video (with Audio) |
 | `MiniMaxAI/MiniMax-H3` | Text-to-Video, First/Last-Frame-to-Video, Reference-to-Video (all with Audio) |
+| `FastVideo/FastVideo-FastH3-8-Step-V2` | Text-to-Video (with Audio, 8-step distilled) |
 | `Qwen/Qwen-Image` | Text-to-Image |
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
 | `Qwen/Qwen-Image-Layered` | Image-to-Image |
@@ -310,6 +311,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 | **LTX-2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No |
 | **LTX-2.3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | No | No | No | No |
 | **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes |
+| **FastH3** | Yes | Yes | No | No | No | No | No | Yes | Yes | No | No | No | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
 | **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No |

docs/source/models/visual-generation.md

@@ -38,6 +38,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode
 | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | Text-to-Video (3-step distilled) |
 | `Lightricks/LTX-2` | Text-to-Video (with Audio), Image-to-Video (with Audio) |
 | `Lightricks/LTX-2.3` | Text-to-Video (with Audio), Video-to-Video (with Audio) |
+| `FastVideo/FastVideo-FastH3-8-Step-V2` | Text-to-Video (with Audio, 8-step distilled) |
 | `Qwen/Qwen-Image` | Text-to-Image |
 | `Qwen/Qwen-Image-2512` | Text-to-Image |
 | `Qwen/Qwen-Image-Layered` | Image-to-Image |
@@ -68,6 +69,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 | **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
 | **LTX-2.3** | Yes | Yes | No | No | No | No | Yes | Yes | Yes | Yes | Yes | No | No | No | No | No |
 | **MiniMax-H3** | Yes | Yes | No | No | Yes | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | No |
+| **FastH3** | Yes | Yes | No | No | No | No | No | No | No | Yes | Yes | No | No | No | No | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | No |
 | **Qwen-Image-Layered** | Yes | No | No | Yes | No | No | No | No | Yes | Yes | Yes | No | No | No | No | No |
 | **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | No |
@@ -190,6 +192,22 @@ are not supported. See the
 - The published [MiniMax-H3 checkpoint license](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE)
   restricts use by territory. Obtain legal approval before downloading or running the weights.
 
+- FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a
+  distilled 8-step V2 checkpoint. Model-specific VSA is the supported inference path
+  for this checkpoint and is selected by default with the TRTLLM backend and the
+  checkpoint's sparsity (80% for V2). Explicit incompatible backends
+  or sparse attention algorithms are rejected. Dense attention omits the trained
+  compression gate and is not equivalent to this checkpoint's VSA inference.
+  The example reads nine sigma-grid points (eight transformer forwards) from
+  `fastvideo_inference.json` and uses its trained DMD ladder with video/audio
+  scheduler shifts of 10/3. Changing only the step count is not equivalent.
+  For single-GPU text-to-video with audio, see
+  `fasth3-bf16-trtllm-vsa-1gpu.yaml`. To use CUTEDSL instead, set
+  `attention_config.backend: CUTEDSL` in that config.
+  Both backends preserve the trained compression gate and dense text/audio routes.
+  CUTEDSL VSA requires supported Blackwell hardware and at most 4096 tiles;
+  unsupported shapes fail explicitly rather than silently running dense attention.
+
 ## Optimizations
 
 ### Quantization

examples/visual_gen/README.md

@@ -29,6 +29,8 @@ python models/ltx2.py --visual_gen_args configs/ltx2-1gpu.yaml
 # 2-GPU DiT tensor parallelism config: configs/minimax-h3-bf16-tp2.yaml.
 # 2-GPU Ulysses config: configs/minimax-h3-bf16-2gpu.yaml.
 python models/minimax_h3.py --model <approved-checkpoint> --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml
+# FastH3 8-Step V2 VSA on Blackwell:
+python models/fasth3.py --visual_gen_args configs/fasth3-bf16-trtllm-vsa-1gpu.yaml
 python models/flux1.py --visual_gen_args configs/flux1-dev-fp4-1gpu.yaml
 python models/flux2.py --visual_gen_args configs/flux2-dev-fp4-1gpu.yaml
 python models/cosmos3/cosmos3.py --visual_gen_args configs/cosmos3-nano-1gpu.yaml --prompt "A robot arm picks fruit in a grocery store"

examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+cuda_graph_config:
+  enable: false
+torch_compile_config:
+  enable: false
+attention_config:
+  backend: TRTLLM
+  sparse_attention_config:
+    algorithm: vsa
+    vsa_sparsity: 0.8
+compilation_config:
+  skip_warmup: true
+  resolutions:
+    - [768, 1344]
+  num_frames:
+    - 124

examples/visual_gen/models/fasth3.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+# 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.
+
+"""FastH3 text-to-video generation with stereo audio."""
+
+import argparse
+
+from tensorrt_llm import VisualGen, VisualGenArgs
+
+FASTH3_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--model",
+        default=FASTH3_MODEL_ID,
+        help="FastH3 checkpoint path or Hugging Face model ID.",
+    )
+    parser.add_argument(
+        "--visual_gen_args",
+        default="configs/fasth3-bf16-trtllm-vsa-1gpu.yaml",
+        help="Path to YAML config (same as trtllm-serve --visual_gen_args).",
+    )
+    parser.add_argument(
+        "--output_path",
+        default="fasth3_t2va_output.mp4",
+        help="Path to save the generated video and audio.",
+    )
+    parser.add_argument("--prompt", default=None, help="Prompt.")
+    args = parser.parse_args()
+
+    extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None
+    visual_gen = VisualGen(model=args.model, args=extra_args)
+    params = visual_gen.default_params
+    params.height = 768
+    params.width = 1344
+    params.num_frames = 124
+    params.seed = 42
+
+    output = visual_gen.generate(
+        inputs=args.prompt
+        or (
+            "A cinematic wide shot of an astronaut walking through a neon-lit "
+            "rainforest at dusk, with soft ambient music and detailed natural sounds."
+        ),
+        params=params,
+    )
+    saved = output.save(args.output_path)
+    print(f"Saved: {saved}")
+
+
+if __name__ == "__main__":
+    main()

tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py

@@ -17,10 +17,12 @@
 
 """TRT-LLM VisualGen pipeline for MiniMax-H3 FL2VA and Ref2VA checkpoints."""
 
+import json
 import time
 from collections.abc import Iterator
 from contextlib import contextmanager, nullcontext
 from io import BytesIO
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
@@ -32,7 +34,10 @@
 from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor
 
 from tensorrt_llm._torch.visual_gen.checkpoints import WeightLoader
-from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig
+from tensorrt_llm._torch.visual_gen.config import (
+    DiffusionPipelineConfig,
+    create_attention_metadata_state,
+)
 from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput
 from tensorrt_llm._torch.visual_gen.pipeline import (
     BasePipeline,
@@ -45,6 +50,8 @@
 from tensorrt_llm.inputs.utils import load_image
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import MpiTopology
+from tensorrt_llm.visual_gen.args import AttentionConfig
+from tensorrt_llm.visual_gen.sparse_attention import VideoSparseAttentionConfig
 
 from .packing import (
     MINIMAX_H3_AUDIO_CHANNELS,
@@ -75,6 +82,7 @@
 from .ref2va import load_references, prepare_references, validate_reference_order
 from .text_encoder_tp import MiniMaxH3TensorParallelTextEncoder
 from .transformer_minimax_h3 import MiniMaxH3Transformer3DModel
+from .vsa import build_h3_vsa_metadata, set_h3_vsa_context
 
 if TYPE_CHECKING:
     from ...sleep import PipelineSleepManager, SleepRestoreMode
@@ -164,10 +172,16 @@ def tp_group_name(self) -> str:
 
 @register_pipeline(
     "MiniMaxH3ModularPipeline",
-    hf_ids=["MiniMaxAI/MiniMax-H3"],
+    hf_ids=[
+        "MiniMaxAI/MiniMax-H3",
+        "FastVideo/FastVideo-FastH3-8-Step-V2",
+    ],
     defaults={"workflow": "fl2va"},
     download_patterns=[
         "modular_model_index.json",
+        "fastvideo_inference.json",
+        "checkpoint_content.json",
+        "checkpoint_metadata.json",
         "LICENSE",
         "README.md",
         "transformer/*",
@@ -203,6 +217,7 @@ def _request_generator(seed: int) -> torch.Generator:
         return torch.Generator().manual_seed(seed)
 
     def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None:
+        self._configure_fasth3_attention(pipeline_config)
         self.text_encoder_tp_size = int(
             getattr(getattr(pipeline_config, "parallel", None), "text_encoder_tp_size", 1)
         )
@@ -371,14 +386,15 @@ def setup_parallel_vae(self) -> None:
 
     @property
     def default_generation_params(self) -> dict:
+        fastvideo_inference = getattr(self, "_fastvideo_inference", {})
         return {
             # MiniMax-H3 derives the default canvas from the first present
             # keyframe, or from 16:9 for T2VA.
             "height": None,
             "width": None,
             "num_frames": 124,
             "frame_rate": float(MINIMAX_H3_FPS),
-            "num_inference_steps": 28,
+            "num_inference_steps": int(fastvideo_inference.get("num_inference_steps", 28)),
         }
 
     @property
@@ -454,6 +470,46 @@ def _run_warmup(
             num_inference_steps=steps,
         )
 
+    @staticmethod
+    def _load_fasth3_manifest(checkpoint_dir: str) -> dict:
+        manifest_path = Path(checkpoint_dir) / "fastvideo_inference.json"
+        if manifest_path.exists():
+            with manifest_path.open() as f:
+                manifest = json.load(f)
+            if str(manifest.get("schema_version", "")).startswith("fasth3-"):
+                return manifest
+        return {}
+
+    @classmethod
+    def _configure_fasth3_attention(cls, config: DiffusionPipelineConfig) -> None:
+        pretrained = getattr(config, "primary_pretrained_config", None)
+        checkpoint_dir = getattr(pretrained, "_name_or_path", None)
+        manifest = cls._load_fasth3_manifest(checkpoint_dir) if checkpoint_dir else {}
+        if not manifest:
+            return
+
+        attention = config.attention
+        backend = attention.backend
+        if backend == "VANILLA" and "backend" not in attention.model_fields_set:
+            backend = "TRTLLM"
+        if backend not in ("TRTLLM", "CUTEDSL"):
+            raise ValueError("FastH3 requires VSA with the TRTLLM or CUTEDSL attention backend.")
+        sparse = attention.sparse_attention_config
+        if sparse is None:
+            sparse = VideoSparseAttentionConfig(vsa_sparsity=manifest.get("vsa_sparsity", 0.8))
+        elif sparse.algorithm != "vsa":
+            raise ValueError("FastH3 requires the vsa sparse attention algorithm.")
+        config.attention = AttentionConfig.model_validate(
+            {**attention.model_dump(), "backend": backend, "sparse_attention_config": sparse}
+        )
+        # Resolve defaults before constructing the transformer and loading its trained gate.
+        if backend == "TRTLLM" and config.attention_metadata_state is None:
+            config.attention_metadata_state = create_attention_metadata_state()
+        for model_config in config.model_configs.values():
+            model_config.attention = config.attention.model_copy(deep=True)
+            if backend == "TRTLLM" and model_config.attention_metadata_state is None:
+                model_config.attention_metadata_state = create_attention_metadata_state()
+
     def _init_transformer(self) -> None:
         self.transformer = MiniMaxH3Transformer3DModel(
             self.pipeline_config.model_configs[PipelineComponent.TRANSFORMER]
@@ -505,6 +561,8 @@ def load_standard_components(
         skip_components = skip_components or []
         offloaded = self.offloader.requested_components()
 
+        self._fastvideo_inference = self._load_fasth3_manifest(checkpoint_dir)
+
         if not _component_skipped(skip_components, PipelineComponent.TOKENIZER):
             self.tokenizer = Qwen2TokenizerFast.from_pretrained(
                 checkpoint_dir,
@@ -563,13 +621,43 @@ def load_standard_components(
 
     def post_load_weights(self) -> None:
         super().post_load_weights()
-        if self.scheduler is not None:
+        # Distilled checkpoints carry their trained shifts in the scheduler configs.
+        if self.scheduler is not None and not getattr(self, "_fastvideo_inference", {}):
             self.scheduler.set_shift(self.VIDEO_SCHEDULER_SHIFT)
-        if self.audio_scheduler is not None:
+        if self.audio_scheduler is not None and not getattr(self, "_fastvideo_inference", {}):
             self.audio_scheduler.set_shift(self.AUDIO_SCHEDULER_SHIFT)
         if self.vae is not None and self.device.type == "cuda":
             _prepare_vae_decoder(self.vae)
 
+    def _set_timesteps(self, num_inference_steps: int) -> None:
+        manifest = getattr(self, "_fastvideo_inference", {})
+        steps = manifest.get("dmd_denoising_steps")
+        if steps is None:
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+            self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device)
+            return
+        if (
+            not steps
+            or any(type(step) is not int or not 0 < step <= 1000 for step in steps)
+            or any(left <= right for left, right in zip(steps, steps[1:]))
+        ):
+            raise ValueError("FastH3 DMD steps must be strictly decreasing integers in (0, 1000].")
+        if num_inference_steps != len(steps) + 1:
+            raise ValueError(
+                f"FastH3 requires {len(steps) + 1} sigma-grid points for {len(steps)} forwards."
+            )
+        base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32)
+        for scheduler, shift_key in (
+            (self.scheduler, "video_scheduler_shift"),
+            (self.audio_scheduler, "audio_scheduler_shift"),
+        ):
+            shift = float(scheduler.shift)
+            if shift_key in manifest and shift != float(manifest[shift_key]):
+                raise ValueError(f"FastH3 {shift_key} disagrees with the checkpoint scheduler.")
+            # Shift each modality's trained noise clock once; explicit sigmas are not shifted.
+            sigmas = shift * base / (1 + (shift - 1) * base)
+            scheduler.set_timesteps(sigmas=sigmas, device=self.device)
+
     def torch_compile(self) -> None:
         super().torch_compile()
         # The ViT VAE decoder is 36 identical blocks run once per clip; compile them like the transformer blocks.
@@ -1211,8 +1299,7 @@ def _forward(
             condition_latents=condition_latents,
         )
 
-        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
-        self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device)
+        self._set_timesteps(num_inference_steps)
         # The denoise loop indexes both schedules with one counter, and
         # set_timesteps collapses float32 duplicates, so unequal lengths would
         # silently truncate a stream instead of failing.
@@ -1257,6 +1344,31 @@ def _forward(
             self.scheduler.timesteps, self.audio_scheduler.timesteps
         )
 
+        _attn_cfg = self.pipeline_config.primary_model_config.attention
+        _sparse_cfg = getattr(_attn_cfg, "sparse_attention_config", None)
+        _vsa_active = (
+            getattr(_attn_cfg, "backend", "VANILLA") in ("CUTEDSL", "TRTLLM")
+            and _sparse_cfg is not None
+            and getattr(_sparse_cfg, "algorithm", None) == "vsa"
+        )
+        _vsa_sparsity = _sparse_cfg.vsa_sparsity if _vsa_active else 0.0
+        vsa_metadata = None
+        if _vsa_active:
+            if self.workflow == "ref2va" or keyframes:
+                raise ValueError("MiniMax-H3 VSA currently supports text-to-video only.")
+            patch = self.transformer.config.patch_size
+            grid = tuple(
+                size // p
+                for size, p in zip(
+                    (num_latent_frames, latent_height, latent_width), patch, strict=True
+                )
+            )
+            vsa_metadata = build_h3_vsa_metadata(
+                (text_indices.numel(), audio_indices.numel(), grid),
+                _vsa_sparsity,
+                self.device,
+            )
+
         def forward_fn(
             video_latents: torch.Tensor,
             extra_stream_latents: dict[str, torch.Tensor],
@@ -1267,7 +1379,7 @@ def forward_fn(
         ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
             del timestep, encoder_hidden_states, extra_tensors
             unique_timesteps, timestep_indices = row_timestep_plan[index]
-            video_velocity, audio_velocity = self.transformer(
+            transformer_kwargs = dict(
                 hidden_states=torch.cat((condition_prefix, video_latents), dim=1),
                 audio_hidden_states=torch.cat(
                     (audio_condition_prefix, extra_stream_latents["audio"]), dim=1
@@ -1284,6 +1396,11 @@ def forward_fn(
                 return_dict=False,
                 static_context=static_context,
             )
+            if _vsa_active:
+                with set_h3_vsa_context(vsa_metadata):
+                    video_velocity, audio_velocity = self.transformer(**transformer_kwargs)
+            else:
+                video_velocity, audio_velocity = self.transformer(**transformer_kwargs)
             # A silent all-zero / non-finite velocity has been observed once in
             # practice: the run completed, reported success, and wrote a black
             # video with silent audio.  The cause was never reproduced, so fail