@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c5200d3565

[Attention][DSA] support dcp for FLASHINFER_MLA_SPARSE (#46076)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Jingyi Yang <girasoleyang@gmail.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: GirasoleY <girasoleyang@gmail.com> Co-authored-by: Jingyi Yang <girasoleyang@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Jiangyun Zhu committed Jul 1, 2026 at 12:32 UTC c5200d3565e67fc4d50f0b91f2391234aa2400db
12 files changed +1996 -98
docs/design/attention_backends.md
+1 -1
@@ -222,7 +222,7 @@ MLA decode backends are selected using the standard
222 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
223 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
224 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
225 -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
225 +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
226 | `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
227 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
228 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
tests/v1/attention/test_indexer_dcp_localize.py new
+951
@@ -0,0 +1,951 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
8 +from vllm.platforms import current_platform
9 +from vllm.utils.import_utils import has_cutedsl
10 +from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata
11 +from vllm.v1.attention.backends.mla.sparse_utils import (
12 + triton_filter_and_convert_dcp_index,
13 +)
14 +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
15 +from vllm.v1.attention.ops.common import CPTritonContext, correct_attn_out
16 +
17 +
18 +def _local_count(length: int, rank: int, world: int, interleave: int) -> int:
19 + return sum(1 for pos in range(length) if (pos // interleave) % world == rank)
20 +
21 +
22 +def _global_to_local_indices(
23 + global_indices: torch.Tensor,
24 + rank: int,
25 + world: int,
26 + interleave: int,
27 +) -> torch.Tensor:
28 + valid = global_indices >= 0
29 + global_i64 = global_indices.to(torch.int64).clamp_min(0)
30 + owner = (global_i64 // interleave) % world
31 + local = (global_i64 // (world * interleave)) * interleave + global_i64 % interleave
32 + return torch.where(valid & (owner == rank), local, -1).to(torch.int64)
33 +
34 +
35 +def _local_to_global_indices(
36 + local_indices: torch.Tensor,
37 + rank: int,
38 + world: int,
39 + interleave: int,
40 +) -> torch.Tensor:
41 + valid = local_indices >= 0
42 + local = local_indices.to(torch.int64).clamp_min(0)
43 + global_indices = (
44 + (local // interleave) * (world * interleave)
45 + + rank * interleave
46 + + local % interleave
47 + )
48 + return torch.where(valid, global_indices, -1).to(torch.int64)
49 +
50 +
51 +def _ref_stable_topk_from_candidates_fp64(
52 + candidate_scores: torch.Tensor,
53 + candidate_token_ids: torch.Tensor,
54 + k: int,
55 +) -> torch.Tensor:
56 + """Pure-PyTorch reference for the CuteDSL stable-topk selector order
57 + (score desc, then lowest global token id). Selects the same SET as the
58 + kernel; only the set is compared in tests."""
59 + num_rows, num_candidates = candidate_scores.shape
60 + device = candidate_scores.device
61 + select_k = min(k, num_candidates)
62 + valid = candidate_token_ids >= 0
63 + bits = (
64 + candidate_scores.to(torch.float32).view(torch.int32).to(torch.int64)
65 + & 0xFFFFFFFF
66 + )
67 + sign = (bits >> 31) & 1
68 + score_key = (
69 + torch.where(sign.bool(), bits ^ 0xFFFFFFFF, bits ^ 0x80000000) & 0xFFFFFFFF
70 + )
71 + id_key = (~candidate_token_ids.to(torch.int64)) & 0xFFFFFFFF
72 + key = (score_key << 32) | id_key
73 + key = torch.where(valid, key, torch.zeros_like(key))
74 + topk_key = key ^ torch.iinfo(torch.int64).min
75 + _, topk_pos = topk_key.topk(select_k, dim=-1)
76 +
77 + selected = candidate_token_ids.gather(1, topk_pos).to(torch.int32)
78 + selected_valid = valid.gather(1, topk_pos)
79 + selected = torch.where(selected_valid, selected, selected.new_full((), -1))
80 + if select_k == k:
81 + return selected
82 + pad = torch.full((num_rows, k - select_k), -1, dtype=torch.int32, device=device)
83 + return torch.cat((selected, pad), dim=1)
84 +
85 +
86 +def _attention_from_indices(
87 + q: torch.Tensor,
88 + k: torch.Tensor,
89 + v: torch.Tensor,
90 + indices: torch.Tensor,
91 +) -> tuple[torch.Tensor, torch.Tensor]:
92 + valid = indices >= 0
93 + safe_indices = indices.clamp_min(0)
94 + selected_k = k[safe_indices]
95 + selected_v = v[safe_indices]
96 + scores = torch.einsum("td,tkd->tk", q, selected_k)
97 + scores = scores.masked_fill(~valid, float("-inf"))
98 + lse = torch.logsumexp(scores, dim=-1)
99 + probs = torch.softmax(scores, dim=-1).masked_fill(~valid, 0.0)
100 + out = torch.einsum("tk,tkd->td", probs, selected_v)
101 + empty_rows = ~valid.any(dim=-1)
102 + out[empty_rows] = 0
103 + lse[empty_rows] = float("-inf")
104 + return out, lse
105 +
106 +
107 +def _dcp_lse_merge(
108 + local_outs: list[torch.Tensor],
109 + local_lses: list[torch.Tensor],
110 +) -> tuple[torch.Tensor, torch.Tensor]:
111 + outs = torch.stack(local_outs, dim=0)
112 + lses = torch.stack(local_lses, dim=0)
113 + merged_lse = torch.logsumexp(lses, dim=0)
114 + weights = torch.exp(lses - merged_lse.unsqueeze(0))
115 + weights = torch.where(torch.isfinite(weights), weights, torch.zeros_like(weights))
116 + merged_out = (outs * weights.unsqueeze(-1)).sum(dim=0)
117 + return merged_out, merged_lse
118 +
119 +
120 +class _FakeDCPGroup:
121 + """Single-process stand-in: ``all_gather`` returns the pre-built
122 + concatenation of every rank's packed ``(score, global_id)`` candidates,
123 + mirroring the one packed all-gather the merge issues."""
124 +
125 + def __init__(self, gathered_packed: torch.Tensor) -> None:
126 + self.gathered_packed = gathered_packed
127 +
128 + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
129 + assert dim == 1
130 + return self.gathered_packed.clone()
131 +
132 +
133 +def _run_decode_topk(
134 + logits: torch.Tensor,
135 + seq_lens: torch.Tensor,
136 + next_n: int,
137 + topk: int,
138 +) -> torch.Tensor:
139 + indices = torch.empty(
140 + (logits.shape[0], topk), dtype=torch.int32, device=logits.device
141 + )
142 + torch.ops._C.top_k_per_row_decode(
143 + logits,
144 + next_n,
145 + seq_lens,
146 + indices,
147 + logits.shape[0],
148 + logits.stride(0),
149 + logits.stride(1),
150 + topk,
151 + )
152 + return indices
153 +
154 +
155 +def _run_persistent_topk(
156 + logits: torch.Tensor,
157 + seq_lens: torch.Tensor,
158 + topk: int,
159 + max_seq_len: int,
160 +) -> torch.Tensor:
161 + indices = torch.empty(
162 + (logits.shape[0], topk), dtype=torch.int32, device=logits.device
163 + )
164 + workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device=logits.device)
165 + torch.ops._C.persistent_topk(
166 + logits,
167 + seq_lens,
168 + indices,
169 + workspace,
170 + topk,
171 + max_seq_len,
172 + )
173 + return indices
174 +
175 +
176 +def _dcp_attention_from_local_topks(
177 + q: torch.Tensor,
178 + k: torch.Tensor,
179 + v: torch.Tensor,
180 + local_topks: list[torch.Tensor],
181 + world: int,
182 + interleave: int,
183 +) -> tuple[torch.Tensor, torch.Tensor]:
184 + local_outs = []
185 + local_lses = []
186 + for rank, local_topk in enumerate(local_topks):
187 + owned = [
188 + pos for pos in range(k.shape[0]) if (pos // interleave) % world == rank
189 + ]
190 + local_out, local_lse = _attention_from_indices(
191 + q, k[owned], v[owned], local_topk.to(torch.int64)
192 + )
193 + local_outs.append(local_out)
194 + local_lses.append(local_lse)
195 + return _dcp_lse_merge(local_outs, local_lses)
196 +
197 +
198 +def _merge_local_topks_global_with_fake_dcp(
199 + local_logits: list[torch.Tensor],
200 + local_topks: list[torch.Tensor],
201 + topk: int,
202 + world: int,
203 + interleave: int,
204 + row_starts: list[torch.Tensor | None] | None = None,
205 +) -> list[torch.Tensor]:
206 + """Run ``_merge_dcp_topk_global`` for every rank against a faked
207 + all-gather and return each rank's global-index result (all ranks should
208 + agree). The fake pre-builds the packed candidate concatenation the merge's
209 + single ``all_gather`` would return."""
210 + # The merge is now CuteDSL-only (no PyTorch fallback), so it runs the real
211 + # Triton pack + CuteDSL selector kernels even behind the faked all-gather.
212 + if not current_platform.is_cuda() or not has_cutedsl():
213 + pytest.skip("DCP merge requires CUDA and CuteDSL")
214 + packed_per_rank = []
215 + for rank, (logits, indices) in enumerate(zip(local_logits, local_topks)):
216 + score_indices = indices.clamp_min(0).to(torch.long)
217 + if row_starts is not None:
218 + rs = row_starts[rank]
219 + assert rs is not None
220 + score_indices = score_indices + rs.to(torch.long).view(-1, 1)
221 + if logits.shape[1] == 0:
222 + scores = torch.full_like(indices, float("-inf"), dtype=torch.float32)
223 + else:
224 + score_indices = score_indices.clamp_max(logits.shape[1] - 1)
225 + scores = logits.gather(1, score_indices).masked_fill(
226 + indices < 0, float("-inf")
227 + )
228 + global_ids = _local_to_global_indices(indices, rank, world, interleave)
229 + packed_per_rank.append(
230 + torch.stack((scores.float(), global_ids.to(torch.float32)), dim=-1)
231 + )
232 +
233 + fake_group = _FakeDCPGroup(torch.cat(packed_per_rank, dim=1).contiguous())
234 + original_get_dcp_group = sparse_indexer.get_dcp_group
235 + sparse_indexer.get_dcp_group = lambda: fake_group
236 + try:
237 + merged = []
238 + for rank, (logits, indices) in enumerate(zip(local_logits, local_topks)):
239 + rank_indices = indices.clone()
240 + result = sparse_indexer._merge_dcp_topk_global(
241 + logits,
242 + rank_indices,
243 + topk,
244 + rank,
245 + world,
246 + interleave,
247 + row_starts=None if row_starts is None else row_starts[rank],
248 + )
249 + assert result is None
250 + merged.append(rank_indices)
251 + return merged
252 + finally:
253 + sparse_indexer.get_dcp_group = original_get_dcp_group
254 +
255 +
256 +@pytest.mark.parametrize("world", [1, 2, 4])
257 +@pytest.mark.parametrize("interleave", [1, 2, 4])
258 +def test_get_dcp_local_seq_lens_matches_naive(world: int, interleave: int):
259 + seq_lens = torch.arange(0, 33, dtype=torch.int32)
260 +
261 + for rank in range(world):
262 + actual = get_dcp_local_seq_lens(seq_lens, world, rank, interleave)
263 + expected = torch.tensor(
264 + [
265 + _local_count(int(seq_len), rank, world, interleave)
266 + for seq_len in seq_lens
267 + ],
268 + dtype=torch.int32,
269 + )
270 + torch.testing.assert_close(actual, expected)
271 +
272 +
273 +def test_get_dcp_local_seq_lens_can_localize_per_token_bounds():
274 + seq_lens = torch.tensor([0, 1, 2, 3, 4, 7, 8, 17], dtype=torch.int32)
275 + world = 4
276 + interleave = 2
277 +
278 + for rank in range(world):
279 + actual = get_dcp_local_seq_lens(seq_lens, world, rank, interleave)
280 + expected = torch.tensor(
281 + [
282 + _local_count(int(seq_len), rank, world, interleave)
283 + for seq_len in seq_lens
284 + ],
285 + dtype=torch.int32,
286 + )
287 + torch.testing.assert_close(actual, expected)
288 +
289 +
290 +def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape():
291 + seq_lens = torch.tensor([[8, 9, 10], [11, 12, 13]], dtype=torch.int32)
292 + world = 2
293 + rank = 1
294 + interleave = 1
295 +
296 + actual = get_dcp_local_seq_lens(seq_lens, world, rank, interleave)
297 + expected = torch.tensor(
298 + [
299 + [_local_count(int(seq_len), rank, world, interleave) for seq_len in row]
300 + for row in seq_lens
301 + ],
302 + dtype=torch.int32,
303 + )
304 +
305 + assert actual.shape == seq_lens.shape
306 + torch.testing.assert_close(actual, expected)
307 +
308 +
309 +def test_get_dcp_local_seq_lens_must_run_after_decode_expansion():
310 + world = 2
311 + rank = 1
312 + interleave = 1
313 + expanded_bounds = torch.tensor([8, 9, 10], dtype=torch.int32)
314 +
315 + localized_after_expansion = get_dcp_local_seq_lens(
316 + expanded_bounds, world, rank, interleave
317 + )
318 + localized_request_len_minus_offsets = get_dcp_local_seq_lens(
319 + torch.tensor([10], dtype=torch.int32), world, rank
320 + ) - torch.tensor([2, 1, 0], dtype=torch.int32)
321 +
322 + assert not torch.equal(
323 + localized_after_expansion, localized_request_len_minus_offsets
324 + )
325 + torch.testing.assert_close(
326 + localized_after_expansion, torch.tensor([4, 4, 5], dtype=torch.int32)
327 + )
328 +
329 +
330 +@pytest.mark.parametrize("interleave", [1, 2])
331 +def test_sparse_dcp_attention_matches_global_topk_attention(interleave: int):
332 + torch.manual_seed(0)
333 + world = 2
334 + topk = 3
335 + num_queries = 4
336 + max_seq_len = 13
337 + head_dim = 8
338 +
339 + q = torch.randn(num_queries, head_dim)
340 + k = torch.randn(max_seq_len, head_dim)
341 + v = torch.randn(max_seq_len, head_dim)
342 + seq_lens = torch.tensor([6, 8, 11, 13], dtype=torch.int64)
343 +
344 + scores = q @ k.T
345 + global_topk = torch.full((num_queries, topk), -1, dtype=torch.int64)
346 + for row, seq_len in enumerate(seq_lens.tolist()):
347 + global_topk[row, :topk] = scores[row, :seq_len].topk(topk).indices
348 +
349 + ref_out, ref_lse = _attention_from_indices(q, k, v, global_topk)
350 +
351 + local_outs = []
352 + local_lses = []
353 + local_topks = []
354 + for rank in range(world):
355 + owned = [
356 + pos for pos in range(max_seq_len) if (pos // interleave) % world == rank
357 + ]
358 + k_local = k[owned]
359 + v_local = v[owned]
360 + local_topk = _global_to_local_indices(global_topk, rank, world, interleave)
361 + local_topks.append(local_topk)
362 + local_out, local_lse = _attention_from_indices(q, k_local, v_local, local_topk)
363 + local_outs.append(local_out)
364 + local_lses.append(local_lse)
365 +
366 + dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
367 +
368 + torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
369 + torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
370 +
371 + gathered_global = torch.cat(
372 + [
373 + _local_to_global_indices(local_topks[rank], rank, world, interleave)
374 + for rank in range(world)
375 + ],
376 + dim=1,
377 + )
378 + assert set(gathered_global[gathered_global >= 0].tolist()) == set(
379 + global_topk.flatten().tolist()
380 + )
381 +
382 +
383 +def test_local_topk_union_is_not_equivalent_to_global_topk_attention():
384 + world = 2
385 + interleave = 1
386 + topk = 2
387 + q = torch.tensor([[1.0]])
388 + k = torch.tensor(
389 + [
390 + [1.00],
391 + [0.90],
392 + [0.95],
393 + [0.85],
394 + ]
395 + )
396 + v = torch.tensor([[0.0], [1000.0], [0.0], [1000.0]])
397 +
398 + scores = q @ k.T
399 + global_topk = scores.topk(topk, dim=-1).indices
400 + ref_out, ref_lse = _attention_from_indices(q, k, v, global_topk)
401 +
402 + local_topks = []
403 + for rank in range(world):
404 + owned = [
405 + pos for pos in range(k.shape[0]) if (pos // interleave) % world == rank
406 + ]
407 + local_topks.append(scores[:, owned].topk(topk, dim=-1).indices)
408 +
409 + local_union_out, local_union_lse = _dcp_attention_from_local_topks(
410 + q, k, v, local_topks, world, interleave
411 + )
412 +
413 + assert not torch.allclose(local_union_out, ref_out)
414 + assert not torch.allclose(local_union_lse, ref_lse)
415 +
416 +
417 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
418 +def test_sparse_decode_dcp_persistent_topk_matches_non_dcp():
419 + torch.manual_seed(3)
420 + device = torch.device("cuda")
421 + world = 2
422 + interleave = 1
423 + topk = 512
424 + num_rows = 2
425 + max_seq_len = 1025
426 + head_dim = 16
427 +
428 + q = torch.randn(num_rows, head_dim, device=device)
429 + k = torch.randn(max_seq_len, head_dim, device=device)
430 + v = torch.randn(max_seq_len, head_dim, device=device)
431 + logits = q @ k.T
432 + seq_lens = torch.tensor([[1024], [1025]], dtype=torch.int32, device=device)
433 +
434 + non_dcp_topk = torch.empty((num_rows, topk), dtype=torch.int64, device=device)
435 + for row, seq_len in enumerate(seq_lens.flatten().tolist()):
436 + non_dcp_topk[row] = logits[row, :seq_len].topk(topk).indices
437 + ref_out, ref_lse = _attention_from_indices(q, k, v, non_dcp_topk)
438 +
439 + local_logits = []
440 + local_topks = []
441 + for rank in range(world):
442 + owned = [
443 + pos for pos in range(max_seq_len) if (pos // interleave) % world == rank
444 + ]
445 + rank_logits = logits[:, owned].contiguous()
446 + rank_seq_lens = get_dcp_local_seq_lens(
447 + seq_lens, world, rank, interleave
448 + ).contiguous()
449 + local_logits.append(rank_logits)
450 + local_topks.append(
451 + _run_persistent_topk(
452 + rank_logits,
453 + rank_seq_lens,
454 + topk,
455 + max_seq_len=rank_logits.shape[1],
456 + )
457 + )
458 +
459 + merged_global_topks = _merge_local_topks_global_with_fake_dcp(
460 + local_logits, local_topks, topk, world, interleave
461 + )
462 + # The radix top-K kernel selects a deterministic SET but writes it in
463 + # nondeterministic (atomicAdd) order; the production path is permutation-
464 + # invariant (compaction + softmax), so all ranks must agree on the set, not
465 + # the array order. (The fp64 fallback happens to return sorted order.)
466 + ref_topk = merged_global_topks[0]
467 + for rank_topk in merged_global_topks[1:]:
468 + for row in range(rank_topk.shape[0]):
469 + assert set(rank_topk[row].tolist()) == set(ref_topk[row].tolist())
470 +
471 + local_outs = []
472 + local_lses = []
473 + for rank, global_topk in enumerate(merged_global_topks):
474 + owned = [
475 + pos for pos in range(max_seq_len) if (pos // interleave) % world == rank
476 + ]
477 + local_topk = _global_to_local_indices(
478 + global_topk.to(torch.int64),
479 + rank,
480 + world,
481 + interleave,
482 + )
483 + local_out, local_lse = _attention_from_indices(
484 + q, k[owned], v[owned], local_topk
485 + )
486 + local_outs.append(local_out)
487 + local_lses.append(local_lse)
488 +
489 + dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
490 + torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
491 + torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
492 +
493 +
494 +@pytest.mark.skipif(
495 + not current_platform.is_cuda() or not has_cutedsl(),
496 + reason="This test requires CUDA and CuteDSL",
497 +)
498 +@pytest.mark.parametrize("use_row_starts", [False, True])
499 +def test_cutedsl_dcp_candidate_pack_and_select_matches_reference(
500 + use_row_starts: bool,
501 +):
502 + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
503 + pack_dcp_topk_candidates_cutedsl,
504 + stable_topk_from_gathered_candidates_cutedsl,
505 + )
506 +
507 + torch.manual_seed(13)
508 + device = torch.device("cuda")
509 + rows = 4
510 + valid_width = 1024
511 + width = valid_width + (8 if use_row_starts else 0)
512 + topk = 512
513 + world = 2
514 + row_starts = (
515 + torch.tensor([0, 2, 4, 1], device=device, dtype=torch.int32)
516 + if use_row_starts
517 + else None
518 + )
519 + row_offsets = (
520 + row_starts
521 + if row_starts is not None
522 + else torch.zeros(rows, device=device, dtype=torch.int32)
523 + )
524 +
525 + packed_by_rank = []
526 + for rank in range(world):
527 + logits = torch.randn((rows, width), device=device, dtype=torch.float32)
528 + local_topks = []
529 + for row in range(rows):
530 + start = int(row_offsets[row].item())
531 + local_topks.append(
532 + logits[row, start : start + valid_width].topk(topk).indices
533 + )
534 + topk_indices = torch.stack(local_topks).to(torch.int32)
535 +
536 + packed = torch.empty((rows, topk, 2), device=device, dtype=torch.float32)
537 + pack_dcp_topk_candidates_cutedsl(
538 + logits,
539 + topk_indices,
540 + packed,
541 + rank,
542 + world,
543 + 1,
544 + row_starts,
545 + )
546 +
547 + expected_scores = logits.gather(
548 + 1, topk_indices.to(torch.long) + row_offsets.to(torch.long).view(-1, 1)
549 + )
550 + expected_ids = (topk_indices * world + rank).to(torch.float32)
551 + torch.testing.assert_close(packed[..., 0], expected_scores)
552 + torch.testing.assert_close(packed[..., 1], expected_ids)
553 + packed_by_rank.append(packed)
554 +
555 + gathered = torch.cat(packed_by_rank, dim=1).contiguous()
556 + actual = torch.empty((rows, topk), device=device, dtype=torch.int32)
557 + returned = stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=actual)
558 + assert returned is actual
559 + expected = _ref_stable_topk_from_candidates_fp64(
560 + gathered[..., 0],
561 + gathered[..., 1].to(torch.int32),
562 + topk,
563 + )
564 +
565 + for row in range(rows):
566 + assert set(actual[row].cpu().tolist()) == set(expected[row].cpu().tolist())
567 +
568 +
569 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
570 +def test_sparse_prefill_dcp_metadata_localizes_causal_bounds():
571 + device = torch.device("cuda")
572 + seq_len = 8
573 +
574 + query_start_loc = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
575 + query_start_loc_cpu = torch.tensor([0, seq_len], dtype=torch.int32)
576 + seq_lens = torch.tensor([seq_len], dtype=torch.int32, device=device)
577 + seq_lens_cpu = torch.tensor([seq_len], dtype=torch.int32)
578 + block_table = torch.zeros((1, 1), dtype=torch.int32, device=device)
579 +
580 + def build(dcp_world_size, dcp_rank, interleave=1):
581 + chunk = build_prefill_chunk_metadata(
582 + start_idx=0,
583 + end_idx=1,
584 + query_start_loc=query_start_loc,
585 + query_start_loc_cpu=query_start_loc_cpu,
586 + uncompressed_seq_lens=seq_lens,
587 + compressed_seq_lens=seq_lens,
588 + compressed_seq_lens_cpu=seq_lens_cpu,
589 + block_table=block_table,
590 + compress_ratio=1,
591 + dcp_rank=dcp_rank,
592 + dcp_world_size=dcp_world_size,
593 + cp_kv_cache_interleave_size=interleave,
594 + )
595 + assert chunk is not None
596 + torch.accelerator.synchronize()
597 + return chunk
598 +
599 + # Non-DCP: local_cu_seq_lens aliases the global cu_seq_lens, and
600 + # cu_seqlen_ks/ke carry the global causal bounds.
601 + chunk = build(dcp_world_size=1, dcp_rank=0)
602 + assert chunk.local_cu_seq_lens is chunk.cu_seq_lens
603 + torch.testing.assert_close(
604 + chunk.cu_seqlen_ks.cpu(),
605 + torch.zeros(seq_len, dtype=torch.int32),
606 + )
607 + torch.testing.assert_close(
608 + chunk.cu_seqlen_ke.cpu(),
609 + torch.arange(1, seq_len + 1, dtype=torch.int32),
610 + )
611 +
612 + # DCP: cu_seqlen_ks/ke are localized in place to this rank's shard.
613 + chunk = build(dcp_world_size=4, dcp_rank=0)
614 + assert chunk.local_cu_seq_lens is not None
615 + torch.testing.assert_close(
616 + chunk.local_cu_seq_lens.cpu(),
617 + torch.tensor([0, 2], dtype=torch.int32),
618 + )
619 + torch.testing.assert_close(
620 + chunk.cu_seqlen_ks.cpu(),
621 + torch.zeros(seq_len, dtype=torch.int32),
622 + )
623 + torch.testing.assert_close(
624 + chunk.cu_seqlen_ke.cpu(),
625 + torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int32),
626 + )
627 +
628 + # DCP with interleave=2: per-token causal bounds localize differently from
629 + # interleave=1 (groups of 2 consecutive tokens are owned together). For
630 + # world=4, rank=0, K=2, per-token global len L=1..8 -> local len
631 + # [1,2,2,2,2,2,2,2] (matches get_dcp_local_seq_lens).
632 + chunk = build(dcp_world_size=4, dcp_rank=0, interleave=2)
633 + assert chunk.local_cu_seq_lens is not None
634 + torch.testing.assert_close(
635 + chunk.cu_seqlen_ks.cpu(),
636 + torch.zeros(seq_len, dtype=torch.int32),
637 + )
638 + torch.testing.assert_close(
639 + chunk.cu_seqlen_ke.cpu(),
640 + torch.tensor([1, 2, 2, 2, 2, 2, 2, 2], dtype=torch.int32),
641 + )
642 +
643 +
644 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
645 +def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():
646 + block_size = 4
647 + num_topk = 128
648 + dcp_size = 2
649 + req_id = torch.zeros(1, dtype=torch.int32, device="cuda")
650 + token_indices = torch.full((1, num_topk), -1, dtype=torch.int32, device="cuda")
651 + token_indices[0, :8] = torch.arange(8, dtype=torch.int32, device="cuda")
652 + block_table = torch.tensor([[10]], dtype=torch.int32, device="cuda")
653 +
654 + out, valid_counts = triton_filter_and_convert_dcp_index(
655 + req_id,
656 + block_table,
657 + token_indices,
658 + dcp_size=dcp_size,
659 + dcp_rank=0,
660 + BLOCK_SIZE=block_size,
661 + NUM_TOPK_TOKENS=num_topk,
662 + return_valid_counts=True,
663 + )
664 +
665 + valid = int(valid_counts.item())
666 + assert valid == 4
667 + assert (out[0, :valid] >= 0).all()
668 + assert (out[0, valid:] == -1).all()
669 + # In-kernel compaction packs valid slots to the front; prefix order is
670 + # unspecified, so compare as a set.
671 + assert set(out[0, :valid].cpu().tolist()) == {40, 41, 42, 43}
672 +
673 +
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."""
683 + device = torch.device("cuda")
684 + torch.manual_seed(7)
685 + dcp_size = 2
686 + block_size = 8
687 + num_topk = 1024 # > BLOCK_N(128) -> multiple tiles per row
688 + num_rows = 5
689 + max_blocks = 64
690 + seq = max_blocks * block_size
691 +
692 + req_id = torch.randint(0, 3, (num_rows,), dtype=torch.int32, device=device)
693 + block_table = torch.randint(
694 + 0, 1000, (3, max_blocks), dtype=torch.int32, device=device
695 + )
696 + # Each row: a dense valid prefix of distinct global token ids, then -1 pad.
697 + token_indices = torch.full(
698 + (num_rows, num_topk), -1, dtype=torch.int32, device=device
699 + )
700 + for r in range(num_rows):
701 + n_valid = int(torch.randint(200, 600, (1,)).item())
702 + perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32)
703 + token_indices[r, :n_valid] = perm
704 +
705 + out, valid_counts = triton_filter_and_convert_dcp_index(
706 + req_id,
707 + block_table,
708 + token_indices,
709 + dcp_size=dcp_size,
710 + dcp_rank=dcp_rank,
711 + cp_kv_cache_interleave_size=interleave,
712 + BLOCK_SIZE=block_size,
713 + NUM_TOPK_TOKENS=num_topk,
714 + return_valid_counts=True,
715 + )
716 +
717 + for r in range(num_rows):
718 + toks = token_indices[r]
719 + toks = toks[toks >= 0]
720 + owner = (toks // interleave) % dcp_size
721 + owned = toks[owner == dcp_rank]
722 + local = (owned // (dcp_size * interleave)) * interleave + owned % interleave
723 + blk = local // block_size
724 + off = local % block_size
725 + expected = (
726 + block_table[int(req_id[r].item()), blk].to(torch.int64) * block_size + off
727 + )
728 + n = int(valid_counts[r].item())
729 + assert n == owned.numel()
730 + assert (out[r, :n] >= 0).all()
731 + assert (out[r, n:] == -1).all()
732 + assert set(out[r, :n].cpu().tolist()) == set(expected.cpu().tolist())
733 +
734 +
735 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
736 +@pytest.mark.parametrize("interleave", [1, 2, 4])
737 +def test_dcp_global_topk_physical_attention_matches_non_dcp(interleave: int):
738 + torch.manual_seed(2)
739 + device = torch.device("cuda")
740 + dcp_size = 2
741 + block_size = 4
742 + num_topk = 128
743 + selected_k = 8
744 + seq_len = 16
745 + head_dim = 16
746 + num_queries = 3
747 +
748 + q = torch.randn(num_queries, head_dim, device=device)
749 + k_global = torch.randn(seq_len, head_dim, device=device)
750 + v_global = torch.randn(seq_len, head_dim, device=device)
751 + scores = q @ k_global.T
752 + global_topk = torch.full(
753 + (num_queries, num_topk), -1, dtype=torch.int32, device=device
754 + )
755 + global_topk[:, :selected_k] = scores.topk(selected_k, dim=-1).indices.to(
756 + torch.int32
757 + )
758 +
759 + ref_out, ref_lse = _attention_from_indices(
760 + q, k_global, v_global, global_topk[:, :selected_k].to(torch.int64)
761 + )
762 +
763 + local_outs = []
764 + local_lses = []
765 + for rank in range(dcp_size):
766 + block_ids = torch.tensor(
767 + [[rank * 10 + 1, rank * 10 + 2]], dtype=torch.int32, device=device
768 + )
769 + num_slots = int((block_ids.max().item() + 1) * block_size)
770 + k_cache = torch.zeros(num_slots, head_dim, device=device)
771 + v_cache = torch.zeros(num_slots, head_dim, device=device)
772 + for global_idx in range(seq_len):
773 + if (global_idx // interleave) % dcp_size != rank:
774 + continue
775 + local_idx = (
776 + global_idx // (dcp_size * interleave)
777 + ) * interleave + global_idx % interleave
778 + block = local_idx // block_size
779 + offset = local_idx % block_size
780 + slot = int(block_ids[0, block].item()) * block_size + offset
781 + k_cache[slot] = k_global[global_idx]
782 + v_cache[slot] = v_global[global_idx]
783 +
784 + slots, valid_counts = triton_filter_and_convert_dcp_index(
785 + torch.zeros(num_queries, dtype=torch.int32, device=device),
786 + block_ids,
787 + global_topk,
788 + dcp_size=dcp_size,
789 + dcp_rank=rank,
790 + cp_kv_cache_interleave_size=interleave,
791 + BLOCK_SIZE=block_size,
792 + NUM_TOPK_TOKENS=num_topk,
793 + return_valid_counts=True,
794 + )
795 + row_ids = torch.arange(num_queries, device=device)
796 + assert (slots[row_ids, valid_counts] == -1).all()
797 + local_out, local_lse = _attention_from_indices(
798 + q, k_cache, v_cache, slots.to(torch.int64)
799 + )
800 + local_outs.append(local_out)
801 + local_lses.append(local_lse)
802 +
803 + dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
804 + torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
805 + torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
806 +
807 +
808 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
809 +@pytest.mark.parametrize("is_lse_base_on_e", [True, False])
810 +def test_correct_attn_out_zeroes_empty_nan_partial(is_lse_base_on_e: bool):
811 + out = torch.full((1, 1, 4), float("nan"), device="cuda")
812 + lses = torch.tensor(
813 + [[[0.0]], [[float("-inf")]]],
814 + dtype=torch.float32,
815 + device="cuda",
816 + )
817 +
818 + corrected, final_lse = correct_attn_out(
819 + out,
820 + lses,
821 + cp_rank=1,
822 + ctx=CPTritonContext(),
823 + is_lse_base_on_e=is_lse_base_on_e,
824 + )
825 + torch.accelerator.synchronize()
826 +
827 + torch.testing.assert_close(corrected, torch.zeros_like(corrected))
828 + torch.testing.assert_close(final_lse, torch.zeros_like(final_lse))
829 +
830 +
831 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
832 +def test_decode_topk_pads_surplus_with_negative_one():
833 + """When a row's valid length < topk the top-k kernel must pad the surplus
834 + slots with -1: the DCP merge (`topk_indices >= 0`) and
835 + `triton_filter_and_convert_dcp_index` (`tok < 0`) both treat <0 as invalid,
836 + so a non-(-1) pad would be silently attended. This is the common case under
837 + DCP, where each rank's local seq_len = global / world is usually << topk.
838 + Checked on the *set* of valid indices (the kernel may order them by score,
839 + not position)."""
840 + torch.manual_seed(0)
841 + device = torch.device("cuda")
842 + topk = 8
843 + next_n = 1
844 + seq_lens_list = [0, 3, 5, 8]
845 + num_rows = len(seq_lens_list)
846 + logits = torch.randn(num_rows, 16, device=device)
847 + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device).view(
848 + num_rows, next_n
849 + )
850 + idx = _run_decode_topk(logits, seq_lens, next_n, topk)
851 + for r, sl in enumerate(seq_lens_list):
852 + valid = idx[r][idx[r] >= 0]
853 + # seq_len <= topk, so top-k selects exactly the whole valid range.
854 + assert valid.numel() == min(sl, topk)
855 + assert set(valid.tolist()) == set(range(sl))
856 + assert (idx[r] == -1).sum().item() == topk - min(sl, topk)
857 +
858 +
859 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
860 +def test_persistent_topk_pads_surplus_with_negative_one():
861 + """Same surplus=-1 invariant for the persistent_topk kernel (k>=512)."""
862 + torch.manual_seed(0)
863 + device = torch.device("cuda")
864 + topk = 512 # persistent_topk requires k in {512, 1024, 2048}
865 + seq_lens_list = [100, 300, 512]
866 + num_rows = len(seq_lens_list)
867 + max_seq_len = 600
868 + logits = torch.randn(num_rows, max_seq_len, device=device)
869 + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device).view(
870 + num_rows, 1
871 + )
872 + idx = _run_persistent_topk(logits, seq_lens, topk, max_seq_len)
873 + for r, sl in enumerate(seq_lens_list):
874 + valid = idx[r][idx[r] >= 0]
875 + assert valid.numel() == min(sl, topk)
876 + assert set(valid.tolist()) == set(range(sl))
877 + assert (idx[r] == -1).sum().item() == topk - min(sl, topk)
878 +
879 +
880 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
881 +def test_sparse_decode_dcp_short_context_matches_non_dcp():
882 + """End-to-end DCP decode where the global seq_len < topk (so every rank's
883 + local top-k is surplus-padded). Exercises the kernel surplus -> merge mask
884 + -> global top-k -> physical localize -> LSE merge chain for the common
885 + short-context decode case, vs the non-DCP reference."""
886 + torch.manual_seed(4)
887 + device = torch.device("cuda")
888 + world = 2
889 + interleave = 1
890 + topk = 512
891 + num_rows = 2
892 + max_seq_len = 300 # < topk -> surplus everywhere
893 + head_dim = 16
894 +
895 + q = torch.randn(num_rows, head_dim, device=device)
896 + k = torch.randn(max_seq_len, head_dim, device=device)
897 + v = torch.randn(max_seq_len, head_dim, device=device)
898 + logits = q @ k.T
899 + seq_lens = torch.tensor([[250], [300]], dtype=torch.int32, device=device)
900 +
901 + non_dcp_topk = torch.empty((num_rows, topk), dtype=torch.int64, device=device)
902 + for row, seq_len in enumerate(seq_lens.flatten().tolist()):
903 + sel = logits[row, :seq_len].topk(min(topk, seq_len)).indices
904 + non_dcp_topk[row, : sel.numel()] = sel
905 + non_dcp_topk[row, sel.numel() :] = -1
906 + ref_out, ref_lse = _attention_from_indices(q, k, v, non_dcp_topk)
907 +
908 + local_logits = []
909 + local_topks = []
910 + for rank in range(world):
911 + owned = [p for p in range(max_seq_len) if (p // interleave) % world == rank]
912 + rank_logits = logits[:, owned].contiguous()
913 + rank_seq_lens = get_dcp_local_seq_lens(seq_lens, world, rank, interleave)
914 + local_logits.append(rank_logits)
915 + local_topks.append(
916 + _run_persistent_topk(
917 + rank_logits,
918 + rank_seq_lens.contiguous(),
919 + topk,
920 + max_seq_len=rank_logits.shape[1],
921 + )
922 + )
923 +
924 + merged_global_topks = _merge_local_topks_global_with_fake_dcp(
925 + local_logits, local_topks, topk, world, interleave
926 + )
927 + # The radix top-K kernel selects a deterministic SET but writes it in
928 + # nondeterministic (atomicAdd) order; the production path is permutation-
929 + # invariant (compaction + softmax), so all ranks must agree on the set, not
930 + # the array order. (The fp64 fallback happens to return sorted order.)
931 + ref_topk = merged_global_topks[0]
932 + for rank_topk in merged_global_topks[1:]:
933 + for row in range(rank_topk.shape[0]):
934 + assert set(rank_topk[row].tolist()) == set(ref_topk[row].tolist())
935 +
936 + local_outs = []
937 + local_lses = []
938 + for rank, global_topk in enumerate(merged_global_topks):
939 + owned = [p for p in range(max_seq_len) if (p // interleave) % world == rank]
940 + local_topk = _global_to_local_indices(
941 + global_topk.to(torch.int64), rank, world, interleave
942 + )
943 + local_out, local_lse = _attention_from_indices(
944 + q, k[owned], v[owned], local_topk
945 + )
946 + local_outs.append(local_out)
947 + local_lses.append(local_lse)
948 +
949 + dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
950 + torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
951 + torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
vllm/model_executor/kernels/attention/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/kernels/attention/dsa/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py new
+420
@@ -0,0 +1,420 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from functools import cache
5 +
6 +import cutlass
7 +import cutlass.cute as cute
8 +import torch
9 +from cuda.bindings.driver import CUstream
10 +from cutlass import Float32, Int32, Uint32, Uint64
11 +from quack.compile_utils import make_fake_tensor
12 +
13 +from vllm.cute_utils import recast_val
14 +from vllm.triton_utils import tl, triton
15 +
16 +
17 +def stable_topk_from_gathered_candidates_cutedsl(
18 + gathered: torch.Tensor,
19 + topk: int,
20 + out: torch.Tensor | None = None,
21 +) -> torch.Tensor:
22 + if out is None:
23 + out = torch.empty(
24 + (gathered.shape[0], topk),
25 + dtype=torch.int32,
26 + device=gathered.device,
27 + )
28 + StableTopKFromGatheredCandidatesKernel.compile(topk, gathered.shape[1])(
29 + gathered, out
30 + )
31 + return out
32 +
33 +
34 +def pack_dcp_topk_candidates_cutedsl(
35 + logits: torch.Tensor,
36 + topk_indices: torch.Tensor,
37 + packed: torch.Tensor,
38 + dcp_rank: int,
39 + dcp_world_size: int,
40 + cp_interleave: int,
41 + row_starts: torch.Tensor | None,
42 +) -> None:
43 + topk = topk_indices.shape[1]
44 + grid = (topk_indices.shape[0], triton.cdiv(topk, 512))
45 + row_starts_arg = row_starts if row_starts is not None else topk_indices
46 + _pack_dcp_topk_candidates_triton_kernel[grid](
47 + logits,
48 + topk_indices,
49 + packed,
50 + row_starts_arg,
51 + logits.stride(0),
52 + logits.stride(1),
53 + topk_indices.stride(0),
54 + topk_indices.stride(1),
55 + packed.stride(0),
56 + packed.stride(1),
57 + packed.stride(2),
58 + logits.shape[1],
59 + DCP_RANK=dcp_rank,
60 + DCP_WORLD_SIZE=dcp_world_size,
61 + CP_INTERLEAVE=cp_interleave,
62 + HAS_ROW_STARTS=row_starts is not None,
63 + TOPK=topk,
64 + BLOCK_SIZE=512,
65 + num_warps=8,
66 + )
67 +
68 +
69 +@triton.jit
70 +def _pack_dcp_topk_candidates_triton_kernel(
71 + logits,
72 + topk_indices,
73 + packed,
74 + row_starts,
75 + logits_stride0: tl.constexpr,
76 + logits_stride1: tl.constexpr,
77 + topk_stride0: tl.constexpr,
78 + topk_stride1: tl.constexpr,
79 + packed_stride0: tl.constexpr,
80 + packed_stride1: tl.constexpr,
81 + packed_stride2: tl.constexpr,
82 + num_cols,
83 + DCP_RANK: tl.constexpr,
84 + DCP_WORLD_SIZE: tl.constexpr,
85 + CP_INTERLEAVE: tl.constexpr,
86 + HAS_ROW_STARTS: tl.constexpr,
87 + TOPK: tl.constexpr,
88 + BLOCK_SIZE: tl.constexpr,
89 +):
90 + row = tl.program_id(0)
91 + tile = tl.program_id(1)
92 + cols = tile * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
93 + mask = cols < TOPK
94 +
95 + local_idx = tl.load(
96 + topk_indices + row * topk_stride0 + cols * topk_stride1,
97 + mask=mask,
98 + other=-1,
99 + )
100 + valid = local_idx >= 0
101 + safe_local_idx = tl.maximum(local_idx, 0)
102 +
103 + row_start = 0
104 + if HAS_ROW_STARTS:
105 + row_start = tl.load(row_starts + row)
106 +
107 + score_col = safe_local_idx + row_start
108 + score_col = tl.minimum(score_col, tl.maximum(num_cols - 1, 0))
109 + score = tl.load(
110 + logits + row * logits_stride0 + score_col * logits_stride1,
111 + mask=mask & valid,
112 + other=-float("inf"),
113 + )
114 +
115 + global_id = (
116 + (safe_local_idx // CP_INTERLEAVE) * (DCP_WORLD_SIZE * CP_INTERLEAVE)
117 + + DCP_RANK * CP_INTERLEAVE
118 + + safe_local_idx % CP_INTERLEAVE
119 + )
120 + global_id = tl.where(valid, global_id, -1)
121 +
122 + packed_base = packed + row * packed_stride0 + cols * packed_stride1
123 + tl.store(packed_base, score, mask=mask)
124 + tl.store(packed_base + packed_stride2, global_id.to(tl.float32), mask=mask)
125 +
126 +
127 +@cute.jit
128 +def _warp_scan_inclusive_i32(val: Int32, lane: Int32) -> Int32:
129 + for i in cutlass.range_constexpr(cute.arch.WARP_SIZE.bit_length() - 1):
130 + offset = 1 << i
131 + partial = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0)
132 + if lane >= offset:
133 + val += partial
134 + return val
135 +
136 +
137 +@cute.jit
138 +def _block_scan_inclusive_i32(
139 + val: Int32,
140 + lane: Int32,
141 + warp_id: Int32,
142 + warp_scratch: cute.Tensor,
143 + warps_per_block: int,
144 +) -> Int32:
145 + prefix = _warp_scan_inclusive_i32(val, lane)
146 + if lane == Int32(cute.arch.WARP_SIZE - 1):
147 + warp_scratch[0, warp_id] = prefix
148 + cute.arch.sync_threads()
149 +
150 + if warp_id == Int32(0):
151 + warp_total = Int32(0)
152 + if lane < Int32(warps_per_block):
153 + warp_total = warp_scratch[0, lane]
154 + warp_prefix = _warp_scan_inclusive_i32(warp_total, lane)
155 + if lane < Int32(warps_per_block):
156 + warp_scratch[0, lane] = warp_prefix - warp_total
157 + cute.arch.sync_threads()
158 +
159 + return prefix + warp_scratch[0, warp_id]
160 +
161 +
162 +class StableTopKFromGatheredCandidatesKernel:
163 + tb_size = 512
164 + hist_bins = 2048
165 + radix_bits = (hist_bins - 1).bit_length()
166 + assert hist_bins == 1 << radix_bits
167 + key_bits = Uint64.width
168 + radix_passes = (key_bits + radix_bits - 1) // radix_bits
169 + final_radix_bits = key_bits - radix_bits * (radix_passes - 1)
170 + hist_chunks = (hist_bins + tb_size - 1) // tb_size
171 + warps_per_block = tb_size // cute.arch.WARP_SIZE
172 +
173 + def __init__(self, topk: int, num_candidates: int):
174 + assert num_candidates % self.tb_size == 0, (
175 + "StableTopKFromGatheredCandidatesKernel requires candidate count "
176 + f"to be a multiple of {self.tb_size}, got {num_candidates}"
177 + )
178 + self.topk = topk
179 + self.keys_per_thread = num_candidates // self.tb_size
180 +
181 + @cute.struct
182 + class SharedStorage:
183 + hist: cute.struct.MemRange[Int32, self.hist_bins]
184 + committed_count: cute.struct.MemRange[Int32, 1]
185 + running_count: cute.struct.MemRange[Int32, 1]
186 + threshold_bin: cute.struct.MemRange[Int32, 1]
187 + threshold_found: cute.struct.MemRange[Int32, 1]
188 + include_threshold_bin: cute.struct.MemRange[Int32, 1]
189 + prefix_s: cute.struct.Align[cute.struct.MemRange[Uint64, 1], 8]
190 + warp_totals: cute.struct.MemRange[Int32, self.warps_per_block]
191 +
192 + self.shared_storage = SharedStorage
193 +
194 + @cute.jit
195 + def __call__(
196 + self,
197 + gathered: cute.Tensor,
198 + out: cute.Tensor,
199 + stream: CUstream,
200 + ):
201 + grid = (gathered.shape[0], 1, 1)
202 + self.kernel(gathered, out).launch(
203 + grid=grid,
204 + block=(self.tb_size, 1, 1),
205 + stream=stream,
206 + )
207 +
208 + @cute.jit
209 + def _stable_key(self, score: Float32, token_id: Int32) -> Uint64:
210 + bits = recast_val(score, Uint32)
211 + mask = Uint32(0x80000000)
212 + if (bits & Uint32(0x80000000)) != Uint32(0):
213 + mask = Uint32(0xFFFFFFFF)
214 + score_key = Uint64(bits ^ mask) << Uint64(32)
215 + id_key = Uint64(~Uint32(token_id))
216 + key = score_key | id_key
217 + if token_id < Int32(0):
218 + key = Uint64(0)
219 + return key
220 +
221 + @cute.jit
222 + def _prefix_matches(
223 + self,
224 + key: Uint64,
225 + prefix: Uint64,
226 + prefix_bits: Int32,
227 + ):
228 + matches = prefix_bits == Int32(0)
229 + if prefix_bits != Int32(0):
230 + shift = Int32(self.key_bits) - prefix_bits
231 + matches = (key >> Uint64(shift)) == (prefix >> Uint64(shift))
232 + return matches
233 +
234 + @cute.jit
235 + def _radix_pass(
236 + self,
237 + keys: cute.Tensor,
238 + output: cute.Tensor,
239 + storage,
240 + tid: Int32,
241 + step: Int32,
242 + bits: int,
243 + is_final_pass: bool,
244 + ):
245 + hist_smem = storage.hist.get_tensor(cute.make_layout((self.hist_bins,)))
246 + committed_count_smem = storage.committed_count.data_ptr()
247 + running_count_smem = storage.running_count.data_ptr()
248 + threshold_bin_smem = storage.threshold_bin.data_ptr()
249 + threshold_found_smem = storage.threshold_found.data_ptr()
250 + include_threshold_bin_smem = storage.include_threshold_bin.data_ptr()
251 + prefix_smem = storage.prefix_s.data_ptr()
252 + warp_totals_smem = storage.warp_totals.get_tensor(
253 + cute.make_layout((1, self.warps_per_block))
254 + )
255 +
256 + prefix_bits = step * Int32(self.radix_bits)
257 + num_bins = 1 << bits
258 + block_scan_iterations = (num_bins + self.tb_size - 1) // self.tb_size
259 + shift = Int32(self.key_bits) - prefix_bits - Int32(bits)
260 + bin_mask = Uint64(num_bins - 1)
261 + prefix = prefix_smem.load()
262 +
263 + for chunk in cutlass.range_constexpr(self.hist_chunks):
264 + hist_smem[tid + Int32(chunk * self.tb_size)] = Int32(0)
265 + if tid == Int32(0):
266 + running_count_smem.store(committed_count_smem.load())
267 + include_threshold_bin_smem.store(Int32(0))
268 + threshold_found_smem.store(Int32(0))
269 + cute.arch.sync_threads()
270 +
271 + for key_idx in cutlass.range_constexpr(self.keys_per_thread):
272 + key = keys[key_idx]
273 + if self._prefix_matches(key, prefix, prefix_bits):
274 + bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
275 + cute.arch.atomic_add(
276 + hist_smem.iterator + bin_idx,
277 + Int32(1),
278 + sem="relaxed",
279 + scope="cta",
280 + )
281 + cute.arch.sync_threads()
282 +
283 + lane = cute.arch.lane_idx()
284 + warp_id = cute.arch.warp_idx()
285 + # Each iteration scans one tb_size-wide slice of bins, high to low.
286 + iter = Int32(0)
287 + threshold_found = threshold_found_smem.load()
288 + while threshold_found == Int32(0) and iter < Int32(block_scan_iterations):
289 + bin_idx = Int32(num_bins - 1) - (iter * Int32(self.tb_size) + tid)
290 + count = hist_smem[bin_idx]
291 + chunk_inclusive = _block_scan_inclusive_i32(
292 + count,
293 + lane,
294 + warp_id,
295 + warp_totals_smem,
296 + self.warps_per_block,
297 + )
298 + running_count = running_count_smem.load()
299 + prior_in_scan_slice = chunk_inclusive - count
300 + remaining = Int32(self.topk) - running_count - prior_in_scan_slice
301 + if count > Int32(0) and remaining > Int32(0) and remaining <= count:
302 + threshold_bin_smem.store(bin_idx)
303 + if count <= remaining or cutlass.const_expr(is_final_pass):
304 + include_threshold_bin_smem.store(Int32(1))
305 + threshold_found_smem.store(Int32(1))
306 + # Barrier: every thread must finish reading running_count for this
307 + # slice before tb_size-1 advances it, else a warp racing ahead to
308 + # the store makes a lagging thread double-count the slice total
309 + # (-> remaining too small -> threshold too high -> under-fill).
310 + cute.arch.sync_threads()
311 + if tid == Int32(self.tb_size - 1):
312 + running_count_smem.store(running_count + chunk_inclusive)
313 + cute.arch.sync_threads()
314 +
315 + threshold_found = threshold_found_smem.load()
316 + iter += Int32(1)
317 +
318 + threshold = threshold_bin_smem.load()
319 + should_include_threshold = include_threshold_bin_smem.load() != Int32(0)
320 + for key_idx in cutlass.range_constexpr(self.keys_per_thread):
321 + key = keys[key_idx]
322 + if self._prefix_matches(key, prefix, prefix_bits):
323 + bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
324 + selected = bin_idx > threshold
325 + if should_include_threshold:
326 + selected = selected or bin_idx == threshold
327 + if selected:
328 + dst = cute.arch.atomic_add(
329 + committed_count_smem,
330 + Int32(1),
331 + sem="relaxed",
332 + scope="cta",
333 + )
334 + if dst < Int32(self.topk):
335 + output[dst] = recast_val(~Uint32(key), Int32)
336 + cute.arch.sync_threads()
337 +
338 + pass_finished = include_threshold_bin_smem.load()
339 + if tid == Int32(0) and pass_finished == Int32(0):
340 + prefix_smem.store(prefix | (Uint64(threshold) << Uint64(shift)))
341 + cute.arch.sync_threads()
342 + return pass_finished
343 +
344 + @cute.kernel
345 + def kernel(
346 + self,
347 + input: cute.Tensor,
348 + out: cute.Tensor,
349 + ):
350 + row, _, _ = cute.arch.block_idx()
351 + tid, _, _ = cute.arch.thread_idx()
352 + input_row = input[row, None, None]
353 + output_row = out[row, None]
354 + keys = cute.make_rmem_tensor((self.keys_per_thread,), Uint64)
355 +
356 + smem = cutlass.utils.SmemAllocator()
357 + storage = smem.allocate(self.shared_storage, 8)
358 + committed_count_smem = storage.committed_count.data_ptr()
359 + prefix_smem = storage.prefix_s.data_ptr()
360 + for i in range(tid, self.topk, self.tb_size):
361 + output_row[i] = Int32(-1)
362 +
363 + for key_idx in cutlass.range_constexpr(self.keys_per_thread):
364 + col = tid + Int32(key_idx * self.tb_size)
365 + score = Float32(input_row[col, 0])
366 + token_id = Int32(input_row[col, 1])
367 + keys[key_idx] = self._stable_key(score, token_id)
368 +
369 + if tid == Int32(0):
370 + committed_count_smem.store(Int32(0))
371 + prefix_smem.store(Uint64(0))
372 + cute.arch.sync_threads()
373 +
374 + step = Int32(0)
375 + finished = Int32(0)
376 + while finished == Int32(0) and step < Int32(self.radix_passes - 1):
377 + finished = self._radix_pass(
378 + keys,
379 + output_row,
380 + storage,
381 + tid,
382 + step,
383 + self.radix_bits,
384 + False,
385 + )
386 + step += Int32(1)
387 +
388 + if finished == Int32(0):
389 + self._radix_pass(
390 + keys,
391 + output_row,
392 + storage,
393 + tid,
394 + Int32(self.radix_passes - 1),
395 + self.final_radix_bits,
396 + True,
397 + )
398 +
399 + @cache
400 + @staticmethod
401 + def compile(topk: int, num_candidates: int):
402 + num_rows = cute.sym_int()
403 +
404 + gathered = cute.runtime.make_fake_tensor(
405 + Float32,
406 + (num_rows, num_candidates, 2),
407 + stride=(cute.sym_int64(divisibility=2), 2, 1),
408 + assumed_align=8,
409 + )
410 + out = make_fake_tensor(Int32, (num_rows, topk), divisibility=1)
411 +
412 + kernel = StableTopKFromGatheredCandidatesKernel(topk, num_candidates)
413 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
414 + return cute.compile(
415 + kernel,
416 + gathered,
417 + out,
418 + stream,
419 + options="--enable-tvm-ffi",
420 + )
vllm/model_executor/layers/sparse_attn_indexer.py
+167 -44
@@ -8,6 +8,8 @@ import vllm.envs as envs
8 from vllm import _custom_ops as ops
9 from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11 +from vllm.config import get_current_vllm_config
12 +from vllm.distributed import get_dcp_group
13 from vllm.forward_context import get_forward_context
14 from vllm.logger import init_logger
15 from vllm.model_executor.custom_op import CustomOp
@@ -21,6 +23,7 @@ from vllm.utils.deep_gemm import (
23 fp8_fp4_paged_mqa_logits,
24 has_deep_gemm,
25 )
26 +from vllm.utils.import_utils import has_cutedsl
27 from vllm.utils.torch_utils import (
28 LayerNameType,
29 _encode_layer_name,
@@ -41,6 +44,85 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
44 MXFP4_BLOCK_SIZE = 32
45
46
47 +def _assert_cutedsl_dcp_merge_supported(
48 + logits: torch.Tensor,
49 + topk_indices: torch.Tensor,
50 + k: int,
51 +) -> None:
52 + # The DCP merge only supports the CuteDSL path (Triton pack kernel + CuteDSL
53 + # stable-topk selector); there is no PyTorch fallback. The first cut targets
54 + # Blackwell/Hopper with index_topk in (512, 1024, 2048) (the selector's radix
55 + # sizing); the Triton pack itself has no shape/topk constraints.
56 + if not has_cutedsl():
57 + raise RuntimeError(
58 + "DCP sparse-indexer merge requires CuteDSL; install it or disable DCP."
59 + )
60 + if logits.device.type != "cuda":
61 + raise RuntimeError("DCP sparse-indexer merge requires CUDA tensors.")
62 + if logits.dtype != torch.float32 or topk_indices.dtype != torch.int32:
63 + raise RuntimeError(
64 + "DCP sparse-indexer merge requires fp32 logits and int32 indices."
65 + )
66 + if k not in (512, 1024, 2048):
67 + raise RuntimeError(
68 + f"DCP sparse-indexer merge requires index_topk in (512, 1024, 2048); "
69 + f"got {k}."
70 + )
71 +
72 +
73 +def _merge_dcp_topk_global(
74 + logits: torch.Tensor,
75 + topk_indices: torch.Tensor,
76 + topk_tokens: int,
77 + dcp_rank: int,
78 + dcp_world_size: int,
79 + cp_interleave: int,
80 + row_starts: torch.Tensor | None = None,
81 +) -> None:
82 + """Merge each DCP rank's local top-K into the global top-K.
83 +
84 + ``topk_indices`` are this rank's local top-K positions into its 1/N KV
85 + shard. A token in the global top-K must also be in its owning rank's local
86 + top-K (at most ``topk_tokens - 1`` tokens rank globally above it, hence at
87 + most that many on its own rank), so exchanging only the per-rank local
88 + candidates is exact -- equivalent to all-gathering the full logit matrix,
89 + but it ships ``dcp_world_size * topk_tokens`` candidates instead of the whole
90 + score row. Overwrites ``topk_indices`` with global token ids (``-1`` for
91 + padding); the attention backend localizes them back to physical slots per
92 + rank.
93 + """
94 + if dcp_world_size <= 1:
95 + return
96 +
97 + # CuteDSL-only path (no PyTorch fallback): Triton-pack each rank's
98 + # (score, global_id) candidates on-device, all-gather, then the CuteDSL
99 + # stable-topk selector.
100 + _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens)
101 + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import (
102 + pack_dcp_topk_candidates_cutedsl,
103 + stable_topk_from_gathered_candidates_cutedsl,
104 + )
105 +
106 + packed = torch.empty(
107 + (*topk_indices.shape, 2),
108 + dtype=torch.float32,
109 + device=topk_indices.device,
110 + )
111 + pack_dcp_topk_candidates_cutedsl(
112 + logits,
113 + topk_indices,
114 + packed,
115 + dcp_rank,
116 + dcp_world_size,
117 + cp_interleave,
118 + row_starts,
119 + )
120 + gathered = get_dcp_group().all_gather(packed, dim=1)
121 + stable_topk_from_gathered_candidates_cutedsl(
122 + gathered, topk_tokens, out=topk_indices
123 + )
124 +
125 +
126 @triton.jit
127 def _fused_indexer_q_rope_quant_kernel(
128 positions,
@@ -227,6 +309,9 @@ def sparse_attn_indexer(
309 topk_indices_buffer: torch.Tensor,
310 skip_k_cache_insert: bool,
311 use_fp4_cache: bool = False,
312 + dcp_rank: int = 0,
313 + dcp_world_size: int = 1,
314 + cp_kv_cache_interleave_size: int = 1,
315 skip_topk_buffer_clear: bool = False,
316 ) -> torch.Tensor:
317 # careful! this will be None in dummy run
@@ -328,16 +413,18 @@ def sparse_attn_indexer(
413 scales_spec,
414 )
415 for chunk in prefill_metadata.chunks:
331 - k_quant = k_quant_full[: chunk.total_seq_lens]
332 - k_scale = k_scale_full[: chunk.total_seq_lens]
333 -
334 - if not chunk.skip_kv_gather:
416 + cu_seqlen_ks = chunk.cu_seqlen_ks
417 + cu_seqlen_ke = chunk.cu_seqlen_ke
418 + assert chunk.local_cu_seq_lens is not None
419 + k_quant = k_quant_full[: chunk.max_local_total_seq_lens]
420 + k_scale = k_scale_full[: chunk.max_local_total_seq_lens]
421 + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0:
422 ops.cp_gather_indexer_k_quant_cache(
423 kv_cache,
424 k_quant,
425 k_scale,
426 chunk.block_table,
340 - chunk.cu_seq_lens,
427 + chunk.local_cu_seq_lens,
428 )
429
430 q_slice = q_quant[chunk.token_start : chunk.token_end]
@@ -346,51 +433,64 @@ def sparse_attn_indexer(
433 if q_scale is not None
434 else None
435 )
349 - # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
350 - # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
351 - if use_fp4_cache:
352 - q_slice_cast = q_slice.view(torch.int8)
353 - k_quant_cast = k_quant.view(torch.int8)
354 - k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
355 - else:
356 - q_slice_cast = q_slice
357 - k_quant_cast = k_quant
358 - k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
359 - if current_platform.is_xpu():
360 - if q_scale_slice is not None:
361 - raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
362 - logits = torch.ops.vllm.xpu_fp8_mqa_logits(
363 - q_slice_cast,
364 - k_quant_cast,
365 - k_scale_cast,
366 - weights[chunk.token_start : chunk.token_end],
367 - chunk.cu_seqlen_ks,
368 - chunk.cu_seqlen_ke,
369 - )
370 - else:
371 - logits = fp8_fp4_mqa_logits(
372 - (q_slice_cast, q_scale_slice),
373 - (k_quant_cast, k_scale_cast),
374 - weights[chunk.token_start : chunk.token_end],
375 - chunk.cu_seqlen_ks,
376 - chunk.cu_seqlen_ke,
377 - clean_logits=False,
378 - )
379 - num_rows = logits.shape[0]
380 -
436 topk_indices = topk_indices_buffer[
437 chunk.token_start : chunk.token_end, :topk_tokens
438 ]
439
385 - ops.top_k_per_row_prefill(
440 + if chunk.local_total_seq_lens == 0:
441 + logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
442 + topk_indices.fill_(-1)
443 + else:
444 + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
445 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
446 + if use_fp4_cache:
447 + q_slice_cast = q_slice.view(torch.int8)
448 + k_quant_cast = k_quant.view(torch.int8)
449 + k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
450 + else:
451 + q_slice_cast = q_slice
452 + k_quant_cast = k_quant
453 + k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
454 + if current_platform.is_xpu():
455 + if q_scale_slice is not None:
456 + raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
457 + logits = torch.ops.vllm.xpu_fp8_mqa_logits(
458 + q_slice_cast,
459 + k_quant_cast,
460 + k_scale_cast,
461 + weights[chunk.token_start : chunk.token_end],
462 + cu_seqlen_ks,
463 + cu_seqlen_ke,
464 + )
465 + else:
466 + logits = fp8_fp4_mqa_logits(
467 + (q_slice_cast, q_scale_slice),
468 + (k_quant_cast, k_scale_cast),
469 + weights[chunk.token_start : chunk.token_end],
470 + cu_seqlen_ks,
471 + cu_seqlen_ke,
472 + clean_logits=False,
473 + )
474 + num_rows = logits.shape[0]
475 + ops.top_k_per_row_prefill(
476 + logits,
477 + cu_seqlen_ks,
478 + cu_seqlen_ke,
479 + topk_indices,
480 + num_rows,
481 + logits.stride(0),
482 + logits.stride(1),
483 + topk_tokens,
484 + )
485 +
486 + _merge_dcp_topk_global(
487 logits,
387 - chunk.cu_seqlen_ks,
388 - chunk.cu_seqlen_ke,
488 topk_indices,
390 - num_rows,
391 - logits.stride(0),
392 - logits.stride(1),
489 topk_tokens,
490 + dcp_rank,
491 + dcp_world_size,
492 + cp_kv_cache_interleave_size,
493 + row_starts=chunk.cu_seqlen_ks,
494 )
495
496 if has_decode:
@@ -508,7 +608,7 @@ def sparse_attn_indexer(
608 topk_indices,
609 topk_workspace,
610 topk_tokens,
511 - attn_metadata_narrowed.max_seq_len,
611 + logits.shape[1],
612 )
613 else:
614 ops.top_k_per_row_decode(
@@ -522,6 +622,16 @@ def sparse_attn_indexer(
622 topk_tokens,
623 )
624
625 + if decode_metadata.global_seq_lens is not None:
626 + _merge_dcp_topk_global(
627 + logits,
628 + topk_indices,
629 + topk_tokens,
630 + dcp_rank,
631 + dcp_world_size,
632 + cp_kv_cache_interleave_size,
633 + )
634 +
635 if decode_metadata.requires_padding:
636 # if padded, we need to unpack
637 # the topk indices removing padded tokens
@@ -553,6 +663,9 @@ def sparse_attn_indexer_fake(
663 topk_indices_buffer: torch.Tensor | None,
664 skip_k_cache_insert: bool,
665 use_fp4_cache: bool = False,
666 + dcp_rank: int = 0,
667 + dcp_world_size: int = 1,
668 + cp_kv_cache_interleave_size: int = 1,
669 skip_topk_buffer_clear: bool = False,
670 ) -> torch.Tensor:
671 return topk_indices_buffer
@@ -604,6 +717,13 @@ class SparseAttnIndexer(CustomOp):
717 self.topk_indices_buffer = topk_indices_buffer
718 self.skip_k_cache_insert = skip_k_cache_insert
719 self.use_fp4_cache = use_fp4_cache
720 + # DCP scalars are constant for the run; resolve them here (config is set
721 + # during model construction) and pass them into the custom op, rather
722 + # than threading them through per-step metadata.
723 + parallel_config = get_current_vllm_config().parallel_config
724 + self.dcp_world_size = parallel_config.decode_context_parallel_size
725 + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
726 + self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
727 if current_platform.is_cuda() and not has_deep_gemm():
728 raise RuntimeError(
729 "Sparse Attention Indexer CUDA op requires DeepGEMM support in "
@@ -657,6 +777,9 @@ class SparseAttnIndexer(CustomOp):
777 self.topk_indices_buffer,
778 self.skip_k_cache_insert,
779 self.use_fp4_cache,
780 + self.dcp_rank,
781 + self.dcp_world_size,
782 + self.cp_kv_cache_interleave_size,
783 )
784
785 def forward_xpu(
vllm/models/deepseek_v32/nvidia/attention.py
+2 -1
@@ -496,7 +496,8 @@ class DeepseekV32Attention(MLAAttention):
496 self.topk_indices_buffer,
497 True, # skip_k_cache_insert
498 False, # use_fp4_cache
499 - True, # skip_topk_buffer_clear (fused_norm_rope already did it)
499 + # fused_norm_rope already cleared the topk buffer this forward.
500 + skip_topk_buffer_clear=True,
501 )
502
503 if attn_metadata is None:
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+104 -17
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, ClassVar
8 import numpy as np
9 import torch
10
11 +from vllm import envs
12 from vllm.config import VllmConfig
13 from vllm.config.cache import CacheDType
14 from vllm.logger import init_logger
@@ -29,8 +30,12 @@ from vllm.v1.attention.backend import (
30 )
31 from vllm.v1.attention.backends.mla.sparse_utils import (
32 triton_convert_req_index_to_global_index,
33 + triton_filter_and_convert_dcp_index,
34 +)
35 +from vllm.v1.attention.backends.utils import (
36 + KVCacheLayoutType,
37 + split_decodes_and_prefills,
38 )
33 -from vllm.v1.attention.backends.utils import KVCacheLayoutType
39 from vllm.v1.kv_cache_interface import AttentionSpec
40
41 if TYPE_CHECKING:
@@ -38,8 +43,6 @@ if TYPE_CHECKING:
43
44 logger = init_logger(__name__)
45
41 -FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
42 -
46
47 class _FlashInferMLASparseBackendBase(AttentionBackend):
48 """Common metadata for concrete FlashInfer sparse MLA backends."""
@@ -252,10 +255,13 @@ class FlashInferMLASparseMetadata(AttentionMetadata):
255
256 # Sequence lengths for all requests (context + query)
257 seq_lens: torch.Tensor
258 + num_decodes: int
259 + num_decode_tokens: int
260
261 # Sparse-specific
262 block_size: int = 64
263 topk_tokens: int = 2048
264 + cp_kv_cache_interleave_size: int = 1
265
266
267 class FlashInferMLASparseMetadataBuilder(
@@ -281,6 +287,12 @@ class FlashInferMLASparseMetadataBuilder(
287 self.mla_dims = get_mla_dims(self.model_config)
288 self.topk_tokens = vllm_config.model_config.hf_config.index_topk
289
290 + self._init_reorder_batch_threshold(
291 + 1,
292 + supports_spec_as_decode=True,
293 + supports_dcp_with_varlen=True,
294 + )
295 +
296 self.req_id_per_token_buffer = torch.empty(
297 (vllm_config.scheduler_config.max_num_batched_tokens,),
298 dtype=torch.int32,
@@ -295,6 +307,12 @@ class FlashInferMLASparseMetadataBuilder(
307 ) -> FlashInferMLASparseMetadata:
308 cm = common_attn_metadata
309 num_tokens = cm.num_actual_tokens
310 + assert self.reorder_batch_threshold is not None
311 + num_decodes, _, num_decode_tokens, _ = split_decodes_and_prefills(
312 + cm,
313 + decode_threshold=self.reorder_batch_threshold,
314 + treat_short_extends_as_decodes=True,
315 + )
316
317 # Build req_id_per_token mapping
318 starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
@@ -320,8 +338,13 @@ class FlashInferMLASparseMetadataBuilder(
338 block_table=cm.block_table_tensor,
339 req_id_per_token=req_id_per_token_tensor,
340 seq_lens=cm.seq_lens,
341 + num_decodes=num_decodes,
342 + num_decode_tokens=num_decode_tokens,
343 block_size=self.kv_cache_spec.block_size,
344 topk_tokens=self.topk_tokens,
345 + cp_kv_cache_interleave_size=(
346 + self.vllm_config.parallel_config.cp_kv_cache_interleave_size
347 + ),
348 )
349
350
@@ -333,7 +356,7 @@ def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
356 global _fi_sparse_workspace
357 if _fi_sparse_workspace is None:
358 _fi_sparse_workspace = torch.zeros(
336 - FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE,
359 + envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
360 dtype=torch.uint8,
361 device=device,
362 )
@@ -347,6 +370,9 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
370 sparse attention computation.
371 """
372
373 + can_return_lse_for_decode: bool = True
374 + lse_base_on_e: bool = False
375 +
376 def __init__(
377 self,
378 num_heads: int,
@@ -421,14 +447,27 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
447 assert self.topk_indices_buffer is not None
448 topk_indices = self.topk_indices_buffer[:num_actual_toks]
449
424 - topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
425 - attn_metadata.req_id_per_token[:num_actual_toks],
426 - attn_metadata.block_table,
427 - topk_indices,
428 - BLOCK_SIZE=attn_metadata.block_size,
429 - NUM_TOPK_TOKENS=topk_indices.shape[1],
430 - return_valid_counts=True,
431 - )
450 + if self.dcp_world_size > 1:
451 + topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index(
452 + attn_metadata.req_id_per_token[:num_actual_toks],
453 + attn_metadata.block_table,
454 + topk_indices,
455 + dcp_size=self.dcp_world_size,
456 + dcp_rank=self.dcp_rank,
457 + cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size),
458 + BLOCK_SIZE=attn_metadata.block_size,
459 + NUM_TOPK_TOKENS=topk_indices.shape[1],
460 + return_valid_counts=True,
461 + )
462 + else:
463 + topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
464 + attn_metadata.req_id_per_token[:num_actual_toks],
465 + attn_metadata.block_table,
466 + topk_indices,
467 + BLOCK_SIZE=attn_metadata.block_size,
468 + NUM_TOPK_TOKENS=topk_indices.shape[1],
469 + return_valid_counts=True,
470 + )
471
472 if self._workspace_buffer is None:
473 self._workspace_buffer = _get_workspace_buffer(q.device)
@@ -444,18 +483,66 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
483
484 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
485
447 - o = trtllm_batch_decode_with_kv_cache_mla(
448 - query=q.unsqueeze(1),
486 + # Single-token sparse decode. trtllm-gen requires the q_len_per_request
487 + # dim, but the sparse attention mask is fully per-token (each query token
488 + # carries its own top-k index row), so unsqueeze is sufficient and
489 + # correct. The MTP/multi-token q_len grouping is a perf-only layout and is
490 + # deferred until MTP is validated end-to-end for this backend.
491 + query = q.unsqueeze(1)
492 + block_tables = topk_indices_physical.unsqueeze(1)
493 + seq_lens_arg = seq_lens
494 +
495 + kernel_out = trtllm_batch_decode_with_kv_cache_mla(
496 + query=query,
497 kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
498 workspace_buffer=self._workspace_buffer,
499 qk_nope_head_dim=self.qk_nope_head_dim,
500 kv_lora_rank=self.kv_lora_rank,
501 qk_rope_head_dim=self.qk_rope_head_dim,
454 - block_tables=topk_indices_physical.unsqueeze(1),
455 - seq_lens=seq_lens,
502 + block_tables=block_tables,
503 + seq_lens=seq_lens_arg,
504 max_seq_len=attn_metadata.topk_tokens,
505 bmm1_scale=self.bmm1_scale,
506 bmm2_scale=self.bmm2_scale,
507 sparse_mla_top_k=attn_metadata.topk_tokens,
508 + return_lse=self.need_to_return_lse_for_decode,
509 )
461 - return o.view(-1, o.shape[-2], o.shape[-1]), None
510 + if self.need_to_return_lse_for_decode:
511 + assert isinstance(kernel_out, tuple)
512 + o, lse = kernel_out
513 + else:
514 + assert isinstance(kernel_out, torch.Tensor)
515 + o = kernel_out
516 + lse = None
517 +
518 + out = o.view(-1, o.shape[-2], o.shape[-1])
519 + if lse is not None:
520 + lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
521 + empty_rows = (topk_indices_physical == -1).all(dim=-1)
522 + out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
523 + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
524 + return out, lse
525 +
526 + @staticmethod
527 + def _normalize_lse(
528 + lse: torch.Tensor,
529 + num_tokens: int,
530 + num_heads: int,
531 + ) -> torch.Tensor:
532 + # FlashInfer returns the decode LSE either as 2D (num_tokens, num_heads)
533 + # or 3D ((num_tokens, num_heads, 1) / (num_tokens, 1, num_heads)).
534 + # Collapse all of these to the (num_tokens, num_heads) the shared DCP
535 + # reducer expects.
536 + if lse.dim() == 3:
537 + if lse.shape[-1] == 1:
538 + lse = lse.squeeze(-1)
539 + elif lse.shape[1] == 1:
540 + lse = lse.squeeze(1)
541 + elif lse.shape[0] * lse.shape[1] == num_tokens:
542 + lse = lse.reshape(num_tokens, lse.shape[-1])
543 + if lse.shape != (num_tokens, num_heads):
544 + raise RuntimeError(
545 + "Unexpected FlashInfer sparse MLA LSE shape: "
546 + f"{tuple(lse.shape)}, expected ({num_tokens}, {num_heads})."
547 + )
548 + return lse
vllm/v1/attention/backends/mla/indexer.py
+176 -12
@@ -6,6 +6,7 @@ import torch
6
7 import vllm.envs as envs
8 from vllm.config import VllmConfig
9 +from vllm.distributed import get_dcp_group
10 from vllm.logger import init_logger
11 from vllm.platforms import current_platform
12 from vllm.triton_utils import tl, triton
@@ -24,6 +25,7 @@ from vllm.v1.attention.backend import (
25 )
26 from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
27 from vllm.v1.attention.backends.utils import (
28 + get_dcp_local_seq_lens,
29 split_decodes_and_prefills,
30 )
31 from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
@@ -168,6 +170,8 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
170 @dataclass
171 class DeepseekV32IndexerPrefillChunkMetadata:
172 block_table: torch.Tensor
173 + # Under DCP (dcp_world_size > 1) these hold this rank's local row bounds;
174 + # otherwise they hold the global bounds.
175 cu_seqlen_ks: torch.Tensor
176 cu_seqlen_ke: torch.Tensor
177 cu_seq_lens: torch.Tensor
@@ -177,6 +181,9 @@ class DeepseekV32IndexerPrefillChunkMetadata:
181 token_end: int
182 num_reqs: int
183 skip_kv_gather: bool = False
184 + local_cu_seq_lens: torch.Tensor | None = None
185 + local_total_seq_lens: int = 0
186 + max_local_total_seq_lens: int = 0
187
188
189 @dataclass
@@ -195,6 +202,7 @@ class DeepSeekV32IndexerDecodeMetadata:
202 decode_lens: torch.Tensor
203 requires_padding: bool
204 schedule_metadata: torch.Tensor
205 + global_seq_lens: torch.Tensor | None = None
206
207
208 @dataclass
@@ -243,6 +251,19 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
251 def __init__(self, *args, **kwargs):
252 super().__init__(*args, **kwargs)
253 scheduler_config = self.vllm_config.scheduler_config
254 + parallel_config = self.vllm_config.parallel_config
255 + self.dcp_world_size = parallel_config.decode_context_parallel_size
256 + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
257 + self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
258 + # The DCP sparse-indexer code is parameterized by interleave size, but
259 + # interleave > 1 is not yet validated end-to-end (gsm8k parity fails),
260 + # so fail closed here rather than silently produce wrong output.
261 + if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size > 1:
262 + raise NotImplementedError(
263 + "DCP sparse indexer currently supports only "
264 + f"cp_kv_cache_interleave_size=1 (got "
265 + f"{self.cp_kv_cache_interleave_size})."
266 + )
267 # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
268 self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
269 self.num_speculative_tokens = (
@@ -300,6 +321,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
321 dtype=torch.int32,
322 device=self.device,
323 )
324 + self.global_decode_seq_lens_buffer = torch.zeros(
325 + (scheduler_config.max_num_batched_tokens,),
326 + dtype=torch.int32,
327 + device=self.device,
328 + )
329 self.arange_buffer = torch.arange(
330 max(
331 scheduler_config.max_num_seqs * next_n,
@@ -331,6 +357,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
357 # Get compress_ratio for DeepseekV4 support
358 if isinstance(self.kv_cache_spec, MLAAttentionSpec):
359 self.compress_ratio = self.kv_cache_spec.compress_ratio
360 + if self.dcp_world_size > 1 and self.compress_ratio > 1:
361 + raise NotImplementedError(
362 + "DCP is not supported with sparse indexer KV compression "
363 + f"(compress_ratio={self.compress_ratio})."
364 + )
365
366 # Pre-allocate buffers for CUDA graph compatibility when
367 if self.compress_ratio > 1:
@@ -348,6 +379,26 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
379 device=self.device,
380 )
381
382 + def _dcp_localize_decode_seq_lens(
383 + self,
384 + seq_lens: torch.Tensor,
385 + num_decodes: int,
386 + seq_lens_is_buffer_view: bool,
387 + ) -> torch.Tensor:
388 + local_seq_lens = get_dcp_local_seq_lens(
389 + seq_lens,
390 + self.dcp_world_size,
391 + self.dcp_rank,
392 + self.cp_kv_cache_interleave_size,
393 + )
394 + if seq_lens_is_buffer_view:
395 + seq_lens.copy_(local_seq_lens)
396 + return seq_lens
397 +
398 + out = self.decode_seq_lens_buffer[:num_decodes]
399 + out.copy_(local_seq_lens)
400 + return out
401 +
402 def _prepare_decode_tensors(
403 self,
404 seq_lens: torch.Tensor,
@@ -468,6 +519,34 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
519 seq_lens = seq_lens_buffer
520 return seq_lens, block_table, decode_lens, num_decodes, requires_padding
521
522 + def _prepare_global_decode_seq_lens(
523 + self,
524 + global_seq_lens: torch.Tensor | None,
525 + decode_lens: torch.Tensor,
526 + decode_lens_cpu: torch.Tensor,
527 + query_start_loc: torch.Tensor,
528 + num_decode_tokens: int,
529 + use_native: bool,
530 + max_decode_len: int,
531 + ) -> torch.Tensor | None:
532 + if global_seq_lens is None:
533 + return None
534 + if use_native or max_decode_len <= 1:
535 + return global_seq_lens
536 +
537 + actual_expanded = int(decode_lens_cpu.sum().item())
538 + if actual_expanded > 0:
539 + expanded_offsets = torch.repeat_interleave(
540 + global_seq_lens - decode_lens - query_start_loc,
541 + decode_lens,
542 + output_size=actual_expanded,
543 + )
544 + self.global_decode_seq_lens_buffer[:actual_expanded] = (
545 + expanded_offsets + self.arange_buffer[:actual_expanded] + 1
546 + )
547 + self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
548 + return self.global_decode_seq_lens_buffer[:num_decode_tokens]
549 +
550 def build(
551 self,
552 common_prefix_len: int,
@@ -481,6 +560,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
560 seq_lens = common_attn_metadata.seq_lens
561 slot_mapping = common_attn_metadata.slot_mapping
562 block_table = common_attn_metadata.block_table_tensor
563 + dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
564
565 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
566 split_decodes_and_prefills(
@@ -549,6 +629,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
629 self.compress_ratio,
630 query_slice=query_slice,
631 skip_kv_gather=query_slice.start > 0,
632 + dcp_rank=self.dcp_rank,
633 + dcp_world_size=self.dcp_world_size,
634 + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
635 )
636 # Skip when total_seq_lens is 0 (i.e., no compressed token).
637 if metadata is not None:
@@ -566,6 +649,17 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
649 common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
650 )
651
652 + # Under DCP the per-token decode bounds must be localized AFTER the
653 + # per-token expansion below, not before. Expanding from a
654 + # request-level localized length subtracts decode offsets in local
655 + # space and yields too-short bounds (e.g. world=2, rank=1, global
656 + # per-token bounds [8, 9, 10] -> [3, 4, 5] instead of [4, 4, 5]), so
657 + # the first decode token would run top-k against too short a local KV
658 + # range and miss valid tokens. Keep the global seq_lens here and
659 + # localize the expanded bounds further down.
660 + global_seq_lens_for_decode: torch.Tensor | None = None
661 + if dcp_local_seq_lens is not None:
662 + global_seq_lens_for_decode = common_attn_metadata.seq_lens[:num_decodes]
663 seq_lens = common_attn_metadata.seq_lens[:num_decodes]
664 block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
665
@@ -573,6 +667,16 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
667 next_n = 1 + self.num_speculative_tokens
668 use_native = not self.use_flattening and max_decode_len <= next_n
669
670 + global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
671 + global_seq_lens=global_seq_lens_for_decode,
672 + decode_lens=decode_lens,
673 + decode_lens_cpu=decode_lens_cpu,
674 + query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
675 + num_decode_tokens=num_decode_tokens,
676 + use_native=use_native,
677 + max_decode_len=max_decode_len,
678 + )
679 +
680 seq_lens, block_table, decode_lens, batch_size, requires_padding = (
681 self._prepare_decode_tensors(
682 seq_lens=seq_lens,
@@ -588,15 +692,22 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
692 )
693 )
694
695 + seq_lens_is_buffer_view = (use_native and next_n > 1) or (
696 + not use_native and max_decode_len > 1
697 + )
698 +
699 + # DCP: localize the now-expanded per-token global bounds to this
700 + # rank's owned KV. Done here (after expansion) so each token's global
701 + # causal length is localized individually; see the comment above.
702 + if dcp_local_seq_lens is not None:
703 + seq_lens = self._dcp_localize_decode_seq_lens(
704 + seq_lens, num_decodes, seq_lens_is_buffer_view
705 + )
706 +
707 # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
708 # compressed tokens. Convert uncompressed seq_lens to compressed.
709 if self.compress_ratio > 1:
594 - # True iff seq_lens aliases decode_seq_lens_buffer (flatten or
595 - # native wrote it); False iff it aliases common_attn_metadata.
596 - seq_lens_is_local_view = (use_native and next_n > 1) or (
597 - not use_native and max_decode_len > 1
598 - )
599 - if seq_lens_is_local_view:
710 + if seq_lens_is_buffer_view:
711 seq_lens //= self.compress_ratio
712 else:
713 # Copy to avoid mutating shared state; keeps CG address stable.
@@ -626,6 +737,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
737 decode_lens=decode_lens,
738 requires_padding=requires_padding,
739 schedule_metadata=self.scheduler_metadata_buffer,
740 + global_seq_lens=global_seq_lens_for_decode,
741 )
742
743 attn_metadata = DeepseekV32IndexerMetadata(
@@ -655,6 +767,9 @@ def build_prefill_chunk_metadata(
767 compress_ratio: int,
768 query_slice: slice | None = None,
769 skip_kv_gather: bool = False,
770 + dcp_rank: int = 0,
771 + dcp_world_size: int = 1,
772 + cp_kv_cache_interleave_size: int = 1,
773 ) -> DeepseekV32IndexerPrefillChunkMetadata | None:
774 total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
775 if total_seq_lens == 0:
@@ -669,6 +784,25 @@ def build_prefill_chunk_metadata(
784 cu_seq_lens[:1] = 0
785 torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
786
787 + local_cu_seq_lens = cu_seq_lens
788 + local_total_seq_lens = total_seq_lens
789 + max_local_total_seq_lens = total_seq_lens
790 + if dcp_world_size > 1:
791 + # Per-rank local KV length under interleave-aware DCP sharding, shape
792 + # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the
793 + # sharding matches the rest of the DCP pipeline (decode/prefill).
794 + local_seq_lens = get_dcp_local_seq_lens(
795 + compressed_seq_lens[start_idx:end_idx],
796 + dcp_world_size,
797 + None,
798 + cp_kv_cache_interleave_size,
799 + )
800 + this_rank_counts = local_seq_lens[:, dcp_rank].to(torch.int32)
801 + local_cu_seq_lens = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
802 + torch.cumsum(this_rank_counts, dim=0, out=local_cu_seq_lens[1:])
803 + local_total_seq_lens = int(local_cu_seq_lens[-1].item())
804 + max_local_total_seq_lens = int(local_seq_lens.sum(dim=0).max().item())
805 +
806 query_start_loc = (
807 query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
808 )
@@ -687,15 +821,21 @@ def build_prefill_chunk_metadata(
821 cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
822 cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
823
824 + # Under DCP the kernel writes this rank's local row bounds into
825 + # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
826 _build_prefill_chunk_metadata_kernel[(num_reqs,)](
827 query_start_loc,
828 uncompressed_seq_lens[start_idx:end_idx],
829 cu_seq_lens,
830 + local_cu_seq_lens,
831 token_to_seq,
832 cu_seq_len_ks,
833 cu_seq_len_ke,
834 qs_start,
835 qs_stop,
836 + dcp_rank,
837 + dcp_world_size,
838 + cp_kv_cache_interleave_size,
839 BLOCK_SIZE=1024,
840 COMPRESS_RATIO=compress_ratio,
841 )
@@ -719,6 +859,9 @@ def build_prefill_chunk_metadata(
859 token_end=token_end,
860 num_reqs=num_reqs,
861 skip_kv_gather=skip_kv_gather,
862 + local_cu_seq_lens=local_cu_seq_lens,
863 + local_total_seq_lens=local_total_seq_lens,
864 + max_local_total_seq_lens=max_local_total_seq_lens,
865 )
866
867
@@ -728,12 +871,18 @@ def _build_prefill_chunk_metadata_kernel(
871 query_start_loc_ptr,
872 uncompressed_seq_lens_ptr,
873 cu_compressed_seq_lens_ptr,
874 + # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
875 + # aliases cu_compressed_seq_lens_ptr otherwise.
876 + row_start_cu_compressed_seq_lens_ptr,
877 # Outputs
878 token_to_seq_ptr,
879 cu_compressed_seq_len_ks_ptr,
880 cu_compressed_seq_len_ke_ptr,
881 query_slice_start,
882 query_slice_stop,
883 + DCP_RANK,
884 + DCP_WORLD,
885 + DCP_INTERLEAVE,
886 BLOCK_SIZE: tl.constexpr,
887 COMPRESS_RATIO: tl.constexpr,
888 ):
@@ -747,6 +896,10 @@ def _build_prefill_chunk_metadata_kernel(
896 seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
897 compressed_seq_len = seq_end - seq_start
898
899 + # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
900 + # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
901 + row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
902 +
903 uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
904 start_pos = uncompressed_seq_len - query_len
905
@@ -760,14 +913,25 @@ def _build_prefill_chunk_metadata_kernel(
913 )
914 out_pos = abs_pos - query_slice_start
915
763 - # Compute cu_seq_len_ks
764 - tl.store(cu_compressed_seq_len_ks_ptr + out_pos, seq_start, mask=mask)
765 -
766 - # Compute cu_seq_len_ke
767 - seq_len_per_token = (start_pos + 1 + offset) // COMPRESS_RATIO
916 + # cu_seq_len_ks: row start in the gathered K buffer.
917 + tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
918 +
919 + # cu_seq_len_ke: row start + per-token context length. Under DCP the
920 + # global per-token length is sharded across ranks.
921 + global_ctx = start_pos + 1 + offset
922 + len_per_token = global_ctx // COMPRESS_RATIO
923 + if DCP_WORLD > 1:
924 + # Per-rank local context length under interleave-aware DCP, matching
925 + # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
926 + base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
927 + remainder = len_per_token - base * DCP_WORLD
928 + remainder = tl.minimum(
929 + tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
930 + )
931 + len_per_token = base + remainder
932 tl.store(
933 cu_compressed_seq_len_ke_ptr + out_pos,
770 - seq_start + seq_len_per_token,
934 + row_start + len_per_token,
935 mask=mask,
936 )
937
vllm/v1/attention/backends/mla/sparse_utils.py
+158 -11
@@ -23,6 +23,16 @@ 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 + # When set, scatter valid slots to a contiguous prefix [0, valid_count) using
27 + # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
28 + # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
29 + # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
30 + # prefix is unspecified (only the selected set matters).
31 + COMPACT_TO_FRONT: tl.constexpr,
32 + # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
33 + DCP_SIZE: tl.constexpr,
34 + DCP_RANK: tl.constexpr,
35 + DCP_INTERLEAVE: tl.constexpr,
36 # strides (in elements)
37 bt_stride0,
38 bt_stride1,
@@ -52,15 +62,27 @@ def _convert_req_index_to_global_index_kernel(
62 if HAS_PREFILL:
63 prefill_req_id = tl.load(prefill_request_id_ptr + token_id)
64 is_prefill = prefill_req_id >= 0
65 +
66 + # DCP de-interleave the global token id into this rank's local slot.
67 + # Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With
68 + # DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never
69 + # remote) and local_idx == tok, so this reduces to the non-DCP path; with
70 + # DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).
71 + owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE
72 + is_remote = owning_rank != DCP_RANK
73 + local_idx = (
74 + tok // (DCP_SIZE * DCP_INTERLEAVE)
75 + ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE
76 +
77 # Compute block id and in-block offset
56 - block_id = tok // BLOCK_SIZE
57 - inblock_off = tok % BLOCK_SIZE
78 + block_id = local_idx // BLOCK_SIZE
79 + inblock_off = local_idx % BLOCK_SIZE
80
81 # Guard block_table access
82 valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
83 bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
62 - is_invalid_tok |= ~valid_block
63 - base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0)
84 + is_invalid_tok |= ~valid_block | is_remote
85 + base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
86 out_val = base * BLOCK_SIZE + inblock_off
87
88 # Override with prefill output if prefill is enabled
@@ -72,14 +94,27 @@ def _convert_req_index_to_global_index_kernel(
94 out_val = tl.where(is_prefill, prefill_out, out_val)
95 out_val = tl.where(is_invalid_tok, -1, out_val)
96
75 - # Store results
76 - out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
77 - tl.store(out_ptr_ij, out_val)
97 + if COMPACT_TO_FRONT:
98 + # Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix
99 + # sum gives each valid lane a distinct local offset; one atomic add of the
100 + # tile's valid count reserves a contiguous base across racing tiles. The
101 + # out buffer is pre-filled with -1, so unwritten tail slots stay -1.
102 + is_valid = (~is_invalid_tok).to(tl.int32)
103 + local_offset = tl.cumsum(is_valid) - is_valid
104 + tile_valid_count = tl.sum(is_valid)
105 + base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
106 + dest = base + local_offset
107 + out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
108 + tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
109 + else:
110 + # Store results in place (input column == output column).
111 + out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
112 + tl.store(out_ptr_ij, out_val)
113
79 - # Count valid indices in this tile and atomically add to row total
80 - if COUNT_VALID:
81 - tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
82 - tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
114 + # Count valid indices in this tile and atomically add to row total
115 + if COUNT_VALID:
116 + tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
117 + tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
118
119
120 def triton_convert_req_index_to_global_index(
@@ -176,6 +211,11 @@ def triton_convert_req_index_to_global_index(
211 BLOCK_N,
212 HAS_PREFILL_WORKSPACE,
213 return_valid_counts,
214 + False, # COMPACT_TO_FRONT: keep input column == output column
215 + # DCP disabled (no-op de-interleave)
216 + 1,
217 + 0,
218 + 1,
219 # strides
220 bt_stride0,
221 bt_stride1,
@@ -189,3 +229,110 @@ def triton_convert_req_index_to_global_index(
229 assert valid_counts is not None
230 return out, valid_counts
231 return out
232 +
233 +
234 +def triton_filter_and_convert_dcp_index(
235 + req_id: torch.Tensor,
236 + block_table: torch.Tensor,
237 + token_indices: torch.Tensor,
238 + dcp_size: int,
239 + dcp_rank: int,
240 + cp_kv_cache_interleave_size: int = 1,
241 + BLOCK_SIZE: int = 64,
242 + NUM_TOPK_TOKENS: int = 2048,
243 + BLOCK_N: int = 128,
244 + return_valid_counts: bool = False,
245 + compact_valid_to_front: bool = True,
246 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
247 + """Filter global per-request indices to this DCP rank's local slots.
248 +
249 + With ``compact_valid_to_front`` (default), the conversion kernel scatters
250 + this rank's owned slots to a contiguous prefix ``[0, valid_count)`` and
251 + leaves the rest ``-1``. DCP filtering marks non-owned slots ``-1`` and so
252 + creates interior gaps; the trtllm-gen sparse kernel reads the first
253 + ``valid_count`` entries of each row, so they must be a contiguous prefix.
254 + Compaction is fused into the kernel (atomic slot allocator) rather than a
255 + separate sort/gather pass. Prefix order is unspecified (only the set matters).
256 + """
257 + assert dcp_size >= 1
258 + assert 0 <= dcp_rank < dcp_size
259 + # Interleave groups must align to KV blocks (globally enforced by
260 + # VllmConfig: block_size % cp_kv_cache_interleave_size == 0); assert the
261 + # local invariant so local_idx // BLOCK_SIZE never straddles a group.
262 + assert BLOCK_SIZE % cp_kv_cache_interleave_size == 0, (
263 + f"BLOCK_SIZE ({BLOCK_SIZE}) must be divisible by "
264 + f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})."
265 + )
266 + assert req_id.dtype == torch.int32
267 + assert block_table.dtype == torch.int32
268 + assert token_indices.dtype == torch.int32
269 + assert token_indices.shape[1] == NUM_TOPK_TOKENS
270 + assert NUM_TOPK_TOKENS % BLOCK_N == 0
271 +
272 + if dcp_size == 1:
273 + return triton_convert_req_index_to_global_index(
274 + req_id,
275 + block_table,
276 + token_indices,
277 + BLOCK_SIZE=BLOCK_SIZE,
278 + NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
279 + BLOCK_N=BLOCK_N,
280 + return_valid_counts=return_valid_counts,
281 + )
282 +
283 + num_tokens = req_id.shape[0]
284 + max_num_blocks_per_req = block_table.shape[1]
285 + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N
286 +
287 + req_id_c = req_id.contiguous()
288 + block_table_c = block_table.contiguous()
289 + token_indices_c = token_indices.contiguous()
290 +
291 + # The compaction uses the valid-count buffer as an atomic slot allocator, so
292 + # it requires counting. Pre-fill out with -1 so the unwritten tail stays -1.
293 + count_valid = return_valid_counts or compact_valid_to_front
294 + if compact_valid_to_front:
295 + out = torch.full_like(token_indices_c, -1)
296 + else:
297 + out = torch.empty_like(token_indices_c)
298 +
299 + valid_counts: torch.Tensor | None = None
300 + if count_valid:
301 + valid_counts = torch.zeros(
302 + num_tokens, dtype=torch.int32, device=token_indices.device
303 + )
304 +
305 + bt_stride0, bt_stride1 = block_table_c.stride()
306 + ti_stride0, ti_stride1 = token_indices_c.stride()
307 + out_stride0, out_stride1 = out.stride()
308 +
309 + _convert_req_index_to_global_index_kernel[(num_tokens, tiles_per_row)](
310 + req_id_c,
311 + block_table_c,
312 + token_indices_c,
313 + out,
314 + valid_counts,
315 + # No prefill workspace on the DCP decode path.
316 + None,
317 + None,
318 + max_num_blocks_per_req,
319 + BLOCK_SIZE,
320 + BLOCK_N,
321 + False, # HAS_PREFILL
322 + count_valid,
323 + compact_valid_to_front,
324 + dcp_size,
325 + dcp_rank,
326 + cp_kv_cache_interleave_size,
327 + bt_stride0,
328 + bt_stride1,
329 + ti_stride0,
330 + ti_stride1,
331 + out_stride0,
332 + out_stride1,
333 + )
334 +
335 + if return_valid_counts:
336 + assert valid_counts is not None
337 + return out, valid_counts
338 + return out
vllm/v1/attention/backends/utils.py
+12 -12
@@ -892,20 +892,20 @@ def get_dcp_local_seq_lens(
892 use this function to calculate split decode seq_lens of each dcp rank.
893 Only consider dcp now, we can extend the case of cp based on this.
894 """
895 - num_requests = seq_lens.size(0)
895 + seq_lens_i32 = seq_lens.to(torch.int32)
896 if dcp_rank is None:
897 - rank_offsets = (
898 - torch.arange(dcp_size, dtype=torch.int32, device=seq_lens.device)
899 - .unsqueeze(0)
900 - .repeat(num_requests, 1)
897 + rank_offsets = torch.arange(
898 + dcp_size,
899 + dtype=torch.int32,
900 + device=seq_lens.device,
901 + ).view(
902 + *((1,) * seq_lens_i32.dim()),
903 + dcp_size,
904 )
905 + seq_lens_tiled = seq_lens_i32.unsqueeze(-1)
906 else:
903 - rank_offsets = torch.tensor(
904 - [[dcp_rank]], dtype=torch.int32, device=seq_lens.device
905 - )
906 - seq_lens_tiled = (
907 - seq_lens.to(torch.int32).unsqueeze(-1).repeat(1, rank_offsets.shape[1])
908 - )
907 + rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device)
908 + seq_lens_tiled = seq_lens_i32
909 base = (
910 seq_lens_tiled
911 // cp_kv_cache_interleave_size
@@ -919,7 +919,7 @@ def get_dcp_local_seq_lens(
919 cp_kv_cache_interleave_size,
920 )
921 dcp_local_seq_lens = base + remainder
922 - return dcp_local_seq_lens.squeeze(1)
922 + return dcp_local_seq_lens
923
924
925 def mamba_get_block_table_tensor(
vllm/v1/attention/ops/common.py
+1
@@ -90,6 +90,7 @@ def _correct_attn_cp_out_kernel(
90 factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
91 output = tl.load(outputs_ptr + output_offsets)
92 output = output * factor
93 + output = tl.where(factor == 0.0, 0.0, output)
94
95 tl.store(new_output_ptr + output_offsets, output)
96