iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: fai <fangzhouai@gmail.com>
This commit is contained in:
Fangzhou Ai
2026-09-15 18:34:24 +00:00
committed by GitHub
parent 2f46a1d9c6
commit dffbb714e4
3 changed files with 406 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/sampler.cu
+205 -58
View File
@@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
}
template <int step, int kNumThreadsPerBlock, int kNumBins, int kNumFinalItems,
bool multipleBlocksPerRow, bool mergeBlocks, typename SmemFinalType,
bool multipleBlocksPerRow, bool mergeBlocks,
bool useTopK512Optimization, typename SmemFinalType,
typename SmemOutputType>
__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<patternShift>(logit, logitPattern)) {
uint32_t binIdx = extractBinIdx<step>(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<step>(-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<patternShift>(-INFINITY, logitPattern) &&
thresholdBinIdx == extractBinIdx<step>(-INFINITY));
auto processBins = [&](float logit, int idx) {
if (isPartialMatch<patternShift>(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 <int kNumThreadsPerBlock, int kNumBins, bool useRadixSort,
bool multipleBlocksPerRow = false, bool mergeBlocks = false,
bool deviceLengthAware = false>
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<float*>(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<float*>(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 <int kNumThreadsPerBlock, bool mergeBlocks = false>
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<int64_t>(rowIdx) * topK;
if constexpr (mergeBlocks) {
if (activeBlocks == 1) return;
const int64_t offset = static_cast<int64_t>(rowIdx) * maxBlocks * topK;
topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, true, false,
true>(outIndicesAux + offset, outLogitsAux + offset, 0,
activeBlocks * topK, outIndices, nullptr, 1, topK);
} else {
if (blockIdx.y >= activeBlocks) return;
logits += static_cast<int64_t>(rowIdx) * stride0;
if (activeBlocks == 1) {
topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, false, false,
true>(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<int64_t>(rowIdx) * maxBlocks + blockIdx.y) * topK;
topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, true, false, false,
true>(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<kNumThreadsPerSplitBlock>
<<<numRows, kNumThreadsPerSplitBlock, 2 * topK * sizeof(int32_t),
stream>>>(logits.const_data_ptr<float>(),
seqLens.const_data_ptr<int>(),
indices.mutable_data_ptr<int>(), nullptr, nullptr,
static_cast<int>(stride0), static_cast<int>(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<kNumThreadsPerSplitBlock>
<<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerSplitBlock,
2 * topK * sizeof(int32_t), stream>>>(
logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
indices.mutable_data_ptr<int>(),
outIndicesAux.mutable_data_ptr<int>(),
outLogitsAux.mutable_data_ptr<float>(), static_cast<int>(stride0),
static_cast<int>(next_n), seqLensIs2D, multipleBlocksPerRowConfig);
constexpr int kNumThreadsPerBlockMerge = 1024;
vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerBlockMerge, true>
<<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
nullptr, seqLens.const_data_ptr<int>(),
indices.mutable_data_ptr<int>(),
outIndicesAux.mutable_data_ptr<int>(),
outLogitsAux.mutable_data_ptr<float>(), 0, static_cast<int>(next_n),
seqLensIs2D, multipleBlocksPerRowConfig);
return;
}
#endif
if (numColumns < kSortingAlgorithmThreshold) {
// Use insertion sort
vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
tests/kernels/test_top_k_per_row.py
+189
View File
@@ -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,
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+12
View File
@@ -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(