From 841166d98f4990f704e7ea4955dac2245cc527c2 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 17 Aug 2026 11:13:36 -0700 Subject: [PATCH] test_mhas_v2: draw sink tokens in the ragged bwd suites Backward + ragged + learnable sink was structurally untestable: the dense bwd suite draws sink tokens but never ragged layout, and the ragged bwd suites never drew sink tokens. This let a dSink corruption in the cuDNN backward dot_do_o pre-kernel ship unnoticed for every ragged config with d_v not in {64,128,256} (bf16/fp16) and all ragged fp8 configs (TransformerEngine issue #3249; fixed in cuDNN backend). Add the same sink-token draw the dense bwd suites use to test_sdpa_random_bwd_ragged_L0 and test_sdpa_fp8_bwd_ragged_L0. Co-Authored-By: Claude Fable 5 --- test/python/test_mhas_v2.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 129313552..9e746ae39 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -495,6 +495,7 @@ def test_sdpa_random_bwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), is_deterministic=RandomChoice({True : 3, False : 1}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), + with_sink_token=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -903,6 +904,7 @@ def test_sdpa_fp8_bwd_ragged_L0(env_info, test_no, request, cudnn_handle): diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT: 1}), is_ragged_or_padded_or_full=RandomChoice({"ragged": 1, "padded": 0, "full": 0}), is_deterministic=RandomChoice({True: 1, False: 1}), + with_sink_token=RandomChoice({True: 1, False: 2}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed)