[ROCm] Fix THD/ragged dQ backward on the bf16 dq_shuffle path (gfx950)

ci/jax.sh

@@ -60,6 +60,10 @@ run_test_config() {
     run 1 test_fused_attn.py
     NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 run_default_fa_lbl "deterministic" 3 test_fused_attn.py -k "TestFusedAttnWithDeterminism"
     NVTE_CK_USES_FWD_V3=0 NVTE_CK_USES_BWD_V3=0 run_default_fa_lbl "v2" 3 test_fused_attn.py # Using FAv2 for forward and backward pass
+    # bf16 atomic dq accumulation (dq_shuffle post-kernel). Default is fp32 (atomic32/dq_convert), so the
+    # bf16 dq_acc path is otherwise never exercised in CI. Scope to THD/RAGGED backward where the
+    # group-mode per-segment dq_acc layout matters (see the equal-dim-128 RAGGED_SELF config).
+    NVTE_CK_IS_V3_ATOMIC_FP32=0 run_default_fa_lbl "atomic16" 3 test_fused_attn.py -k "test_backward and RAGGED"
     run_default_fa 1 test_layer.py # it effectively always uses unfused attention
     run_default_fa 1 test_sanity_import.py
     run_default_fa 1 test_softmax.py

tests/jax/test_fused_attn.py

@@ -1336,6 +1336,49 @@ def check_dqkv(primitive, reference, pad, idx):
             QKVLayout.THD_THD_THD,
             id="2-1024-2048-12-6-128-64-FP16-CROSS-GQA-RAGGED_SEPARATE",
         ),
+        # equal head_dim 128 + bf16 + gqa + self-attn (llama3-like). Previously there was no
+        # equal-dim-128 self-attn coverage; the THD/RAGGED case exercises the group-mode dq_acc
+        # backward (the bf16 dq_shuffle path when NVTE_CK_IS_V3_ATOMIC_FP32=0), which was silently
+        # producing wrong dQ for every packed segment past cu_seqlens offset 0.
+        pytest.param(
+            2,
+            2048,
+            2048,
+            12,
+            6,
+            128,
+            128,
+            jnp.bfloat16,
+            QKVLayout.BSHD_BSHD_BSHD,
+            id="2-2048-2048-12-6-128-128-BF16-GQA-SELF",
+        ),
+        pytest.param(
+            2,
+            2048,
+            2048,
+            12,
+            6,
+            128,
+            128,
+            jnp.bfloat16,
+            QKVLayout.THD_THD_THD,
+            id="2-2048-2048-12-6-128-128-BF16-GQA-RAGGED_SELF",
+        ),
+        # non-16-multiple seqlen: exercises the atomic16 dq_acc seqlen padding (pad16(s_q)) on the bf16
+        # dq_shuffle THD backward path; make sure the CK flow handles a max_seqlen that is not a multiple of
+        # kV3DqAccSeqAlign (2044 % 16 == 12).
+        pytest.param(
+            2,
+            2044,
+            2044,
+            12,
+            6,
+            128,
+            128,
+            jnp.bfloat16,
+            QKVLayout.THD_THD_THD,
+            id="2-2044-2044-12-6-128-128-BF16-GQA-RAGGED_SELF-UNALIGNED",
+        ),
         pytest.param(
             10,
             4096,

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

@@ -14,6 +14,15 @@
 
 namespace ck_fused_attn{
 
+// Fixed head dim of AITER's v3 asm bf16 dq_shuffle backward path. AITER's atomic16 dq_acc scratch is
+// padded to this head dim regardless of the actual d_qk (see AITER csrc/py_itfs_cu/asm_mha_varlen_bwd.cu:
+// "padding dq_accum seqlen to 16x of max_seqlen_q, head dim to 128"), and the only ragged/group dq_shuffle
+// kernel is bwd_hd128_dq_shuffle_group. Both the dq_acc sizing (fused_attn_ck.cpp) and the stride override
+// (ck_fused_attn_bwd.cpp) reference this so the two sites can never drift apart.
+static constexpr int64_t kV3DqAccHeadDim = 128;
+// AITER pads the dq_acc seqlen up to a multiple of this in the atomic16 layout (same source reference).
+static constexpr int64_t kV3DqAccSeqAlign = 16;
+
 // input qkv dtypes
 enum class DType {
   kFloat16    = 0,  /*!< 16-bit float (E5M10) */

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp

@@ -605,6 +605,33 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){
   const bool has_dbias = built.has_dbias;
   const BiasShape bias_shape = built.bias_shape;
 
+  // dq_acc scratch layout (mirrors AITER's asm_mha_varlen_bwd wiring). build_bwd_fmha_args sets the
+  // fp32-packed layout (nsplits, H, total_q, d_qk) with batch_stride_dq_acc=0 -- correct for the fp32
+  // dq_convert post-kernel. But the bf16 dq_shuffle post-kernel (is_v3_atomic_fp32=0) expects a
+  // per-segment padded layout (nsplits, B, H, pad16(s_q), kV3DqAccHeadDim); with batch_stride=0 it
+  // mis-indexes every ragged segment past cu_seqlens offset 0, corrupting dQ. Switch to the per-segment
+  // layout only when the v3 asm path actually runs (ck_attn_bwd_uses_v3 probe) -- v2/fallback keeps the
+  // fp32-packed layout. The dq_acc buffer is sized to hold either layout in fused_attn_ck.cpp.
+  // Layout + padding reference: AITER csrc/py_itfs_cu/asm_mha_varlen_bwd.cu (atomic16 dq_accum =
+  // torch::zeros({1, batch, nhead, pad16(max_seqlen_q), 128})).
+  if (args.is_group_mode() && !args.is_v3_atomic_fp32 && ck_attn_bwd_uses_v3(args)) {
+    // The only ragged/group dq_shuffle kernel is bwd_hd128_dq_shuffle_group, so ck_attn_bwd_uses_v3 can
+    // only be true here for d_qk == kV3DqAccHeadDim. Assert it so a future hd64/hd192 dq_shuffle kernel
+    // (which would need a different dq_acc head dim) fails loudly instead of silently corrupting dQ.
+    if (args.d_qk != kV3DqAccHeadDim) {
+      throw std::runtime_error("ck_fused_attn bwd: bf16 dq_shuffle dq_acc layout assumes d_qk == 128.");
+    }
+    const int64_t padded_sq = ((static_cast<int64_t>(args.s_q) + kV3DqAccSeqAlign - 1)
+                               / kV3DqAccSeqAlign) * kV3DqAccSeqAlign;
+    fmha_args.stride_dq_acc = kV3DqAccHeadDim;
+    fmha_args.nhead_stride_dq_acc = padded_sq * kV3DqAccHeadDim;
+    fmha_args.batch_stride_dq_acc = static_cast<int64_t>(args.h) * padded_sq * kV3DqAccHeadDim;
+    // split_stride is irrelevant in practice: the v3 asm path has no deterministic mode, so nsplits==1
+    // and the split index is always 0. Set it consistently anyway.
+    fmha_args.split_stride_dq_acc = static_cast<int>(static_cast<int64_t>(args.b) * args.h
+                                                     * padded_sq * kV3DqAccHeadDim);
+  }
+
   bool ck_log_config = false;
   if (const char* env_p = std::getenv("CK_FUSED_ATTN_LOG_CONFIG") ) {
     if (env_p != nullptr && std::string(env_p) == "1")

transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp

@@ -749,8 +749,27 @@ void fused_attn_ck_bwd_impl(
   // First h*max_tokens_q*sizeof(float) is the lse-d buffer (passed as softmax_lsed)
   void* lse_workspace = planner.allocate(h*max_tokens_q*sizeof(float));
 
-  // CK requires dq_acc ptr; size depends on deterministic mode
-  void* dq_acc_ptr = planner.allocate(nsplits*h*max_tokens_q*d_qk*sizeof(float));
+  // CK requires a dq_acc scratch buffer. ck_attn_bwd picks the layout at launch time: fp32-packed
+  // (nsplits, H, total_q, d_qk), or -- when the v3 asm dq_shuffle path runs -- bf16 per-segment padded
+  // (nsplits, B, H, pad16(s_q), kV3DqAccHeadDim). See AITER csrc/py_itfs_cu/asm_mha_varlen_bwd.cu for the
+  // atomic16 layout. Whether v3 runs needs a probe (ck_attn_bwd_uses_v3) that depends on args only built in
+  // the execution pass, so we cannot call it here in the workspace-sizing pass; instead we size the buffer
+  // to hold EITHER layout. The bf16 group dq_shuffle kernel is hd128-only (kV3DqAccHeadDim), so it is only
+  // ever reachable when d_qk == kV3DqAccHeadDim -- gate on that (rather than probing) to avoid over-sizing
+  // hd64/hd192, and to keep this conservative but never-under-sized regardless of the launch-time decision.
+  // Element size is 2: the atomic16 dq_acc buffer is bf16 (16-bit), the only input dtype the v3 asm path
+  // admits -- do not tie it to nvte_dtype_size(dtype).
+  const size_t dq_acc_fp32_bytes = nsplits*h*max_tokens_q*d_qk*sizeof(float);
+  const bool dq_acc_bf16_possible = is_ragged && !nvte_ck_is_v3_atomic_fp32
+                                    && d_qk == static_cast<uint64_t>(ck_fused_attn::kV3DqAccHeadDim);
+  const size_t dq_acc_padded_sq =
+      ((s_q + ck_fused_attn::kV3DqAccSeqAlign - 1) / ck_fused_attn::kV3DqAccSeqAlign)
+      * ck_fused_attn::kV3DqAccSeqAlign;
+  const size_t dq_acc_bf16_bytes =
+      nsplits*b*h*dq_acc_padded_sq*ck_fused_attn::kV3DqAccHeadDim*sizeof(uint16_t);  // atomic16 = bf16
+  const size_t dq_acc_bytes = (dq_acc_bf16_possible && dq_acc_bf16_bytes > dq_acc_fp32_bytes)
+      ? dq_acc_bf16_bytes : dq_acc_fp32_bytes;
+  void* dq_acc_ptr = planner.allocate(dq_acc_bytes);
 
   void* dk_expanded_ptr = nullptr;
   void* dv_expanded_ptr = nullptr;
@@ -892,8 +911,9 @@ void fused_attn_ck_bwd_impl(
   }
 
   // Initialize workspace buffers.
-  // dq_acc is of shape (nsplits, B, S, H, D_qk); CK requires zeroing
-  NVTE_CHECK_CUDA(cudaMemsetAsync(dq_acc_ptr, 0, sizeof(float)*nsplits*h*max_tokens_q*d_qk, stream));
+  // dq_acc is sized to hold either layout (see allocation above); ck_attn_bwd picks the strides.
+  // CK requires zeroing the whole buffer.
+  NVTE_CHECK_CUDA(cudaMemsetAsync(dq_acc_ptr, 0, dq_acc_bytes, stream));
   if(devPtrAlibiSlope){
     dim3 block, grid;
     block.x = 1024;