Scope FastH3 metadata loading to the MiniMax-H3 pipeline

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

docs/source/models/visual-generation.md

@@ -173,8 +173,9 @@ Tile geometry and blending follow the loaded Diffusers VAE.
 - 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

examples/visual_gen/README.md

@@ -26,10 +26,8 @@ python models/ltx2.py --visual_gen_args configs/ltx2-1gpu.yaml
 # FP8 blockwise config: configs/minimax-h3-fp8-blockwise-1gpu.yaml.
 # 8-GPU Ulysses config: configs/minimax-h3-bf16-8gpu.yaml.
 python models/minimax_h3.py --model <approved-checkpoint> --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_ti2v.py --visual_gen_args configs/cosmos3-nano-1gpu.yaml --prompt "A robot arm picks fruit in a grocery store"

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

@@ -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(

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

@@ -635,14 +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
-
-
         h3_workflow = resolved_pipeline_config.get("workflow")
         if h3_workflow is not None:
             if h3_workflow not in ("fl2va", "ref2va"):

examples/visual_gen/models/fasth3.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
@@ -307,9 +309,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.
@@ -413,6 +413,14 @@ def load_standard_components(
         del kwargs
         skip_components = skip_components or []
 
+        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,

tensorrt_llm/_torch/visual_gen/config.py

@@ -468,11 +468,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(
@@ -564,7 +579,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(

tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.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,