[TRTLLM-16806][feat] Enable TRTLLM Ulysses attention for MiniMax-H3

docs/source/models/visual-generation.md

@@ -174,7 +174,7 @@ Tile geometry and blending follow the loaded Diffusers VAE.
 Tensor parallelism shards the DiT attention and gated MLP weights via
 `parallel_config.tp_size`. The text encoder, AdaLN modulation, and media
 input/output projections remain replicated. TP can be combined with Ulysses;
-Ulysses requires VANILLA or FA4 attention. CFG, ring, and Attention2D parallelism
+Ulysses supports VANILLA, FA4, TRTLLM, and FLASHINFER attention. CFG, ring, and Attention2D parallelism
 are not supported. See the
 [2-GPU TP config](../../../examples/visual_gen/configs/minimax-h3-bf16-tp2.yaml).
 

tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py

@@ -444,6 +444,79 @@ def _compact_qkv(
         v = v.reshape(batch_size * kv_seq_len, -1).contiguous()
         return q, k, v
 
+    @torch.compiler.disable
+    def _forward_padded_self_attention(
+        self,
+        q: torch.Tensor,
+        k: Optional[torch.Tensor],
+        v: Optional[torch.Tensor],
+        key_padding_mask: torch.Tensor,
+        timestep: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Attend to a valid prefix; callers explicitly discard padded query rows.
+
+        Ulysses still needs the padded output shape for its return collective.
+        Cache prefix lengths across layers sharing the same request mask to avoid
+        repeating the device-to-host synchronization in every attention block.
+        """
+        if key_padding_mask.dtype != torch.bool or key_padding_mask.shape != q.shape[:2]:
+            raise ValueError("TRTLLM padding mask must be boolean with shape [batch, sequence].")
+        state = self.metadata._metadata_state
+        cached = state.get("padding_prefix_lengths")
+        version = None if torch.is_inference(key_padding_mask) else key_padding_mask._version
+        if (
+            version is None
+            or cached is None
+            or cached[0] is not key_padding_mask
+            or cached[1] != version
+        ):
+            mask = key_padding_mask.cpu()
+            lengths = mask.sum(dim=1).tolist()
+            expected = torch.arange(mask.shape[1]).unsqueeze(0) < torch.tensor(lengths).unsqueeze(1)
+            if not torch.equal(mask, expected):
+                raise ValueError("TRTLLM padded self-attention requires trailing padding.")
+            state["padding_prefix_lengths"] = (key_padding_mask, version, lengths)
+        else:
+            lengths = cached[2]
+        batch_size, seq_len = q.shape[:2]
+        if all(length == seq_len for length in lengths):
+            return self.forward(q, k, v, batch_size, seq_len, timestep=timestep)
+
+        groups: dict[int, list[int]] = {}
+        for batch_idx, length in enumerate(lengths):
+            groups.setdefault(length, []).append(batch_idx)
+        if len(groups) == 1:
+            length = lengths[0]
+            if length == 0:
+                return q.new_zeros((batch_size, seq_len, self.num_heads * self.head_dim))
+            output = self.forward(
+                q=q[:, :length].contiguous(),
+                k=k[:, :length].contiguous() if k is not None else None,
+                v=v[:, :length].contiguous() if v is not None else None,
+                batch_size=batch_size,
+                seq_len=length,
+                timestep=timestep,
+            )
+            return torch.nn.functional.pad(output, (0, 0, 0, seq_len - length))
+
+        output = q.new_zeros((batch_size, seq_len, self.num_heads * self.head_dim))
+        for length, batch_indices in groups.items():
+            if length == 0:
+                continue
+            indices = torch.tensor(batch_indices, device=q.device)
+            group_output = self.forward(
+                q=q[:, :length].index_select(0, indices),
+                k=k[:, :length].index_select(0, indices) if k is not None else None,
+                v=v[:, :length].index_select(0, indices) if v is not None else None,
+                batch_size=len(batch_indices),
+                seq_len=length,
+                timestep=timestep,
+            )
+            output.index_copy_(
+                0, indices, torch.nn.functional.pad(group_output, (0, 0, 0, seq_len - length))
+            )
+        return output
+
     def forward(
         self,
         q: torch.Tensor,
@@ -456,6 +529,8 @@ def forward(
         sparse_backend_args: Optional[SparseBackendForwardArgs] = None,
         timestep: Optional[torch.Tensor] = None,
         kv: Optional[ForwardPlan] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        discard_padded_queries: bool = False,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -485,22 +560,38 @@ def forward(
                 prediction hooks. A ``block_sparse_inputs`` payload selects the
                 generic block-sparse FMHA.
             timestep: Denoising timestep consumed by the sparse prediction hooks.
-            **kwargs: Backend-specific keyword arguments forwarded by the attention
-                module, such as ``key_padding_mask``; ignored here, as by the other
-                backends. The TRTLLM kernels have no key-padding input, so callers
-                that need padded keys must guard at the model level or select the
-                ``VANILLA`` backend.
             kv: A ``ForwardPlan`` from a ``BlockCausalKVCacheManager``. When given,
                 ``k``/``v`` are the new tokens only: the fused kernel writes them at
                 ``staging_offset`` and attends over everything cached before them plus
                 themselves, full within a causal block and causal across causal blocks.
                 The first ``kv.num_tokens`` rows are real; rows past them are padding,
                 neither cached nor attended, and output zero. ``batch_size`` must be 1.
                 Sparse attention is not supported on this path.
+            key_padding_mask: Boolean valid-prefix mask, honored only when
+                discard_padded_queries=True. Otherwise retains the legacy TRTLLM
+                behavior of ignoring this backend-specific keyword.
+            discard_padded_queries: Opt in to valid-prefix dense self-attention.
+                Caller guarantees masked query rows are unused; their outputs
+                are zero-filled. Callers requiring general masking must select
+                a backend that supports it, such as VANILLA.
+            **kwargs: Additional backend-specific arguments.
 
         Returns:
             Output tensor [B, S, H*D]
         """
+        if key_padding_mask is not None and discard_padded_queries:
+            if (
+                kv is not None
+                or attention_mask != PredefinedAttentionMask.FULL
+                or seq_len_kv not in (None, seq_len)
+                or self.quant_attention_config is not None
+                or sparse_backend_args is not None
+                or self.sparse_params is not None
+            ):
+                raise ValueError(
+                    "TRTLLM padding requires dense self-attention with discarded padded queries."
+                )
+            return self._forward_padded_self_attention(q, k, v, key_padding_mask, timestep)
         if kv is not None:
             if sparse_backend_args is not None:
                 raise NotImplementedError("Sparse attention is not supported over a K/V cache.")

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

@@ -229,9 +229,14 @@ def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None:
                     "cfg_size=1, ring_size=1, attn2d_size=(1, 1), and "
                     "tp_size * ulysses_size=world_size."
                 )
-            if vgm.ulysses_size > 1 and pipeline_config.attention.backend not in ("VANILLA", "FA4"):
+            if vgm.ulysses_size > 1 and pipeline_config.attention.backend not in (
+                "VANILLA",
+                "FA4",
+                "TRTLLM",
+                "FLASHINFER",
+            ):
                 raise NotImplementedError(
-                    "MiniMax-H3 Ulysses requires VANILLA or FA4 attention. "
+                    "MiniMax-H3 Ulysses requires VANILLA, FA4, TRTLLM, or FLASHINFER attention. "
                     "Other attention backends do not support the required "
                     "packed-row padding mask."
                 )

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

@@ -245,6 +245,7 @@ def forward(
             key,
             value,
             key_padding_mask=key_padding_mask,
+            discard_padded_queries=True,
             timestep=timestep,
         )
         return self.to_out[0](hidden_states)
@@ -285,6 +286,7 @@ def compute_v() -> torch.Tensor:
             compute_k,
             compute_v,
             key_padding_mask=key_padding_mask,
+            discard_padded_queries=True,
             timestep=timestep,
         )
         batch_size, local_seq_len = out_4d.shape[:2]
