WIP: add thd support

Signed-off-by: Charlene Yang <[email protected]>

Author: cyanguwa

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu

@@ -324,18 +324,31 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
       }
 
       std::shared_ptr<fe::graph::Tensor_attributes> Max, Sum_Exp;
+      if (is_ragged_q && cudnn_runtime_version >= 90600) {
+        offset_stats =
+            mha_graph->tensor(fe::graph::Tensor_attributes()
+                                  .set_name("offset_stats")
+                                  .set_dim({b + 1, 1, 1, 1})
+                                  .set_stride({1, 1, 1, 1})
+                                  .set_data_type(get_cudnn_fe_dtype(ragged_offset_type)));
+      }
       if (return_max_score) {
         Max = mha_graph->tensor(fe::graph::Tensor_attributes()
                                     .set_name("Max")
                                     .set_dim({b, h, s_q, 1})
-                                    .set_stride({h * s_q, s_q, 1, 1})
                                     .set_data_type(fe::DataType_t::FLOAT));
-        sdpa_options.set_logit_max(Max);
         Sum_Exp = mha_graph->tensor(fe::graph::Tensor_attributes()
                                         .set_name("Sum_Exp")
                                         .set_dim({b, h, s_q, 1})
-                                        .set_stride({h * s_q, s_q, 1, 1})
                                         .set_data_type(fe::DataType_t::FLOAT));
+        if (is_ragged_q && cudnn_runtime_version >= 90600) {
+          Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats);
+          Sum_Exp->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats);
+        } else {
+          Max->set_stride({h * s_q, s_q, 1, 1});
+          Sum_Exp->set_stride({h * s_q, s_q, 1, 1});
+        }
+        sdpa_options.set_logit_max(Max);
         sdpa_options.set_score_sum_exp(Sum_Exp);
       }
 
@@ -357,12 +370,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl(
       if (!return_max_score) {
         Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1});
         if (is_ragged_q && cudnn_runtime_version >= 90600) {
-          offset_stats =
-              mha_graph->tensor(fe::graph::Tensor_attributes()
-                                    .set_name("offset_stats")
-                                    .set_dim({b + 1, 1, 1, 1})
-                                    .set_stride({1, 1, 1, 1})
-                                    .set_data_type(get_cudnn_fe_dtype(ragged_offset_type)));
           Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats);
         } else {
           Stats->set_stride({h * s_q, s_q, 1, 1});

transformer_engine/common/fused_attn/utils.h

@@ -115,23 +115,22 @@ struct FADescriptor_v1 {
   cudnn_frontend::DataType_t o_tensor_type;
   cudnn_frontend::DataType_t do_tensor_type;
   cudnn_frontend::DataType_t dqkv_tensor_type;
-  bool generate_stats;
   bool generate_max_sum_exp;
 
   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,
                     page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h,
                     attnScale, isTraining, dropoutProbability, layout, mask_type, softmax_type,
                     window_size_left, window_size_right, deterministic, bias_type, qkv_tensor_type,
-                    o_tensor_type, do_tensor_type, dqkv_tensor_type, generate_stats,
+                    o_tensor_type, do_tensor_type, dqkv_tensor_type,
                     generate_max_sum_exp) <
            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.attnScale, rhs.isTraining,
                     rhs.dropoutProbability, rhs.layout, rhs.mask_type, rhs.softmax_type,
                     rhs.window_size_left, rhs.window_size_right, rhs.deterministic, rhs.bias_type,
                     rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type,
-                    rhs.dqkv_tensor_type, rhs.generate_stats, rhs.generate_max_sum_exp);
+                    rhs.dqkv_tensor_type, rhs.generate_max_sum_exp);
   }
 };
 

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -1350,7 +1350,7 @@ def forward(
         return out_ret
 
     @staticmethod
-    def backward(ctx, d_out, *args):
+    def backward(ctx, d_out, *_args):
         # pylint: disable=missing-function-docstring
 
         # d_out is expected to be in FP8 if is_output_fp8=True,

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -1845,7 +1845,7 @@ def forward(
         return out_ret
 
     @staticmethod
-    def backward(ctx, dout, *args):
+    def backward(ctx, dout, *_args):
         # pylint: disable=missing-function-docstring
 
         # add NVTX range
@@ -2864,7 +2864,7 @@ def forward(
         return out
 
     @staticmethod
-    def backward(ctx, dout, *args):
+    def backward(ctx, dout, *_args):
         # pylint: disable=missing-function-docstring
         nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVAllGather.backward")
         cp_size = get_distributed_world_size(ctx.cp_group)
@@ -3425,7 +3425,7 @@ def forward(
         return out_ret
 
     @staticmethod
-    def backward(ctx, dout, *args):
+    def backward(ctx, dout, *_args):
         # pylint: disable=missing-function-docstring
         nvtx_range_push("transformer_engine.AttnFuncWithCPAndQKVOA2A.backward")
         cp_size = get_distributed_world_size(ctx.cp_group)

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -486,9 +486,6 @@ def get_attention_backend(
         if use_flash_attention:
             use_flash_attention = False
             logger.debug("Disabling FlashAttention for max_score")
-        if use_fused_attention and qkv_format == "thd":
-            logger.debug("Disabling FusedAttention for max_score and qkv_format = thd")
-            use_fused_attention = False
         if fp8 and fp8_meta["recipe"].fp8_dpa:
             use_flash_attention = False
             use_fused_attention = False

transformer_engine/pytorch/cpp_extensions/fused_attn.py

@@ -323,10 +323,14 @@ def fused_attn_fwd(
     )
 
     if return_max_score:
-        # output_tensors: out [b, sq, h, d] or [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1]
+        qkv_format = qkv_layout.replace("3","").replace("2","").split("_")[0]
+        # thd:  output_tensors: out [tq, h, d],    Max [tq, h, 1],    Sum_Exp [tq, h, 1]
+        # bshd: output_tensors: out [b, sq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1]
+        # sbhd: output_tensors: out [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1]
         stats = output_tensors[1] + torch.log(output_tensors[2])
-        # Max [b, h, sq, 1] -> max_score [h]
-        max_score = torch.amax(output_tensors[1], dim=(0, 2, 3)).to(dtype=output_tensors[0].dtype)
+        amax_dims = (0, 2) if qkv_format == "thd" else (0, 2, 3)
+        # Max -> max_score [h]
+        max_score = torch.amax(output_tensors[1], dim=amax_dims).to(dtype=output_tensors[0].dtype)
         aux_ctx_tensors = [stats]
         aux_ctx_tensors.extend(output_tensors[3:])
         return output_tensors[0], aux_ctx_tensors, max_score

transformer_engine/pytorch/csrc/extensions/attention.cpp

@@ -253,7 +253,7 @@ std::vector<py::object> fused_attn_fwd(
   // f16_max512   : S [b, h, sq, skv]
   // f16_arbitrary:
   // return_max_score=false: S [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1]
-  // return_max_score=true: Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv]
+  // return_max_score=true: Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1]
   // fp8          : M [b, h, sq, 1], ZInv [b, h, sq, 1], rng_state [2]
   size_t i = 0;
   at::Tensor output_tensor;
@@ -262,7 +262,7 @@ std::vector<py::object> fused_attn_fwd(
       allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])),
                     static_cast<DType>(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false);
   set_tensor_param(i++, output_tensor);
-  // fp8 has an additional softmax stats tensor, ZInv
+  // fp8 has an additional softmax stats tensor, ZInv; return_max_score=true has an additional Sum_Exp tensor
   if (return_max_score || qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) {
     output_tensor =
         allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])),