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

[wenchenvincent]

Summary

On ROCm/gfx950 the CK fused-attention backward silently corrupts dQ for THD/ragged (packed / document-masked) attention on the default bf16 dq-accumulation path (NVTE_CK_IS_V3_ATOMIC_FP32=0). Forward, dK and dV are correct.

Root cause

build_bwd_fmha_args allocates the dq_acc scratch as fp32-packed (nsplits, H, total_q, d_qk) with batch_stride_dq_acc = 0. That matches the fp32 dq_convert post-kernel, but the bf16 dq_shuffle post-kernel expects a per-segment padded layout — with batch_stride=0 every ragged segment past cu_seqlens offset 0 is written to the offset-0 slot, so dQ is correct only for the first segment and garbage for the rest.

Fix

When the v3 asm path actually runs — probed via ck_attn_bwd_uses_v3 (added in #662) — switch dq_acc to AITER's per-segment bf16 layout (nsplits, B, H, pad16(s_q), 128) with matching stride/nhead/batch/split strides. The v2/fallback path and non-ragged/atomic_fp32=1 keep the fp32-packed layout. fused_attn_ck.cpp sizes the scratch to hold either layout (the layout choice needs a GPU probe that must not run in the workspace-sizing pass).

Tests

• Add an equal-dim-128 bf16 GQA self-attn config (BSHD + THD/RAGGED) to test_fused_attn.py. Previously there was no symmetric 128/128 THD backward coverage (only asymmetric 128/64) — exactly the llama3 shape.
• Add an NVTE_CK_IS_V3_ATOMIC_FP32=0 (atomic16 / dq_shuffle) CI run of the THD backward tests in ci/jax.sh. The default fp32 path never exercised the bf16 dq_acc path.

Verification (MI355X / gfx950)

TE THD padding_causal backward vs a pure-JAX block-diagonal reference:

• dQ cosine 0.0036 → 1.0000 with atomic_fp32=0; per-segment all correct across [256], [100,80,76], [128,128], [64×4], [200,56], for both the v3 (dq_shuffle) and v2 paths.
• Still uses the fast bwd_hd128_dq_shuffle_group kernel (fix is on the fast path, not a fallback).
• atomic_fp32=1 unchanged; end-to-end MaxText llama3.1-8B block-diagonal THD packing converges to eval-loss 3.3 (previously stalled ~6.1 and NaN'd).

Builds on the V3 API check mechanism from #662. Supersedes #657 (which targeted an earlier base before #662).

🤖 Generated with Claude Code

[github-actions]

Nit / maintainability: the 128 here is the fixed head-dim of the v3 asm bwd_hd128_dq_shuffle_group kernel, and it is repeated (uncommented) at fused_attn_ck.cpp:759 in the sizing pass. The whole branch silently assumes that whenever ck_attn_bwd_uses_v3(args) returns true, args.d_qk == 128. That's true today, but if AITER ever adds an hd64 or hd192 v3 variant, the strides here and the buffer size on the other side will both be wrong in the same way — undetectable without a repro.

Two low-cost hardening options:

• Name it (static constexpr int64_t kV3DqAccHeadDim = 128;) in a shared header so the two sites reference the same constant, or
• Add a NVTE_CHECK(args.d_qk == 128, ...) inside this branch so a future kernel-support expansion fails loudly instead of silently corrupting dQ again.

Not blocking — the fix as written is correct for the current kernel set.

[github-actions]

Two small points on the bf16 layout sizing:

1. The 128 here mirrors the hard-coded 128 in ck_fused_attn_bwd.cpp:617-620 (v3 bwd_hd128_dq_shuffle_group head dim). See the sibling comment there — sharing a named constant, or asserting d_qk == 128 on the CK side, would make the coupling explicit.

2. nvte_dtype_size(dtype) "happens to" give 2 here because the v3 atomic16 path is only reached for 16-bit input dtypes. If some future path admits fp8/fp4 backwards through this branch, the buffer would be under-sized (strides in ck_fused_attn_bwd.cpp are computed assuming the element size is 2 — element strides are in units of dtype). Using sizeof(uint16_t) or a 2 literal with a // atomic16 comment would make the layout invariant local and more robust than tying it to the input dtype.

Non-blocking nits; correctness is fine for the current set of supported dtypes.

[github-actions]

Reviewed the THD/ragged dQ backward fix for the bf16 dq_shuffle path (4 files, +62/-4).

Verdict: Looks good. Root cause and fix are well-explained inline; the sizing/stride switch is correctly gated on is_group_mode && !is_v3_atomic_fp32 && ck_attn_bwd_uses_v3(args), and the workspace buffer is conservatively sized as the max of both layouts since the v3 probe can't run in the sizing pass. Test coverage (equal-dim-128 BF16 GQA self-attn, both BSHD and THD) and the new NVTE_CK_IS_V3_ATOMIC_FP32=0 CI leg close a real coverage gap.

Nits (non-blocking): Two inline comments on the hard-coded 128 in ck_fused_attn_bwd.cpp:617-620 and fused_attn_ck.cpp:759 — the v3 kernel's fixed head-dim is repeated at both sites and coupled to the input dtype size in the sizing calc. A named constant / explicit d_qk == 128 assert would make the assumption self-documenting.

Copyright headers: OK — all four modified files carry AMD headers ending in 2026.

[wangye805]

Is this only for hdim 128 cases?

[wangye805]

1). I presume this pad16(s_q) is coming from aiter source codes, better put a code pointer here for future reference
2). previously we never test cases with sq not multiple of 16. So if we wanted to put this in, we better add a testcase for non 16 multiple sq and also make sure ck flow does not break

