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

[Perf][Kernel] Add sampled filtering for persistent top-k (#56346)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Misha Goin committed Sep 15, 2026 at 17:15 UTC c6fa1f05d1bc627171c3a58c05f23a1c41202ac1
5 files changed +556 -35
benchmarks/kernels/benchmark_persistent_topk.py new
+134
@@ -0,0 +1,134 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows.
4 +
5 +Run on a reserved GPU, for example:
6 + chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py
7 +
8 +Use --capacity to emulate graphs captured at a larger maximum context length.
9 +Run the same command on the base and candidate builds to compare JSON results.
10 +"""
11 +
12 +import argparse
13 +import itertools
14 +import json
15 +import statistics
16 +from pathlib import Path
17 +
18 +import torch
19 +from flashinfer.testing import bench_gpu_time_with_cupti
20 +
21 +import vllm._custom_ops # noqa: F401
22 +from vllm.platforms import current_platform
23 +
24 +
25 +def benchmark(rows: int, width: int, k: int, args) -> dict:
26 + capacity = args.capacity or width
27 + if capacity < width:
28 + raise ValueError("Capacity must be at least the valid-length upper bound")
29 + torch.manual_seed(args.seed)
30 + logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
31 + if args.full_length:
32 + lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32)
33 + else:
34 + lengths = torch.randint(
35 + int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32
36 + )
37 + invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None]
38 + logits.masked_fill_(invalid, float("nan"))
39 + out = torch.empty((rows, k), device="cuda", dtype=torch.int32)
40 + workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8)
41 + cooperative = (
42 + args.backend == "auto"
43 + and rows <= 64
44 + and capacity % 4 == 0
45 + and current_platform.has_device_capability(90)
46 + and not current_platform.is_device_capability_family(120)
47 + )
48 + op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk
49 +
50 + def run():
51 + op(logits, lengths, out, workspace, k, capacity)
52 +
53 + run()
54 + valid = torch.arange(k, device="cuda")[None] < lengths[:, None]
55 + assert torch.all(out[~valid] == -1)
56 + assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid)
57 + sorted_indices = out.sort(dim=1).values
58 + assert torch.all(
59 + (sorted_indices[:, 1:] != sorted_indices[:, :-1])
60 + | (sorted_indices[:, 1:] == -1)
61 + )
62 + selected = logits.gather(1, out.clamp_min(0).long()).masked_fill(
63 + ~valid, -float("inf")
64 + )
65 + reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values
66 + torch.testing.assert_close(
67 + selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0
68 + )
69 + trials = [
70 + statistics.median(
71 + bench_gpu_time_with_cupti(
72 + run,
73 + use_cuda_graph=True,
74 + cold_l2_cache=not args.warm,
75 + dry_run_iters=5,
76 + repeat_iters=args.repeat,
77 + )
78 + )
79 + * 1000
80 + for _ in range(args.trials)
81 + ]
82 + us = statistics.median(trials)
83 + return dict(
84 + rows=rows,
85 + width=width,
86 + capacity=capacity,
87 + k=k,
88 + backend="cooperative" if cooperative else "persistent",
89 + us=us,
90 + trials_us=trials,
91 + # Minimum useful traffic: one valid-score read and one index write.
92 + effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000),
93 + )
94 +
95 +
96 +def main():
97 + parser = argparse.ArgumentParser(description=__doc__)
98 + parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024])
99 + parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840])
100 + parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048])
101 + parser.add_argument("--capacity", type=int)
102 + parser.add_argument(
103 + "--full-length", action="store_true", help="Use exactly --widths valid scores"
104 + )
105 + parser.add_argument("--backend", choices=["auto", "persistent"], default="auto")
106 + parser.add_argument(
107 + "--warm", action="store_true", help="Keep L2 warm between replays"
108 + )
109 + parser.add_argument("--repeat", type=int, default=60)
110 + parser.add_argument("--trials", type=int, default=3)
111 + parser.add_argument("--seed", type=int, default=123)
112 + parser.add_argument("--output", type=Path)
113 + args = parser.parse_args()
114 + metadata = dict(
115 + gpu=str(torch.cuda.get_device_properties(0)),
116 + torch=torch.__version__,
117 + cuda=torch.version.cuda,
118 + args=vars(args),
119 + )
120 + print(json.dumps(metadata, default=str))
121 + results = []
122 + for rows, width, k in itertools.product(args.rows, args.widths, args.top_k):
123 + result = benchmark(rows, width, k, args)
124 + results.append(result)
125 + print(json.dumps(result), flush=True)
126 + if args.output:
127 + args.output.write_text(
128 + json.dumps(dict(metadata=metadata, results=results), indent=2, default=str)
129 + + "\n"
130 + )
131 +
132 +
133 +if __name__ == "__main__":
134 + main()
csrc/libtorch_stable/persistent_topk.cuh
+52 -34
@@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
1030 constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
1031 sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB
1032
1033 -/*!
1034 - * \brief Filtered Top-K kernel for ragged sequences.
1035 - *
1036 - * \tparam DType Data type (float, half, nv_bfloat16)
1037 - * \tparam IdType Index type (int32_t)
1038 - * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
1039 - */
1040 -template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
1041 - bool UsePredicatedShortLoads = false>
1042 -__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1043 - FilteredTopKUnifiedKernel(const DType* __restrict__ input,
1044 - IdType* __restrict__ output,
1045 - const IdType* __restrict__ lengths,
1046 - uint32_t num_rows, uint32_t top_k,
1047 - uint32_t max_len) {
1033 +template <uint32_t MAX_K>
1034 +struct FilteredTopKStorage {
1035 + alignas(128) int histogram[2][256 + 128];
1036 + alignas(128) int counter;
1037 + alignas(128) int threshold_bin;
1038 + alignas(128) int num_input[2];
1039 + alignas(128) int indices[MAX_K];
1040 + int last_remain;
1041 +};
1042 +
1043 +// With CheckOverflow, return false before using a truncated stash so the caller
1044 +// can retry with exact full-row selection.
1045 +template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K,
1046 + bool UsePredicatedShortLoads, bool CheckOverflow = false>
1047 +__device__ __forceinline__ bool filtered_topk_row(
1048 + const DType* score, IdType* dst, int length, uint32_t top_k,
1049 + FilteredTopKStorage<MAX_K>& storage) {
1050 constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS;
1051 constexpr int RADIX = 256;
1052 constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE;
1051 -
1052 - const uint32_t bid = blockIdx.x;
1053 const int tx = threadIdx.x;
1054
1055 - if (bid >= num_rows) return;
1056 -
1057 - const int length =
1058 - (lengths != nullptr) ? lengths[bid] : static_cast<int>(max_len);
1059 - const DType* score = input + bid * max_len;
1060 - IdType* dst = output + bid * top_k;
1061 -
1055 // Trivial case: length <= top_k
1056 if (length <= static_cast<int>(top_k)) {
1057 for (int i = tx; i < static_cast<int>(top_k); i += BLOCK_SIZE) {
1058 dst[i] = (i < length) ? static_cast<IdType>(i) : static_cast<IdType>(-1);
1059 }
1067 - return;
1060 + return true;
1061 }
1062
1063 // Short path
@@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1070 hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
1071 _smem_reg);
1072 }
1080 - return;
1073 + return true;
1074 }
1075
1083 - // Static shared memory
1084 - alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
1085 - alignas(128) __shared__ int s_counter;
1086 - alignas(128) __shared__ int s_threshold_bin_id;
1087 - alignas(128) __shared__ int s_num_input[2];
1088 - alignas(128) __shared__ int s_indices[MAX_K];
1089 -
1076 + auto& s_histogram_buf = storage.histogram;
1077 + auto& s_counter = storage.counter;
1078 + auto& s_threshold_bin_id = storage.threshold_bin;
1079 + auto& s_num_input = storage.num_input;
1080 + auto& s_indices = storage.indices;
1081 auto& s_histogram = s_histogram_buf[0];
1082
1083 // Dynamic shared memory for input double buffer
@@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1204 // Stage 2: refine with 8bit radix passes
1205 #pragma unroll
1206 for (int round = 0; round < NUM_ROUNDS; ++round) {
1216 - __shared__ int s_last_remain;
1207 + auto& s_last_remain = storage.last_remain;
1208 const auto r_idx = round % 2;
1209
1210 const auto _raw_num_input = s_num_input[r_idx];
1211 + if constexpr (CheckOverflow) {
1212 + if (_raw_num_input > SMEM_INPUT_SIZE) return false;
1213 + }
1214 const auto num_input =
1215 (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE;
1216
@@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1278 const int idx = s_indices[base];
1279 dst[base] = static_cast<IdType>(idx);
1280 }
1281 + return true;
1282 +}
1283 +
1284 +/*!
1285 + * \brief Filtered Top-K kernel for ragged sequences.
1286 + *
1287 + * \tparam DType Data type (float, half, nv_bfloat16)
1288 + * \tparam IdType Index type (int32_t)
1289 + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
1290 + */
1291 +template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
1292 + bool UsePredicatedShortLoads = false>
1293 +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1294 + FilteredTopKUnifiedKernel(const DType* __restrict__ input,
1295 + IdType* __restrict__ output,
1296 + const IdType* __restrict__ lengths,
1297 + uint32_t num_rows, uint32_t top_k,
1298 + uint32_t max_len) {
1299 + const uint32_t bid = blockIdx.x;
1300 + if (bid >= num_rows) return;
1301 + const int length = lengths ? lengths[bid] : static_cast<int>(max_len);
1302 + __shared__ FilteredTopKStorage<MAX_K> storage;
1303 + filtered_topk_row<DType, IdType, VEC_SIZE, MAX_K, UsePredicatedShortLoads>(
1304 + input + bid * max_len, output + bid * top_k, length, top_k, storage);
1305 }
1306
1307 // Helper to compute GCD for VEC_SIZE selection
csrc/libtorch_stable/sampled_topk.cuh new
+253
@@ -0,0 +1,253 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect:
5 +// https://github.com/deepseek-ai/DeepSelect
6 +// A sample estimates a cutoff, then survivors are compacted in shared memory
7 +// for exact FP32 selection. Too few or too many candidates trigger an exact
8 +// full-row fallback.
9 +
10 +#pragma once
11 +
12 +#include <cub/cub.cuh>
13 +#include <cuda_runtime.h>
14 +#include <cstdint>
15 +
16 +#include "persistent_topk.cuh"
17 +
18 +namespace vllm::sampled_topk {
19 +
20 +constexpr int kThreads = 1024;
21 +// Conservative crossover bounds from the B300 batch/length sweep.
22 +template <int K>
23 +constexpr int kMinSampledLength = K == 512 ? 98304 : 65536;
24 +// Target half the buffer to leave room for sampling error.
25 +constexpr int kSample = 4096;
26 +constexpr int kCapacity = 8192;
27 +
28 +struct Storage {
29 + uint2 candidates[kCapacity];
30 + int histogram[2048];
31 + typename cub::BlockScan<int, kThreads>::TempStorage scan;
32 + int count;
33 + int remaining;
34 + int bin;
35 + uint32_t prefix;
36 +};
37 +
38 +__device__ __forceinline__ uint32_t ordered(float value) {
39 + const uint32_t bits = __float_as_uint(value);
40 + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
41 +}
42 +
43 +__device__ __forceinline__ uint32_t coarse(uint32_t key) {
44 + const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key;
45 + return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits));
46 +}
47 +
48 +// Select the bin containing the rank and update the rank within that bin.
49 +template <int Bins>
50 +__device__ void select_bin(Storage& s) {
51 + constexpr int kItems = (Bins + kThreads - 1) / kThreads;
52 + int sums[kItems];
53 +#pragma unroll
54 + for (int j = 0; j < kItems; ++j) {
55 + const int bin = Bins - 1 - (threadIdx.x * kItems + j);
56 + sums[j] = bin >= 0 ? s.histogram[bin] : 0;
57 + }
58 + cub::BlockScan<int, kThreads>(s.scan).InclusiveSum(sums, sums);
59 + const int remaining = s.remaining;
60 + __syncthreads();
61 +#pragma unroll
62 + for (int j = 0; j < kItems; ++j) {
63 + const int bin = Bins - 1 - (threadIdx.x * kItems + j);
64 + if (bin >= 0 && sums[j] >= remaining &&
65 + sums[j] - s.histogram[bin] < remaining) {
66 + s.bin = bin;
67 + s.remaining = remaining - (sums[j] - s.histogram[bin]);
68 + }
69 + }
70 + __syncthreads();
71 +}
72 +
73 +// Find a coarse bin, then refine its full FP32 keys one byte at a time.
74 +// Read shared candidates normally, or the full row for exact fallback.
75 +template <bool Buffered>
76 +__device__ uint32_t threshold(Storage& s, const float* row, int count,
77 + int rank) {
78 + const int tid = threadIdx.x;
79 + if (tid == 0) {
80 + s.remaining = rank;
81 + s.prefix = 0;
82 + }
83 + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
84 + __syncthreads();
85 + for (int i = tid; i < count; i += kThreads) {
86 + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
87 + atomicAdd(&s.histogram[coarse(key)], 1);
88 + }
89 + __syncthreads();
90 + select_bin<2048>(s);
91 + const int coarse_bin = s.bin;
92 + uint32_t mask = 0;
93 + for (int shift = 24; shift >= 0; shift -= 8) {
94 + if (tid < 256) s.histogram[tid] = 0;
95 + __syncthreads();
96 + const uint32_t prefix = s.prefix;
97 + for (int i = tid; i < count; i += kThreads) {
98 + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
99 + if (coarse(key) == coarse_bin && (key & mask) == prefix) {
100 + atomicAdd(&s.histogram[(key >> shift) & 255], 1);
101 + }
102 + }
103 + __syncthreads();
104 + select_bin<256>(s);
105 + if (tid == 0) s.prefix = prefix | (static_cast<uint32_t>(s.bin) << shift);
106 + mask |= 255u << shift;
107 + __syncthreads();
108 + }
109 + return s.prefix;
110 +}
111 +
112 +__device__ __forceinline__ int reserve(bool hit, int* counter) {
113 + const uint32_t mask = __ballot_sync(0xffffffffu, hit);
114 + if (mask == 0) return 0;
115 + const int lane = threadIdx.x & 31;
116 + int base = 0;
117 + if (lane == 0) base = atomicAdd(counter, __popc(mask));
118 + base = __shfl_sync(0xffffffffu, base, 0);
119 + return base + __popc(mask & ((1u << lane) - 1));
120 +}
121 +
122 +template <bool Buffered>
123 +__device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
124 + uint32_t cutoff) {
125 + const int tid = threadIdx.x;
126 + if (tid == 0) s.count = 0;
127 + __syncthreads();
128 + for (int base = 0; base < count; base += kThreads) {
129 + const int i = base + tid;
130 + const uint2 pair = i < count ? (Buffered ? s.candidates[i]
131 + : make_uint2(i, ordered(row[i])))
132 + : make_uint2(0, 0);
133 + bool keep = i < count && pair.y > cutoff;
134 + if (i < count && pair.y == cutoff) {
135 + keep = atomicSub(&s.remaining, 1) > 0;
136 + }
137 + const int offset = reserve(keep, &s.count);
138 + if (keep) dst[offset] = pair.x;
139 + }
140 +}
141 +
142 +template <int K>
143 +__global__ void __launch_bounds__(kThreads)
144 + sampled_topk_kernel(const float* __restrict__ input,
145 + const int32_t* __restrict__ lengths,
146 + int32_t* __restrict__ output, int64_t stride,
147 + int max_length) {
148 + extern __shared__ __align__(16) unsigned char smem[];
149 + auto& s = *reinterpret_cast<Storage*>(smem);
150 + const int tid = threadIdx.x;
151 + const int length = max(0, min(lengths[blockIdx.x], max_length));
152 + const float* row = input + blockIdx.x * stride;
153 + int32_t* dst = output + blockIdx.x * K;
154 + // Graph allocations can hold rows much shorter than the host length bound.
155 + if (length < kMinSampledLength<K>) {
156 + __shared__ filtered_topk::FilteredTopKStorage<K> fallback_storage;
157 + bool selected;
158 + if ((stride & 3) == 0 && (reinterpret_cast<uintptr_t>(row) & 15) == 0) {
159 + selected =
160 + filtered_topk::filtered_topk_row<float, int32_t, 4, K, false, true>(
161 + row, dst, length, K, fallback_storage);
162 + } else {
163 + selected =
164 + filtered_topk::filtered_topk_row<float, int32_t, 1, K, true, true>(
165 + row, dst, length, K, fallback_storage);
166 + }
167 + if (!selected) {
168 + const uint32_t exact_cutoff = threshold<false>(s, row, length, K);
169 + emit<false>(s, row, dst, length, exact_cutoff);
170 + }
171 + return;
172 + }
173 + if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length);
174 + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
175 + __syncthreads();
176 + // Sample contiguous warps spread across the row to keep loads coalesced.
177 + for (int i = tid; i < kSample; i += kThreads) {
178 + const int index = (i / 32) * (length / (kSample / 32)) + i % 32;
179 + atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1);
180 + }
181 + __syncthreads();
182 + select_bin<2048>(s);
183 + const uint16_t half_key = s.bin << 5;
184 + const uint16_t half_bits =
185 + (half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key;
186 + const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits)));
187 + if (tid == 0) s.count = 0;
188 + __syncthreads();
189 + constexpr int kItems = 16;
190 + // Compact survivors with one reservation per warp per tile. Keep counting
191 + // beyond capacity so overflow triggers fallback instead of using a subset.
192 + for (int base = 0; base < length; base += kItems * kThreads) {
193 + uint32_t keys[kItems];
194 + uint32_t hit_mask = 0;
195 + if ((reinterpret_cast<uintptr_t>(row) & 15) == 0 &&
196 + base + kItems * kThreads <= length) {
197 +#pragma unroll
198 + for (int j = 0; j < kItems / 4; ++j) {
199 + const float4 values =
200 + reinterpret_cast<const float4*>(row + base)[tid + j * kThreads];
201 + keys[4 * j] = ordered(values.x);
202 + keys[4 * j + 1] = ordered(values.y);
203 + keys[4 * j + 2] = ordered(values.z);
204 + keys[4 * j + 3] = ordered(values.w);
205 + }
206 + } else {
207 +#pragma unroll
208 + for (int j = 0; j < kItems; ++j) {
209 + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
210 + keys[j] = i < length ? ordered(row[i]) : 0;
211 + }
212 + }
213 +#pragma unroll
214 + for (int j = 0; j < kItems; ++j) {
215 + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
216 + if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j;
217 + }
218 + const int lane = tid & 31;
219 + const int hits = __popc(hit_mask);
220 + int inclusive = hits;
221 +#pragma unroll
222 + for (int delta = 1; delta < 32; delta *= 2) {
223 + const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta);
224 + if (lane >= delta) inclusive += previous;
225 + }
226 + const int total = __shfl_sync(0xffffffffu, inclusive, 31);
227 + int warp_base = 0;
228 + if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total);
229 + warp_base = __shfl_sync(0xffffffffu, warp_base, 31);
230 + int offset = warp_base + inclusive - hits;
231 +#pragma unroll
232 + for (int j = 0; j < kItems; ++j) {
233 + if (hit_mask & (1u << j)) {
234 + if (offset < kCapacity) {
235 + s.candidates[offset] = make_uint2(
236 + base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]);
237 + }
238 + ++offset;
239 + }
240 + }
241 + }
242 + __syncthreads();
243 + const int count = s.count;
244 + if (count >= K && count <= kCapacity) {
245 + const uint32_t exact_cutoff = threshold<true>(s, row, count, K);
246 + emit<true>(s, row, dst, count, exact_cutoff);
247 + } else {
248 + const uint32_t exact_cutoff = threshold<false>(s, row, length, K);
249 + emit<false>(s, row, dst, length, exact_cutoff);
250 + }
251 +}
252 +
253 +} // namespace vllm::sampled_topk
csrc/libtorch_stable/topk.cu
+17 -1
@@ -9,6 +9,7 @@
9
10 #ifndef USE_ROCM
11 #include "persistent_topk.cuh"
12 + #include "sampled_topk.cuh"
13 #endif
14
15 namespace {
@@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
37 max_smem_per_block = device_prop->sharedMemPerBlockOptin;
38 }
39
39 - if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
40 + // Allow static fallback storage in addition to the 128 KiB dynamic buffer.
41 + if (num_rows > 64 &&
42 + max_seq_len >= vllm::sampled_topk::kMinSampledLength<TopK> &&
43 + max_smem_per_block >= 144 * 1024) {
44 + auto kernel = vllm::sampled_topk::sampled_topk_kernel<TopK>;
45 + constexpr size_t smem_size =
46 + vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC;
47 + cudaError_t status = cudaFuncSetAttribute(
48 + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
49 + STD_TORCH_CHECK(status == cudaSuccess,
50 + "sampled_topk smem failed: ", cudaGetErrorString(status));
51 + kernel<<<num_rows, vllm::sampled_topk::kThreads, smem_size, stream>>>(
52 + logits.const_data_ptr<float>(), lengths.const_data_ptr<int32_t>(),
53 + output.mutable_data_ptr<int32_t>(), stride,
54 + static_cast<int>(std::min(max_seq_len, logits.size(1))));
55 + } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
56 cudaError_t status =
57 vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
58 logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
tests/kernels/test_top_k_per_row.py
+100
@@ -932,6 +932,106 @@ def test_deepseek_workspace_topk(
932 )
933
934
935 +@pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+")
936 +@pytest.mark.parametrize("rows", [65, 128])
937 +@pytest.mark.parametrize("top_k", [512, 1024, 2048])
938 +@pytest.mark.parametrize(
939 + "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"]
940 +)
941 +@torch.inference_mode()
942 +def test_persistent_topk_sampled_graph(
943 + rows: int, top_k: int, distribution: str
944 +) -> None:
945 + """Sampling and both fallbacks preserve exact values across graph replays."""
946 + if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024:
947 + pytest.skip("Sampled top-k requires at least 144 KiB of shared memory")
948 + set_random_seed(42)
949 + min_sampled_length = 98304 if top_k == 512 else 65536
950 + width = min_sampled_length + 3
951 + lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda")
952 + logits = create_random_logits(
953 + torch.zeros_like(lengths),
954 + lengths,
955 + torch.float32,
956 + 42,
957 + False,
958 + "10LSBits" if distribution == "10LSBits" else "random",
959 + )
960 + if distribution == "ascending":
961 + logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32))
962 + elif distribution == "constant":
963 + logits.fill_(1.0)
964 + elif distribution == "sampled_peaks":
965 + # A biased sample leaves fewer than k survivors and must fall back.
966 + logits.zero_()
967 + sample = torch.arange(4096, device="cuda")
968 + positions = (sample // 32) * (width // 128) + sample % 32
969 + logits[:, positions] = 1 + sample.float() / 4096
970 + original = logits.clone()
971 + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda")
972 + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda")
973 +
974 + def run() -> None:
975 + torch.ops._C.persistent_topk(
976 + logits,
977 + lengths.view(-1, 4 if rows == 128 else 1),
978 + indices,
979 + workspace,
980 + top_k,
981 + width,
982 + )
983 +
984 + run()
985 + graph = torch.cuda.CUDAGraph()
986 + with torch.cuda.graph(graph):
987 + run()
988 + positions = torch.arange(width, device="cuda")
989 + slots = torch.arange(top_k, device="cuda")
990 + for step in range(3):
991 + bounds = torch.tensor(
992 + [
993 + -1,
994 + 0,
995 + 1,
996 + top_k - 1,
997 + top_k,
998 + 32769,
999 + min_sampled_length - 1,
1000 + min_sampled_length,
1001 + width,
1002 + width + 17,
1003 + ],
1004 + dtype=torch.int32,
1005 + device="cuda",
1006 + )
1007 + lengths.copy_(
1008 + bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()]
1009 + )
1010 + logits.copy_(original if step != 1 else original.flip(1))
1011 + logits.masked_fill_(positions[None] >= lengths[:, None], float("nan"))
1012 + indices.fill_(-2)
1013 + graph.replay()
1014 + valid = slots[None] < lengths.clamp(max=top_k)[:, None]
1015 + assert torch.all(indices[~valid] == -1)
1016 + assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid)
1017 + ordered_indices = indices.sort(dim=1).values
1018 + assert torch.all(
1019 + (ordered_indices[:, 1:] != ordered_indices[:, :-1])
1020 + | (ordered_indices[:, 1:] == -1)
1021 + )
1022 + selected = logits.gather(1, indices.clamp_min(0).long())
1023 + selected.masked_fill_(~valid, -float("inf"))
1024 + expected = logits.masked_fill(
1025 + positions[None] >= lengths[:, None], -float("inf")
1026 + )
1027 + torch.testing.assert_close(
1028 + selected.sort(dim=1, descending=True).values,
1029 + expected.topk(top_k, dim=1).values,
1030 + atol=0,
1031 + rtol=0,
1032 + )
1033 +
1034 +
1035 def run_large_context_topk_test(
1036 batch_size: int,
1037 seq_lens: list[int],