[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 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_impl always allocates the dq_acc scratch as fp32-packed (nsplits, H, total_q, d_qk) and ck_fused_attn_bwd sets batch_stride_dq_acc = 0 in group mode. That layout only matches the fp32 dq_convert post-kernel (atomic_fp32=1). On the bf16 dq_shuffle post-kernel (atomic_fp32=0), the 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

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. The fp32 dq_convert path and non-ragged BSHD/SBHD are unchanged.

Verification (MI355X / gfx950)

TE THD padding_causal backward vs a pure-JAX block-diagonal reference, NVTE_CK_IS_V3_ATOMIC_FP32=0:

• dQ cosine 0.0036 → 1.0000; per-segment all +1.0000 across [256], [100,80,76], [128,128], [64×4], [200,56].
• Still uses the fast bwd_hd128_dq_shuffle_group kernel (fix is on the fast path, not a fallback).
• atomic_fp32=1 path unchanged (no regression).
• End-to-end: MaxText llama3.1-8B block-diagonal THD packing now trains correctly (previously stalled ~6.1 loss and NaN'd).

Note: fused_attn_ck_hip.cpp is generated by hipify from fused_attn_ck.cpp, so only the source is changed here.

🤖 Generated with Claude Code

[ipanfilo]

For safety it should probably also check nvte_ck_uses_bwd_v3

[wangye805]

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]

So we may need adding extra call with v3_api_check

[wenchenvincent]

Superseded by #671, which rebases this fix onto dev on top of the V3 API check mechanism (#662). #671 uses the new ck_attn_bwd_uses_v3() probe to gate the bf16 per-segment dq_acc layout (so v2/fallback correctly keeps fp32-packed), drops the standalone QoLA aiter patch in favor of dev's dry-run patch, and adds the same tests. Closing in favor of #671.