@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 184076c3fe
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 21:15 UTC

[DeepSeek v3.2] Make top-k work for any logit values. (#27568)

Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Daniel Cámpora committed Dec 8, 2025 at 15:55 UTC 184076c3fecf8c322648f36f69a4de836d7f519b
5 files changed +629 -210
csrc/ops.h
+8 -5
@@ -102,13 +102,16 @@ void apply_repetition_penalties_(torch::Tensor& logits,
102 const torch::Tensor& output_mask,
103 const torch::Tensor& repetition_penalties);
104
105 -void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
106 - const torch::Tensor& rowEnds, torch::Tensor& indices,
107 - int64_t numRows, int64_t stride0, int64_t stride1);
105 +void top_k_per_row_prefill(const torch::Tensor& logits,
106 + const torch::Tensor& rowStarts,
107 + const torch::Tensor& rowEnds, torch::Tensor& indices,
108 + int64_t numRows, int64_t stride0, int64_t stride1,
109 + int64_t topK);
110
111 void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
110 - const torch::Tensor& seq_lens, torch::Tensor& indices,
111 - int64_t numRows, int64_t stride0, int64_t stride1);
112 + const torch::Tensor& seqLens, torch::Tensor& indices,
113 + int64_t numRows, int64_t stride0, int64_t stride1,
114 + int64_t topK);
115
116 void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
117 torch::Tensor& weight, torch::Tensor& scale,
csrc/sampler.cu
+536 -179
@@ -44,41 +44,300 @@ __global__ void apply_repetition_penalties_kernel(
44 }
45 }
46
47 -static inline __device__ uint16_t extractBinIdx(float x) {
48 - union {
49 - __half h;
50 - uint16_t u16;
51 - } tmp;
52 - tmp.h = __float2half_rn(x);
53 - tmp.u16 = (x < 0.f) ? (~tmp.u16 & 0xffff) : (tmp.u16 | 0x8000);
54 - return 511 - (tmp.u16 >> 7);
47 +__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
48 + uint32_t bits = __float_as_uint(x);
49 + return (bits & 0x80000000) ? bits : ~bits & 0x7fffffff;
50 }
51
57 -template <int kNumThreadsPerBlock = 512, int kNumBins = 512, int kTopK = 2048>
58 -__device__ void topKPerRowJob(const float* logits, const int rowStart,
59 - const int rowEnd, const int rowIdx,
60 - int* outIndices, int stride0, int stride1) {
61 - // The number of elements per thread for the final top-k sort.
62 - static constexpr int kNumTopKItemsPerThread = kTopK / kNumThreadsPerBlock;
63 - // The class to sort the elements during the final top-k sort.
64 - using TopKSort = cub::BlockRadixSort<float, kNumThreadsPerBlock,
65 - kNumTopKItemsPerThread, int>;
52 +template <int step>
53 +static inline __device__ uint32_t extractBinIdx(float x) {
54 + if constexpr (step == 0) {
55 + __half hx = __float2half(x);
56 + uint16_t bits = __half_as_ushort(hx);
57 + bits = (bits & 0x8000) ? bits : ~bits & 0x7fff;
58 + return bits >> 5;
59 + } else {
60 + uint32_t bits = __float_as_uint(x);
61 + bits = (bits & 0x80000000) ? bits : ~bits & 0x7fffffff;
62 +
63 + if constexpr (step == 1) {
64 + return bits >> 21;
65 + } else if constexpr (step == 2) {
66 + return (bits >> 10) & 0x7ff;
67 + } else if constexpr (step == 3) {
68 + return bits & 0x3ff;
69 + }
70 + }
71 +}
72 +
73 +template <int shift>
74 +static inline __device__ bool isPartialMatch(float x, uint32_t pattern) {
75 + if constexpr (shift == 0) {
76 + return true;
77 + }
78 + uint32_t bits = __float_as_uint(x);
79 + bits = (bits & 0x80000000) ? bits : ~bits & 0x7fffffff;
80 + return (bits ^ pattern) >> shift == 0;
81 +}
82 +
83 +/**
84 + * Map a Func over the input data, using vectorized load instructions if
85 + * possible.
86 + *
87 + * @tparam T element type
88 + * @tparam IdxT indexing type
89 + * @tparam Func void (T x, IdxT idx)
90 + *
91 + * @param thread_rank rank of the calling thread among all participating threads
92 + * @param num_threads number of the threads that participate in processing
93 + * @param in the input data
94 + * @param len the number of elements to read
95 + * @param f the lambda taking two arguments (T x, IdxT idx)
96 + */
97 +template <typename T, typename idxT, typename Func>
98 +__device__ void vectorized_process(size_t thread_rank, size_t num_threads,
99 + const T* in, idxT len, Func f) {
100 + constexpr int WARP_SIZE = 32;
101 + using WideT = float4;
102 + if constexpr (sizeof(T) >= sizeof(WideT)) {
103 + for (idxT i = thread_rank; i < len; i += num_threads) {
104 + f(in[i], i);
105 + }
106 + } else {
107 + static_assert(sizeof(WideT) % sizeof(T) == 0);
108 + constexpr int items_per_scalar = sizeof(WideT) / sizeof(T);
109 + // TODO: it's UB
110 + union {
111 + WideT scalar;
112 + T array[items_per_scalar];
113 + } wide;
114 +
115 + int skip_cnt =
116 + (reinterpret_cast<size_t>(in) % sizeof(WideT))
117 + ? ((sizeof(WideT) - reinterpret_cast<size_t>(in) % sizeof(WideT)) /
118 + sizeof(T))
119 + : 0;
120 + if (skip_cnt > len) {
121 + skip_cnt = len;
122 + }
123 + const WideT* in_cast = reinterpret_cast<decltype(in_cast)>(in + skip_cnt);
124 + const idxT len_cast = (len - skip_cnt) / items_per_scalar;
125 +
126 + for (idxT i = thread_rank; i < len_cast; i += num_threads) {
127 + wide.scalar = in_cast[i];
128 + const idxT real_i = skip_cnt + i * items_per_scalar;
129 +#pragma unroll
130 + for (int j = 0; j < items_per_scalar; ++j) {
131 + f(wide.array[j], real_i + j);
132 + }
133 + }
134 +
135 + static_assert(WARP_SIZE >= items_per_scalar);
136 + // and because items_per_scalar > skip_cnt, WARP_SIZE > skip_cnt
137 + // no need to use loop
138 + if (thread_rank < skip_cnt) {
139 + f(in[thread_rank], thread_rank);
140 + }
141 + // because len_cast = (len - skip_cnt) / items_per_scalar,
142 + // len_cast * items_per_scalar + items_per_scalar > len - skip_cnt;
143 + // and so
144 + // len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <=
145 + // WARP_SIZE no need to use loop
146 + const idxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank;
147 + if (remain_i < len) {
148 + f(in[remain_i], remain_i);
149 + }
150 + }
151 +}
152 +
153 +template <int step, int kNumThreadsPerBlock, int kNumBins, int kNumFinalItems,
154 + bool multipleBlocksPerRow, bool mergeBlocks, typename SmemFinalType,
155 + typename SmemOutputType>
156 +__device__ bool processHistogramStep(
157 + const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern,
158 + int& thresholdBinIdx, SmemOutputType& smemOutput, int* smemThresholdBinIdx,
159 + int* smemFinalDstIdx, int* smemFinalBinSize, int* smemFoundTopKValues,
160 + SmemFinalType& smemFinal, int stride1, int rowStart, int topK) {
161 + // Clear the histogram.
162 +#pragma unroll
163 + for (int idx = threadIdx.x; idx < kNumBins; idx += kNumThreadsPerBlock) {
164 + smemFinal.histo.data[idx] = 0;
165 + }
166 +
167 + // Make sure the histogram is ready.
168 + __syncthreads();
169 +
170 + // Update pattern
171 + constexpr auto patternShift = step < 2 ? 0 : step == 2 ? 21 : 10;
172 + if constexpr (step == 2) {
173 + logitPattern = static_cast<uint32_t>(thresholdBinIdx & 0x7ff)
174 + << patternShift;
175 + } else if constexpr (step == 3) {
176 + logitPattern |= static_cast<uint32_t>(thresholdBinIdx & 0x7ff)
177 + << patternShift;
178 + }
179 +
180 + auto distributeToBins = [&](float logit, int /* idx */ = 0) {
181 + if (isPartialMatch<patternShift>(logit, logitPattern)) {
182 + uint32_t binIdx = extractBinIdx<step>(logit);
183 + atomicAdd(&smemFinal.histo.data[binIdx], 1);
184 + }
185 + };
186 +
187 + // Distribute the elements to the histogram bins.
188 + if (stride1 == 1) {
189 + vectorized_process(threadIdx.x, kNumThreadsPerBlock, logits + rowStart,
190 + rowEnd - rowStart, distributeToBins);
191 + } else {
192 + for (int idx = rowStart + threadIdx.x; idx < rowEnd;
193 + idx += kNumThreadsPerBlock) {
194 + float logit = logits[idx * stride1];
195 + distributeToBins(logit, idx);
196 + }
197 + }
198 + // Make sure the histogram is ready.
199 + __syncthreads();
200 +
201 + // Reads the value of the starting position in the smemOutput array
202 + int lastValue = smemFoundTopKValues[0];
203 +
204 + for (int round = 0; round < kNumBins / kNumThreadsPerBlock; round++) {
205 + // Read the values from SMEM.
206 + int idx = threadIdx.x + kNumThreadsPerBlock * round;
207 + int binCount{0};
208 + binCount = smemFinal.histo.data[idx];
209 +
210 + // Make sure each thread has read its value.
211 + __syncthreads();
212 +
213 + // Compute the prefix sum.
214 + int prefixSum{0}, totalSum{0};
215 + using Scan = cub::BlockScan<int, kNumThreadsPerBlock>;
216 + Scan(smemFinal.histo.scan).ExclusiveSum(binCount, prefixSum, totalSum);
217 +
218 + // Update the histogram with the prefix sums.
219 + prefixSum += lastValue;
220 + totalSum += lastValue;
221 + smemFinal.histo.data[idx] = prefixSum;
222 +
223 + // Make sure the data is in shared memory.
224 + __syncthreads();
225 +
226 + // Find the last valid bin.
227 + bool foundThreshold = false;
228 + if (prefixSum < topK) {
229 + int nextPrefixSum = threadIdx.x == kNumThreadsPerBlock - 1
230 + ? totalSum
231 + : smemFinal.histo.data[idx + 1];
232 +
233 + if (nextPrefixSum >= topK) {
234 + smemThresholdBinIdx[0] = idx;
235 + smemFinalBinSize[0] = nextPrefixSum - prefixSum;
236 + foundThreshold = true;
237 + }
238 + }
239 +
240 + // Early exit: if any thread found the threshold, we can skip remaining
241 + // rounds
242 + if (__syncthreads_or(foundThreshold)) {
243 + break;
244 + }
245 +
246 + lastValue = totalSum;
247 + }
248 +
249 + // Make sure the data is in shared memory.
250 + __syncthreads();
251 +
252 + // The threshold bin.
253 + thresholdBinIdx = smemThresholdBinIdx[0];
254 +
255 + auto processBins = [&](float logit, int idx) {
256 + if (isPartialMatch<patternShift>(logit, logitPattern)) {
257 + uint32_t binIdx = extractBinIdx<step>(logit);
258 + if (binIdx < thresholdBinIdx) {
259 + // The element is part of the top-k selection
260 + int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
261 +
262 + if constexpr (mergeBlocks) {
263 + smemOutput[dstIdx] = indices[idx];
264 + } else if constexpr (multipleBlocksPerRow) {
265 + smemOutput[dstIdx] = idx + rowStart;
266 + reinterpret_cast<float*>(smemOutput + topK)[dstIdx] = logit;
267 + } else {
268 + smemOutput[dstIdx] = idx;
269 + }
270 + }
271 + if constexpr (step < 3) {
272 + // Only fill the final items for sorting if the threshold bin fits
273 + if (binIdx == thresholdBinIdx &&
274 + smemFinalBinSize[0] <= kNumFinalItems) {
275 + int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
276 + smemFinal.items.logits[dstIdx] = logit;
277 + if constexpr (mergeBlocks) {
278 + smemFinal.items.indices[dstIdx] = indices[idx];
279 + } else if constexpr (multipleBlocksPerRow) {
280 + smemFinal.items.indices[dstIdx] = idx + rowStart;
281 + } else {
282 + smemFinal.items.indices[dstIdx] = idx;
283 + }
284 + }
285 + } else {
286 + if (binIdx == thresholdBinIdx) {
287 + // The elements in the threshold bin share the same 32 bits at step 3
288 + int dstIdx = atomicAdd(&smemFinal.histo.data[binIdx], 1);
289 + if (dstIdx < topK) {
290 + if constexpr (mergeBlocks) {
291 + smemOutput[dstIdx] = indices[idx];
292 + } else if constexpr (multipleBlocksPerRow) {
293 + smemOutput[dstIdx] = idx + rowStart;
294 + reinterpret_cast<float*>(smemOutput + topK)[dstIdx] = logit;
295 + } else {
296 + smemOutput[dstIdx] = idx;
297 + }
298 + }
299 + }
300 + }
301 + }
302 + };
303 +
304 + if (stride1 == 1) {
305 + vectorized_process(threadIdx.x, kNumThreadsPerBlock, logits + rowStart,
306 + rowEnd - rowStart, processBins);
307 + } else {
308 + for (int idx = rowStart + threadIdx.x; idx < rowEnd;
309 + idx += kNumThreadsPerBlock) {
310 + float logit = logits[idx * stride1];
311 + processBins(logit, idx);
312 + }
313 + }
314 +
315 + // Make sure the elements are in shared memory.
316 + __syncthreads();
317 +
318 + // Check if we should continue to next step
319 + return smemFinalBinSize[0] > kNumFinalItems;
320 +}
321
322 +// Follows half - 11 - 11 - 10 bit iterations
323 +template <int kNumThreadsPerBlock, int kNumBins, bool useRadixSort,
324 + bool multipleBlocksPerRow = false, bool mergeBlocks = false>
325 +static __device__ void topKPerRowJob(const int* indices, const float* logits,
326 + int rowStart, int rowEnd, int* outIndices,
327 + float* outLogits, int stride1, int topK) {
328 // The number of slots for the final pass.
68 - static constexpr int kNumFinalItems = 3072;
329 + static constexpr int kNumFinalItems = 2048;
330 // The number of elements per thread for the final sort.
331 static constexpr int kNumFinalItemsPerThread =
332 kNumFinalItems / kNumThreadsPerBlock;
333 // The class to sort the elements during the final pass.
334 using FinalSort = cub::BlockRadixSort<float, kNumThreadsPerBlock,
335 kNumFinalItemsPerThread, int>;
75 -
336 + using FinalSortTempStorage =
337 + std::conditional_t<useRadixSort, typename FinalSort::TempStorage, int>;
338 // The class to compute the inclusive prefix-sum over the histogram.
339 using Scan = cub::BlockScan<int, kNumThreadsPerBlock>;
340
79 - // Shared memory to compute the block scan.
80 - __shared__ typename Scan::TempStorage smemScan;
81 -
341 // The structure to store the final items (for the final pass).
342 struct FinalItems {
343 // Shared memory to store the indices for the final pass.
@@ -87,200 +346,225 @@ __device__ void topKPerRowJob(const float* logits, const int rowStart,
346 float logits[kNumFinalItems];
347 };
348
349 + struct Histogram {
350 + typename Scan::TempStorage scan;
351 + int data[kNumBins];
352 + };
353 +
354 // Shared memory to compute the block sort.
355 __shared__ union {
356 FinalItems items;
93 - typename FinalSort::TempStorage finalSort;
94 - typename TopKSort::TempStorage topKSort;
357 + FinalSortTempStorage finalSort;
358 + Histogram histo;
359 } smemFinal;
360
97 - // Shared memory to store the histogram.
98 - __shared__ int smemHistogram[kNumBins];
361 // Shared memory to store the selected indices.
100 - __shared__ int smemIndices[kTopK];
362 + // If we are processing using multiple blocks, we need to store the logits and
363 + // indices.
364 + extern __shared__ int32_t smemOutput[];
365 +
366 // Shared memory to store the threshold bin.
367 __shared__ int smemThresholdBinIdx[1];
368 // Shared memory counter to register the candidates for the final phase.
369 __shared__ int smemFinalDstIdx[1];
370 + // Shared memory to determine if the threshold bin fits in the final items.
371 + __shared__ int smemFinalBinSize[1];
372 + // Shared memory to keep track of the top-k values found so far by the
373 + // previous iterations
374 + __shared__ int smemFoundTopKValues[1];
375
376 // The length of the row.
377 int rowLen = rowEnd - rowStart;
378
379 // Shortcut if the length of the row is smaller than Top-K. Indices are not
380 // sorted by their corresponding logit.
111 - if (rowLen <= kTopK) {
381 + if (rowLen <= topK) {
382 for (int rowIt = threadIdx.x; rowIt < rowLen;
383 rowIt += kNumThreadsPerBlock) {
114 - int idx = rowStart + rowIt;
115 - outIndices[rowIdx * kTopK + rowIt] = idx - rowStart;
384 + if constexpr (multipleBlocksPerRow) {
385 + outIndices[rowIt] = rowIt + rowStart;
386 + outLogits[rowIt] = logits[rowIt + rowStart];
387 + } else {
388 + outIndices[rowIt] = rowIt;
389 + }
390 }
117 - for (int rowIt = rowLen + threadIdx.x; rowIt < kTopK;
391 + for (int rowIt = rowLen + threadIdx.x; rowIt < topK;
392 rowIt += kNumThreadsPerBlock) {
119 - outIndices[rowIdx * kTopK + rowIt] = -1;
393 + outIndices[rowIt] = -1;
394 + if constexpr (multipleBlocksPerRow) {
395 + outLogits[rowIt] = -FLT_MAX;
396 + }
397 }
121 - return;
122 - }
123 -
124 - // Clear the histogram.
125 - if (threadIdx.x < kNumBins) {
126 - smemHistogram[threadIdx.x] = 0;
127 - }
128 -
129 - // Make sure the histogram is ready.
130 - __syncthreads();
131 -
132 - // Fetch elements one-by-one.
133 - for (int rowIt = rowStart + threadIdx.x; rowIt < rowEnd;
134 - rowIt += kNumThreadsPerBlock) {
135 - uint16_t idx = extractBinIdx(logits[rowIdx * stride0 + rowIt * stride1]);
136 - atomicAdd(&smemHistogram[idx], 1);
137 - }
138 -
139 - // Make sure the histogram is ready.
140 - __syncthreads();
141 -
142 - // Read the values from SMEM.
143 - int binCount{0};
144 - if (threadIdx.x < kNumBins) {
145 - binCount = smemHistogram[threadIdx.x];
146 - }
147 -
148 - // Make sure each thread has read its value.
149 - __syncthreads();
150 -
151 - // Compute the prefix sum.
152 - int prefixSum{0}, totalSum{0};
153 - Scan(smemScan).ExclusiveSum(binCount, prefixSum, totalSum);
154 -
155 - // Update the histogram with the prefix sums.
156 - if (threadIdx.x < kNumBins) {
157 - smemHistogram[threadIdx.x] = prefixSum;
158 - }
398
160 - // Make sure the data is in shared memory.
161 - __syncthreads();
162 -
163 - // Find the last valid bin.
164 - if (threadIdx.x < kNumBins) {
165 - int nextPrefixSum =
166 - threadIdx.x == kNumBins - 1 ? totalSum : smemHistogram[threadIdx.x + 1];
167 - if (prefixSum < kTopK && nextPrefixSum >= kTopK) {
168 - smemThresholdBinIdx[0] = threadIdx.x;
169 - }
399 + return;
400 }
171 -
172 - // Clear the counter to store the items for the final phase.
401 + // Initialize values
402 if (threadIdx.x == 0) {
403 smemFinalDstIdx[0] = 0;
404 + smemFoundTopKValues[0] = 0;
405 }
176 -
177 - // Make sure the data is in shared memory.
406 __syncthreads();
407 + int thresholdBinIdx = -1;
408 + uint32_t logitPattern = 0;
409 +
410 + // Step 0: Process first 11 bits of half representation
411 + bool continueToNextStep =
412 + processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
413 + multipleBlocksPerRow, mergeBlocks>(
414 + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
415 + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
416 + smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
417 +
418 + if (continueToNextStep) {
419 + // Step 1: Process next 11 bits
420 + continueToNextStep =
421 + processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
422 + multipleBlocksPerRow, mergeBlocks>(
423 + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
424 + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
425 + smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
426 + }
427
180 - // The threshold bin.
181 - int thresholdBinIdx = smemThresholdBinIdx[0];
182 -
183 - // Fetch elements one-by-one and populate the shared memory buffers.
184 - for (int rowIt = rowStart + threadIdx.x; rowIt < rowEnd;
185 - rowIt += kNumThreadsPerBlock) {
186 - float logit = logits[rowIdx * stride0 + rowIt * stride1];
187 - uint16_t idx = extractBinIdx(logit);
188 - if (idx < thresholdBinIdx) {
189 - int dstIdx = atomicAdd(&smemHistogram[idx], 1);
190 - smemIndices[dstIdx] = rowIt;
191 - } else if (idx == thresholdBinIdx) {
192 - int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
193 - if (dstIdx < kNumFinalItems) {
194 - smemFinal.items.logits[dstIdx] = logit;
195 - smemFinal.items.indices[dstIdx] = rowIt;
196 - }
197 - }
428 + if (continueToNextStep) {
429 + // Step 2: Process next 11 bits
430 + continueToNextStep =
431 + processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
432 + multipleBlocksPerRow, mergeBlocks>(
433 + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
434 + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
435 + smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
436 }
437
200 - // Make sure the elements are in shared memory.
201 - __syncthreads();
438 + if (continueToNextStep) {
439 + // Step 3: Process last 10 bits
440 + processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
441 + multipleBlocksPerRow, mergeBlocks>(
442 + indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
443 + smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
444 + smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
445 + }
446
203 - // The logits of the elements to be sorted in the final pass.
204 - float finalLogits[kNumFinalItemsPerThread];
205 - // The indices of the elements to be sorted in the final pass.
206 - int finalIndices[kNumFinalItemsPerThread];
447 + if (!continueToNextStep) {
448 + // The histogram did not proceed to the final 10 bits, therefore we need to
449 + // sort the final items The logits of the elements to be sorted in the final
450 + // pass.
451 + if constexpr (useRadixSort) {
452 + // Sorting with radix sort
453 + float finalLogits[kNumFinalItemsPerThread];
454 + // The indices of the elements to be sorted in the final pass.
455 + int finalIndices[kNumFinalItemsPerThread];
456
208 -// Init.
457 #pragma unroll
210 - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
211 - finalLogits[ii] = -FLT_MAX;
212 - }
458 + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
459 + finalLogits[ii] = -FLT_MAX;
460 + }
461
214 -// Read the elements from SMEM.
462 + // Read the elements from SMEM.
463 #pragma unroll
216 - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
217 - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
218 - if (srcIdx < smemFinalDstIdx[0]) {
219 - finalLogits[ii] = smemFinal.items.logits[srcIdx];
220 - finalIndices[ii] = smemFinal.items.indices[srcIdx];
221 - }
222 - }
464 + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
465 + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
466 + if (srcIdx < smemFinalDstIdx[0]) {
467 + finalLogits[ii] = smemFinal.items.logits[srcIdx];
468 + finalIndices[ii] = smemFinal.items.indices[srcIdx];
469 + }
470 + }
471 + // Make sure the shared memory has been read.
472 + __syncthreads();
473
224 - // Make sure the shared memory has been read.
225 - __syncthreads();
474 + // Sort the elements.
475 + FinalSort(smemFinal.finalSort)
476 + .SortDescendingBlockedToStriped(finalLogits, finalIndices);
477
227 - // Sort the elements.
228 - FinalSort(smemFinal.finalSort)
229 - .SortDescendingBlockedToStriped(finalLogits, finalIndices);
478 + // Copy the data back to the shared memory storage.
479 + int baseIdx = smemFoundTopKValues[0];
480
231 - // Copy the data back to the shared memory storage.
232 - int baseIdx = thresholdBinIdx > 0 ? smemHistogram[thresholdBinIdx - 1] : 0;
481 #pragma unroll
234 - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
235 - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
236 - int dstIdx = baseIdx + srcIdx;
237 - if (dstIdx < kTopK) {
238 - smemIndices[dstIdx] = finalIndices[ii];
482 + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
483 + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
484 + int dstIdx = baseIdx + srcIdx;
485 +
486 + if (dstIdx < topK) {
487 + smemOutput[dstIdx] = finalIndices[ii];
488 + if constexpr (multipleBlocksPerRow) {
489 + reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
490 + finalLogits[ii];
491 + }
492 + }
493 + }
494 + } else {
495 + // Sorting with insertion sort
496 + auto baseIdx = smemFoundTopKValues[0];
497 + for (int i = threadIdx.x; i < smemFinalDstIdx[0];
498 + i += kNumThreadsPerBlock) {
499 + int outIndex = 0;
500 + auto logit = smemFinal.items.logits[i];
501 + for (int j = 0; j < smemFinalDstIdx[0]; j++) {
502 + auto otherLogit = smemFinal.items.logits[j];
503 + if (logit < otherLogit || (logit == otherLogit && i < j)) {
504 + outIndex++;
505 + }
506 + }
507 + // Store if outIndex is in bounds
508 + if (outIndex + baseIdx < topK) {
509 + smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i];
510 + if constexpr (multipleBlocksPerRow) {
511 + reinterpret_cast<float*>(smemOutput + topK)[outIndex + baseIdx] =
512 + smemFinal.items.logits[i];
513 + }
514 + }
515 + }
516 }
517 + __syncthreads();
518 }
519
242 - // Make sure the data is in shared memory.
243 - __syncthreads();
244 -
245 -// Store to global memory.
246 -#pragma unroll
247 - for (int ii = 0; ii < kNumTopKItemsPerThread; ++ii) {
248 - int offset = rowIdx * kTopK + ii * kNumThreadsPerBlock + threadIdx.x;
249 - outIndices[offset] =
250 - smemIndices[ii * kNumThreadsPerBlock + threadIdx.x] - rowStart;
520 + // Store to global memory.
521 + for (int i = threadIdx.x; i < topK; i += kNumThreadsPerBlock) {
522 + if constexpr (multipleBlocksPerRow) {
523 + outIndices[i] = smemOutput[i];
524 + outLogits[i] = reinterpret_cast<float*>(smemOutput + topK)[i];
525 + } else {
526 + if (stride1 == 1) {
527 + // stride1 == 1 will use vectorized_process, which indexes already skip
528 + // the rowStart.
529 + outIndices[i] = smemOutput[i];
530 + } else {
531 + outIndices[i] = smemOutput[i] - rowStart;
532 + }
533 + }
534 }
535 }
536
254 -template <int kNumThreadsPerBlock = 512>
255 -static __global__ void topKPerRow(const float* logits, const int* rowStarts,
256 - const int* rowEnds, int* outIndices,
257 - int stride0, int stride1) {
537 +template <int kNumThreadsPerBlock, bool useRadixSort>
538 +static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(
539 + const float* logits, const int* rowStarts, const int* rowEnds,
540 + int* outIndices, int stride0, int stride1, const int topK,
541 + const int offsetIndex) {
542 // The number of bins in the histogram.
259 - static constexpr int kNumBins = 512;
260 -
261 - // The top-k width.
262 - static constexpr int kTopK = 2048;
543 + static constexpr int kNumBins = 2048;
544
545 // The row computed by this block.
265 - int rowIdx = blockIdx.x;
546 + int rowIdx = blockIdx.x + offsetIndex;
547
548 // The range of logits within the row.
549 int rowStart = rowStarts[rowIdx];
550 int rowEnd = rowEnds[rowIdx];
551
271 - topKPerRowJob<kNumThreadsPerBlock, kNumBins, kTopK>(
272 - logits, rowStart, rowEnd, rowIdx, outIndices, stride0, stride1);
552 + // Local pointers to this block
553 + outIndices += rowIdx * topK;
554 + logits += rowIdx * stride0;
555 +
556 + topKPerRowJob<kNumThreadsPerBlock, kNumBins, useRadixSort>(
557 + nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK);
558 }
559
275 -template <int kNumThreadsPerBlock = 512>
276 -static __global__ void topKPerRowDecode(const float* logits, const int* seqLens,
277 - int* outIndices, int stride0,
278 - int stride1, int next_n) {
560 +template <int kNumThreadsPerBlock, bool useRadixSort,
561 + bool multipleBlocksPerRow = false, bool mergeBlocks = false>
562 +static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
563 + const float* logits, const int* seqLens, int* outIndices, int stride0,
564 + int stride1, const int topK, int next_n, float* outLogits = nullptr,
565 + const int numBlocksToMerge = 0, const int* indices = nullptr) {
566 // The number of bins in the histogram.
280 - static constexpr int kNumBins = 512;
281 -
282 - // The top-k width.
283 - static constexpr int kTopK = 2048;
567 + static constexpr int kNumBins = 2048;
568
569 // The row computed by this block.
570 int rowIdx = blockIdx.x;
@@ -290,8 +574,25 @@ static __global__ void topKPerRowDecode(const float* logits, const int* seqLens,
574 int seq_len = seqLens[rowIdx / next_n];
575 int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;
576
293 - topKPerRowJob<kNumThreadsPerBlock, kNumBins, kTopK>(
294 - logits, rowStart, rowEnd, rowIdx, outIndices, stride0, stride1);
577 + // Local pointers to this block
578 + if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
579 + outIndices += rowIdx * topK;
580 + } else if constexpr (multipleBlocksPerRow) {
581 + const auto blockSize = rowEnd / gridDim.y; // 16384 / 2 = 8192
582 + rowStart = blockSize * blockIdx.y; // 8192 * 1 = 8192
583 + rowEnd = gridDim.y == blockIdx.y + 1 ? rowEnd : rowStart + blockSize;
584 + outIndices += rowIdx * gridDim.y * topK + blockIdx.y * topK;
585 + outLogits += rowIdx * gridDim.y * topK + blockIdx.y * topK;
586 + } else if constexpr (mergeBlocks) {
587 + rowEnd = numBlocksToMerge * topK;
588 + indices += rowIdx * numBlocksToMerge * topK;
589 + outIndices += rowIdx * topK;
590 + }
591 + logits += rowIdx * stride0;
592 +
593 + topKPerRowJob<kNumThreadsPerBlock, kNumBins, useRadixSort,
594 + multipleBlocksPerRow, mergeBlocks>(
595 + indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK);
596 }
597
598 } // namespace vllm
@@ -339,28 +640,84 @@ void apply_repetition_penalties_(
640
641 void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
642 const torch::Tensor& seqLens, torch::Tensor& indices,
342 - int64_t numRows, int64_t stride0, int64_t stride1) {
343 - // Compute the results on the device.
643 + int64_t numRows, int64_t stride0, int64_t stride1,
644 + int64_t topK) {
645 + constexpr int kSortingAlgorithmThreshold = 12288;
646 + constexpr int kSplitWorkThreshold = 200 * 1000;
647 constexpr int kNumThreadsPerBlock = 512;
648 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
346 -
347 - vllm::topKPerRowDecode<kNumThreadsPerBlock>
348 - <<<numRows, kNumThreadsPerBlock, 0, stream>>>(
349 - logits.data_ptr<float>(), seqLens.data_ptr<int>(),
350 - indices.data_ptr<int>(), static_cast<int>(stride0),
351 - static_cast<int>(stride1), static_cast<int>(next_n));
649 + const auto numColumns = logits.size(1);
650 +
651 + if (numColumns < kSortingAlgorithmThreshold) {
652 + // Use insertion sort
653 + vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
654 + <<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
655 + logits.data_ptr<float>(), seqLens.data_ptr<int>(),
656 + indices.data_ptr<int>(), static_cast<int>(stride0),
657 + static_cast<int>(stride1), static_cast<int>(topK),
658 + static_cast<int>(next_n));
659 + } else if (numColumns < kSplitWorkThreshold) {
660 + // From this threshold, use radix sort instead
661 + vllm::topKPerRowDecode<kNumThreadsPerBlock, true>
662 + <<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
663 + logits.data_ptr<float>(), seqLens.data_ptr<int>(),
664 + indices.data_ptr<int>(), static_cast<int>(stride0),
665 + static_cast<int>(stride1), static_cast<int>(topK),
666 + static_cast<int>(next_n));
667 + } else {
668 + // Long sequences are run in two steps
669 + constexpr auto multipleBlocksPerRowConfig = 10;
670 +
671 + const auto outIndicesAux =
672 + torch::empty({numRows, multipleBlocksPerRowConfig, topK},
673 + torch::dtype(torch::kInt32).device(logits.device()));
674 + const auto outLogitsAux =
675 + torch::empty({numRows, multipleBlocksPerRowConfig, topK},
676 + torch::dtype(torch::kFloat).device(logits.device()));
677 +
678 + vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
679 + <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
680 + 2 * topK * sizeof(int32_t), stream>>>(
681 + logits.data_ptr<float>(), seqLens.data_ptr<int>(),
682 + outIndicesAux.data_ptr<int>(), static_cast<int>(stride0),
683 + static_cast<int>(stride1), static_cast<int>(topK),
684 + static_cast<int>(next_n), outLogitsAux.data_ptr<float>());
685 +
686 + constexpr int kNumThreadsPerBlockMerge = 1024;
687 + vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
688 + <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
689 + outLogitsAux.data_ptr<float>(), seqLens.data_ptr<int>(),
690 + indices.data_ptr<int>(), multipleBlocksPerRowConfig * topK, 1,
691 + static_cast<int>(topK), static_cast<int>(next_n), nullptr,
692 + multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());
693 + }
694 }
695
354 -void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
355 - const torch::Tensor& rowEnds, torch::Tensor& indices,
356 - int64_t numRows, int64_t stride0, int64_t stride1) {
357 - // Compute the results on the device.
696 +void top_k_per_row_prefill(const torch::Tensor& logits,
697 + const torch::Tensor& rowStarts,
698 + const torch::Tensor& rowEnds, torch::Tensor& indices,
699 + int64_t numRows, int64_t stride0, int64_t stride1,
700 + int64_t topK) {
701 + constexpr int kSortingAlgorithmThreshold = 12288;
702 constexpr int kNumThreadsPerBlock = 512;
703 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
704
361 - vllm::topKPerRow<kNumThreadsPerBlock>
362 - <<<numRows, kNumThreadsPerBlock, 0, stream>>>(
363 - logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
364 - rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
365 - static_cast<int>(stride0), static_cast<int>(stride1));
705 + int numInsertionBlocks =
706 + std::min(static_cast<int>(numRows), kSortingAlgorithmThreshold);
707 + vllm::topKPerRowPrefill<kNumThreadsPerBlock, false>
708 + <<<numInsertionBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
709 + stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
710 + rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
711 + static_cast<int>(stride0), static_cast<int>(stride1),
712 + static_cast<int>(topK), 0);
713 +
714 + if (numRows > kSortingAlgorithmThreshold) {
715 + int numRadixBlocks = numRows - kSortingAlgorithmThreshold;
716 + vllm::topKPerRowPrefill<kNumThreadsPerBlock, true>
717 + <<<numRadixBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
718 + stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
719 + rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
720 + static_cast<int>(stride0), static_cast<int>(stride1),
721 + static_cast<int>(topK), kSortingAlgorithmThreshold);
722 + }
723 }
csrc/torch_bindings.cpp
+5 -5
@@ -179,15 +179,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
179
180 // Optimized top-k per row operation
181 ops.def(
182 - "top_k_per_row(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
182 + "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
183 "Tensor! indices, int numRows, int stride0, "
184 - "int stride1) -> ()");
185 - ops.impl("top_k_per_row", torch::kCUDA, &top_k_per_row);
184 + "int stride1, int topK) -> ()");
185 + ops.impl("top_k_per_row_prefill", torch::kCUDA, &top_k_per_row_prefill);
186
187 ops.def(
188 "top_k_per_row_decode(Tensor logits, int next_n, "
189 - "Tensor seq_lens, Tensor! indices, int numRows, "
190 - "int stride0, int stride1) -> ()");
189 + "Tensor seq_lens, Tensor! indices, "
190 + "int numRows, int stride0, int stride1, int topK) -> ()");
191 ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
192
193 // Layernorm-quant
tests/kernels/test_top_k_per_row.py
+77 -18
@@ -9,23 +9,45 @@ from vllm.platforms import current_platform
9
10 # Test parameters
11 NUM_ROWS = [1, 32, 2050]
12 -TOP_K_VALUES = [2048]
13 -BATCH_SIZE = [1, 2, 4, 2048, 4096]
14 -NEXT_N = [1, 2, 4, 8]
12 +TOP_K_VALUES = [2048, 3000]
13 +BATCH_SIZE = [1, 2, 2048]
14 +NEXT_N = [1, 8]
15 +DATA_GENERATION = ["random", "10LSBits"]
16
17
18 def create_random_logits(
19 row_starts: torch.Tensor,
20 row_ends: torch.Tensor,
20 - vocab_size: int,
21 dtype: torch.dtype,
22 seed: int,
23 + data_generation: str,
24 ) -> torch.Tensor:
25 """Create random logits tensor for testing."""
26 torch.manual_seed(seed)
27 np.random.seed(seed)
28 # Generate logits with some structure to make testing more meaningful
28 - logits = torch.randn(row_starts.shape[0], max(row_ends), dtype=dtype, device="cuda")
29 + if data_generation == "random":
30 + logits = torch.randn(
31 + row_starts.shape[0], max(row_ends), dtype=dtype, device="cuda"
32 + )
33 + elif data_generation == "10LSBits":
34 + top_22_bits_mask = 0xFFFFFC00
35 + last_10_bits_mask = 0x000003FF
36 + fixed_top_22_bits = 0x3F900000
37 + # Generate random bits for the last 10 bits
38 + random_bottom_bits = torch.randint(
39 + 0,
40 + 2**10,
41 + (row_starts.shape[0], max(row_ends)),
42 + dtype=torch.int32,
43 + device="cuda",
44 + )
45 + # Combine: fixed top 22 bits with random last 10 bits
46 + logits_bits = (fixed_top_22_bits & top_22_bits_mask) | (
47 + random_bottom_bits & last_10_bits_mask
48 + )
49 + logits = logits_bits.view(dtype)
50 +
51 for i, end in enumerate(row_ends):
52 logits[i, end:] = float("-inf")
53 return logits
@@ -113,13 +135,13 @@ def test_top_k_per_row(
135 # Create test data
136 vocab_size = 20000
137 row_starts, row_ends = create_row_boundaries(num_rows, vocab_size)
116 - logits = create_random_logits(row_starts, row_ends, vocab_size, torch.float32, 42)
138 + logits = create_random_logits(row_starts, row_ends, torch.float32, 42, "random")
139
140 # Create output tensors
141 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
142
143 # Run CUDA implementation
122 - torch.ops._C.top_k_per_row(
144 + torch.ops._C.top_k_per_row_prefill(
145 logits,
146 row_starts,
147 row_ends,
@@ -127,6 +149,7 @@ def test_top_k_per_row(
149 num_rows,
150 logits.stride(0),
151 logits.stride(1),
152 + top_k,
153 )
154
155 # Run reference implementation
@@ -139,27 +162,23 @@ def test_top_k_per_row(
162 # Compare results
163 assert compare_top_k_results(
164 logits, indices, torch_indices, row_starts, row_ends, top_k
142 - ), "CUDA top_k_per_row results don't match torch.topk"
165 + ), "CUDA top_k_per_row_prefill results don't match torch.topk"
166
167
145 -@pytest.mark.parametrize("top_k", TOP_K_VALUES)
146 -@pytest.mark.parametrize("batch_size", BATCH_SIZE)
147 -@pytest.mark.parametrize("next_n", NEXT_N)
148 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
149 -@torch.inference_mode()
150 -def test_top_k_per_row_decode(
168 +def _run_top_k_per_row_decode_test(
169 top_k: int,
170 batch_size: int,
171 next_n: int,
172 + vocab_size: int,
173 + data_generation: str,
174 ) -> None:
175 """
156 - Test top_k_per_row with seq_lens tensor.
176 + Helper function to run top_k_per_row_decode test with given parameters.
177 """
178 torch.set_default_device("cuda:0")
179
180 # Create test data
181 num_rows = batch_size * next_n
162 - vocab_size = 20000
182 seq_lens = torch.randint(
183 vocab_size, (batch_size,), dtype=torch.int32, device="cuda"
184 )
@@ -167,7 +186,9 @@ def test_top_k_per_row_decode(
186 row_indices = torch.arange(num_rows, device="cuda") // next_n
187 next_n_offset = torch.arange(num_rows, device="cuda") % next_n
188 row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1
170 - logits = create_random_logits(row_starts, row_ends, vocab_size, torch.float32, 42)
189 + logits = create_random_logits(
190 + row_starts, row_ends, torch.float32, 42, data_generation
191 + )
192
193 # Create output tensors
194 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
@@ -181,6 +202,7 @@ def test_top_k_per_row_decode(
202 num_rows,
203 logits.stride(0),
204 logits.stride(1),
205 + top_k,
206 )
207
208 torch.cuda.synchronize()
@@ -195,4 +217,41 @@ def test_top_k_per_row_decode(
217 # Compare results
218 assert compare_top_k_results(
219 logits, indices, torch_indices, row_starts, row_ends, top_k
198 - ), "CUDA top_k_per_row results don't match torch.topk"
220 + ), "CUDA top_k_per_row_decode results don't match torch.topk"
221 +
222 +
223 +@pytest.mark.parametrize("top_k", TOP_K_VALUES)
224 +@pytest.mark.parametrize("batch_size", BATCH_SIZE)
225 +@pytest.mark.parametrize("next_n", NEXT_N)
226 +@pytest.mark.parametrize("data_generation", DATA_GENERATION)
227 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
228 +@torch.inference_mode()
229 +def test_top_k_per_row_decode(
230 + top_k: int,
231 + batch_size: int,
232 + next_n: int,
233 + data_generation: str,
234 +) -> None:
235 + """
236 + Test top_k_per_row with seq_lens tensor.
237 + """
238 + vocab_size = 20000
239 + _run_top_k_per_row_decode_test(
240 + top_k, batch_size, next_n, vocab_size, data_generation
241 + )
242 +
243 +
244 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
245 +@torch.inference_mode()
246 +def test_top_k_per_row_decode_large_vocab_size() -> None:
247 + """
248 + Test top_k_per_row_decode with large vocabulary size.
249 + """
250 + top_k = 2048
251 + batch_size = 2
252 + next_n = 2
253 + vocab_size = 300000
254 + data_generation = "random"
255 + _run_top_k_per_row_decode_test(
256 + top_k, batch_size, next_n, vocab_size, data_generation
257 + )
vllm/model_executor/models/deepseek_v2.py
+3 -3
@@ -684,11 +684,10 @@ def sparse_attn_indexer(
684 chunk.cu_seqlen_ke,
685 )
686 num_rows = logits.shape[0]
687 - assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
687 topk_indices = topk_indices_buffer[
688 chunk.token_start : chunk.token_end, :topk_tokens
689 ]
691 - torch.ops._C.top_k_per_row(
690 + torch.ops._C.top_k_per_row_prefill(
691 logits,
692 chunk.cu_seqlen_ks,
693 chunk.cu_seqlen_ke,
@@ -696,6 +695,7 @@ def sparse_attn_indexer(
695 num_rows,
696 logits.stride(0),
697 logits.stride(1),
698 + topk_tokens,
699 )
700
701 if has_decode:
@@ -738,7 +738,6 @@ def sparse_attn_indexer(
738 max_model_len=max_model_len,
739 )
740 num_rows = logits.shape[0]
741 - assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
741 topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]
742
743 torch.ops._C.top_k_per_row_decode(
@@ -749,6 +748,7 @@ def sparse_attn_indexer(
748 num_rows,
749 logits.stride(0),
750 logits.stride(1),
751 + topk_tokens,
752 )
753 if decode_metadata.requires_padding:
754 # if padded, we need to unpack