From fd0bbcee8f8988ebb90f434d3e5bf8b7f1987988 Mon Sep 17 00:00:00 2001 From: Wen Chen Date: Thu, 16 Jul 2026 00:55:34 +0000 Subject: [PATCH 1/2] [ROCm] Fix THD/ragged dQ backward on the bf16 dq_shuffle path (gfx950) 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, 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. 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 strides. v2/fallback and non-ragged/atomic_fp32 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 -- and an NVTE_CK_IS_V3_ATOMIC_FP32=0 (atomic16 / dq_shuffle) CI run of the THD backward tests, which the default fp32 path never exercised. Verified on MI355X (gfx950): THD padding_causal backward dQ cos 0.0036 -> 1.0000 with atomic_fp32=0; per-segment all correct across [256], [100,80,76], [128,128], [64x4], [200,56], for both the v3 (dq_shuffle) and v2 paths; atomic_fp32=1 unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- ci/jax.sh | 4 +++ tests/jax/test_fused_attn.py | 28 +++++++++++++++++++ .../ck_fused_attn/src/ck_fused_attn_bwd.cpp | 15 ++++++++++ .../common/fused_attn_rocm/fused_attn_ck.cpp | 19 ++++++++++--- 4 files changed, 62 insertions(+), 4 deletions(-) diff --git a/ci/jax.sh b/ci/jax.sh index 648e910fc..d58cbe021 100755 --- a/ci/jax.sh +++ b/ci/jax.sh @@ -60,6 +60,10 @@ run_test_config() { run 1 test_fused_attn.py NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 run_default_fa_lbl "deterministic" 3 test_fused_attn.py -k "TestFusedAttnWithDeterminism" NVTE_CK_USES_FWD_V3=0 NVTE_CK_USES_BWD_V3=0 run_default_fa_lbl "v2" 3 test_fused_attn.py # Using FAv2 for forward and backward pass + # bf16 atomic dq accumulation (dq_shuffle post-kernel). Default is fp32 (atomic32/dq_convert), so the + # bf16 dq_acc path is otherwise never exercised in CI. Scope to THD/RAGGED backward where the + # group-mode per-segment dq_acc layout matters (see the equal-dim-128 RAGGED_SELF config). + NVTE_CK_IS_V3_ATOMIC_FP32=0 run_default_fa_lbl "atomic16" 3 test_fused_attn.py -k "test_backward and RAGGED" run_default_fa 1 test_layer.py # it effectively always uses unfused attention run_default_fa 1 test_sanity_import.py run_default_fa 1 test_softmax.py diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 78fb2f103..36b1b808d 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1336,6 +1336,34 @@ def check_dqkv(primitive, reference, pad, idx): QKVLayout.THD_THD_THD, id="2-1024-2048-12-6-128-64-FP16-CROSS-GQA-RAGGED_SEPARATE", ), + # equal head_dim 128 + bf16 + gqa + self-attn (llama3-like). Previously there was no + # equal-dim-128 self-attn coverage; the THD/RAGGED case exercises the group-mode dq_acc + # backward (the bf16 dq_shuffle path when NVTE_CK_IS_V3_ATOMIC_FP32=0), which was silently + # producing wrong dQ for every packed segment past cu_seqlens offset 0. + pytest.param( + 2, + 2048, + 2048, + 12, + 6, + 128, + 128, + jnp.bfloat16, + QKVLayout.BSHD_BSHD_BSHD, + id="2-2048-2048-12-6-128-128-BF16-GQA-SELF", + ), + pytest.param( + 2, + 2048, + 2048, + 12, + 6, + 128, + 128, + jnp.bfloat16, + QKVLayout.THD_THD_THD, + id="2-2048-2048-12-6-128-128-BF16-GQA-RAGGED_SELF", + ), pytest.param( 10, 4096, 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 a7425be98..2eaa1bcca 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 @@ -605,6 +605,21 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ const bool has_dbias = built.has_dbias; const BiasShape bias_shape = built.bias_shape; + // dq_acc scratch layout (mirrors AITER's asm_mha_varlen_bwd wiring). build_bwd_fmha_args sets the + // fp32-packed layout (nsplits, H, total_q, d_qk) with batch_stride_dq_acc=0 -- correct for the fp32 + // dq_convert post-kernel. But the bf16 dq_shuffle post-kernel (is_v3_atomic_fp32=0) expects a + // per-segment padded layout (nsplits, B, H, pad16(s_q), 128); with batch_stride=0 it mis-indexes every + // ragged segment past cu_seqlens offset 0, corrupting dQ. Switch to the per-segment layout only when + // the v3 asm path actually runs (ck_attn_bwd_uses_v3 probe) -- v2/fallback keeps the fp32-packed layout. + // The dq_acc buffer is sized to hold either layout in fused_attn_ck.cpp. + if (args.is_group_mode() && !args.is_v3_atomic_fp32 && ck_attn_bwd_uses_v3(args)) { + const int64_t padded_sq = static_cast(((args.s_q + 15) / 16) * 16); + fmha_args.stride_dq_acc = 128; + fmha_args.nhead_stride_dq_acc = static_cast(padded_sq * 128); + fmha_args.batch_stride_dq_acc = static_cast(args.h * padded_sq * 128); + fmha_args.split_stride_dq_acc = static_cast(args.b * args.h * padded_sq * 128); + } + bool ck_log_config = false; if (const char* env_p = std::getenv("CK_FUSED_ATTN_LOG_CONFIG") ) { if (env_p != nullptr && std::string(env_p) == "1") 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 0c7d80a6a..ef79e4dbe 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,18 @@ 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 a dq_acc scratch buffer. ck_attn_bwd picks the layout at launch time: fp32-packed + // (nsplits, H, total_q, d_qk), or -- when the v3 asm dq_shuffle path runs (is_v3_atomic_fp32=0, ragged) + // -- bf16 per-segment padded (nsplits, B, H, pad16(s_q), 128). The v3 choice needs a GPU probe that must + // not run in the workspace-sizing pass, so size the buffer to hold EITHER layout for the ragged/non-atomic + // case; ck_attn_bwd sets the strides. Non-ragged / atomic_fp32 only ever use the fp32-packed layout. + const size_t dq_acc_fp32_bytes = nsplits*h*max_tokens_q*d_qk*sizeof(float); + const size_t dq_acc_padded_sq = ((s_q + 15) / 16) * 16; + const size_t dq_acc_bf16_bytes = nsplits*b*h*dq_acc_padded_sq*128*nvte_dtype_size(dtype); + const bool dq_acc_max_of_both = is_ragged && !nvte_ck_is_v3_atomic_fp32; + const size_t dq_acc_bytes = (dq_acc_max_of_both && dq_acc_bf16_bytes > dq_acc_fp32_bytes) + ? dq_acc_bf16_bytes : dq_acc_fp32_bytes; + void* dq_acc_ptr = planner.allocate(dq_acc_bytes); void* dk_expanded_ptr = nullptr; void* dv_expanded_ptr = nullptr; @@ -892,8 +902,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 is sized to hold either layout (see allocation above); ck_attn_bwd picks the strides. + // CK requires zeroing the whole buffer. + NVTE_CHECK_CUDA(cudaMemsetAsync(dq_acc_ptr, 0, dq_acc_bytes, stream)); if(devPtrAlibiSlope){ dim3 block, grid; block.x = 1024; From 60f2d1271da9442211e39b077dd13d69fdba9a29 Mon Sep 17 00:00:00 2001 From: Wen Chen Date: Tue, 21 Jul 2026 05:19:25 +0000 Subject: [PATCH 2/2] [ROCm] Address review: harden dq_acc hd128 assumption + non-16 seqlen test Review feedback on the THD/ragged dq_shuffle dq_acc layout (wangye805): - Name the fixed v3 dq_shuffle head dim as ck_fused_attn::kV3DqAccHeadDim (128) and the seqlen alignment as kV3DqAccSeqAlign (16) in the shared header, and reference them from both the dq_acc sizing (fused_attn_ck.cpp) and the stride override (ck_fused_attn_bwd.cpp) so the two sites can't drift. - Add a code pointer to AITER's asm_mha_varlen_bwd.cu (the atomic16 dq_acc layout / pad16 origin) at both sites. - Assert d_qk == kV3DqAccHeadDim in the stride override: the only ragged dq_shuffle kernel is bwd_hd128_dq_shuffle_group, so a future hd64/hd192 dq_shuffle would otherwise silently corrupt dQ again -- now it fails loudly. - Gate the bf16 dq_acc sizing on d_qk == kV3DqAccHeadDim (the exact hd128-only condition) instead of just ragged/non-atomic, so it no longer over-sizes for hd64/hd192. Use sizeof(uint16_t) with an // atomic16 comment for the element size rather than tying it to the input dtype. - Note that split_stride_dq_acc is a no-op here (v3 has no deterministic mode, so nsplits==1); set it consistently anyway. - Add a non-16-multiple seqlen THD config (2044) to test_fused_attn.py to exercise the pad16(s_q) rounding and confirm the CK flow does not break. Verified on MI355X (gfx950): dQ cos 1.0 for v3/v2/atomic_fp32=1 and per-segment across [256],[100,80,76],[128,128],[64x4],[200,56]; non-16 seqlen (segments 103/77/61, S=241 -> pad16 256) dQ cos 1.0, finite. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/jax/test_fused_attn.py | 15 ++++++++++ .../include/ck_fused_attn/ck_fused_attn.hpp | 9 ++++++ .../ck_fused_attn/src/ck_fused_attn_bwd.cpp | 30 +++++++++++++------ .../common/fused_attn_rocm/fused_attn_ck.cpp | 25 +++++++++++----- 4 files changed, 62 insertions(+), 17 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 36b1b808d..0ec7556a7 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1364,6 +1364,21 @@ def check_dqkv(primitive, reference, pad, idx): QKVLayout.THD_THD_THD, id="2-2048-2048-12-6-128-128-BF16-GQA-RAGGED_SELF", ), + # non-16-multiple seqlen: exercises the atomic16 dq_acc seqlen padding (pad16(s_q)) on the bf16 + # dq_shuffle THD backward path; make sure the CK flow handles a max_seqlen that is not a multiple of + # kV3DqAccSeqAlign (2044 % 16 == 12). + pytest.param( + 2, + 2044, + 2044, + 12, + 6, + 128, + 128, + jnp.bfloat16, + QKVLayout.THD_THD_THD, + id="2-2044-2044-12-6-128-128-BF16-GQA-RAGGED_SELF-UNALIGNED", + ), pytest.param( 10, 4096, diff --git a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp index 150c8bfb3..ee4cd6fe1 100644 --- a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp +++ b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp @@ -14,6 +14,15 @@ namespace ck_fused_attn{ +// Fixed head dim of AITER's v3 asm bf16 dq_shuffle backward path. AITER's atomic16 dq_acc scratch is +// padded to this head dim regardless of the actual d_qk (see AITER csrc/py_itfs_cu/asm_mha_varlen_bwd.cu: +// "padding dq_accum seqlen to 16x of max_seqlen_q, head dim to 128"), and the only ragged/group dq_shuffle +// kernel is bwd_hd128_dq_shuffle_group. Both the dq_acc sizing (fused_attn_ck.cpp) and the stride override +// (ck_fused_attn_bwd.cpp) reference this so the two sites can never drift apart. +static constexpr int64_t kV3DqAccHeadDim = 128; +// AITER pads the dq_acc seqlen up to a multiple of this in the atomic16 layout (same source reference). +static constexpr int64_t kV3DqAccSeqAlign = 16; + // input qkv dtypes enum class DType { kFloat16 = 0, /*!< 16-bit float (E5M10) */ 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 2eaa1bcca..8dccfd170 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 @@ -608,16 +608,28 @@ hipError_t ck_attn_bwd(const CkAttnBwdArgs& args, hipStream_t stream){ // dq_acc scratch layout (mirrors AITER's asm_mha_varlen_bwd wiring). build_bwd_fmha_args sets the // fp32-packed layout (nsplits, H, total_q, d_qk) with batch_stride_dq_acc=0 -- correct for the fp32 // dq_convert post-kernel. But the bf16 dq_shuffle post-kernel (is_v3_atomic_fp32=0) expects a - // per-segment padded layout (nsplits, B, H, pad16(s_q), 128); with batch_stride=0 it mis-indexes every - // ragged segment past cu_seqlens offset 0, corrupting dQ. Switch to the per-segment layout only when - // the v3 asm path actually runs (ck_attn_bwd_uses_v3 probe) -- v2/fallback keeps the fp32-packed layout. - // The dq_acc buffer is sized to hold either layout in fused_attn_ck.cpp. + // per-segment padded layout (nsplits, B, H, pad16(s_q), kV3DqAccHeadDim); with batch_stride=0 it + // mis-indexes every ragged segment past cu_seqlens offset 0, corrupting dQ. Switch to the per-segment + // layout only when the v3 asm path actually runs (ck_attn_bwd_uses_v3 probe) -- v2/fallback keeps the + // fp32-packed layout. The dq_acc buffer is sized to hold either layout in fused_attn_ck.cpp. + // Layout + padding reference: AITER csrc/py_itfs_cu/asm_mha_varlen_bwd.cu (atomic16 dq_accum = + // torch::zeros({1, batch, nhead, pad16(max_seqlen_q), 128})). if (args.is_group_mode() && !args.is_v3_atomic_fp32 && ck_attn_bwd_uses_v3(args)) { - const int64_t padded_sq = static_cast(((args.s_q + 15) / 16) * 16); - fmha_args.stride_dq_acc = 128; - fmha_args.nhead_stride_dq_acc = static_cast(padded_sq * 128); - fmha_args.batch_stride_dq_acc = static_cast(args.h * padded_sq * 128); - fmha_args.split_stride_dq_acc = static_cast(args.b * args.h * padded_sq * 128); + // The only ragged/group dq_shuffle kernel is bwd_hd128_dq_shuffle_group, so ck_attn_bwd_uses_v3 can + // only be true here for d_qk == kV3DqAccHeadDim. Assert it so a future hd64/hd192 dq_shuffle kernel + // (which would need a different dq_acc head dim) fails loudly instead of silently corrupting dQ. + if (args.d_qk != kV3DqAccHeadDim) { + throw std::runtime_error("ck_fused_attn bwd: bf16 dq_shuffle dq_acc layout assumes d_qk == 128."); + } + const int64_t padded_sq = ((static_cast(args.s_q) + kV3DqAccSeqAlign - 1) + / kV3DqAccSeqAlign) * kV3DqAccSeqAlign; + fmha_args.stride_dq_acc = kV3DqAccHeadDim; + fmha_args.nhead_stride_dq_acc = padded_sq * kV3DqAccHeadDim; + fmha_args.batch_stride_dq_acc = static_cast(args.h) * padded_sq * kV3DqAccHeadDim; + // split_stride is irrelevant in practice: the v3 asm path has no deterministic mode, so nsplits==1 + // and the split index is always 0. Set it consistently anyway. + fmha_args.split_stride_dq_acc = static_cast(static_cast(args.b) * args.h + * padded_sq * kV3DqAccHeadDim); } bool ck_log_config = false; 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 ef79e4dbe..5675f0b25 100644 --- a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp +++ b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp @@ -750,15 +750,24 @@ void fused_attn_ck_bwd_impl( void* lse_workspace = planner.allocate(h*max_tokens_q*sizeof(float)); // CK requires a dq_acc scratch buffer. ck_attn_bwd picks the layout at launch time: fp32-packed - // (nsplits, H, total_q, d_qk), or -- when the v3 asm dq_shuffle path runs (is_v3_atomic_fp32=0, ragged) - // -- bf16 per-segment padded (nsplits, B, H, pad16(s_q), 128). The v3 choice needs a GPU probe that must - // not run in the workspace-sizing pass, so size the buffer to hold EITHER layout for the ragged/non-atomic - // case; ck_attn_bwd sets the strides. Non-ragged / atomic_fp32 only ever use the fp32-packed layout. + // (nsplits, H, total_q, d_qk), or -- when the v3 asm dq_shuffle path runs -- bf16 per-segment padded + // (nsplits, B, H, pad16(s_q), kV3DqAccHeadDim). See AITER csrc/py_itfs_cu/asm_mha_varlen_bwd.cu for the + // atomic16 layout. Whether v3 runs needs a probe (ck_attn_bwd_uses_v3) that depends on args only built in + // the execution pass, so we cannot call it here in the workspace-sizing pass; instead we size the buffer + // to hold EITHER layout. The bf16 group dq_shuffle kernel is hd128-only (kV3DqAccHeadDim), so it is only + // ever reachable when d_qk == kV3DqAccHeadDim -- gate on that (rather than probing) to avoid over-sizing + // hd64/hd192, and to keep this conservative but never-under-sized regardless of the launch-time decision. + // Element size is 2: the atomic16 dq_acc buffer is bf16 (16-bit), the only input dtype the v3 asm path + // admits -- do not tie it to nvte_dtype_size(dtype). const size_t dq_acc_fp32_bytes = nsplits*h*max_tokens_q*d_qk*sizeof(float); - const size_t dq_acc_padded_sq = ((s_q + 15) / 16) * 16; - const size_t dq_acc_bf16_bytes = nsplits*b*h*dq_acc_padded_sq*128*nvte_dtype_size(dtype); - const bool dq_acc_max_of_both = is_ragged && !nvte_ck_is_v3_atomic_fp32; - const size_t dq_acc_bytes = (dq_acc_max_of_both && dq_acc_bf16_bytes > dq_acc_fp32_bytes) + const bool dq_acc_bf16_possible = is_ragged && !nvte_ck_is_v3_atomic_fp32 + && d_qk == static_cast(ck_fused_attn::kV3DqAccHeadDim); + const size_t dq_acc_padded_sq = + ((s_q + ck_fused_attn::kV3DqAccSeqAlign - 1) / ck_fused_attn::kV3DqAccSeqAlign) + * ck_fused_attn::kV3DqAccSeqAlign; + const size_t dq_acc_bf16_bytes = + nsplits*b*h*dq_acc_padded_sq*ck_fused_attn::kV3DqAccHeadDim*sizeof(uint16_t); // atomic16 = bf16 + const size_t dq_acc_bytes = (dq_acc_bf16_possible && dq_acc_bf16_bytes > dq_acc_fp32_bytes) ? dq_acc_bf16_bytes : dq_acc_fp32_bytes; void* dq_acc_ptr = planner.allocate(dq_acc_bytes);