[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>
Nick Hill committed
Aug 8, 2026 at 16:32 UTC
9e6be4a72bd29ecf2168a41f659aebfe2cdeda4a
3 files changed
+87
-33
tests/v1/attention/test_indexer_dcp_localize.py
+13
-8
@@ -674,17 +674,20 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():
674
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
675
@pytest.mark.parametrize("interleave", [1, 2])
676
@pytest.mark.parametrize("dcp_rank", [0, 1])
677
-def test_dcp_filter_compaction_matches_reference(interleave: int, dcp_rank: int):
678
- """In-kernel compaction (atomic slot allocator across multiple column tiles)
679
- must, for every row, produce exactly the rank-owned physical slots packed
680
- into [0, valid_count) with -1 in the tail -- the same SET a reference filter
681
- + sort/gather produces. Uses wide rows (> BLOCK_N valid slots) so the
682
- cross-tile atomic allocation is exercised, with interior -1 gaps."""
677
+# 384 is not a power of two, so it exercises the multi-tile atomic allocator
678
+# rather than the single-tile path 1024 takes.
679
+@pytest.mark.parametrize("num_topk", [1024, 384])
680
+def test_dcp_filter_compaction_matches_reference(
681
+ interleave: int, dcp_rank: int, num_topk: int
682
+):
683
+ """In-kernel compaction must, for every row, produce exactly the rank-owned
684
+ physical slots packed into [0, valid_count) with -1 in the tail -- the same
685
+ SET a reference filter + sort/gather produces. Uses wide rows (> BLOCK_N
686
+ valid slots), with interior -1 gaps."""
687
device = torch.device("cuda")
688
torch.manual_seed(7)
689
dcp_size = 2
690
block_size = 8
687
- num_topk = 1024 # > BLOCK_N(128) -> multiple tiles per row
691
num_rows = 5
692
max_blocks = 64
693
seq = max_blocks * block_size
@@ -698,7 +701,9 @@ def test_dcp_filter_compaction_matches_reference(interleave: int, dcp_rank: int)
701
(num_rows, num_topk), -1, dtype=torch.int32, device=device
702
)
703
for r in range(num_rows):
701
- n_valid = int(torch.randint(200, 600, (1,)).item())
704
+ # Ids are distinct, so a row holds at most `seq` valid entries.
705
+ hi = min(num_topk * 3 // 4, seq)
706
+ n_valid = int(torch.randint(num_topk // 4, hi, (1,)).item())
707
perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32)
708
token_indices[r, :n_valid] = perm
709
tests/v1/attention/test_sparse_mla_backends.py
+13
-9
@@ -1250,32 +1250,36 @@ def test_split_indexer_prefill_chunks_single_request_overflow():
1250
assert out == expected
1251
1252
1253
-def test_triton_convert_returns_valid_counts():
1253
+# 384 is not a power of two, so it counts via the tiled atomic accumulation
1254
+# rather than the single-tile path 128 takes.
1255
+@pytest.mark.parametrize("num_topk_tokens", [128, 384])
1256
+def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1257
"""Test that return_valid_counts correctly counts non-negative indices."""
1258
device = torch.device(DEVICE_TYPE)
1259
num_tokens = 8
1260
num_requests = 2
1261
max_blocks_per_req = 10
1262
block_size = 64
1260
- num_topk_tokens = 128
1263
1264
req_id = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device)
1265
block_table = torch.arange(
1266
num_requests * max_blocks_per_req, dtype=torch.int32, device=device
1267
).view(num_requests, max_blocks_per_req)
1268
1267
- # Create token indices with varying numbers of valid entries
1268
- # Token 0: 64 valid, 64 invalid (-1)
1269
- # Token 1: 32 valid, 96 invalid
1270
- # Token 2: 128 valid (all)
1271
- # Token 3: 1 valid, 127 invalid
1272
- # etc.
1269
+ # Create token indices with varying numbers of valid entries: half the row,
1270
+ # a quarter of it, the whole row, then a single valid entry -- twice over.
1271
token_indices = torch.full(
1272
(num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device
1273
)
1274
+ valid_counts_per_token = [
1275
+ num_topk_tokens // 2,
1276
+ num_topk_tokens // 4,
1277
+ num_topk_tokens,
1278
+ 1,
1279
+ ] * 2
1280
expected_valid = []
1281
for i in range(num_tokens):
1278
- num_valid = [64, 32, 128, 1, 64, 32, 128, 1][i]
1282
+ num_valid = valid_counts_per_token[i]
1283
token_indices[i, :num_valid] = torch.arange(
1284
num_valid, dtype=torch.int32, device=device
1285
) % (block_size * max_blocks_per_req)
vllm/v1/attention/backends/mla/sparse_utils.py
+61
-16
@@ -23,6 +23,9 @@ def _convert_req_index_to_global_index_kernel(
23
BLOCK_N: tl.constexpr, # tile width along columns
24
HAS_PREFILL: tl.constexpr,
25
COUNT_VALID: tl.constexpr, # whether to count valid indices
26
+ # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
27
+ # is an in-register reduction and needs no atomic.
28
+ SINGLE_TILE: tl.constexpr,
29
# When set, scatter valid slots to a contiguous prefix [0, valid_count) using
30
# valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
31
# -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(
102
# sum gives each valid lane a distinct local offset; one atomic add of the
103
# tile's valid count reserves a contiguous base across racing tiles. The
104
# out buffer is pre-filled with -1, so unwritten tail slots stay -1.
105
+ # With no racing tiles the base is 0 and the allocator becomes a store.
106
is_valid = (~is_invalid_tok).to(tl.int32)
107
local_offset = tl.cumsum(is_valid) - is_valid
108
tile_valid_count = tl.sum(is_valid)
105
- base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
109
+ if SINGLE_TILE:
110
+ base = 0
111
+ tl.store(valid_count_ptr + token_id, tile_valid_count)
112
+ else:
113
+ base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
114
dest = base + local_offset
115
out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
116
tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
@@ -111,10 +119,36 @@ def _convert_req_index_to_global_index_kernel(
119
out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
120
tl.store(out_ptr_ij, out_val)
121
114
- # Count valid indices in this tile and atomically add to row total
122
+ # Accumulate the tile's valid count into the row total; a single tile's
123
+ # reduction *is* the total.
124
if COUNT_VALID:
125
tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
117
- tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
126
+ if SINGLE_TILE:
127
+ tl.store(valid_count_ptr + token_id, tile_valid_count)
128
+ else:
129
+ tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
130
+
131
+
132
+def _remap_tiling(
133
+ NUM_TOPK_TOKENS: int, BLOCK_N: int, count_valid: bool
134
+) -> tuple[bool, int, int, int]:
135
+ """Pick the column tiling for the index remap kernel.
136
+
137
+ Counting the valid slots per row is the only reason the column tiles have to
138
+ talk to each other, so when counting give one program the whole row: the
139
+ count becomes an in-register reduction plus a plain store, needing neither
140
+ atomics nor a zero-initialized counter. The row is one ``tl.arange``, so this
141
+ needs a power-of-two width; other top-k sizes stay tiled and atomic.
142
+
143
+ Returns:
144
+ (single_tile, block_n, tiles_per_row, num_warps)
145
+ """
146
+ single_tile = (
147
+ count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS
148
+ )
149
+ if single_tile:
150
+ return True, NUM_TOPK_TOKENS, 1, 8
151
+ return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4
152
153
154
def triton_convert_req_index_to_global_index(
@@ -172,7 +206,10 @@ def triton_convert_req_index_to_global_index(
206
207
num_tokens = req_id.shape[0]
208
max_num_blocks_per_req = block_table.shape[1]
175
- tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N
209
+
210
+ single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
211
+ NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts
212
+ )
213
214
# Ensure contiguous tensors on the same device
215
req_id_c = req_id.contiguous()
@@ -180,12 +217,11 @@ def triton_convert_req_index_to_global_index(
217
token_indices_c = token_indices.contiguous()
218
out = torch.empty_like(token_indices_c)
219
183
- # Allocate valid count buffer if needed (must be zero-initialized for atomics)
220
valid_counts: torch.Tensor | None = None
221
if return_valid_counts:
186
- valid_counts = torch.zeros(
187
- num_tokens, dtype=torch.int32, device=token_indices.device
188
- )
222
+ # Zero-init only matters for the atomic accumulation path.
223
+ alloc = torch.empty if single_tile else torch.zeros
224
+ valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
225
226
# Strides in elements
227
bt_stride0, bt_stride1 = block_table_c.stride()
@@ -213,9 +249,10 @@ def triton_convert_req_index_to_global_index(
249
# shapes / constexprs
250
max_num_blocks_per_req,
251
BLOCK_SIZE,
216
- BLOCK_N,
252
+ block_n,
253
HAS_PREFILL_WORKSPACE,
254
return_valid_counts,
255
+ single_tile,
256
False, # COMPACT_TO_FRONT: keep input column == output column
257
# DCP disabled (no-op de-interleave)
258
1,
@@ -228,6 +265,7 @@ def triton_convert_req_index_to_global_index(
265
ti_stride1,
266
out_stride0,
267
out_stride1,
268
+ num_warps=num_warps,
269
)
270
271
if return_valid_counts:
@@ -287,15 +325,20 @@ def triton_filter_and_convert_dcp_index(
325
326
num_tokens = req_id.shape[0]
327
max_num_blocks_per_req = block_table.shape[1]
290
- tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N
328
329
req_id_c = req_id.contiguous()
330
block_table_c = block_table.contiguous()
331
token_indices_c = token_indices.contiguous()
332
296
- # The compaction uses the valid-count buffer as an atomic slot allocator, so
297
- # it requires counting. Pre-fill out with -1 so the unwritten tail stays -1.
333
+ # The compaction uses the valid-count buffer as a slot allocator, so it
334
+ # requires counting. Pre-fill out with -1 so the unwritten tail stays -1.
335
count_valid = return_valid_counts or compact_valid_to_front
336
+
337
+ # The compaction builds on the counting, so it shares the tiling.
338
+ single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
339
+ NUM_TOPK_TOKENS, BLOCK_N, count_valid
340
+ )
341
+
342
if compact_valid_to_front:
343
out = torch.full_like(token_indices_c, -1)
344
else:
@@ -303,9 +346,9 @@ def triton_filter_and_convert_dcp_index(
346
347
valid_counts: torch.Tensor | None = None
348
if count_valid:
306
- valid_counts = torch.zeros(
307
- num_tokens, dtype=torch.int32, device=token_indices.device
308
- )
349
+ # Zero-init only matters for the atomic accumulation path.
350
+ alloc = torch.empty if single_tile else torch.zeros
351
+ valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
352
353
bt_stride0, bt_stride1 = block_table_c.stride()
354
ti_stride0, ti_stride1 = token_indices_c.stride()
@@ -322,9 +365,10 @@ def triton_filter_and_convert_dcp_index(
365
None,
366
max_num_blocks_per_req,
367
BLOCK_SIZE,
325
- BLOCK_N,
368
+ block_n,
369
False, # HAS_PREFILL
370
count_valid,
371
+ single_tile,
372
compact_valid_to_front,
373
dcp_size,
374
dcp_rank,
@@ -335,6 +379,7 @@ def triton_filter_and_convert_dcp_index(
379
ti_stride1,
380
out_stride0,
381
out_stride1,
382
+ num_warps=num_warps,
383
)
384
385
if return_valid_counts: