[All] Refactor fused attention APIs with cuDNN-frontend support checks and opaque config/params handles

docs/envvars.rst

@@ -190,18 +190,18 @@ backend-selection overview.
    :Default: ``1``
    :Description: Enable or disable UnfusedDotProductAttention backend (native PyTorch). When set to ``0``, UnfusedDotProductAttention will not be used.
 
-.. envvar:: NVTE_FUSED_ATTN_BACKEND
-
-   :Type: ``int`` (1 or 2)
-   :Default: Auto-selected
-   :Description: Request a cuDNN FusedAttention backend when that request is supported by the active fused-attention path. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. BF16/FP16 attention uses sub-backend ``1`` when eligible. FP8 attention uses sub-backend ``2`` when FP8 DPA is enabled and supported by the architecture, cuDNN version, and input configuration.
-
 .. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD
 
    :Type: ``int`` (0 or 1)
    :Default: ``0``
    :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention.
 
+.. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG
+
+   :Type: ``int`` (0, 1 or 2), optionally followed by ``:<ranks>``
+   :Default: ``0``
+   :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints an end-of-run summary of the cache counters and the mean time of each cuDNN build stage. ``2`` additionally traces every event as it happens: each graph built, each graph cuDNN accepts and the cache keeps, each lookup and whether it hit or missed, each first execution that compiles kernels, and each execution. When the launcher exports a rank, only rank 0 logs; append ``:<ranks>`` to override, as in ``1:all`` for level 1 on every rank or ``2:0,3`` for level 2 on ranks 0 and 3.
+
 .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO
 
    :Type: ``int`` (0 or 1)

docs/examples/attention/attention.ipynb

@@ -249,11 +249,7 @@
     "NVTE_DEBUG       = 0/1   # disables/enables debugging\n",
     "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n",
     "```\n",
