From dffbb714e4e8e4b95ccc888df98d47c7d2cef78d Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:34:24 -0400 Subject: [PATCH] [ROCm][Perf] Optimize DSV4.1 K=512 decode top-k on gfx950 (#56743) Signed-off-by: fai --- csrc/libtorch_stable/sampler.cu | 263 ++++++++++++++---- tests/kernels/test_top_k_per_row.py | 189 +++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 12 + 3 files changed, 406 insertions(+), 58 deletions(-) diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index ada1c4ba39..b8d64f7af1 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads, } template __device__ bool processHistogramStep( const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern, @@ -180,10 +181,15 @@ __device__ bool processHistogramStep( << patternShift; } + int negativeInfinityCount = 0; auto distributeToBins = [&](float logit, int /* idx */ = 0) { if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - atomicAdd(&smemFinal.histo.data[binIdx], 1); + if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) { + ++negativeInfinityCount; + } else { + atomicAdd(&smemFinal.histo.data[binIdx], 1); + } } }; @@ -198,6 +204,12 @@ __device__ bool processHistogramStep( distributeToBins(logit, idx); } } + if constexpr (useTopK512Optimization) { + if (negativeInfinityCount > 0) { + const auto bin = extractBinIdx(-INFINITY); + atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount); + } + } // Make sure the histogram is ready. __syncthreads(); @@ -254,6 +266,12 @@ __device__ bool processHistogramStep( // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; + // Resolve a -inf cutoff with exact tie emission, keeping sort padding out. + const bool finalBinFits = + smemFinalBinSize[0] <= kNumFinalItems && + !(useTopK512Optimization && step < 3 && + isPartialMatch(-INFINITY, logitPattern) && + thresholdBinIdx == extractBinIdx(-INFINITY)); auto processBins = [&](float logit, int idx) { if (isPartialMatch(logit, logitPattern)) { @@ -262,8 +280,7 @@ __device__ bool processHistogramStep( // 1. This is step 0 and the threshold bin is small enough (no step 1) // 2. This is step >= 1 (where pattern matching filters correctly) // This prevents duplicates when step 0 and step 1 both run. - bool shouldWriteDirectly = - (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1); + bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1); if (binIdx < thresholdBinIdx && shouldWriteDirectly) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -279,8 +296,7 @@ __device__ bool processHistogramStep( } if constexpr (step < 3) { // Only fill the final items for sorting if the threshold bin fits - if (binIdx == thresholdBinIdx && - smemFinalBinSize[0] <= kNumFinalItems) { + if (binIdx == thresholdBinIdx && finalBinFits) { int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1); smemFinal.items.logits[dstIdx] = logit; if constexpr (mergeBlocks) { @@ -325,20 +341,21 @@ __device__ bool processHistogramStep( __syncthreads(); // Check if we should continue to next step - return smemFinalBinSize[0] > kNumFinalItems; + return !finalBinFits; } // Follows half - 11 - 11 - 10 bit iterations // Keep the adaptive kernel's device instantiation separate from legacy calls. template + bool deviceLengthAware = false, bool useTopK512Optimization = false> static __device__ void topKPerRowJob(const int* indices, const float* logits, int rowStart, int rowEnd, int* outIndices, float* outLogits, int stride1, int topK) { static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks); + static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024); // The number of slots for the final pass. - static constexpr int kNumFinalItems = 2048; + static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048; // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; @@ -404,7 +421,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, rowIt += kNumThreadsPerBlock) { outIndices[rowIt] = -1; if constexpr (multipleBlocksPerRow) { - outLogits[rowIt] = -FLT_MAX; + outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX; } } @@ -422,7 +439,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 0: Process first 11 bits of half representation bool continueToNextStep = processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -431,7 +449,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 1: Process next 11 bits continueToNextStep = processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -441,7 +460,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 2: Process next 11 bits continueToNextStep = processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -450,7 +470,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, if (continueToNextStep) { // Step 3: Process last 10 bits processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -460,50 +481,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // The histogram did not proceed to the final 10 bits, therefore we need to // sort the final items The logits of the elements to be sorted in the final // pass. - if constexpr (useRadixSort) { - // Sorting with radix sort - float finalLogits[kNumFinalItemsPerThread]; - // The indices of the elements to be sorted in the final pass. - int finalIndices[kNumFinalItemsPerThread]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - finalLogits[ii] = -FLT_MAX; - } - - // Read the elements from SMEM. -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < smemFinalDstIdx[0]) { - finalLogits[ii] = smemFinal.items.logits[srcIdx]; - finalIndices[ii] = smemFinal.items.indices[srcIdx]; - } - } - // Make sure the shared memory has been read. - __syncthreads(); - - // Sort the elements. - FinalSort(smemFinal.finalSort) - .SortDescendingBlockedToStriped(finalLogits, finalIndices); - - // Copy the data back to the shared memory storage. - int baseIdx = smemFoundTopKValues[0]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - int dstIdx = baseIdx + srcIdx; - - if (dstIdx < topK) { - smemOutput[dstIdx] = finalIndices[ii]; - if constexpr (multipleBlocksPerRow) { - reinterpret_cast(smemOutput + topK)[dstIdx] = - finalLogits[ii]; - } - } - } - } else { + auto insertionSort = [&]() { // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; for (int i = threadIdx.x; i < smemFinalDstIdx[0]; @@ -525,6 +503,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, } } } + }; + if constexpr (useRadixSort) { + if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) { + insertionSort(); + } else { + // Sorting with radix sort + float finalLogits[kNumFinalItemsPerThread]; + // The indices of the elements to be sorted in the final pass. + int finalIndices[kNumFinalItemsPerThread]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX; + if constexpr (useTopK512Optimization) finalIndices[ii] = -1; + } + + // Read the elements from SMEM. +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + if (srcIdx < smemFinalDstIdx[0]) { + finalLogits[ii] = smemFinal.items.logits[srcIdx]; + finalIndices[ii] = smemFinal.items.indices[srcIdx]; + } + } + // Make sure the shared memory has been read. + __syncthreads(); + + // Sort the elements. + FinalSort(smemFinal.finalSort) + .SortDescendingBlockedToStriped(finalLogits, finalIndices); + + // Copy the data back to the shared memory storage. + int baseIdx = smemFoundTopKValues[0]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + int dstIdx = baseIdx + srcIdx; + + if (dstIdx < topK) { + smemOutput[dstIdx] = finalIndices[ii]; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[dstIdx] = + finalLogits[ii]; + } + } + } + } + } else { + insertionSort(); } __syncthreads(); } @@ -669,6 +698,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware( stride1, topK); } +#ifdef USE_ROCM +// The grid and workspace stay fixed while graph replays change device lengths. +__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows, + int maxBlocks) { + int activeBlocks = 1; + if (numRows <= 32 && rowLength > 16 * 1024) { + activeBlocks = rowLength <= 128 * 1024 ? 8 + : rowLength <= 256 * 1024 ? 10 + : 16; + } else if (numRows <= 128 && rowLength > 64 * 1024) { + activeBlocks = rowLength <= 512 * 1024 ? 4 : 10; + } else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) { + activeBlocks = numRows <= 160 ? 3 : 2; + } + // Every block in a merged row must contain at least k real indices. + return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512)); +} + +template +static __global__ +__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware( + const float* logits, const int* seqLens, int* outIndices, + int* outIndicesAux, float* outLogitsAux, int stride0, int next_n, + int seqLensIs2D, int maxBlocks) { + #if defined(__gfx950__) + constexpr int topK = 512; + constexpr int kNumBins = 2048; + const int rowIdx = blockIdx.x; + const int rowEnd = + seqLensIs2D + ? max(0, seqLens[rowIdx]) + : max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1); + const int activeBlocks = + gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks); + outIndices += static_cast(rowIdx) * topK; + if constexpr (mergeBlocks) { + if (activeBlocks == 1) return; + const int64_t offset = static_cast(rowIdx) * maxBlocks * topK; + topKPerRowJob(outIndicesAux + offset, outLogitsAux + offset, 0, + activeBlocks * topK, outIndices, nullptr, 1, topK); + } else { + if (blockIdx.y >= activeBlocks) return; + logits += static_cast(rowIdx) * stride0; + if (activeBlocks == 1) { + topKPerRowJob(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1, + topK); + return; + } + const int blockSize = rowEnd / activeBlocks; + const int rowStart = blockSize * blockIdx.y; + const int blockRowEnd = + blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize; + const int64_t offset = + (static_cast(rowIdx) * maxBlocks + blockIdx.y) * topK; + topKPerRowJob(nullptr, logits, rowStart, blockRowEnd, + outIndicesAux + offset, outLogitsAux + offset, 1, topK); + } + #endif +} +#endif + } // namespace vllm void apply_repetition_penalties_( @@ -731,6 +824,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, // the same seq_len and the kernel computes the per-row offset itself. int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0; +#ifdef USE_ROCM + const bool useGfx950DeviceLengthAwareTopK512 = + topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 && + numColumns <= 1024 * 1024 && + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0; + if (useGfx950DeviceLengthAwareTopK512) { + int multipleBlocksPerRowConfig = 1; + if (numRows <= 32 && numColumns > 16 * 1024) { + multipleBlocksPerRowConfig = 16; + } else if (numRows <= 128 && numColumns > 64 * 1024) { + multipleBlocksPerRowConfig = 10; + } else if (numRows <= 160 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 3; + } else if (numRows <= 256 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 2; + } + + constexpr int kNumThreadsPerSplitBlock = 1024; + if (multipleBlocksPerRowConfig == 1) { + vllm::topKPerRowDecode512DeviceLengthAware + <<>>(logits.const_data_ptr(), + seqLens.const_data_ptr(), + indices.mutable_data_ptr(), nullptr, nullptr, + static_cast(stride0), static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } + const auto outIndicesAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Int, std::nullopt, logits.device()); + const auto outLogitsAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Float, std::nullopt, logits.device()); + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + logits.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), static_cast(stride0), + static_cast(next_n), seqLensIs2D, multipleBlocksPerRowConfig); + constexpr int kNumThreadsPerBlockMerge = 1024; + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + nullptr, seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), 0, static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } +#endif + if (numColumns < kSortingAlgorithmThreshold) { // Use insertion sort vllm::topKPerRowDecode diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index c629ae15de..591362020e 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -22,11 +22,22 @@ def _has_device_capability(major: int) -> bool: return current_platform.is_cuda() and current_platform.has_device_capability(major) +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + # DeepSelect is compiled for sm_100a/sm_103a only. requires_sm100 = pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="DeepSelect requires SM100a/SM103a", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), reason="This test exercises the gfx950 launch configuration" +) COOPERATIVE_TOPK_BACKEND = pytest.param( @@ -487,6 +498,150 @@ def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None: ) +def _assert_exact_topk( + logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor +) -> None: + ends = row_ends.reshape(-1, 1).clamp_min(0) + valid = (indices >= 0) & (indices < ends) + assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten()) + assert torch.all(valid | (indices == -1)) + ordered = indices.sort(1).values + assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any() + + width = int(ends.max()) + visible = logits[:, :width].clone() + columns = torch.arange(width, device=logits.device) + visible.masked_fill_(columns[None, :] >= ends, float("-inf")) + expected = visible.topk(min(indices.shape[1], width), dim=1).values + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, float("-inf")) + selected = selected.sort(1, descending=True).values + assert torch.equal(selected[:, : expected.shape[1]], expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None: + """Honor causal offsets and exact compressed lengths, including short rows.""" + next_n, width, top_k = 6, 16_384, 512 + lengths = torch.tensor( + [0, 1, 511, 512, 513, 1023, 1024, 16_000], + dtype=torch.int32, + device="cuda", + ) + offsets = torch.arange(next_n, dtype=torch.int32, device="cuda") + row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten() + rows = row_ends.numel() + storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda") + logits = storage[:, 1 : width + 1] + logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda") + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + _assert_exact_topk(logits, indices, row_ends) + + +@pytest.mark.parametrize( + ("rows", "row_length"), + [ + pytest.param(128, 65_536, id="single-before-64k"), + pytest.param(128, 65_537, id="four-splits-after-64k"), + pytest.param(129, 262_144, id="single-before-three-split-length"), + pytest.param(129, 262_145, id="three-splits-after-256k"), + pytest.param(160, 262_145, id="three-splits-upper-row-bound"), + pytest.param(161, 262_144, id="single-before-two-split-length"), + pytest.param(161, 262_145, id="two-splits-after-256k"), + pytest.param(256, 262_145, id="two-splits-upper-row-bound"), + pytest.param(257, 262_145, id="large-row-single-block"), + pytest.param(384, 262_145, id="optimized-row-limit"), + pytest.param(385, 262_145, id="native-fallback"), + ], +) +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries( + rows: int, row_length: int +) -> None: + width, top_k = 262_145, 512 + row = torch.arange(width, dtype=torch.float32, device="cuda") + logits = row.expand(rows, -1) + lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda") + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + + _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width) + + expected = torch.arange( + row_length - top_k, row_length, dtype=torch.int32, device="cuda" + ).expand_as(indices) + assert torch.equal(indices.sort(1).values, expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: + """Preserve visible top-k scores when masks and lengths change on replay.""" + rows, next_n, width, top_k = 24, 6, 262_145, 512 + logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda") + lengths = torch.full( + (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda" + ) + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + logits[:, :70_000] = float("-inf") + bin_sizes = [128, 129, 1024, 1025] + for row in range(rows): + pattern = row % 9 + if pattern < 4: + count = [0, 1, 511, 513][pattern] + logits[row, :count] = torch.arange( + count, dtype=torch.float32, device="cuda" + ) + elif pattern == 4: + logits[row, :511] = 0 + logits[row, 511] = torch.finfo(torch.float32).min + else: + count = bin_sizes[pattern - 5] + higher = max(0, top_k - count // 2) + logits[row, :higher] = 2 + logits[row, higher : higher + count] = 1 + logits[row, higher] = 1 + 2**-23 + + def run() -> None: + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + + run() + _assert_exact_topk(logits, indices, lengths) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + bounds = [ + 0, + 511, + 512, + 513, + 16_384, + 16_385, + 131_072, + 131_073, + 262_144, + 262_145, + ] + row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda") + repeats = (rows + len(bounds) - 1) // len(bounds) + lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows]) + logits[:] = torch.arange(width, dtype=torch.float32, device="cuda") + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @torch.inference_mode() def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: @@ -620,6 +775,31 @@ def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: compress_ratio=4, num_rows=257, max_valid_seq_len=250_000, + num_columns=524_288, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=385, + max_valid_seq_len=250_000, + num_columns=524_288, + topk_tokens=512, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=50_000, + num_columns=1_048_577, + topk_tokens=512, on_gfx950=True, ), decode_kernel, @@ -638,6 +818,15 @@ def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: ) ineligible_cases = [ dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True), + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=1_048_576, + num_columns=1_048_576, + topk_tokens=512, + on_gfx950=True, + ), dict( is_prefill=False, compress_ratio=4, diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 44b6017a05..0e342cb0e0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -58,6 +58,7 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 _GFX950_C4A_NATIVE_MAX_ROWS = 256 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 @@ -69,6 +70,8 @@ def _get_aiter_top_k_kernel( compress_ratio: int, num_rows: int, max_valid_seq_len: int | None = None, + num_columns: int | None = None, + topk_tokens: int = 1024, on_gfx950: bool = _ON_GFX950, ) -> Callable[..., None] | None: if compress_ratio <= 1 or not on_gfx950: @@ -76,6 +79,13 @@ def _get_aiter_top_k_kernel( if not is_prefill: assert max_valid_seq_len is not None + if ( + topk_tokens == 512 + and 0 < num_rows <= 384 + and num_columns is not None + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS + ): + return None # AITER v0.1.19 decode is one-block only. This measured gfx950 # FP32/k=1024 compressed-row boundary is independent of the native # split-count boundary in sampler.cu. @@ -1261,6 +1271,8 @@ def rocm_aiter_sparse_attn_indexer( compress_ratio=compress_ratio, num_rows=num_rows, max_valid_seq_len=max_compressed_seq_len, + num_columns=logits.shape[1], + topk_tokens=topk_tokens, ) if aiter_topk_kernel is not None: _launch_aiter_top_k_per_row_decode(