| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import numpy as np |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.torch_utils import set_random_seed |
| 11 | |
| 12 | # Test parameters |
| 13 | NUM_ROWS = [1, 32, 2050] |
| 14 | TOP_K_VALUES = [2048, 3000] |
| 15 | BATCH_SIZE = [1, 2, 2048] |
| 16 | NEXT_N = [1, 8] |
| 17 | DATA_GENERATION = ["random", "10LSBits"] |
| 18 | RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 |
| 19 | |
| 20 | |
| 21 | def _has_device_capability(major: int) -> bool: |
| 22 | return current_platform.is_cuda() and current_platform.has_device_capability(major) |
| 23 | |
| 24 | |
| 25 | def _on_gfx950() -> bool: |
| 26 | if not current_platform.is_rocm(): |
| 27 | return False |
| 28 | from vllm.platforms.rocm import on_gfx950 |
| 29 | |
| 30 | return on_gfx950() |
| 31 | |
| 32 | |
| 33 | # DeepSelect is compiled for sm_100a/sm_103a only. |
| 34 | requires_sm100 = pytest.mark.skipif( |
| 35 | not current_platform.is_device_capability_family(100), |
| 36 | reason="DeepSelect requires SM100a/SM103a", |
| 37 | ) |
| 38 | requires_gfx950 = pytest.mark.skipif( |
| 39 | not _on_gfx950(), reason="This test exercises the gfx950 launch configuration" |
| 40 | ) |
| 41 | |
| 42 | |
| 43 | COOPERATIVE_TOPK_BACKEND = pytest.param( |
| 44 | "cooperative_topk", |
| 45 | marks=pytest.mark.skipif( |
| 46 | not _has_device_capability(90), |
| 47 | reason="cooperative_topk requires SM90+", |
| 48 | ), |
| 49 | ) |
| 50 | WORKSPACE_TOPK_BACKENDS = ["persistent_topk", COOPERATIVE_TOPK_BACKEND] |
| 51 | TOPK_BACKENDS = ["top_k_per_row_decode", *WORKSPACE_TOPK_BACKENDS] |
| 52 | |
| 53 | |
| 54 | def _run_topk_backend( |
| 55 | backend: str, |
| 56 | logits: torch.Tensor, |
| 57 | lengths: torch.Tensor, |
| 58 | indices: torch.Tensor, |
| 59 | top_k: int, |
| 60 | max_seq_len: int, |
| 61 | next_n: int = 1, |
| 62 | ) -> None: |
| 63 | if backend == "top_k_per_row_decode": |
| 64 | torch.ops._C.top_k_per_row_decode( |
| 65 | logits, |
| 66 | next_n, |
| 67 | lengths, |
| 68 | indices, |
| 69 | indices.shape[0], |
| 70 | logits.stride(0), |
| 71 | logits.stride(1), |
| 72 | top_k, |
| 73 | ) |
| 74 | elif backend == "persistent_topk": |
| 75 | workspace = torch.empty( |
| 76 | RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda" |
| 77 | ) |
| 78 | torch.ops._C.persistent_topk( |
| 79 | logits, lengths, indices, workspace, top_k, max_seq_len |
| 80 | ) |
| 81 | elif backend == "cooperative_topk": |
| 82 | if indices.shape[0] > 64: |
| 83 | pytest.skip( |
| 84 | "cooperative_topk supports <=64 rows; " |
| 85 | "persistent_topk covers larger batches" |
| 86 | ) |
| 87 | if logits.stride(0) % 4 != 0: |
| 88 | pytest.skip( |
| 89 | "cooperative_topk requires row stride divisible by 4; " |
| 90 | "persistent_topk covers unaligned strides" |
| 91 | ) |
| 92 | workspace = torch.empty( |
| 93 | RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda" |
| 94 | ) |
| 95 | torch.ops._C.cooperative_topk( |
| 96 | logits, lengths, indices, workspace, top_k, max_seq_len |
| 97 | ) |
| 98 | else: |
| 99 | raise ValueError(f"Unknown top-k backend: {backend}") |
| 100 | |
| 101 | |
| 102 | def create_random_logits( |
| 103 | row_starts: torch.Tensor, |
| 104 | row_ends: torch.Tensor, |
| 105 | dtype: torch.dtype, |
| 106 | seed: int, |
| 107 | clean_logits: bool, |
| 108 | data_generation: str, |
| 109 | ) -> torch.Tensor: |
| 110 | """Create random logits tensor for testing.""" |
| 111 | torch.manual_seed(seed) |
| 112 | np.random.seed(seed) |
| 113 | # Generate logits with some structure to make testing more meaningful |
| 114 | if data_generation == "random": |
| 115 | logits = torch.randn( |
| 116 | row_starts.shape[0], max(row_ends), dtype=dtype, device="cuda" |
| 117 | ) |
| 118 | elif data_generation == "10LSBits": |
| 119 | top_22_bits_mask = 0xFFFFFC00 |
| 120 | last_10_bits_mask = 0x000003FF |
| 121 | fixed_top_22_bits = 0x3F900000 |
| 122 | # Generate random bits for the last 10 bits |
| 123 | random_bottom_bits = torch.randint( |
| 124 | 0, |
| 125 | 2**10, |
| 126 | (row_starts.shape[0], max(row_ends)), |
| 127 | dtype=torch.int32, |
| 128 | device="cuda", |
| 129 | ) |
| 130 | # Combine: fixed top 22 bits with random last 10 bits |
| 131 | logits_bits = (fixed_top_22_bits & top_22_bits_mask) | ( |
| 132 | random_bottom_bits & last_10_bits_mask |
| 133 | ) |
| 134 | logits = logits_bits.view(dtype) |
| 135 | |
| 136 | if clean_logits: |
| 137 | for i, end in enumerate(row_ends): |
| 138 | logits[i, end:] = float("-inf") |
| 139 | return logits |
| 140 | |
| 141 | |
| 142 | def create_row_boundaries( |
| 143 | seq_len: int, vocab_size: int |
| 144 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 145 | """Create row start and end indices for testing.""" |
| 146 | row_starts = torch.zeros(seq_len, dtype=torch.int32, device="cuda") |
| 147 | row_ends = torch.arange(1, seq_len + 1, device="cuda", dtype=torch.int32) |
| 148 | return row_starts, row_ends |
| 149 | |
| 150 | |
| 151 | def compare_top_k_results( |
| 152 | logits: torch.Tensor, |
| 153 | cuda_indices: torch.Tensor, |
| 154 | torch_indices: torch.Tensor, |
| 155 | row_starts: torch.Tensor, |
| 156 | row_ends: torch.Tensor, |
| 157 | top_k: int, |
| 158 | tolerance: float = 1e-5, |
| 159 | ) -> bool: |
| 160 | """Compare results from CUDA top_k_per_row with torch.topk. |
| 161 | Both results should be sorted and contain the same top-k elements. |
| 162 | """ |
| 163 | num_rows = cuda_indices.shape[0] |
| 164 | |
| 165 | for row_idx in range(num_rows): |
| 166 | # Get valid elements using row boundaries |
| 167 | row_start = row_starts[row_idx].item() |
| 168 | row_end = row_ends[row_idx].item() |
| 169 | row_length = row_end - row_start |
| 170 | num_valid = min(top_k, row_length) |
| 171 | cuda_row_indices = cuda_indices[row_idx][:num_valid].cpu() |
| 172 | torch_row_indices = torch_indices[row_idx][:num_valid].cpu() |
| 173 | |
| 174 | # Compare the sets of indices first |
| 175 | cuda_set = set(cuda_row_indices.tolist()) |
| 176 | torch_set = set(torch_row_indices.tolist()) |
| 177 | if cuda_set == torch_set: |
| 178 | continue |
| 179 | |
| 180 | # Any difference in elements, compare the values |
| 181 | logits_row = logits[row_idx] |
| 182 | cuda_row_values = [logits_row[i] for i in cuda_row_indices] |
| 183 | torch_row_values = [logits_row[i] for i in torch_row_indices] |
| 184 | |
| 185 | cuda_only_values, torch_only_values = [], [] |
| 186 | for idx in cuda_set - torch_set: |
| 187 | cuda_pos = (cuda_row_indices == idx).nonzero(as_tuple=True)[0] |
| 188 | cuda_only_values.append(cuda_row_values[cuda_pos[0]]) |
| 189 | |
| 190 | for idx in torch_set - cuda_set: |
| 191 | torch_pos = (torch_row_indices == idx).nonzero(as_tuple=True)[0] |
| 192 | torch_only_values.append(torch_row_values[torch_pos[0]]) |
| 193 | |
| 194 | if len(cuda_only_values) != len(torch_only_values): |
| 195 | return False |
| 196 | if not torch.allclose( |
| 197 | torch.tensor(cuda_only_values), |
| 198 | torch.tensor(torch_only_values), |
| 199 | rtol=tolerance, |
| 200 | atol=tolerance, |
| 201 | ): |
| 202 | return False |
| 203 | |
| 204 | return True |
| 205 | |
| 206 | |
| 207 | def validate_topk_against_reference( |
| 208 | logits: torch.Tensor, |
| 209 | cuda_indices: torch.Tensor, |
| 210 | row_starts: torch.Tensor, |
| 211 | row_ends: torch.Tensor, |
| 212 | top_k: int, |
| 213 | kernel_name: str, |
| 214 | ) -> None: |
| 215 | """Validate CUDA top-k results against PyTorch reference implementation. |
| 216 | |
| 217 | Args: |
| 218 | logits: Input logits tensor |
| 219 | cuda_indices: CUDA kernel output indices |
| 220 | row_starts: Row start positions |
| 221 | row_ends: Row end positions |
| 222 | top_k: Number of top elements to select |
| 223 | kernel_name: Name of the kernel being tested (for error messages) |
| 224 | |
| 225 | """ |
| 226 | num_rows = cuda_indices.shape[0] |
| 227 | torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 228 | |
| 229 | for i in range(num_rows): |
| 230 | row_end = int(row_ends[i]) |
| 231 | k_i = min(top_k, row_end) |
| 232 | idx = logits[i, :row_end].topk(k_i, dim=-1)[1] |
| 233 | torch_indices[i, :k_i] = idx |
| 234 | |
| 235 | assert compare_top_k_results( |
| 236 | logits, cuda_indices, torch_indices, row_starts, row_ends, top_k |
| 237 | ), f"{kernel_name} results don't match torch.topk" |
| 238 | |
| 239 | |
| 240 | @pytest.mark.parametrize("num_rows", NUM_ROWS) |
| 241 | @pytest.mark.parametrize("top_k", TOP_K_VALUES) |
| 242 | @pytest.mark.parametrize("clean_logits", [True, False]) |
| 243 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 244 | @torch.inference_mode() |
| 245 | def test_top_k_per_row( |
| 246 | num_rows: int, |
| 247 | top_k: int, |
| 248 | clean_logits: bool, |
| 249 | ) -> None: |
| 250 | """Test top_k_per_row.""" |
| 251 | set_random_seed(0) |
| 252 | torch.set_default_device("cuda:0") |
| 253 | |
| 254 | # Create test data |
| 255 | vocab_size = 20000 |
| 256 | row_starts, row_ends = create_row_boundaries(num_rows, vocab_size) |
| 257 | logits = create_random_logits( |
| 258 | row_starts, row_ends, torch.float32, 42, clean_logits, "random" |
| 259 | ) |
| 260 | |
| 261 | # Create output tensors |
| 262 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 263 | |
| 264 | # Run CUDA implementation |
| 265 | torch.ops._C.top_k_per_row_prefill( |
| 266 | logits, |
| 267 | row_starts, |
| 268 | row_ends, |
| 269 | indices, |
| 270 | num_rows, |
| 271 | logits.stride(0), |
| 272 | logits.stride(1), |
| 273 | top_k, |
| 274 | ) |
| 275 | |
| 276 | # Run reference implementation |
| 277 | torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 278 | for i in range(num_rows): |
| 279 | row_end = int(row_ends[i]) |
| 280 | k_i = min(top_k, row_end) |
| 281 | idx = logits[i, :row_end].topk(k_i, dim=-1)[1] |
| 282 | torch_indices[i, :k_i] = idx |
| 283 | |
| 284 | # Compare results |
| 285 | assert compare_top_k_results( |
| 286 | logits, indices, torch_indices, row_starts, row_ends, top_k |
| 287 | ), "CUDA top_k_per_row_prefill results don't match torch.topk" |
| 288 | |
| 289 | |
| 290 | def _run_top_k_per_row_decode_test( |
| 291 | top_k: int, |
| 292 | batch_size: int, |
| 293 | next_n: int, |
| 294 | vocab_size: int, |
| 295 | clean_logits: bool, |
| 296 | data_generation: str, |
| 297 | ) -> None: |
| 298 | """Helper function to run top_k_per_row_decode test with given parameters.""" |
| 299 | torch.set_default_device("cuda:0") |
| 300 | |
| 301 | # Create test data |
| 302 | num_rows = batch_size * next_n |
| 303 | seq_lens = torch.randint( |
| 304 | low=next_n, |
| 305 | high=vocab_size, |
| 306 | size=(batch_size,), |
| 307 | dtype=torch.int32, |
| 308 | device="cuda", |
| 309 | ) |
| 310 | row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") |
| 311 | row_indices = torch.arange(num_rows, device="cuda") // next_n |
| 312 | next_n_offset = torch.arange(num_rows, device="cuda") % next_n |
| 313 | row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 |
| 314 | logits = create_random_logits( |
| 315 | row_starts, row_ends, torch.float32, 42, clean_logits, data_generation |
| 316 | ) |
| 317 | |
| 318 | # Create output tensors |
| 319 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 320 | |
| 321 | # Run CUDA implementation |
| 322 | torch.ops._C.top_k_per_row_decode( |
| 323 | logits, |
| 324 | next_n, |
| 325 | seq_lens, |
| 326 | indices, |
| 327 | num_rows, |
| 328 | logits.stride(0), |
| 329 | logits.stride(1), |
| 330 | top_k, |
| 331 | ) |
| 332 | |
| 333 | torch.accelerator.synchronize() |
| 334 | |
| 335 | # Run reference implementation |
| 336 | torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 337 | for i in range(num_rows): |
| 338 | row_end = int(row_ends[i]) |
| 339 | k_i = min(top_k, row_end) |
| 340 | idx = logits[i, :row_end].topk(k_i, dim=-1)[1] |
| 341 | torch_indices[i, :k_i] = idx |
| 342 | |
| 343 | # Compare results |
| 344 | assert compare_top_k_results( |
| 345 | logits, indices, torch_indices, row_starts, row_ends, top_k |
| 346 | ), "CUDA top_k_per_row_decode results don't match torch.topk" |
| 347 | |
| 348 | |
| 349 | @pytest.mark.parametrize("top_k", TOP_K_VALUES) |
| 350 | @pytest.mark.parametrize("batch_size", BATCH_SIZE) |
| 351 | @pytest.mark.parametrize("next_n", NEXT_N) |
| 352 | @pytest.mark.parametrize("clean_logits", [True, False]) |
| 353 | @pytest.mark.parametrize("data_generation", DATA_GENERATION) |
| 354 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 355 | @torch.inference_mode() |
| 356 | def test_top_k_per_row_decode( |
| 357 | top_k: int, |
| 358 | batch_size: int, |
| 359 | next_n: int, |
| 360 | clean_logits: bool, |
| 361 | data_generation: str, |
| 362 | ) -> None: |
| 363 | """Test top_k_per_row with seq_lens tensor.""" |
| 364 | set_random_seed(0) |
| 365 | vocab_size = 20000 |
| 366 | _run_top_k_per_row_decode_test( |
| 367 | top_k, batch_size, next_n, vocab_size, clean_logits, data_generation |
| 368 | ) |
| 369 | |
| 370 | |
| 371 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 372 | @pytest.mark.parametrize("clean_logits", [True, False]) |
| 373 | @torch.inference_mode() |
| 374 | def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None: |
| 375 | """Test top_k_per_row_decode with large vocabulary size.""" |
| 376 | set_random_seed(0) |
| 377 | top_k = 2048 |
| 378 | batch_size = 2 |
| 379 | next_n = 2 |
| 380 | vocab_size = 300000 |
| 381 | data_generation = "random" |
| 382 | _run_top_k_per_row_decode_test( |
| 383 | top_k, batch_size, next_n, vocab_size, clean_logits, data_generation |
| 384 | ) |
| 385 | |
| 386 | |
| 387 | @pytest.mark.parametrize( |
| 388 | "batch_size,num_speculative_tokens,min_seq_len,max_seq_len", |
| 389 | [ |
| 390 | pytest.param(1, 2, 125_000, 250_000, id="eight-splits"), |
| 391 | pytest.param(7, 4, 125_000, 250_000, id="four-splits"), |
| 392 | pytest.param(16, 4, 2_500, 2_501, id="dynamic-short-rows"), |
| 393 | pytest.param(13, 4, 25_000, 25_001, id="dynamic-three-splits"), |
| 394 | pytest.param(16, 4, 65_535, 65_536, id="dynamic-boundary-mixed"), |
| 395 | pytest.param(16, 4, 125_000, 125_001, id="dynamic-two-splits"), |
| 396 | pytest.param(64, 3, 125_000, 125_001, id="dynamic-two-splits-max-rows"), |
| 397 | pytest.param(64, 4, 125_000, 125_001, id="baseline-fallback"), |
| 398 | ], |
| 399 | ) |
| 400 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") |
| 401 | @torch.inference_mode() |
| 402 | def test_top_k_per_row_decode_gfx950_long_c4a( |
| 403 | batch_size: int, |
| 404 | num_speculative_tokens: int, |
| 405 | min_seq_len: int, |
| 406 | max_seq_len: int, |
| 407 | ) -> None: |
| 408 | properties = torch.cuda.get_device_properties(0) |
| 409 | if not properties.gcnArchName.startswith("gfx950"): |
| 410 | pytest.skip("This test exercises the gfx950 launch configuration") |
| 411 | |
| 412 | query_tokens = num_speculative_tokens + 1 |
| 413 | num_rows = batch_size * query_tokens |
| 414 | stride = 262_144 |
| 415 | top_k = 1024 |
| 416 | offsets = torch.arange(num_rows, dtype=torch.int32, device="cuda") |
| 417 | seq_lens = min_seq_len + offsets * (max_seq_len - min_seq_len) // max( |
| 418 | num_rows - 1, 1 |
| 419 | ) |
| 420 | seq_lens = seq_lens.reshape(batch_size, query_tokens) |
| 421 | logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda") |
| 422 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 423 | |
| 424 | torch.ops._C.top_k_per_row_decode( |
| 425 | logits, |
| 426 | query_tokens, |
| 427 | seq_lens, |
| 428 | indices, |
| 429 | num_rows, |
| 430 | logits.stride(0), |
| 431 | logits.stride(1), |
| 432 | top_k, |
| 433 | ) |
| 434 | |
| 435 | row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") |
| 436 | validate_topk_against_reference( |
| 437 | logits, |
| 438 | indices, |
| 439 | row_starts, |
| 440 | seq_lens.reshape(-1), |
| 441 | top_k, |
| 442 | "gfx950 long C4A top-k", |
| 443 | ) |
| 444 | |
| 445 | |
| 446 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") |
| 447 | @torch.inference_mode() |
| 448 | def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None: |
| 449 | properties = torch.cuda.get_device_properties(0) |
| 450 | if not properties.gcnArchName.startswith("gfx950"): |
| 451 | pytest.skip("This test exercises the gfx950 launch configuration") |
| 452 | |
| 453 | batch_size = 20 |
| 454 | query_tokens = 4 |
| 455 | num_rows = batch_size * query_tokens |
| 456 | stride = 262_144 |
| 457 | top_k = 1024 |
| 458 | seq_lens = torch.linspace( |
| 459 | 125_003, |
| 460 | 250_000, |
| 461 | batch_size, |
| 462 | dtype=torch.int32, |
| 463 | device="cuda", |
| 464 | ) |
| 465 | row_offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda") |
| 466 | row_ends = (seq_lens[:, None] - query_tokens + row_offsets + 1).reshape(-1) |
| 467 | logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda") |
| 468 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 469 | |
| 470 | torch.ops._C.top_k_per_row_decode( |
| 471 | logits, |
| 472 | query_tokens, |
| 473 | seq_lens, |
| 474 | indices, |
| 475 | num_rows, |
| 476 | logits.stride(0), |
| 477 | logits.stride(1), |
| 478 | top_k, |
| 479 | ) |
| 480 | |
| 481 | row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") |
| 482 | validate_topk_against_reference( |
| 483 | logits, |
| 484 | indices, |
| 485 | row_starts, |
| 486 | row_ends, |
| 487 | top_k, |
| 488 | "gfx950 long C4A top-k with 1-D sequence lengths", |
| 489 | ) |
| 490 | |
| 491 | |
| 492 | def _assert_exact_topk( |
| 493 | logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor |
| 494 | ) -> None: |
| 495 | ends = row_ends.reshape(-1, 1).clamp_min(0) |
| 496 | valid = (indices >= 0) & (indices < ends) |
| 497 | assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten()) |
| 498 | assert torch.all(valid | (indices == -1)) |
| 499 | ordered = indices.sort(1).values |
| 500 | assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any() |
| 501 | |
| 502 | width = int(ends.max()) |
| 503 | visible = logits[:, :width].clone() |
| 504 | columns = torch.arange(width, device=logits.device) |
| 505 | visible.masked_fill_(columns[None, :] >= ends, float("-inf")) |
| 506 | expected = visible.topk(min(indices.shape[1], width), dim=1).values |
| 507 | selected = logits.gather(1, indices.clamp_min(0).long()) |
| 508 | selected.masked_fill_(~valid, float("-inf")) |
| 509 | selected = selected.sort(1, descending=True).values |
| 510 | assert torch.equal(selected[:, : expected.shape[1]], expected) |
| 511 | |
| 512 | |
| 513 | @requires_gfx950 |
| 514 | @torch.inference_mode() |
| 515 | def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None: |
| 516 | """Honor causal offsets and exact compressed lengths, including short rows.""" |
| 517 | next_n, width, top_k = 6, 16_384, 512 |
| 518 | lengths = torch.tensor( |
| 519 | [0, 1, 511, 512, 513, 1023, 1024, 16_000], |
| 520 | dtype=torch.int32, |
| 521 | device="cuda", |
| 522 | ) |
| 523 | offsets = torch.arange(next_n, dtype=torch.int32, device="cuda") |
| 524 | row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten() |
| 525 | rows = row_ends.numel() |
| 526 | storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda") |
| 527 | logits = storage[:, 1 : width + 1] |
| 528 | logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda") |
| 529 | indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") |
| 530 | _run_topk_backend( |
| 531 | "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n |
| 532 | ) |
| 533 | _assert_exact_topk(logits, indices, row_ends) |
| 534 | |
| 535 | |
| 536 | @pytest.mark.parametrize( |
| 537 | ("rows", "row_length"), |
| 538 | [ |
| 539 | pytest.param(128, 65_536, id="single-before-64k"), |
| 540 | pytest.param(128, 65_537, id="four-splits-after-64k"), |
| 541 | pytest.param(129, 262_144, id="single-before-three-split-length"), |
| 542 | pytest.param(129, 262_145, id="three-splits-after-256k"), |
| 543 | pytest.param(160, 262_145, id="three-splits-upper-row-bound"), |
| 544 | pytest.param(161, 262_144, id="single-before-two-split-length"), |
| 545 | pytest.param(161, 262_145, id="two-splits-after-256k"), |
| 546 | pytest.param(256, 262_145, id="two-splits-upper-row-bound"), |
| 547 | pytest.param(257, 262_145, id="large-row-single-block"), |
| 548 | pytest.param(384, 262_145, id="optimized-row-limit"), |
| 549 | pytest.param(385, 262_145, id="native-fallback"), |
| 550 | ], |
| 551 | ) |
| 552 | @requires_gfx950 |
| 553 | @torch.inference_mode() |
| 554 | def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries( |
| 555 | rows: int, row_length: int |
| 556 | ) -> None: |
| 557 | width, top_k = 262_145, 512 |
| 558 | row = torch.arange(width, dtype=torch.float32, device="cuda") |
| 559 | logits = row.expand(rows, -1) |
| 560 | lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda") |
| 561 | indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") |
| 562 | |
| 563 | _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width) |
| 564 | |
| 565 | expected = torch.arange( |
| 566 | row_length - top_k, row_length, dtype=torch.int32, device="cuda" |
| 567 | ).expand_as(indices) |
| 568 | assert torch.equal(indices.sort(1).values, expected) |
| 569 | |
| 570 | |
| 571 | @requires_gfx950 |
| 572 | @torch.inference_mode() |
| 573 | def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: |
| 574 | """Preserve visible top-k scores when masks and lengths change on replay.""" |
| 575 | rows, next_n, width, top_k = 24, 6, 262_145, 512 |
| 576 | logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda") |
| 577 | lengths = torch.full( |
| 578 | (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda" |
| 579 | ) |
| 580 | indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") |
| 581 | logits[:, :70_000] = float("-inf") |
| 582 | bin_sizes = [128, 129, 1024, 1025] |
| 583 | for row in range(rows): |
| 584 | pattern = row % 9 |
| 585 | if pattern < 4: |
| 586 | count = [0, 1, 511, 513][pattern] |
| 587 | logits[row, :count] = torch.arange( |
| 588 | count, dtype=torch.float32, device="cuda" |
| 589 | ) |
| 590 | elif pattern == 4: |
| 591 | logits[row, :511] = 0 |
| 592 | logits[row, 511] = torch.finfo(torch.float32).min |
| 593 | else: |
| 594 | count = bin_sizes[pattern - 5] |
| 595 | higher = max(0, top_k - count // 2) |
| 596 | logits[row, :higher] = 2 |
| 597 | logits[row, higher : higher + count] = 1 |
| 598 | logits[row, higher] = 1 + 2**-23 |
| 599 | |
| 600 | def run() -> None: |
| 601 | _run_topk_backend( |
| 602 | "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n |
| 603 | ) |
| 604 | |
| 605 | run() |
| 606 | _assert_exact_topk(logits, indices, lengths) |
| 607 | torch.accelerator.synchronize() |
| 608 | graph = torch.cuda.CUDAGraph() |
| 609 | with torch.cuda.graph(graph): |
| 610 | run() |
| 611 | indices.fill_(-777) |
| 612 | graph.replay() |
| 613 | _assert_exact_topk(logits, indices, lengths) |
| 614 | |
| 615 | bounds = [ |
| 616 | 0, |
| 617 | 511, |
| 618 | 512, |
| 619 | 513, |
| 620 | 16_384, |
| 621 | 16_385, |
| 622 | 131_072, |
| 623 | 131_073, |
| 624 | 262_144, |
| 625 | 262_145, |
| 626 | ] |
| 627 | row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda") |
| 628 | repeats = (rows + len(bounds) - 1) // len(bounds) |
| 629 | lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows]) |
| 630 | logits[:] = torch.arange(width, dtype=torch.float32, device="cuda") |
| 631 | indices.fill_(-777) |
| 632 | graph.replay() |
| 633 | _assert_exact_topk(logits, indices, lengths) |
| 634 | |
| 635 | |
| 636 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") |
| 637 | @torch.inference_mode() |
| 638 | def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: |
| 639 | from vllm._aiter_ops import rocm_aiter_ops |
| 640 | from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( |
| 641 | _launch_aiter_top_k_per_row_prefill, |
| 642 | ) |
| 643 | |
| 644 | row_starts = torch.tensor([1, 6, 20], dtype=torch.int32, device="cuda") |
| 645 | row_ends = torch.tensor([3, 18, 23], dtype=torch.int32, device="cuda") |
| 646 | logits = torch.arange(96, dtype=torch.float32, device="cuda").reshape(3, 32) |
| 647 | top_k = 4 |
| 648 | indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda") |
| 649 | |
| 650 | if not rocm_aiter_ops.is_indexer_top_k_supported( |
| 651 | is_prefill=True, |
| 652 | compress_ratio=4, |
| 653 | num_rows=logits.shape[0], |
| 654 | ): |
| 655 | pytest.skip("AITER top-k is unavailable") |
| 656 | assert ( |
| 657 | _launch_aiter_top_k_per_row_prefill( |
| 658 | logits, |
| 659 | row_starts, |
| 660 | row_ends, |
| 661 | indices, |
| 662 | top_k, |
| 663 | ) |
| 664 | is None |
| 665 | ) |
| 666 | |
| 667 | for row_idx, (row_start, row_end) in enumerate( |
| 668 | zip(row_starts.tolist(), row_ends.tolist()) |
| 669 | ): |
| 670 | num_valid = min(top_k, row_end - row_start) |
| 671 | expected = logits[row_idx, row_start:row_end].topk(num_valid).indices |
| 672 | actual = indices[row_idx, :num_valid] |
| 673 | assert set(actual.tolist()) == set(expected.tolist()) |
| 674 | assert torch.all(indices[row_idx, num_valid:] == -1) |
| 675 | |
| 676 | |
| 677 | @pytest.mark.parametrize("num_speculative_tokens", [2, 3, 4]) |
| 678 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") |
| 679 | @torch.inference_mode() |
| 680 | def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( |
| 681 | num_speculative_tokens: int, |
| 682 | ) -> None: |
| 683 | from vllm._aiter_ops import rocm_aiter_ops |
| 684 | |
| 685 | query_tokens = num_speculative_tokens + 1 |
| 686 | offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda") |
| 687 | final_uncompressed_lens = torch.tensor([43, 59], dtype=torch.int32, device="cuda") |
| 688 | seq_lens = (final_uncompressed_lens[:, None] - query_tokens + offsets + 1) // 4 |
| 689 | seq_lens[0, 0] = 0 |
| 690 | |
| 691 | num_rows = seq_lens.numel() |
| 692 | logits = torch.arange(num_rows * 16, dtype=torch.float32, device="cuda").reshape( |
| 693 | num_rows, 16 |
| 694 | ) |
| 695 | top_k = 4 |
| 696 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 697 | |
| 698 | if not rocm_aiter_ops.is_indexer_top_k_supported( |
| 699 | is_prefill=False, |
| 700 | compress_ratio=4, |
| 701 | num_rows=num_rows, |
| 702 | max_valid_seq_len=seq_lens.max().item(), |
| 703 | ): |
| 704 | pytest.skip("AITER top-k is unavailable") |
| 705 | assert ( |
| 706 | rocm_aiter_ops.indexer_top_k_decode( |
| 707 | logits, |
| 708 | 1, |
| 709 | seq_lens.reshape(-1), |
| 710 | indices, |
| 711 | top_k, |
| 712 | ) |
| 713 | is None |
| 714 | ) |
| 715 | |
| 716 | for row_idx, row_end in enumerate(seq_lens.reshape(-1).tolist()): |
| 717 | num_valid = min(top_k, row_end) |
| 718 | expected = logits[row_idx, :row_end].topk(num_valid).indices |
| 719 | actual = indices[row_idx, :num_valid] |
| 720 | assert set(actual.tolist()) == set(expected.tolist()) |
| 721 | assert torch.all(indices[row_idx, num_valid:] == -1) |
| 722 | |
| 723 | |
| 724 | @requires_gfx950 |
| 725 | @torch.inference_mode() |
| 726 | def test_aiter_decode_topk_matches_per_row_under_mtp() -> None: |
| 727 | """The aiter backend expands (B, 1) seq_lens into per-row ends before |
| 728 | launching, so next_n > 1 must agree with the in-tree kernel, including |
| 729 | the rows that clamp to a zero-length prefix. |
| 730 | """ |
| 731 | from vllm.model_executor.layers.indexer_topk import get_indexer_topk |
| 732 | |
| 733 | next_n = 4 |
| 734 | top_k = 512 |
| 735 | num_columns = 1024 |
| 736 | seq_lens = torch.tensor([[2], [513], [700]], dtype=torch.int32, device="cuda") |
| 737 | num_rows = seq_lens.shape[0] * next_n |
| 738 | |
| 739 | torch.manual_seed(0) |
| 740 | logits = torch.randn(num_rows, num_columns, dtype=torch.float32, device="cuda") |
| 741 | |
| 742 | outputs = {} |
| 743 | for backend in ("aiter", "per_row"): |
| 744 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 745 | get_indexer_topk(backend)(logits, seq_lens, next_n, indices, top_k, num_columns) |
| 746 | outputs[backend] = indices |
| 747 | |
| 748 | row_ends = ( |
| 749 | (seq_lens - next_n + 1 + torch.arange(next_n, device="cuda")) |
| 750 | .clamp_(min=0) |
| 751 | .reshape(-1) |
| 752 | .tolist() |
| 753 | ) |
| 754 | assert row_ends[:next_n] == [0, 0, 1, 2] |
| 755 | |
| 756 | for row_idx, row_end in enumerate(row_ends): |
| 757 | num_valid = min(top_k, row_end) |
| 758 | aiter_row = outputs["aiter"][row_idx] |
| 759 | assert set(aiter_row[:num_valid].tolist()) == set( |
| 760 | outputs["per_row"][row_idx][:num_valid].tolist() |
| 761 | ) |
| 762 | assert torch.all(aiter_row[num_valid:] == -1) |
| 763 | |
| 764 | |
| 765 | def _enable_aiter_topk(monkeypatch, enabled: bool = True) -> None: |
| 766 | from vllm._aiter_ops import rocm_aiter_ops |
| 767 | |
| 768 | monkeypatch.setattr( |
| 769 | rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: enabled, raising=True |
| 770 | ) |
| 771 | |
| 772 | |
| 773 | def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: |
| 774 | from vllm._aiter_ops import rocm_aiter_ops |
| 775 | |
| 776 | _enable_aiter_topk(monkeypatch) |
| 777 | is_supported = rocm_aiter_ops.is_indexer_top_k_supported |
| 778 | |
| 779 | eligible_cases = [ |
| 780 | dict(is_prefill=True, compress_ratio=4, num_rows=1), |
| 781 | dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_536), |
| 782 | dict( |
| 783 | is_prefill=False, |
| 784 | compress_ratio=4, |
| 785 | num_rows=257, |
| 786 | max_valid_seq_len=250_000, |
| 787 | ), |
| 788 | dict( |
| 789 | is_prefill=False, |
| 790 | compress_ratio=2, |
| 791 | num_rows=384, |
| 792 | max_valid_seq_len=1_048_576, |
| 793 | ), |
| 794 | ] |
| 795 | for kwargs in eligible_cases: |
| 796 | assert is_supported(**kwargs) is True |
| 797 | |
| 798 | ineligible_cases = [ |
| 799 | dict(is_prefill=True, compress_ratio=1, num_rows=1), |
| 800 | dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_537), |
| 801 | dict( |
| 802 | is_prefill=False, |
| 803 | compress_ratio=4, |
| 804 | num_rows=256, |
| 805 | max_valid_seq_len=125_000, |
| 806 | ), |
| 807 | ] |
| 808 | for kwargs in ineligible_cases: |
| 809 | assert is_supported(**kwargs) is False |
| 810 | |
| 811 | _enable_aiter_topk(monkeypatch, enabled=False) |
| 812 | assert is_supported(is_prefill=True, compress_ratio=4, num_rows=1) is False |
| 813 | |
| 814 | |
| 815 | def test_dsv4_native_top_k_window() -> None: |
| 816 | """The in-tree decode kernel's measured advantage band over AITER, which |
| 817 | applies to the DSV4 indexer only.""" |
| 818 | from vllm._aiter_ops import rocm_aiter_ops |
| 819 | |
| 820 | prefers_native = rocm_aiter_ops.dsv4_indexer_prefers_native_top_k |
| 821 | assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=512) is True |
| 822 | assert prefers_native(num_rows=385, num_columns=524_288, topk_tokens=512) is False |
| 823 | assert prefers_native(num_rows=128, num_columns=1_048_577, topk_tokens=512) is False |
| 824 | assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=1024) is False |
| 825 | |
| 826 | |
| 827 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 828 | @pytest.mark.parametrize( |
| 829 | "seq_len_range,test_id", |
| 830 | [ |
| 831 | pytest.param((4000, 8000), "short_sequences", id="short"), |
| 832 | pytest.param((8000, 32000), "medium_sequences", id="medium"), |
| 833 | pytest.param((32000, 163840), "long_sequences", id="long"), |
| 834 | ], |
| 835 | ) |
| 836 | @pytest.mark.parametrize("clean_logits", [True, False]) |
| 837 | @pytest.mark.parametrize("top_k", [2048]) |
| 838 | @pytest.mark.parametrize("next_n", [1, 4]) |
| 839 | @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) |
| 840 | @torch.inference_mode() |
| 841 | def test_deepseek_workspace_topk( |
| 842 | seq_len_range: tuple[int, int], |
| 843 | test_id: str, |
| 844 | clean_logits: bool, |
| 845 | top_k: int, |
| 846 | next_n: int, |
| 847 | backend: str, |
| 848 | ) -> None: |
| 849 | """Test workspace top-k backends with varying sequence lengths and speculative |
| 850 | decoding. |
| 851 | Supports speculative decoding with next_n > 1. |
| 852 | """ |
| 853 | set_random_seed(42 if test_id == "short_sequences" else 43) |
| 854 | torch.set_default_device("cuda:0") |
| 855 | |
| 856 | batch_size = 4 |
| 857 | num_rows = batch_size * next_n |
| 858 | |
| 859 | seq_lens = torch.randint( |
| 860 | seq_len_range[0], |
| 861 | seq_len_range[1], |
| 862 | (batch_size,), |
| 863 | dtype=torch.int32, |
| 864 | device="cuda", |
| 865 | ) |
| 866 | seq_lens = (seq_lens + 3) & ~3 # align to 4 for TMA |
| 867 | |
| 868 | # Compute row boundaries for speculative decoding |
| 869 | row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") |
| 870 | row_indices = torch.arange(num_rows, device="cuda") // next_n |
| 871 | next_n_offset = torch.arange(num_rows, device="cuda") % next_n |
| 872 | row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 |
| 873 | |
| 874 | logits = create_random_logits( |
| 875 | row_starts, row_ends, torch.float32, 42, clean_logits, "random" |
| 876 | ) |
| 877 | |
| 878 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 879 | |
| 880 | if next_n == 1: |
| 881 | lengths = seq_lens |
| 882 | else: |
| 883 | offsets = torch.arange(next_n, device=logits.device, dtype=torch.int32) |
| 884 | lengths = (seq_lens.unsqueeze(1) - next_n + 1 + offsets).flatten() |
| 885 | |
| 886 | max_seq_len = int(seq_lens.max().item()) |
| 887 | _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len, next_n) |
| 888 | |
| 889 | validate_topk_against_reference( |
| 890 | logits, indices, row_starts, row_ends, top_k, f"{backend} ({test_id})" |
| 891 | ) |
| 892 | |
| 893 | |
| 894 | @pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+") |
| 895 | @pytest.mark.parametrize("rows", [65, 128]) |
| 896 | @pytest.mark.parametrize("top_k", [512, 1024, 2048]) |
| 897 | @pytest.mark.parametrize( |
| 898 | "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"] |
| 899 | ) |
| 900 | @torch.inference_mode() |
| 901 | def test_persistent_topk_sampled_graph( |
| 902 | rows: int, top_k: int, distribution: str |
| 903 | ) -> None: |
| 904 | """Sampling and both fallbacks preserve exact values across graph replays.""" |
| 905 | if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024: |
| 906 | pytest.skip("Sampled top-k requires at least 144 KiB of shared memory") |
| 907 | set_random_seed(42) |
| 908 | min_sampled_length = 98304 if top_k == 512 else 65536 |
| 909 | width = min_sampled_length + 3 |
| 910 | lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda") |
| 911 | logits = create_random_logits( |
| 912 | torch.zeros_like(lengths), |
| 913 | lengths, |
| 914 | torch.float32, |
| 915 | 42, |
| 916 | False, |
| 917 | "10LSBits" if distribution == "10LSBits" else "random", |
| 918 | ) |
| 919 | if distribution == "ascending": |
| 920 | logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32)) |
| 921 | elif distribution == "constant": |
| 922 | logits.fill_(1.0) |
| 923 | elif distribution == "sampled_peaks": |
| 924 | # A biased sample leaves fewer than k survivors and must fall back. |
| 925 | logits.zero_() |
| 926 | sample = torch.arange(4096, device="cuda") |
| 927 | positions = (sample // 32) * (width // 128) + sample % 32 |
| 928 | logits[:, positions] = 1 + sample.float() / 4096 |
| 929 | original = logits.clone() |
| 930 | indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") |
| 931 | workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") |
| 932 | |
| 933 | def run() -> None: |
| 934 | torch.ops._C.persistent_topk( |
| 935 | logits, |
| 936 | lengths.view(-1, 4 if rows == 128 else 1), |
| 937 | indices, |
| 938 | workspace, |
| 939 | top_k, |
| 940 | width, |
| 941 | ) |
| 942 | |
| 943 | run() |
| 944 | graph = torch.cuda.CUDAGraph() |
| 945 | with torch.cuda.graph(graph): |
| 946 | run() |
| 947 | positions = torch.arange(width, device="cuda") |
| 948 | slots = torch.arange(top_k, device="cuda") |
| 949 | for step in range(3): |
| 950 | bounds = torch.tensor( |
| 951 | [ |
| 952 | -1, |
| 953 | 0, |
| 954 | 1, |
| 955 | top_k - 1, |
| 956 | top_k, |
| 957 | 32769, |
| 958 | min_sampled_length - 1, |
| 959 | min_sampled_length, |
| 960 | width, |
| 961 | width + 17, |
| 962 | ], |
| 963 | dtype=torch.int32, |
| 964 | device="cuda", |
| 965 | ) |
| 966 | lengths.copy_( |
| 967 | bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()] |
| 968 | ) |
| 969 | logits.copy_(original if step != 1 else original.flip(1)) |
| 970 | logits.masked_fill_(positions[None] >= lengths[:, None], float("nan")) |
| 971 | indices.fill_(-2) |
| 972 | graph.replay() |
| 973 | valid = slots[None] < lengths.clamp(max=top_k)[:, None] |
| 974 | assert torch.all(indices[~valid] == -1) |
| 975 | assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid) |
| 976 | ordered_indices = indices.sort(dim=1).values |
| 977 | assert torch.all( |
| 978 | (ordered_indices[:, 1:] != ordered_indices[:, :-1]) |
| 979 | | (ordered_indices[:, 1:] == -1) |
| 980 | ) |
| 981 | selected = logits.gather(1, indices.clamp_min(0).long()) |
| 982 | selected.masked_fill_(~valid, -float("inf")) |
| 983 | expected = logits.masked_fill( |
| 984 | positions[None] >= lengths[:, None], -float("inf") |
| 985 | ) |
| 986 | torch.testing.assert_close( |
| 987 | selected.sort(dim=1, descending=True).values, |
| 988 | expected.topk(top_k, dim=1).values, |
| 989 | atol=0, |
| 990 | rtol=0, |
| 991 | ) |
| 992 | |
| 993 | |
| 994 | def run_large_context_topk_test( |
| 995 | batch_size: int, |
| 996 | seq_lens: list[int], |
| 997 | top_k: int, |
| 998 | data_type: str = "random", |
| 999 | seed: int = 42, |
| 1000 | backend: str = "cooperative_topk", |
| 1001 | ) -> None: |
| 1002 | """Helper to run a top-k backend test with given parameters. |
| 1003 | |
| 1004 | Args: |
| 1005 | batch_size: Number of rows/sequences |
| 1006 | seq_lens: List of sequence lengths (one per row) |
| 1007 | top_k: Number of top elements to select |
| 1008 | data_type: Type of test data to generate |
| 1009 | seed: Random seed for reproducibility |
| 1010 | backend: Top-k backend to test |
| 1011 | |
| 1012 | """ |
| 1013 | torch.set_default_device("cuda:0") |
| 1014 | set_random_seed(seed) |
| 1015 | |
| 1016 | # Create test data |
| 1017 | num_rows = batch_size |
| 1018 | max_len = max(seq_lens) |
| 1019 | lengths = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") |
| 1020 | |
| 1021 | if data_type == "random": |
| 1022 | logits = torch.randn(num_rows, max_len, dtype=torch.float32, device="cuda") |
| 1023 | elif data_type == "sorted_asc": |
| 1024 | # Each row gets its own ascending sequence based on its length |
| 1025 | logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda") |
| 1026 | for i, length in enumerate(seq_lens): |
| 1027 | logits[i, :length] = torch.arange( |
| 1028 | length, dtype=torch.float32, device="cuda" |
| 1029 | ) |
| 1030 | if length < max_len: |
| 1031 | logits[i, length:] = float("-inf") |
| 1032 | elif data_type == "sorted_desc": |
| 1033 | # Each row gets its own descending sequence based on its length |
| 1034 | logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda") |
| 1035 | for i, length in enumerate(seq_lens): |
| 1036 | logits[i, :length] = torch.arange( |
| 1037 | length, 0, -1, dtype=torch.float32, device="cuda" |
| 1038 | ) |
| 1039 | if length < max_len: |
| 1040 | logits[i, length:] = float("-inf") |
| 1041 | elif data_type == "all_same": |
| 1042 | logits = torch.ones(num_rows, max_len, dtype=torch.float32, device="cuda") |
| 1043 | for i, length in enumerate(seq_lens): |
| 1044 | if length < max_len: |
| 1045 | logits[i, length:] = float("-inf") |
| 1046 | elif data_type == "many_ties": |
| 1047 | # Only 10 unique values, many duplicates |
| 1048 | logits = torch.randint(0, 10, (num_rows, max_len), device="cuda").float() / 10.0 |
| 1049 | for i, length in enumerate(seq_lens): |
| 1050 | if length < max_len: |
| 1051 | logits[i, length:] = float("-inf") |
| 1052 | elif data_type == "small_differences": |
| 1053 | # Very small differences to test float precision |
| 1054 | base = torch.randn(num_rows, max_len, dtype=torch.float32, device="cuda") |
| 1055 | noise = ( |
| 1056 | torch.randn(num_rows, max_len, dtype=torch.float32, device="cuda") * 1e-6 |
| 1057 | ) |
| 1058 | logits = base + noise |
| 1059 | for i, length in enumerate(seq_lens): |
| 1060 | if length < max_len: |
| 1061 | logits[i, length:] = float("-inf") |
| 1062 | else: |
| 1063 | raise ValueError(f"Unknown data_type: {data_type}") |
| 1064 | |
| 1065 | # Create output tensor |
| 1066 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 1067 | |
| 1068 | max_seq_len = max(seq_lens) |
| 1069 | _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len) |
| 1070 | |
| 1071 | torch.accelerator.synchronize() |
| 1072 | |
| 1073 | torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 1074 | for i in range(num_rows): |
| 1075 | length = seq_lens[i] |
| 1076 | k_i = min(top_k, length) |
| 1077 | if k_i > 0: |
| 1078 | idx = logits[i, :length].topk(k_i, dim=-1)[1] |
| 1079 | torch_indices[i, :k_i] = idx |
| 1080 | if k_i < top_k: |
| 1081 | torch_indices[i, k_i:] = -1 |
| 1082 | else: |
| 1083 | torch_indices[i, :] = -1 |
| 1084 | |
| 1085 | # Compare results |
| 1086 | for i in range(num_rows): |
| 1087 | length = seq_lens[i] |
| 1088 | k_i = min(top_k, length) |
| 1089 | |
| 1090 | if k_i == 0: |
| 1091 | continue |
| 1092 | |
| 1093 | cuda_row = indices[i, :k_i].cpu() |
| 1094 | torch_row = torch_indices[i, :k_i].cpu() |
| 1095 | |
| 1096 | # Filter out -1 padding values from cuda_row |
| 1097 | valid_mask = cuda_row >= 0 |
| 1098 | cuda_row = cuda_row[valid_mask] |
| 1099 | |
| 1100 | # Compare sets (order may differ for ties) |
| 1101 | cuda_set = set(cuda_row.tolist()) |
| 1102 | torch_set = set(torch_row.tolist()) |
| 1103 | |
| 1104 | if cuda_set == torch_set: |
| 1105 | continue |
| 1106 | |
| 1107 | # If sets differ, check if it's due to equal values (ties) |
| 1108 | cuda_vals = logits[i, cuda_row].cpu() |
| 1109 | torch_vals = logits[i, torch_row].cpu() |
| 1110 | |
| 1111 | # Check that min CUDA value >= max of values NOT in top-k |
| 1112 | if k_i < length: |
| 1113 | non_topk_indices = torch.tensor( |
| 1114 | list(set(range(length)) - cuda_set), dtype=torch.int32 |
| 1115 | ) |
| 1116 | if len(non_topk_indices) > 0: |
| 1117 | non_topk_vals = logits[i, non_topk_indices].cpu() |
| 1118 | min_cuda_val = cuda_vals.min() |
| 1119 | max_non_topk = non_topk_vals.max() |
| 1120 | |
| 1121 | # Allow small tolerance for floating point errors |
| 1122 | assert min_cuda_val >= max_non_topk - 1e-4, ( |
| 1123 | f"Row {i}: CUDA top-k contains values smaller than non-top-k. " |
| 1124 | f"Min CUDA: {min_cuda_val}, Max non-top-k: {max_non_topk}, " |
| 1125 | f"Length: {length}, k: {k_i}, CUDA indices: {sorted(cuda_set)[:10]}..., " # noqa: E501 |
| 1126 | f"Expected indices: {sorted(torch_set)[:10]}..." |
| 1127 | ) |
| 1128 | |
| 1129 | # For ties, verify the values are close |
| 1130 | assert torch.allclose( |
| 1131 | cuda_vals.sort(descending=True)[0], |
| 1132 | torch_vals.sort(descending=True)[0], |
| 1133 | rtol=1e-4, |
| 1134 | atol=1e-4, |
| 1135 | ), f"""Row {i}: Top-k values don't match. |
| 1136 | CUDA: {cuda_vals.sort(descending=True)[0][:10]}, |
| 1137 | Torch: {torch_vals.sort(descending=True)[0][:10]}""" |
| 1138 | |
| 1139 | |
| 1140 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1141 | @pytest.mark.parametrize("top_k", [512, 1024, 2048]) |
| 1142 | def test_cooperative_topk_cs2(top_k: int) -> None: |
| 1143 | """The 64-row dispatch uses the two-CTA cooperative kernel.""" |
| 1144 | run_large_context_topk_test( |
| 1145 | batch_size=64, |
| 1146 | seq_lens=[65536] * 64, |
| 1147 | top_k=top_k, |
| 1148 | backend="cooperative_topk", |
| 1149 | ) |
| 1150 | |
| 1151 | |
| 1152 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1153 | @pytest.mark.parametrize( |
| 1154 | "test_config", |
| 1155 | [ |
| 1156 | # ==================== CATEGORY: Sequence Length Edge Cases ==================== |
| 1157 | pytest.param( |
| 1158 | {"seq_lens": [1, 10, 100, 2048], "top_k": 2048, "data_type": "random"}, |
| 1159 | id="seq_len_edge_very_small_to_medium", |
| 1160 | ), |
| 1161 | pytest.param( |
| 1162 | { |
| 1163 | "seq_lens": [2049, 2100, 2500, 3000], |
| 1164 | "top_k": 2048, |
| 1165 | "data_type": "random", |
| 1166 | }, |
| 1167 | id="seq_len_edge_above_k", |
| 1168 | ), |
| 1169 | pytest.param( |
| 1170 | {"seq_lens": [8000, 16384, 20000], "top_k": 2048, "data_type": "random"}, |
| 1171 | id="algo_transition_filtered_radix", |
| 1172 | ), |
| 1173 | # ==================== CATEGORY: Data Distributions ==================== |
| 1174 | pytest.param( |
| 1175 | {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "sorted_asc"}, |
| 1176 | id="data_sorted_ascending", |
| 1177 | ), |
| 1178 | pytest.param( |
| 1179 | {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "sorted_desc"}, |
| 1180 | id="data_sorted_descending", |
| 1181 | ), |
| 1182 | pytest.param( |
| 1183 | {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "all_same"}, |
| 1184 | id="data_all_same", |
| 1185 | ), |
| 1186 | pytest.param( |
| 1187 | {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "many_ties"}, |
| 1188 | id="data_many_ties", |
| 1189 | ), |
| 1190 | pytest.param( |
| 1191 | { |
| 1192 | "seq_lens": [5000, 10000], |
| 1193 | "top_k": 2048, |
| 1194 | "data_type": "small_differences", |
| 1195 | }, |
| 1196 | id="data_float_precision", |
| 1197 | ), |
| 1198 | # ==================== CATEGORY: Alignment / Vectorization ==================== |
| 1199 | pytest.param( |
| 1200 | { |
| 1201 | "seq_lens": [2055, 2056, 2057, 2063], |
| 1202 | "top_k": 2048, |
| 1203 | "data_type": "random", |
| 1204 | }, |
| 1205 | id="align_vec_boundaries_low", |
| 1206 | ), |
| 1207 | pytest.param( |
| 1208 | { |
| 1209 | "seq_lens": [4095, 4096, 4097, 4102], |
| 1210 | "top_k": 2048, |
| 1211 | "data_type": "random", |
| 1212 | }, |
| 1213 | id="align_4k_boundary", |
| 1214 | ), |
| 1215 | pytest.param( |
| 1216 | { |
| 1217 | "seq_lens": [8191, 8192, 8193, 8198], |
| 1218 | "top_k": 2048, |
| 1219 | "data_type": "random", |
| 1220 | }, |
| 1221 | id="align_8k_boundary", |
| 1222 | ), |
| 1223 | pytest.param( |
| 1224 | { |
| 1225 | "seq_lens": [16383, 16384, 16385, 16390], |
| 1226 | "top_k": 2048, |
| 1227 | "data_type": "random", |
| 1228 | }, |
| 1229 | id="align_16k_boundary", |
| 1230 | ), |
| 1231 | ], |
| 1232 | ) |
| 1233 | @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) |
| 1234 | @torch.inference_mode() |
| 1235 | def test_workspace_topk_correctness(test_config: dict, backend: str) -> None: |
| 1236 | """Comprehensive correctness tests covering: |
| 1237 | - Sequence length edge cases (trivial, boundary, varied) |
| 1238 | - Very small sequences (< 100 elements) |
| 1239 | - Mixed sequence lengths in same batch |
| 1240 | - Data distributions (sorted, ties, precision) |
| 1241 | - Memory alignment / vectorization boundaries |
| 1242 | """ |
| 1243 | run_large_context_topk_test( |
| 1244 | batch_size=len(test_config["seq_lens"]), |
| 1245 | seq_lens=test_config["seq_lens"], |
| 1246 | top_k=test_config["top_k"], |
| 1247 | data_type=test_config.get("data_type", "random"), |
| 1248 | backend=backend, |
| 1249 | ) |
| 1250 | |
| 1251 | |
| 1252 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1253 | @pytest.mark.parametrize( |
| 1254 | "test_config", |
| 1255 | [ |
| 1256 | # ==================== CATEGORY: Batch Size Scalability ==================== |
| 1257 | pytest.param( |
| 1258 | {"batch_size": 1, "seq_len": 5000, "top_k": 2048}, |
| 1259 | id="batch_1", |
| 1260 | ), |
| 1261 | pytest.param( |
| 1262 | {"batch_size": 4, "seq_len": 5000, "top_k": 2048}, |
| 1263 | id="batch_4", |
| 1264 | ), |
| 1265 | pytest.param( |
| 1266 | {"batch_size": 32, "seq_len": 5000, "top_k": 2048}, |
| 1267 | id="batch_32", |
| 1268 | ), |
| 1269 | pytest.param( |
| 1270 | {"batch_size": 256, "seq_len": 5000, "top_k": 2048}, |
| 1271 | id="batch_256", |
| 1272 | ), |
| 1273 | # ==================== CATEGORY: Single-CTA vs Multi-CTA ==================== |
| 1274 | pytest.param( |
| 1275 | {"batch_size": 2, "seq_len": 4096, "top_k": 2048}, |
| 1276 | id="single_cta_4k", |
| 1277 | ), |
| 1278 | pytest.param( |
| 1279 | {"batch_size": 2, "seq_len": 8192, "top_k": 2048}, |
| 1280 | id="single_cta_8k", |
| 1281 | ), |
| 1282 | pytest.param( |
| 1283 | {"batch_size": 2, "seq_len": 163840, "top_k": 2048}, |
| 1284 | id="multi_cta_163840_dsv3_max", |
| 1285 | ), |
| 1286 | # ==================== CATEGORY: Extreme Cases ==================== |
| 1287 | pytest.param( |
| 1288 | {"batch_size": 512, "seq_len": 5000, "top_k": 2048}, |
| 1289 | id="extreme_large_batch", |
| 1290 | ), |
| 1291 | pytest.param( |
| 1292 | {"batch_size": 2, "seq_len": 163840, "top_k": 2048}, |
| 1293 | id="extreme_dsv3_max_context", |
| 1294 | ), |
| 1295 | ], |
| 1296 | ) |
| 1297 | @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) |
| 1298 | @torch.inference_mode() |
| 1299 | def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None: |
| 1300 | """Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2): |
| 1301 | - Batch size scalability (1, 4, 32, 256) |
| 1302 | - Single-CTA vs Multi-CTA execution |
| 1303 | - Extreme configurations (large batch, max context length) |
| 1304 | """ |
| 1305 | run_large_context_topk_test( |
| 1306 | batch_size=test_config["batch_size"], |
| 1307 | seq_lens=[test_config["seq_len"]] * test_config["batch_size"], |
| 1308 | top_k=test_config["top_k"], |
| 1309 | backend=backend, |
| 1310 | ) |
| 1311 | |
| 1312 | |
| 1313 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1314 | @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) |
| 1315 | @torch.inference_mode() |
| 1316 | def test_workspace_topk_stress(backend: str) -> None: |
| 1317 | """Stress test with random configurations to catch edge cases. |
| 1318 | Capped at 163840 (DeepSeek V3.2 max context) for realistic testing. |
| 1319 | """ |
| 1320 | torch.set_default_device("cuda:0") |
| 1321 | top_k = 2048 |
| 1322 | |
| 1323 | for seed in range(3): |
| 1324 | set_random_seed(seed) |
| 1325 | |
| 1326 | # Random batch size (limited for speed) |
| 1327 | batch_size = torch.randint(1, 32, (1,)).item() |
| 1328 | |
| 1329 | # Random sequence lengths capped at DeepSeek V3.2 max context |
| 1330 | seq_lens_tensor = torch.randint(100, 163840, (batch_size,)) |
| 1331 | if backend == "cooperative_topk": |
| 1332 | seq_lens = ((seq_lens_tensor + 3) & ~3).tolist() |
| 1333 | else: |
| 1334 | seq_lens = seq_lens_tensor.tolist() |
| 1335 | |
| 1336 | run_large_context_topk_test( |
| 1337 | batch_size=batch_size, |
| 1338 | seq_lens=seq_lens, |
| 1339 | top_k=top_k, |
| 1340 | seed=seed, |
| 1341 | backend=backend, |
| 1342 | ) |
| 1343 | |
| 1344 | |
| 1345 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1346 | @pytest.mark.parametrize("backend", TOPK_BACKENDS) |
| 1347 | @pytest.mark.parametrize("top_k", [512, 1024, 2048]) |
| 1348 | @torch.inference_mode() |
| 1349 | def test_deepseek_topk_backends_no_error_and_reference( |
| 1350 | backend: str, |
| 1351 | top_k: int, |
| 1352 | ) -> None: |
| 1353 | """Exercise every production top-k backend on the same inputs.""" |
| 1354 | run_large_context_topk_test( |
| 1355 | batch_size=4, |
| 1356 | seq_lens=[2049, 4097, 8191, 12000], |
| 1357 | top_k=top_k, |
| 1358 | data_type="random", |
| 1359 | seed=123, |
| 1360 | backend=backend, |
| 1361 | ) |
| 1362 | |
| 1363 | |
| 1364 | @pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+") |
| 1365 | @torch.inference_mode() |
| 1366 | def test_cooperative_topk_512_tie_workspace_is_per_row() -> None: |
| 1367 | """Regression test for TopK=512 tie workspace row overlap.""" |
| 1368 | torch.set_default_device("cuda:0") |
| 1369 | |
| 1370 | top_k = 512 |
| 1371 | num_rows = 2 |
| 1372 | stride = 65536 |
| 1373 | lengths = torch.tensor([40960, 65536], dtype=torch.int32, device="cuda") |
| 1374 | logits = torch.full( |
| 1375 | (num_rows, stride), float("-inf"), dtype=torch.float32, device="cuda" |
| 1376 | ) |
| 1377 | |
| 1378 | # Row 0 must never select these low indices: many better row-0 ties exist. |
| 1379 | logits[0, :2048] = -10.0 |
| 1380 | logits[0, 2048 : lengths[0]] = 1.0 |
| 1381 | # Row 1 has higher exact tie scores. With the old row * TopK tie_ws stride, |
| 1382 | # these row-1 ties could overwrite row 0's TopK=512 refinement workspace. |
| 1383 | logits[1, : lengths[1]] = 2.0 |
| 1384 | |
| 1385 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 1386 | workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") |
| 1387 | torch.ops._C.cooperative_topk(logits, lengths, indices, workspace, top_k, stride) |
| 1388 | torch.accelerator.synchronize() |
| 1389 | |
| 1390 | row0 = indices[0].cpu() |
| 1391 | assert torch.all(row0 >= 2048), ( |
| 1392 | "cooperative_topk TopK=512 selected row-0 low-score indices, likely " |
| 1393 | "from overlapping tie_ws rows" |
| 1394 | ) |
| 1395 | |
| 1396 | |
| 1397 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1398 | @pytest.mark.parametrize( |
| 1399 | "test_config", |
| 1400 | [ |
| 1401 | # Mixed batch: rows spanning all four paths (trivial, decode, medium, large) |
| 1402 | pytest.param( |
| 1403 | { |
| 1404 | "seq_lens": [2000, 6000, 30000, 80000], |
| 1405 | "data_type": "random", |
| 1406 | }, |
| 1407 | id="mixed_all_paths", |
| 1408 | ), |
| 1409 | # All decode/medium rows (typical decode scenario) |
| 1410 | pytest.param( |
| 1411 | { |
| 1412 | "seq_lens": [2048, 4096, 8192, 16000], |
| 1413 | "data_type": "random", |
| 1414 | }, |
| 1415 | id="all_decode_medium", |
| 1416 | ), |
| 1417 | # All large rows |
| 1418 | pytest.param( |
| 1419 | { |
| 1420 | "seq_lens": [70000, 100000, 163840], |
| 1421 | "data_type": "random", |
| 1422 | }, |
| 1423 | id="all_large", |
| 1424 | ), |
| 1425 | # Boundary around LARGE_THRESHOLD (32K) |
| 1426 | pytest.param( |
| 1427 | { |
| 1428 | "seq_lens": [32767, 32768, 32769, 32772], |
| 1429 | "data_type": "random", |
| 1430 | }, |
| 1431 | id="large_threshold_boundary", |
| 1432 | ), |
| 1433 | # Single row medium |
| 1434 | pytest.param( |
| 1435 | { |
| 1436 | "seq_lens": [5000], |
| 1437 | "data_type": "random", |
| 1438 | }, |
| 1439 | id="single_row_medium", |
| 1440 | ), |
| 1441 | # Single row large |
| 1442 | pytest.param( |
| 1443 | { |
| 1444 | "seq_lens": [100000], |
| 1445 | "top_k": 2048, |
| 1446 | "data_type": "random", |
| 1447 | }, |
| 1448 | id="single_row_large", |
| 1449 | ), |
| 1450 | # Trivial rows mixed with medium and large |
| 1451 | pytest.param( |
| 1452 | { |
| 1453 | "seq_lens": [100, 2048, 10000, 80000], |
| 1454 | "data_type": "random", |
| 1455 | }, |
| 1456 | id="trivial_medium_large_mix", |
| 1457 | ), |
| 1458 | ], |
| 1459 | ) |
| 1460 | @pytest.mark.parametrize("top_k", [512, 2048]) |
| 1461 | @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) |
| 1462 | @torch.inference_mode() |
| 1463 | def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None: |
| 1464 | """Tests specific to workspace top-k backends: |
| 1465 | - Mixed medium/large rows in the same batch (dynamic per-row dispatch) |
| 1466 | - Boundary around LARGE_THRESHOLD (32K) |
| 1467 | - Trivial + medium + large rows in a single batch |
| 1468 | """ |
| 1469 | run_large_context_topk_test( |
| 1470 | batch_size=len(test_config["seq_lens"]), |
| 1471 | seq_lens=test_config["seq_lens"], |
| 1472 | top_k=top_k, |
| 1473 | data_type=test_config.get("data_type", "random"), |
| 1474 | backend=backend, |
| 1475 | ) |
| 1476 | |
| 1477 | |
| 1478 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1479 | @torch.inference_mode() |
| 1480 | def test_persistent_topk_reused_group_after_short_row() -> None: |
| 1481 | """A short row must not advance a group's radix histogram ring.""" |
| 1482 | torch.set_default_device("cuda:0") |
| 1483 | set_random_seed(0) |
| 1484 | |
| 1485 | top_k = 2048 |
| 1486 | long_seq_len = 32769 |
| 1487 | radix = 256 |
| 1488 | fixed_smem = ((radix + radix + 5) * 4 + 15) & ~15 |
| 1489 | props = torch.cuda.get_device_properties(0) |
| 1490 | max_smem = props.shared_memory_per_block_optin |
| 1491 | if max_smem <= props.shared_memory_per_multiprocessor // 2: |
| 1492 | pytest.skip("Cannot force one persistent_topk CTA per SM") |
| 1493 | |
| 1494 | max_chunk = ((max_smem - fixed_smem) // 4 // 4) * 4 |
| 1495 | ctas_per_group = max( |
| 1496 | (props.multi_processor_count - 1 + 9) // 10, |
| 1497 | (long_seq_len + max_chunk - 1) // max_chunk, |
| 1498 | ) |
| 1499 | if ctas_per_group >= props.multi_processor_count: |
| 1500 | pytest.skip("Not enough SMs to construct a reused CTA group") |
| 1501 | |
| 1502 | stride = ctas_per_group * max_chunk |
| 1503 | num_groups = max(1, (props.multi_processor_count - 1) // ctas_per_group) |
| 1504 | num_rows = 3 * num_groups |
| 1505 | lengths = torch.full((num_rows,), top_k, dtype=torch.int32, device="cuda") |
| 1506 | target_row = 2 * num_groups |
| 1507 | lengths[0] = long_seq_len |
| 1508 | lengths[num_groups] = long_seq_len - 1 |
| 1509 | lengths[target_row] = long_seq_len |
| 1510 | |
| 1511 | logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda") |
| 1512 | indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 1513 | _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, stride) |
| 1514 | torch.accelerator.synchronize() |
| 1515 | |
| 1516 | expected = logits[target_row, :long_seq_len].topk(top_k).indices |
| 1517 | assert set(indices[target_row].cpu().tolist()) == set(expected.cpu().tolist()) |
| 1518 | |
| 1519 | |
| 1520 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 1521 | @pytest.mark.parametrize("top_k", [512, 2048]) |
| 1522 | @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) |
| 1523 | @torch.inference_mode() |
| 1524 | def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None: |
| 1525 | """Test workspace top-k backends with padded logits (large stride, small seq_len) |
| 1526 | to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits |
| 1527 | returns [B, max_model_len] with max_model_len=163840. |
| 1528 | """ |
| 1529 | set_random_seed(42) |
| 1530 | torch.set_default_device("cuda:0") |
| 1531 | |
| 1532 | batch_size = 4 |
| 1533 | padded_stride = 163840 # DeepSeek-V3.2 max_model_len |
| 1534 | actual_seq_lens = [3000, 5000, 8000, 12000] |
| 1535 | |
| 1536 | # Create padded logits tensor (like fp8_paged_mqa_logits output) |
| 1537 | logits = torch.full( |
| 1538 | (batch_size, padded_stride), |
| 1539 | float("-inf"), |
| 1540 | dtype=torch.float32, |
| 1541 | device="cuda", |
| 1542 | ) |
| 1543 | for i, sl in enumerate(actual_seq_lens): |
| 1544 | logits[i, :sl] = torch.randn(sl, dtype=torch.float32, device="cuda") |
| 1545 | |
| 1546 | lengths = torch.tensor(actual_seq_lens, dtype=torch.int32, device="cuda") |
| 1547 | indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda") |
| 1548 | _run_topk_backend(backend, logits, lengths, indices, top_k, max(actual_seq_lens)) |
| 1549 | torch.accelerator.synchronize() |
| 1550 | |
| 1551 | # Validate against torch.topk |
| 1552 | for i in range(batch_size): |
| 1553 | sl = actual_seq_lens[i] |
| 1554 | k_i = min(top_k, sl) |
| 1555 | expected = logits[i, :sl].topk(k_i, dim=-1)[1].cpu() |
| 1556 | actual = indices[i, :k_i].cpu() |
| 1557 | |
| 1558 | expected_set = set(expected.tolist()) |
| 1559 | actual_set = set(actual.tolist()) |
| 1560 | |
| 1561 | if expected_set != actual_set: |
| 1562 | # Allow ties |
| 1563 | expected_vals = logits[i, expected].cpu().sort(descending=True)[0] |
| 1564 | actual_vals = logits[i, actual].cpu().sort(descending=True)[0] |
| 1565 | assert torch.allclose(expected_vals, actual_vals, rtol=1e-4, atol=1e-4), ( |
| 1566 | f"Row {i}: {backend} with padded stride doesn't match. " |
| 1567 | f"seq_len={sl}, stride={padded_stride}" |
| 1568 | ) |
| 1569 | |
| 1570 | |
| 1571 | @requires_sm100 |
| 1572 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 1573 | @pytest.mark.parametrize("batch_size", [1, 8, 64]) |
| 1574 | @pytest.mark.parametrize("vocab_size", [4096, 131072]) |
| 1575 | @pytest.mark.parametrize("top_k", [512, 2048]) |
| 1576 | @torch.inference_mode() |
| 1577 | def test_deep_select_topk( |
| 1578 | dtype: torch.dtype, |
| 1579 | batch_size: int, |
| 1580 | vocab_size: int, |
| 1581 | top_k: int, |
| 1582 | ) -> None: |
| 1583 | """DeepSelect wrapper vs torch.topk with variable per-row ends.""" |
| 1584 | from vllm.model_executor.layers import indexer_topk |
| 1585 | |
| 1586 | set_random_seed(0) |
| 1587 | torch.set_default_device("cuda:0") |
| 1588 | |
| 1589 | row_ends = torch.randint( |
| 1590 | 1, vocab_size + 1, (batch_size,), dtype=torch.int32, device="cuda" |
| 1591 | ) |
| 1592 | # Force one short row to exercise the -1 out-of-bounds fill. |
| 1593 | row_ends[0] = min(top_k - 1, vocab_size) |
| 1594 | logits = torch.randn(batch_size, vocab_size, dtype=dtype, device="cuda") |
| 1595 | # Values beyond each row's end must never be read by the kernel: fill |
| 1596 | # them with NaN (production tails are uninitialized; DeepSelect traps on |
| 1597 | # NaN with abort_when_nan_found=True, so a wrong read fails loudly). |
| 1598 | col_idx = torch.arange(vocab_size, device="cuda") |
| 1599 | logits[col_idx[None, :] >= row_ends[:, None]] = float("nan") |
| 1600 | |
| 1601 | indices = indexer_topk.deep_select_topk(logits, top_k, end=row_ends) |
| 1602 | torch.accelerator.synchronize() |
| 1603 | |
| 1604 | assert indices.shape == (batch_size, top_k) |
| 1605 | row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda") |
| 1606 | torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda") |
| 1607 | for i in range(batch_size): |
| 1608 | row_end = int(row_ends[i]) |
| 1609 | k_i = min(top_k, row_end) |
| 1610 | idx = logits[i, :row_end].topk(k_i, dim=-1)[1] |
| 1611 | torch_indices[i, :k_i] = idx |
| 1612 | assert torch.all(indices[i, k_i:] == -1), ( |
| 1613 | f"Row {i}: expected -1 fill after {k_i} valid indices" |
| 1614 | ) |
| 1615 | # All selected indices must lie inside [0, row_end). |
| 1616 | assert torch.all(indices[i, :k_i] < row_end) |
| 1617 | |
| 1618 | assert compare_top_k_results( |
| 1619 | logits, indices, torch_indices, row_starts, row_ends, top_k |
| 1620 | ), "DeepSelect topk results don't match torch.topk" |
| 1621 | |
| 1622 | |
| 1623 | @requires_sm100 |
| 1624 | @torch.inference_mode() |
| 1625 | def test_deep_select_topk_preallocated_output() -> None: |
| 1626 | """DeepSelect writes into a preallocated aligned buffer (no `end`).""" |
| 1627 | from vllm.model_executor.layers import indexer_topk |
| 1628 | |
| 1629 | set_random_seed(0) |
| 1630 | torch.set_default_device("cuda:0") |
| 1631 | |
| 1632 | batch_size = 8 |
| 1633 | vocab_size = 131072 |
| 1634 | top_k = 2048 |
| 1635 | logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda") |
| 1636 | |
| 1637 | # Mimic vLLM's topk_indices_buffer: slice of a wider contiguous buffer, |
| 1638 | # whose stride(0) is 32B-aligned. |
| 1639 | buffer = torch.full((batch_size, 4096), -2, dtype=torch.int32, device="cuda") |
| 1640 | output_idx = buffer[:, :top_k] |
| 1641 | result = indexer_topk.deep_select_topk(logits, top_k, output_idx=output_idx) |
| 1642 | torch.accelerator.synchronize() |
| 1643 | |
| 1644 | assert result.data_ptr() == output_idx.data_ptr() |
| 1645 | assert torch.all(buffer[:, top_k:] == -2), "kernel wrote out of bounds" |
| 1646 | |
| 1647 | row_ends = torch.full((batch_size,), vocab_size, dtype=torch.int32, device="cuda") |
| 1648 | row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda") |
| 1649 | torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda") |
| 1650 | for i in range(batch_size): |
| 1651 | torch_indices[i] = logits[i].topk(top_k, dim=-1)[1] |
| 1652 | |
| 1653 | assert compare_top_k_results( |
| 1654 | logits, output_idx, torch_indices, row_starts, row_ends, top_k |
| 1655 | ), "DeepSelect topk (preallocated output) results don't match torch.topk" |
| 1656 | |
| 1657 | |
| 1658 | def _has_flashinfer_topk() -> bool: |
| 1659 | if not current_platform.is_cuda(): |
| 1660 | return False |
| 1661 | # find_spec only: importing flashinfer initializes CUDA at import time. |
| 1662 | import importlib.util |
| 1663 | |
| 1664 | return importlib.util.find_spec("flashinfer") is not None |
| 1665 | |
| 1666 | |
| 1667 | def _has_cooperative_topk() -> bool: |
| 1668 | return _has_device_capability( |
| 1669 | 90 |
| 1670 | ) and not current_platform.is_device_capability_family(120) |
| 1671 | |
| 1672 | |
| 1673 | SPARSE_INDEXER_EXPLICIT_BACKENDS = [ |
| 1674 | pytest.param( |
| 1675 | "per_row", |
| 1676 | marks=pytest.mark.skipif( |
| 1677 | not current_platform.is_cuda(), reason="requires CUDA" |
| 1678 | ), |
| 1679 | ), |
| 1680 | "torch", |
| 1681 | pytest.param( |
| 1682 | "persistent", |
| 1683 | marks=pytest.mark.skipif( |
| 1684 | not current_platform.is_cuda(), reason="requires CUDA" |
| 1685 | ), |
| 1686 | ), |
| 1687 | pytest.param( |
| 1688 | "cooperative", |
| 1689 | marks=pytest.mark.skipif( |
| 1690 | not _has_cooperative_topk(), |
| 1691 | reason="cooperative_topk requires SM90+ (non-SM12x)", |
| 1692 | ), |
| 1693 | ), |
| 1694 | pytest.param( |
| 1695 | "flashinfer", |
| 1696 | marks=pytest.mark.skipif( |
| 1697 | not _has_flashinfer_topk(), reason="requires flashinfer top-k" |
| 1698 | ), |
| 1699 | ), |
| 1700 | pytest.param("deep_select", marks=requires_sm100), |
| 1701 | ] |
| 1702 | |
| 1703 | |
| 1704 | @pytest.mark.parametrize("next_n", [1, 4]) |
| 1705 | @pytest.mark.parametrize("backend", SPARSE_INDEXER_EXPLICIT_BACKENDS) |
| 1706 | @torch.inference_mode() |
| 1707 | def test_sparse_indexer_decode_topk_explicit_backends( |
| 1708 | backend: str, next_n: int, workspace_init |
| 1709 | ) -> None: |
| 1710 | """Every explicit sparse-indexer decode top-k backend must restrict each |
| 1711 | row to its seq_len (dirty data past the end must never be selected) and |
| 1712 | match torch.topk on the valid region.""" |
| 1713 | from vllm.config import VllmConfig, set_current_vllm_config |
| 1714 | from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk |
| 1715 | |
| 1716 | set_random_seed(0) |
| 1717 | torch.set_default_device("cuda:0") |
| 1718 | |
| 1719 | top_k = 2048 |
| 1720 | batch_size = 4 |
| 1721 | num_rows = batch_size * next_n |
| 1722 | |
| 1723 | # Per-row effective lens in the (B, next_n) "native spec decode" form; |
| 1724 | # cooperative/persistent_topk require lengths.numel() == num_rows. |
| 1725 | row_ends = torch.randint(4000, 16000, (num_rows,), dtype=torch.int32, device="cuda") |
| 1726 | seq_lens = row_ends.reshape(batch_size, next_n) |
| 1727 | max_seq_len = int(row_ends.max()) |
| 1728 | # Pad the stride: multiple of 4 for cooperative_topk (TMA) and |
| 1729 | # 1024B-aligned (256 fp32) for DeepSelect. |
| 1730 | stride = (max_seq_len + 255) // 256 * 256 |
| 1731 | |
| 1732 | row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") |
| 1733 | |
| 1734 | logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda") |
| 1735 | col_idx = torch.arange(stride, device="cuda") |
| 1736 | logits[col_idx[None, :] >= row_ends[:, None]] = 1e30 # dirty tail |
| 1737 | |
| 1738 | indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda") |
| 1739 | cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) |
| 1740 | with set_current_vllm_config(cfg): |
| 1741 | SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, max_seq_len) |
| 1742 | torch.accelerator.synchronize() |
| 1743 | |
| 1744 | # k_i == top_k for every row here (row_ends >= 3997 > top_k). |
| 1745 | assert torch.all(indices >= 0) |
| 1746 | torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") |
| 1747 | for i in range(num_rows): |
| 1748 | row_end = int(row_ends[i]) |
| 1749 | torch_indices[i] = logits[i, :row_end].topk(top_k, dim=-1)[1] |
| 1750 | assert torch.all(indices[i] < row_end), ( |
| 1751 | f"{backend}: row {i} selected indices past its seq_len" |
| 1752 | ) |
| 1753 | |
| 1754 | assert compare_top_k_results( |
| 1755 | logits, indices, torch_indices, row_starts, row_ends, top_k |
| 1756 | ), f"{backend} results don't match torch.topk" |
| 1757 | |
| 1758 | |
| 1759 | @pytest.mark.parametrize( |
| 1760 | "backend", |
| 1761 | [ |
| 1762 | "torch", |
| 1763 | pytest.param( |
| 1764 | "per_row", |
| 1765 | marks=pytest.mark.skipif( |
| 1766 | not current_platform.is_cuda(), reason="requires CUDA" |
| 1767 | ), |
| 1768 | ), |
| 1769 | pytest.param( |
| 1770 | "flashinfer", |
| 1771 | marks=pytest.mark.skipif( |
| 1772 | not _has_flashinfer_topk(), reason="requires flashinfer top-k" |
| 1773 | ), |
| 1774 | ), |
| 1775 | pytest.param("deep_select", marks=requires_sm100), |
| 1776 | ], |
| 1777 | ) |
| 1778 | @torch.inference_mode() |
| 1779 | def test_sparse_indexer_decode_topk_short_seq_lens( |
| 1780 | backend: str, workspace_init |
| 1781 | ) -> None: |
| 1782 | """seq_len < next_n: derived per-row ends must clamp at 0 (the reference |
| 1783 | kernels do max(0, ...)); negative ends are OOB for the ragged kernels.""" |
| 1784 | from vllm.config import VllmConfig, set_current_vllm_config |
| 1785 | from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk |
| 1786 | |
| 1787 | set_random_seed(0) |
| 1788 | next_n = 4 |
| 1789 | top_k = 512 |
| 1790 | vocab = 8192 # 1024B-aligned rows for DeepSelect |
| 1791 | # 1D (B,) seq_lens: rows of a request share its len and each kernel |
| 1792 | # derives per-row ends (with clamping) itself. |
| 1793 | seq_lens = torch.tensor([2, 6], dtype=torch.int32, device="cuda") |
| 1794 | num_rows = 2 * next_n |
| 1795 | logits = torch.randn(num_rows, vocab, dtype=torch.float32, device="cuda") |
| 1796 | indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda") |
| 1797 | |
| 1798 | cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) |
| 1799 | with set_current_vllm_config(cfg): |
| 1800 | SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, vocab) |
| 1801 | torch.accelerator.synchronize() |
| 1802 | |
| 1803 | # Derived ends: [2,6] - next_n + 1 + arange(next_n) -> [-1,0,1,2] and |
| 1804 | # [3,4,5,6], clamped to [0,0,1,2] and [3,4,5,6]. |
| 1805 | for row, end in enumerate([0, 0, 1, 2, 3, 4, 5, 6]): |
| 1806 | k_i = min(top_k, end) |
| 1807 | assert torch.all(indices[row, k_i:] == -1), f"{backend}: row {row}" |
| 1808 | if k_i: |
| 1809 | ref = logits[row, :end].topk(k_i).indices.sort().values |
| 1810 | got = indices[row, :k_i].sort().values |
| 1811 | assert torch.equal(got, ref), f"{backend}: row {row}" |
| 1812 | |
| 1813 | |
| 1814 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 1815 | @torch.inference_mode() |
| 1816 | def test_sparse_indexer_topk_backend_resolution() -> None: |
| 1817 | """Auto heuristic chain and fail-fast validation of explicit backends.""" |
| 1818 | from vllm.config import VllmConfig, set_current_vllm_config |
| 1819 | from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk |
| 1820 | |
| 1821 | # 1024B-aligned stride: 131072 * 4 bytes. |
| 1822 | logits = torch.randn(64, 131072, dtype=torch.float32, device="cuda") |
| 1823 | # 4-divisible (cooperative TMA ok) but not 1024B-aligned (DeepSelect no): |
| 1824 | # 131076 * 4 = 524304 bytes. |
| 1825 | unaligned_logits = torch.randn(128, 131076, dtype=torch.float32, device="cuda")[ |
| 1826 | :, :131072 |
| 1827 | ] |
| 1828 | |
| 1829 | def resolve( |
| 1830 | backend: str, |
| 1831 | t: torch.Tensor = logits, |
| 1832 | k: int = 2048, |
| 1833 | num_rows: int = 64, |
| 1834 | ) -> str: |
| 1835 | cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) |
| 1836 | with set_current_vllm_config(cfg): |
| 1837 | return SparseIndexerTopk().resolve_backend(t, k, num_rows) |
| 1838 | |
| 1839 | is_sm100 = current_platform.is_device_capability_family(100) |
| 1840 | has_coop = _has_cooperative_topk() |
| 1841 | has_fi = _has_flashinfer_topk() |
| 1842 | |
| 1843 | # "auto" is the chain cooperative -> persistent -> per_row, with the |
| 1844 | # cooperative/persistent switch at cooperative_topk's cluster-wave row |
| 1845 | # limit. It must never select the opt-in backends by itself. |
| 1846 | if has_coop: |
| 1847 | from vllm.model_executor.layers.indexer_topk import ( |
| 1848 | AUTO_COOPERATIVE_MAX_ROWS, |
| 1849 | ) |
| 1850 | |
| 1851 | # cooperative for every batch within its cluster-wave row limit... |
| 1852 | assert ( |
| 1853 | resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative" |
| 1854 | ) |
| 1855 | assert resolve("auto", k=512, num_rows=8) == "cooperative" |
| 1856 | # ...persistent past the row limit, even where DeepSelect would be |
| 1857 | # applicable. |
| 1858 | assert ( |
| 1859 | resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1) |
| 1860 | == "persistent" |
| 1861 | ) |
| 1862 | assert resolve("auto", k=512, num_rows=128) == "persistent" |
| 1863 | # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the |
| 1864 | # same row-limit logic applies. |
| 1865 | assert ( |
| 1866 | resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS) |
| 1867 | == "cooperative" |
| 1868 | ) |
| 1869 | # ...and persistent past the row limit. |
| 1870 | assert resolve("auto", unaligned_logits, num_rows=128) == "persistent" |
| 1871 | # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels) |
| 1872 | # -> per_row fallback. |
| 1873 | assert resolve("auto", k=5000) == "per_row" |
| 1874 | |
| 1875 | # Explicit values are returned as-is when constraints are met. |
| 1876 | assert resolve("per_row") == "per_row" |
| 1877 | assert resolve("torch") == "torch" |
| 1878 | assert resolve("persistent") == "persistent" |
| 1879 | if has_coop: |
| 1880 | assert resolve("cooperative") == "cooperative" |
| 1881 | if has_fi: |
| 1882 | assert resolve("flashinfer") == "flashinfer" |
| 1883 | if is_sm100: |
| 1884 | # Explicit deep_select has no row-count requirement. |
| 1885 | assert resolve("deep_select", num_rows=1) == "deep_select" |
| 1886 | with pytest.raises(RuntimeError, match="constraints"): |
| 1887 | resolve("deep_select", unaligned_logits) |
| 1888 | |
| 1889 | # Fail-fast on unmet constraints. |
| 1890 | with pytest.raises(RuntimeError, match="num_rows must be <= 64"): |
| 1891 | resolve("cooperative", num_rows=128) |
| 1892 | with pytest.raises(RuntimeError, match="topk_tokens must be in"): |
| 1893 | resolve("persistent", k=3000) |
| 1894 | with pytest.raises(RuntimeError, match="topk_tokens must be in"): |
| 1895 | resolve("cooperative", k=3000) |
| 1896 | if not is_sm100: |
| 1897 | with pytest.raises(RuntimeError, match="SM100a/SM103a"): |
| 1898 | resolve("deep_select") |
| 1899 | |
| 1900 | |
| 1901 | def _patch_aiter_topk(monkeypatch, decode_kernel, enabled: bool = True) -> None: |
| 1902 | from vllm._aiter_ops import rocm_aiter_ops |
| 1903 | |
| 1904 | _enable_aiter_topk(monkeypatch, enabled=enabled) |
| 1905 | monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: enabled, raising=True) |
| 1906 | monkeypatch.setattr(rocm_aiter_ops, "indexer_top_k_decode", decode_kernel) |
| 1907 | |
| 1908 | |
| 1909 | def _make_topk(backend: str): |
| 1910 | from vllm.config import VllmConfig, set_current_vllm_config |
| 1911 | from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk |
| 1912 | |
| 1913 | cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) |
| 1914 | with set_current_vllm_config(cfg): |
| 1915 | return SparseIndexerTopk() |
| 1916 | |
| 1917 | |
| 1918 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 1919 | def test_sparse_indexer_topk_aiter_requires_rocm() -> None: |
| 1920 | """Explicitly requesting the ROCm-only backend must fail fast elsewhere.""" |
| 1921 | logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda") |
| 1922 | with pytest.raises(RuntimeError, match="ROCm"): |
| 1923 | _make_topk("aiter").resolve_backend(logits, 2048, 8) |
| 1924 | |
| 1925 | |
| 1926 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") |
| 1927 | def test_sparse_indexer_topk_backend_resolution_aiter(monkeypatch) -> None: |
| 1928 | """Aiter is opt-in only: "auto" never selects it (the shape-dependent perf |
| 1929 | gate lives at the call site), and an explicit request validates hard |
| 1930 | capability only.""" |
| 1931 | |
| 1932 | def decode_kernel(*args, **kwargs) -> None: |
| 1933 | pass |
| 1934 | |
| 1935 | _patch_aiter_topk(monkeypatch, decode_kernel) |
| 1936 | logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda") |
| 1937 | |
| 1938 | assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row" |
| 1939 | assert _make_topk("aiter").resolve_backend(logits, 2048, 8) == "aiter" |
| 1940 | # Shapes the call-site perf gate would decline still resolve. |
| 1941 | assert _make_topk("aiter").resolve_backend(logits, 512, 8) == "aiter" |
| 1942 | with pytest.raises(RuntimeError, match="topk_tokens must be in"): |
| 1943 | _make_topk("aiter").resolve_backend(logits, 3000, 8) |
| 1944 | |
| 1945 | _enable_aiter_topk(monkeypatch, enabled=False) |
| 1946 | with pytest.raises(RuntimeError, match="no tuned indexer top-k kernels"): |
| 1947 | _make_topk("aiter").resolve_backend(logits, 2048, 8) |
| 1948 | |
| 1949 | _patch_aiter_topk(monkeypatch, decode_kernel, enabled=False) |
| 1950 | assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row" |
| 1951 | with pytest.raises(RuntimeError, match="VLLM_ROCM_USE_AITER"): |
| 1952 | _make_topk("aiter").resolve_backend(logits, 2048, 8) |
| 1953 | |
| 1954 | |
| 1955 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") |
| 1956 | @torch.inference_mode() |
| 1957 | def test_sparse_indexer_topk_aiter_applies_no_shape_gate(monkeypatch) -> None: |
| 1958 | """Shape-dependent perf gating belongs to the call site: once "aiter" is |
| 1959 | requested the dispatcher must run it for every supported k, including the |
| 1960 | k=512 band the in-tree kernel wins on DSV4's compressed-KV logits.""" |
| 1961 | from vllm.model_executor.layers import indexer_topk |
| 1962 | |
| 1963 | calls: list[tuple] = [] |
| 1964 | _patch_aiter_topk(monkeypatch, lambda *args: calls.append(args)) |
| 1965 | monkeypatch.setattr( |
| 1966 | indexer_topk.ops, |
| 1967 | "top_k_per_row_decode", |
| 1968 | lambda *args: pytest.fail("an explicit aiter request must not fall back"), |
| 1969 | ) |
| 1970 | |
| 1971 | num_rows = 8 |
| 1972 | logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda") |
| 1973 | seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda") |
| 1974 | indices = torch.empty((num_rows, 512), dtype=torch.int32, device="cuda") |
| 1975 | op = _make_topk("aiter") |
| 1976 | |
| 1977 | op(logits, seq_lens, 1, indices, 512, 4096) |
| 1978 | # Past the compressed-context boundary the call-site gate declines on. |
| 1979 | op(logits, seq_lens, 1, indices, 512, 4 * (64 * 1024 + 1)) |
| 1980 | assert len(calls) == 2 |
| 1981 | |
| 1982 | |
| 1983 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") |
| 1984 | @torch.inference_mode() |
| 1985 | def test_sparse_indexer_topk_auto_stays_on_per_row(monkeypatch) -> None: |
| 1986 | """The "auto" chain must not reach AITER on its own: the decode top-k runs |
| 1987 | through top_k_per_row_decode unless the caller asks for the aiter backend.""" |
| 1988 | from vllm.model_executor.layers import indexer_topk |
| 1989 | |
| 1990 | def decode_kernel(*args, **kwargs) -> None: |
| 1991 | pytest.fail("AITER is opt-in only") |
| 1992 | |
| 1993 | _patch_aiter_topk(monkeypatch, decode_kernel) |
| 1994 | calls: list[tuple] = [] |
| 1995 | monkeypatch.setattr( |
| 1996 | indexer_topk.ops, "top_k_per_row_decode", lambda *args: calls.append(args) |
| 1997 | ) |
| 1998 | |
| 1999 | num_rows = 8 |
| 2000 | logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda") |
| 2001 | seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda") |
| 2002 | indices = torch.empty((num_rows, 1024), dtype=torch.int32, device="cuda") |
| 2003 | |
| 2004 | _make_topk("auto")(logits, seq_lens, 1, indices, 1024, 4096) |
| 2005 | assert len(calls) == 1 |
| 2006 | |
| 2007 | |
| 2008 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") |
| 2009 | @torch.inference_mode() |
| 2010 | def test_sparse_indexer_topk_aiter_expands_mtp_row_ends(monkeypatch) -> None: |
| 2011 | """AITER's decode kernel only implements the next_n == 1 form, so spec |
| 2012 | decode batches must be handed the equivalent expanded per-row ends.""" |
| 2013 | from vllm.model_executor.layers import indexer_topk |
| 2014 | |
| 2015 | recorded: dict = {} |
| 2016 | |
| 2017 | def decode_kernel(logits, next_n, seq_lens, *args, **kwargs) -> None: |
| 2018 | recorded["next_n"] = next_n |
| 2019 | recorded["row_ends"] = seq_lens.tolist() |
| 2020 | |
| 2021 | _patch_aiter_topk(monkeypatch, decode_kernel) |
| 2022 | monkeypatch.setattr( |
| 2023 | indexer_topk.ops, |
| 2024 | "top_k_per_row_decode", |
| 2025 | lambda *args: pytest.fail("an explicit aiter request must not fall back"), |
| 2026 | ) |
| 2027 | |
| 2028 | next_n = 2 |
| 2029 | logits = torch.randn(4, 4096, dtype=torch.float32, device="cuda") |
| 2030 | seq_lens = torch.tensor([[100], [200]], dtype=torch.int32, device="cuda") |
| 2031 | indices = torch.empty((4, 1024), dtype=torch.int32, device="cuda") |
| 2032 | |
| 2033 | _make_topk("aiter")(logits, seq_lens, next_n, indices, 1024, 800) |
| 2034 | |
| 2035 | assert recorded["next_n"] == 1 |
| 2036 | assert recorded["row_ends"] == [99, 100, 199, 200] |