From 891f9f73a40dda479aa53614686361ba373c1432 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Thu, 23 Jul 2026 21:54:14 -0700 Subject: [PATCH 1/3] Remove alibi, score_max/sum_exp, and dropout randomization from test_mhas_v2 Drop is_alibi, with_score_max, with_score_sum_exp, and is_dropout from the randomized SDPA test configurations. The plumbing in the sdpa test package is kept; only the test file stops exercising these options. Co-Authored-By: Claude Fable 5 --- test/python/test_mhas_v2.py | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 746cf2d1d..e3af2ff5a 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -116,8 +116,6 @@ def test_sdpa_random_fwd_L0(env_info, test_no, request, cudnn_handle): with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -148,18 +146,13 @@ def test_sdpa_random_fwd_unified_L1(env_info, test_no, request, cudnn_handle): with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_bias=RandomChoice({True : 1, False : 3}), - is_alibi=RandomChoice({True : 1, False : 3}), is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "cu_padded" : 1, "full" : 1}), with_unfuse_fma=RandomChoice({True : 1, False : 1}), # Randomly enable unfuse_fma for SM100 - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), - is_dropout=RandomChoice({True : 1, False : 3}), with_rope=RandomChoice({True : 1, False : 3}), # RoPE at end to preserve existing test distributions ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) - test.cfg.dropout_prob = 0.1 if test.cfg.is_dropout else 0.0 test.cfg.implementation = getattr(cudnn.attention_implementation, request.config.getoption("--implementation") or "", cudnn.attention_implementation.UNIFIED) test.showConfig(test_no, request) @@ -235,8 +228,6 @@ def test_sdpa_random_sq1_L0(env_info, test_no, request, cudnn_handle): with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 0, "full" : 1}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), # sink_token not supported with s_q==1 # dropout not supported with s_q==1 ) as randomization_ctx: @@ -268,8 +259,6 @@ def test_sdpa_random_sq1_unified_L1(env_info, test_no, request, cudnn_handle): with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 0, "full" : 1}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), # sink_token not supported with s_q==1 # dropout not supported with s_q==1 ) as randomization_ctx: @@ -306,8 +295,6 @@ def test_sdpa_random_lean_attn_L0(env_info, test_no, request, cudnn_handle): with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), # sink_token not supported with s_q==1 # dropout not supported with s_q==1 ) as randomization_ctx: @@ -339,8 +326,6 @@ def test_sdpa_random_lean_attn_unified_L1(env_info, test_no, request, cudnn_hand with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), # sink_token not supported with s_q==1 # dropout not supported with s_q==1 ) as randomization_ctx: @@ -376,8 +361,6 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle): with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -408,14 +391,10 @@ def test_sdpa_random_fwd_ragged_unified_L1(env_info, test_no, request, cudnn_han with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), # Modified from non-unified test diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "cu_ragged" : 1, "padded" : 0, "full" : 0}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), - is_dropout=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) - test.cfg.dropout_prob = 0.1 if test.cfg.is_dropout else 0.0 test.cfg.implementation = getattr(cudnn.attention_implementation, request.config.getoption("--implementation") or "", cudnn.attention_implementation.UNIFIED) test.showConfig(test_no, request) @@ -450,8 +429,6 @@ def test_sdpa_random_fwd_ragged_offset_multiplier_unified_L1(env_info, test_no, with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), is_ragged_or_padded_or_full=RandomChoice({"ragged_mult" : 1, "cu_ragged_mult" : 1}), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -522,8 +499,6 @@ def test_sdpa_fwd_paged_L0(env_info, test_no, request, cudnn_handle): diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 0}), block_size=RandomBlockSize(min=1, max=1024, with_high_probability=[1,32,128]), - with_score_max=RandomChoice({True : 1, False : 3}), - with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) From 81c7e373243a1d9ef5cb88c7e008abb985080ed3 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 24 Jul 2026 14:05:10 -0700 Subject: [PATCH 2/3] Add dedicated tests for rarely-used SDPA features Cover ALiBi, score_max/score_sum_exp, and dropout with a few deterministic fwd/bwd configs in test_sdpa_unpopular_features.py, now that test_mhas_v2.py no longer randomizes over these options. Co-Authored-By: Claude Fable 5 --- test/python/test_sdpa_unpopular_features.py | 86 +++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 test/python/test_sdpa_unpopular_features.py diff --git a/test/python/test_sdpa_unpopular_features.py b/test/python/test_sdpa_unpopular_features.py new file mode 100644 index 000000000..85117e53e --- /dev/null +++ b/test/python/test_sdpa_unpopular_features.py @@ -0,0 +1,86 @@ +""" +Dedicated coverage for rarely-used SDPA features that are no longer exercised +by the randomized configurations in test_mhas_v2.py: ALiBi masking, the +score_max/score_sum_exp softmax outputs, and dropout. A few deterministic +configs verify that basic graph lowering and numerics keep working. +""" + +import cudnn +import pytest +import torch + +from sdpa.fp16 import exec_sdpa +from sdpa.random_config import ExecConfig + + +def make_config( + *, + data_type=torch.bfloat16, + is_infer=True, + is_alibi=False, + is_dropout=False, + dropout_prob=0.0, + with_score_max=False, + with_score_sum_exp=False, + right_bound=None, +): + cfg = ExecConfig( + data_type=data_type, + rng_data_seed=1234, + rng_geom_seed=5678, + is_alibi=is_alibi, + is_infer=is_infer, + is_paged=False, + is_bias=False, + is_block_mask=False, + is_padding=False, + is_cu_seq_len=False, + is_ragged=False, + is_dropout=is_dropout, + is_determin=False, + with_score_max=with_score_max, + with_score_sum_exp=with_score_sum_exp, + batches=2, + d_qk=64, + d_v=64, + s_q=512, + s_kv=512, + h_q=4, + h_k=4, + h_v=4, + diag_align=cudnn.diagonal_alignment.TOP_LEFT, + left_bound=None, + right_bound=right_bound, + dropout_prob=dropout_prob, + implementation=cudnn.attention_implementation.AUTO, + ) + cfg.fill_derived_fields() + return cfg + + +@pytest.mark.L0 +@pytest.mark.parametrize("is_infer", [True, False], ids=["fwd", "bwd"]) +def test_sdpa_alibi(is_infer, request, cudnn_handle): + # ALiBi requires a causal mask (right_bound=0, top-left alignment). + cfg = make_config(is_infer=is_infer, is_alibi=True, right_bound=0) + exec_sdpa(cfg, request, cudnn_handle) + + +@pytest.mark.L0 +@pytest.mark.parametrize("is_infer", [True, False], ids=["fwd", "bwd"]) +def test_sdpa_score_max_sum_exp(is_infer, request, cudnn_handle): + cfg = make_config(is_infer=is_infer, with_score_max=True, with_score_sum_exp=True) + exec_sdpa(cfg, request, cudnn_handle) + + +@pytest.mark.L0 +@pytest.mark.parametrize("is_infer", [True, False], ids=["fwd", "bwd"]) +def test_sdpa_dropout(is_infer, request, cudnn_handle): + cfg = make_config( + data_type=torch.float16, + is_infer=is_infer, + is_dropout=True, + dropout_prob=0.1, + right_bound=0, + ) + exec_sdpa(cfg, request, cudnn_handle) From e9c7238a21dc84064abf6a10282e1c566b8e0064 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 24 Jul 2026 19:35:26 -0700 Subject: [PATCH 3/3] Rename test_sdpa_unpopular_features.py to test_sdpa_custom_features.py Co-Authored-By: Claude Fable 5 --- ...st_sdpa_unpopular_features.py => test_sdpa_custom_features.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename test/python/{test_sdpa_unpopular_features.py => test_sdpa_custom_features.py} (100%) diff --git a/test/python/test_sdpa_unpopular_features.py b/test/python/test_sdpa_custom_features.py similarity index 100% rename from test/python/test_sdpa_unpopular_features.py rename to test/python/test_sdpa_custom_features.py