From 4fc35a51e47cdf5e5ed651c3bbe81f61c037c6ab Mon Sep 17 00:00:00 2001 From: Kshitij Janardan Lakhani Date: Thu, 3 Sep 2026 00:07:54 -0700 Subject: [PATCH 1/2] Guard THD learnable dSink on older cuDNN Signed-off-by: Kshitij Janardan Lakhani --- .../common/fused_attn/fused_attn.cpp | 8 ++++++++ .../attention/dot_product_attention/utils.py | 17 ++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 0e773fdc84..11794330b5 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 33d612b4f4..fe8c81edc6 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/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 ( From d930129e6663beaf3792dde863761a379cfed15a Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani Date: Fri, 4 Sep 2026 14:57:43 -0700 Subject: [PATCH 2/2] Log a warning on the jax framework side for dsink cudnn bug Signed-off-by: Kshitij Lakhani --- transformer_engine/jax/flax/transformer.py | 33 ++++++++++++++++------ 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 4b497826cc..9b74458711 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/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: