diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp index b0ed9ae6e5..6c8b7a3536 100644 --- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp +++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp @@ -529,8 +529,15 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ fmha_args.stride_o = args.stride_s_o; fmha_args.stride_randval = args.s_kv; fmha_args.stride_do = args.stride_s_do; + // dq_acc layout mirrors aiter's asm_mha_varlen_bwd: + // - fp32 dq_convert path: fp32-packed (nsplits, H, total_q, d_qk), batch_stride_dq_acc = 0 in group mode. + // - bf16 dq_shuffle path in ragged/group mode: per-segment padded (nsplits, B, H, pad16(s_q), 128) with + // a fixed nonzero batch stride, so each ragged segment lands in its own slot (the old batch_stride=0 + // packed layout mis-indexed every segment past cu_seqlens offset 0 -> corrupt dQ). + const bool dq_acc_bf16_ragged = args.is_group_mode() && !args.is_v3_atomic_fp32; + const int64_t dq_acc_padded_sq = static_cast(((args.s_q + 15) / 16) * 16); //dq_acc of shape (nsplits, B, H, S, D) - fmha_args.stride_dq_acc = args.d_qk; + fmha_args.stride_dq_acc = dq_acc_bf16_ragged ? 128 : args.d_qk; fmha_args.stride_dq = args.stride_s_dq; fmha_args.stride_dk = is_mqa_gqa? args.stride_s_dk_expanded : args.stride_s_dk; fmha_args.stride_dv = is_mqa_gqa? args.stride_s_dv_expanded : args.stride_s_dv; @@ -548,7 +555,9 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ fmha_args.nhead_stride_randval = args.is_group_mode() ? 0 : args.s_q * args.s_kv; fmha_args.nhead_stride_do = args.stride_h_do; fmha_args.nhead_stride_lsed = args.is_group_mode() ? args.max_tokens_q : args.s_q; - fmha_args.nhead_stride_dq_acc = static_cast((args.is_group_mode() ? args.max_tokens_q : args.s_q) * args.d_qk); + fmha_args.nhead_stride_dq_acc = dq_acc_bf16_ragged + ? static_cast(dq_acc_padded_sq * 128) + : static_cast((args.is_group_mode() ? args.max_tokens_q : args.s_q) * args.d_qk); fmha_args.nhead_stride_dq = args.stride_h_dq; fmha_args.nhead_stride_dk = is_mqa_gqa? args.stride_h_dk_expanded : args.stride_h_dk; fmha_args.nhead_stride_dv = is_mqa_gqa? args.stride_h_dv_expanded : args.stride_h_dv; @@ -564,13 +573,17 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ fmha_args.batch_stride_randval = args.is_group_mode() ? 0 : args.h * args.s_q * args.s_kv; fmha_args.batch_stride_do = args.is_group_mode() ? 0 : args.stride_b_do; fmha_args.batch_stride_lsed = args.is_group_mode() ? 0 : args.h * args.s_q; - fmha_args.batch_stride_dq_acc = args.is_group_mode() ? 0 : static_cast(args.h * args.s_q * args.d_qk); + fmha_args.batch_stride_dq_acc = dq_acc_bf16_ragged + ? static_cast(args.h * dq_acc_padded_sq * 128) + : (args.is_group_mode() ? 0 : static_cast(args.h * args.s_q * args.d_qk)); fmha_args.batch_stride_dq = args.is_group_mode() ? 0 : args.stride_b_dq; fmha_args.batch_stride_dk = args.is_group_mode() ? 0 : (is_mqa_gqa? args.stride_b_dk_expanded : args.stride_b_dk); fmha_args.batch_stride_dv = args.is_group_mode() ? 0 : (is_mqa_gqa? args.stride_b_dv_expanded : args.stride_b_dv); // for dbias, use h since h can be different from bias_h fmha_args.batch_stride_dbias = args.is_group_mode() ? 0 : args.h * args.s_q * args.s_kv; - fmha_args.split_stride_dq_acc = static_cast(args.is_group_mode() ? (args.max_tokens_q * args.h * args.d_qk) : (args.b * args.h * args.s_q * args.d_qk)); + fmha_args.split_stride_dq_acc = dq_acc_bf16_ragged + ? static_cast(args.b * args.h * dq_acc_padded_sq * 128) + : static_cast(args.is_group_mode() ? (args.max_tokens_q * args.h * args.d_qk) : (args.b * args.h * args.s_q * args.d_qk)); fmha_args.window_size_left = args.window_size_left; fmha_args.window_size_right = args.window_size_right; diff --git a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp index 0c7d80a6a6..cc279471ad 100644 --- a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp +++ b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp @@ -749,8 +749,20 @@ 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 dq_acc ptr; size/dtype/layout depend on the dq post-processing path (mirrors aiter's + // asm_mha_varlen_bwd wiring): + // - fp32 dq_convert (is_v3_atomic_fp32=1): fp32-packed (nsplits, H, total_q, d_qk). + // - bf16 dq_shuffle (is_v3_atomic_fp32=0) in ragged/group mode: per-segment padded bf16 + // (nsplits, B, H, pad16(s_q), 128) so each ragged segment has a fixed-stride slot (batch_stride!=0). + // TE previously always used the fp32-packed layout with batch_stride_dq_acc=0, which the dq_shuffle + // kernel mis-indexes for segments past cu_seqlens offset 0 -> corrupt dQ. See ck_fused_attn_bwd.cpp. + const bool dq_acc_bf16_ragged = is_ragged && !nvte_ck_is_v3_atomic_fp32; + const size_t dq_acc_padded_sq = ((s_q + 15) / 16) * 16; + const size_t dq_acc_elems = dq_acc_bf16_ragged + ? (nsplits * b * h * dq_acc_padded_sq * 128) + : (nsplits * h * max_tokens_q * d_qk); + const size_t dq_acc_elem_size = dq_acc_bf16_ragged ? nvte_dtype_size(dtype) : sizeof(float); + void* dq_acc_ptr = planner.allocate(dq_acc_elems * dq_acc_elem_size); void* dk_expanded_ptr = nullptr; void* dv_expanded_ptr = nullptr; @@ -892,8 +904,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 layout/dtype set at allocation (fp32-packed, or bf16 per-segment padded for ragged dq_shuffle); + // CK requires zeroing the whole buffer. + NVTE_CHECK_CUDA(cudaMemsetAsync(dq_acc_ptr, 0, dq_acc_elems * dq_acc_elem_size, stream)); if(devPtrAlibiSlope){ dim3 block, grid; block.x = 1024;