From a40bd4e3ead27c8173ab71f2fa19fb1fd755e2fb Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 14:13:28 -0700 Subject: [PATCH] [Common] Gate FP8 THD sink backward on cuDNN 9.26 cuDNN versions before 9.26 can misindex ragged stats when FP8 THD uses a sink token during backward. Preserve FP8 THD inference and vanilla-softmax training on older supported versions while avoiding the affected sink backward path. Signed-off-by: Sudhakar Singh --- transformer_engine/common/fused_attn/fused_attn.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 544c02aceb..b58282f764 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -294,6 +294,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || ((cudnn_runtime_version >= 92300 && (sm_arch_ >= 100 || (sm_arch_ >= 90 && !is_training))) && qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size && + // cuDNN before 9.26 misindexes ragged FP8 Stats during sink-token backward. + (!is_training || softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX || + cudnn_runtime_version >= 92600) && (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) &&