diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 58aeea645067..1f9023e2c82d 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-8-Step-V2` | Text-to-Video (with Audio, 8-step distilled) | | `Qwen/Qwen-Image` | Text-to-Image | | `Qwen/Qwen-Image-2512` | Text-to-Image | | `Qwen/Qwen-Image-Layered` | Image-to-Image | @@ -310,6 +311,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page | **LTX-2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | | **LTX-2.3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | No | No | No | No | | **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | +| **FastH3** | Yes | Yes | No | No | No | No | No | Yes | Yes | No | No | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | | **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 430995eda3f7..e67115044f84 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-8-Step-V2` | Text-to-Video (with Audio, 8-step distilled) | | `Qwen/Qwen-Image` | Text-to-Image | | `Qwen/Qwen-Image-2512` | Text-to-Image | | `Qwen/Qwen-Image-Layered` | Image-to-Image | @@ -68,6 +69,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | | **LTX-2.3** | Yes | Yes | No | No | No | No | Yes | Yes | Yes | Yes | Yes | No | No | No | No | No | | **MiniMax-H3** | Yes | Yes | No | No | Yes | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | No | +| **FastH3** | Yes | Yes | No | No | No | No | No | No | No | Yes | Yes | No | No | No | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | No | | **Qwen-Image-Layered** | Yes | No | No | Yes | No | No | No | No | Yes | Yes | Yes | No | No | No | No | No | | **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | No | @@ -190,6 +192,22 @@ are not supported. See the - The published [MiniMax-H3 checkpoint license](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) restricts use by territory. Obtain legal approval before downloading or running the weights. +- FastH3 uses the MiniMax-H3 pipeline and transformer architecture with a + distilled 8-step V2 checkpoint. Model-specific VSA is the supported inference path + for this checkpoint and is selected by default with the TRTLLM backend and the + checkpoint's sparsity (80% for V2). Explicit incompatible backends + or sparse attention algorithms are rejected. Dense attention omits the trained + compression gate and is not equivalent to this checkpoint's VSA inference. + The example reads nine sigma-grid points (eight transformer forwards) from + `fastvideo_inference.json` and uses its trained DMD ladder with video/audio + scheduler shifts of 10/3. Changing only the step count is not equivalent. + For single-GPU text-to-video with audio, see + `fasth3-bf16-trtllm-vsa-1gpu.yaml`. To use CUTEDSL instead, set + `attention_config.backend: CUTEDSL` in that config. + Both backends preserve the trained compression gate and dense text/audio routes. + CUTEDSL VSA requires supported Blackwell hardware and at most 4096 tiles; + unsupported shapes fail explicitly rather than silently running dense attention. + ## Optimizations ### Quantization diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 1e3532a44033..3dcf5dc00b47 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -29,6 +29,8 @@ python models/ltx2.py --visual_gen_args configs/ltx2-1gpu.yaml # 2-GPU DiT tensor parallelism config: configs/minimax-h3-bf16-tp2.yaml. # 2-GPU Ulysses config: configs/minimax-h3-bf16-2gpu.yaml. python models/minimax_h3.py --model --visual_gen_args configs/minimax-h3-bf16-1gpu.yaml +# FastH3 8-Step V2 VSA on Blackwell: +python models/fasth3.py --visual_gen_args configs/fasth3-bf16-trtllm-vsa-1gpu.yaml python models/flux1.py --visual_gen_args configs/flux1-dev-fp4-1gpu.yaml python models/flux2.py --visual_gen_args configs/flux2-dev-fp4-1gpu.yaml python models/cosmos3/cosmos3.py --visual_gen_args configs/cosmos3-nano-1gpu.yaml --prompt "A robot arm picks fruit in a grocery store" 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..11ceec547293 --- /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.8 +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 new file mode 100644 index 000000000000..6c5f289cb881 --- /dev/null +++ b/examples/visual_gen/models/fasth3.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FastH3 text-to-video generation with stereo audio.""" + +import argparse + +from tensorrt_llm import VisualGen, VisualGenArgs + +FASTH3_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + default=FASTH3_MODEL_ID, + help="FastH3 checkpoint path or Hugging Face model ID.", + ) + parser.add_argument( + "--visual_gen_args", + default="configs/fasth3-bf16-trtllm-vsa-1gpu.yaml", + help="Path to YAML config (same as trtllm-serve --visual_gen_args).", + ) + parser.add_argument( + "--output_path", + default="fasth3_t2va_output.mp4", + help="Path to save the generated video and audio.", + ) + parser.add_argument("--prompt", default=None, help="Prompt.") + args = parser.parse_args() + + extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None + visual_gen = VisualGen(model=args.model, args=extra_args) + params = visual_gen.default_params + params.height = 768 + params.width = 1344 + params.num_frames = 124 + params.seed = 42 + + output = visual_gen.generate( + inputs=args.prompt + or ( + "A cinematic wide shot of an astronaut walking through a neon-lit " + "rainforest at dusk, with soft ambient music and detailed natural sounds." + ), + params=params, + ) + saved = output.save(args.output_path) + print(f"Saved: {saved}") + + +if __name__ == "__main__": + main() 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..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 @@ -17,10 +17,12 @@ """TRT-LLM VisualGen pipeline for MiniMax-H3 FL2VA and Ref2VA checkpoints.""" +import json import time from collections.abc import Iterator from contextlib import contextmanager, nullcontext from io import BytesIO +from pathlib import Path from typing import TYPE_CHECKING, Any, Optional import numpy as np @@ -32,7 +34,10 @@ from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor from tensorrt_llm._torch.visual_gen.checkpoints import WeightLoader -from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionPipelineConfig, + create_attention_metadata_state, +) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import ( BasePipeline, @@ -45,6 +50,8 @@ from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger from tensorrt_llm.mapping import MpiTopology +from tensorrt_llm.visual_gen.args import AttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import VideoSparseAttentionConfig from .packing import ( MINIMAX_H3_AUDIO_CHANNELS, @@ -75,6 +82,7 @@ from .ref2va import load_references, prepare_references, validate_reference_order from .text_encoder_tp import MiniMaxH3TensorParallelTextEncoder from .transformer_minimax_h3 import MiniMaxH3Transformer3DModel +from .vsa import build_h3_vsa_metadata, set_h3_vsa_context if TYPE_CHECKING: from ...sleep import PipelineSleepManager, SleepRestoreMode @@ -164,10 +172,16 @@ def tp_group_name(self) -> str: @register_pipeline( "MiniMaxH3ModularPipeline", - hf_ids=["MiniMaxAI/MiniMax-H3"], + hf_ids=[ + "MiniMaxAI/MiniMax-H3", + "FastVideo/FastVideo-FastH3-8-Step-V2", + ], defaults={"workflow": "fl2va"}, download_patterns=[ "modular_model_index.json", + "fastvideo_inference.json", + "checkpoint_content.json", + "checkpoint_metadata.json", "LICENSE", "README.md", "transformer/*", @@ -203,6 +217,7 @@ def _request_generator(seed: int) -> torch.Generator: return torch.Generator().manual_seed(seed) def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None: + self._configure_fasth3_attention(pipeline_config) self.text_encoder_tp_size = int( getattr(getattr(pipeline_config, "parallel", None), "text_encoder_tp_size", 1) ) @@ -371,6 +386,7 @@ def setup_parallel_vae(self) -> None: @property def default_generation_params(self) -> dict: + fastvideo_inference = getattr(self, "_fastvideo_inference", {}) return { # MiniMax-H3 derives the default canvas from the first present # keyframe, or from 16:9 for T2VA. @@ -378,7 +394,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 @@ -454,6 +470,46 @@ def _run_warmup( num_inference_steps=steps, ) + @staticmethod + def _load_fasth3_manifest(checkpoint_dir: str) -> dict: + manifest_path = Path(checkpoint_dir) / "fastvideo_inference.json" + if manifest_path.exists(): + with manifest_path.open() as f: + manifest = json.load(f) + if str(manifest.get("schema_version", "")).startswith("fasth3-"): + return manifest + return {} + + @classmethod + def _configure_fasth3_attention(cls, config: DiffusionPipelineConfig) -> None: + pretrained = getattr(config, "primary_pretrained_config", None) + checkpoint_dir = getattr(pretrained, "_name_or_path", None) + manifest = cls._load_fasth3_manifest(checkpoint_dir) if checkpoint_dir else {} + if not manifest: + return + + attention = config.attention + backend = attention.backend + if backend == "VANILLA" and "backend" not in attention.model_fields_set: + backend = "TRTLLM" + if backend not in ("TRTLLM", "CUTEDSL"): + raise ValueError("FastH3 requires VSA with the TRTLLM or CUTEDSL attention backend.") + sparse = attention.sparse_attention_config + if sparse is None: + sparse = VideoSparseAttentionConfig(vsa_sparsity=manifest.get("vsa_sparsity", 0.8)) + elif sparse.algorithm != "vsa": + raise ValueError("FastH3 requires the vsa sparse attention algorithm.") + config.attention = AttentionConfig.model_validate( + {**attention.model_dump(), "backend": backend, "sparse_attention_config": sparse} + ) + # Resolve defaults before constructing the transformer and loading its trained gate. + if backend == "TRTLLM" and config.attention_metadata_state is None: + config.attention_metadata_state = create_attention_metadata_state() + for model_config in config.model_configs.values(): + model_config.attention = config.attention.model_copy(deep=True) + if backend == "TRTLLM" and model_config.attention_metadata_state is None: + model_config.attention_metadata_state = create_attention_metadata_state() + def _init_transformer(self) -> None: self.transformer = MiniMaxH3Transformer3DModel( self.pipeline_config.model_configs[PipelineComponent.TRANSFORMER] @@ -505,6 +561,8 @@ def load_standard_components( skip_components = skip_components or [] offloaded = self.offloader.requested_components() + self._fastvideo_inference = self._load_fasth3_manifest(checkpoint_dir) + if not _component_skipped(skip_components, PipelineComponent.TOKENIZER): self.tokenizer = Qwen2TokenizerFast.from_pretrained( checkpoint_dir, @@ -563,13 +621,43 @@ def load_standard_components( def post_load_weights(self) -> None: super().post_load_weights() - if self.scheduler is not None: + # Distilled checkpoints carry their trained shifts in the scheduler configs. + if self.scheduler is not None and not getattr(self, "_fastvideo_inference", {}): self.scheduler.set_shift(self.VIDEO_SCHEDULER_SHIFT) - if self.audio_scheduler is not None: + if self.audio_scheduler is not None and not getattr(self, "_fastvideo_inference", {}): self.audio_scheduler.set_shift(self.AUDIO_SCHEDULER_SHIFT) if self.vae is not None and self.device.type == "cuda": _prepare_vae_decoder(self.vae) + def _set_timesteps(self, num_inference_steps: int) -> None: + manifest = getattr(self, "_fastvideo_inference", {}) + steps = manifest.get("dmd_denoising_steps") + if steps is None: + self.scheduler.set_timesteps(num_inference_steps, device=self.device) + self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) + return + if ( + not steps + or any(type(step) is not int or not 0 < step <= 1000 for step in steps) + or any(left <= right for left, right in zip(steps, steps[1:])) + ): + raise ValueError("FastH3 DMD steps must be strictly decreasing integers in (0, 1000].") + if num_inference_steps != len(steps) + 1: + raise ValueError( + f"FastH3 requires {len(steps) + 1} sigma-grid points for {len(steps)} forwards." + ) + base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32) + for scheduler, shift_key in ( + (self.scheduler, "video_scheduler_shift"), + (self.audio_scheduler, "audio_scheduler_shift"), + ): + shift = float(scheduler.shift) + if shift_key in manifest and shift != float(manifest[shift_key]): + raise ValueError(f"FastH3 {shift_key} disagrees with the checkpoint scheduler.") + # Shift each modality's trained noise clock once; explicit sigmas are not shifted. + sigmas = shift * base / (1 + (shift - 1) * base) + scheduler.set_timesteps(sigmas=sigmas, device=self.device) + def torch_compile(self) -> None: super().torch_compile() # The ViT VAE decoder is 36 identical blocks run once per clip; compile them like the transformer blocks. @@ -1211,8 +1299,7 @@ def _forward( condition_latents=condition_latents, ) - self.scheduler.set_timesteps(num_inference_steps, device=self.device) - self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) + self._set_timesteps(num_inference_steps) # The denoise loop indexes both schedules with one counter, and # set_timesteps collapses float32 duplicates, so unequal lengths would # silently truncate a stream instead of failing. @@ -1257,6 +1344,31 @@ def _forward( self.scheduler.timesteps, self.audio_scheduler.timesteps ) + _attn_cfg = self.pipeline_config.primary_model_config.attention + _sparse_cfg = getattr(_attn_cfg, "sparse_attention_config", None) + _vsa_active = ( + getattr(_attn_cfg, "backend", "VANILLA") in ("CUTEDSL", "TRTLLM") + and _sparse_cfg is not None + and getattr(_sparse_cfg, "algorithm", None) == "vsa" + ) + _vsa_sparsity = _sparse_cfg.vsa_sparsity if _vsa_active else 0.0 + vsa_metadata = None + if _vsa_active: + if self.workflow == "ref2va" or keyframes: + raise ValueError("MiniMax-H3 VSA currently supports text-to-video only.") + patch = self.transformer.config.patch_size + grid = tuple( + size // p + for size, p in zip( + (num_latent_frames, latent_height, latent_width), patch, strict=True + ) + ) + vsa_metadata = build_h3_vsa_metadata( + (text_indices.numel(), audio_indices.numel(), grid), + _vsa_sparsity, + self.device, + ) + def forward_fn( video_latents: torch.Tensor, extra_stream_latents: dict[str, torch.Tensor], @@ -1267,7 +1379,7 @@ def forward_fn( ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: del timestep, encoder_hidden_states, extra_tensors unique_timesteps, timestep_indices = row_timestep_plan[index] - video_velocity, audio_velocity = self.transformer( + transformer_kwargs = dict( hidden_states=torch.cat((condition_prefix, video_latents), dim=1), audio_hidden_states=torch.cat( (audio_condition_prefix, extra_stream_latents["audio"]), dim=1 @@ -1284,6 +1396,11 @@ def forward_fn( return_dict=False, static_context=static_context, ) + if _vsa_active: + with set_h3_vsa_context(vsa_metadata): + video_velocity, audio_velocity = self.transformer(**transformer_kwargs) + else: + video_velocity, audio_velocity = self.transformer(**transformer_kwargs) # A silent all-zero / non-finite velocity has been observed once in # practice: the run completed, reported success, and wrote a black # video with silent audio. The cause was never reproduced, so fail 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..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 @@ -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,29 @@ 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 + self._vsa_backend = model_config.attention.backend + if sparse is not None and sparse.algorithm == "vsa": + 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( + 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 +270,25 @@ 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, + backend=self._vsa_backend, + ) + return self.to_out[0](output.flatten(2)) + hidden_states = self._attn_impl( query, key, @@ -423,7 +472,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..b05c9bae640e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/vsa.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +""" + +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 _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) + + 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 + 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 + 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/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/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..e74385c948f5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -494,6 +494,7 @@ 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 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_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index 8f43b29c4c11..61cbcb05dbd0 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 @@ -229,7 +234,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( @@ -544,6 +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-8-Step-V2" in VisualGen.supported_models() def test_modular_manifest_preserves_pipeline_metadata(tmp_path: Path) -> None: @@ -1075,11 +1086,70 @@ def test_default_generation_steps_match_reference_app() -> None: assert pipeline.default_generation_params["num_inference_steps"] == 28 +@pytest.mark.parametrize( + "manifest, expected_steps", + [ + (None, 28), + ({"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), + ], +) +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.load_standard_components( + str(tmp_path), torch.device("cpu"), skip_components=list(PipelineComponent) + ) + + 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-8-Step-V2" + revision = "3da2ddf" def _capture_partial_download( model: str, @@ -1095,7 +1165,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 +1176,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 +1200,88 @@ def _reject_full_download(*args: object, **kwargs: object) -> None: "scheduler/*", "audio_scheduler/*", ], - "revision": "5d9b308a59ab12e67147f191e184baf704185bd1", + "revision": revision, } +@pytest.mark.parametrize( + "attention_kwargs,expected_backend,expected_sparsity", + [ + ({}, "TRTLLM", 0.8), + ({"backend": "TRTLLM"}, "TRTLLM", 0.8), + ({"backend": "CUTEDSL"}, "CUTEDSL", 0.8), + ( + { + "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}), + 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-8-Step-V2", + "num_inference_steps": 9, + "vsa_sparsity": 0.8, + } + ), + encoding="utf-8", + ) + + 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() + 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"] == 9 + + 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, ) -> None: 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..7748d3823181 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_vsa.py @@ -0,0 +1,88 @@ +# 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.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.models.minimax_h3.vsa import ( + build_h3_vsa_metadata, + h3_vsa_attention, +) + + +@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 + 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.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_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(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, + 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(batch, seq, 2, 128) for x in (qt, kt, vt)] + actual = backend.forward( + *tensors, batch_size=batch, seq_len=seq, sparse_backend_args=sparse_backend_args + ) + routes = sparse_backend_args.block_sparse_inputs + 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() + 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(batch, -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) + 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)