[ROCm] Fix THD/ragged dQ backward on the bf16 dq_shuffle path (gfx950)#657
Open
wenchenvincent wants to merge 1 commit into
Open
[ROCm] Fix THD/ragged dQ backward on the bf16 dq_shuffle path (gfx950)#657wenchenvincent wants to merge 1 commit into
wenchenvincent wants to merge 1 commit into
Conversation
The CK backward's dq_acc scratch was always allocated fp32-packed (nsplits, H, total_q, d_qk) with batch_stride_dq_acc=0 in group mode. That layout only matches the fp32 dq_convert post-kernel (NVTE_CK_IS_V3_ATOMIC_FP32=1). On the default bf16 dq_shuffle path (atomic_fp32=0, gfx950) the post-kernel expects a per-segment padded layout, so with batch_stride=0 every ragged segment past cu_seqlens offset 0 is written to the offset-0 slot -- dQ is correct only for the first segment and garbage for the rest, silently corrupting gradients for packed/THD training (forward, dK, dV are unaffected). Mirror aiter's asm_mha_varlen_bwd wiring: in ragged/group mode with atomic_fp32=0, allocate dq_acc as bf16 (nsplits, B, H, pad16(s_q), 128) and set stride_dq_acc=128, nhead_stride=pad16(s_q)*128, batch_stride=H*pad16(s_q)*128, split_stride=B*H*pad16(s_q)*128, so each segment lands in its own fixed-stride slot. The fp32 dq_convert path and non-ragged BSHD/SBHD are unchanged. Verified on MI355X (gfx950): te thd padding_causal backward dQ cos 0.0036 -> 1.0000 with atomic_fp32=0, all segments correct across [256], [100,80,76], [128,128], [64x4], [200,56]; still uses the fast bwd_hd128_dq_shuffle_group kernel; atomic_fp32=1 unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ipanfilo
reviewed
Jul 1, 2026
| // (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; |
Collaborator
There was a problem hiding this comment.
For safety it should probably also check nvte_ck_uses_bwd_v3
Collaborator
There was a problem hiding this comment.
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
Collaborator
There was a problem hiding this comment.
So we may need adding extra call with v3_api_check
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On ROCm/gfx950 the CK fused-attention backward silently corrupts dQ for THD/ragged (packed / document-masked) attention when using the default bf16 accumulation path (
NVTE_CK_IS_V3_ATOMIC_FP32=0). Forward, dK and dV are correct.Root cause
fused_attn_ck_bwd_implalways allocates thedq_accscratch as fp32-packed(nsplits, H, total_q, d_qk)andck_fused_attn_bwdsetsbatch_stride_dq_acc = 0in group mode. That layout only matches the fp32dq_convertpost-kernel (atomic_fp32=1). On the bf16dq_shufflepost-kernel (atomic_fp32=0), the kernel expects a per-segment padded layout — withbatch_stride=0every ragged segment pastcu_seqlensoffset 0 is written to the offset-0 slot, so dQ is correct only for the first segment and garbage for the rest.Fix
Mirror aiter's
asm_mha_varlen_bwdwiring. In ragged/group mode withatomic_fp32=0, allocatedq_accas bf16(nsplits, B, H, pad16(s_q), 128)and setstride_dq_acc=128,nhead_stride=pad16(s_q)*128,batch_stride=H*pad16(s_q)*128,split_stride=B*H*pad16(s_q)*128. The fp32dq_convertpath and non-ragged BSHD/SBHD are unchanged.Verification (MI355X / gfx950)
TE THD
padding_causalbackward vs a pure-JAX block-diagonal reference,NVTE_CK_IS_V3_ATOMIC_FP32=0:+1.0000across[256],[100,80,76],[128,128],[64×4],[200,56].bwd_hd128_dq_shuffle_groupkernel (fix is on the fast path, not a fallback).atomic_fp32=1path unchanged (no regression).Note:
fused_attn_ck_hip.cppis generated by hipify fromfused_attn_ck.cpp, so only the source is changed here.🤖 Generated with Claude Code