| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import numpy as np |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.import_utils import has_cutedsl |
| 11 | from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata |
| 12 | from vllm.v1.attention.backends.mla.sparse_utils import ( |
| 13 | triton_filter_and_convert_dcp_index, |
| 14 | ) |
| 15 | from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens |
| 16 | from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out |
| 17 | |
| 18 | |
| 19 | def _local_count(length: int, rank: int, world: int, interleave: int) -> int: |
| 20 | return sum(1 for pos in range(length) if (pos // interleave) % world == rank) |
| 21 | |
| 22 | |
| 23 | def _global_to_local_indices( |
| 24 | global_indices: torch.Tensor, |
| 25 | rank: int, |
| 26 | world: int, |
| 27 | interleave: int, |
| 28 | ) -> torch.Tensor: |
| 29 | valid = global_indices >= 0 |
| 30 | global_i64 = global_indices.to(torch.int64).clamp_min(0) |
| 31 | owner = (global_i64 // interleave) % world |
| 32 | local = (global_i64 // (world * interleave)) * interleave + global_i64 % interleave |
| 33 | return torch.where(valid & (owner == rank), local, -1).to(torch.int64) |
| 34 | |
| 35 | |
| 36 | def _local_to_global_indices( |
| 37 | local_indices: torch.Tensor, |
| 38 | rank: int, |
| 39 | world: int, |
| 40 | interleave: int, |
| 41 | ) -> torch.Tensor: |
| 42 | valid = local_indices >= 0 |
| 43 | local = local_indices.to(torch.int64).clamp_min(0) |
| 44 | global_indices = ( |
| 45 | (local // interleave) * (world * interleave) |
| 46 | + rank * interleave |
| 47 | + local % interleave |
| 48 | ) |
| 49 | return torch.where(valid, global_indices, -1).to(torch.int64) |
| 50 | |
| 51 | |
| 52 | def _ref_stable_topk_from_candidates_fp64( |
| 53 | candidate_scores: torch.Tensor, |
| 54 | candidate_token_ids: torch.Tensor, |
| 55 | k: int, |
| 56 | ) -> torch.Tensor: |
| 57 | """Pure-PyTorch reference for the CuteDSL stable-topk selector order |
| 58 | (score desc, then lowest global token id). Selects the same SET as the |
| 59 | kernel; only the set is compared in tests.""" |
| 60 | num_rows, num_candidates = candidate_scores.shape |
| 61 | device = candidate_scores.device |
| 62 | select_k = min(k, num_candidates) |
| 63 | valid = candidate_token_ids >= 0 |
| 64 | bits = ( |
| 65 | candidate_scores.to(torch.float32).view(torch.int32).to(torch.int64) |
| 66 | & 0xFFFFFFFF |
| 67 | ) |
| 68 | sign = (bits >> 31) & 1 |
| 69 | score_key = ( |
| 70 | torch.where(sign.bool(), bits ^ 0xFFFFFFFF, bits ^ 0x80000000) & 0xFFFFFFFF |
| 71 | ) |
| 72 | id_key = (~candidate_token_ids.to(torch.int64)) & 0xFFFFFFFF |
| 73 | key = (score_key << 32) | id_key |
| 74 | key = torch.where(valid, key, torch.zeros_like(key)) |
| 75 | topk_key = key ^ torch.iinfo(torch.int64).min |
| 76 | _, topk_pos = topk_key.topk(select_k, dim=-1) |
| 77 | |
| 78 | selected = candidate_token_ids.gather(1, topk_pos).to(torch.int32) |
| 79 | selected_valid = valid.gather(1, topk_pos) |
| 80 | selected = torch.where(selected_valid, selected, selected.new_full((), -1)) |
| 81 | if select_k == k: |
| 82 | return selected |
| 83 | pad = torch.full((num_rows, k - select_k), -1, dtype=torch.int32, device=device) |
| 84 | return torch.cat((selected, pad), dim=1) |
| 85 | |
| 86 | |
| 87 | def _attention_from_indices( |
| 88 | q: torch.Tensor, |
| 89 | k: torch.Tensor, |
| 90 | v: torch.Tensor, |
| 91 | indices: torch.Tensor, |
| 92 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 93 | valid = indices >= 0 |
| 94 | safe_indices = indices.clamp_min(0) |
| 95 | selected_k = k[safe_indices] |
| 96 | selected_v = v[safe_indices] |
| 97 | scores = torch.einsum("td,tkd->tk", q, selected_k) |
| 98 | scores = scores.masked_fill(~valid, float("-inf")) |
| 99 | lse = torch.logsumexp(scores, dim=-1) |
| 100 | probs = torch.softmax(scores, dim=-1).masked_fill(~valid, 0.0) |
| 101 | out = torch.einsum("tk,tkd->td", probs, selected_v) |
| 102 | empty_rows = ~valid.any(dim=-1) |
| 103 | out[empty_rows] = 0 |
| 104 | lse[empty_rows] = float("-inf") |
| 105 | return out, lse |
| 106 | |
| 107 | |
| 108 | def _dcp_lse_merge( |
| 109 | local_outs: list[torch.Tensor], |
| 110 | local_lses: list[torch.Tensor], |
| 111 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 112 | outs = torch.stack(local_outs, dim=0) |
| 113 | lses = torch.stack(local_lses, dim=0) |
| 114 | merged_lse = torch.logsumexp(lses, dim=0) |
| 115 | weights = torch.exp(lses - merged_lse.unsqueeze(0)) |
| 116 | weights = torch.where(torch.isfinite(weights), weights, torch.zeros_like(weights)) |
| 117 | merged_out = (outs * weights.unsqueeze(-1)).sum(dim=0) |
| 118 | return merged_out, merged_lse |
| 119 | |
| 120 | |
| 121 | class _FakeDCPGroup: |
| 122 | """Single-process stand-in: ``all_gather`` returns the pre-built |
| 123 | concatenation of every rank's packed ``(score, global_id)`` candidates, |
| 124 | mirroring the one packed all-gather the merge issues.""" |
| 125 | |
| 126 | def __init__(self, gathered_packed: torch.Tensor) -> None: |
| 127 | self.gathered_packed = gathered_packed |
| 128 | |
| 129 | def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 130 | assert dim == 1 |
| 131 | return self.gathered_packed.clone() |
| 132 | |
| 133 | |
| 134 | def _run_decode_topk( |
| 135 | logits: torch.Tensor, |
| 136 | seq_lens: torch.Tensor, |
| 137 | next_n: int, |
| 138 | topk: int, |
| 139 | ) -> torch.Tensor: |
| 140 | indices = torch.empty( |
| 141 | (logits.shape[0], topk), dtype=torch.int32, device=logits.device |
| 142 | ) |
| 143 | torch.ops._C.top_k_per_row_decode( |
| 144 | logits, |
| 145 | next_n, |
| 146 | seq_lens, |
| 147 | indices, |
| 148 | logits.shape[0], |
| 149 | logits.stride(0), |
| 150 | logits.stride(1), |
| 151 | topk, |
| 152 | ) |
| 153 | return indices |
| 154 | |
| 155 | |
| 156 | def _run_persistent_topk( |
| 157 | logits: torch.Tensor, |
| 158 | seq_lens: torch.Tensor, |
| 159 | topk: int, |
| 160 | max_seq_len: int, |
| 161 | ) -> torch.Tensor: |
| 162 | indices = torch.empty( |
| 163 | (logits.shape[0], topk), dtype=torch.int32, device=logits.device |
| 164 | ) |
| 165 | workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device=logits.device) |
| 166 | torch.ops._C.persistent_topk( |
| 167 | logits, |
| 168 | seq_lens, |
| 169 | indices, |
| 170 | workspace, |
| 171 | topk, |
| 172 | max_seq_len, |
| 173 | ) |
| 174 | return indices |
| 175 | |
| 176 | |
| 177 | def _dcp_attention_from_local_topks( |
| 178 | q: torch.Tensor, |
| 179 | k: torch.Tensor, |
| 180 | v: torch.Tensor, |
| 181 | local_topks: list[torch.Tensor], |
| 182 | world: int, |
| 183 | interleave: int, |
| 184 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 185 | local_outs = [] |
| 186 | local_lses = [] |
| 187 | for rank, local_topk in enumerate(local_topks): |
| 188 | owned = [ |
| 189 | pos for pos in range(k.shape[0]) if (pos // interleave) % world == rank |
| 190 | ] |
| 191 | local_out, local_lse = _attention_from_indices( |
| 192 | q, k[owned], v[owned], local_topk.to(torch.int64) |
| 193 | ) |
| 194 | local_outs.append(local_out) |
| 195 | local_lses.append(local_lse) |
| 196 | return _dcp_lse_merge(local_outs, local_lses) |
| 197 | |
| 198 | |
| 199 | def _merge_local_topks_global_with_fake_dcp( |
| 200 | local_logits: list[torch.Tensor], |
| 201 | local_topks: list[torch.Tensor], |
| 202 | topk: int, |
| 203 | world: int, |
| 204 | interleave: int, |
| 205 | row_starts: list[torch.Tensor | None] | None = None, |
| 206 | ) -> list[torch.Tensor]: |
| 207 | """Run ``_merge_dcp_topk_global`` for every rank against a faked |
| 208 | all-gather and return each rank's global-index result (all ranks should |
| 209 | agree). The fake pre-builds the packed candidate concatenation the merge's |
| 210 | single ``all_gather`` would return.""" |
| 211 | # The merge is now CuteDSL-only (no PyTorch fallback), so it runs the real |
| 212 | # Triton pack + CuteDSL selector kernels even behind the faked all-gather. |
| 213 | if not current_platform.is_cuda() or not has_cutedsl(): |
| 214 | pytest.skip("DCP merge requires CUDA and CuteDSL") |
| 215 | packed_per_rank = [] |
| 216 | for rank, (logits, indices) in enumerate(zip(local_logits, local_topks)): |
| 217 | score_indices = indices.clamp_min(0).to(torch.long) |
| 218 | if row_starts is not None: |
| 219 | rs = row_starts[rank] |
| 220 | assert rs is not None |
| 221 | score_indices = score_indices + rs.to(torch.long).view(-1, 1) |
| 222 | if logits.shape[1] == 0: |
| 223 | scores = torch.full_like(indices, float("-inf"), dtype=torch.float32) |
| 224 | else: |
| 225 | score_indices = score_indices.clamp_max(logits.shape[1] - 1) |
| 226 | scores = logits.gather(1, score_indices).masked_fill( |
| 227 | indices < 0, float("-inf") |
| 228 | ) |
| 229 | global_ids = _local_to_global_indices(indices, rank, world, interleave) |
| 230 | packed_per_rank.append( |
| 231 | torch.stack((scores.float(), global_ids.to(torch.float32)), dim=-1) |
| 232 | ) |
| 233 | |
| 234 | fake_group = _FakeDCPGroup(torch.cat(packed_per_rank, dim=1).contiguous()) |
| 235 | original_get_dcp_group = sparse_indexer.get_dcp_group |
| 236 | sparse_indexer.get_dcp_group = lambda: fake_group |
| 237 | try: |
| 238 | merged = [] |
| 239 | for rank, (logits, indices) in enumerate(zip(local_logits, local_topks)): |
| 240 | rank_indices = indices.clone() |
| 241 | result = sparse_indexer._merge_dcp_topk_global( |
| 242 | logits, |
| 243 | rank_indices, |
| 244 | topk, |
| 245 | rank, |
| 246 | world, |
| 247 | interleave, |
| 248 | row_starts=None if row_starts is None else row_starts[rank], |
| 249 | ) |
| 250 | assert result is None |
| 251 | merged.append(rank_indices) |
| 252 | return merged |
| 253 | finally: |
| 254 | sparse_indexer.get_dcp_group = original_get_dcp_group |
| 255 | |
| 256 | |
| 257 | @pytest.mark.parametrize("world", [1, 2, 4]) |
| 258 | @pytest.mark.parametrize("interleave", [1, 2, 4]) |
| 259 | def test_get_dcp_local_seq_lens_matches_naive(world: int, interleave: int): |
| 260 | seq_lens = torch.arange(0, 33, dtype=torch.int32) |
| 261 | |
| 262 | for rank in range(world): |
| 263 | actual = get_dcp_local_seq_lens(seq_lens, world, rank, interleave) |
| 264 | expected = torch.tensor( |
| 265 | [ |
| 266 | _local_count(int(seq_len), rank, world, interleave) |
| 267 | for seq_len in seq_lens |
| 268 | ], |
| 269 | dtype=torch.int32, |
| 270 | ) |
| 271 | torch.testing.assert_close(actual, expected) |
| 272 | |
| 273 | |
| 274 | def test_get_dcp_local_seq_lens_can_localize_per_token_bounds(): |
| 275 | seq_lens = torch.tensor([0, 1, 2, 3, 4, 7, 8, 17], dtype=torch.int32) |
| 276 | world = 4 |
| 277 | interleave = 2 |
| 278 | |
| 279 | for rank in range(world): |
| 280 | actual = get_dcp_local_seq_lens(seq_lens, world, rank, interleave) |
| 281 | expected = torch.tensor( |
| 282 | [ |
| 283 | _local_count(int(seq_len), rank, world, interleave) |
| 284 | for seq_len in seq_lens |
| 285 | ], |
| 286 | dtype=torch.int32, |
| 287 | ) |
| 288 | torch.testing.assert_close(actual, expected) |
| 289 | |
| 290 | |
| 291 | def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape(): |
| 292 | seq_lens = torch.tensor([[8, 9, 10], [11, 12, 13]], dtype=torch.int32) |
| 293 | world = 2 |
| 294 | rank = 1 |
| 295 | interleave = 1 |
| 296 | |
| 297 | actual = get_dcp_local_seq_lens(seq_lens, world, rank, interleave) |
| 298 | expected = torch.tensor( |
| 299 | [ |
| 300 | [_local_count(int(seq_len), rank, world, interleave) for seq_len in row] |
| 301 | for row in seq_lens |
| 302 | ], |
| 303 | dtype=torch.int32, |
| 304 | ) |
| 305 | |
| 306 | assert actual.shape == seq_lens.shape |
| 307 | torch.testing.assert_close(actual, expected) |
| 308 | |
| 309 | |
| 310 | def test_get_dcp_local_seq_lens_reuses_device_rank_tensor(monkeypatch): |
| 311 | seq_lens = torch.tensor([0, 1, 7, 8, 17], dtype=torch.int32) |
| 312 | world = 4 |
| 313 | rank = 2 |
| 314 | rank_tensor = torch.tensor(rank, dtype=torch.int32) |
| 315 | expected = get_dcp_local_seq_lens(seq_lens, world, rank, 2) |
| 316 | |
| 317 | def fail_tensor_construction(*args, **kwargs): |
| 318 | raise AssertionError("rank tensor must be reused") |
| 319 | |
| 320 | with monkeypatch.context() as patch: |
| 321 | patch.setattr(torch, "tensor", fail_tensor_construction) |
| 322 | actual = get_dcp_local_seq_lens(seq_lens, world, rank_tensor, 2) |
| 323 | |
| 324 | torch.testing.assert_close(actual, expected) |
| 325 | |
| 326 | |
| 327 | def test_get_dcp_local_seq_lens_must_run_after_decode_expansion(): |
| 328 | world = 2 |
| 329 | rank = 1 |
| 330 | interleave = 1 |
| 331 | expanded_bounds = torch.tensor([8, 9, 10], dtype=torch.int32) |
| 332 | |
| 333 | localized_after_expansion = get_dcp_local_seq_lens( |
| 334 | expanded_bounds, world, rank, interleave |
| 335 | ) |
| 336 | localized_request_len_minus_offsets = get_dcp_local_seq_lens( |
| 337 | torch.tensor([10], dtype=torch.int32), world, rank |
| 338 | ) - torch.tensor([2, 1, 0], dtype=torch.int32) |
| 339 | |
| 340 | assert not torch.equal( |
| 341 | localized_after_expansion, localized_request_len_minus_offsets |
| 342 | ) |
| 343 | torch.testing.assert_close( |
| 344 | localized_after_expansion, torch.tensor([4, 4, 5], dtype=torch.int32) |
| 345 | ) |
| 346 | |
| 347 | |
| 348 | @pytest.mark.parametrize("interleave", [1, 2]) |
| 349 | def test_sparse_dcp_attention_matches_global_topk_attention(interleave: int): |
| 350 | torch.manual_seed(0) |
| 351 | world = 2 |
| 352 | topk = 3 |
| 353 | num_queries = 4 |
| 354 | max_seq_len = 13 |
| 355 | head_dim = 8 |
| 356 | |
| 357 | q = torch.randn(num_queries, head_dim) |
| 358 | k = torch.randn(max_seq_len, head_dim) |
| 359 | v = torch.randn(max_seq_len, head_dim) |
| 360 | seq_lens = torch.tensor([6, 8, 11, 13], dtype=torch.int64) |
| 361 | |
| 362 | scores = q @ k.T |
| 363 | global_topk = torch.full((num_queries, topk), -1, dtype=torch.int64) |
| 364 | for row, seq_len in enumerate(seq_lens.tolist()): |
| 365 | global_topk[row, :topk] = scores[row, :seq_len].topk(topk).indices |
| 366 | |
| 367 | ref_out, ref_lse = _attention_from_indices(q, k, v, global_topk) |
| 368 | |
| 369 | local_outs = [] |
| 370 | local_lses = [] |
| 371 | local_topks = [] |
| 372 | for rank in range(world): |
| 373 | owned = [ |
| 374 | pos for pos in range(max_seq_len) if (pos // interleave) % world == rank |
| 375 | ] |
| 376 | k_local = k[owned] |
| 377 | v_local = v[owned] |
| 378 | local_topk = _global_to_local_indices(global_topk, rank, world, interleave) |
| 379 | local_topks.append(local_topk) |
| 380 | local_out, local_lse = _attention_from_indices(q, k_local, v_local, local_topk) |
| 381 | local_outs.append(local_out) |
| 382 | local_lses.append(local_lse) |
| 383 | |
| 384 | dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses) |
| 385 | |
| 386 | torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5) |
| 387 | torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5) |
| 388 | |
| 389 | gathered_global = torch.cat( |
| 390 | [ |
| 391 | _local_to_global_indices(local_topks[rank], rank, world, interleave) |
| 392 | for rank in range(world) |
| 393 | ], |
| 394 | dim=1, |
| 395 | ) |
| 396 | assert set(gathered_global[gathered_global >= 0].tolist()) == set( |
| 397 | global_topk.flatten().tolist() |
| 398 | ) |
| 399 | |
| 400 | |
| 401 | def test_local_topk_union_is_not_equivalent_to_global_topk_attention(): |
| 402 | world = 2 |
| 403 | interleave = 1 |
| 404 | topk = 2 |
| 405 | q = torch.tensor([[1.0]]) |
| 406 | k = torch.tensor( |
| 407 | [ |
| 408 | [1.00], |
| 409 | [0.90], |
| 410 | [0.95], |
| 411 | [0.85], |
| 412 | ] |
| 413 | ) |
| 414 | v = torch.tensor([[0.0], [1000.0], [0.0], [1000.0]]) |
| 415 | |
| 416 | scores = q @ k.T |
| 417 | global_topk = scores.topk(topk, dim=-1).indices |
| 418 | ref_out, ref_lse = _attention_from_indices(q, k, v, global_topk) |
| 419 | |
| 420 | local_topks = [] |
| 421 | for rank in range(world): |
| 422 | owned = [ |
| 423 | pos for pos in range(k.shape[0]) if (pos // interleave) % world == rank |
| 424 | ] |
| 425 | local_topks.append(scores[:, owned].topk(topk, dim=-1).indices) |
| 426 | |
| 427 | local_union_out, local_union_lse = _dcp_attention_from_local_topks( |
| 428 | q, k, v, local_topks, world, interleave |
| 429 | ) |
| 430 | |
| 431 | assert not torch.allclose(local_union_out, ref_out) |
| 432 | assert not torch.allclose(local_union_lse, ref_lse) |
| 433 | |
| 434 | |
| 435 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 436 | def test_sparse_decode_dcp_persistent_topk_matches_non_dcp(): |
| 437 | torch.manual_seed(3) |
| 438 | device = torch.device("cuda") |
| 439 | world = 2 |
| 440 | interleave = 1 |
| 441 | topk = 512 |
| 442 | num_rows = 2 |
| 443 | max_seq_len = 1025 |
| 444 | head_dim = 16 |
| 445 | |
| 446 | q = torch.randn(num_rows, head_dim, device=device) |
| 447 | k = torch.randn(max_seq_len, head_dim, device=device) |
| 448 | v = torch.randn(max_seq_len, head_dim, device=device) |
| 449 | logits = q @ k.T |
| 450 | seq_lens = torch.tensor([[1024], [1025]], dtype=torch.int32, device=device) |
| 451 | |
| 452 | non_dcp_topk = torch.empty((num_rows, topk), dtype=torch.int64, device=device) |
| 453 | for row, seq_len in enumerate(seq_lens.flatten().tolist()): |
| 454 | non_dcp_topk[row] = logits[row, :seq_len].topk(topk).indices |
| 455 | ref_out, ref_lse = _attention_from_indices(q, k, v, non_dcp_topk) |
| 456 | |
| 457 | local_logits = [] |
| 458 | local_topks = [] |
| 459 | for rank in range(world): |
| 460 | owned = [ |
| 461 | pos for pos in range(max_seq_len) if (pos // interleave) % world == rank |
| 462 | ] |
| 463 | rank_logits = logits[:, owned].contiguous() |
| 464 | rank_seq_lens = get_dcp_local_seq_lens( |
| 465 | seq_lens, world, rank, interleave |
| 466 | ).contiguous() |
| 467 | local_logits.append(rank_logits) |
| 468 | local_topks.append( |
| 469 | _run_persistent_topk( |
| 470 | rank_logits, |
| 471 | rank_seq_lens, |
| 472 | topk, |
| 473 | max_seq_len=rank_logits.shape[1], |
| 474 | ) |
| 475 | ) |
| 476 | |
| 477 | merged_global_topks = _merge_local_topks_global_with_fake_dcp( |
| 478 | local_logits, local_topks, topk, world, interleave |
| 479 | ) |
| 480 | # The radix top-K kernel selects a deterministic SET but writes it in |
| 481 | # nondeterministic (atomicAdd) order; the production path is permutation- |
| 482 | # invariant (compaction + softmax), so all ranks must agree on the set, not |
| 483 | # the array order. (The fp64 fallback happens to return sorted order.) |
| 484 | ref_topk = merged_global_topks[0] |
| 485 | for rank_topk in merged_global_topks[1:]: |
| 486 | for row in range(rank_topk.shape[0]): |
| 487 | assert set(rank_topk[row].tolist()) == set(ref_topk[row].tolist()) |
| 488 | |
| 489 | local_outs = [] |
| 490 | local_lses = [] |
| 491 | for rank, global_topk in enumerate(merged_global_topks): |
| 492 | owned = [ |
| 493 | pos for pos in range(max_seq_len) if (pos // interleave) % world == rank |
| 494 | ] |
| 495 | local_topk = _global_to_local_indices( |
| 496 | global_topk.to(torch.int64), |
| 497 | rank, |
| 498 | world, |
| 499 | interleave, |
| 500 | ) |
| 501 | local_out, local_lse = _attention_from_indices( |
| 502 | q, k[owned], v[owned], local_topk |
| 503 | ) |
| 504 | local_outs.append(local_out) |
| 505 | local_lses.append(local_lse) |
| 506 | |
| 507 | dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses) |
| 508 | torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5) |
| 509 | torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5) |
| 510 | |
| 511 | |
| 512 | @pytest.mark.skipif( |
| 513 | not current_platform.is_cuda() or not has_cutedsl(), |
| 514 | reason="This test requires CUDA and CuteDSL", |
| 515 | ) |
| 516 | @pytest.mark.parametrize("use_row_starts", [False, True]) |
| 517 | def test_cutedsl_dcp_candidate_pack_and_select_matches_reference( |
| 518 | use_row_starts: bool, |
| 519 | ): |
| 520 | from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( |
| 521 | pack_dcp_topk_candidates_cutedsl, |
| 522 | stable_topk_from_gathered_candidates_cutedsl, |
| 523 | ) |
| 524 | |
| 525 | torch.manual_seed(13) |
| 526 | device = torch.device("cuda") |
| 527 | rows = 4 |
| 528 | valid_width = 1024 |
| 529 | width = valid_width + (8 if use_row_starts else 0) |
| 530 | topk = 512 |
| 531 | world = 2 |
| 532 | row_starts = ( |
| 533 | torch.tensor([0, 2, 4, 1], device=device, dtype=torch.int32) |
| 534 | if use_row_starts |
| 535 | else None |
| 536 | ) |
| 537 | row_offsets = ( |
| 538 | row_starts |
| 539 | if row_starts is not None |
| 540 | else torch.zeros(rows, device=device, dtype=torch.int32) |
| 541 | ) |
| 542 | |
| 543 | packed_by_rank = [] |
| 544 | for rank in range(world): |
| 545 | logits = torch.randn((rows, width), device=device, dtype=torch.float32) |
| 546 | local_topks = [] |
| 547 | for row in range(rows): |
| 548 | start = int(row_offsets[row].item()) |
| 549 | local_topks.append( |
| 550 | logits[row, start : start + valid_width].topk(topk).indices |
| 551 | ) |
| 552 | topk_indices = torch.stack(local_topks).to(torch.int32) |
| 553 | |
| 554 | packed = torch.empty((rows, topk, 2), device=device, dtype=torch.float32) |
| 555 | pack_dcp_topk_candidates_cutedsl( |
| 556 | logits, |
| 557 | topk_indices, |
| 558 | packed, |
| 559 | rank, |
| 560 | world, |
| 561 | 1, |
| 562 | row_starts, |
| 563 | ) |
| 564 | |
| 565 | expected_scores = logits.gather( |
| 566 | 1, topk_indices.to(torch.long) + row_offsets.to(torch.long).view(-1, 1) |
| 567 | ) |
| 568 | expected_ids = (topk_indices * world + rank).to(torch.float32) |
| 569 | torch.testing.assert_close(packed[..., 0], expected_scores) |
| 570 | torch.testing.assert_close(packed[..., 1], expected_ids) |
| 571 | packed_by_rank.append(packed) |
| 572 | |
| 573 | gathered = torch.cat(packed_by_rank, dim=1).contiguous() |
| 574 | actual = torch.empty((rows, topk), device=device, dtype=torch.int32) |
| 575 | returned = stable_topk_from_gathered_candidates_cutedsl(gathered, topk, out=actual) |
| 576 | assert returned is actual |
| 577 | expected = _ref_stable_topk_from_candidates_fp64( |
| 578 | gathered[..., 0], |
| 579 | gathered[..., 1].to(torch.int32), |
| 580 | topk, |
| 581 | ) |
| 582 | |
| 583 | for row in range(rows): |
| 584 | assert set(actual[row].cpu().tolist()) == set(expected[row].cpu().tolist()) |
| 585 | |
| 586 | |
| 587 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 588 | def test_sparse_prefill_dcp_metadata_localizes_causal_bounds(): |
| 589 | device = torch.device("cuda") |
| 590 | seq_len = 8 |
| 591 | |
| 592 | query_start_loc = torch.tensor([0, seq_len], dtype=torch.int32, device=device) |
| 593 | query_start_loc_cpu = torch.tensor([0, seq_len], dtype=torch.int32) |
| 594 | seq_lens = torch.tensor([seq_len], dtype=torch.int32, device=device) |
| 595 | seq_lens_cpu = torch.tensor([seq_len], dtype=torch.int32) |
| 596 | block_table = torch.zeros((1, 1), dtype=torch.int32, device=device) |
| 597 | |
| 598 | def build(dcp_world_size, dcp_rank, interleave=1): |
| 599 | chunk = build_prefill_chunk_metadata( |
| 600 | start_idx=0, |
| 601 | end_idx=1, |
| 602 | query_start_loc=query_start_loc, |
| 603 | query_start_loc_cpu=query_start_loc_cpu, |
| 604 | uncompressed_seq_lens=seq_lens, |
| 605 | compressed_seq_lens=seq_lens, |
| 606 | compressed_seq_lens_cpu=seq_lens_cpu, |
| 607 | block_table=block_table, |
| 608 | compress_ratio=1, |
| 609 | dcp_rank=dcp_rank, |
| 610 | dcp_world_size=dcp_world_size, |
| 611 | cp_kv_cache_interleave_size=interleave, |
| 612 | ) |
| 613 | assert chunk is not None |
| 614 | torch.accelerator.synchronize() |
| 615 | return chunk |
| 616 | |
| 617 | # Non-DCP: local_cu_seq_lens aliases the global cu_seq_lens, and |
| 618 | # cu_seqlen_ks/ke carry the global causal bounds. |
| 619 | chunk = build(dcp_world_size=1, dcp_rank=0) |
| 620 | assert chunk.local_cu_seq_lens is chunk.cu_seq_lens |
| 621 | torch.testing.assert_close( |
| 622 | chunk.cu_seqlen_ks.cpu(), |
| 623 | torch.zeros(seq_len, dtype=torch.int32), |
| 624 | ) |
| 625 | torch.testing.assert_close( |
| 626 | chunk.cu_seqlen_ke.cpu(), |
| 627 | torch.arange(1, seq_len + 1, dtype=torch.int32), |
| 628 | ) |
| 629 | |
| 630 | # DCP: cu_seqlen_ks/ke are localized in place to this rank's shard. |
| 631 | chunk = build(dcp_world_size=4, dcp_rank=0) |
| 632 | assert chunk.local_cu_seq_lens is not None |
| 633 | torch.testing.assert_close( |
| 634 | chunk.local_cu_seq_lens.cpu(), |
| 635 | torch.tensor([0, 2], dtype=torch.int32), |
| 636 | ) |
| 637 | torch.testing.assert_close( |
| 638 | chunk.cu_seqlen_ks.cpu(), |
| 639 | torch.zeros(seq_len, dtype=torch.int32), |
| 640 | ) |
| 641 | torch.testing.assert_close( |
| 642 | chunk.cu_seqlen_ke.cpu(), |
| 643 | torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int32), |
| 644 | ) |
| 645 | |
| 646 | # DCP with interleave=2: per-token causal bounds localize differently from |
| 647 | # interleave=1 (groups of 2 consecutive tokens are owned together). For |
| 648 | # world=4, rank=0, K=2, per-token global len L=1..8 -> local len |
| 649 | # [1,2,2,2,2,2,2,2] (matches get_dcp_local_seq_lens). |
| 650 | chunk = build(dcp_world_size=4, dcp_rank=0, interleave=2) |
| 651 | assert chunk.local_cu_seq_lens is not None |
| 652 | torch.testing.assert_close( |
| 653 | chunk.cu_seqlen_ks.cpu(), |
| 654 | torch.zeros(seq_len, dtype=torch.int32), |
| 655 | ) |
| 656 | torch.testing.assert_close( |
| 657 | chunk.cu_seqlen_ke.cpu(), |
| 658 | torch.tensor([1, 2, 2, 2, 2, 2, 2, 2], dtype=torch.int32), |
| 659 | ) |
| 660 | |
| 661 | |
| 662 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 663 | @pytest.mark.parametrize("block_stride_rows", [None, 12]) |
| 664 | def test_dcp_filter_compacts_valid_slots_for_sparse_kernel( |
| 665 | block_stride_rows: int | None, |
| 666 | ): |
| 667 | block_size = 4 |
| 668 | num_topk = 128 |
| 669 | dcp_size = 2 |
| 670 | req_id = torch.zeros(1, dtype=torch.int32, device="cuda") |
| 671 | token_indices = torch.full((1, num_topk), -1, dtype=torch.int32, device="cuda") |
| 672 | token_indices[0, :8] = torch.arange(8, dtype=torch.int32, device="cuda") |
| 673 | block_table = torch.tensor([[10]], dtype=torch.int32, device="cuda") |
| 674 | |
| 675 | out, valid_counts = triton_filter_and_convert_dcp_index( |
| 676 | req_id, |
| 677 | block_table, |
| 678 | token_indices, |
| 679 | dcp_size=dcp_size, |
| 680 | dcp_rank=0, |
| 681 | BLOCK_SIZE=block_size, |
| 682 | BLOCK_STRIDE_ROWS=block_stride_rows, |
| 683 | NUM_TOPK_TOKENS=num_topk, |
| 684 | return_valid_counts=True, |
| 685 | ) |
| 686 | |
| 687 | valid = int(valid_counts.item()) |
| 688 | assert valid == 4 |
| 689 | assert (out[0, :valid] >= 0).all() |
| 690 | assert (out[0, valid:] == -1).all() |
| 691 | # In-kernel compaction packs valid slots to the front; prefix order is |
| 692 | # unspecified, so compare as a set. |
| 693 | block_stride_rows = block_stride_rows or block_size |
| 694 | expected_base = 10 * block_stride_rows |
| 695 | assert set(out[0, :valid].cpu().tolist()) == set( |
| 696 | range(expected_base, expected_base + block_size) |
| 697 | ) |
| 698 | |
| 699 | |
| 700 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 701 | @pytest.mark.parametrize("interleave", [1, 2]) |
| 702 | @pytest.mark.parametrize("dcp_rank", [0, 1]) |
| 703 | # Include padded single-tile compaction and the multi-tile atomic fallback. |
| 704 | @pytest.mark.parametrize("num_topk", [1024, 2176, 4224]) |
| 705 | def test_dcp_filter_compaction_matches_reference( |
| 706 | interleave: int, dcp_rank: int, num_topk: int |
| 707 | ): |
| 708 | """In-kernel compaction must, for every row, produce exactly the rank-owned |
| 709 | physical slots packed into [0, valid_count) with -1 in the tail -- the same |
| 710 | SET a reference filter + sort/gather produces. Uses wide rows (> BLOCK_N |
| 711 | valid slots), with interior -1 gaps.""" |
| 712 | device = torch.device("cuda") |
| 713 | torch.manual_seed(7) |
| 714 | dcp_size = 2 |
| 715 | block_size = 8 |
| 716 | num_rows = 5 |
| 717 | max_blocks = 64 |
| 718 | seq = max_blocks * block_size |
| 719 | |
| 720 | req_id = torch.randint(0, 3, (num_rows,), dtype=torch.int32, device=device) |
| 721 | block_table = torch.randint( |
| 722 | 0, 1000, (3, max_blocks), dtype=torch.int32, device=device |
| 723 | ) |
| 724 | # Each row: a dense valid prefix of distinct global token ids, then -1 pad. |
| 725 | token_indices = torch.full( |
| 726 | (num_rows, num_topk), -1, dtype=torch.int32, device=device |
| 727 | ) |
| 728 | for r in range(num_rows): |
| 729 | # Ids are distinct, so a row holds at most `seq` valid entries. |
| 730 | hi = min(num_topk * 3 // 4, seq) |
| 731 | n_valid = int(torch.randint(hi // 3, hi, (1,)).item()) |
| 732 | perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32) |
| 733 | token_indices[r, :n_valid] = perm |
| 734 | |
| 735 | out, valid_counts = triton_filter_and_convert_dcp_index( |
| 736 | req_id, |
| 737 | block_table, |
| 738 | token_indices, |
| 739 | dcp_size=dcp_size, |
| 740 | dcp_rank=dcp_rank, |
| 741 | cp_kv_cache_interleave_size=interleave, |
| 742 | BLOCK_SIZE=block_size, |
| 743 | NUM_TOPK_TOKENS=num_topk, |
| 744 | return_valid_counts=True, |
| 745 | ) |
| 746 | |
| 747 | for r in range(num_rows): |
| 748 | toks = token_indices[r] |
| 749 | toks = toks[toks >= 0] |
| 750 | owner = (toks // interleave) % dcp_size |
| 751 | owned = toks[owner == dcp_rank] |
| 752 | local = (owned // (dcp_size * interleave)) * interleave + owned % interleave |
| 753 | blk = local // block_size |
| 754 | off = local % block_size |
| 755 | expected = ( |
| 756 | block_table[int(req_id[r].item()), blk].to(torch.int64) * block_size + off |
| 757 | ) |
| 758 | n = int(valid_counts[r].item()) |
| 759 | assert n == owned.numel() |
| 760 | assert (out[r, :n] >= 0).all() |
| 761 | assert (out[r, n:] == -1).all() |
| 762 | assert set(out[r, :n].cpu().tolist()) == set(expected.cpu().tolist()) |
| 763 | |
| 764 | |
| 765 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 766 | @pytest.mark.parametrize("interleave", [1, 2, 4]) |
| 767 | def test_dcp_global_topk_physical_attention_matches_non_dcp(interleave: int): |
| 768 | torch.manual_seed(2) |
| 769 | device = torch.device("cuda") |
| 770 | dcp_size = 2 |
| 771 | block_size = 4 |
| 772 | num_topk = 128 |
| 773 | selected_k = 8 |
| 774 | seq_len = 16 |
| 775 | head_dim = 16 |
| 776 | num_queries = 3 |
| 777 | |
| 778 | q = torch.randn(num_queries, head_dim, device=device) |
| 779 | k_global = torch.randn(seq_len, head_dim, device=device) |
| 780 | v_global = torch.randn(seq_len, head_dim, device=device) |
| 781 | scores = q @ k_global.T |
| 782 | global_topk = torch.full( |
| 783 | (num_queries, num_topk), -1, dtype=torch.int32, device=device |
| 784 | ) |
| 785 | global_topk[:, :selected_k] = scores.topk(selected_k, dim=-1).indices.to( |
| 786 | torch.int32 |
| 787 | ) |
| 788 | |
| 789 | ref_out, ref_lse = _attention_from_indices( |
| 790 | q, k_global, v_global, global_topk[:, :selected_k].to(torch.int64) |
| 791 | ) |
| 792 | |
| 793 | local_outs = [] |
| 794 | local_lses = [] |
| 795 | for rank in range(dcp_size): |
| 796 | block_ids = torch.tensor( |
| 797 | [[rank * 10 + 1, rank * 10 + 2]], dtype=torch.int32, device=device |
| 798 | ) |
| 799 | num_slots = int((block_ids.max().item() + 1) * block_size) |
| 800 | k_cache = torch.zeros(num_slots, head_dim, device=device) |
| 801 | v_cache = torch.zeros(num_slots, head_dim, device=device) |
| 802 | for global_idx in range(seq_len): |
| 803 | if (global_idx // interleave) % dcp_size != rank: |
| 804 | continue |
| 805 | local_idx = ( |
| 806 | global_idx // (dcp_size * interleave) |
| 807 | ) * interleave + global_idx % interleave |
| 808 | block = local_idx // block_size |
| 809 | offset = local_idx % block_size |
| 810 | slot = int(block_ids[0, block].item()) * block_size + offset |
| 811 | k_cache[slot] = k_global[global_idx] |
| 812 | v_cache[slot] = v_global[global_idx] |
| 813 | |
| 814 | slots, valid_counts = triton_filter_and_convert_dcp_index( |
| 815 | torch.zeros(num_queries, dtype=torch.int32, device=device), |
| 816 | block_ids, |
| 817 | global_topk, |
| 818 | dcp_size=dcp_size, |
| 819 | dcp_rank=rank, |
| 820 | cp_kv_cache_interleave_size=interleave, |
| 821 | BLOCK_SIZE=block_size, |
| 822 | NUM_TOPK_TOKENS=num_topk, |
| 823 | return_valid_counts=True, |
| 824 | ) |
| 825 | row_ids = torch.arange(num_queries, device=device) |
| 826 | assert (slots[row_ids, valid_counts] == -1).all() |
| 827 | local_out, local_lse = _attention_from_indices( |
| 828 | q, k_cache, v_cache, slots.to(torch.int64) |
| 829 | ) |
| 830 | local_outs.append(local_out) |
| 831 | local_lses.append(local_lse) |
| 832 | |
| 833 | dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses) |
| 834 | torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5) |
| 835 | torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5) |
| 836 | |
| 837 | |
| 838 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 839 | @pytest.mark.parametrize("is_lse_base_on_e", [True, False]) |
| 840 | def test_correct_attn_out_zeroes_empty_nan_partial(is_lse_base_on_e: bool): |
| 841 | out = torch.full((1, 1, 4), float("nan"), device="cuda") |
| 842 | lses = torch.tensor( |
| 843 | [[[0.0]], [[float("-inf")]]], |
| 844 | dtype=torch.float32, |
| 845 | device="cuda", |
| 846 | ) |
| 847 | |
| 848 | corrected, final_lse = correct_attn_out( |
| 849 | out, |
| 850 | lses, |
| 851 | cp_rank=1, |
| 852 | ctx=CPTritonContext(), |
| 853 | is_lse_base_on_e=is_lse_base_on_e, |
| 854 | ) |
| 855 | torch.accelerator.synchronize() |
| 856 | |
| 857 | torch.testing.assert_close(corrected, torch.zeros_like(corrected)) |
| 858 | torch.testing.assert_close(final_lse, torch.zeros_like(final_lse)) |
| 859 | |
| 860 | |
| 861 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 862 | def test_decode_topk_pads_surplus_with_negative_one(): |
| 863 | """When a row's valid length < topk the top-k kernel must pad the surplus |
| 864 | slots with -1: the DCP merge (`topk_indices >= 0`) and |
| 865 | `triton_filter_and_convert_dcp_index` (`tok < 0`) both treat <0 as invalid, |
| 866 | so a non-(-1) pad would be silently attended. This is the common case under |
| 867 | DCP, where each rank's local seq_len = global / world is usually << topk. |
| 868 | Checked on the *set* of valid indices (the kernel may order them by score, |
| 869 | not position).""" |
| 870 | torch.manual_seed(0) |
| 871 | device = torch.device("cuda") |
| 872 | topk = 8 |
| 873 | next_n = 1 |
| 874 | seq_lens_list = [0, 3, 5, 8] |
| 875 | num_rows = len(seq_lens_list) |
| 876 | logits = torch.randn(num_rows, 16, device=device) |
| 877 | seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device).view( |
| 878 | num_rows, next_n |
| 879 | ) |
| 880 | idx = _run_decode_topk(logits, seq_lens, next_n, topk) |
| 881 | for r, sl in enumerate(seq_lens_list): |
| 882 | valid = idx[r][idx[r] >= 0] |
| 883 | # seq_len <= topk, so top-k selects exactly the whole valid range. |
| 884 | assert valid.numel() == min(sl, topk) |
| 885 | assert set(valid.tolist()) == set(range(sl)) |
| 886 | assert (idx[r] == -1).sum().item() == topk - min(sl, topk) |
| 887 | |
| 888 | |
| 889 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 890 | def test_persistent_topk_pads_surplus_with_negative_one(): |
| 891 | """Same surplus=-1 invariant for the persistent_topk kernel (k>=512).""" |
| 892 | torch.manual_seed(0) |
| 893 | device = torch.device("cuda") |
| 894 | topk = 512 # persistent_topk requires k in {512, 1024, 2048} |
| 895 | seq_lens_list = [100, 300, 512] |
| 896 | num_rows = len(seq_lens_list) |
| 897 | max_seq_len = 600 |
| 898 | logits = torch.randn(num_rows, max_seq_len, device=device) |
| 899 | seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device).view( |
| 900 | num_rows, 1 |
| 901 | ) |
| 902 | idx = _run_persistent_topk(logits, seq_lens, topk, max_seq_len) |
| 903 | for r, sl in enumerate(seq_lens_list): |
| 904 | valid = idx[r][idx[r] >= 0] |
| 905 | assert valid.numel() == min(sl, topk) |
| 906 | assert set(valid.tolist()) == set(range(sl)) |
| 907 | assert (idx[r] == -1).sum().item() == topk - min(sl, topk) |
| 908 | |
| 909 | |
| 910 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 911 | def test_sparse_decode_dcp_short_context_matches_non_dcp(): |
| 912 | """End-to-end DCP decode where the global seq_len < topk (so every rank's |
| 913 | local top-k is surplus-padded). Exercises the kernel surplus -> merge mask |
| 914 | -> global top-k -> physical localize -> LSE merge chain for the common |
| 915 | short-context decode case, vs the non-DCP reference.""" |
| 916 | torch.manual_seed(4) |
| 917 | device = torch.device("cuda") |
| 918 | world = 2 |
| 919 | interleave = 1 |
| 920 | topk = 512 |
| 921 | num_rows = 2 |
| 922 | max_seq_len = 300 # < topk -> surplus everywhere |
| 923 | head_dim = 16 |
| 924 | |
| 925 | q = torch.randn(num_rows, head_dim, device=device) |
| 926 | k = torch.randn(max_seq_len, head_dim, device=device) |
| 927 | v = torch.randn(max_seq_len, head_dim, device=device) |
| 928 | logits = q @ k.T |
| 929 | seq_lens = torch.tensor([[250], [300]], dtype=torch.int32, device=device) |
| 930 | |
| 931 | non_dcp_topk = torch.empty((num_rows, topk), dtype=torch.int64, device=device) |
| 932 | for row, seq_len in enumerate(seq_lens.flatten().tolist()): |
| 933 | sel = logits[row, :seq_len].topk(min(topk, seq_len)).indices |
| 934 | non_dcp_topk[row, : sel.numel()] = sel |
| 935 | non_dcp_topk[row, sel.numel() :] = -1 |
| 936 | ref_out, ref_lse = _attention_from_indices(q, k, v, non_dcp_topk) |
| 937 | |
| 938 | local_logits = [] |
| 939 | local_topks = [] |
| 940 | for rank in range(world): |
| 941 | owned = [p for p in range(max_seq_len) if (p // interleave) % world == rank] |
| 942 | rank_logits = logits[:, owned].contiguous() |
| 943 | rank_seq_lens = get_dcp_local_seq_lens(seq_lens, world, rank, interleave) |
| 944 | local_logits.append(rank_logits) |
| 945 | local_topks.append( |
| 946 | _run_persistent_topk( |
| 947 | rank_logits, |
| 948 | rank_seq_lens.contiguous(), |
| 949 | topk, |
| 950 | max_seq_len=rank_logits.shape[1], |
| 951 | ) |
| 952 | ) |
| 953 | |
| 954 | merged_global_topks = _merge_local_topks_global_with_fake_dcp( |
| 955 | local_logits, local_topks, topk, world, interleave |
| 956 | ) |
| 957 | # The radix top-K kernel selects a deterministic SET but writes it in |
| 958 | # nondeterministic (atomicAdd) order; the production path is permutation- |
| 959 | # invariant (compaction + softmax), so all ranks must agree on the set, not |
| 960 | # the array order. (The fp64 fallback happens to return sorted order.) |
| 961 | ref_topk = merged_global_topks[0] |
| 962 | for rank_topk in merged_global_topks[1:]: |
| 963 | for row in range(rank_topk.shape[0]): |
| 964 | assert set(rank_topk[row].tolist()) == set(ref_topk[row].tolist()) |
| 965 | |
| 966 | local_outs = [] |
| 967 | local_lses = [] |
| 968 | for rank, global_topk in enumerate(merged_global_topks): |
| 969 | owned = [p for p in range(max_seq_len) if (p // interleave) % world == rank] |
| 970 | local_topk = _global_to_local_indices( |
| 971 | global_topk.to(torch.int64), rank, world, interleave |
| 972 | ) |
| 973 | local_out, local_lse = _attention_from_indices( |
| 974 | q, k[owned], v[owned], local_topk |
| 975 | ) |
| 976 | local_outs.append(local_out) |
| 977 | local_lses.append(local_lse) |
| 978 | |
| 979 | dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses) |
| 980 | torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5) |
| 981 | torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5) |
| 982 | |
| 983 | |
| 984 | @pytest.mark.parametrize("dcp_world_size", [2, 4, 8]) |
| 985 | @pytest.mark.parametrize("interleave", [1, 64]) |
| 986 | @pytest.mark.parametrize("rows_per_req", [1, 2]) |
| 987 | @pytest.mark.parametrize( |
| 988 | "req_lens", [[7], [8, 8], [1, 9], [63, 64, 65], [255, 256, 257], [256, 1, 2730]] |
| 989 | ) |
| 990 | def test_pcp_plan_deinterleave_restores_global_order( |
| 991 | monkeypatch, dcp_world_size, interleave, rows_per_req, req_lens |
| 992 | ): |
| 993 | """Gathering block-interleaved shards must restore each request's token order.""" |
| 994 | from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan |
| 995 | |
| 996 | monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False) |
| 997 | monkeypatch.setattr( |
| 998 | "vllm.v1.attention.backends.mla.indexer.async_tensor_h2d", |
| 999 | lambda x, device: x.to(device), |
| 1000 | ) |
| 1001 | scheduled = np.array(req_lens, dtype=np.int64) |
| 1002 | shard_rows = np.array( |
| 1003 | [ |
| 1004 | max( |
| 1005 | _local_count(g, r, dcp_world_size, interleave) |
| 1006 | for r in range(dcp_world_size) |
| 1007 | ) |
| 1008 | for g in req_lens |
| 1009 | ] |
| 1010 | ) |
| 1011 | rows = np.repeat(np.arange(len(scheduled)), rows_per_req) |
| 1012 | plan = build_pcp_global_chunk_plan( |
| 1013 | rows, shard_rows[rows], dcp_world_size, torch.device("cpu"), interleave |
| 1014 | ) |
| 1015 | |
| 1016 | starts = np.concatenate([[0], np.cumsum(scheduled)]) |
| 1017 | |
| 1018 | # Build each rank's padded shard exactly as the cache gather would. |
| 1019 | padded_cu = plan.padded_local_cu[::rows_per_req].tolist() |
| 1020 | shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1) |
| 1021 | for r in range(dcp_world_size): |
| 1022 | for i, g in enumerate(req_lens): |
| 1023 | owned = [t for t in range(g) if (t // interleave) % dcp_world_size == r] |
| 1024 | for local, t in enumerate(owned): |
| 1025 | shards[r, padded_cu[i] + local, 0] = starts[i] + t |
| 1026 | |
| 1027 | gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1) |
| 1028 | restored = gathered[plan.deinterleave_idx] |
| 1029 | # The layout is padded per request; only the real positions are read. |
| 1030 | row_start = plan.row_start_cu.tolist() |
| 1031 | for row, i in enumerate(rows): |
| 1032 | g = req_lens[i] |
| 1033 | torch.testing.assert_close( |
| 1034 | restored[row_start[row] : row_start[row] + g, 0], |
| 1035 | torch.arange(starts[i], starts[i] + g, dtype=torch.float32), |
| 1036 | ) |