diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 2809c17c4d..60468cd337 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -674,17 +674,20 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel(): @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize("interleave", [1, 2]) @pytest.mark.parametrize("dcp_rank", [0, 1]) -def test_dcp_filter_compaction_matches_reference(interleave: int, dcp_rank: int): - """In-kernel compaction (atomic slot allocator across multiple column tiles) - must, for every row, produce exactly the rank-owned physical slots packed - into [0, valid_count) with -1 in the tail -- the same SET a reference filter - + sort/gather produces. Uses wide rows (> BLOCK_N valid slots) so the - cross-tile atomic allocation is exercised, with interior -1 gaps.""" +# 384 is not a power of two, so it exercises the multi-tile atomic allocator +# rather than the single-tile path 1024 takes. +@pytest.mark.parametrize("num_topk", [1024, 384]) +def test_dcp_filter_compaction_matches_reference( + interleave: int, dcp_rank: int, num_topk: int +): + """In-kernel compaction must, for every row, produce exactly the rank-owned + physical slots packed into [0, valid_count) with -1 in the tail -- the same + SET a reference filter + sort/gather produces. Uses wide rows (> BLOCK_N + valid slots), with interior -1 gaps.""" device = torch.device("cuda") torch.manual_seed(7) dcp_size = 2 block_size = 8 - num_topk = 1024 # > BLOCK_N(128) -> multiple tiles per row num_rows = 5 max_blocks = 64 seq = max_blocks * block_size @@ -698,7 +701,9 @@ def test_dcp_filter_compaction_matches_reference(interleave: int, dcp_rank: int) (num_rows, num_topk), -1, dtype=torch.int32, device=device ) for r in range(num_rows): - n_valid = int(torch.randint(200, 600, (1,)).item()) + # Ids are distinct, so a row holds at most `seq` valid entries. + hi = min(num_topk * 3 // 4, seq) + n_valid = int(torch.randint(num_topk // 4, hi, (1,)).item()) perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32) token_indices[r, :n_valid] = perm diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index a30549d5f8..dec43fd874 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1250,32 +1250,36 @@ def test_split_indexer_prefill_chunks_single_request_overflow(): assert out == expected -def test_triton_convert_returns_valid_counts(): +# 384 is not a power of two, so it counts via the tiled atomic accumulation +# rather than the single-tile path 128 takes. +@pytest.mark.parametrize("num_topk_tokens", [128, 384]) +def test_triton_convert_returns_valid_counts(num_topk_tokens: int): """Test that return_valid_counts correctly counts non-negative indices.""" device = torch.device(DEVICE_TYPE) num_tokens = 8 num_requests = 2 max_blocks_per_req = 10 block_size = 64 - num_topk_tokens = 128 req_id = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device) block_table = torch.arange( num_requests * max_blocks_per_req, dtype=torch.int32, device=device ).view(num_requests, max_blocks_per_req) - # Create token indices with varying numbers of valid entries - # Token 0: 64 valid, 64 invalid (-1) - # Token 1: 32 valid, 96 invalid - # Token 2: 128 valid (all) - # Token 3: 1 valid, 127 invalid - # etc. + # Create token indices with varying numbers of valid entries: half the row, + # a quarter of it, the whole row, then a single valid entry -- twice over. token_indices = torch.full( (num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device ) + valid_counts_per_token = [ + num_topk_tokens // 2, + num_topk_tokens // 4, + num_topk_tokens, + 1, + ] * 2 expected_valid = [] for i in range(num_tokens): - num_valid = [64, 32, 128, 1, 64, 32, 128, 1][i] + num_valid = valid_counts_per_token[i] token_indices[i, :num_valid] = torch.arange( num_valid, dtype=torch.int32, device=device ) % (block_size * max_blocks_per_req) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 681ab7cfd9..fc340a2246 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -23,6 +23,9 @@ def _convert_req_index_to_global_index_kernel( BLOCK_N: tl.constexpr, # tile width along columns HAS_PREFILL: tl.constexpr, COUNT_VALID: tl.constexpr, # whether to count valid indices + # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count + # is an in-register reduction and needs no atomic. + SINGLE_TILE: tl.constexpr, # When set, scatter valid slots to a contiguous prefix [0, valid_count) using # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries). @@ -99,10 +102,15 @@ def _convert_req_index_to_global_index_kernel( # sum gives each valid lane a distinct local offset; one atomic add of the # tile's valid count reserves a contiguous base across racing tiles. The # out buffer is pre-filled with -1, so unwritten tail slots stay -1. + # With no racing tiles the base is 0 and the allocator becomes a store. is_valid = (~is_invalid_tok).to(tl.int32) local_offset = tl.cumsum(is_valid) - is_valid tile_valid_count = tl.sum(is_valid) - base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + if SINGLE_TILE: + base = 0 + tl.store(valid_count_ptr + token_id, tile_valid_count) + else: + base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) dest = base + local_offset out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1 tl.store(out_ptr_dest, out_val, mask=is_valid == 1) @@ -111,10 +119,36 @@ def _convert_req_index_to_global_index_kernel( out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 tl.store(out_ptr_ij, out_val) - # Count valid indices in this tile and atomically add to row total + # Accumulate the tile's valid count into the row total; a single tile's + # reduction *is* the total. if COUNT_VALID: tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32)) - tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + if SINGLE_TILE: + tl.store(valid_count_ptr + token_id, tile_valid_count) + else: + tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + + +def _remap_tiling( + NUM_TOPK_TOKENS: int, BLOCK_N: int, count_valid: bool +) -> tuple[bool, int, int, int]: + """Pick the column tiling for the index remap kernel. + + Counting the valid slots per row is the only reason the column tiles have to + talk to each other, so when counting give one program the whole row: the + count becomes an in-register reduction plus a plain store, needing neither + atomics nor a zero-initialized counter. The row is one ``tl.arange``, so this + needs a power-of-two width; other top-k sizes stay tiled and atomic. + + Returns: + (single_tile, block_n, tiles_per_row, num_warps) + """ + single_tile = ( + count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS + ) + if single_tile: + return True, NUM_TOPK_TOKENS, 1, 8 + return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4 def triton_convert_req_index_to_global_index( @@ -172,7 +206,10 @@ def triton_convert_req_index_to_global_index( num_tokens = req_id.shape[0] max_num_blocks_per_req = block_table.shape[1] - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + single_tile, block_n, tiles_per_row, num_warps = _remap_tiling( + NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts + ) # Ensure contiguous tensors on the same device req_id_c = req_id.contiguous() @@ -180,12 +217,11 @@ def triton_convert_req_index_to_global_index( token_indices_c = token_indices.contiguous() out = torch.empty_like(token_indices_c) - # Allocate valid count buffer if needed (must be zero-initialized for atomics) valid_counts: torch.Tensor | None = None if return_valid_counts: - valid_counts = torch.zeros( - num_tokens, dtype=torch.int32, device=token_indices.device - ) + # Zero-init only matters for the atomic accumulation path. + alloc = torch.empty if single_tile else torch.zeros + valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) # Strides in elements bt_stride0, bt_stride1 = block_table_c.stride() @@ -213,9 +249,10 @@ def triton_convert_req_index_to_global_index( # shapes / constexprs max_num_blocks_per_req, BLOCK_SIZE, - BLOCK_N, + block_n, HAS_PREFILL_WORKSPACE, return_valid_counts, + single_tile, False, # COMPACT_TO_FRONT: keep input column == output column # DCP disabled (no-op de-interleave) 1, @@ -228,6 +265,7 @@ def triton_convert_req_index_to_global_index( ti_stride1, out_stride0, out_stride1, + num_warps=num_warps, ) if return_valid_counts: @@ -287,15 +325,20 @@ def triton_filter_and_convert_dcp_index( num_tokens = req_id.shape[0] max_num_blocks_per_req = block_table.shape[1] - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N req_id_c = req_id.contiguous() block_table_c = block_table.contiguous() token_indices_c = token_indices.contiguous() - # The compaction uses the valid-count buffer as an atomic slot allocator, so - # it requires counting. Pre-fill out with -1 so the unwritten tail stays -1. + # The compaction uses the valid-count buffer as a slot allocator, so it + # requires counting. Pre-fill out with -1 so the unwritten tail stays -1. count_valid = return_valid_counts or compact_valid_to_front + + # The compaction builds on the counting, so it shares the tiling. + single_tile, block_n, tiles_per_row, num_warps = _remap_tiling( + NUM_TOPK_TOKENS, BLOCK_N, count_valid + ) + if compact_valid_to_front: out = torch.full_like(token_indices_c, -1) else: @@ -303,9 +346,9 @@ def triton_filter_and_convert_dcp_index( valid_counts: torch.Tensor | None = None if count_valid: - valid_counts = torch.zeros( - num_tokens, dtype=torch.int32, device=token_indices.device - ) + # Zero-init only matters for the atomic accumulation path. + alloc = torch.empty if single_tile else torch.zeros + valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) bt_stride0, bt_stride1 = block_table_c.stride() ti_stride0, ti_stride1 = token_indices_c.stride() @@ -322,9 +365,10 @@ def triton_filter_and_convert_dcp_index( None, max_num_blocks_per_req, BLOCK_SIZE, - BLOCK_N, + block_n, False, # HAS_PREFILL count_valid, + single_tile, compact_valid_to_front, dcp_size, dcp_rank, @@ -335,6 +379,7 @@ def triton_filter_and_convert_dcp_index( ti_stride1, out_stride0, out_stride1, + num_warps=num_warps, ) if return_valid_counts: