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

[ROCm][Perf] Optimize DSV4.1 K=512 decode top-k on gfx950 (#56743)

Signed-off-by: fai <fangzhouai@gmail.com>

Fangzhou Ai committed Sep 15, 2026 at 14:34 UTC dffbb714e4e8e4b95ccc888df98d47c7d2cef78d
3 files changed +406 -58
csrc/libtorch_stable/sampler.cu
+205 -58
@@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
154 }
155
156 template <int step, int kNumThreadsPerBlock, int kNumBins, int kNumFinalItems,
157 - bool multipleBlocksPerRow, bool mergeBlocks, typename SmemFinalType,
157 + bool multipleBlocksPerRow, bool mergeBlocks,
158 + bool useTopK512Optimization, typename SmemFinalType,
159 typename SmemOutputType>
160 __device__ bool processHistogramStep(
161 const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern,
@@ -180,10 +181,15 @@ __device__ bool processHistogramStep(
181 << patternShift;
182 }
183
184 + int negativeInfinityCount = 0;
185 auto distributeToBins = [&](float logit, int /* idx */ = 0) {
186 if (isPartialMatch<patternShift>(logit, logitPattern)) {
187 uint32_t binIdx = extractBinIdx<step>(logit);
186 - atomicAdd(&smemFinal.histo.data[binIdx], 1);
188 + if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) {
189 + ++negativeInfinityCount;
190 + } else {
191 + atomicAdd(&smemFinal.histo.data[binIdx], 1);
192 + }
193 }
194 };
195
@@ -198,6 +204,12 @@ __device__ bool processHistogramStep(
204 distributeToBins(logit, idx);
205 }
206 }
207 + if constexpr (useTopK512Optimization) {
208 + if (negativeInfinityCount > 0) {
209 + const auto bin = extractBinIdx<step>(-INFINITY);
210 + atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount);
211 + }
212 + }
213 // Make sure the histogram is ready.
214 __syncthreads();
215
@@ -254,6 +266,12 @@ __device__ bool processHistogramStep(
266
267 // The threshold bin.
268 thresholdBinIdx = smemThresholdBinIdx[0];
269 + // Resolve a -inf cutoff with exact tie emission, keeping sort padding out.
270 + const bool finalBinFits =
271 + smemFinalBinSize[0] <= kNumFinalItems &&
272 + !(useTopK512Optimization && step < 3 &&
273 + isPartialMatch<patternShift>(-INFINITY, logitPattern) &&
274 + thresholdBinIdx == extractBinIdx<step>(-INFINITY));
275
276 auto processBins = [&](float logit, int idx) {
277 if (isPartialMatch<patternShift>(logit, logitPattern)) {
@@ -262,8 +280,7 @@ __device__ bool processHistogramStep(
280 // 1. This is step 0 and the threshold bin is small enough (no step 1)
281 // 2. This is step >= 1 (where pattern matching filters correctly)
282 // This prevents duplicates when step 0 and step 1 both run.
265 - bool shouldWriteDirectly =
266 - (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
283 + bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1);
284 if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
285 // The element is part of the top-k selection
286 int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
@@ -279,8 +296,7 @@ __device__ bool processHistogramStep(
296 }
297 if constexpr (step < 3) {
298 // Only fill the final items for sorting if the threshold bin fits
282 - if (binIdx == thresholdBinIdx &&
283 - smemFinalBinSize[0] <= kNumFinalItems) {
299 + if (binIdx == thresholdBinIdx && finalBinFits) {
300 int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
301 smemFinal.items.logits[dstIdx] = logit;
302 if constexpr (mergeBlocks) {
@@ -325,20 +341,21 @@ __device__ bool processHistogramStep(
341 __syncthreads();
342
343 // Check if we should continue to next step
328 - return smemFinalBinSize[0] > kNumFinalItems;
344 + return !finalBinFits;
345 }
346
347 // Follows half - 11 - 11 - 10 bit iterations
348 // Keep the adaptive kernel's device instantiation separate from legacy calls.
349 template <int kNumThreadsPerBlock, int kNumBins, bool useRadixSort,
350 bool multipleBlocksPerRow = false, bool mergeBlocks = false,
335 - bool deviceLengthAware = false>
351 + bool deviceLengthAware = false, bool useTopK512Optimization = false>
352 static __device__ void topKPerRowJob(const int* indices, const float* logits,
353 int rowStart, int rowEnd, int* outIndices,
354 float* outLogits, int stride1, int topK) {
355 static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks);
356 + static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024);
357 // The number of slots for the final pass.
341 - static constexpr int kNumFinalItems = 2048;
358 + static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048;
359 // The number of elements per thread for the final sort.
360 static constexpr int kNumFinalItemsPerThread =
361 kNumFinalItems / kNumThreadsPerBlock;
@@ -404,7 +421,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
421 rowIt += kNumThreadsPerBlock) {
422 outIndices[rowIt] = -1;
423 if constexpr (multipleBlocksPerRow) {
407 - outLogits[rowIt] = -FLT_MAX;
424 + outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX;
425 }
426 }
427
@@ -422,7 +439,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
439 // Step 0: Process first 11 bits of half representation
440 bool continueToNextStep =
441 processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
425 - multipleBlocksPerRow, mergeBlocks>(
442 + multipleBlocksPerRow, mergeBlocks,
443 + useTopK512Optimization>(
444 indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
445 smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
446 smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -431,7 +449,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
449 // Step 1: Process next 11 bits
450 continueToNextStep =
451 processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
434 - multipleBlocksPerRow, mergeBlocks>(
452 + multipleBlocksPerRow, mergeBlocks,
453 + useTopK512Optimization>(
454 indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
455 smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
456 smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -441,7 +460,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
460 // Step 2: Process next 11 bits
461 continueToNextStep =
462 processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
444 - multipleBlocksPerRow, mergeBlocks>(
463 + multipleBlocksPerRow, mergeBlocks,
464 + useTopK512Optimization>(
465 indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
466 smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
467 smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -450,7 +470,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
470 if (continueToNextStep) {
471 // Step 3: Process last 10 bits
472 processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
453 - multipleBlocksPerRow, mergeBlocks>(
473 + multipleBlocksPerRow, mergeBlocks,
474 + useTopK512Optimization>(
475 indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
476 smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
477 smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -460,50 +481,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
481 // The histogram did not proceed to the final 10 bits, therefore we need to
482 // sort the final items The logits of the elements to be sorted in the final
483 // pass.
463 - if constexpr (useRadixSort) {
464 - // Sorting with radix sort
465 - float finalLogits[kNumFinalItemsPerThread];
466 - // The indices of the elements to be sorted in the final pass.
467 - int finalIndices[kNumFinalItemsPerThread];
468 -
469 -#pragma unroll
470 - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
471 - finalLogits[ii] = -FLT_MAX;
472 - }
473 -
474 - // Read the elements from SMEM.
475 -#pragma unroll
476 - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
477 - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
478 - if (srcIdx < smemFinalDstIdx[0]) {
479 - finalLogits[ii] = smemFinal.items.logits[srcIdx];
480 - finalIndices[ii] = smemFinal.items.indices[srcIdx];
481 - }
482 - }
483 - // Make sure the shared memory has been read.
484 - __syncthreads();
485 -
486 - // Sort the elements.
487 - FinalSort(smemFinal.finalSort)
488 - .SortDescendingBlockedToStriped(finalLogits, finalIndices);
489 -
490 - // Copy the data back to the shared memory storage.
491 - int baseIdx = smemFoundTopKValues[0];
492 -
493 -#pragma unroll
494 - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
495 - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
496 - int dstIdx = baseIdx + srcIdx;
497 -
498 - if (dstIdx < topK) {
499 - smemOutput[dstIdx] = finalIndices[ii];
500 - if constexpr (multipleBlocksPerRow) {
501 - reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
502 - finalLogits[ii];
503 - }
504 - }
505 - }
506 - } else {
484 + auto insertionSort = [&]() {
485 // Sorting with insertion sort
486 auto baseIdx = smemFoundTopKValues[0];
487 for (int i = threadIdx.x; i < smemFinalDstIdx[0];
@@ -525,6 +503,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
503 }
504 }
505 }
506 + };
507 + if constexpr (useRadixSort) {
508 + if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) {
509 + insertionSort();
510 + } else {
511 + // Sorting with radix sort
512 + float finalLogits[kNumFinalItemsPerThread];
513 + // The indices of the elements to be sorted in the final pass.
514 + int finalIndices[kNumFinalItemsPerThread];
515 +
516 +#pragma unroll
517 + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
518 + finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX;
519 + if constexpr (useTopK512Optimization) finalIndices[ii] = -1;
520 + }
521 +
522 + // Read the elements from SMEM.
523 +#pragma unroll
524 + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
525 + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
526 + if (srcIdx < smemFinalDstIdx[0]) {
527 + finalLogits[ii] = smemFinal.items.logits[srcIdx];
528 + finalIndices[ii] = smemFinal.items.indices[srcIdx];
529 + }
530 + }
531 + // Make sure the shared memory has been read.
532 + __syncthreads();
533 +
534 + // Sort the elements.
535 + FinalSort(smemFinal.finalSort)
536 + .SortDescendingBlockedToStriped(finalLogits, finalIndices);
537 +
538 + // Copy the data back to the shared memory storage.
539 + int baseIdx = smemFoundTopKValues[0];
540 +
541 +#pragma unroll
542 + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
543 + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
544 + int dstIdx = baseIdx + srcIdx;
545 +
546 + if (dstIdx < topK) {
547 + smemOutput[dstIdx] = finalIndices[ii];
548 + if constexpr (multipleBlocksPerRow) {
549 + reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
550 + finalLogits[ii];
551 + }
552 + }
553 + }
554 + }
555 + } else {
556 + insertionSort();
557 }
558 __syncthreads();
559 }
@@ -669,6 +698,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware(
698 stride1, topK);
699 }
700
701 +#ifdef USE_ROCM
702 +// The grid and workspace stay fixed while graph replays change device lengths.
703 +__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows,
704 + int maxBlocks) {
705 + int activeBlocks = 1;
706 + if (numRows <= 32 && rowLength > 16 * 1024) {
707 + activeBlocks = rowLength <= 128 * 1024 ? 8
708 + : rowLength <= 256 * 1024 ? 10
709 + : 16;
710 + } else if (numRows <= 128 && rowLength > 64 * 1024) {
711 + activeBlocks = rowLength <= 512 * 1024 ? 4 : 10;
712 + } else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) {
713 + activeBlocks = numRows <= 160 ? 3 : 2;
714 + }
715 + // Every block in a merged row must contain at least k real indices.
716 + return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512));
717 +}
718 +
719 +template <int kNumThreadsPerBlock, bool mergeBlocks = false>
720 +static __global__
721 +__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware(
722 + const float* logits, const int* seqLens, int* outIndices,
723 + int* outIndicesAux, float* outLogitsAux, int stride0, int next_n,
724 + int seqLensIs2D, int maxBlocks) {
725 + #if defined(__gfx950__)
726 + constexpr int topK = 512;
727 + constexpr int kNumBins = 2048;
728 + const int rowIdx = blockIdx.x;
729 + const int rowEnd =
730 + seqLensIs2D
731 + ? max(0, seqLens[rowIdx])
732 + : max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1);
733 + const int activeBlocks =
734 + gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks);
735 + outIndices += static_cast<int64_t>(rowIdx) * topK;
736 + if constexpr (mergeBlocks) {
737 + if (activeBlocks == 1) return;
738 + const int64_t offset = static_cast<int64_t>(rowIdx) * maxBlocks * topK;
739 + topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, true, false,
740 + true>(outIndicesAux + offset, outLogitsAux + offset, 0,
741 + activeBlocks * topK, outIndices, nullptr, 1, topK);
742 + } else {
743 + if (blockIdx.y >= activeBlocks) return;
744 + logits += static_cast<int64_t>(rowIdx) * stride0;
745 + if (activeBlocks == 1) {
746 + topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, false, false,
747 + true>(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1,
748 + topK);
749 + return;
750 + }
751 + const int blockSize = rowEnd / activeBlocks;
752 + const int rowStart = blockSize * blockIdx.y;
753 + const int blockRowEnd =
754 + blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize;
755 + const int64_t offset =
756 + (static_cast<int64_t>(rowIdx) * maxBlocks + blockIdx.y) * topK;
757 + topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, true, false, false,
758 + true>(nullptr, logits, rowStart, blockRowEnd,
759 + outIndicesAux + offset, outLogitsAux + offset, 1, topK);
760 + }
761 + #endif
762 +}
763 +#endif
764 +
765 } // namespace vllm
766
767 void apply_repetition_penalties_(
@@ -731,6 +824,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
824 // the same seq_len and the kernel computes the per-row offset itself.
825 int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0;
826
827 +#ifdef USE_ROCM
828 + const bool useGfx950DeviceLengthAwareTopK512 =
829 + topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 &&
830 + numColumns <= 1024 * 1024 &&
831 + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0;
832 + if (useGfx950DeviceLengthAwareTopK512) {
833 + int multipleBlocksPerRowConfig = 1;
834 + if (numRows <= 32 && numColumns > 16 * 1024) {
835 + multipleBlocksPerRowConfig = 16;
836 + } else if (numRows <= 128 && numColumns > 64 * 1024) {
837 + multipleBlocksPerRowConfig = 10;
838 + } else if (numRows <= 160 && numColumns > 256 * 1024) {
839 + multipleBlocksPerRowConfig = 3;
840 + } else if (numRows <= 256 && numColumns > 256 * 1024) {
841 + multipleBlocksPerRowConfig = 2;
842 + }
843 +
844 + constexpr int kNumThreadsPerSplitBlock = 1024;
845 + if (multipleBlocksPerRowConfig == 1) {
846 + vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerSplitBlock>
847 + <<<numRows, kNumThreadsPerSplitBlock, 2 * topK * sizeof(int32_t),
848 + stream>>>(logits.const_data_ptr<float>(),
849 + seqLens.const_data_ptr<int>(),
850 + indices.mutable_data_ptr<int>(), nullptr, nullptr,
851 + static_cast<int>(stride0), static_cast<int>(next_n),
852 + seqLensIs2D, multipleBlocksPerRowConfig);
853 + return;
854 + }
855 + const auto outIndicesAux = torch::stable::empty(
856 + {numRows, multipleBlocksPerRowConfig, topK},
857 + torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
858 + const auto outLogitsAux = torch::stable::empty(
859 + {numRows, multipleBlocksPerRowConfig, topK},
860 + torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
861 + vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerSplitBlock>
862 + <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerSplitBlock,
863 + 2 * topK * sizeof(int32_t), stream>>>(
864 + logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
865 + indices.mutable_data_ptr<int>(),
866 + outIndicesAux.mutable_data_ptr<int>(),
867 + outLogitsAux.mutable_data_ptr<float>(), static_cast<int>(stride0),
868 + static_cast<int>(next_n), seqLensIs2D, multipleBlocksPerRowConfig);
869 + constexpr int kNumThreadsPerBlockMerge = 1024;
870 + vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerBlockMerge, true>
871 + <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
872 + nullptr, seqLens.const_data_ptr<int>(),
873 + indices.mutable_data_ptr<int>(),
874 + outIndicesAux.mutable_data_ptr<int>(),
875 + outLogitsAux.mutable_data_ptr<float>(), 0, static_cast<int>(next_n),
876 + seqLensIs2D, multipleBlocksPerRowConfig);
877 + return;
878 + }
879 +#endif
880 +
881 if (numColumns < kSortingAlgorithmThreshold) {
882 // Use insertion sort
883 vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
tests/kernels/test_top_k_per_row.py
+189
@@ -22,11 +22,22 @@ 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(
@@ -487,6 +498,150 @@ def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None:
498 )
499
500
501 +def _assert_exact_topk(
502 + logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor
503 +) -> None:
504 + ends = row_ends.reshape(-1, 1).clamp_min(0)
505 + valid = (indices >= 0) & (indices < ends)
506 + assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten())
507 + assert torch.all(valid | (indices == -1))
508 + ordered = indices.sort(1).values
509 + assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any()
510 +
511 + width = int(ends.max())
512 + visible = logits[:, :width].clone()
513 + columns = torch.arange(width, device=logits.device)
514 + visible.masked_fill_(columns[None, :] >= ends, float("-inf"))
515 + expected = visible.topk(min(indices.shape[1], width), dim=1).values
516 + selected = logits.gather(1, indices.clamp_min(0).long())
517 + selected.masked_fill_(~valid, float("-inf"))
518 + selected = selected.sort(1, descending=True).values
519 + assert torch.equal(selected[:, : expected.shape[1]], expected)
520 +
521 +
522 +@requires_gfx950
523 +@torch.inference_mode()
524 +def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None:
525 + """Honor causal offsets and exact compressed lengths, including short rows."""
526 + next_n, width, top_k = 6, 16_384, 512
527 + lengths = torch.tensor(
528 + [0, 1, 511, 512, 513, 1023, 1024, 16_000],
529 + dtype=torch.int32,
530 + device="cuda",
531 + )
532 + offsets = torch.arange(next_n, dtype=torch.int32, device="cuda")
533 + row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten()
534 + rows = row_ends.numel()
535 + storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda")
536 + logits = storage[:, 1 : width + 1]
537 + logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda")
538 + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda")
539 + _run_topk_backend(
540 + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n
541 + )
542 + _assert_exact_topk(logits, indices, row_ends)
543 +
544 +
545 +@pytest.mark.parametrize(
546 + ("rows", "row_length"),
547 + [
548 + pytest.param(128, 65_536, id="single-before-64k"),
549 + pytest.param(128, 65_537, id="four-splits-after-64k"),
550 + pytest.param(129, 262_144, id="single-before-three-split-length"),
551 + pytest.param(129, 262_145, id="three-splits-after-256k"),
552 + pytest.param(160, 262_145, id="three-splits-upper-row-bound"),
553 + pytest.param(161, 262_144, id="single-before-two-split-length"),
554 + pytest.param(161, 262_145, id="two-splits-after-256k"),
555 + pytest.param(256, 262_145, id="two-splits-upper-row-bound"),
556 + pytest.param(257, 262_145, id="large-row-single-block"),
557 + pytest.param(384, 262_145, id="optimized-row-limit"),
558 + pytest.param(385, 262_145, id="native-fallback"),
559 + ],
560 +)
561 +@requires_gfx950
562 +@torch.inference_mode()
563 +def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries(
564 + rows: int, row_length: int
565 +) -> None:
566 + width, top_k = 262_145, 512
567 + row = torch.arange(width, dtype=torch.float32, device="cuda")
568 + logits = row.expand(rows, -1)
569 + lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda")
570 + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda")
571 +
572 + _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width)
573 +
574 + expected = torch.arange(
575 + row_length - top_k, row_length, dtype=torch.int32, device="cuda"
576 + ).expand_as(indices)
577 + assert torch.equal(indices.sort(1).values, expected)
578 +
579 +
580 +@requires_gfx950
581 +@torch.inference_mode()
582 +def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None:
583 + """Preserve visible top-k scores when masks and lengths change on replay."""
584 + rows, next_n, width, top_k = 24, 6, 262_145, 512
585 + logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda")
586 + lengths = torch.full(
587 + (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda"
588 + )
589 + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda")
590 + logits[:, :70_000] = float("-inf")
591 + bin_sizes = [128, 129, 1024, 1025]
592 + for row in range(rows):
593 + pattern = row % 9
594 + if pattern < 4:
595 + count = [0, 1, 511, 513][pattern]
596 + logits[row, :count] = torch.arange(
597 + count, dtype=torch.float32, device="cuda"
598 + )
599 + elif pattern == 4:
600 + logits[row, :511] = 0
601 + logits[row, 511] = torch.finfo(torch.float32).min
602 + else:
603 + count = bin_sizes[pattern - 5]
604 + higher = max(0, top_k - count // 2)
605 + logits[row, :higher] = 2
606 + logits[row, higher : higher + count] = 1
607 + logits[row, higher] = 1 + 2**-23
608 +
609 + def run() -> None:
610 + _run_topk_backend(
611 + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n
612 + )
613 +
614 + run()
615 + _assert_exact_topk(logits, indices, lengths)
616 + torch.accelerator.synchronize()
617 + graph = torch.cuda.CUDAGraph()
618 + with torch.cuda.graph(graph):
619 + run()
620 + indices.fill_(-777)
621 + graph.replay()
622 + _assert_exact_topk(logits, indices, lengths)
623 +
624 + bounds = [
625 + 0,
626 + 511,
627 + 512,
628 + 513,
629 + 16_384,
630 + 16_385,
631 + 131_072,
632 + 131_073,
633 + 262_144,
634 + 262_145,
635 + ]
636 + row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda")
637 + repeats = (rows + len(bounds) - 1) // len(bounds)
638 + lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows])
639 + logits[:] = torch.arange(width, dtype=torch.float32, device="cuda")
640 + indices.fill_(-777)
641 + graph.replay()
642 + _assert_exact_topk(logits, indices, lengths)
643 +
644 +
645 @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
646 @torch.inference_mode()
647 def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
@@ -620,6 +775,31 @@ def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None:
775 compress_ratio=4,
776 num_rows=257,
777 max_valid_seq_len=250_000,
778 + num_columns=524_288,
779 + on_gfx950=True,
780 + ),
781 + decode_kernel,
782 + ),
783 + (
784 + dict(
785 + is_prefill=False,
786 + compress_ratio=2,
787 + num_rows=385,
788 + max_valid_seq_len=250_000,
789 + num_columns=524_288,
790 + topk_tokens=512,
791 + on_gfx950=True,
792 + ),
793 + decode_kernel,
794 + ),
795 + (
796 + dict(
797 + is_prefill=False,
798 + compress_ratio=2,
799 + num_rows=384,
800 + max_valid_seq_len=50_000,
801 + num_columns=1_048_577,
802 + topk_tokens=512,
803 on_gfx950=True,
804 ),
805 decode_kernel,
@@ -638,6 +818,15 @@ def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None:
818 )
819 ineligible_cases = [
820 dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True),
821 + dict(
822 + is_prefill=False,
823 + compress_ratio=2,
824 + num_rows=384,
825 + max_valid_seq_len=1_048_576,
826 + num_columns=1_048_576,
827 + topk_tokens=512,
828 + on_gfx950=True,
829 + ),
830 dict(
831 is_prefill=False,
832 compress_ratio=4,
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+12
@@ -58,6 +58,7 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
58
59 _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
60 _GFX950_C4A_NATIVE_MAX_ROWS = 256
61 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024
62 # Conservative perf gate, not a correctness bound: OPUS is correct for any query
63 # count, but Triton stays faster below this measured crossover.
64 _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024
@@ -69,6 +70,8 @@ def _get_aiter_top_k_kernel(
70 compress_ratio: int,
71 num_rows: int,
72 max_valid_seq_len: int | None = None,
73 + num_columns: int | None = None,
74 + topk_tokens: int = 1024,
75 on_gfx950: bool = _ON_GFX950,
76 ) -> Callable[..., None] | None:
77 if compress_ratio <= 1 or not on_gfx950:
@@ -76,6 +79,13 @@ def _get_aiter_top_k_kernel(
79
80 if not is_prefill:
81 assert max_valid_seq_len is not None
82 + if (
83 + topk_tokens == 512
84 + and 0 < num_rows <= 384
85 + and num_columns is not None
86 + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS
87 + ):
88 + return None
89 # AITER v0.1.19 decode is one-block only. This measured gfx950
90 # FP32/k=1024 compressed-row boundary is independent of the native
91 # split-count boundary in sampler.cu.
@@ -1261,6 +1271,8 @@ def rocm_aiter_sparse_attn_indexer(
1271 compress_ratio=compress_ratio,
1272 num_rows=num_rows,
1273 max_valid_seq_len=max_compressed_seq_len,
1274 + num_columns=logits.shape[1],
1275 + topk_tokens=topk_tokens,
1276 )
1277 if aiter_topk_kernel is not None:
1278 _launch_aiter_top_k_per_row_decode(