[wangye805]

Actually, aiter do not have deterministic path so split stride should not affect the results. (their nsplits=1)

[wangye805]

This will over allocate for hdim 64 and under allocate for hdim 192 cases

[wangye805]

I recall we have a is_v3_api_check which we can use now

[wenchenvincent]

Thanks @wangye805 — addressed in 60f2d12:

• "Only for hdim 128?" Yes. The only ragged/group dq_shuffle kernel is bwd_hd128_dq_shuffle_group, so ck_attn_bwd_uses_v3 can only return true here for d_qk==128. Made it explicit: named kV3DqAccHeadDim=128 / kV3DqAccSeqAlign=16 constants in the shared header (referenced at both the sizing and stride sites), and a hard d_qk == kV3DqAccHeadDim assert in the stride branch so a future hd64/hd192 dq_shuffle fails loudly instead of silently corrupting dQ.
• pad16 code pointer + non-16 seqlen test. Added the AITER asm_mha_varlen_bwd.cu reference (atomic16 dq_accum = {1, batch, nhead, pad16(max_seqlen_q), 128}) at both sites, and a non-16-multiple THD config (s_q=2044) to test_fused_attn.py. Verified locally with segments 103/77/61 (S=241 → pad16 256): dQ cos 1.0, finite — the CK flow handles it.
• split_stride is a no-op (nsplits=1). Correct — v3 has no deterministic path, so nsplits==1 and the split index is always 0. Noted it in a comment; set it consistently anyway.
• Over/under allocation for hd64/hd192. Fixed: the bf16 sizing is now gated on d_qk == kV3DqAccHeadDim, so hd64/hd192 use the fp32-packed size (no over-sizing). Also switched the element size to sizeof(uint16_t) with an // atomic16 comment rather than nvte_dtype_size(dtype).
• "use is_v3_api_check in the sizing pass." I looked at this — the catch is that ck_attn_bwd_uses_v3 needs a CkAttnBwdArgs, but the dq_acc allocation happens in the workspace-sizing pass which early-returns before ck_args is built (in the exec pass). Building a second minimal args just for a sizing-pass probe would reintroduce exactly the probe-vs-launch divergence that build_bwd_fmha_args (Added AITER V3 API check mechanism #662) exists to prevent. So I instead gate the sizing on d_qk == kV3DqAccHeadDim — the exact and only condition under which the bf16 group dq_shuffle runs — which is conservative (never under-sizes) and needs no probe in the sizing pass. The strides still use the real ck_attn_bwd_uses_v3 probe at launch. Happy to revisit if you'd prefer the reorder instead.