[All] Guard THD learnable dSink on older cuDNN

transformer_engine/common/fused_attn/fused_attn.cpp

@@ -253,6 +253,12 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
                                      max_seqlen_kv, head_dim_qk, head_dim_v) == DType::kInt64);
   const bool supported_ragged_offset_size =
       (!requires_64bit_ragged_offset || cudnn_runtime_version >= 90500);
+  // Before cuDNN 9.26, the generic F16/BF16 sink backward kernel indexes ragged Stats as
+  // dense. The specialized kernels for these value head dimensions are not affected.
+  const bool has_f16_ragged_dsink_bug = is_training && qkv_format == NVTE_QKV_Format::NVTE_THD &&
+                                        softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX &&
+                                        cudnn_runtime_version < 92600 && head_dim_v != 64 &&
+                                        head_dim_v != 128 && head_dim_v != 256;
 
   if ((q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2) &&
       sm_arch_ >= 90 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS &&
@@ -475,6 +481,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
             dropout == 0.0)))) &&
         // check 64-bit ragged offset support
         (supported_ragged_offset_size) &&
+        // check for incorrect dSink or an out-of-bounds Stats read in cuDNN before 9.26
+        (!has_f16_ragged_dsink_bug) &&
         // 9.10.0/9.10.1: known bugs with SDPA F16
         (cudnn_runtime_version != 91000) && (cudnn_runtime_version != 91001) &&
         // softmax type

transformer_engine/jax/flax/transformer.py

@@ -33,6 +33,7 @@
 from ..attention import is_fused_attn_kernel_available, make_swa_mask, canonicalize_attn_mask_type
 from ..attention import fused_attn
 from ..attention import CPStrategy
+from ..cpp_extensions.misc import get_cudnn_version
 from ..softmax import SoftmaxFusionType
 from ..sharding import num_of_devices
 from ..sharding import get_sharding_map_logic_axis_to_mesh_axis
@@ -843,14 +844,30 @@ def __call__(
         use_fused_attn = enable_fused_attn and has_fused_attn_kernel
 
         if enable_fused_attn and not has_fused_attn_kernel:
-            warnings.warn(
-                "Fused attention is not enabled because there is no available kernel.\n"
-                "Fall back to the unfused attention.\n"
-                "Please try to update the cuDNN and TE to the latest version.\n"
-                f"{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n"
-                f"{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n"
-                f"{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\n"
-            )
+            if (
+                input_dtype in (jnp.float16, jnp.bfloat16)
+                and not deterministic
+                and qkv_layout.is_thd()
+                and softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX
+                and get_cudnn_version() < (9, 26, 0)
+                and head_dim_v not in (64, 128, 256)
+            ):
+                warnings.warn(
+                    "Disabling fused attention due to a known cuDNN issue with THD learnable "
+                    f"softmax backward and head_dim_v = {head_dim_v}. Upgrade to cuDNN 9.26 or "
+                    "later to use fused attention for this configuration. Falling back to "
+                    "unfused attention."
+                )
+            else:
+                warnings.warn(
+                    "Fused attention is not enabled because there is no available kernel.\n"
+                    "Fall back to the unfused attention.\n"
+                    "Please try to update the cuDNN and TE to the latest version.\n"
+                    f"{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n"
+                    f"{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n"
+                    f"{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n"
+                    f"{head_dim_v=}\n"
+                )
 
         dropout_rng = None
         if not deterministic and self.attention_dropout > 0.0:

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -1491,7 +1491,22 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             deterministic,
         )
         if fused_attention_backend == FusedAttnBackend.No_Backend.value:
-            logger.debug("Disabling FusedAttention as no backend supports the provided input")
+            if (
+                q_type in (TE_DType[torch.float16], TE_DType[torch.bfloat16])
+                and is_training
+                and qkv_format == "thd"
+                and softmax_type == "learnable"
+                and cudnn_version < (9, 26, 0)
+                and head_dim_v not in (64, 128, 256)
+            ):
+                logger.warning(
+                    "Disabling FusedAttention due to a known cuDNN issue with THD learnable "
+                    "softmax backward and head_dim_v = %s. Upgrade to cuDNN 9.26 or later to "
+                    "use FusedAttention for this configuration.",
+                    head_dim_v,
+                )
+            else:
+                logger.debug("Disabling FusedAttention as no backend supports the provided input")
             use_fused_attention = False
             fused_attention_backend = None
         elif (