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

docs/envvars.rst

@@ -190,18 +190,32 @@ 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: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). Output goes to stderr, prefixed ``[FUSED-ATTN-CACHE]``.
+
+      ``1`` emits one line per event that happens once per distinct cache key -- ``CREATE_GRAPH`` when a graph is constructed, ``CACHE_GRAPH`` when cuDNN has agreed to run it, ``BUILD_PLANS`` when its kernels are compiled on first execution -- plus an end-of-run summary block (one row per build site, per thread and in total, plus a row per pass across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries.
+
+      Every line names the thread and device it came from, then the build site behind it -- ``f16`` or ``fp8``, then the pass -- and carries the counters of that site alone, so a process that uses both backends can still tell which of them built what. One line is one pass; the forward and the backward read as adjacent rows. A summary row is an event line without the event name, the block is delimited by ``===== summary begin =====`` and ``===== summary end =====``, and ``tid=all dev=all`` marks the totals. A build site the run never reached is left out rather than shown as a row of zeros.
+
+      ``hit`` and ``miss`` account for every lookup. A miss builds a graph, counted in ``create_graph``, and keeps it only if cuDNN agrees to run it, counted in ``cache_graph``; so the columns fall ``miss`` >= ``create_graph`` >= ``cache_graph``. ``create_graph`` minus ``cache_graph`` is graphs cuDNN refused to run. Nothing is cached for them, so this counts rejected builds rather than rejected configurations: a configuration that is queried again is built and rejected again. A site whose ``cache_graph`` stays put while ``miss`` climbs never runs fused and keeps paying to find that out, which makes these the columns to read when attention is slower than expected and nothing raised an error -- and at level 1, a ``CREATE_GRAPH`` line with no ``CACHE_GRAPH`` after it is one such rejection as it happens. The reason cuDNN gave is not logged here; it reaches the framework as the message explaining why the fused backend was not selected. ``miss`` and ``create_graph`` should agree: configurations FusedAttention itself does not serve are refused before any graph is built, and reported as that same message, so a gap between those two columns means a graph build failed where none was expected to.
+
+      ``2`` additionally emits a per-lookup ``HIT``/``MISS`` line carrying the full cache key, and a per-execution ``EXECUTE`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing.
+
+      Each line is written after the cache lock is released rather than under it, so that no thread waits on stderr while holding the cache. With several threads active this means the lines can appear in a different order than the lookups they report; the counters each line carries still increase in event order, and lines from a single thread are still in that thread's order.
+
+      By default only rank 0 emits, so that output does not scale with the world size. Append ``:<ranks>`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations.
+
+      Has negligible overhead when unset.
+
 .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO
 
    :Type: ``int`` (0 or 1)

docs/examples/attention/attention.ipynb

@@ -253,7 +253,21 @@
     "<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>"
+    "</div>\n",
+    "\n",
+    "Once cuDNN attention has been selected, a separate flag reports on the cuDNN graph cache underneath it. Transformer Engine builds a cuDNN graph per distinct attention configuration and reuses it, so a workload that unexpectedly rebuilds graphs pays for it on every step.\n",
+    "```\n",
+    "NVTE_FUSED_ATTN_CACHE_DEBUG = 0/1/2 # disables/enables graph cache diagnostics\n",
+    "```\n",
+    "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]`, names the thread and device it came from and then the build site behind it -- `f16` or `fp8`, then the pass -- and there is one per event that happens once per configuration: `CREATE_GRAPH` when a graph is constructed, and `BUILD_PLANS` when its kernels are compiled on first execution. Each event name is also the counter column it increments. A configuration cuDNN declines raises a miss and no build, so `miss` minus `create_graph` counts refused builds; nothing is cached for a refusal, so a configuration that is queried again is built and refused again. An end-of-run summary block gives one row per build site, per thread and in total (`tid=all dev=all`), followed by where the build time went:\n",
+    "```\n",
+    "[FUSED-ATTN-CACHE] tid=0   dev=0   | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, ...\n",
+    "[FUSED-ATTN-CACHE] tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, ...\n",
+    "[FUSED-ATTN-CACHE] f16 fwd check_support          | calls=1 | time=   42.135 ms/call\n",
+    "```\n",
+    "The number to read first is `create_graph`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n",
+    "\n",
+    "This flag is supported in both PyTorch and JAX, since the cache it reports on lives in the common C++ layer. By default only rank 0 emits; see [NVTE_FUSED_ATTN_CACHE_DEBUG](../../envvars.rst) for selecting other ranks."
    ]
   },
   {
@@ -346,17 +360,11 @@
     "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",
+    "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",
     "</div>\n",
     "\n",
     "### 2.3 Example Tests\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,
@@ -479,6 +441,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

@@ -53,6 +53,9 @@
 # Get determinism
 _deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")))
 
+# CI test level
+_TEST_LEVEL = os.getenv("NVTE_JAX_UNITTEST_LEVEL", "L0")
+
 
 @pytest.fixture(autouse=True, scope="module")
 def init():
@@ -469,6 +472,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.")
@@ -571,8 +595,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,
@@ -587,9 +624,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,
+            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

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,17 @@ 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 FusedAttnHelper: no fused attention backend available for this configuration",
+            )
 
     def fake_fused_attn(
         qkv,
@@ -454,11 +462,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
 
@@ -533,7 +537,7 @@ def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch):
     assert captured["attn_bias_type"] is AttnBiasType.NO_BIAS
     assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD
     assert captured["softmax_type"] is AttnSoftmaxType.VANILLA_SOFTMAX
-    assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD
+    assert captured["kernel_checks"][0][0][4] is QKVLayout.BSHD_BSHD_BSHD
 
 
 def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch):
@@ -557,7 +561,7 @@ def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monke
     assert captured["qkv"][0].shape == (1, 8, 1, 16)
     assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD
     assert captured["score_mod"] is _identity_score_mod
-    assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD
+    assert captured["kernel_checks"][0][0][4] is QKVLayout.BSHD_BSHD_BSHD
 
 
 def test_multi_head_attention_plumbs_score_mod_to_dot_product_attention(monkeypatch):