Merge branch 'main' into nvegesna/gemma2-softcap-core

Three conflicts, all two sides adding different things at the same line.

tests/pytorch/utils.py: #2964 adds bottom_right_diagonal to ModelConfig and to
the AttentionParams call where this branch adds softcap. AttentionParams carries
both fields, so both are kept.

dot_product_attention.py: the attention_params_kwargs dict. #2964 tightens the
alibi_slopes_shape guard to require core_attention_bias_type == "alibi". Kept
that and kept the "softcap" key, which get_attention_backend needs to see in
order to disqualify the backends that cannot honour a cap. Dropping it here
would leave softcap silently ignored during backend selection with no test
failure, which is how it was lost in an earlier merge on this branch.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

Author: nvegesna-netizen

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,23 @@ 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,
+        max_segments_per_seq=max_segments_per_seq,
     )
     if not has_kernel:
         return False, "no fused attention kernel for the THD SWA shape"

docs/examples/jax/test_attention.py

@@ -87,21 +87,23 @@ 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),
+        max_segments_per_seq=4,
     )
     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,
@@ -488,9 +450,17 @@ def impl_test_context_parallel_attn(
             cp_load_balanced=load_balanced,
         )
 
+        # Mirror _FusedAttnCPWithAllGatherHelper.get_adjusted_max_segments_per_seq()
+        runner_segments = runner._get_max_segments_per_sequence()
+        if stripe_size and cp_strategy in (CPStrategy.DEFAULT, CPStrategy.ALL_GATHER):
+            max_segments_per_seq = runner_segments + seqlen // (stripe_size * cp_size)
+        else:
+            max_segments_per_seq = runner_segments
+
         def check_has_backend_for_mask(mask_type):
             return is_fused_attn_kernel_available(
                 is_training,
+                batch,
                 dtype,
                 dtype,
                 qkv_layout,
@@ -504,8 +474,9 @@ def check_has_backend_for_mask(mask_type):
                 seqlen,
                 hidden,
                 hidden,
-                None,
-            )  # no SWA for CP
+                None,  # no SWA for CP
+                max_segments_per_seq=max_segments_per_seq,
+            )
 
         # For causal masking we depend on having bottom right support also.
         # The API does not check this and instead we rely on lower level checks to raise