Support dense THD stats graphs on SM8x

transformer_engine/common/fused_attn/config_and_params.cpp

@@ -112,7 +112,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const {
   if (cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) {
     const auto cudnn_runtime_version = cudnnGetVersion();
     const int sm_arch_ = cuda::sm_arch(cuda::current_device());
-    if (cudnn_runtime_version >= 90600 && sm_arch_ != 120) {
+    if (supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_)) {
       if (cache_cfg.is_ragged_q) {
         cache_cfg.max_seqlen_q = cache_cfg.bucketed_num_tokens_q;
       }

transformer_engine/common/fused_attn/config_and_params.h

@@ -19,6 +19,12 @@
 namespace transformer_engine {
 namespace fused_attn {
 
+// Packed THD graph dimensions and ragged Stats/LSE are not supported on SM8x or SM120.
+// Those architectures require dense, BHSD-like graph dimensions for the auxiliary tensors.
+inline constexpr bool supports_packed_ragged_graph(size_t cudnn_runtime_version, int sm_arch) {
+  return cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120;
+}
+
 struct FusedAttnConfig {
   // basic attention settings
   bool is_training = true;

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu

@@ -81,7 +81,9 @@ 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;
+  const bool use_packed_ragged_graph =
+      supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_);
+  const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph;
 
   NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout);
   bool is_paged_kv = cfg.is_paged_kv;
@@ -98,10 +100,9 @@ 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]
-    // 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) {
+    // SM8x and SM120 require BHSD-like dimensions/strides with max_seqlen at plan build.
+    // Other supported architectures use token-count dimensions for graph reuse.
+    if (use_packed_ragged_graph) {
       // 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:
@@ -660,14 +661,16 @@ 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 use_packed_ragged_graph =
+      supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_);
+  const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph;
 
   // keep original batch size because cu_seqlens are created with [b+1] shape
   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) {
+    // SM8x and SM120 require BHSD-like strides with max_seqlen (see fwd).
+    if (use_packed_ragged_graph) {
       // 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 = bucketed_batch_size;
@@ -835,7 +838,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 && use_packed_ragged_graph) {
         sdpa_backward_options.set_max_total_seq_len_kv(s_kv);
       }
 
@@ -1134,12 +1137,10 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i
   const size_t max_seqlen_q = cfg.max_seqlen_q;
   const size_t num_tokens_q = cfg.num_tokens_q;
   const bool return_max_logit = cfg.return_max_logit;
-  const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout;
   const NVTE_Bias_Type bias_type = cfg.bias_type;
   const NVTE_Softmax_Type softmax_type = cfg.softmax_type;
 
   const auto QKV_type = input_Q->data.dtype;
-  NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout);
   void *devPtrQ = input_Q->data.dptr;
   void *devPtrK = input_K->data.dptr;
   void *devPtrV = input_V->data.dptr;
@@ -1171,13 +1172,13 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i
   size_t i = 0;
   if (Aux_CTX_Tensors->size == 0) {
     const auto cudnn_runtime_version = cudnnGetVersion();
+    const bool use_ragged_stats =
+        graph_cfg.is_ragged_q && supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_);
 
     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
-    // [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)) {
+    // Match the packed or dense shape declared by the graph (same as Max below).
+    if (use_ragged_stats) {
       output_S->data.shape = {num_tokens_q, num_attn_heads, 1};
     } else {
       output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1};
@@ -1187,8 +1188,7 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i
     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 (use_ragged_stats) {
         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/pytorch/attention/dot_product_attention/utils.py

@@ -1111,13 +1111,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
                 use_flash_attention_4 = False
 
     # Filter: QKV layout
-    if qkv_format == "thd":
+    if "thd" in (q_format, kv_format):
         if pad_between_seqs:
             if (  # pylint: disable=too-many-boolean-expressions
                 use_flash_attention_2 and FlashAttentionUtils.is_installed
             ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed):
                 logger.debug(
-                    "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is "
+                    "Disabling FlashAttention 2 and 4 when Q or KV uses THD and there is "
                     "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]"
                 )
             use_flash_attention_2 = False
@@ -1130,7 +1130,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             if cudnn_version < (9, 18, 1):
                 if use_fused_attention:
                     logger.debug(
-                        "Disabling FusedAttention as qkv_format = thd is"
+                        "Disabling FusedAttention when Q or KV uses THD because it is"
                         " not supported for compute capability = sm120 and cuDNN version < 9.18.1"
                     )
                 use_fused_attention = False