From d5f018fe19a526c053744e3e58558331c8f9003c Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:36:19 +0000 Subject: [PATCH 1/9] Add FastH3 visual generation support Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/supported-models.md | 2 + docs/source/models/visual-generation.md | 6 ++ examples/visual_gen/README.md | 1 + .../visual_gen/configs/fasth3-bf16-1gpu.yaml | 17 +++++ examples/visual_gen/models/fasth3.py | 69 +++++++++++++++++++ tensorrt_llm/_torch/visual_gen/config.py | 7 ++ .../models/minimax_h3/pipeline_minimax_h3.py | 13 +++- .../visual_gen/test_minimax_h3_pipeline.py | 62 ++++++++++++++--- 8 files changed, 164 insertions(+), 13 deletions(-) create mode 100644 examples/visual_gen/configs/fasth3-bf16-1gpu.yaml create mode 100644 examples/visual_gen/models/fasth3.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 58aeea645067..4bc6ecf82dd3 100644 --- a/docs/source/models/supported-models.md +++ b/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-4-step-Preview-v1-VSA-DataFree` | Text-to-Video (with Audio, 4-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 | diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 430995eda3f7..b67fe643e882 100644 --- a/docs/source/models/visual-generation.md +++ b/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-4-step-Preview-v1-VSA-DataFree` | Text-to-Video (with Audio, 4-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 | @@ -178,6 +180,10 @@ Ulysses requires VANILLA or FA4 attention. CFG, ring, and Attention2D parallelis are not supported. See the [2-GPU TP config](../../../examples/visual_gen/configs/minimax-h3-bf16-tp2.yaml). +- FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a + distilled 4-step checkpoint. TensorRT-LLM currently runs it through the dense + MiniMax-H3 attention path; model-specific VSA gate/metadata wiring is not yet + implemented for the H3 transformer. - Text-to-video (T2VA), first/last-frame-to-video (FL2VA), and reference-to-video with audio (Ref2VA) are supported on one GPU. Enable Ref2VA with `pipeline_config: {workflow: ref2va}`; it requires `transformer_ref/` weights. diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 1e3532a44033..d7eb1609a91d 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -29,6 +29,7 @@ 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 --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml +python models/fasth3.py --visual_gen_args configs/fasth3-bf16-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" diff --git a/examples/visual_gen/configs/fasth3-bf16-1gpu.yaml b/examples/visual_gen/configs/fasth3-bf16-1gpu.yaml new file mode 100644 index 000000000000..1b5985a56366 --- /dev/null +++ b/examples/visual_gen/configs/fasth3-bf16-1gpu.yaml @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# 1-GPU BF16 FastH3 VisualGen configuration. FastH3 is a distilled +# MiniMax-H3 checkpoint; this config uses the dense MiniMax-H3 attention path. +cuda_graph_config: + enable: false + +attention_config: + backend: VANILLA + +compilation_config: + skip_warmup: false + resolutions: + - [768, 1344] + num_frames: + - 124 diff --git a/examples/visual_gen/models/fasth3.py b/examples/visual_gen/models/fasth3.py new file mode 100644 index 000000000000..8e12ea35e15a --- /dev/null +++ b/examples/visual_gen/models/fasth3.py @@ -0,0 +1,69 @@ +#!/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-4-step-Preview-v1-VSA-DataFree" + + +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-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.num_inference_steps = 5 + 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() diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 8c0b8aab1936..bf9ce7282a7b 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -635,6 +635,13 @@ def from_pretrained( resolved_pipeline_config = dict(args.pipeline_config) if args else {} extra_attrs.update(resolved_pipeline_config) + fastvideo_inference_path = checkpoint_path / "fastvideo_inference.json" + if fastvideo_inference_path.exists(): + with open(fastvideo_inference_path) as f: + fastvideo_inference = json.load(f) + if str(fastvideo_inference.get("schema_version", "")).startswith("fasth3-"): + extra_attrs["fastvideo_inference"] = fastvideo_inference + model_index_path = checkpoint_path / "model_index.json" if not model_index_path.exists(): model_index_path = checkpoint_path / "modular_model_index.json" diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index 559a7d4d6233..caa7057eecab 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py @@ -164,10 +164,16 @@ def tp_group_name(self) -> str: @register_pipeline( "MiniMaxH3ModularPipeline", - hf_ids=["MiniMaxAI/MiniMax-H3"], + hf_ids=[ + "MiniMaxAI/MiniMax-H3", + "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", + ], defaults={"workflow": "fl2va"}, download_patterns=[ "modular_model_index.json", + "fastvideo_inference.json", + "checkpoint_content.json", + "checkpoint_metadata.json", "LICENSE", "README.md", "transformer/*", @@ -371,6 +377,9 @@ def setup_parallel_vae(self) -> None: @property def default_generation_params(self) -> dict: + fastvideo_inference = getattr(self.pipeline_config, "extra_attrs", {}).get( + "fastvideo_inference", {} + ) return { # MiniMax-H3 derives the default canvas from the first present # keyframe, or from 16:9 for T2VA. @@ -378,7 +387,7 @@ def default_generation_params(self) -> dict: "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 diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index 8f43b29c4c11..fa759dc65c94 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -544,6 +544,9 @@ def test_public_visual_gen_lists_minimax_h3() -> None: from tensorrt_llm import VisualGen assert "MiniMaxAI/MiniMax-H3" in VisualGen.supported_models() + assert ( + "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree" in VisualGen.supported_models() + ) def test_modular_manifest_preserves_pipeline_metadata(tmp_path: Path) -> None: @@ -1075,11 +1078,20 @@ def test_default_generation_steps_match_reference_app() -> None: assert pipeline.default_generation_params["num_inference_steps"] == 28 +def test_fasth3_generation_steps_follow_inference_contract() -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.extra_attrs = {"fastvideo_inference": {"num_inference_steps": 5}} + + assert pipeline.default_generation_params["num_inference_steps"] == 5 + + def test_hf_download_is_scoped_to_the_supported_h3_components( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: captured: dict[str, object] = {} + model = "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree" + revision = "5ea076f" def _capture_partial_download( model: str, @@ -1095,7 +1107,7 @@ def _capture_partial_download( def _reject_full_download(*args: object, **kwargs: object) -> None: del args, kwargs - raise AssertionError("MiniMax-H3 must not download the full 354 GB repository") + raise AssertionError("MiniMax-H3 checkpoints must not download the full repository") monkeypatch.setattr( "tensorrt_llm._torch.visual_gen.pipeline_loader.download_hf_partial", @@ -1106,21 +1118,18 @@ def _reject_full_download(*args: object, **kwargs: object) -> None: "tensorrt_llm._torch.visual_gen.pipeline_loader.download_hf_model", _reject_full_download, ) - loader = PipelineLoader( - VisualGenArgs( - model="MiniMaxAI/MiniMax-H3", - revision="5d9b308a59ab12e67147f191e184baf704185bd1", - ), - device="cpu", - ) + loader = PipelineLoader(VisualGenArgs(model=model, revision=revision), device="cpu") - resolved = loader._resolve_checkpoint_dir("MiniMaxAI/MiniMax-H3") + resolved = loader._resolve_checkpoint_dir(model) assert resolved == str(tmp_path) assert captured == { - "model": "MiniMaxAI/MiniMax-H3", + "model": model, "allow_patterns": [ "modular_model_index.json", + "fastvideo_inference.json", + "checkpoint_content.json", + "checkpoint_metadata.json", "LICENSE", "README.md", "transformer/*", @@ -1133,10 +1142,41 @@ def _reject_full_download(*args: object, **kwargs: object) -> None: "scheduler/*", "audio_scheduler/*", ], - "revision": "5d9b308a59ab12e67147f191e184baf704185bd1", + "revision": revision, } +def test_fasth3_manifest_metadata_controls_pipeline_defaults(tmp_path: Path) -> None: + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer" / "config.json").write_text( + json.dumps({"hidden_size": 32}), + encoding="utf-8", + ) + (tmp_path / "modular_model_index.json").write_text( + json.dumps( + { + "_class_name": "MiniMaxH3ModularPipeline", + "transformer": ["diffusers", "MiniMaxH3Transformer3DModel"], + } + ), + encoding="utf-8", + ) + (tmp_path / "fastvideo_inference.json").write_text( + json.dumps( + { + "schema_version": "fasth3-inference-contract-v1", + "model_id": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", + "num_inference_steps": 5, + } + ), + encoding="utf-8", + ) + + config = DiffusionPipelineConfig.from_pretrained(str(tmp_path)) + + assert config.extra_attrs["fastvideo_inference"]["num_inference_steps"] == 5 + + def test_infer_supports_a_last_frame_without_a_first_frame( monkeypatch: pytest.MonkeyPatch, ) -> None: From 323bdcee5a866cd2034c401502d3600a29365e1b Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:54:56 +0000 Subject: [PATCH 2/9] Add FastH3 TRTLLM VSA attention support Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../configs/fasth3-bf16-trtllm-vsa-1gpu.yaml | 17 +++ examples/visual_gen/models/fasth3.py | 1 - .../models/minimax_h3/pipeline_minimax_h3.py | 33 ++++- .../minimax_h3/transformer_minimax_h3.py | 46 +++++- .../visual_gen/models/minimax_h3/vsa.py | 138 ++++++++++++++++++ .../_torch/visual_gen/test_minimax_h3_vsa.py | 79 ++++++++++ 6 files changed, 310 insertions(+), 4 deletions(-) create mode 100644 examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml create mode 100644 tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py create mode 100644 tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py diff --git a/examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml b/examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml new file mode 100644 index 000000000000..aafc87cf4036 --- /dev/null +++ b/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.9 +compilation_config: + skip_warmup: true + resolutions: + - [768, 1344] + num_frames: + - 124 diff --git a/examples/visual_gen/models/fasth3.py b/examples/visual_gen/models/fasth3.py index 8e12ea35e15a..0ef6354e2529 100644 --- a/examples/visual_gen/models/fasth3.py +++ b/examples/visual_gen/models/fasth3.py @@ -20,7 +20,6 @@ from tensorrt_llm import VisualGen, VisualGenArgs - FASTH3_MODEL_ID = "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree" diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index caa7057eecab..07a30f61a6a4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py @@ -75,6 +75,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 @@ -1266,6 +1267,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], @@ -1276,7 +1302,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 @@ -1293,6 +1319,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 diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py index e09d334f917a..527f5cdb1da2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py @@ -47,6 +47,7 @@ from tensorrt_llm.quantization.mode import QuantAlgo from .fused_adaln import h3_adaln_supported, h3_gate_res_norm_mod, h3_norm_mod +from .vsa import get_h3_vsa_context, h3_vsa_attention MINIMAX_H3_MODALITY_NUM = 3 @@ -96,6 +97,19 @@ def apply_minimax_h3_rotary_emb( return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() +def _dense_attention_model_config(model_config: DiffusionModelConfig) -> DiffusionModelConfig: + """Return a shallow config copy with VSA disabled for non-video attention.""" + + sparse_config = getattr(model_config.attention, "sparse_attention_config", None) + if sparse_config is None or getattr(sparse_config, "algorithm", None) != "vsa": + return model_config + + dense_config = model_config.model_copy(deep=False) + dense_config.attention = model_config.attention.model_copy(deep=True) + dense_config.attention.sparse_attention_config = None + return dense_config + + class MiniMaxH3RotaryPosEmbed(nn.Module): """Three-axis RoPE over packed ``(time, height, width)`` coordinates.""" @@ -152,13 +166,28 @@ def __init__( # (fused_rope.py) through the override of forward() below. fuse_qk_norm_rope=True, interleave=False, - config=model_config, + config=_dense_attention_model_config(model_config), layer_idx=layer_idx, module_name=module_name, async_ulysses=use_async_ulysses, separate_qkv_is_self_attention=True, ) self.use_async_ulysses = use_async_ulysses + sparse = model_config.attention.sparse_attention_config + self.to_gate_compress = None + if sparse is not None and sparse.algorithm == "vsa": + if model_config.attention.backend != "TRTLLM": + raise ValueError("MiniMax-H3 VSA currently requires the TRTLLM backend.") + if model_config.mapping.world_size != 1: + raise ValueError("MiniMax-H3 VSA currently supports a single GPU.") + self.to_gate_compress = Linear( + hidden_size, + num_attention_heads * attention_head_dim, + bias=False, + dtype=model_config.torch_dtype, + mapping=model_config.mapping, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) def _use_fused_qk_norm_rope( self, @@ -240,6 +269,19 @@ def forward( key = apply_minimax_h3_rotary_emb(key, *rotary_emb) query = query.flatten(2) key = key.flatten(2) + + if self.to_gate_compress is not None: + query = query.view( + batch_size, sequence_length, self.local_num_attention_heads, self.head_dim + ) + key = key.view_as(query) + metadata = get_h3_vsa_context() + gate = self.to_gate_compress(hidden_states).view_as(query) + output = h3_vsa_attention( + query, key, value.view_as(query), gate, metadata, self._attn_impl + ) + return self.to_out[0](output.flatten(2)) + hidden_states = self._attn_impl( query, key, @@ -423,7 +465,7 @@ def __init__( num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, qk_norm_eps=qk_norm_eps, - model_config=model_config, + model_config=_dense_attention_model_config(model_config), layer_idx=layer_idx, module_name=f"token_refiner.refiner_blocks.{layer_idx}.attn", ) diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py new file mode 100644 index 000000000000..1541d6a29a22 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""H3 mixed-modality VSA routing for TRTLLM block-sparse attention. + +The FastH3 inference contract uses 64-token (4, 4, 4) video tiles, dense +text/audio segments, and a learned additive compression gate. +""" + +import math +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseBackendForwardArgs, +) + + +@dataclass +class H3VSAMetadata: + source: torch.Tensor + inverse: torch.Tensor + sizes: torch.Tensor + valid_bits: torch.Tensor + prefix_tiles: int + sparsity: float + + +_context = ContextVar("h3_vsa", default=None) + + +@contextmanager +def set_h3_vsa_context(metadata): + token = _context.set(metadata) + try: + yield + finally: + _context.reset(token) + + +def get_h3_vsa_context(): + metadata = _context.get() + if metadata is None: + raise RuntimeError("MiniMax-H3 VSA requires packed-sequence metadata.") + return metadata + + +def build_h3_vsa_metadata(segments, sparsity, device): + if not 0 <= sparsity < 1: + raise ValueError("VSA sparsity must be in [0, 1).") + tiles = [] + cursor = 0 + prefix_tiles = 0 + for segment in segments: + if isinstance(segment, int): + for start in range(0, segment, 64): + tiles.append(torch.arange(cursor + start, cursor + min(start + 64, segment))) + prefix_tiles += 1 + cursor += segment + else: + t, h, w = segment + grid = torch.arange(cursor, cursor + t * h * w).reshape(t, h, w) + for ti in range(0, t, 4): + for hi in range(0, h, 4): + for wi in range(0, w, 4): + tiles.append(grid[ti : ti + 4, hi : hi + 4, wi : wi + 4].flatten()) + cursor += t * h * w + source = torch.full((len(tiles), 64), -1, dtype=torch.long) + sizes = torch.tensor([tile.numel() for tile in tiles]) + for index, tile in enumerate(tiles): + source[index, : tile.numel()] = tile + source = source.flatten() + valid = source >= 0 + inverse = torch.empty(cursor, dtype=torch.long) + inverse[source[valid]] = torch.arange(source.numel())[valid] + words = (valid.reshape(-1, 32).long() * (1 << torch.arange(32))).sum(-1) + return H3VSAMetadata( + source.to(device), + inverse.to(device), + sizes.to(device), + words.to(device=device, dtype=torch.uint32), + prefix_tiles, + sparsity, + ) + + +def h3_vsa_attention(q, k, v, gate, metadata, attention): + """Execute trained H3 selection, sparse fine attention and coarse merge.""" + if q.ndim != 4 or q.shape != k.shape or q.shape != v.shape or gate.shape != q.shape: + raise ValueError("H3 VSA requires matching BSHD Q/K/V and compression gate.") + if q.shape[1] != metadata.inverse.numel(): + raise ValueError("H3 VSA metadata does not match the packed sequence.") + batch, _, heads, dim = q.shape + ntiles = metadata.sizes.numel() + valid = (metadata.source >= 0).view(1, -1, 1, 1) + + def tile_pool(x): + tiled = x.index_select(1, metadata.source.clamp_min(0)).masked_fill(~valid, 0) + pooled = tiled.view(batch, ntiles, 64, heads, dim).sum(2, dtype=torch.float32) + pooled = pooled / metadata.sizes.view(1, -1, 1, 1) + return tiled, pooled.transpose(1, 2) + + qt, qp = tile_pool(q) + kt, kp = tile_pool(k) + vt, vp = tile_pool(v) + scores = torch.matmul(qp, kp.transpose(-2, -1)) * dim**-0.5 + coarse = torch.matmul(scores.softmax(-1), vp).transpose(1, 2).to(q.dtype) + prefix = metadata.prefix_tiles + topk = max(1, math.ceil((1 - metadata.sparsity) * (ntiles - prefix))) + selected = scores[..., prefix:].topk(topk, dim=-1).indices + prefix + mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, selected, True) + mask[..., :prefix] = True + mask[:, :, :prefix, :] = True + counts = mask.sum(-1, dtype=torch.int32) + ends = counts.flatten().cumsum(0, dtype=torch.int32).view(batch, heads, -1) + starts = ends - counts + indptr = torch.cat((starts, ends[..., -1:]), dim=-1).contiguous() + indices = mask.nonzero()[:, -1].to(torch.int32).contiguous() + routes = BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=ntiles, + block_indptr=indptr, + block_indices=indices, + kv_valid_bits=metadata.valid_bits[None].expand(batch, -1).contiguous(), + ) + fine = attention( + qt.flatten(2), + kt.flatten(2), + vt.flatten(2), + sparse_backend_args=SparseBackendForwardArgs(block_sparse_inputs=routes), + ).view(batch, ntiles * 64, heads, dim) + fine = fine.index_select(1, metadata.inverse) + coarse = coarse.index_select(1, metadata.inverse // 64) + return torch.addcmul(fine, coarse, gate) diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py new file mode 100644 index 000000000000..65b5b19b9db7 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.models.minimax_h3.vsa import ( + build_h3_vsa_metadata, + h3_vsa_attention, +) + + +def test_h3_vsa_geometry(): + metadata = build_h3_vsa_metadata((7, 5, (5, 5, 7)), 0.9, "cpu") + valid = metadata.source >= 0 + assert metadata.sizes.sum() == 187 + assert metadata.prefix_tiles == 2 + torch.testing.assert_close(metadata.source[metadata.inverse], torch.arange(187)) + assert torch.unique(metadata.source[valid]).numel() == 187 + # Partial text/audio tiles must not be pooled together. + assert metadata.sizes[:2].tolist() == [7, 5] + + +@pytest.mark.parametrize("sparsity", [0.0, 0.5, 0.9]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_h3_vsa_trtllm_masked_reference(sparsity): + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("PrimTS block-sparse test requires Blackwell") + torch.manual_seed(42) + metadata = build_h3_vsa_metadata((7, 5, (5, 5, 7)), sparsity, "cuda") + q, k, v = [torch.randn(1, 187, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + gate = torch.randn_like(q) * 0.1 + backend = TrtllmAttention( + num_heads=2, + num_kv_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_metadata_state={}, + ) + reference = [] + + def run(qt, kt, vt, sparse_backend_args): + seq = qt.shape[1] + tensors = [x.view(1, seq, 2, 128) for x in (qt, kt, vt)] + actual = backend.forward( + *tensors, batch_size=1, seq_len=seq, sparse_backend_args=sparse_backend_args + ) + routes = sparse_backend_args.block_sparse_inputs + blocks = torch.zeros(1, 2, seq // 64, seq // 64, dtype=torch.bool, device="cuda") + for head in range(2): + for row in range(seq // 64): + start, end = routes.block_indptr[0, head, row : row + 2].tolist() + blocks[0, head, row, routes.block_indices[start:end].long()] = True + prefix = metadata.prefix_tiles + assert blocks[:, :, :prefix].all() + assert blocks[..., :prefix].all() + mask = blocks.repeat_interleave(64, -2).repeat_interleave(64, -1) + mask &= (metadata.source >= 0)[None, None, None, :] + expected = torch.nn.functional.scaled_dot_product_attention( + *[x.transpose(1, 2) for x in tensors], + attn_mask=mask, + ).transpose(1, 2) + actual = actual.view_as(expected) + valid = metadata.source >= 0 + torch.testing.assert_close(actual[:, valid], expected[:, valid], atol=0.02, rtol=0.02) + pooled = [ + x.view(1, -1, 64, 2, 128).sum(2, dtype=torch.float32) + / metadata.sizes[None, :, None, None] + for x in tensors + ] + qp, kp, vp = [x.transpose(1, 2) for x in pooled] + coarse = (qp @ kp.transpose(-2, -1) / 128**0.5).softmax(-1) @ vp + coarse = coarse.transpose(1, 2).to(q.dtype)[:, metadata.inverse // 64] + reference.append(expected[:, metadata.inverse] + coarse * gate) + return actual.flatten(2) + + output = h3_vsa_attention(q, k, v, gate, metadata, run) + torch.testing.assert_close(output, reference[0], atol=0.02, rtol=0.02) From 6cf999d9d321982ee781fb3138629a5f412974bd Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:45:13 +0000 Subject: [PATCH 3/9] Support CUTEDSL VSA for FastH3 Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/visual-generation.md | 9 ++- examples/visual_gen/README.md | 3 + .../configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml | 17 ++++++ .../minimax_h3/transformer_minimax_h3.py | 13 +++- .../visual_gen/models/minimax_h3/vsa.py | 59 ++++++++++++++++++- .../_torch/visual_gen/test_minimax_h3_vsa.py | 30 ++++++---- 6 files changed, 112 insertions(+), 19 deletions(-) create mode 100644 examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index b67fe643e882..0c9574be2456 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -181,9 +181,12 @@ are not supported. See the [2-GPU TP config](../../../examples/visual_gen/configs/minimax-h3-bf16-tp2.yaml). - FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a - distilled 4-step checkpoint. TensorRT-LLM currently runs it through the dense - MiniMax-H3 attention path; model-specific VSA gate/metadata wiring is not yet - implemented for the H3 transformer. + distilled 4-step checkpoint. Dense attention and model-specific VSA are supported. + For single-GPU text-to-video with audio, enable VSA using + `fasth3-bf16-trtllm-vsa-1gpu.yaml` or `fasth3-bf16-cutedsl-vsa-1gpu.yaml`. + Both 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. - Text-to-video (T2VA), first/last-frame-to-video (FL2VA), and reference-to-video with audio (Ref2VA) are supported on one GPU. Enable Ref2VA with `pipeline_config: {workflow: ref2va}`; it requires `transformer_ref/` weights. diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index d7eb1609a91d..89cca6afe554 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -30,6 +30,9 @@ python models/ltx2.py --visual_gen_args configs/ltx2-1gpu.yaml # 2-GPU Ulysses config: configs/minimax-h3-bf16-2gpu.yaml. python models/minimax_h3.py --model --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml python models/fasth3.py --visual_gen_args configs/fasth3-bf16-1gpu.yaml +# FastH3 VSA on Blackwell, with either attention backend: +python models/fasth3.py --visual_gen_args configs/fasth3-bf16-trtllm-vsa-1gpu.yaml +python models/fasth3.py --visual_gen_args configs/fasth3-bf16-cutedsl-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" diff --git a/examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml b/examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml new file mode 100644 index 000000000000..d2e42d130079 --- /dev/null +++ b/examples/visual_gen/configs/fasth3-bf16-cutedsl-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: CUTEDSL + sparse_attention_config: + algorithm: vsa + vsa_sparsity: 0.9 +compilation_config: + skip_warmup: true + resolutions: + - [768, 1344] + num_frames: + - 124 diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py index 527f5cdb1da2..74723b690b8f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py @@ -175,9 +175,10 @@ def __init__( self.use_async_ulysses = use_async_ulysses sparse = model_config.attention.sparse_attention_config self.to_gate_compress = None + self._vsa_backend = model_config.attention.backend if sparse is not None and sparse.algorithm == "vsa": - if model_config.attention.backend != "TRTLLM": - raise ValueError("MiniMax-H3 VSA currently requires the TRTLLM backend.") + if self._vsa_backend not in ("TRTLLM", "CUTEDSL"): + raise ValueError("MiniMax-H3 VSA requires the TRTLLM or CUTEDSL backend.") if model_config.mapping.world_size != 1: raise ValueError("MiniMax-H3 VSA currently supports a single GPU.") self.to_gate_compress = Linear( @@ -278,7 +279,13 @@ def forward( metadata = get_h3_vsa_context() gate = self.to_gate_compress(hidden_states).view_as(query) output = h3_vsa_attention( - query, key, value.view_as(query), gate, metadata, self._attn_impl + query, + key, + value.view_as(query), + gate, + metadata, + self._attn_impl, + backend=self._vsa_backend, ) return self.to_out[0](output.flatten(2)) diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py index 1541d6a29a22..b05c9bae640e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""H3 mixed-modality VSA routing for TRTLLM block-sparse attention. +"""H3 mixed-modality VSA routing for TRTLLM and CuTe DSL sparse attention. The FastH3 inference contract uses 64-token (4, 4, 4) video tiles, dense text/audio segments, and a learned additive compression gate. @@ -87,12 +87,62 @@ def build_h3_vsa_metadata(segments, sparsity, device): ) -def h3_vsa_attention(q, k, v, gate, metadata, attention): +def _cutedsl_fine(q, k, v, selected, metadata): + from ...cute_dsl_kernels.blackwell.video_sparse_attention.interface import ( + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not is_cute_supported(q) or q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError( + "H3 CUTEDSL VSA requires matching FP16/BF16 tensors with head_dim=128 on Blackwell." + ) + batch, _, heads, _ = q.shape + ntiles = metadata.sizes.numel() + prefix = metadata.prefix_tiles + # Prefix queries are dense; video queries retain the prefix plus selected video tiles. + indices = torch.zeros((batch, heads, ntiles, ntiles), device=q.device, dtype=torch.int32) + indices[:, :, :prefix] = torch.arange(ntiles, device=q.device, dtype=torch.int32) + indices[:, :, prefix:, :prefix] = torch.arange(prefix, device=q.device, dtype=torch.int32) + indices[:, :, prefix:, prefix : prefix + selected.shape[-1]] = selected[:, :, prefix:] + counts = torch.full( + (batch, heads, ntiles), prefix + selected.shape[-1], device=q.device, dtype=torch.int32 + ) + counts[:, :, :prefix] = ntiles + # The kernel shares a route count across each pair of query tiles. Duplicate + # boundary queries so dense and sparse groups each have an even tile count. + query_tiles = list(range(prefix)) + if prefix % 2: + query_tiles.append(prefix - 1) + query_tiles.extend(range(prefix, ntiles)) + if (ntiles - prefix) % 2: + query_tiles.append(ntiles - 1) + query_tiles = torch.tensor(query_tiles, device=q.device) + restore = torch.arange(ntiles, device=q.device) + restore[prefix:] += prefix % 2 + queries = q.view(batch, ntiles, 64, heads, -1).index_select(1, query_tiles).flatten(1, 2) + fine, _ = block_sparse_attn_from_indices_cute( + *[x.transpose(1, 2).contiguous() for x in (queries, k, v)], + q2k_idx=indices.index_select(2, query_tiles), + q2k_num=counts.index_select(2, query_tiles), + variable_block_sizes=metadata.sizes.to(torch.int32), + ) + return ( + fine.transpose(1, 2) + .reshape(batch, -1, 64, heads, q.shape[-1]) + .index_select(1, restore) + .flatten(1, 2) + ) + + +def h3_vsa_attention(q, k, v, gate, metadata, attention, *, backend="TRTLLM"): """Execute trained H3 selection, sparse fine attention and coarse merge.""" if q.ndim != 4 or q.shape != k.shape or q.shape != v.shape or gate.shape != q.shape: raise ValueError("H3 VSA requires matching BSHD Q/K/V and compression gate.") if q.shape[1] != metadata.inverse.numel(): raise ValueError("H3 VSA metadata does not match the packed sequence.") + if backend not in ("TRTLLM", "CUTEDSL"): + raise ValueError(f"Unsupported H3 VSA backend: {backend}") batch, _, heads, dim = q.shape ntiles = metadata.sizes.numel() valid = (metadata.source >= 0).view(1, -1, 1, 1) @@ -111,6 +161,11 @@ def tile_pool(x): prefix = metadata.prefix_tiles topk = max(1, math.ceil((1 - metadata.sparsity) * (ntiles - prefix))) selected = scores[..., prefix:].topk(topk, dim=-1).indices + prefix + if backend == "CUTEDSL": + fine = _cutedsl_fine(qt, kt, vt, selected, metadata) + fine = fine.index_select(1, metadata.inverse) + coarse = coarse.index_select(1, metadata.inverse // 64) + return torch.addcmul(fine, coarse, gate) mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, selected, True) mask[..., :prefix] = True mask[:, :, :prefix, :] = True diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py index 65b5b19b9db7..b918ac537ccb 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py @@ -23,13 +23,18 @@ def test_h3_vsa_geometry(): @pytest.mark.parametrize("sparsity", [0.0, 0.5, 0.9]) +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("segments", [(7, 5, (5, 5, 7)), (7, (4, 4, 12))]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_h3_vsa_trtllm_masked_reference(sparsity): +def test_h3_vsa_backends_masked_reference(sparsity, batch, segments): if torch.cuda.get_device_capability()[0] != 10: pytest.skip("PrimTS block-sparse test requires Blackwell") torch.manual_seed(42) - metadata = build_h3_vsa_metadata((7, 5, (5, 5, 7)), sparsity, "cuda") - q, k, v = [torch.randn(1, 187, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + metadata = build_h3_vsa_metadata(segments, sparsity, "cuda") + seq_len = metadata.inverse.numel() + q, k, v = [ + torch.randn(batch, seq_len, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3) + ] gate = torch.randn_like(q) * 0.1 backend = TrtllmAttention( num_heads=2, @@ -42,16 +47,17 @@ def test_h3_vsa_trtllm_masked_reference(sparsity): def run(qt, kt, vt, sparse_backend_args): seq = qt.shape[1] - tensors = [x.view(1, seq, 2, 128) for x in (qt, kt, vt)] + tensors = [x.view(batch, seq, 2, 128) for x in (qt, kt, vt)] actual = backend.forward( - *tensors, batch_size=1, seq_len=seq, sparse_backend_args=sparse_backend_args + *tensors, batch_size=batch, seq_len=seq, sparse_backend_args=sparse_backend_args ) routes = sparse_backend_args.block_sparse_inputs - blocks = torch.zeros(1, 2, seq // 64, seq // 64, dtype=torch.bool, device="cuda") - for head in range(2): - for row in range(seq // 64): - start, end = routes.block_indptr[0, head, row : row + 2].tolist() - blocks[0, head, row, routes.block_indices[start:end].long()] = True + blocks = torch.zeros(batch, 2, seq // 64, seq // 64, dtype=torch.bool, device="cuda") + for b in range(batch): + for head in range(2): + for row in range(seq // 64): + start, end = routes.block_indptr[b, head, row : row + 2].tolist() + blocks[b, head, row, routes.block_indices[start:end].long()] = True prefix = metadata.prefix_tiles assert blocks[:, :, :prefix].all() assert blocks[..., :prefix].all() @@ -65,7 +71,7 @@ def run(qt, kt, vt, sparse_backend_args): valid = metadata.source >= 0 torch.testing.assert_close(actual[:, valid], expected[:, valid], atol=0.02, rtol=0.02) pooled = [ - x.view(1, -1, 64, 2, 128).sum(2, dtype=torch.float32) + x.view(batch, -1, 64, 2, 128).sum(2, dtype=torch.float32) / metadata.sizes[None, :, None, None] for x in tensors ] @@ -77,3 +83,5 @@ def run(qt, kt, vt, sparse_backend_args): output = h3_vsa_attention(q, k, v, gate, metadata, run) torch.testing.assert_close(output, reference[0], atol=0.02, rtol=0.02) + cute_output = h3_vsa_attention(q, k, v, gate, metadata, None, backend="CUTEDSL") + torch.testing.assert_close(cute_output, reference[0], atol=0.02, rtol=0.02) From f6fb081bc8628389ba8edaaaebff4c107d9e7a86 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:13:28 +0000 Subject: [PATCH 4/9] Scope FastH3 metadata loading to the MiniMax-H3 pipeline Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/visual-generation.md | 5 ++-- examples/visual_gen/README.md | 4 +-- .../visual_gen/configs/fasth3-bf16-1gpu.yaml | 17 ----------- .../configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml | 17 ----------- examples/visual_gen/models/fasth3.py | 2 +- tensorrt_llm/_torch/visual_gen/config.py | 7 ----- .../models/minimax_h3/pipeline_minimax_h3.py | 14 +++++++-- .../visual_gen/test_minimax_h3_pipeline.py | 29 ++++++++++++++++--- .../_torch/visual_gen/test_minimax_h3_vsa.py | 2 +- 9 files changed, 42 insertions(+), 55 deletions(-) delete mode 100644 examples/visual_gen/configs/fasth3-bf16-1gpu.yaml delete mode 100644 examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 0c9574be2456..9ab9dff7853d 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -183,8 +183,9 @@ are not supported. See the - FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a distilled 4-step checkpoint. Dense attention and model-specific VSA are supported. For single-GPU text-to-video with audio, enable VSA using - `fasth3-bf16-trtllm-vsa-1gpu.yaml` or `fasth3-bf16-cutedsl-vsa-1gpu.yaml`. - Both preserve the trained compression gate and dense text/audio routes. + `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. - Text-to-video (T2VA), first/last-frame-to-video (FL2VA), and reference-to-video diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 89cca6afe554..85ccbdc3762f 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -29,10 +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 --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml -python models/fasth3.py --visual_gen_args configs/fasth3-bf16-1gpu.yaml -# FastH3 VSA on Blackwell, with either attention backend: +# FastH3 VSA on Blackwell: python models/fasth3.py --visual_gen_args configs/fasth3-bf16-trtllm-vsa-1gpu.yaml -python models/fasth3.py --visual_gen_args configs/fasth3-bf16-cutedsl-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" diff --git a/examples/visual_gen/configs/fasth3-bf16-1gpu.yaml b/examples/visual_gen/configs/fasth3-bf16-1gpu.yaml deleted file mode 100644 index 1b5985a56366..000000000000 --- a/examples/visual_gen/configs/fasth3-bf16-1gpu.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# 1-GPU BF16 FastH3 VisualGen configuration. FastH3 is a distilled -# MiniMax-H3 checkpoint; this config uses the dense MiniMax-H3 attention path. -cuda_graph_config: - enable: false - -attention_config: - backend: VANILLA - -compilation_config: - skip_warmup: false - resolutions: - - [768, 1344] - num_frames: - - 124 diff --git a/examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml b/examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml deleted file mode 100644 index d2e42d130079..000000000000 --- a/examples/visual_gen/configs/fasth3-bf16-cutedsl-vsa-1gpu.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# 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: CUTEDSL - sparse_attention_config: - algorithm: vsa - vsa_sparsity: 0.9 -compilation_config: - skip_warmup: true - resolutions: - - [768, 1344] - num_frames: - - 124 diff --git a/examples/visual_gen/models/fasth3.py b/examples/visual_gen/models/fasth3.py index 0ef6354e2529..15da5b129cc7 100644 --- a/examples/visual_gen/models/fasth3.py +++ b/examples/visual_gen/models/fasth3.py @@ -32,7 +32,7 @@ def main() -> None: ) parser.add_argument( "--visual_gen_args", - default="configs/fasth3-bf16-1gpu.yaml", + default="configs/fasth3-bf16-trtllm-vsa-1gpu.yaml", help="Path to YAML config (same as trtllm-serve --visual_gen_args).", ) parser.add_argument( diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index bf9ce7282a7b..8c0b8aab1936 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -635,13 +635,6 @@ def from_pretrained( resolved_pipeline_config = dict(args.pipeline_config) if args else {} extra_attrs.update(resolved_pipeline_config) - fastvideo_inference_path = checkpoint_path / "fastvideo_inference.json" - if fastvideo_inference_path.exists(): - with open(fastvideo_inference_path) as f: - fastvideo_inference = json.load(f) - if str(fastvideo_inference.get("schema_version", "")).startswith("fasth3-"): - extra_attrs["fastvideo_inference"] = fastvideo_inference - model_index_path = checkpoint_path / "model_index.json" if not model_index_path.exists(): model_index_path = checkpoint_path / "modular_model_index.json" diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index 07a30f61a6a4..0b11a8619c78 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/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 @@ -378,9 +380,7 @@ def setup_parallel_vae(self) -> None: @property def default_generation_params(self) -> dict: - fastvideo_inference = getattr(self.pipeline_config, "extra_attrs", {}).get( - "fastvideo_inference", {} - ) + fastvideo_inference = getattr(self, "_fastvideo_inference", {}) return { # MiniMax-H3 derives the default canvas from the first present # keyframe, or from 16:9 for T2VA. @@ -515,6 +515,14 @@ def load_standard_components( skip_components = skip_components or [] offloaded = self.offloader.requested_components() + self._fastvideo_inference = {} + fastvideo_inference_path = Path(checkpoint_dir) / "fastvideo_inference.json" + if fastvideo_inference_path.exists(): + with fastvideo_inference_path.open() as f: + fastvideo_inference = json.load(f) + if str(fastvideo_inference.get("schema_version", "")).startswith("fasth3-"): + self._fastvideo_inference = fastvideo_inference + if not _component_skipped(skip_components, PipelineComponent.TOKENIZER): self.tokenizer = Qwen2TokenizerFast.from_pretrained( checkpoint_dir, diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index fa759dc65c94..f71ef53de8c1 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -1078,11 +1078,26 @@ def test_default_generation_steps_match_reference_app() -> None: assert pipeline.default_generation_params["num_inference_steps"] == 28 -def test_fasth3_generation_steps_follow_inference_contract() -> None: +@pytest.mark.parametrize( + "manifest, expected_steps", + [ + (None, 28), + ({"schema_version": "fasth3-inference-contract-v1", "num_inference_steps": 5}, 5), + ({"schema_version": "fasth3-inference-contract-v1"}, 28), + ({"schema_version": "other-model-v1", "num_inference_steps": 5}, 28), + ], +) +def test_fasth3_generation_steps_follow_inference_contract( + tmp_path: Path, manifest: dict | None, expected_steps: int +) -> None: + if manifest is not None: + (tmp_path / "fastvideo_inference.json").write_text(json.dumps(manifest)) pipeline = _SyntheticMiniMaxH3Pipeline() - pipeline.pipeline_config.extra_attrs = {"fastvideo_inference": {"num_inference_steps": 5}} + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=list(PipelineComponent) + ) - assert pipeline.default_generation_params["num_inference_steps"] == 5 + assert pipeline.default_generation_params["num_inference_steps"] == expected_steps def test_hf_download_is_scoped_to_the_supported_h3_components( @@ -1174,7 +1189,13 @@ def test_fasth3_manifest_metadata_controls_pipeline_defaults(tmp_path: Path) -> config = DiffusionPipelineConfig.from_pretrained(str(tmp_path)) - assert config.extra_attrs["fastvideo_inference"]["num_inference_steps"] == 5 + assert "fastvideo_inference" not in config.extra_attrs + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config = config + pipeline.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=list(PipelineComponent) + ) + assert pipeline.default_generation_params["num_inference_steps"] == 5 def test_infer_supports_a_last_frame_without_a_first_frame( diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py index b918ac537ccb..47b2f9f08f82 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py @@ -4,7 +4,7 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention.backends.trtllm import TrtllmAttention from tensorrt_llm._torch.visual_gen.models.minimax_h3.vsa import ( build_h3_vsa_metadata, h3_vsa_attention, From 2511a453de7127b0f51e348be217d54cbb6a90c7 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:23:20 +0000 Subject: [PATCH 5/9] Add FastH3 golden-video regression coverage and CI enrollment Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/visual-generation.md | 4 +- .../fasth3_lpips_golden_video.json | 27 +++++++++ ...sual_gen_lpips_golden_media_minimax_h3.zip | 3 + .../visual_gen/test_visual_gen_fasth3.py | 60 +++++++++++++++++++ .../visual_gen/visual_gen_test_utils.py | 2 + .../test_lists/test-db/l0_b200.yml | 2 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../_torch/visual_gen/test_minimax_h3_vsa.py | 1 + 8 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fasth3_lpips_golden_video.json create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_minimax_h3.zip create mode 100644 tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 9ab9dff7853d..b85673c86e5f 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -181,7 +181,9 @@ are not supported. See the [2-GPU TP config](../../../examples/visual_gen/configs/minimax-h3-bf16-tp2.yaml). - FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a - distilled 4-step checkpoint. Dense attention and model-specific VSA are supported. + distilled 4-step checkpoint. Model-specific VSA is the supported inference path + for this checkpoint. Dense attention is a comparison-only path: it does not use + the trained compression gate and is not equivalent to the checkpoint's VSA inference. For single-GPU text-to-video with audio, enable VSA using `fasth3-bf16-trtllm-vsa-1gpu.yaml`. To use CUTEDSL instead, set `attention_config.backend: CUTEDSL` in that config. diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fasth3_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fasth3_lpips_golden_video.json new file mode 100644 index 000000000000..1a4fbcd72ab5 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/fasth3_lpips_golden_video.json @@ -0,0 +1,27 @@ +{ + "video": "fasth3_lpips_golden_video.mp4", + "model": "FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", + "model_id": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", + "model_revision": "5ea076f35b84da4c3c82217112fa733d8eea2ae1", + "source": "TRTLLM VisualGen regression baseline, not an independent FastVideo reference", + "source_commit": "ba6c890ae8551a6a5baf9630ceb6d4d212ec0761", + "generated_at": "2026-10-27", + "torch_version": "2.14.0a0+b2c75dd062.nv26.09", + "gpu": "NVIDIA B200", + "backend": "TRTLLM", + "dtype": "bfloat16", + "vsa_sparsity": 0.9, + "cuda_graph": false, + "torch_compile": false, + "prompt": "A cinematic wide shot of an astronaut walking through a neon-lit rainforest at dusk, with soft ambient music and detailed natural sounds.", + "height": 768, + "width": 1344, + "num_frames": 124, + "num_inference_steps": 5, + "transformer_forwards": 4, + "seed": 42, + "frame_rate": 24.0, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "threshold_note": "Provisional regression bound matching existing VisualGen golden tests; not a calibrated FastVideo-reference acceptance score." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_minimax_h3.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_minimax_h3.zip new file mode 100644 index 000000000000..d235e295f31b --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media_minimax_h3.zip @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2392dd03a11811e7e18cb03fc451a1c0a10d45c87a799ad7aeb16ce41d99e3e8 +size 2184018 diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py new file mode 100644 index 000000000000..6c50bd023c3b --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FastH3 example regression against a recorded TRTLLM VSA video. + +This golden detects output drift; it does not establish FastVideo reference parity. +Stage the checkpoint revision recorded in the golden JSON under LLM_MODELS_ROOT. +""" + +import json +from pathlib import Path + +import pytest +import torch +from defs.common import venv_check_call +from defs.examples.visual_gen.visual_gen_test_utils import ( + VISUAL_GEN_LPIPS_GOLDEN_DIR, + _assert_lpips_below_threshold, + _golden_media_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, +) +from test_common.llm_data import get_checkpoint + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_fasth3_example_lpips_against_golden( + request, tmp_path, _visual_gen_deps, llm_root, llm_venv +): + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("FastH3 TRTLLM VSA requires Blackwell") + metadata_path = Path(VISUAL_GEN_LPIPS_GOLDEN_DIR) / "fasth3_lpips_golden_video.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + golden_path = _golden_media_path(tmp_path, metadata["video"], "FastH3 TRTLLM regression golden") + model_path = get_checkpoint(metadata["model"]) + output_path = tmp_path / "fasth3_output.mp4" + script_path = Path(llm_root) / "examples/visual_gen/models/fasth3.py" + config_path = Path(llm_root) / "examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml" + # Exercise the example's published schedule, seed and geometry defaults. + venv_check_call( + llm_venv, + [ + str(script_path), + "--model", + model_path, + "--visual_gen_args", + str(config_path), + "--prompt", + metadata["prompt"], + "--output_path", + str(output_path), + ], + ) + assert output_path.is_file(), "FastH3 example did not save a video" + score = _run_lpips_eval( + tmp_path, "fasth3", "video", metadata["prompt"], golden_path, output_path + ) + threshold = metadata["lpips_threshold"] + _preserve_lpips_candidate_on_failure(request, score, threshold, output_path, metadata["video"]) + _assert_lpips_below_threshold(score, threshold) diff --git a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py index 844c050f08a6..5ab16bcce08b 100644 --- a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py +++ b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py @@ -483,6 +483,8 @@ def _require_exists(path, label, is_dir=False): def _golden_media_family(media_name): + if media_name.startswith("fasth3_"): + return "minimax_h3" if media_name.startswith("cosmos3_"): return "cosmos3" if media_name.startswith(("fastwan_", "wan21_", "wan22_")): diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 30c3af0b9e9a..036ecf6f8b60 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -494,10 +494,12 @@ l0_b200: - unittest/_torch/visual_gen/test_glm_image_transformer.py - unittest/_torch/visual_gen/test_glm_image_pipeline.py - unittest/_torch/visual_gen/test_minimax_h3_transformer.py + - unittest/_torch/visual_gen/test_minimax_h3_vsa.py - unittest/_torch/visual_gen/test_minimax_h3_fused_rope.py - unittest/_torch/visual_gen/test_minimax_h3_fused_adaln.py - unittest/_torch/visual_gen/test_minimax_h3_vae_decode.py - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_public_visual_gen_api_smoke TIMEOUT (60) + - examples/visual_gen/test_visual_gen_fasth3.py::test_fasth3_example_lpips_against_golden TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[t2va] TIMEOUT (60) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[fl2va] TIMEOUT (60) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[ref2va] TIMEOUT (30) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index cee949636708..128b72b91cc7 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -183,6 +183,7 @@ l0_cpu: - unittest/_torch/visual_gen/test_ltx23_pipeline_config.py - unittest/_torch/visual_gen/test_ltx23_retake_initialization.py - unittest/_torch/visual_gen/test_minimax_h3_pipeline.py + - unittest/_torch/visual_gen/test_minimax_h3_vsa.py - unittest/_torch/visual_gen/test_minimax_h3_tiled_vae.py - unittest/_torch/visual_gen/test_minimax_h3_scheduler_packing.py - unittest/_torch/visual_gen/test_model_loader.py diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py index 47b2f9f08f82..7748d3823181 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py @@ -11,6 +11,7 @@ ) +@pytest.mark.cpu_only def test_h3_vsa_geometry(): metadata = build_h3_vsa_metadata((7, 5, (5, 5, 7)), 0.9, "cpu") valid = metadata.source >= 0 From bd66da3bb05894c1e7f2e9259ca6234a87edda22 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:17:39 +0000 Subject: [PATCH 6/9] Fix MiniMax-H3 synthetic pipeline attention config Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../unittest/_torch/visual_gen/test_minimax_h3_pipeline.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index f71ef53de8c1..615839722442 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -229,7 +229,12 @@ def __init__(self) -> None: # that BasePipeline.device reads. self.workflow = "fl2va" self._device = torch.device("cpu") - self.pipeline_config = SimpleNamespace(visual_gen_mapping=None) + self.pipeline_config = SimpleNamespace( + visual_gen_mapping=None, + primary_model_config=SimpleNamespace( + attention=SimpleNamespace(backend="VANILLA", sparse_attention_config=None) + ), + ) self.offloader = PipelineOffloader(self) self.transformer = _FakeMiniMaxH3Transformer() self.vae = SimpleNamespace( From 159704b007990054286e7c81ce45ab3eac6bc700 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:55:44 +0000 Subject: [PATCH 7/9] Address FastH3 documentation and default VSA review comments Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/visual-generation.md | 23 ++++---- .../models/minimax_h3/pipeline_minimax_h3.py | 56 ++++++++++++++++--- .../visual_gen/test_minimax_h3_pipeline.py | 49 +++++++++++++++- 3 files changed, 108 insertions(+), 20 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index b85673c86e5f..df285bbaa9a1 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -180,16 +180,6 @@ Ulysses requires VANILLA or FA4 attention. CFG, ring, and Attention2D parallelis are not supported. See the [2-GPU TP config](../../../examples/visual_gen/configs/minimax-h3-bf16-tp2.yaml). -- FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a - distilled 4-step checkpoint. Model-specific VSA is the supported inference path - for this checkpoint. Dense attention is a comparison-only path: it does not use - the trained compression gate and is not equivalent to the checkpoint's VSA inference. - For single-GPU text-to-video with audio, enable VSA using - `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. - Text-to-video (T2VA), first/last-frame-to-video (FL2VA), and reference-to-video with audio (Ref2VA) are supported on one GPU. Enable Ref2VA with `pipeline_config: {workflow: ref2va}`; it requires `transformer_ref/` weights. @@ -202,6 +192,19 @@ 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 4-step 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 (90% for this release). 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. + 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 diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index 0b11a8619c78..e0dac95011c6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py @@ -34,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, @@ -47,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, @@ -212,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) ) @@ -464,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.9)) + 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] @@ -515,13 +561,7 @@ def load_standard_components( skip_components = skip_components or [] offloaded = self.offloader.requested_components() - self._fastvideo_inference = {} - fastvideo_inference_path = Path(checkpoint_dir) / "fastvideo_inference.json" - if fastvideo_inference_path.exists(): - with fastvideo_inference_path.open() as f: - fastvideo_inference = json.load(f) - if str(fastvideo_inference.get("schema_version", "")).startswith("fasth3-"): - self._fastvideo_inference = fastvideo_inference + self._fastvideo_inference = self._load_fasth3_manifest(checkpoint_dir) if not _component_skipped(skip_components, PipelineComponent.TOKENIZER): self.tokenizer = Qwen2TokenizerFast.from_pretrained( diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index 615839722442..601207114d41 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -46,7 +46,12 @@ from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm._torch.visual_gen.pipeline_registry import AutoPipeline, PipelineComponent from tensorrt_llm.mapping import Mapping -from tensorrt_llm.visual_gen.args import ParallelConfig, TorchCompileConfig, VisualGenArgs +from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + ParallelConfig, + TorchCompileConfig, + VisualGenArgs, +) pytestmark = pytest.mark.cpu_only @@ -1166,7 +1171,25 @@ def _reject_full_download(*args: object, **kwargs: object) -> None: } -def test_fasth3_manifest_metadata_controls_pipeline_defaults(tmp_path: Path) -> None: +@pytest.mark.parametrize( + "attention_kwargs,expected_backend,expected_sparsity", + [ + ({}, "TRTLLM", 0.9), + ({"backend": "TRTLLM"}, "TRTLLM", 0.9), + ({"backend": "CUTEDSL"}, "CUTEDSL", 0.9), + ( + { + "backend": "CUTEDSL", + "sparse_attention_config": {"algorithm": "vsa", "vsa_sparsity": 0.5}, + }, + "CUTEDSL", + 0.5, + ), + ], +) +def test_fasth3_manifest_metadata_controls_pipeline_defaults( + tmp_path: Path, attention_kwargs: dict, expected_backend: str, expected_sparsity: float +) -> None: (tmp_path / "transformer").mkdir() (tmp_path / "transformer" / "config.json").write_text( json.dumps({"hidden_size": 32}), @@ -1193,6 +1216,19 @@ def test_fasth3_manifest_metadata_controls_pipeline_defaults(tmp_path: Path) -> ) config = DiffusionPipelineConfig.from_pretrained(str(tmp_path)) + config.attention = AttentionConfig(**attention_kwargs) + MiniMaxH3Pipeline._configure_fasth3_attention(config) + assert config.attention.backend == expected_backend + assert config.attention.sparse_attention_config.algorithm == "vsa" + assert config.attention.sparse_attention_config.vsa_sparsity == expected_sparsity + assert config.primary_model_config.attention == config.attention + if expected_backend == "TRTLLM": + assert config.attention_metadata_state == {"metadata_cache": {}} + assert config.primary_model_config.attention_metadata_state == {"metadata_cache": {}} + assert ( + config.primary_model_config.attention_metadata_state + is not config.attention_metadata_state + ) assert "fastvideo_inference" not in config.extra_attrs pipeline = _SyntheticMiniMaxH3Pipeline() @@ -1202,6 +1238,15 @@ def test_fasth3_manifest_metadata_controls_pipeline_defaults(tmp_path: Path) -> ) assert pipeline.default_generation_params["num_inference_steps"] == 5 + config.attention = AttentionConfig(backend="VANILLA") + with pytest.raises(ValueError, match="FastH3 requires VSA"): + MiniMaxH3Pipeline._configure_fasth3_attention(config) + + (tmp_path / "fastvideo_inference.json").unlink() + MiniMaxH3Pipeline._configure_fasth3_attention(config) + assert config.attention.backend == "VANILLA" + assert config.attention.sparse_attention_config is None + def test_infer_supports_a_last_frame_without_a_first_frame( monkeypatch: pytest.MonkeyPatch, From 8270ba81f561335675ba90fe387be689fe6b351a Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:13:43 +0000 Subject: [PATCH 8/9] Defer FastH3 E2E CI enrollment until checkpoint is staged Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 036ecf6f8b60..e74385c948f5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -499,7 +499,6 @@ l0_b200: - unittest/_torch/visual_gen/test_minimax_h3_fused_adaln.py - unittest/_torch/visual_gen/test_minimax_h3_vae_decode.py - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_public_visual_gen_api_smoke TIMEOUT (60) - - examples/visual_gen/test_visual_gen_fasth3.py::test_fasth3_example_lpips_against_golden TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[t2va] TIMEOUT (60) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[fl2va] TIMEOUT (60) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[ref2va] TIMEOUT (30) From f7e085a5d37478ba6ba6ef5cfe51b165c8abaad5 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Tue, 3 Nov 2026 18:33:15 +0000 Subject: [PATCH 9/9] Update FastH3 support to the V2 checkpoint Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/supported-models.md | 2 +- docs/source/models/visual-generation.md | 9 ++- examples/visual_gen/README.md | 2 +- .../configs/fasth3-bf16-trtllm-vsa-1gpu.yaml | 2 +- examples/visual_gen/models/fasth3.py | 3 +- .../models/minimax_h3/pipeline_minimax_h3.py | 41 +++++++++++-- .../visual_gen/test_visual_gen_fasth3.py | 60 ------------------- .../visual_gen/test_minimax_h3_pipeline.py | 58 ++++++++++++++---- 8 files changed, 91 insertions(+), 86 deletions(-) delete mode 100644 tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 4bc6ecf82dd3..1f9023e2c82d 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -289,7 +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-4-step-Preview-v1-VSA-DataFree` | Text-to-Video (with Audio, 4-step distilled) | +| `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 | diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index df285bbaa9a1..e67115044f84 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -38,7 +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-4-step-Preview-v1-VSA-DataFree` | Text-to-Video (with Audio, 4-step distilled) | +| `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 | @@ -193,11 +193,14 @@ are not supported. See the 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 4-step checkpoint. Model-specific VSA is the supported inference path + 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 (90% for this release). Explicit incompatible backends + 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. diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 85ccbdc3762f..3dcf5dc00b47 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -29,7 +29,7 @@ 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 --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml -# FastH3 VSA on Blackwell: +# 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 diff --git a/examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml b/examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml index aafc87cf4036..11ceec547293 100644 --- a/examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml +++ b/examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml @@ -8,7 +8,7 @@ attention_config: backend: TRTLLM sparse_attention_config: algorithm: vsa - vsa_sparsity: 0.9 + vsa_sparsity: 0.8 compilation_config: skip_warmup: true resolutions: diff --git a/examples/visual_gen/models/fasth3.py b/examples/visual_gen/models/fasth3.py index 15da5b129cc7..6c5f289cb881 100644 --- a/examples/visual_gen/models/fasth3.py +++ b/examples/visual_gen/models/fasth3.py @@ -20,7 +20,7 @@ from tensorrt_llm import VisualGen, VisualGenArgs -FASTH3_MODEL_ID = "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree" +FASTH3_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2" def main() -> None: @@ -49,7 +49,6 @@ def main() -> None: params.height = 768 params.width = 1344 params.num_frames = 124 - params.num_inference_steps = 5 params.seed = 42 output = visual_gen.generate( diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index e0dac95011c6..bae7ef7bf9f6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py @@ -174,7 +174,7 @@ def tp_group_name(self) -> str: "MiniMaxH3ModularPipeline", hf_ids=[ "MiniMaxAI/MiniMax-H3", - "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", + "FastVideo/FastVideo-FastH3-8-Step-V2", ], defaults={"workflow": "fl2va"}, download_patterns=[ @@ -496,7 +496,7 @@ def _configure_fasth3_attention(cls, config: DiffusionPipelineConfig) -> None: 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.9)) + 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( @@ -621,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. @@ -1269,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. diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py deleted file mode 100644 index 6c50bd023c3b..000000000000 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_fasth3.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FastH3 example regression against a recorded TRTLLM VSA video. - -This golden detects output drift; it does not establish FastVideo reference parity. -Stage the checkpoint revision recorded in the golden JSON under LLM_MODELS_ROOT. -""" - -import json -from pathlib import Path - -import pytest -import torch -from defs.common import venv_check_call -from defs.examples.visual_gen.visual_gen_test_utils import ( - VISUAL_GEN_LPIPS_GOLDEN_DIR, - _assert_lpips_below_threshold, - _golden_media_path, - _preserve_lpips_candidate_on_failure, - _run_lpips_eval, -) -from test_common.llm_data import get_checkpoint - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -def test_fasth3_example_lpips_against_golden( - request, tmp_path, _visual_gen_deps, llm_root, llm_venv -): - if torch.cuda.get_device_capability()[0] != 10: - pytest.skip("FastH3 TRTLLM VSA requires Blackwell") - metadata_path = Path(VISUAL_GEN_LPIPS_GOLDEN_DIR) / "fasth3_lpips_golden_video.json" - metadata = json.loads(metadata_path.read_text(encoding="utf-8")) - golden_path = _golden_media_path(tmp_path, metadata["video"], "FastH3 TRTLLM regression golden") - model_path = get_checkpoint(metadata["model"]) - output_path = tmp_path / "fasth3_output.mp4" - script_path = Path(llm_root) / "examples/visual_gen/models/fasth3.py" - config_path = Path(llm_root) / "examples/visual_gen/configs/fasth3-bf16-trtllm-vsa-1gpu.yaml" - # Exercise the example's published schedule, seed and geometry defaults. - venv_check_call( - llm_venv, - [ - str(script_path), - "--model", - model_path, - "--visual_gen_args", - str(config_path), - "--prompt", - metadata["prompt"], - "--output_path", - str(output_path), - ], - ) - assert output_path.is_file(), "FastH3 example did not save a video" - score = _run_lpips_eval( - tmp_path, "fasth3", "video", metadata["prompt"], golden_path, output_path - ) - threshold = metadata["lpips_threshold"] - _preserve_lpips_candidate_on_failure(request, score, threshold, output_path, metadata["video"]) - _assert_lpips_below_threshold(score, threshold) diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index 601207114d41..61cbcb05dbd0 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -554,9 +554,7 @@ def test_public_visual_gen_lists_minimax_h3() -> None: from tensorrt_llm import VisualGen assert "MiniMaxAI/MiniMax-H3" in VisualGen.supported_models() - assert ( - "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree" in VisualGen.supported_models() - ) + assert "FastVideo/FastVideo-FastH3-8-Step-V2" in VisualGen.supported_models() def test_modular_manifest_preserves_pipeline_metadata(tmp_path: Path) -> None: @@ -1092,7 +1090,7 @@ def test_default_generation_steps_match_reference_app() -> None: "manifest, expected_steps", [ (None, 28), - ({"schema_version": "fasth3-inference-contract-v1", "num_inference_steps": 5}, 5), + ({"schema_version": "fasth3-inference-contract-v1", "num_inference_steps": 9}, 9), ({"schema_version": "fasth3-inference-contract-v1"}, 28), ({"schema_version": "other-model-v1", "num_inference_steps": 5}, 28), ], @@ -1110,13 +1108,48 @@ def test_fasth3_generation_steps_follow_inference_contract( assert pipeline.default_generation_params["num_inference_steps"] == expected_steps +def test_fasth3_v2_trained_schedule() -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.scheduler = MiniMaxH3Scheduler(shift=10.0) + pipeline.audio_scheduler = MiniMaxH3Scheduler(shift=3.0) + pipeline._fastvideo_inference = { + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0, + } + pipeline.post_load_weights() + pipeline._set_timesteps(9) + base = torch.tensor([0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0]) + for scheduler, shift in ((pipeline.scheduler, 10.0), (pipeline.audio_scheduler, 3.0)): + expected = shift * base / (1 + (shift - 1) * base) + torch.testing.assert_close(scheduler.sigmas.cpu(), expected, rtol=0, atol=0) + torch.testing.assert_close(scheduler.timesteps.cpu(), 1 - expected[:-1], rtol=0, atol=0) + assert len(scheduler.timesteps) == 8 + with pytest.raises(ValueError, match="9 sigma-grid points"): + pipeline._set_timesteps(8) + pipeline.scheduler.set_shift(12.0) + with pytest.raises(ValueError, match="video_scheduler_shift"): + pipeline._set_timesteps(9) + pipeline._fastvideo_inference = {} + pipeline._set_timesteps(28) + assert len(pipeline.scheduler.timesteps) == 27 + + +@pytest.mark.parametrize("steps", [[], [999, 999], [999, 1000], [1001, 500], [999, 0]]) +def test_fasth3_rejects_invalid_dmd_steps(steps: list[int]) -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline._fastvideo_inference = {"dmd_denoising_steps": steps} + with pytest.raises(ValueError, match="strictly decreasing"): + pipeline._set_timesteps(len(steps) + 1) + + def test_hf_download_is_scoped_to_the_supported_h3_components( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: captured: dict[str, object] = {} - model = "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree" - revision = "5ea076f" + model = "FastVideo/FastVideo-FastH3-8-Step-V2" + revision = "3da2ddf" def _capture_partial_download( model: str, @@ -1174,9 +1207,9 @@ def _reject_full_download(*args: object, **kwargs: object) -> None: @pytest.mark.parametrize( "attention_kwargs,expected_backend,expected_sparsity", [ - ({}, "TRTLLM", 0.9), - ({"backend": "TRTLLM"}, "TRTLLM", 0.9), - ({"backend": "CUTEDSL"}, "CUTEDSL", 0.9), + ({}, "TRTLLM", 0.8), + ({"backend": "TRTLLM"}, "TRTLLM", 0.8), + ({"backend": "CUTEDSL"}, "CUTEDSL", 0.8), ( { "backend": "CUTEDSL", @@ -1208,8 +1241,9 @@ def test_fasth3_manifest_metadata_controls_pipeline_defaults( json.dumps( { "schema_version": "fasth3-inference-contract-v1", - "model_id": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", - "num_inference_steps": 5, + "model_id": "FastVideo/FastVideo-FastH3-8-Step-V2", + "num_inference_steps": 9, + "vsa_sparsity": 0.8, } ), encoding="utf-8", @@ -1236,7 +1270,7 @@ def test_fasth3_manifest_metadata_controls_pipeline_defaults( pipeline.load_standard_components( str(tmp_path), torch.device("cpu"), skip_components=list(PipelineComponent) ) - assert pipeline.default_generation_params["num_inference_steps"] == 5 + assert pipeline.default_generation_params["num_inference_steps"] == 9 config.attention = AttentionConfig(backend="VANILLA") with pytest.raises(ValueError, match="FastH3 requires VSA"):