[ROCm] Add THD/ragged bf16 (atomic16) backward test coverage

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

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,