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

[ROCm][Perf] Optimize DeepSeek V4 C4A top-k with AITER (#52882)

Fangzhou Ai committed Aug 21, 2026 at 13:00 UTC fe76112ff2981e9373765d76687d53f94880f38e
5 files changed +612 -31
csrc/libtorch_stable/sampler.cu
+132 -11
@@ -2,6 +2,8 @@
2 #include "dispatch_utils.h"
3 #include "torch_utils.h"
4
5 +#include <cstring>
6 +
7 #ifndef USE_ROCM
8 #include <cub/cub.cuh>
9 #else
@@ -327,11 +329,14 @@ __device__ bool processHistogramStep(
329 }
330
331 // Follows half - 11 - 11 - 10 bit iterations
332 +// Keep the adaptive kernel's device instantiation separate from legacy calls.
333 template <int kNumThreadsPerBlock, int kNumBins, bool useRadixSort,
331 - bool multipleBlocksPerRow = false, bool mergeBlocks = false>
334 + bool multipleBlocksPerRow = false, bool mergeBlocks = false,
335 + bool deviceLengthAware = false>
336 static __device__ void topKPerRowJob(const int* indices, const float* logits,
337 int rowStart, int rowEnd, int* outIndices,
338 float* outLogits, int stride1, int topK) {
339 + static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks);
340 // The number of slots for the final pass.
341 static constexpr int kNumFinalItems = 2048;
342 // The number of elements per thread for the final sort.
@@ -613,6 +618,57 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
618 indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK);
619 }
620
621 +// Measured native-kernel policy for gfx950, FP32, and k=1024. This
622 +// active-column split crossover is independent of the Python AITER/native
623 +// backend crossover; the two policies can be retuned separately. At exactly
624 +// 64K, an eager call may select AITER while a captured native replay uses two
625 +// splits.
626 +constexpr int kGfx950C4AShortRowNumBlocks = 3;
627 +constexpr int kGfx950C4ALongRowNumBlocks = 2;
628 +constexpr int kGfx950C4ATwoSplitMinRowLength = 64 * 1024;
629 +
630 +template <int kNumThreadsPerBlock, bool mergeBlocks = false>
631 +static __global__
632 +__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware(
633 + const float* logits, const int* seqLens, int* outIndices, int stride0,
634 + int stride1, const int topK, int next_n, int seqLensIs2D = 0,
635 + float* outLogits = nullptr, const int* indices = nullptr) {
636 + static constexpr int kNumBins = 2048;
637 + int rowIdx = blockIdx.x;
638 + int batch_idx = rowIdx / next_n;
639 + int next_n_idx = rowIdx % next_n;
640 + int seq_len = seqLensIs2D ? seqLens[rowIdx] : seqLens[batch_idx];
641 + int rowEnd =
642 + seqLensIs2D ? max(0, seq_len) : max(0, seq_len - next_n + next_n_idx + 1);
643 + const int activeBlocks = rowEnd >= kGfx950C4ATwoSplitMinRowLength
644 + ? kGfx950C4ALongRowNumBlocks
645 + : kGfx950C4AShortRowNumBlocks;
646 + int rowStart = 0;
647 +
648 + if constexpr (!mergeBlocks) {
649 + if (blockIdx.y >= activeBlocks) return;
650 + const auto blockSize = rowEnd / activeBlocks;
651 + rowStart = blockSize * blockIdx.y;
652 + rowEnd = activeBlocks == blockIdx.y + 1 ? rowEnd : rowStart + blockSize;
653 + outIndices +=
654 + static_cast<int64_t>(rowIdx) * kGfx950C4AShortRowNumBlocks * topK +
655 + blockIdx.y * topK;
656 + outLogits +=
657 + static_cast<int64_t>(rowIdx) * kGfx950C4AShortRowNumBlocks * topK +
658 + blockIdx.y * topK;
659 + } else {
660 + rowEnd = activeBlocks * topK;
661 + indices +=
662 + static_cast<int64_t>(rowIdx) * kGfx950C4AShortRowNumBlocks * topK;
663 + outIndices += static_cast<int64_t>(rowIdx) * topK;
664 + }
665 + logits += static_cast<int64_t>(rowIdx) * stride0;
666 +
667 + topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, !mergeBlocks, mergeBlocks,
668 + true>(indices, logits, rowStart, rowEnd, outIndices, outLogits,
669 + stride1, topK);
670 +}
671 +
672 } // namespace vllm
673
674 void apply_repetition_penalties_(
@@ -693,7 +749,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
749 static_cast<int>(next_n), seqLensIs2D);
750 } else {
751 // Long sequences are run in two steps
696 - constexpr auto multipleBlocksPerRowConfig = 10;
752 +#ifdef USE_ROCM
753 + const bool useGfx950DeviceLengthAwareTopK1024 =
754 + topK == 1024 && numRows > 64 && numRows <= 256 &&
755 + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0;
756 + if (useGfx950DeviceLengthAwareTopK1024) {
757 + // Keep three blocks for short FULL-graph replays, but let each row use
758 + // two blocks once its live device length reaches the long-context
759 + // crossover. The grid and workspace shapes remain capture-stable.
760 + constexpr int multipleBlocksPerRowConfig =
761 + vllm::kGfx950C4AShortRowNumBlocks;
762 + const auto outIndicesAux = torch::stable::empty(
763 + {numRows, multipleBlocksPerRowConfig, topK},
764 + torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
765 + const auto outLogitsAux = torch::stable::empty(
766 + {numRows, multipleBlocksPerRowConfig, topK},
767 + torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
768 +
769 + constexpr int kNumThreadsPerSplitBlock = 1024;
770 + vllm::topKPerRowDecodeDeviceLengthAware<kNumThreadsPerSplitBlock>
771 + <<<dim3(numRows, multipleBlocksPerRowConfig),
772 + kNumThreadsPerSplitBlock, 2 * topK * sizeof(int32_t), stream>>>(
773 + logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
774 + outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
775 + static_cast<int>(stride1), static_cast<int>(topK),
776 + static_cast<int>(next_n), seqLensIs2D,
777 + outLogitsAux.mutable_data_ptr<float>());
778 + constexpr int kNumThreadsPerBlockMerge = 1024;
779 + vllm::topKPerRowDecodeDeviceLengthAware<kNumThreadsPerBlockMerge, true>
780 + <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t),
781 + stream>>>(
782 + outLogitsAux.const_data_ptr<float>(),
783 + seqLens.const_data_ptr<int>(), indices.mutable_data_ptr<int>(),
784 + multipleBlocksPerRowConfig * topK, 1, static_cast<int>(topK),
785 + static_cast<int>(next_n), seqLensIs2D, nullptr,
786 + outIndicesAux.const_data_ptr<int>());
787 + return;
788 + }
789 +#endif
790 +
791 +#ifdef USE_ROCM
792 + const bool useGfx950TopK1024 =
793 + topK == 1024 && numRows <= 64 &&
794 + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0;
795 +#else
796 + constexpr bool useGfx950TopK1024 = false;
797 +#endif
798 + int multipleBlocksPerRowConfig = 10;
799 + if (useGfx950TopK1024) {
800 + if (numRows <= 32) {
801 + multipleBlocksPerRowConfig = 8;
802 + } else {
803 + multipleBlocksPerRowConfig = 4;
804 + }
805 + }
806
807 const auto outIndicesAux = torch::stable::empty(
808 {numRows, multipleBlocksPerRowConfig, topK},
@@ -702,14 +811,26 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
811 {numRows, multipleBlocksPerRowConfig, topK},
812 torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
813
705 - vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
706 - <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
707 - 2 * topK * sizeof(int32_t), stream>>>(
708 - logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
709 - outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
710 - static_cast<int>(stride1), static_cast<int>(topK),
711 - static_cast<int>(next_n), seqLensIs2D,
712 - outLogitsAux.mutable_data_ptr<float>());
814 + if (useGfx950TopK1024) {
815 + constexpr int kNumThreadsPerSplitBlock = 1024;
816 + vllm::topKPerRowDecode<kNumThreadsPerSplitBlock, true, true>
817 + <<<dim3(numRows, multipleBlocksPerRowConfig),
818 + kNumThreadsPerSplitBlock, 2 * topK * sizeof(int32_t), stream>>>(
819 + logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
820 + outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
821 + static_cast<int>(stride1), static_cast<int>(topK),
822 + static_cast<int>(next_n), seqLensIs2D,
823 + outLogitsAux.mutable_data_ptr<float>());
824 + } else {
825 + vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
826 + <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
827 + 2 * topK * sizeof(int32_t), stream>>>(
828 + logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
829 + outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
830 + static_cast<int>(stride1), static_cast<int>(topK),
831 + static_cast<int>(next_n), seqLensIs2D,
832 + outLogitsAux.mutable_data_ptr<float>());
833 + }
834
835 constexpr int kNumThreadsPerBlockMerge = 1024;
836 vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
@@ -753,4 +874,4 @@ void top_k_per_row_prefill(const torch::stable::Tensor& logits,
874 static_cast<int>(stride0), static_cast<int>(stride1),
875 static_cast<int>(topK), kSortingAlgorithmThreshold);
876 }
756 -}
\ No newline at end of file
877 +}
tests/kernels/test_top_k_per_row.py
+293
@@ -374,6 +374,299 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
374 )
375
376
377 +@pytest.mark.parametrize(
378 + "batch_size,num_speculative_tokens,min_seq_len,max_seq_len",
379 + [
380 + pytest.param(1, 2, 125_000, 250_000, id="eight-splits"),
381 + pytest.param(7, 4, 125_000, 250_000, id="four-splits"),
382 + pytest.param(16, 4, 2_500, 2_501, id="dynamic-short-rows"),
383 + pytest.param(13, 4, 25_000, 25_001, id="dynamic-three-splits"),
384 + pytest.param(16, 4, 65_535, 65_536, id="dynamic-boundary-mixed"),
385 + pytest.param(16, 4, 125_000, 125_001, id="dynamic-two-splits"),
386 + pytest.param(64, 3, 125_000, 125_001, id="dynamic-two-splits-max-rows"),
387 + pytest.param(64, 4, 125_000, 125_001, id="baseline-fallback"),
388 + ],
389 +)
390 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
391 +@torch.inference_mode()
392 +def test_top_k_per_row_decode_gfx950_long_c4a(
393 + batch_size: int,
394 + num_speculative_tokens: int,
395 + min_seq_len: int,
396 + max_seq_len: int,
397 +) -> None:
398 + properties = torch.cuda.get_device_properties(0)
399 + if not properties.gcnArchName.startswith("gfx950"):
400 + pytest.skip("This test exercises the gfx950 launch configuration")
401 +
402 + query_tokens = num_speculative_tokens + 1
403 + num_rows = batch_size * query_tokens
404 + stride = 262_144
405 + top_k = 1024
406 + offsets = torch.arange(num_rows, dtype=torch.int32, device="cuda")
407 + seq_lens = min_seq_len + offsets * (max_seq_len - min_seq_len) // max(
408 + num_rows - 1, 1
409 + )
410 + seq_lens = seq_lens.reshape(batch_size, query_tokens)
411 + logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
412 + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
413 +
414 + torch.ops._C.top_k_per_row_decode(
415 + logits,
416 + query_tokens,
417 + seq_lens,
418 + indices,
419 + num_rows,
420 + logits.stride(0),
421 + logits.stride(1),
422 + top_k,
423 + )
424 +
425 + row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
426 + validate_topk_against_reference(
427 + logits,
428 + indices,
429 + row_starts,
430 + seq_lens.reshape(-1),
431 + top_k,
432 + "gfx950 long C4A top-k",
433 + )
434 +
435 +
436 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
437 +@torch.inference_mode()
438 +def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None:
439 + properties = torch.cuda.get_device_properties(0)
440 + if not properties.gcnArchName.startswith("gfx950"):
441 + pytest.skip("This test exercises the gfx950 launch configuration")
442 +
443 + batch_size = 20
444 + query_tokens = 4
445 + num_rows = batch_size * query_tokens
446 + stride = 262_144
447 + top_k = 1024
448 + seq_lens = torch.linspace(
449 + 125_003,
450 + 250_000,
451 + batch_size,
452 + dtype=torch.int32,
453 + device="cuda",
454 + )
455 + row_offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda")
456 + row_ends = (seq_lens[:, None] - query_tokens + row_offsets + 1).reshape(-1)
457 + logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
458 + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
459 +
460 + torch.ops._C.top_k_per_row_decode(
461 + logits,
462 + query_tokens,
463 + seq_lens,
464 + indices,
465 + num_rows,
466 + logits.stride(0),
467 + logits.stride(1),
468 + top_k,
469 + )
470 +
471 + row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
472 + validate_topk_against_reference(
473 + logits,
474 + indices,
475 + row_starts,
476 + row_ends,
477 + top_k,
478 + "gfx950 long C4A top-k with 1-D sequence lengths",
479 + )
480 +
481 +
482 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
483 +@torch.inference_mode()
484 +def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
485 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
486 + _get_aiter_top_k_kernel,
487 + _launch_aiter_top_k_per_row_prefill,
488 + )
489 +
490 + row_starts = torch.tensor([1, 6, 20], dtype=torch.int32, device="cuda")
491 + row_ends = torch.tensor([3, 18, 23], dtype=torch.int32, device="cuda")
492 + logits = torch.arange(96, dtype=torch.float32, device="cuda").reshape(3, 32)
493 + top_k = 4
494 + indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda")
495 +
496 + aiter_topk_kernel = _get_aiter_top_k_kernel(
497 + is_prefill=True,
498 + compress_ratio=4,
499 + num_rows=logits.shape[0],
500 + on_gfx950=True,
501 + )
502 + if aiter_topk_kernel is None:
503 + pytest.skip("AITER top-k is unavailable")
504 + assert (
505 + _launch_aiter_top_k_per_row_prefill(
506 + aiter_topk_kernel,
507 + logits,
508 + row_starts,
509 + row_ends,
510 + indices,
511 + top_k,
512 + )
513 + is None
514 + )
515 +
516 + for row_idx, (row_start, row_end) in enumerate(
517 + zip(row_starts.tolist(), row_ends.tolist())
518 + ):
519 + num_valid = min(top_k, row_end - row_start)
520 + expected = logits[row_idx, row_start:row_end].topk(num_valid).indices
521 + actual = indices[row_idx, :num_valid]
522 + assert set(actual.tolist()) == set(expected.tolist())
523 + assert torch.all(indices[row_idx, num_valid:] == -1)
524 +
525 +
526 +@pytest.mark.parametrize("num_speculative_tokens", [2, 3, 4])
527 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
528 +@torch.inference_mode()
529 +def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths(
530 + num_speculative_tokens: int,
531 +) -> None:
532 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
533 + _get_aiter_top_k_kernel,
534 + _launch_aiter_top_k_per_row_decode,
535 + )
536 +
537 + query_tokens = num_speculative_tokens + 1
538 + offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda")
539 + final_uncompressed_lens = torch.tensor([43, 59], dtype=torch.int32, device="cuda")
540 + seq_lens = (final_uncompressed_lens[:, None] - query_tokens + offsets + 1) // 4
541 + seq_lens[0, 0] = 0
542 +
543 + num_rows = seq_lens.numel()
544 + logits = torch.arange(num_rows * 16, dtype=torch.float32, device="cuda").reshape(
545 + num_rows, 16
546 + )
547 + top_k = 4
548 + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
549 +
550 + aiter_topk_kernel = _get_aiter_top_k_kernel(
551 + is_prefill=False,
552 + compress_ratio=4,
553 + num_rows=num_rows,
554 + max_valid_seq_len=seq_lens.max().item(),
555 + on_gfx950=True,
556 + )
557 + if aiter_topk_kernel is None:
558 + pytest.skip("AITER top-k is unavailable")
559 + assert (
560 + _launch_aiter_top_k_per_row_decode(
561 + aiter_topk_kernel,
562 + logits,
563 + seq_lens,
564 + indices,
565 + top_k,
566 + )
567 + is None
568 + )
569 +
570 + for row_idx, row_end in enumerate(seq_lens.reshape(-1).tolist()):
571 + num_valid = min(top_k, row_end)
572 + expected = logits[row_idx, :row_end].topk(num_valid).indices
573 + actual = indices[row_idx, :num_valid]
574 + assert set(actual.tolist()) == set(expected.tolist())
575 + assert torch.all(indices[row_idx, num_valid:] == -1)
576 +
577 +
578 +def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None:
579 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse
580 +
581 + def prefill_kernel(*args, **kwargs) -> None:
582 + pass
583 +
584 + def decode_kernel(*args, **kwargs) -> None:
585 + pass
586 +
587 + monkeypatch.setattr(
588 + rocm_aiter_mla_sparse,
589 + "_get_aiter_topk_ops",
590 + lambda: (prefill_kernel, decode_kernel),
591 + )
592 + get_kernel = rocm_aiter_mla_sparse._get_aiter_top_k_kernel
593 +
594 + eligible_cases = [
595 + (
596 + dict(is_prefill=True, compress_ratio=4, num_rows=1, on_gfx950=True),
597 + prefill_kernel,
598 + ),
599 + (
600 + dict(
601 + is_prefill=False,
602 + compress_ratio=4,
603 + num_rows=1,
604 + max_valid_seq_len=65_536,
605 + on_gfx950=True,
606 + ),
607 + decode_kernel,
608 + ),
609 + (
610 + dict(
611 + is_prefill=False,
612 + compress_ratio=4,
613 + num_rows=257,
614 + max_valid_seq_len=250_000,
615 + on_gfx950=True,
616 + ),
617 + decode_kernel,
618 + ),
619 + ]
620 + for kwargs, expected_kernel in eligible_cases:
621 + assert get_kernel(**kwargs) is expected_kernel
622 +
623 + def fail_if_imported() -> None:
624 + pytest.fail("ineligible shapes must not import AITER top-k")
625 +
626 + monkeypatch.setattr(
627 + rocm_aiter_mla_sparse,
628 + "_get_aiter_topk_ops",
629 + fail_if_imported,
630 + )
631 + ineligible_cases = [
632 + dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True),
633 + dict(
634 + is_prefill=False,
635 + compress_ratio=4,
636 + num_rows=1,
637 + max_valid_seq_len=65_537,
638 + on_gfx950=True,
639 + ),
640 + dict(
641 + is_prefill=False,
642 + compress_ratio=4,
643 + num_rows=256,
644 + max_valid_seq_len=125_000,
645 + on_gfx950=True,
646 + ),
647 + dict(
648 + is_prefill=False,
649 + compress_ratio=4,
650 + num_rows=320,
651 + max_valid_seq_len=2_500,
652 + on_gfx950=False,
653 + ),
654 + ]
655 + for kwargs in ineligible_cases:
656 + assert get_kernel(**kwargs) is None
657 +
658 + monkeypatch.setattr(rocm_aiter_mla_sparse, "_get_aiter_topk_ops", lambda: None)
659 + assert (
660 + get_kernel(
661 + is_prefill=True,
662 + compress_ratio=4,
663 + num_rows=1,
664 + on_gfx950=True,
665 + )
666 + is None
667 + )
668 +
669 +
670 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
671 @pytest.mark.parametrize(
672 "seq_len_range,test_id",
vllm/model_executor/layers/sparse_attn_indexer.py
+3
@@ -750,6 +750,7 @@ class SparseAttnIndexer(CustomOp):
750 topk_indices_buffer: torch.Tensor,
751 skip_k_cache_insert: bool = False,
752 use_fp4_cache: bool = False,
753 + compress_ratio: int = 1,
754 ):
755 super().__init__()
756 self.k_cache = k_cache
@@ -762,6 +763,7 @@ class SparseAttnIndexer(CustomOp):
763 self.topk_indices_buffer = topk_indices_buffer
764 self.skip_k_cache_insert = skip_k_cache_insert
765 self.use_fp4_cache = use_fp4_cache
766 + self.compress_ratio = compress_ratio
767 self.dense_mha_metadata_layer_name = ""
768 # DCP scalars are constant for the run; resolve them here (config is set
769 # during model construction) and pass them into the custom op, rather
@@ -873,6 +875,7 @@ class SparseAttnIndexer(CustomOp):
875 self.max_total_seq_len,
876 self.topk_indices_buffer,
877 skip_k_cache_insert=self.skip_k_cache_insert,
878 + compress_ratio=self.compress_ratio,
879 )
880 raise RuntimeError(
881 "Sparse attention indexer ROCm path is only supported on AITER. "
vllm/models/deepseek_v4/attention.py
+1
@@ -854,6 +854,7 @@ class DeepseekV4Indexer(nn.Module):
854 self.topk_indices_buffer,
855 skip_k_cache_insert=True,
856 use_fp4_cache=self.use_fp4_kv,
857 + compress_ratio=self.compress_ratio,
858 )
859
860 # None on ROCm — maybe_execute_in_parallel falls back to sequential.
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+183 -20
@@ -3,6 +3,7 @@
3 import functools
4 import importlib
5 import math
6 +from collections.abc import Callable
7 from importlib.util import find_spec
8
9 import torch
@@ -10,6 +11,7 @@ import torch.nn.functional as F
11
12 import vllm.envs as envs
13 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
14 +from vllm.config import CUDAGraphMode
15 from vllm.forward_context import get_forward_context
16 from vllm.platforms import current_platform
17 from vllm.triton_utils import tl, triton
@@ -25,6 +27,127 @@ else:
27 _ON_GFX950 = False
28
29
30 +@functools.cache
31 +def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None:
32 + try:
33 + from aiter.ops.topk import (
34 + top_k_per_row_decode,
35 + top_k_per_row_prefill,
36 + )
37 + except ImportError:
38 + return None
39 + return top_k_per_row_prefill, top_k_per_row_decode
40 +
41 +
42 +_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
43 +_GFX950_C4A_NATIVE_MAX_ROWS = 256
44 +
45 +
46 +def _get_aiter_top_k_kernel(
47 + *,
48 + is_prefill: bool,
49 + compress_ratio: int,
50 + num_rows: int,
51 + max_valid_seq_len: int | None = None,
52 + on_gfx950: bool = _ON_GFX950,
53 +) -> Callable[..., None] | None:
54 + if compress_ratio <= 1 or not on_gfx950:
55 + return None
56 +
57 + if not is_prefill:
58 + assert max_valid_seq_len is not None
59 + # AITER v0.1.19 decode is one-block only. This measured gfx950
60 + # FP32/k=1024 compressed-row boundary is independent of the native
61 + # split-count boundary in sampler.cu.
62 + if (
63 + num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS
64 + and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN
65 + ):
66 + return None
67 +
68 + topk_ops = _get_aiter_topk_ops()
69 + if topk_ops is None:
70 + return None
71 + return topk_ops[0] if is_prefill else topk_ops[1]
72 +
73 +
74 +@triton.jit
75 +def _localize_aiter_prefill_topk_kernel(
76 + indices_ptr,
77 + row_starts_ptr,
78 + indices_stride,
79 + num_topk,
80 + BLOCK_SIZE: tl.constexpr,
81 +):
82 + row_idx = tl.program_id(0)
83 + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
84 + mask = offsets < num_topk
85 + indices = tl.load(indices_ptr + row_idx * indices_stride + offsets, mask=mask)
86 + row_start = tl.load(row_starts_ptr + row_idx)
87 + localized = tl.where(indices < 0, indices, indices - row_start)
88 + tl.store(
89 + indices_ptr + row_idx * indices_stride + offsets,
90 + localized,
91 + mask=mask,
92 + )
93 +
94 +
95 +def _localize_aiter_prefill_topk(
96 + indices: torch.Tensor,
97 + row_starts: torch.Tensor,
98 +) -> None:
99 + num_rows, num_topk = indices.shape
100 + block_size = 256
101 + _localize_aiter_prefill_topk_kernel[(num_rows, triton.cdiv(num_topk, block_size))](
102 + indices,
103 + row_starts,
104 + indices.stride(0),
105 + num_topk,
106 + BLOCK_SIZE=block_size,
107 + )
108 +
109 +
110 +def _launch_aiter_top_k_per_row_prefill(
111 + top_k_per_row_prefill: Callable[..., None],
112 + logits: torch.Tensor,
113 + row_starts: torch.Tensor,
114 + row_ends: torch.Tensor,
115 + indices: torch.Tensor,
116 + topk_tokens: int,
117 +) -> None:
118 + top_k_per_row_prefill(
119 + logits,
120 + row_starts,
121 + row_ends,
122 + indices,
123 + None,
124 + logits.shape[0],
125 + logits.stride(0),
126 + logits.stride(1),
127 + k=topk_tokens,
128 + )
129 + _localize_aiter_prefill_topk(indices, row_starts)
130 +
131 +
132 +def _launch_aiter_top_k_per_row_decode(
133 + top_k_per_row_decode: Callable[..., None],
134 + logits: torch.Tensor,
135 + seq_lens: torch.Tensor,
136 + indices: torch.Tensor,
137 + topk_tokens: int,
138 +) -> None:
139 + top_k_per_row_decode(
140 + logits,
141 + 1,
142 + seq_lens.reshape(-1),
143 + indices,
144 + logits.shape[0],
145 + logits.stride(0),
146 + logits.stride(1),
147 + k=topk_tokens,
148 + )
149 +
150 +
151 @triton.jit
152 def _indexer_k_quant_and_cache_kernel(
153 k_ptr, # [num_tokens, head_dim]
@@ -692,6 +815,7 @@ def rocm_aiter_sparse_attn_indexer_fake(
815 total_seq_lens: int,
816 topk_indices_buffer: torch.Tensor | None,
817 skip_k_cache_insert: bool = False,
818 + compress_ratio: int = 1,
819 ) -> torch.Tensor:
820 return topk_indices_buffer
821
@@ -712,9 +836,11 @@ def rocm_aiter_sparse_attn_indexer(
836 total_seq_lens: int,
837 topk_indices_buffer: torch.Tensor | None,
838 skip_k_cache_insert: bool = False,
839 + compress_ratio: int = 1,
840 ) -> torch.Tensor:
841 # careful! this will be None in dummy run
717 - attn_metadata = get_forward_context().attn_metadata
842 + forward_context = get_forward_context()
843 + attn_metadata = forward_context.attn_metadata
844 fp8_dtype = current_platform.fp8_dtype()
845 from vllm.utils.torch_utils import _resolve_layer_name
846
@@ -773,6 +899,7 @@ def rocm_aiter_sparse_attn_indexer(
899 total_seq_lens,
900 topk_indices_buffer,
901 skip_k_cache_insert,
902 + compress_ratio,
903 )
904 layer_attn_metadata = attn_metadata[k_cache_prefix]
905 assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata)
@@ -833,16 +960,31 @@ def rocm_aiter_sparse_attn_indexer(
960
961 num_rows = logits.shape[0]
962
836 - torch.ops._C.top_k_per_row_prefill(
837 - logits,
838 - chunk.cu_seqlen_ks,
839 - chunk.cu_seqlen_ke,
840 - topk_indices,
841 - num_rows,
842 - logits.stride(0),
843 - logits.stride(1),
844 - topk_tokens,
963 + aiter_topk_kernel = _get_aiter_top_k_kernel(
964 + is_prefill=True,
965 + compress_ratio=compress_ratio,
966 + num_rows=num_rows,
967 )
968 + if aiter_topk_kernel is not None:
969 + _launch_aiter_top_k_per_row_prefill(
970 + aiter_topk_kernel,
971 + logits,
972 + chunk.cu_seqlen_ks,
973 + chunk.cu_seqlen_ke,
974 + topk_indices,
975 + topk_tokens,
976 + )
977 + else:
978 + torch.ops._C.top_k_per_row_prefill(
979 + logits,
980 + chunk.cu_seqlen_ks,
981 + chunk.cu_seqlen_ke,
982 + topk_indices,
983 + num_rows,
984 + logits.stride(0),
985 + logits.stride(1),
986 + topk_tokens,
987 + )
988
989 if has_decode:
990 decode_metadata = layer_attn_metadata.decode
@@ -882,16 +1024,37 @@ def rocm_aiter_sparse_attn_indexer(
1024 topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
1025 num_rows = logits.shape[0]
1026
885 - torch.ops._C.top_k_per_row_decode(
886 - logits,
887 - next_n,
888 - decode_metadata.seq_lens,
889 - topk_indices,
890 - num_rows,
891 - logits.stride(0),
892 - logits.stride(1),
893 - topk_tokens,
894 - )
1027 + # FULL graphs are not keyed by context length, so use a replay-safe
1028 + # upper bound when choosing the captured kernel.
1029 + if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL:
1030 + max_compressed_seq_len = max_model_len
1031 + else:
1032 + max_compressed_seq_len = layer_attn_metadata.max_seq_len // compress_ratio
1033 + aiter_topk_kernel = _get_aiter_top_k_kernel(
1034 + is_prefill=False,
1035 + compress_ratio=compress_ratio,
1036 + num_rows=num_rows,
1037 + max_valid_seq_len=max_compressed_seq_len,
1038 + )
1039 + if aiter_topk_kernel is not None:
1040 + _launch_aiter_top_k_per_row_decode(
1041 + aiter_topk_kernel,
1042 + logits,
1043 + decode_metadata.seq_lens,
1044 + topk_indices,
1045 + topk_tokens,
1046 + )
1047 + else:
1048 + torch.ops._C.top_k_per_row_decode(
1049 + logits,
1050 + next_n,
1051 + decode_metadata.seq_lens,
1052 + topk_indices,
1053 + num_rows,
1054 + logits.stride(0),
1055 + logits.stride(1),
1056 + topk_tokens,
1057 + )
1058
1059 if decode_metadata.requires_padding:
1060 # if padded, we need to unpack