-    "<div class=\"alert alert-info\">\n",
-    "<b>Note:</b>\n",
-    "    \n",
-    "These flags are supported in PyTorch only as of Transformer Engine 2.0. JAX support is expected to be added in the future.\n",
-    "</div>"
+    "These flags are available in both PyTorch and Jax in TE v2.20 onwards."
    ]
   },
   {
@@ -330,6 +326,51 @@
     "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "fc0c9e3d",
+   "metadata": {},
+   "source": [
+    "For the FusedAttention backend, cuDNN graphs are cached and reused, and the following flag helps understand the activities taking place in the cache. `NVTE_FUSED_ATTN_CACHE_DEBUG` is supported in both PyTorch and JAX.\n",
+    "```\n",
+    "NVTE_FUSED_ATTN_CACHE_DEBUG = 0         # disables cache diagnostics\n",
+    "NVTE_FUSED_ATTN_CACHE_DEBUG = 1/2:ranks # enables cache diagnostics (events)/(tracing):select ranks for information collection\n",
+    "```\n",
+    "At level 1, these cache events are collected and reported in the end-of-run summary: `create_graph` (number of graphs created), `cache_graph` (number of graphs cached), `hit` (number of cache hits), `miss` (number of cache misses), `build_plans` (number of graphs whose plans are built), and `execute` (number of plan executions). These event counters are tallied up per backend (f16 and fp8), per pass (forward and backward), and per thread. They help paint a picture of the graph reuse rate and caching effectiveness. Another part of the end-of-run summary are the build times of various stages of the cuDNN graph. Out of the 5 cuDNN calls, `validate`, `build_operation_graph`, `create_execution_plans`, `check_support`, and `build_plans`, `build_plans` is the most expensive one, because it is where the kernels get compiled. It happens once per eligible graph and is expected to execute many times in a real-life run.\n",
+    "\n",
+    "For more verbose diagnostics, level 2 enables tracing for all the events taking place in the cache. For example, for every `miss` event, a \"MISS\" line is printed in the logs, as well as the cache key that triggered it. This helps identify if the cache is processing the right config as we intended.\n",
+    "\n",
+    "By default, only rank 0 is enabled for both level 1 and level 2. Users can use the \"level:ranks\" format to specify a select set of ranks to collect diagnostics on. For example, \"2:0,3\" means to enable level 2 diagnostics on ranks 0 and 3. An example of the level 2 diagnostics, which also includes the level 1 summary, is as follows, for a `tests/pytorch/attention/test_attention.py::test_dot_product_attention` test.\n",
+    "```\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd MISS         | train=1 det=0 cg=0 maxlogit=0 mask=0 bias=0 wl=-1 wr=-1 brd=0 softmax=0 scale_mode=0 dropout=0 attn_scale=1 qkv_dt=5 o_dt=5 do_dt=6 dqkv_dt=6 qkv_lay=0 o_fmt=0 do_fmt=8 dqkv_lay=26 qkv_sif=8 do_sif=8 b=8 h=16 hg=16 dqk=64 dv=64 sq=128 skv=128 tq=0 tkv=0 bb=0 btq=0 btkv=0 npk=0 npv=0 psk=0 psv=0 mppk=0 mppv=0 bias_b=0 bias_h=0 bias_sq=0 bias_skv=0\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd CREATE_GRAPH | hit=   0, miss=   1, create_graph=   1, cache_graph=   0, build_plans=   0, execute=   0\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd CACHE_GRAPH  | hit=   0, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   0, execute=   0\n",
+    "...\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd HIT          | train=1 det=0 cg=0 maxlogit=0 mask=0 bias=0 wl=-1 wr=-1 brd=0 softmax=0 scale_mode=0 dropout=0 attn_scale=1 qkv_dt=5 o_dt=5 do_dt=6 dqkv_dt=6 qkv_lay=0 o_fmt=0 do_fmt=8 dqkv_lay=26 qkv_sif=8 do_sif=8 b=8 h=16 hg=16 dqk=64 dv=64 sq=128 skv=128 tq=0 tkv=0 bb=0 btq=0 btkv=0 npk=0 npv=0 psk=0 psv=0 mppk=0 mppv=0 bias_b=0 bias_h=0 bias_sq=0 bias_skv=0\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd BUILD_PLANS  | hit=   3, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   1, execute=   0\n",
+    "...\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd EXECUTE      | hit=   5, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   1, execute=   1\n",
+    "...\n",
+    "[FUSED-ATTN-CACHE] rank=0 | ===== summary begin =====\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 fwd | hit=   5, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   1, execute=   1\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=0   dev=0   | f16 bwd | hit=   1, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   0, execute=   0\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=1   dev=0   | f16 bwd | hit=   4, miss=   0, create_graph=   0, cache_graph=   0, build_plans=   1, execute=   1\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=all dev=all | f16 fwd | hit=   5, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   1, execute=   1\n",
+    "[FUSED-ATTN-CACHE] rank=0 | tid=all dev=all | f16 bwd | hit=   5, miss=   1, create_graph=   1, cache_graph=   1, build_plans=   1, execute=   1\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 fwd validate               | calls=1 | time=    0.168 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 fwd build_operation_graph  | calls=1 | time=   11.405 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 fwd create_execution_plans | calls=1 | time=    2.054 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 fwd check_support          | calls=1 | time=    0.044 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 fwd build_plans            | calls=1 | time=  327.083 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 bwd validate               | calls=1 | time=    0.009 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 bwd build_operation_graph  | calls=1 | time=    1.026 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 bwd create_execution_plans | calls=1 | time=    2.874 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 bwd check_support          | calls=1 | time=    0.016 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | f16 bwd build_plans            | calls=1 | time=  514.599 ms/call\n",
+    "[FUSED-ATTN-CACHE] rank=0 | ===== summary end =====\n",
+    "```"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "611d8fdb",
@@ -346,18 +387,7 @@
     "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n",
     "```\n",
     "\n",
-    "**cuDNN attention sub-backends:**\n",
-    "This environment variable allows users to express their preference of cuDNN attention sub-backends. However, the elected sub-backend will only be used *if* it is eligible, i.e. if it has support for the provided inputs and runtime environment.\n",
-    "```\n",
-    "NVTE_FUSED_ATTN_BACKEND = 1/2 # user preference of cuDNN sub-backend\n",
-    "```\n",
-    "\n",
-    "```\n",
-    "<div class=\"alert alert-info\">\n",
-    "<b>Note</b>\n",
-    "    \n",
-    "Environment variables <code>NVTE_FLASH_ATTN</code>, <code>NVTE_UNFUSED_ATTN</code>, <code>NVTE_FUSED_ATTN_BACKEND</code>, and <code>NVTE_FUSED_ATTN_USE_FAv2_BWD</code> are supported in PyTorch. <code>NVTE_FUSED_ATTN</code> and <code>NVTE_ALLOW_NONDETERMINISTIC_ALGO</code> are supported in both PyTorch and JAX.\n",
-    "</div>\n",
+    "Environment variables <code>NVTE_FLASH_ATTN</code>, <code>NVTE_UNFUSED_ATTN</code>, and <code>NVTE_FUSED_ATTN_USE_FAv2_BWD</code> are supported in PyTorch. <code>NVTE_FUSED_ATTN</code> and <code>NVTE_ALLOW_NONDETERMINISTIC_ALGO</code> are supported in both PyTorch and JAX.\n",
     "\n",
     "### 2.3 Example Tests\n",
     "\n",

docs/examples/jax/attention_context_parallel.py

@@ -245,21 +245,22 @@ def context_parallel_supported() -> Tuple[bool, str]:
         return False, f"needs {cp_size} GPUs"
 
     has_kernel = is_fused_attn_kernel_available(
-        True,
-        dtype,
-        dtype,
-        QKVLayout.THD_THD_THD,
-        AttnBiasType.NO_BIAS,
-        AttnMaskType.PADDING_CAUSAL_MASK,
-        AttnSoftmaxType.VANILLA_SOFTMAX,
-        0.0,
-        num_query_heads,
-        num_kv_heads,
-        seq,
-        seq,
-        head_dim,
-        head_dim,
-        window_size,
+        is_training=True,
+        batch_size=batch,
+        q_dtype=dtype,
+        kv_dtype=dtype,
+        qkv_layout=QKVLayout.THD_THD_THD,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_probability=0.0,
+        q_num_heads=num_query_heads,
+        kv_num_heads=num_kv_heads,
+        q_max_seqlen=seq,
+        kv_max_seqlen=seq,
+        head_dim_qk=head_dim,
+        head_dim_v=head_dim,
+        window_size=window_size,
     )
     if not has_kernel:
         return False, "no fused attention kernel for the THD SWA shape"

docs/examples/jax/test_attention.py

@@ -87,21 +87,22 @@ def _context_parallel_supported():
         return False, f"needs {cp_size} GPUs"
 
     has_kernel = is_fused_attn_kernel_available(
-        True,
-        jnp.bfloat16,
-        jnp.bfloat16,
-        QKVLayout.THD_THD_THD,
-        AttnBiasType.NO_BIAS,
-        AttnMaskType.PADDING_CAUSAL_MASK,
-        AttnSoftmaxType.VANILLA_SOFTMAX,
-        0.0,
-        128,
-        8,
-        65536,
-        65536,
-        128,
-        128,
-        (8192, 0),
+        is_training=True,
+        batch_size=2,
+        q_dtype=jnp.bfloat16,
+        kv_dtype=jnp.bfloat16,
+        qkv_layout=QKVLayout.THD_THD_THD,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_probability=0.0,
+        q_num_heads=128,
+        kv_num_heads=8,
+        q_max_seqlen=65536,
+        kv_max_seqlen=65536,
+        head_dim_qk=128,
+        head_dim_v=128,
+        window_size=(8192, 0),
     )
     if not has_kernel:
         return False, "no fused attention kernel for the THD SWA shape"

tests/jax/test_distributed_fused_attn.py

@@ -82,25 +82,6 @@ def impl_test_self_attn(
         is_training = True
         batch, seqlen, num_head, hidden = data_shape
 
-        if not is_fused_attn_kernel_available(
-            is_training,
-            dtype,
-            dtype,
-            QKVLayout.BS3HD,
-            attn_bias_type,
-            attn_mask_type,
-            softmax_type,
-            dropout_prob,
-            num_head,
-            num_head,
-            seqlen,
-            seqlen,
-            hidden,
-            hidden,
-            None,  # no window
-        ):
-            pytest.skip("No FusedAttn backend found")
-
         col_ref = self.generate_collectives_count_ref(
             mesh_shape,
             mesh_axes,
@@ -234,25 +215,6 @@ def test_cross_attn(
 
         batch, seqlen, num_head, hidden = data_shape
 
-        if not is_fused_attn_kernel_available(
-            is_training,
-            dtype,
-            dtype,
-            QKVLayout.BSHD_BS2HD,
-            attn_bias_type,
-            attn_mask_type,
-            softmax_type,
-            dropout_prob,
-            num_head,
-            num_head,
-            seqlen,
-            seqlen,
-            hidden,
-            hidden,
-            None,  # no window
-        ):
-            pytest.skip("No FusedAttn backend found")
-
         col_ref = self.generate_collectives_count_ref()
         runner = FusedAttnRunner(
             batch,
@@ -491,6 +453,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

@@ -4,7 +4,7 @@
 """Tests for fused attention"""
 import os
 from enum import Enum, auto
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from functools import partial
 from math import sqrt
 from typing import Any, Callable, Mapping, Tuple, Optional, Dict
@@ -48,11 +48,14 @@
 )
 
 from distributed_test_base import assert_equal_collectives
-from utils import assert_allclose, print_debug_tensor_stats
+from utils import assert_allclose, get_test_level, print_debug_tensor_stats
 
 # Get determinism
 _deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")))
 
+# CI test level
+_TEST_LEVEL = get_test_level()
+
 
 @pytest.fixture(autouse=True, scope="module")
 def init():
@@ -383,6 +386,48 @@ def score_mod(_graph, score, _tensors):
         )
 
 
+def test_fused_attn_backend_message():
+    """Test the error messaging of the fused attention backend query."""
+    baseline = FusedAttnHelper(
+        is_training=True,
+        batch_size=2,
+        q_dtype=jnp.bfloat16,
+        kv_dtype=jnp.bfloat16,
+        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.NO_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_probability=0.0,
+        q_num_heads=8,
+        kv_num_heads=8,
+        q_max_seqlen=128,
+        kv_max_seqlen=128,
+        head_dim_qk=64,
+        head_dim_v=64,
+        window_size=(-1, -1),
+        attn_scale=0.125,
+    )
+
+    # One of TE's rules is violated and the error message is surfaced
+    backend, message = replace(
+        baseline, attn_bias_type=AttnBiasType.PRE_SCALE_BIAS
+    ).get_fused_attn_backend()
+    assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend
+    assert message == "Fused attention does not support pre-scale bias."
+
+    # No error message if supported; otherwise skip the test
+    backend, message = baseline.get_fused_attn_backend()
+    if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend:
+        pytest.skip(f"FusedAttention does not support the baseline config: {message}")
+    assert message == ""
+
+    # All TE rules have cleared; now gets rejected by cuDNN's support check
+    # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string
+    backend, message = replace(baseline, head_dim_qk=1024, head_dim_v=1024).get_fused_attn_backend()
+    assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend
+    assert message != ""
+
+
 class BiasShape(Enum):
     """
     Enum class to represent the different bias shapes used in the fused attention.
@@ -482,6 +527,27 @@ def _get_max_segments_per_sequence(self):
             return 1
 
     def _check_configs(self):
+        # Trim SWA configs for L0 and L1 to reduce test time; need to trim more in future test refactoring.
+        if self.window_size is not None and (
+            self.dropout_prob != 0.0 or self.attn_bias_type is not AttnBiasType.NO_BIAS
+        ):
+            if _TEST_LEVEL == "L0" and (
+                self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX
+                or self.dtype != jnp.bfloat16
+                or self.attn_bias_type is not AttnBiasType.POST_SCALE_BIAS
+                or self.attn_mask_type is not AttnMaskType.NO_MASK
+            ):
+                pytest.skip(
+                    "Trimmed SWA+bias/dropout config: only vanilla-softmax + bf16 + post_scale_bias"
+                    " + no-mask runs at L0"
+                )
+            if _TEST_LEVEL == "L1" and (
+                self.dtype != jnp.float16 or self.softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX
+            ):
+                pytest.skip(
+                    "Trimmed SWA+bias/dropout config: only float16 + learnable-softmax runs at L1"
+                )
+
         # TODO(KshitijLakhani): probably add/move this to is_fused_attn_available
         if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding():
             pytest.skip("THD format requires padding masks.")
@@ -584,8 +650,21 @@ def _check_configs(self):
                 "is either BSHD_BSHD_BSHD or THD_THD_THD"
             )
 
-        self.backend = FusedAttnHelper(
+        bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None
+        if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS:
+            if self.bias_shape == BiasShape._1HSS:
+                bias_batch, bias_heads = 1, self.num_heads_q
+            elif self.bias_shape == BiasShape._B1SS:
+                bias_batch, bias_heads = self.batch_size, 1
+            elif self.bias_shape == BiasShape._BHSS:
+                bias_batch, bias_heads = self.batch_size, self.num_heads_q
+            elif self.bias_shape == BiasShape._11SS:
+                bias_batch, bias_heads = 1, 1
+            bias_seqlen_q, bias_seqlen_kv = self.max_seqlen_q, self.max_seqlen_kv
+
+        self.backend, message = FusedAttnHelper(
             self.is_training,
+            self.batch_size,
             self.dtype,
             self.dtype,
             self.qkv_layout,
@@ -600,9 +679,14 @@ def _check_configs(self):
             self.head_dim_qk,
             self.head_dim_v,
             (-1, -1) if self.window_size is None else self.window_size,
+            bottom_right_diagonal=self.attn_mask_type.is_bottom_right(),
+            bias_batch=bias_batch,
+            bias_heads=bias_heads,
+            bias_seqlen_q=bias_seqlen_q,
+            bias_seqlen_kv=bias_seqlen_kv,
         ).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