Fix cuDNN SDPA score_mod cache collision, sm_12x arch gates, per-device plan cache

tests/jax/test_fused_attn_score_mod.py

@@ -394,6 +394,21 @@ def _identity_score_mod(_graph, score, _tensors):
     return score
 
 
+# Soft-cap / relative-bias score_mods are idiomatically written as module-level lambdas, which
+# all share __qualname__ == "<lambda>"; keep them in a tuple so the cache key must tell them apart.
+_MODULE_LEVEL_LAMBDA_SCORE_MODS = (
+    lambda graph, score, tensors: graph.mul(a=score, b=tensors["scale"]),
+    lambda graph, score, tensors: graph.add(a=score, b=tensors["scale"]),
+)
+
+# Same body, differing only in a default argument. The constant lives in __defaults__, not in
+# the code object, so keying on __code__ alone is not enough to tell these apart.
+_DEFAULT_ARG_LAMBDA_SCORE_MODS = (
+    lambda graph, score, tensors, cap=10.0: graph.mul(a=score, b=cap),
+    lambda graph, score, tensors, cap=50.0: graph.mul(a=score, b=cap),
+)
+
+
 def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True):
     captured = {}
 
@@ -730,6 +745,46 @@ def forward(self, _graph, score, _tensors):
     assert tex_attention._graph_cache_key("fwd", config_1, ()) is None
 
 
+def test_fused_attn_score_mod_config_separates_module_level_lambdas():
+    """Distinct module-level lambdas must not share one compiled cuDNN graph."""
+    first, second = _MODULE_LEVEL_LAMBDA_SCORE_MODS
+    assert first.__qualname__ == second.__qualname__ == "<lambda>"
+
+    key_1 = tex_attention._score_mod_callback_cache_key(first)
+    key_2 = tex_attention._score_mod_callback_cache_key(second)
+    assert key_1 != key_2
+    assert key_1 == tex_attention._score_mod_callback_cache_key(first)
+
+    # Named module-level functions stay cacheable and stable.
+    named_key = tex_attention._score_mod_callback_cache_key(_identity_score_mod)
+    assert not tex_attention._score_mod_key_is_uncacheable(named_key)
+    assert named_key == tex_attention._score_mod_callback_cache_key(_identity_score_mod)
+
+    config_1, _, _ = make_fused_attn_score_mod_config(
+        first, None, {"scale": 1.0}, None, _CONFIG_TEST_SCALING_FACTOR, True
+    )
+    config_2, _, _ = make_fused_attn_score_mod_config(
+        second, None, {"scale": 1.0}, None, _CONFIG_TEST_SCALING_FACTOR, True
+    )
+    assert config_1 != config_2
+    assert tex_attention._graph_cache_key("fwd", config_1, ()) != tex_attention._graph_cache_key(
+        "fwd", config_2, ()
+    )
+
+
+def test_fused_attn_score_mod_config_separates_default_arg_lambdas():
+    """Lambdas differing only in a default argument must not share one compiled graph."""
+    first, second = _DEFAULT_ARG_LAMBDA_SCORE_MODS
+    # Same code object: the differing constant is in __defaults__, not co_consts.
+    assert first.__code__.co_consts == second.__code__.co_consts
+    assert first.__defaults__ != second.__defaults__
+
+    key_1 = tex_attention._score_mod_callback_cache_key(first)
+    key_2 = tex_attention._score_mod_callback_cache_key(second)
+    assert key_1 != key_2
+    assert not tex_attention._score_mod_key_is_uncacheable(key_1)
+
+
 @pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required")
 def test_fused_attn_score_mod_post_scale_bias_optional_bprop():
     """Post-scale-bias score_mod matches the JAX reference without explicit bprop."""

transformer_engine/common/fused_attn/fused_attn.cpp

@@ -502,23 +502,25 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
                    "Please upgrade your cuDNN version if possible."
                 << std::endl;
     }
