[Perf][Sparse MLA] Drop the atomic contention in the index remap (#50365)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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
 

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)

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: