diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 19553ecab5a4..3080fb71900a 100644 --- a/docs/source/models/visual-generation.md +++ b/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) diff --git a/tensorrt_llm/_torch/visual_gen/attention/attention.py b/tensorrt_llm/_torch/visual_gen/attention/attention.py index eb32925a0882..abce3afe489b 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/attention.py +++ b/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: diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention/backends/flash_attn4.py index fa34d17570e4..e57e365d2da0 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/flash_attn4.py +++ b/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( 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..a1ab154bc103 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 @@ -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) ] diff --git a/tests/unittest/_torch/visual_gen/test_attention_fa4.py b/tests/unittest/_torch/visual_gen/test_attention_fa4.py index db6a0c0ea37d..8b21dce02888 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_fa4.py +++ b/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) diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py index 2ed43d61e447..62cf2f38008e 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py +++ b/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")