-    if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && sm_arch_ == 120) {
+    // cudnn-frontend gates SDPA on the compute-capability major version, so sm_121 shares these.
+    const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130);
+    if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && is_sm12x) {
       if (cudnn_runtime_version < 91801) {
         backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend;
-        std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < "
+        std::cout << "Warning: Given combination of sm_arch_ == 12x and cudnn_runtime_version < "
                      "91801 is not supported. "
                   << " Please upgrade your cuDNN version if possible." << std::endl;
       } else if (deterministic && is_training) {
         backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend;
-        std::cout << "Warning: Deterministic fused attention on SM120 is not supported."
+        std::cout << "Warning: Deterministic fused attention on SM12x is not supported."
                   << std::endl;
       } else {
-        // Known missing support for T3HD/TH3D layouts on SM120
+        // Known missing support for T3HD/TH3D layouts on SM12x
         const bool is_t3hd_or_th3d =
             (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD || qkv_layout == NVTE_QKV_Layout::NVTE_TH3D);
         if (is_t3hd_or_th3d) {
           backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend;
-          std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM120 is not supported. "
+          std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM12x is not supported. "
                     << " Please consider using other THD layouts if possible." << std::endl;
         }
       }

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu

@@ -86,7 +86,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
   const auto cudnn_runtime_version = cudnnGetVersion();
   const int device_id = cuda::current_device();
   const int sm_arch_ = cuda::sm_arch(device_id);
-  bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120;
+  // cuDNN gates these behaviours on the compute capability family, so sm_121 etc. behave
+  // like sm_120 and must take the same workarounds.
+  const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130);
+  bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && !is_sm12x;
 
   NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout);
   bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD);
@@ -112,10 +115,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
   int64_t actual_b = b;
   if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) {
     NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!");
-    // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3]
+    // On SM 12x, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3]
     // as interleaved and rejects them. Use BHSD-like dimensions/strides with max_seqlen at plan build
     // so the check passes; ragged offset still provides variable-length boundaries.
-    if (sm_arch_ != 120) {
+    if (!is_sm12x) {
       // replace batch size and maximum sequence lengths with maximum token counts
       // for query and key/value so the graph is static within each quantization bucket.
       // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size:
@@ -179,6 +182,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
         cudnn_frontend::DataType_t::NOT_SET,
         cudnn_frontend::DataType_t::NOT_SET,
         return_max_logit,
+        device_id,
     };
 
     namespace fe = cudnn_frontend;
@@ -678,7 +682,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl(
   const auto cudnn_runtime_version = cudnnGetVersion();
   const int device_id = cuda::current_device();
   const int sm_arch_ = cuda::sm_arch(device_id);
-  bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120;
+  const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130);
+  bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && !is_sm12x;
 
   NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout);
   bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD);
@@ -690,8 +695,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl(
   int64_t actual_b = b;
   if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) {
     NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!");
-    // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd).
-    if (sm_arch_ != 120) {
+    // On SM 12x, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd).
+    if (!is_sm12x) {
       // replace batch size and maximum sequence lengths with maximum token counts
       // for query and key/value so the graph is static within each quantization bucket
       b = max_b;
@@ -743,6 +748,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl(
         cudnn_frontend::DataType_t::NOT_SET,
         cudnn_frontend::DataType_t::NOT_SET,
         false,
+        device_id,
     };
 
     namespace fe = cudnn_frontend;
@@ -893,7 +899,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl(
       if (use_ragged_stats) {
         sdpa_backward_options.set_max_total_seq_len_q(s_q);
       }
-      if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) {
+      if (is_ragged_kv && cudnn_runtime_version >= 90600 && !is_sm12x) {
         sdpa_backward_options.set_max_total_seq_len_kv(s_kv);
       }
 
@@ -1204,6 +1210,7 @@ void fused_attn_arbitrary_seqlen_fwd(
 
   const int device_id = cuda::current_device();
   const int sm_arch_ = cuda::sm_arch(device_id);
+  const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130);
 
   void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr;
   void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr;
@@ -1231,10 +1238,9 @@ void fused_attn_arbitrary_seqlen_fwd(
 
     Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]);
     output_S->data.dptr = nullptr;
-    // sm120 does not use ragged stats: the graph declares a dense
+    // sm12x does not use ragged stats: the graph declares a dense
     // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below).
-    if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) &&
-        (sm_arch_ != 120)) {
+    if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && !is_sm12x) {
       output_S->data.shape = {num_tokens_q, num_attn_heads, 1};
     } else {
       output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1};
@@ -1244,8 +1250,7 @@ void fused_attn_arbitrary_seqlen_fwd(
     if (return_max_logit) {
       Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]);
       output_Max->data.dptr = nullptr;
-      if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) &&
-          (sm_arch_ != 120)) {
+      if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && !is_sm12x) {
         output_Max->data.shape = {num_tokens_q, num_attn_heads, 1};
       } else {
         output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1};

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -6,6 +6,7 @@
 
 #include "../common.h"
 #include "../cudnn_utils.h"
