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,