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..0ec7556a7 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1336,6 +1336,49 @@ 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", + ), + # 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 a7425be98..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 @@ -605,6 +605,33 @@ 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), 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)) { + // 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; 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..5675f0b25 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,27 @@ 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 -- 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 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); void* dk_expanded_ptr = nullptr; void* dv_expanded_ptr = nullptr; @@ -892,8 +911,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;