From f75f7155464d8adb31ff35c50125512cc7ee5729 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 26 Aug 2026 19:44:35 +0100 Subject: [PATCH] test(attention): follow the sdnq block mask helpers upstream get_block_mask_input takes only the mask now, so the helper pads to 4d itself, and the standalone nesting filter is inlined into prune_configs. The raise for a tile list that cannot nest moves into the prune_configs row, pinned with from_small so the fallback to the small configs cannot refill the list before the filter runs. --- test/test-attention-sdnq-sparse.py | 42 ++++++++++++------------------ 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/test/test-attention-sdnq-sparse.py b/test/test-attention-sdnq-sparse.py index c6995aabf..35d30fbd3 100644 --- a/test/test-attention-sdnq-sparse.py +++ b/test/test-attention-sdnq-sparse.py @@ -174,15 +174,17 @@ def expand(keep: torch.Tensor, seq_q: int, seq_kv: int): return keep.repeat_interleave(BLOCK_M, dim=-2).repeat_interleave(BLOCK_N, dim=-1)[..., :seq_q, :seq_kv].contiguous() -def ordered(keep: torch.Tensor, batch: int, heads: int): - """The count and ascending index list per query block that the kernel walks, built the way the entry builds them: size-1 batch and head dims kept, padded to the descriptor chunks.""" - return atten_module.get_block_mask_input(keep, batch, heads) +def ordered(keep: torch.Tensor): + """The count and ascending index list per query block that the kernel walks, built the way the entry builds them: left-padded to 4d, size-1 batch and head dims kept, padded to the descriptor chunks.""" + while keep.ndim < 4: + keep = keep.unsqueeze(0) + return atten_module.get_block_mask_input(keep) def raw(q, k, v, attn_mask=None, block_mask=None, block_count=None, block_index=None): """The triton op without the input prep: nothing is quantized, so the mask path is the only difference between arms.""" if block_mask is not None: - block_count, block_index = ordered(block_mask, q.shape[0], q.shape[1]) + block_count, block_index = ordered(block_mask) sparse = block_count is not None return atten_module.sdnq_triton_atten_fwd( q, k, v, None, None, None, @@ -262,7 +264,7 @@ def test_block_lists_must_be_contiguous_and_padded(): return None q, k, v = qkv(seq=1100) # 9 query blocks and 18 kv blocks, so both lists get padded (to 12 and 32) keep = random_keep(2, 4, 1100, 1100) - count, index = ordered(keep, 2, 4) + count, index = ordered(keep) nq, nk = blocks(1100, 1100) assert count.is_contiguous() and index.is_contiguous() assert count.shape[-1] % atten_module.block_count_chunk == 0 and index.shape[-1] % atten_module.block_index_chunk == 0 @@ -287,7 +289,7 @@ def test_block_lists_must_be_contiguous_and_padded(): raise AssertionError(f'{label} was accepted') assert torch.equal(raw(q, k, v, block_mask=keep), reference) shared_keep = keep[:1, :1] - shared_count, shared_index = ordered(shared_keep, 2, 4) + shared_count, shared_index = ordered(shared_keep) assert tuple(shared_count.shape[:2]) == (1, 1) and tuple(shared_index.shape[:2]) == (1, 1), 'size-1 batch and head dims stay as given, the kernel broadcasts by shape' assert torch.equal(raw(q, k, v, block_count=shared_count, block_index=shared_index), raw(q, k, v, block_mask=shared_keep.expand(2, 4, -1, -1).contiguous())) return True @@ -365,8 +367,8 @@ def test_launcher_validates_the_block_lists(): return None q, k, v = qkv() ones = torch.ones(1, 1, *blocks(1000, 1000), dtype=torch.int8, device=device) - count, index = ordered(ones, 2, 4) - wide_count, wide_index = ordered(ones.expand(5, 4, -1, -1).contiguous(), 5, 4) + count, index = ordered(ones) + wide_count, wide_index = ordered(ones.expand(5, 4, -1, -1).contiguous()) bad = ( ('a count without an index', {'block_count': count}), ('an index without a count', {'block_index': index}), @@ -392,22 +394,8 @@ def test_launcher_validates_the_block_lists(): # Autotune nesting # ============================================================ -def test_nesting_filter_keeps_only_nesting_tiles(): - configs = configs_for((32, 64, 128, 256), (16, 32, 64, 128)) - args = {'do_block_mask': 1, 'BLOCK_MASK_M': BLOCK_M, 'BLOCK_MASK_N': BLOCK_N} - kept = atten_module.nest_block_mask_configs(configs, args) - assert {(c.kwargs['BLOCK_SIZE_M'], c.kwargs['BLOCK_SIZE_N']) for c in kept} == {(m, n) for m in (32, 64, 128) for n in (16, 32, 64)} - assert atten_module.nest_block_mask_configs(configs, {'do_block_mask': 0}) is configs - assert atten_module.nest_block_mask_configs(configs, {}) is configs - try: - atten_module.nest_block_mask_configs(configs_for((256,), (128,)), args) - except ValueError as e: - assert 'SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST' in str(e), e - return True - raise AssertionError('a tile list with nothing nesting did not raise') - - def test_prune_configs_applies_the_filter(): + """The nesting filter keeps exactly the tiles a mask block divides by, leaves a dense launch alone, and raises rather than letting a tile list that cannot nest run.""" if not kernel_available: return None q, k, v = qkv() @@ -420,7 +408,12 @@ def test_prune_configs_applies_the_filter(): without = atten_module.prune_configs(configs_for((32, 64, 128, 256), (16, 32, 64, 128)), dict(args, do_block_mask=0)) assert with_mask, 'nothing survived' assert [c.kwargs for c in with_mask] == [c.kwargs for c in without if nests(c)], ([c.kwargs for c in with_mask], [c.kwargs for c in without]) - return True + try: # from_small is the terminal arm, so the fallback to the small configs cannot refill the list before the filter reaches it + atten_module.prune_configs(configs_for((256,), (128,)), args, from_small=True) + except ValueError as e: + assert 'SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST' in str(e), e + return True + raise AssertionError('a tile list with nothing nesting did not raise') # ============================================================ @@ -641,7 +634,6 @@ def run_all(): log.warning('=== autotune nesting ===') cat = category('autotune') for fn in [ - test_nesting_filter_keeps_only_nesting_tiles, test_prune_configs_applies_the_filter, ]: run_test(cat, fn)