[All] Refactor nvte_get_fused_attn_backend with cudnn-frontend calls

tests/jax/test_distributed_fused_attn.py

@@ -83,6 +83,7 @@ def impl_test_self_attn(
 
         if not is_fused_attn_kernel_available(
             is_training,
+            batch,
             dtype,
             dtype,
             QKVLayout.BS3HD,
@@ -235,6 +236,7 @@ def test_cross_attn(
 
         if not is_fused_attn_kernel_available(
             is_training,
+            batch,
             dtype,
             dtype,
             QKVLayout.BSHD_BS2HD,
@@ -425,6 +427,7 @@ def impl_test_context_parallel_attn(
         def check_has_backend_for_mask(mask_type):
             return is_fused_attn_kernel_available(
                 is_training,
+                batch,
                 dtype,
                 dtype,
                 qkv_layout,

tests/jax/test_fused_attn.py

@@ -520,8 +520,9 @@ def _check_configs(self):
                 "is either BSHD_BSHD_BSHD or THD_THD_THD"
             )
 
-        self.backend = FusedAttnHelper(
+        self.backend, message = FusedAttnHelper(
             self.is_training,
+            self.batch_size,
             self.dtype,
             self.dtype,
             self.qkv_layout,
@@ -536,9 +537,10 @@ def _check_configs(self):
             self.head_dim_qk,
             self.head_dim_v,
             (-1, -1) if self.window_size is None else self.window_size,
+            self.attn_mask_type.is_bottom_right(),
         ).get_fused_attn_backend()
         if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen:
-            pytest.skip("Unsupported inputs combination or device compute capability.")
+            pytest.skip(message)
 
         if (
             self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS

tests/jax/test_fused_attn_score_mod.py

@@ -18,7 +18,7 @@
 )
 from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config
 from transformer_engine.jax.flax import transformer as flax_transformer
-from transformer_engine_jax import get_device_compute_capability
+from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend
 from test_fused_attn import FusedAttnRunner, SeqDescFormat
 
 
@@ -397,9 +397,14 @@ def _identity_score_mod(_graph, score, _tensors):
 def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True):
     captured = {}
 
-    def fake_fused_attn_kernel_check(*args, **kwargs):
-        captured.setdefault("kernel_checks", []).append((args, kwargs))
-        return kernel_available
+    class FakeFusedAttnHelper:
+        def __init__(self, *args, **kwargs):
+            captured.setdefault("kernel_checks", []).append((args, kwargs))
+
+        def get_fused_attn_backend(self):
+            if kernel_available:
+                return NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, ""
+            return NVTE_Fused_Attn_Backend.NVTE_No_Backend, "fake: no backend"
 
     def fake_fused_attn(
         qkv,
@@ -454,11 +459,7 @@ def fake_fused_attn(
         )
         return qkv[0]
 
-    monkeypatch.setattr(
-        flax_transformer,
-        "is_fused_attn_kernel_available",
-        fake_fused_attn_kernel_check,
-    )
+    monkeypatch.setattr(flax_transformer, "FusedAttnHelper", FakeFusedAttnHelper)
     monkeypatch.setattr(flax_transformer, "fused_attn", fake_fused_attn)
     return captured
 

tests/pytorch/attention/test_attention.py

@@ -1772,12 +1772,23 @@ def test_dpa_fp8_extra_state(model, dtype):
     config = model_configs_fp8_extra_state[model]
     # Test backend availability
     is_training = True
+    fp8_recipe = recipe.DelayedScaling(
+        margin=0,
+        fp8_format=recipe.Format.HYBRID,
+        amax_history_len=1,
+        amax_compute_algo="most_recent",
+        fp8_dpa=True,
+    )
+    fp8_meta = {}
+    fp8_meta["recipe"] = fp8_recipe
     available_backends, _, fused_attn_backends = get_available_attention_backends(
         config,
         qkv_dtype=torch.float8_e4m3fn,
         qkv_layout="sb3hd",
         is_training=is_training,
         deterministic=_deterministic,
+        fp8=True,
+        fp8_meta=fp8_meta,
     )
     flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends
     if not fused_attn_supported and not flash_attn_supported:
@@ -2566,13 +2577,25 @@ def test_custom_mha_fp8_vs_f16(dtype, model):
     Both paths take F16 input and output. QKV layout is bs3hd"""
 
     config = model_configs_fp8[model]
+    os.environ["NVTE_UnfusedDPA_Emulate_FP8"] = "1"
 
     # Test backend availability
     is_training = True
+    fp8_meta = {}
+    fp8_recipe = recipe.DelayedScaling(
+        margin=0,
+        fp8_format=recipe.Format.HYBRID,
+        amax_history_len=1,
+        amax_compute_algo="most_recent",
+        fp8_dpa=True,
+    )
+    fp8_meta["recipe"] = fp8_recipe
     available_backends, _, fused_attn_backends = get_available_attention_backends(
         config,
         qkv_dtype=torch.float8_e4m3fn,
         qkv_layout="bs3hd",
+        fp8=True,
+        fp8_meta=fp8_meta,
         is_training=is_training,
         deterministic=_deterministic,
     )
@@ -2650,6 +2673,7 @@ def _run_custom_mha_fp8(dtype, config, backend):
         fp8_format=recipe.Format.HYBRID,
         amax_history_len=1,
         amax_compute_algo="most_recent",
+        fp8_dpa=True,
     )
 
     mha = Custom_MHA_FP8(config).to(dtype=dtype, device="cuda")

tests/pytorch/utils.py

@@ -311,6 +311,10 @@ def __init__(
         self.attn_type = "self" if (self.max_seqlen_q == self.max_seqlen_kv) else "cross"
         self.bias_shape = bias_shape
         self.window_size = check_set_window_size(self.attn_mask_type, window_size)
+        self.bottom_right_diagonal = self.attn_mask_type not in {
+            "causal",
+            "padding_causal",
+        }
         self.context_parallel = context_parallel
         self.cp_comm_type = cp_comm_type
         self.return_max_logit = return_max_logit
@@ -357,9 +361,15 @@ def get_available_attention_backends(
         if config.bias_shape == "bhss":
             alibi_slopes_shape = [config.batch_size, config.num_heads]
 
-    core_attention_bias_shape = (
-        config.bias_shape if config.attn_bias_type == "post_scale_bias" else None
-    )
+    core_attention_bias_shape = None
+    if config.attn_bias_type == "post_scale_bias":
+        b_dim, h_dim, sq_dim, skv_dim = config.bias_shape
+        core_attention_bias_shape = (
+            config.batch_size if b_dim == "b" else 1,
+            config.num_heads if h_dim == "h" else 1,
+            config.max_seqlen_q if sq_dim == "s" else 1,
+            config.max_seqlen_kv if skv_dim == "s" else 1,
+        )
     core_attention_bias_requires_grad = False
     # d=256 is supported by cuDNN 9.0+ for inference but not training
     if (
@@ -368,7 +378,7 @@ def get_available_attention_backends(
         and config.head_dim_v <= 128
     ):
         # TODO(KshitijLakhani): Remove this guard when cuDNN starts support dbias calculation for bias shape 111s
-        if core_attention_bias_shape != "111s":
+        if config.bias_shape != "111s":
             core_attention_bias_requires_grad = True
 
     fused_attn_backends = []
@@ -389,6 +399,7 @@ def test():
             head_dim_v=config.head_dim_v,
             attn_mask_type=config.attn_mask_type,
             window_size=config.window_size,
+            bottom_right_diagonal=config.bottom_right_diagonal,
             alibi_slopes_shape=alibi_slopes_shape,
             core_attention_bias_type=config.attn_bias_type,
             core_attention_bias_shape=core_attention_bias_shape,

transformer_engine/common/CMakeLists.txt

@@ -181,6 +181,7 @@ list(APPEND transformer_engine_cpp_sources
      cudnn_utils.cpp
      transformer_engine.cpp
      fused_attn/fused_attn.cpp
+     fused_attn/config_and_params.cpp
      gemm/config.cpp
      normalization/common.cpp
      normalization/rtc_dispatch.cpp