+#include "../util/cuda_runtime.h"
 #include "../util/system.h"
 #include "fused_attn_fp8.h"
 #include "utils.h"
@@ -116,7 +117,8 @@ void fused_attn_fp8_fwd_impl(
                                o_tensor_type,
                                cudnn_frontend::DataType_t::NOT_SET,
                                cudnn_frontend::DataType_t::NOT_SET,
-                               false};
+                               false,
+                               cuda::current_device()};
 
     namespace fe = cudnn_frontend;
     using graph_and_tensors =
@@ -584,7 +586,8 @@ void fused_attn_fp8_bwd_impl(
                                o_tensor_type,
                                do_tensor_type,
                                dqkv_tensor_type,
-                               false};
+                               false,
+                               cuda::current_device()};
 
     namespace fe = cudnn_frontend;
     using graph_and_tensors =

transformer_engine/common/fused_attn/utils.h

@@ -288,6 +288,9 @@ struct FADescriptor_v1 {
   cudnn_frontend::DataType_t do_tensor_type;
   cudnn_frontend::DataType_t dqkv_tensor_type;
   bool return_max_logit;
+  // The plan caches are keyed by this descriptor but the cuDNN handle is per-device, so a
+  // plan built on one device must never be replayed on another.
+  int device_id;
 
   bool operator<(const FADescriptor_v1 &rhs) const {
     return std::tie(b, h, hg, s_q, s_kv, d_qk, d_v, num_pages_k, num_pages_v, page_size_k,
@@ -296,7 +299,7 @@ struct FADescriptor_v1 {
                     do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, mask_type,
                     softmax_type, window_size_left, window_size_right, bottom_right_diagonal,
                     deterministic, bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type,
-                    dqkv_tensor_type, return_max_logit) <
+                    dqkv_tensor_type, return_max_logit, device_id) <
            std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k,
                     rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k,
                     rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.bias_sq, rhs.bias_skv,
@@ -305,7 +308,7 @@ struct FADescriptor_v1 {
                     rhs.do_scale_inv_format, rhs.mask_type, rhs.softmax_type, rhs.window_size_left,
                     rhs.window_size_right, rhs.bottom_right_diagonal, rhs.deterministic,
                     rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type,
-                    rhs.dqkv_tensor_type, rhs.return_max_logit);
+                    rhs.dqkv_tensor_type, rhs.return_max_logit, rhs.device_id);
   }
 };
 

transformer_engine/jax/cpp_extensions/flex_attention.py

@@ -202,7 +202,27 @@ def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any:
         and callback.__closure__ is None
         and "<locals>" not in callback.__qualname__
     ):
-        return ("function", callback.__module__, callback.__qualname__)
+        # __qualname__ is "<lambda>" for every module-level lambda, so it does not identify
+        # the body; include the code object so distinct score_mods get distinct graphs.
+        # Defaults live outside the code object and must be keyed separately: baking a
+        # scalar in via a default argument is the usual way to write a soft-cap score_mod.
+        code = callback.__code__
+        kwdefaults = callback.__kwdefaults__
+        key = (
+            "function",
+            callback.__module__,
+            callback.__qualname__,
+            code.co_code,
+            code.co_consts,
+            code.co_names,
+            callback.__defaults__,
+            tuple(sorted(kwdefaults.items())) if kwdefaults else None,
+        )
+        try:
+            hash(key)
+        except TypeError:
+            return _UncacheableScoreModKey()
+        return key
 
     return _UncacheableScoreModKey()
 

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -1616,11 +1616,11 @@ def forward(
         softmax_lse_in_packed_format = False
         if qkv_format == "thd":
             if use_fused_attention:
-                softmax_lse_in_packed_format = get_cudnn_version() >= (
-                    9,
-                    6,
-                    0,
-                ) and get_device_compute_capability() != (12, 0)
+                # sm12x (sm_120, sm_121, ...) keeps the dense BHS1 stats layout; cuDNN
+                # classifies by compute-capability family, so match on the major version.
+                softmax_lse_in_packed_format = (
+                    get_cudnn_version() >= (9, 6, 0) and get_device_compute_capability()[0] != 12
+                )
             else:
                 softmax_lse_in_packed_format = fa_utils.v2_6_0_plus or use_flash_attn_3