@@ -477,6 +479,8 @@ def __init__(
         self._supports_key_padding_mask = model_config.attention.backend.upper() in (
             "VANILLA",
             "FA4",
+            "TRTLLM",
+            "FLASHINFER",
         )
         self.sharder = SequenceSharder.from_vgm(
             model_config.visual_gen_mapping,
@@ -512,10 +516,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         if pad:
             if not self._supports_key_padding_mask:
                 raise NotImplementedError(
-                    "Padded token-refiner sequences require key_padding_mask support; use VANILLA or FA4."
+                    "Padded token-refiner sequences require key_padding_mask support; "
+                    "use VANILLA, FA4, TRTLLM or FLASHINFER."
                 )
-            valid = hidden_states.new_ones((hidden_states.shape[0], seq_len), dtype=torch.bool)
-            key_padding_mask = _pad_tensor_dim(valid, 1, pad, False)
+            # Backends cache mask plans using the tensor mutation counter.
+            with torch.inference_mode(False):
+                valid = hidden_states.new_ones((hidden_states.shape[0], seq_len), dtype=torch.bool)
+                key_padding_mask = _pad_tensor_dim(valid, 1, pad, False)
             hidden_states = _pad_tensor_dim(hidden_states, 1, pad)
         hidden_states = self.sharder.shard(hidden_states, dim=1)
         for block in self.refiner_blocks:
@@ -666,7 +673,12 @@ def __init__(self, model_config: DiffusionModelConfig) -> None:
         super().__init__(model_config)
         self._attn_backend = getattr(model_config.attention, "backend", "VANILLA")
         # Only VANILLA supports arbitrary padding between packed modalities.
-        self._supports_key_padding_mask = self._attn_backend in ("VANILLA", "FA4")
+        self._supports_key_padding_mask = self._attn_backend in (
+            "VANILLA",
+            "FA4",
+            "TRTLLM",
+            "FLASHINFER",
+        )
         quant_algo = model_config.quant_config.quant_algo
         # TODO: Support calibrated checkpoints in load_weights before enabling static quantization.
         if quant_algo is not None and not (
@@ -997,19 +1009,23 @@ def forward(
                     "Padded packed sequences (negative token_tags) need an attention "
                     "backend that honours key_padding_mask; the "
                     f"{self._attn_backend} backend cannot express it. Use "
-                    "VANILLA or FA4."
+                    "VANILLA, FA4, TRTLLM or FLASHINFER."
                 )
-            # FA4 expresses padding through seqused_k, which requires valid
-            # tokens to form a prefix. Ulysses appends padding in this layout.
-            if self._attn_backend == "FA4" and bool(
+            # FA4 seqused_k and TRTLLM prefix trimming require trailing padding.
+            # Ulysses appends padding in this layout.
+            if self._attn_backend in ("FA4", "TRTLLM") and bool(
                 ((token_tags_for_mask[:-1] < 0) & (token_tags_for_mask[1:] >= 0)).any()
             ):
                 raise NotImplementedError(
-                    "FA4 requires trailing packed-row padding; use VANILLA for interior padding."
+                    f"{self._attn_backend} requires trailing packed-row padding; use VANILLA for interior padding."
+                )
+            # Preserve a mutation counter for shared backend mask-plan caches.
+            with torch.inference_mode(False):
+                key_padding_mask = (
+                    (token_tags_for_mask >= 0)
+                    .unsqueeze(0)
+                    .expand(packed_hidden_states.shape[0], -1)
                 )
-            key_padding_mask = (
-                (token_tags_for_mask >= 0).unsqueeze(0).expand(packed_hidden_states.shape[0], -1)
-            )
 
         padded_sequence_length = packed_hidden_states.shape[1]
         packed_hidden_states = self.sharder.shard(packed_hidden_states, dim=1)

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

@@ -506,7 +506,8 @@ def test_minimax_h3_text_encoder_tp_lpips(tmp_path, monkeypatch, world_size):
 
 
 @pytest.mark.parametrize(
-    "ulysses_size,vae_size,backend", [(2, 2, "VANILLA"), (4, 3, "VANILLA"), (2, 2, "FA4")]
+    "ulysses_size,vae_size,backend",
+    [(2, 2, "VANILLA"), (4, 3, "VANILLA"), (2, 2, "FA4"), (2, 2, "TRTLLM"), (2, 2, "FLASHINFER")],
 )
 def test_minimax_h3_ulysses_parallel_vae_lpips(tmp_path, ulysses_size, vae_size, backend):
     """Compare real Ulysses + tiled VAE collectives with a fresh one-GPU baseline."""

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

@@ -41,12 +41,16 @@ l0_dgx_b200:
   - unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-2-1]
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-2-1]
+  - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[TRTLLM-2-1]
+  - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FLASHINFER-2-1]
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_ulysses_and_parallel_vae
   - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_tp2_example_lpips TIMEOUT (30)
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py::test_minimax_h3_text_encoder_tp[tp2_world2]
   - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_text_encoder_tp_lpips[2] TIMEOUT (30)
   - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-VANILLA] TIMEOUT (30)
   - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-FA4] TIMEOUT (30)
