fa4 support

Signed-off-by: Xin Yao <xiny@nvidia.com>

Author: yaox12

tests/pytorch/attention/test_attention.py

@@ -53,7 +53,7 @@
 )
 
 _current_file = pathlib.Path(__file__).resolve()
-sys.path.append(str(_current_file.parent.parent))
+sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import (
     reset_rng_states,
     compare_and_assert,
@@ -362,6 +362,141 @@ def test_dpa_num_splits(dtype, model_configs, model):
     )
 
 
+# ==============================
+# Flash Attention 4 (FA4) tests
+# ==============================
+
+model_configs_fa4_base = {
+    # test: ModelConfig(b, sq, hq, dqk)
+    # Standard head dims
+    "fa4_base_1": ModelConfig(4, 128, 16, 64),
+    "fa4_base_2": ModelConfig(2, 2048, 24, 128, attn_mask_type="causal"),
+    "fa4_base_3": ModelConfig(2, 1024, 8, 96, attn_mask_type="causal"),
+    # GQA
+    "fa4_gqa_1": ModelConfig(2, 1024, 32, 128, num_gqa_groups=8, attn_mask_type="causal"),
+    "fa4_gqa_2": ModelConfig(2, 1024, 16, 128, num_gqa_groups=1, attn_mask_type="causal"),
+    # num_splits
+    "fa4_splits_1": ModelConfig(2, 2048, 24, 128, num_splits=2),
+    "fa4_splits_2": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096, num_splits=4),
+}
+
+
+@pytest.mark.skipif(
+    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
+)
+@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.")
+@pytest.mark.parametrize("dtype", param_types_lean)
+@pytest.mark.parametrize("model_configs", [model_configs_fa4_base])
+@pytest.mark.parametrize("model", model_configs_fa4_base.keys())
+def test_dpa_fa4_base(dtype, model_configs, model):
+    """Test DotProductAttention with FA4: base configs, extended head dims, GQA, num_splits"""
+    test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False)
+
+
+model_configs_fa4_mla = {
+    # test: ModelConfig(b, sq, hq, dqk, head_dim_v=dv)
+    "fa4_mla_1": ModelConfig(4, 128, 16, 128, head_dim_v=64),
+    "fa4_mla_2": ModelConfig(2, 128, 16, 64, max_seqlen_kv=256, head_dim_v=128),
+    "fa4_mla_3": ModelConfig(2, 1024, 16, 96, head_dim_v=64, attn_mask_type="causal"),
+    # dqk=128, dv=96: FA4 SM100 backward has dK_reduce_ncol misalignment for dV;
+    # the backend filter should reject FA4 and fall back to another backend.
+    "fa4_mla_4": ModelConfig(2, 1024, 16, 128, head_dim_v=96, attn_mask_type="causal"),
+    # DeepSeek-style MLA: dqk=192, dv=128 (supported on SM100 as special case)
+    "fa4_mla_deepseek": ModelConfig(2, 1024, 16, 192, head_dim_v=128, attn_mask_type="causal"),
+}
+
+
+@pytest.mark.skipif(
+    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
+)
+@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.")
+@pytest.mark.parametrize("dtype", param_types_lean)
+@pytest.mark.parametrize("model_configs", [model_configs_fa4_mla])
+@pytest.mark.parametrize("model", model_configs_fa4_mla.keys())
+def test_dpa_fa4_mla(dtype, model_configs, model):
+    """Test DotProductAttention with FA4: MLA (head_dim_qk != head_dim_v)"""
+    test_dot_product_attention(
+        dtype, model_configs, model, False, True, "bshd_bshd_bshd", False, False
+    )
+
+
+model_configs_fa4_swa = {
+    # test: ModelConfig(b, sq, hq, dqk, window_size=(left, right))
+    "fa4_swa_1": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal", window_size=(128, 0)),
+    "fa4_swa_2": ModelConfig(2, 2048, 24, 64, attn_mask_type="causal", window_size=(64, 0)),
+    "fa4_swa_3": ModelConfig(
+        2, 2048, 16, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(256, 0)
+    ),
+    "fa4_swa_4": ModelConfig(
+        2, 2048, 16, 128, attn_mask_type="padding_causal", window_size=(128, 0)
+    ),
+}
+
+
+@pytest.mark.skipif(
+    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
+)
+@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.")
+@pytest.mark.parametrize("dtype", param_types_lean)
+@pytest.mark.parametrize("model_configs", [model_configs_fa4_swa])
+@pytest.mark.parametrize("model", model_configs_fa4_swa.keys())
+@pytest.mark.parametrize("qkv_layout", ["sbhd_sbhd_sbhd", "bshd_bshd_bshd"])
+def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout):
+    """Test DotProductAttention with FA4: sliding window attention"""
+    test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, True, False)
+
+
+model_configs_fa4_varlen = {
+    # test: ModelConfig(b, sq, hq, dqk)
+    "fa4_varlen_1": ModelConfig(4, 128, 16, 64, attn_mask_type="padding"),
+    "fa4_varlen_2": ModelConfig(2, 2048, 24, 128, attn_mask_type="padding_causal"),
+    "fa4_varlen_3": ModelConfig(
+        2, 2048, 24, 128, num_gqa_groups=4, attn_mask_type="padding_causal"
+    ),
+    "fa4_varlen_4": ModelConfig(
+        2, 128, 16, 64, max_seqlen_kv=256, attn_mask_type="padding"
+    ),
+}
+
+
+@pytest.mark.skipif(
+    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
+)
+@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.")
+@pytest.mark.parametrize("dtype", param_types_lean)
+@pytest.mark.parametrize("model_configs", [model_configs_fa4_varlen])
+@pytest.mark.parametrize("model", model_configs_fa4_varlen.keys())
+@pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "bshd_bshd_bshd"])
+def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout):
+    """Test DotProductAttention with FA4: variable-length sequences (varlen/thd)"""
+    test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False)
+
+
+model_configs_fa4_mask = {
+    # test: ModelConfig(b, sq, hq, dqk)
+    "fa4_mask_no_mask": ModelConfig(2, 1024, 16, 128),
+    "fa4_mask_causal": ModelConfig(2, 1024, 16, 128, attn_mask_type="causal"),
+    "fa4_mask_causal_br": ModelConfig(2, 1024, 16, 128, attn_mask_type="causal_bottom_right"),
+    "fa4_mask_padding": ModelConfig(2, 1024, 16, 128, attn_mask_type="padding"),
+    "fa4_mask_padding_causal": ModelConfig(2, 1024, 16, 128, attn_mask_type="padding_causal"),
+    "fa4_mask_padding_causal_br": ModelConfig(
+        2, 1024, 16, 128, attn_mask_type="padding_causal_bottom_right"
+    ),
+}
+
+
+@pytest.mark.skipif(
+    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
+)
+@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.")
+@pytest.mark.parametrize("dtype", param_types_lean)
+@pytest.mark.parametrize("model_configs", [model_configs_fa4_mask])
+@pytest.mark.parametrize("model", model_configs_fa4_mask.keys())
+def test_dpa_fa4_mask(dtype, model_configs, model):
+    """Test DotProductAttention with FA4: various attention mask types"""
+    test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False)
+
+
 model_configs_softmax = {
     # test: ModelConfig(b, sq, hq, dqk)
     "softmax_1_0": ModelConfig(2, 2048, 64, 64, num_gqa_groups=8),

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -156,20 +156,19 @@
 
 # Try to import Flash Attention v4
 try:
-    fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-cute"))
+    fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-4"))
 except PackageNotFoundError:
     flash_attn_func_v4 = None
     flash_attn_varlen_func_v4 = None
-    flash_attn_with_kvcache_v4 = None
+    flash_attn_combine_v4 = None
     _flash_attn_fwd_v4 = None
     _flash_attn_bwd_v4 = None
-    # pass  # only print warning if use_flash_attention_4 = True in get_attention_backend
 else:
     from flash_attn.cute.interface import flash_attn_func as flash_attn_func_v4
     from flash_attn.cute.interface import flash_attn_varlen_func as flash_attn_varlen_func_v4
+    from flash_attn.cute.interface import flash_attn_combine as flash_attn_combine_v4
     from flash_attn.cute.interface import _flash_attn_fwd as _flash_attn_fwd_v4
     from flash_attn.cute.interface import _flash_attn_bwd as _flash_attn_bwd_v4
-    # flash_attn_with_kvcache_v4 = None  # FA4 does not support kvcache yet
     fa_utils.set_flash_attention_4_params()
 
 # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16
@@ -942,7 +941,7 @@ def forward(
         if flash_attention_backend is not None and flash_attention_backend > PkgVersion("3.0.0b"):
             use_flash_attn_3 = True
         use_flash_attn_4 = False
-        if flash_attention_backend is not None and str(flash_attention_backend).endswith("cute"):
+        if flash_attention_backend is not None and flash_attention_backend > PkgVersion("4.0.0b"):
             use_flash_attn_4 = True
         if context_parallel and all(
             not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer]
@@ -996,6 +995,9 @@ def forward(
                 #       |                         | thd + padding
                 #       | flash_attn_with_kvcache | KV cache (not-paged/paged), i.e.
                 #       |                         |     bshd/sbhd/thd + padding
+                # FA v4 | flash_attn_func         | bshd/sbhd + not padding
+                #       | flash_attn_varlen_func  | bshd/sbhd + padding
+                #       |                         | thd + padding
                 fa_optional_forward_args_thd = []
                 if qkv_format in ["bshd", "sbhd"] and "padding" not in attn_mask_type:
                     func = None
@@ -1006,24 +1008,33 @@ def forward(
                     else:
                         func = flash_attn_func
                 else:
-                    if not use_flash_attn_3:
+                    if use_flash_attn_4:
+                        func = flash_attn_varlen_func_v4
+                    elif not use_flash_attn_3:
                         func = flash_attn_varlen_func
                     elif inference_params is None:
                         func = flash_attn_varlen_func_v3  # pylint: disable=possibly-used-before-assignment
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
-                    if not use_flash_attn_3 or inference_params is None:
+                    if not use_flash_attn_4 and (
+                        not use_flash_attn_3 or inference_params is None
+                    ):
                         fa_optional_forward_args_thd.append(cu_seqlens_q)
                         fa_optional_forward_args_thd.append(cu_seqlens_kv)
                         fa_optional_forward_args_thd.append(max_seqlen_q)
                         fa_optional_forward_args_thd.append(max_seqlen_kv)
                 if use_flash_attn_4:
                     fa_4_optional_forward_kwargs = {
-                        # "window_size": window_size,
+                        "window_size": window_size,
                         "num_splits": num_splits,
                     }
                     if inference_params is None:
                         fa_4_optional_forward_kwargs["deterministic"] = self.deterministic
+                    if func is flash_attn_varlen_func_v4:
+                        fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q
+                        fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv
+                        fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q
+                        fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv
                     output = func(
                         query_layer,
                         key_layer,

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -145,8 +145,7 @@ class FlashAttentionUtils:
     fa4_version = PkgVersion("0")
     use_v4 = False
     v4_installation_steps = """\
-(1) git clone https://github.com/Dao-AILab/flash-attention.git
-(2) pip install flash-attention/flash_attn/cute"""
+(1) pip install flash-attn-4"""
     v4_warning_printed = False
 
     @staticmethod
@@ -460,13 +459,10 @@ def get_attention_backend(
         if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
             logger.debug("Disabling FlashAttention 3 for compute capability != sm90")
         use_flash_attention_3 = False
-    # TODO: Other compute capabilities support:
-    #  SM80: not enabled
-    #  SM90: has bugs
-    #  SM120: WIP
-    if device_compute_capability != (10, 0):
+    # FA4 supports SM80, SM90, SM100, SM120
+    if device_compute_capability < (8, 0):
         if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
-            logger.debug("Disabling FlashAttention 4 for compute capability != sm100")
+            logger.debug("Disabling FlashAttention 4 for compute capability < sm80")
         use_flash_attention_4 = False
 
     # Filter: Data type
@@ -588,7 +584,7 @@ def get_attention_backend(
     # Flash v2 | FP16/BF16      | non-paged/paged | sm80+        | bshd,sbhd,thd | >= 256
     # Flash v3 | FP16/BF16      | non-paged/paged | sm90         | bshd,sbhd,thd | >= 1
     #          | FP8            | non-paged/paged | sm90         | thd           | >= 1
-    # Flash v4 | N/A            | N/A             | N/A          | N/A           | N/A
+    # Flash v4 | FP16/BF16      | TODO            | sm80+        | bshd,sbhd,thd | TODO
     # Unfused  | FP32/FP16/BF16 | non-paged/paged | all          | bshd,sbhd,thd | >= 1
     if inference_params is not None:
         # Temporarily disabling fused attention for kv caching for sm89 irrespective of cuDNN version
@@ -642,9 +638,6 @@ def get_attention_backend(
         if use_flash_attention_2 and FlashAttentionUtils.is_installed:
             logger.debug("Disabling FlashAttention 2 as it does not support MLA.")
             use_flash_attention_2 = False
-        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
-            logger.debug("Disabling FlashAttention 4 as it does not support MLA.")
-            use_flash_attention_4 = False
 
         qkv_layout_group = qkv_layout.replace("b", "").replace("s", "").replace("t", "")
         if use_fused_attention and qkv_layout_group != "hd_hd_hd":
@@ -717,17 +710,50 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             )
             use_flash_attention_3 = False
 
-    if (
-        use_flash_attention_4
-        and FlashAttentionUtils.v4_is_installed
-        and (head_dim_qk != head_dim_v or head_dim_qk not in [64, 96, 128])
-    ):
-        logger.debug(
-            "Disabling FlashAttention 4 due to unsupported head_dim_qk and head_dim_v. "
-            "Supported: head_dim_qk == head_dim_v, head_dim_qk in [64, 96, 128]. "
-            f"Found: head_dim_qk = {head_dim_qk}, head_dim_v = {head_dim_v}."
-        )
-        use_flash_attention_4 = False
+    if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
+        # FA4 head dimension support is architecture-dependent
+        # (matches _validate_head_dims in flash_attn.cute.interface):
+        #   SM90:      head_dim <= 256 and head_dim_v <= 256
+        #   SM100/110: head_dim <= 128 and head_dim_v <= 128,
+        #              OR DeepSeek MLA shape (head_dim=192, head_dim_v=128)
+        #   SM80/120:  constrained by shared memory (~256 max in practice)
+        _fa4_hdim_ok = True
+        if device_compute_capability >= (10, 0) and device_compute_capability < (12, 0):
+            _is_standard = head_dim_qk <= 128 and head_dim_v <= 128
+            _is_deepseek = head_dim_qk == 192 and head_dim_v == 128
+            _fa4_hdim_ok = _is_standard or _is_deepseek
+        else:
+            _fa4_hdim_ok = head_dim_qk <= 256 and head_dim_v <= 256
+        if not _fa4_hdim_ok:
+            logger.debug(
+                "Disabling FlashAttention 4 due to unsupported head dimensions. "
+                f"Found: head_dim_qk = {head_dim_qk}, head_dim_v = {head_dim_v}, "
+                f"on sm{device_compute_capability[0] * 10 + device_compute_capability[1]}."
+            )
+            use_flash_attention_4 = False
+        # Workaround: SM100 backward kernel bug when MLA + 2CTA (head_dim_qk >= 128).
+        # FlashAttentionBackwardSm100 computes dK_reduce_ncol = gcd(32, tile_hdim // 2)
+        # based on Q/K head_dim but reuses it for dV TMEM load atoms. When
+        # (tile_hdimv // 2) % dK_reduce_ncol != 0, dV reads are misaligned.
+        # See: flash_attn/cute/flash_bwd_sm100.py, line ~262 and ~3890.
+        elif (
+            _fa4_hdim_ok
+            and is_training
+            and head_dim_qk != head_dim_v
+            and head_dim_qk >= 128
+            and device_compute_capability >= (10, 0)
+            and device_compute_capability < (12, 0)
+        ):
+            _tile_hdim = math.ceil(head_dim_qk / 16) * 16
+            _tile_hdimv = math.ceil(head_dim_v / 16) * 16
+            _dk_reduce_ncol = math.gcd(32, _tile_hdim // 2)
+            if (_tile_hdimv // 2) % _dk_reduce_ncol != 0:
+                logger.debug(
+                    "Disabling FlashAttention 4 for training due to SM100 backward kernel "
+                    "bug with MLA head dimensions (dK_reduce_ncol misalignment for dV). "
+                    f"Found: head_dim_qk = {head_dim_qk}, head_dim_v = {head_dim_v}."
+                )
+                use_flash_attention_4 = False
 
     # Filter: QKV layout
     if qkv_format == "thd":
@@ -749,10 +775,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
                     " not supported for compute capability = sm120"
                 )
             use_fused_attention = False
-        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
-            logger.debug("Disabling FlashAttention 4 for qkv_format = thd")
-            use_flash_attention_4 = False
-
     # Filter: Dropout
     if attention_dropout != 0.0:
         if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
@@ -816,6 +838,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             "Disabling UnfusedDotProductAttention as it does not support context parallelism"
         )
         use_unfused_attention = False
+    if context_parallel and use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
+        logger.debug(
+            "Disabling FlashAttention 4 as it does not support context parallelism yet"
+        )
+        use_flash_attention_4 = False
     if context_parallel and (
         use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
     ):
@@ -1228,10 +1255,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         flash_attention_backend = FlashAttentionUtils.version
     if use_flash_attention_3:
         flash_attention_backend = FlashAttentionUtils.fa3_version
-    # FA4 is released with the package name "flash-attn-cute" and version starting from 0.1.0
-    # We need to add the ".cute" suffix to the version number to distinguish.
     if use_flash_attention_4:
-        flash_attention_backend = PkgVersion(f"{str(FlashAttentionUtils.fa4_version)}+cute")
+        flash_attention_backend = FlashAttentionUtils.fa4_version
 
     logger.debug(
         "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s,"
@@ -1248,12 +1273,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
     )
 
     # Select FusedAttention for performance
+    # FA4 is preferred over FusedAttention when available
     if use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0):
-        logger.debug(
-            "Disabling FlashAttention to give FusedAttention preference on Hopper+ "
-            "for performance reasons"
-        )
-        use_flash_attention = False
+        if not use_flash_attention_4:
+            logger.debug(
+                "Disabling FlashAttention to give FusedAttention preference on Hopper+ "
+                "for performance reasons"
+            )
+            use_flash_attention = False
 
     # Selected backend
     if use_flash_attention: