[ROCm] Jax Add softmax sink (learnable off-by-one) support for the ROCm/CK fused attention backend

tests/jax/test_distributed_fused_attn.py

@@ -52,7 +52,7 @@
 class TestDistributedSelfAttn:
 
     def generate_collectives_count_ref(
-        self, mesh_shape, mesh_axes, mesh_resource, with_bias, shape, dtype
+        self, mesh_shape, mesh_axes, mesh_resource, with_bias, shape, dtype, softmax_type
     ):
         jax_dtype = jax.dtypes.canonicalize_dtype(dtype)
         _, seqlen, heads, _ = shape
@@ -64,8 +64,13 @@ def generate_collectives_count_ref(
 
         all_reduce_loss_bytes = 4  # 1 * FP32
         bias_bytes = int(with_bias) * (heads // tp_size) * seqlen * seqlen * jax_dtype.itemsize
-        allreduce_total_bytes = all_reduce_loss_bytes + (bias_bytes * is_dp_enabled)
-        # for loss and dbias
+        # dsoftmax_offset is [1, heads, 1, 1] and always FP32, regardless of the QKV dtype
+        with_softmax_offset = softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX
+        softmax_offset_bytes = int(with_softmax_offset) * (heads // tp_size) * 4
+        allreduce_total_bytes = all_reduce_loss_bytes + (
+            (bias_bytes + softmax_offset_bytes) * is_dp_enabled
+        )
+        # for loss, dbias and dsoftmax_offset
         return generate_collectives_count(allreduce=allreduce_total_bytes, allgather=0, other=0)
 
     def impl_test_self_attn(
@@ -111,6 +116,7 @@ def impl_test_self_attn(
             attn_bias_type != AttnBiasType.NO_BIAS,
             data_shape,
             dtype,
+            softmax_type,
         )
         runner = FusedAttnRunner(
             batch,
@@ -200,10 +206,23 @@ def test_self_attn(
 
 class TestDistributedCrossAttn:
 
-    def generate_collectives_count_ref(self):
-        # for loss
+    def generate_collectives_count_ref(
+        self, mesh_shape, mesh_axes, mesh_resource, shape, softmax_type
+    ):
+        _, _, heads, _ = shape
+        is_dp_enabled = mesh_resource.dp_resource is not None
+        tp_size = 1
+        if mesh_resource.tpsp_resource is not None:
+            idx = mesh_axes.index(mesh_resource.tpsp_resource)
+            tp_size = mesh_shape[idx]
+
         all_reduce_loss_bytes = 4  # 1 * FP32
-        return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0)
+        # dsoftmax_offset is [1, heads, 1, 1] and always FP32, regardless of the QKV dtype
+        with_softmax_offset = softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX
+        softmax_offset_bytes = int(with_softmax_offset) * (heads // tp_size) * 4
+        allreduce_total_bytes = all_reduce_loss_bytes + (softmax_offset_bytes * is_dp_enabled)
+        # for loss and dsoftmax_offset
+        return generate_collectives_count(allreduce=allreduce_total_bytes, allgather=0, other=0)
 
     @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
     @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_CROSS_ATTN_DATA_SHAPES)
@@ -256,7 +275,13 @@ def test_cross_attn(
         ):
             pytest.skip("No FusedAttn backend found")
 
-        col_ref = self.generate_collectives_count_ref()
+        col_ref = self.generate_collectives_count_ref(
+            mesh_shape,
+            mesh_axes,
+            mesh_resource,
+            data_shape,
+            softmax_type,
+        )
         runner = FusedAttnRunner(
             batch,
             seqlen,

tests/jax/test_fused_attn.py

@@ -1142,18 +1142,26 @@ def grad_func(
         }
         reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference}
 
+        arg_nums = (0, 1, 2)
+        grad_shardings = (self.qkvo_sharding, self.qkvo_sharding, self.qkvo_sharding)
+
+        optional_dgrad_idx = 3
+
         # We can compute dBias only for the [1, h, s, s] layout
-        if self.bias_shape == BiasShape._1HSS:
-            arg_nums = (0, 1, 2, 3)
-            grad_shardings = (
-                self.qkvo_sharding,
-                self.qkvo_sharding,
-                self.qkvo_sharding,
-                self.bias_sharding,
-            )
-        else:
-            arg_nums = (0, 1, 2)
-            grad_shardings = (self.qkvo_sharding, self.qkvo_sharding, self.qkvo_sharding)
+        compute_dbias = self.bias_shape == BiasShape._1HSS
+        if compute_dbias:
+            arg_nums += (3,)
+            grad_shardings += (self.bias_sharding,)
+            dgrad_idx_dbias = optional_dgrad_idx
+            optional_dgrad_idx += 1
+
+        # dsoftmax_offset is only meaningful for the learnable softmax variant
+        compute_dsoftmax_offset = self.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX
+        if compute_dsoftmax_offset:
+            arg_nums += (4,)
+            grad_shardings += (self.softmax_offset_sharding,)
+            dgrad_idx_dsoftmax_offset = optional_dgrad_idx
+            optional_dgrad_idx += 1
 
         # Use FP16/BF16 to sum the results may cause overflow, use FP32 for the summation
         jitted_primitive = jit(
@@ -1264,11 +1272,11 @@ def check_dqkv(primitive, reference, pad, idx):
         check_dqkv(primitive_dk, reference_dk, self.pad_kv, 1)
         check_dqkv(primitive_dv, reference_dv, self.pad_kv, 2)
 
-        if self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape == BiasShape._1HSS:
+        if self.attn_bias_type != AttnBiasType.NO_BIAS and compute_dbias:
             # TODO(mgoldfarb-nvidia): Inverse reorder bias once supported by a CP implementation.
 
-            primitive_dbias = primitive_dgrad[3]
-            reference_dbias = reference_dgrad[3]
+            primitive_dbias = primitive_dgrad[dgrad_idx_dbias]
+            reference_dbias = reference_dgrad[dgrad_idx_dbias]
 
             # Assume all batch has the same actual_seqlen, probably needs to extend the tests
             bias_mask = self.mask[0, 0]
@@ -1300,6 +1308,33 @@ def check_dqkv(primitive, reference, pad, idx):
                 dtype=self.dtype,
             )
 
+        if compute_dsoftmax_offset:
+            primitive_dsoftmax_offset = primitive_dgrad[dgrad_idx_dsoftmax_offset]
+            reference_dsoftmax_offset = reference_dgrad[dgrad_idx_dsoftmax_offset]
+
+            print_debug_tensor_stats("primitive_dsoftmax_offset", primitive_dsoftmax_offset)
+            print_debug_tensor_stats("reference_dsoftmax_offset", reference_dsoftmax_offset)
+            print_debug_tensor_stats(
+                "diff_dsoftmax_offset",
+                jnp.abs(primitive_dsoftmax_offset - reference_dsoftmax_offset),
+            )
+
+            if is_hip_extension():
+                assert not jnp.any(
+                    jnp.isnan(primitive_dsoftmax_offset)
+                ), "Fused dsoftmax_offset contains NaN"
+                assert not jnp.any(
+                    jnp.isinf(primitive_dsoftmax_offset)
+                ), "Fused dsoftmax_offset contains Inf"
+
+            # softmax_offset is always fp32, but its gradient is only as accurate as the
+            # attention math that produced it, so tolerance follows the compute dtype.
+            assert_allclose(
+                primitive_dsoftmax_offset,
+                reference_dsoftmax_offset,
+                dtype=self.dtype,
+            )
+
         if self.coll_count_ref is not None:
             with jax.set_mesh(self.mesh), autocast(mesh_resource=self.mesh_resource):
                 target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text()

transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp

@@ -79,6 +79,10 @@ struct CKAttnCommonArgs {
   void* philox_seed_ptr = nullptr;
   void* philox_offset_ptr = nullptr;
 
+  // Softmax sink (learnable / off-by-one)
+  const void* sink_ptr = nullptr;
+  bool has_sink = false;
+
   // O layout (o_ptr lives in derived because fwd writes it / bwd reads it)
   uint64_t stride_b_o = 0, stride_h_o = 0, stride_s_o = 0;
 
@@ -134,6 +138,9 @@ struct CkAttnBwdArgs : CKAttnCommonArgs {
   void* dbias_expanded_ptr = nullptr;
   void* dbias_ptr = nullptr;
 
+  // Softmax sink gradient
+  void* d_sink_ptr = nullptr;
+
   // Workspace shared with forward LSE
   void* lse_workspace_ptr = nullptr;
 

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

@@ -368,6 +368,8 @@ void log_bwd_config(const char* func_name, const aiter::mha_bwd_args& fmha_args,
   log_value(log_file, "dk_ptr", fmha_args.dk_ptr);
   log_value(log_file, "dv_ptr", fmha_args.dv_ptr);
   log_value(log_file, "dbias_ptr", fmha_args.dbias_ptr);
+  log_value(log_file, "sink_ptr", fmha_args.sink_ptr);
+  log_value(log_file, "d_sink_ptr", fmha_args.d_sink_ptr);
 
   log_value(log_file, "seqstart_q_ptr", fmha_args.seqstart_q_ptr);
   log_value(log_file, "seqstart_k_ptr", fmha_args.seqstart_k_ptr);
@@ -529,8 +531,8 @@ BwdFmhaArgs build_bwd_fmha_args(const CkAttnBwdArgs& args){
   }
 
   aiter::mha_bwd_args fmha_args{};
-  fmha_args.sink_ptr = nullptr;
-  fmha_args.d_sink_ptr = nullptr;
+  fmha_args.sink_ptr = args.sink_ptr;
+  fmha_args.d_sink_ptr = args.d_sink_ptr;
   fmha_args.mask_type = static_cast<int>(static_cast<mask_enum>(args.attn_mask_type));
   // Mirrors AITER's small-seqlen guard at aiter/ops/mha.py:1689.
   fmha_args.use_asm_v3 = (args.s_q < 16) ? false : args.uses_bwd_v3;

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

@@ -90,6 +90,8 @@ void log_fwd_config(const char* func_name, bool has_dropout, const aiter::mha_fw
 
   log_value(log_file, "dropout_seed_ptr", std::get<0>(std::get<std::pair<const void*, const void*>>(fmha_args.drop_seed_offset)));
   log_value(log_file, "dropout_offset_ptr", std::get<1>(std::get<std::pair<const void*, const void*>>(fmha_args.drop_seed_offset)));
+  log_value(log_file, "has_sink", fmha_args.has_sink);
+  log_value(log_file, "sink_ptr", fmha_args.sink_ptr);
 }
 
 void dump_fwd_timings(const char* dump_path, float average_runtime){
@@ -153,7 +155,7 @@ aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){
 
   fmha_args.block_scale_seqstart_q_ptr = nullptr;
   fmha_args.block_scale_seqstart_k_ptr = nullptr;
-  fmha_args.sink_ptr = nullptr;
+  fmha_args.sink_ptr = args.sink_ptr;
   fmha_args.seqlen_k     = args.s_kv; // unused in group mode (or kvcache enabled)
   fmha_args.max_seqlen_q = args.s_q;
 
@@ -200,10 +202,15 @@ aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){
   fmha_args.bias_type       = static_cast<int>(bias_type);
   fmha_args.has_lse         = args.lse_ptr!=nullptr;
   fmha_args.qscale_type     = static_cast<int>(quant_scale_enum::no_scale);
-  fmha_args.has_sink        = false;
+  fmha_args.has_sink        = args.has_sink;
   fmha_args.q_descale_ptr    = nullptr;
   fmha_args.k_descale_ptr    = nullptr;
   fmha_args.v_descale_ptr    = nullptr;
+  // sink_size is CK's StreamingLLM sink *prefix width* in key columns, which is a
+  // different feature from the learnable softmax offset NVTE asks for here. The
+  // offset only needs has_sink + sink_ptr (CK folds it into the softmax
+  // denominator).
+  // (aiter's mha_bwd_args has no sink_size at all). Keep it at 0.
   fmha_args.sink_size        = 0;
   fmha_args.min_seqlen_q     = 0;
   fmha_args.block_scale_size_q  = 0;

transformer_engine/common/fused_attn_rocm/fused_attn.cpp

@@ -309,7 +309,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
         qkv_layout,
         bias_type,
         attn_mask_type,
-        softmax_type,
         dropout,
         num_attn_heads, num_gqa_groups,
         max_seqlen_q, max_seqlen_kv,
@@ -370,6 +369,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso
   const Tensor *input_K = convertNVTETensorCheck(K);
   const Tensor *input_V = convertNVTETensorCheck(V);
   const Tensor *input_Bias = convertNVTETensorCheck(Bias);
+  const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(SoftmaxOffset);
   Tensor *output_O = convertNVTETensorCheck(O);
   Tensor *wkspace = convertNVTETensorCheck(workspace);
 
@@ -401,9 +401,9 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso
     fused_attn_ck_fwd(
       b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v,
       is_training, attn_scale, dropout, 
-      qkv_layout, bias_type, attn_mask_type,
+      qkv_layout, bias_type, attn_mask_type, softmax_type,
       window_size_left, window_size_right,
-      input_Q, input_K, input_V, input_Bias, 
+      input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset,
       output_O, Aux_CTX_Tensors,
       input_cu_seqlens_q,
       input_cu_seqlens_kv,
@@ -474,6 +474,8 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso
   const Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); //softmax lse
   const Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[1]);
   Tensor *input_Bias = nullptr;
+  Tensor *input_SoftmaxOffset = nullptr;
+  Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(dSoftmaxOffset);
 
   auto ndim = input_Q->data.shape.size();
   size_t b = input_cu_seqlens_q->data.shape[0] - 1;
@@ -500,18 +502,26 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso
       false, deterministic);
 
   if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_CK) {
+    size_t ctx_next_id = 2;
     if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) {
-      input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[2]);
+      input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[ctx_next_id++]);
+    }
+    if (softmax_type != NVTE_VANILLA_SOFTMAX) {
+      input_SoftmaxOffset =
+          convertNVTETensorCheck(Aux_CTX_Tensors->tensors[ctx_next_id++]);
     }
     fused_attn_ck_bwd(
       b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v,
       attn_scale, dropout, 
       qkv_layout, bias_type, attn_mask_type,
+      softmax_type,
       window_size_left, window_size_right,
       deterministic,
       input_Q, input_K, input_V, input_O, input_dO, input_Bias, 
+      input_SoftmaxOffset,
       output_S,
       output_dQ, output_dK, output_dV, output_dBias,
+      output_dSoftmaxOffset,
       input_cu_seqlens_q,
       input_cu_seqlens_kv,
       input_cu_seqlens_q_padded,