+  - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-TRTLLM] TIMEOUT (30)
+  - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-FLASHINFER] TIMEOUT (30)
   - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxParallelism::test_ulysses_2gpu_correctness
   - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxCombinedOptimizations::test_all_optimizations_combined
 - condition:
@@ -67,8 +71,12 @@ l0_dgx_b200:
   - unittest/disaggregated/test_transfer_lifecycle_gpu.py
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-4-1]
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-4-1]
+  - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[TRTLLM-4-1]
+  - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FLASHINFER-4-1]
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-2-2]
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-2-2]
+  - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[TRTLLM-2-2]
+  - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FLASHINFER-2-2]
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_partial_vae_groups
   - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py::test_minimax_h3_text_encoder_tp[tp2_world4]
   - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_text_encoder_tp_lpips[4] TIMEOUT (30)

tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py

@@ -38,8 +38,8 @@ def _make_model_config(
     return DiffusionModelConfig(
         pretrained_config=SimpleNamespace(
             num_attention_heads=2,
-            attention_head_dim=128 if backend == "FA4" else 16,
-            hidden_size=256 if backend == "FA4" else 32,
+            attention_head_dim=128 if backend in ("FA4", "TRTLLM", "FLASHINFER") else 16,
+            hidden_size=256 if backend in ("FA4", "TRTLLM", "FLASHINFER") else 32,
             num_layers=2,
             num_refiner_layers=1,
             ffn_dim=32,
@@ -190,8 +190,10 @@ def _logic_h3_ulysses_and_parallel_vae(
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two CUDA GPUs")
 @pytest.mark.parametrize("async_ulysses", [False, True], ids=["nccl", "async"])
-@pytest.mark.parametrize("backend", ["VANILLA", "FA4"])
+@pytest.mark.parametrize("backend", ["VANILLA", "FA4", "TRTLLM", "FLASHINFER"])
 def test_h3_ulysses_and_parallel_vae(backend: str, async_ulysses: bool) -> None:
+    if backend == "TRTLLM" and torch.cuda.get_device_capability() not in ((10, 0), (10, 3)):
+        pytest.skip("TRTLLM H3 requires SM100 or SM103")
     if backend == "FA4" and torch.cuda.get_device_capability()[0] < 10:
         pytest.skip("FA4 requires SM100 or newer")
     if torch.cuda.device_count() < 2:
@@ -261,11 +263,13 @@ def test_h3_partial_vae_groups() -> None:
 
 
 @pytest.mark.parametrize("tp_size,ulysses_size", [(2, 1), (4, 1), (2, 2)])
-@pytest.mark.parametrize("backend", ["VANILLA", "FA4"])
+@pytest.mark.parametrize("backend", ["VANILLA", "FA4", "TRTLLM", "FLASHINFER"])
 def test_h3_tensor_parallel(tp_size: int, ulysses_size: int, backend: str) -> None:
     world_size = tp_size * ulysses_size
     if torch.cuda.device_count() < world_size:
         pytest.skip(f"Requires {world_size} CUDA GPUs")
+    if backend == "TRTLLM" and torch.cuda.get_device_capability() not in ((10, 0), (10, 3)):
+        pytest.skip("TRTLLM H3 requires SM100 or SM103")
     if backend == "FA4" and torch.cuda.get_device_capability()[0] < 10:
         pytest.skip("FA4 requires SM100 or newer")
     run_distributed(

tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py

@@ -1649,7 +1649,7 @@ def test_pipeline_rejects_unsupported_ulysses_topology(field: str, value: int) -
         MiniMaxH3Pipeline(config)
 
 
-@pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL"])
+@pytest.mark.parametrize("backend", ["CUTEDSL", "CUDNN"])
 def test_pipeline_rejects_unvalidated_ulysses_attention(backend: str) -> None:
     config = _ulysses_pipeline_config()
     config.attention.backend = backend
@@ -1780,3 +1780,16 @@ def initialize(self, config):
     config.visual_gen_mapping.ulysses_size = ulysses_size
     MiniMaxH3Pipeline(config)
     assert initialized == [config]
+
+
+@pytest.mark.parametrize("backend", ["TRTLLM", "FLASHINFER"])
+def test_pipeline_accepts_masked_ulysses_backends(
+    monkeypatch: pytest.MonkeyPatch, backend: str
+) -> None:
+    monkeypatch.setattr(
+        h3_pipeline.BasePipeline, "__init__", lambda self, config: torch.nn.Module.__init__(self)
+    )
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0))
+    config = _ulysses_pipeline_config(2)
+    config.attention.backend = backend
+    MiniMaxH3Pipeline(config)