Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wenchenvincent wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int64_t>(((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;
Expand All @@ -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<int64_t>((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<int64_t>(dq_acc_padded_sq * 128)
: static_cast<int64_t>((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;
Expand All @@ -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<int64_t>(args.h * args.s_q * args.d_qk);
fmha_args.batch_stride_dq_acc = dq_acc_bf16_ragged
? static_cast<int64_t>(args.h * dq_acc_padded_sq * 128)
: (args.is_group_mode() ? 0 : static_cast<int64_t>(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<int>(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<int>(args.b * args.h * dq_acc_padded_sq * 128)
: static_cast<int>(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;
Expand Down
21 changes: 17 additions & 4 deletions transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

@ipanfilo ipanfilo Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For safety it should probably also check nvte_ck_uses_bwd_v3

@wangye805 wangye805 Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Emm, even if checking nvte_ck_uses_bwd_v3, I still have one concern: we don't have a way to know whether we fall back to v2 path. Currently v3 asm path coverage is still limited. Once we fall back to v2, there is no bf16 atomic path... @wenchenvincent

@ipanfilo ipanfilo Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So we may need adding extra call with v3_api_check

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;
Expand Down Expand Up @@ -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;
Expand Down