[None][perf] Request FA4's 2-CTA kernel for MiniMax-H3 on SM103 and SM107

docs/source/models/visual-generation.md

@@ -183,6 +183,10 @@ are not supported. See the
   role `reference`. Audio requires an image or video reference. See the
   [Ref2VA example](../../../examples/visual_gen/README.md#minimax-h3-reference-to-video-and-audio)
   for input limits, ordering, and media dependencies.
+- With the FA4 backend, MiniMax-H3 requests FA4's 2-CTA forward kernel for its transformer
+  blocks on SM103 and SM107, where it is the faster BF16 attention at the clip lengths H3
+  runs; SM100 keeps FA4's default. An explicit `FA_DISABLE_2CTA` in the environment overrides
+  the request either way.
 - MiniMax-H3 currently restricts TRTLLM attention to SM100 or SM103. This is a
   model-specific restriction, not a general VisualGen backend requirement.
 - The published [MiniMax-H3 checkpoint license](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE)

tensorrt_llm/_torch/visual_gen/attention/attention.py

@@ -73,7 +73,14 @@ def __init__(
         enable_sequence_parallel: bool = True,
         async_ulysses: bool = False,
         separate_qkv_is_self_attention: bool = False,
+        backend_kwargs: Optional[dict] = None,
     ):
+        """
+        Args:
+            backend_kwargs: Extra keyword arguments for the compute backend's constructor
+                (``create_attention``), for options a model sets per backend such as FA4's
+                ``use_2cta``. The caller is responsible for matching them to the backend.
+        """
         super().__init__()
 
         config = config or DiffusionModelConfig()
@@ -286,6 +293,7 @@ def __init__(
             attention_config=config.attention,
             attention_metadata_state=attention_metadata_state,
             sparse_params=sparse_params,
+            **(backend_kwargs or {}),
         )
 
         if enable_sequence_parallel and is_separate_qkv_cross_attention and vgm is not None:

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

@@ -19,16 +19,23 @@
 Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16.
 """
 
+import contextlib
 import math
 import os
-from typing import Optional, Tuple
+from typing import Iterator, Optional, Tuple
 
 import torch
 
+from tensorrt_llm.logger import logger
+
 from ....attention.backends.interface import PredefinedAttentionMask
 from .interface import AttentionBackend, AttentionTensorLayout
 
-# 2CTA regresses long sequences; must precede the import below.
+# FA4 folds FA_DISABLE_2CTA into a module flag when flash_attn.cute.utils is imported below,
+# so the process-wide default is fixed here: the 2-CTA kernel regresses long sequences on
+# SM 100. A model that has measured otherwise on its part asks for the kernel per backend
+# instance through ``use_2cta``; an explicit FA_DISABLE_2CTA in the environment wins over it.
+_FA_DISABLE_2CTA_FROM_ENV = "FA_DISABLE_2CTA" in os.environ
 os.environ.setdefault("FA_DISABLE_2CTA", "1")
 
 
@@ -108,6 +115,37 @@ def is_valid_tile(self) -> Boolean:
     _flash_attn_fwd_import_error = e
 
 
+def fa4_2cta_requestable() -> bool:
+    """Whether a backend-level ``use_2cta`` request takes effect in this process.
+
+    False when FA_DISABLE_2CTA was set explicitly in the environment (the user's choice wins)
+    or when FA4 is not importable.
+    """
+    return _flash_attn_fwd is not None and not _FA_DISABLE_2CTA_FROM_ENV
+
+
+@contextlib.contextmanager
+def _fa4_2cta_request(use_2cta: Optional[bool]) -> Iterator[None]:
+    """Steer FA4's 2-CTA kernel choice for the enclosed forward.
+
+    ``flash_attn.cute.utils`` reads its ``_fa_disable_2cta_enabled`` flag on every forward, so a
+    request is applied by setting that flag around the call and restoring it afterwards. FA4's
+    own opt-out on CUDA 12 toolkits (``_fa_disable_2cta_cuda12``) stays in force. ``None`` leaves
+    the process default untouched.
+    """
+    if use_2cta is None or not fa4_2cta_requestable():
+        yield
+        return
+    from flash_attn.cute import utils as fa4_utils
+
+    previous = fa4_utils._fa_disable_2cta_enabled
+    fa4_utils._fa_disable_2cta_enabled = not use_2cta
+    try:
+        yield
+    finally:
+        fa4_utils._fa_disable_2cta_enabled = previous
+
+
 class FlashAttn4Attention(AttentionBackend):
     """
     Flash Attention 4 backend for diffusion models.
@@ -125,14 +163,28 @@ def __init__(
         head_dim: int = 64,
         num_kv_heads: Optional[int] = None,
         dtype: Optional[torch.dtype] = None,
+        use_2cta: Optional[bool] = None,
         **kwargs,
     ):
+        """
+        Args:
+            use_2cta: Request (True) or refuse (False) FA4's 2-CTA forward kernel for this
+                backend instance; ``None`` keeps FA4's process default. Ignored when
+                FA_DISABLE_2CTA is set explicitly in the environment. FA4 still applies its
+                own eligibility rules (non-causal, head_dim 128 / 192, long enough sequence).
+        """
         self.layer_idx = layer_idx
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.num_kv_heads = num_kv_heads or num_heads
         self.dtype = dtype
         self.scale = 1.0 / math.sqrt(head_dim)
+        self.use_2cta = use_2cta if fa4_2cta_requestable() else None
+        if use_2cta is not None and self.use_2cta is None and layer_idx == 0:
+            logger.info(
+                f"FA4 use_2cta={use_2cta} requested but FA_DISABLE_2CTA is set in the "
+                "environment; keeping the environment's choice."
+            )
 
         # FA4 expects [B, S, H, D] format
         self._preferred_layout = AttentionTensorLayout.NHD
@@ -148,23 +200,24 @@ def _fwd(
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Calls _flash_attn_fwd with torch.compile disabled. Returns (output, lse)."""
         # FA4's private forward API may append diagnostics that this backend does not consume.
-        output, lse, *_ = _flash_attn_fwd(
-            q,
-            k,
-            v,
-            seqused_k=seqused_k,
-            softmax_scale=self.scale,
-            causal=causal,
-            window_size_left=None,
-            window_size_right=None,
-            learnable_sink=None,
-            softcap=0.0,
-            pack_gqa=None,
-            mask_mod=None,
-            block_sparse_tensors=None,
-            return_lse=True,
-            num_splits=0,
-        )
+        with _fa4_2cta_request(self.use_2cta):
+            output, lse, *_ = _flash_attn_fwd(
+                q,
+                k,
+                v,
+                seqused_k=seqused_k,
+                softmax_scale=self.scale,
+                causal=causal,
+                window_size_left=None,
+                window_size_right=None,
+                learnable_sink=None,
+                softcap=0.0,
+                pack_gqa=None,
+                mask_mod=None,
+                block_sparse_tensors=None,
+                return_lse=True,
+                num_splits=0,
+            )
         return output, lse
 
     def _prepare_inputs(

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

@@ -119,6 +119,23 @@ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
         return freqs.cos(), freqs.sin()
 
 
+# Compute capabilities on which FA4's 2-CTA forward kernel is the faster BF16 attention for
+# H3's shapes (56 heads x 128, 19k to 104k packed tokens): SM 103 (B300) and SM 107 (R200).
+# On SM 100 the process default stays, where the same kernel regresses long sequences.
+FA4_2CTA_CAPABILITIES = frozenset({(10, 3), (10, 7)})
+
+
+def fa4_use_2cta_default(attn_backend: str) -> Optional[bool]:
+    """Return the H3 default for FA4's 2-CTA kernel on the current device.
+
+    ``True`` requests the kernel on the parts in ``FA4_2CTA_CAPABILITIES``; ``None`` keeps FA4's
+    process default everywhere else, for other backends and without a CUDA device.
+    """
+    if attn_backend.upper() != "FA4" or not torch.cuda.is_available():
+        return None
+    return True if torch.cuda.get_device_capability() in FA4_2CTA_CAPABILITIES else None
+
+
 class MiniMaxH3Attention(Attention):
     """Full non-causal attention with per-head RMSNorm and partial 3D RoPE."""
 
@@ -132,12 +149,16 @@ def __init__(
         model_config: DiffusionModelConfig,
         layer_idx: int,
         module_name: str,
+        fa4_use_2cta: Optional[bool] = None,
     ) -> None:
         vgm = model_config.visual_gen_mapping
         ulysses_size = vgm.ulysses_size if vgm is not None else 1
         # With ulysses_size > 1 and parallel.async_ulysses, use separate Q/K/V projections so
         # each one is pushed to the peers while the next is computed (forward_async).
         use_async_ulysses = bool(ulysses_size > 1 and model_config.parallel.async_ulysses)
+        backend_kwargs = None
+        if fa4_use_2cta is not None and model_config.attention.backend.upper() == "FA4":
+            backend_kwargs = {"use_2cta": fa4_use_2cta}
         super().__init__(
             hidden_size=hidden_size,
             num_attention_heads=num_attention_heads,
@@ -157,6 +178,7 @@ def __init__(
             module_name=module_name,
             async_ulysses=use_async_ulysses,
             separate_qkv_is_self_attention=True,
+            backend_kwargs=backend_kwargs,
         )
         self.use_async_ulysses = use_async_ulysses
 
@@ -549,6 +571,7 @@ def __init__(
         model_config: DiffusionModelConfig,
         layer_idx: int,
         fuse_bf16_gate_up_swiglu: bool = False,
+        fa4_use_2cta: Optional[bool] = None,
     ) -> None:
         super().__init__()
         self.norm1 = RMSNorm(
@@ -562,6 +585,7 @@ def __init__(
             model_config=model_config,
             layer_idx=layer_idx,
             module_name=f"transformer_blocks.{layer_idx}.attn",
+            fa4_use_2cta=fa4_use_2cta,
         )
         self.norm2 = RMSNorm(
             hidden_size=hidden_size, eps=norm_eps, dtype=model_config.torch_dtype, has_weights=True
@@ -667,6 +691,9 @@ def __init__(self, model_config: DiffusionModelConfig) -> None:
         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")
+        # The transformer blocks' attention runs at the packed clip length; the token refiner
+        # keeps FA4's process default.
+        self.fa4_use_2cta = fa4_use_2cta_default(self._attn_backend)
         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 (
@@ -824,6 +851,7 @@ def __init__(self, model_config: DiffusionModelConfig) -> None:
                     model_config=model_config,
                     layer_idx=layer_idx,
                     fuse_bf16_gate_up_swiglu=fuse_bf16_gate_up_swiglu,
+                    fa4_use_2cta=self.fa4_use_2cta,
                 )
                 for layer_idx in range(num_layers)
             ]

tests/unittest/_torch/visual_gen/test_attention_fa4.py

@@ -194,3 +194,42 @@ def fa4(splits):
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])
+
+
+def test_use_2cta_request_toggles_fa4_flag_for_the_call_only():
+    """``use_2cta`` flips FA4's module flag around the forward and restores it afterwards."""
+    from flash_attn.cute import utils as fa4_utils
+
+    from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import (
+        _fa4_2cta_request,
+        fa4_2cta_requestable,
+    )
+
+    before = fa4_utils._fa_disable_2cta_enabled
+    with _fa4_2cta_request(True):
+        assert fa4_utils._fa_disable_2cta_enabled is (
+            before if not fa4_2cta_requestable() else False
+        )
+    assert fa4_utils._fa_disable_2cta_enabled is before
+    with _fa4_2cta_request(False):
+        assert fa4_utils._fa_disable_2cta_enabled is (
+            before if not fa4_2cta_requestable() else True
+        )
+    assert fa4_utils._fa_disable_2cta_enabled is before
+    with _fa4_2cta_request(None):
+        assert fa4_utils._fa_disable_2cta_enabled is before
+
+
+@pytest.mark.parametrize("use_2cta", [True, False])
+def test_use_2cta_forward_matches_default_kernel(use_2cta):
+    """Both FA4 kernel choices give the same attention output on a 2-CTA-eligible shape."""
+    from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import fa4_2cta_requestable
+
+    if not fa4_2cta_requestable():
+        pytest.skip("FA_DISABLE_2CTA set explicitly in the environment")
+    torch.manual_seed(0)
+    # 2-CTA eligibility needs head_dim 128 and seqlen_q > 2 x tile_m.
+    q, k, v = (torch.randn(1, 2048, 8, 128, dtype=torch.bfloat16, device="cuda") for _ in range(3))
+    reference = FlashAttn4Attention(num_heads=8, head_dim=128).forward(q, k, v)
+    candidate = FlashAttn4Attention(num_heads=8, head_dim=128, use_2cta=use_2cta).forward(q, k, v)
+    torch.testing.assert_close(candidate, reference, rtol=2e-2, atol=2e-2)

tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py

@@ -85,6 +85,7 @@ def _make_model_config(
     quant_config: QuantConfig | None = None,
     dynamic_weight_quant: bool = False,
     force_dynamic_quantization: bool = False,
+    attention_backend: str = "VANILLA",
     **overrides: object,
 ) -> DiffusionModelConfig:
     pretrained_config_dict = _TINY_CONFIG | overrides
@@ -97,7 +98,7 @@ def _make_model_config(
         dynamic_weight_quant=dynamic_weight_quant,
         force_dynamic_quantization=force_dynamic_quantization,
         mapping=Mapping(),
-        attention=AttentionConfig(backend="VANILLA"),
+        attention=AttentionConfig(backend=attention_backend),
     )
 
 
@@ -1692,3 +1693,65 @@ def test_load_weights_accepts_diffusers_layout_with_async_ulysses(
     # Every parameter, including the re-ordered gate_up_proj and the renamed down_proj.
     for name, parameter in model.named_parameters():
         torch.testing.assert_close(parameter, expected[name].to(parameter.dtype), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "backend,capability,expected",
+    [
+        ("FA4", (10, 0), None),
+        ("FA4", (10, 3), True),
+        ("FA4", (10, 7), True),
+        ("FA4", (9, 0), None),
+        ("VANILLA", (10, 3), None),
+        ("CUTEDSL", (10, 7), None),
+    ],
+)
+def test_fa4_2cta_default_follows_device_capability(
+    monkeypatch: pytest.MonkeyPatch,
+    backend: str,
+    capability: tuple[int, int],
+    expected: bool | None,
+) -> None:
+    """H3 asks for FA4's 2-CTA kernel on SM 103 and SM 107 only; elsewhere FA4's default stands."""
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability)
+    assert h3.fa4_use_2cta_default(backend) is expected
+
+
+def test_fa4_2cta_default_without_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert h3.fa4_use_2cta_default("FA4") is None
+
+
+def test_fa4_2cta_request_reaches_transformer_block_backends(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The request lands on every transformer block's FA4 backend, not on the token refiner's."""
+    from tensorrt_llm._torch.visual_gen.attention.backends.flash_attn4 import (
+        FlashAttn4Attention,
+        fa4_2cta_requestable,
+    )
+
+    monkeypatch.setattr(h3, "fa4_use_2cta_default", lambda attn_backend: True)
+    model = h3.MiniMaxH3Transformer3DModel(
+        _make_model_config(attention_backend="FA4", num_layers=1, num_refiner_layers=1)
+    )
+
+    assert model.fa4_use_2cta is True
+    expected = True if fa4_2cta_requestable() else None
+    for block in model.transformer_blocks:
+        backend = block.attn.attn
+        assert isinstance(backend, FlashAttn4Attention)
+        assert backend.use_2cta is expected
+    refiner_backend = model.token_refiner.refiner_blocks[0].attn.attn
+    assert isinstance(refiner_backend, FlashAttn4Attention)
+    assert refiner_backend.use_2cta is None
+
+
+def test_fa4_2cta_request_is_not_sent_to_other_backends(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A VANILLA backend never sees FA4's keyword, even when the default would request it."""
+    monkeypatch.setattr(h3, "fa4_use_2cta_default", lambda attn_backend: True)
+    model = h3.MiniMaxH3Transformer3DModel(
+        _make_model_config(attention_backend="VANILLA", num_layers=1, num_refiner_layers=1)
+    )
+    assert not hasattr(model.transformer_blocks[0].attn.attn, "use_2cta")