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 5, 11:00 UTC
[Perf][Kernel] Persistent TopK scheduler: unified CUDAGraph-safe kernel with dynamic per-row dispatch - DeepSeek-V3.2 DSA decode (#37421)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Roberto L. Castro committed
Apr 8, 2026 at 19:35 UTC
b55d830ec782b47f89c8457ba453f521fa092660
9 files changed
+2039
-483
.buildkite/test_areas/kernels.yaml
+3
-2
@@ -18,10 +18,9 @@ steps:
18
source_file_dependencies:
19
- csrc/
20
- tests/kernels/core
21
- - tests/kernels/test_top_k_per_row.py
21
- tests/kernels/test_concat_mla_q.py
22
commands:
24
- - pytest -v -s kernels/core kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py
23
+ - pytest -v -s kernels/core kernels/test_concat_mla_q.py
24
25
- label: Kernels Attention Test %N
26
timeout_in_minutes: 35
@@ -107,6 +106,7 @@ steps:
106
- vllm/v1/attention/backends/mla/flashinfer_mla.py
107
- vllm/v1/attention/selector.py
108
- vllm/platforms/cuda.py
109
+ - tests/kernels/test_top_k_per_row.py
110
commands:
111
- nvidia-smi
112
- python3 examples/basic/offline_inference/chat.py
@@ -117,6 +117,7 @@ steps:
117
- pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
118
- pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
119
- pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
120
+ - pytest -v -s tests/kernels/test_top_k_per_row.py
121
# Quantization
122
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
123
- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
csrc/ops.h
+3
-3
@@ -114,9 +114,9 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
114
int64_t numRows, int64_t stride0, int64_t stride1,
115
int64_t topK);
116
117
-void large_context_topk(const torch::Tensor& score, torch::Tensor& indices,
118
- const torch::Tensor& lengths,
119
- std::optional<torch::Tensor> row_starts_opt);
117
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
118
+ torch::Tensor& output, torch::Tensor& workspace, int64_t k,
119
+ int64_t max_seq_len);
120
121
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
122
torch::Tensor& weight, torch::Tensor& scale,
csrc/persistent_topk.cuh
new
+1321
@@ -0,0 +1,1321 @@
1
+/*
2
+ * Persistent TopK Scheduler for DSA Indexer
3
+ */
4
+
5
+#ifndef PERSISTENT_TOPK_CUH_
6
+#define PERSISTENT_TOPK_CUH_
7
+
8
+#include <cuda.h>
9
+#include <cuda_fp16.h>
10
+#include <cuda_runtime.h>
11
+#include <cub/cub.cuh>
12
+#include <cstdint>
13
+
14
+namespace vllm {
15
+namespace persistent {
16
+
17
+// ============================================================================
18
+// Constants
19
+// ============================================================================
20
+
21
+constexpr int TopK = 2048;
22
+constexpr int kThreadsPerBlock = 1024;
23
+constexpr int RADIX = 256;
24
+
25
+// Medium path: all shared state in dynamic smem (no static __shared__,
26
+// which would inflate the kernel's smem footprint and kill occupancy
27
+// for the decode/trivial paths).
28
+constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072
29
+constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20
30
+constexpr size_t kMediumHeaderSize =
31
+ (kMediumHistBytes + kMediumScalarsBytes + 127) & ~size_t(127); // 3200
32
+constexpr int MAX_BUFFERED_ITEMS = 4096;
33
+constexpr size_t kSmemMedium =
34
+ kMediumHeaderSize + 2 * MAX_BUFFERED_ITEMS * sizeof(int); // 35968
35
+constexpr uint32_t RADIX_THRESHOLD = 32768;
36
+
37
+// Decode path constants
38
+constexpr int kDecodeBins = 2048;
39
+constexpr uint32_t HIST2048_THRESHOLD = 8192;
40
+
41
+// Large path: fixed shared memory for histograms + scalars
42
+constexpr size_t kFixedSmemLarge =
43
+ ((RADIX + RADIX + 5) * sizeof(uint32_t) + 15) & ~size_t(15);
44
+
45
+// ============================================================================
46
+// Common helpers
47
+// ============================================================================
48
+
49
+__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
50
+ uint32_t bits = __float_as_uint(x);
51
+ return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
52
+}
53
+
54
+__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
55
+ __half h = __float2half_rn(x);
56
+ uint16_t bits = __half_as_ushort(h);
57
+ uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
58
+ : static_cast<uint16_t>(bits | 0x8000);
59
+ return static_cast<uint8_t>(key >> 8);
60
+}
61
+
62
+// ============================================================================
63
+// Vectorized load helpers
64
+// ============================================================================
65
+
66
+// Unconditional float4 load with cache hint (.cg = cache at global level only).
67
+__device__ __forceinline__ void load_float4(const float* ptr, float& v0,
68
+ float& v1, float& v2, float& v3) {
69
+ uint32_t r0, r1, r2, r3;
70
+ asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];\n"
71
+ : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
72
+ : "l"(ptr));
73
+ v0 = __uint_as_float(r0);
74
+ v1 = __uint_as_float(r1);
75
+ v2 = __uint_as_float(r2);
76
+ v3 = __uint_as_float(r3);
77
+}
78
+
79
+// Per-element predicated scalar loads with -inf default.
80
+__device__ __forceinline__ void load_float4_predicated(const float* ptr,
81
+ int base, int seq_len,
82
+ float& v0, float& v1,
83
+ float& v2, float& v3) {
84
+ uint32_t r0, r1, r2, r3;
85
+ int p0 = (base < seq_len);
86
+ int p1 = (base + 1 < seq_len);
87
+ int p2 = (base + 2 < seq_len);
88
+ int p3 = (base + 3 < seq_len);
89
+ asm volatile(
90
+ "{\n"
91
+ " .reg .pred pr0, pr1, pr2, pr3;\n"
92
+ " setp.ne.u32 pr0, %4, 0;\n"
93
+ " setp.ne.u32 pr1, %5, 0;\n"
94
+ " setp.ne.u32 pr2, %6, 0;\n"
95
+ " setp.ne.u32 pr3, %7, 0;\n"
96
+ " mov.u32 %0, 0xFF800000;\n"
97
+ " mov.u32 %1, 0xFF800000;\n"
98
+ " mov.u32 %2, 0xFF800000;\n"
99
+ " mov.u32 %3, 0xFF800000;\n"
100
+ " @pr0 ld.global.cg.u32 %0, [%8];\n"
101
+ " @pr1 ld.global.cg.u32 %1, [%8+4];\n"
102
+ " @pr2 ld.global.cg.u32 %2, [%8+8];\n"
103
+ " @pr3 ld.global.cg.u32 %3, [%8+12];\n"
104
+ "}\n"
105
+ : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
106
+ : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr));
107
+ v0 = __uint_as_float(r0);
108
+ v1 = __uint_as_float(r1);
109
+ v2 = __uint_as_float(r2);
110
+ v3 = __uint_as_float(r3);
111
+}
112
+
113
+// ============================================================================
114
+// Large path: inter-CTA coordination state (one per group)
115
+// ============================================================================
116
+
117
+struct RadixRowState {
118
+ uint32_t histogram[3][256]; // Triple-buffered histograms
119
+ uint32_t remaining_k;
120
+ uint32_t prefix;
121
+ int arrival_counter;
122
+ int output_counter;
123
+};
124
+
125
+// ============================================================================
126
+// Kernel parameters
127
+// ============================================================================
128
+
129
+struct PersistentTopKParams {
130
+ const float* __restrict__ input; // [num_rows, stride]
131
+ int32_t* __restrict__ output; // [num_rows, TopK]
132
+ int32_t* __restrict__ lengths; // [num_rows]
133
+ RadixRowState* row_states; // large path: per-group state
134
+ uint32_t num_rows;
135
+ uint32_t stride;
136
+ uint32_t chunk_size; // large path: elements per CTA
137
+ uint32_t ctas_per_group; // 1=medium, >1=large
138
+ uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit)
139
+};
140
+
141
+// ============================================================================
142
+// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192)
143
+// Uses 11-bit half-precision bins for fine granularity.
144
+// One histogram pass typically suffices since 8192/2048 = 4 elements/bin avg.
145
+// ============================================================================
146
+
147
+// 11-bit bin from half-precision representation (ascending: high values -> high
148
+// bins)
149
+__device__ __forceinline__ uint32_t decode_bin(float x) {
150
+ __half hx = __float2half(x);
151
+ uint16_t bits = __half_as_ushort(hx);
152
+ uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
153
+ : static_cast<uint16_t>(bits | 0x8000);
154
+ return key >> 5;
155
+}
156
+
157
+__device__ __noinline__ void histogram_2048_topk(
158
+ const float* __restrict__ logits, int32_t* __restrict__ output_indices,
159
+ int32_t seq_len) {
160
+ extern __shared__ int decode_smem[];
161
+ const int tx = threadIdx.x;
162
+ const int lane = tx & 31;
163
+
164
+ // ---- Layout constants ----
165
+ constexpr int SBASE = 8192 - 8; // 8184
166
+ constexpr int RHIST = RADIX + 128; // 384
167
+ constexpr int BOFF = 2 * RHIST; // 768
168
+ constexpr int DBUF = (SBASE - BOFF) / 2; // 3708
169
+ constexpr int MAX_ITEMS_PER_THREAD =
170
+ (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock;
171
+
172
+ enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 };
173
+
174
+ // ---- Initialize scalars (prevents stale data from prior rows) ----
175
+ if (tx < 8) {
176
+ decode_smem[SBASE + tx] = 0;
177
+ }
178
+
179
+ // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ----
180
+ int* histo = decode_smem;
181
+ uint16_t reg_bins[MAX_ITEMS_PER_THREAD];
182
+ int nitems = 0;
183
+
184
+ for (int i = tx; i < kDecodeBins; i += kThreadsPerBlock) {
185
+ histo[i] = 0;
186
+ }
187
+ __syncthreads();
188
+
189
+ const int n_vec = (seq_len + 3) >> 2;
190
+ const bool row_aligned = ((reinterpret_cast<uintptr_t>(logits) & 15) == 0);
191
+
192
+ for (int i = tx; i < n_vec; i += kThreadsPerBlock) {
193
+ const int base = i << 2;
194
+ float v0, v1, v2, v3;
195
+
196
+ if (row_aligned && base + 3 < seq_len) {
197
+ load_float4(logits + base, v0, v1, v2, v3);
198
+ } else {
199
+ load_float4_predicated(logits + base, base, seq_len, v0, v1, v2, v3);
200
+ }
201
+
202
+ const uint16_t b0 = static_cast<uint16_t>(decode_bin(v0));
203
+ const uint16_t b1 = static_cast<uint16_t>(decode_bin(v1));
204
+ const uint16_t b2 = static_cast<uint16_t>(decode_bin(v2));
205
+ const uint16_t b3 = static_cast<uint16_t>(decode_bin(v3));
206
+ reg_bins[nitems++] = b0;
207
+ reg_bins[nitems++] = b1;
208
+ reg_bins[nitems++] = b2;
209
+ reg_bins[nitems++] = b3;
210
+ atomicAdd(&histo[b0], 1);
211
+ atomicAdd(&histo[b1], 1);
212
+ atomicAdd(&histo[b2], 1);
213
+ atomicAdd(&histo[b3], 1);
214
+ }
215
+ __syncthreads();
216
+
217
+ // ---- CUB suffix sum ----
218
+ using BlockScanT = cub::BlockScan<int, kThreadsPerBlock>;
219
+ const int h0 = histo[2 * tx];
220
+ const int pair_sum = h0 + histo[2 * tx + 1];
221
+
222
+ auto& scan_storage = *reinterpret_cast<typename BlockScanT::TempStorage*>(
223
+ decode_smem + kDecodeBins);
224
+
225
+ int pair_prefix, total;
226
+ BlockScanT(scan_storage).ExclusiveSum(pair_sum, pair_prefix, total);
227
+
228
+ // Find threshold bin purely from registers
229
+ const int pair_suffix = total - pair_prefix;
230
+
231
+ if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) {
232
+ decode_smem[SBASE + sTHR] = 2 * tx;
233
+ }
234
+ {
235
+ const int right_suf = pair_suffix - h0;
236
+ const int next_suf = pair_suffix - pair_sum;
237
+ if (right_suf >= TopK && next_suf < TopK) {
238
+ decode_smem[SBASE + sTHR] = 2 * tx + 1;
239
+ }
240
+ }
241
+ __syncthreads();
242
+
243
+ const int threshold = decode_smem[SBASE + sTHR];
244
+
245
+ // ---- Phase 2: Collection with warp-aggregated atomicAdds ----
246
+ int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF};
247
+ const int sOUT_abs = SBASE + sOUT;
248
+ const int sBUF0_abs = SBASE + sBUF0;
249
+
250
+ {
251
+ const uint32_t uthr = static_cast<uint32_t>(threshold);
252
+ int item = 0;
253
+ const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock;
254
+
255
+ for (int iter = 0; iter < n_vec_iters; iter++) {
256
+ const int i = tx + iter * kThreadsPerBlock;
257
+ const bool vec_valid = (i < n_vec);
258
+ const int base_idx = i << 2;
259
+
260
+#pragma unroll 4
261
+ for (int sub = 0; sub < 4; sub++) {
262
+ const int elem_idx = base_idx + sub;
263
+ uint32_t bin = 0;
264
+ if (vec_valid) bin = reg_bins[item++];
265
+ const bool is_above = vec_valid && (bin > uthr);
266
+ const bool is_equal = vec_valid && (bin == uthr);
267
+
268
+ const uint32_t above_mask = __ballot_sync(0xffffffff, is_above);
269
+ if (above_mask) {
270
+ const int above_count = __popc(above_mask);
271
+ const int above_rank = __popc(above_mask & ((1u << lane) - 1));
272
+ int above_base;
273
+ if (lane == 0) {
274
+ above_base = atomicAdd(&decode_smem[sOUT_abs], above_count);
275
+ }
276
+ above_base = __shfl_sync(0xffffffff, above_base, 0);
277
+ if (is_above) {
278
+ output_indices[above_base + above_rank] = elem_idx;
279
+ }
280
+ }
281
+
282
+ const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal);
283
+ if (equal_mask) {
284
+ const int equal_count = __popc(equal_mask);
285
+ const int equal_rank = __popc(equal_mask & ((1u << lane) - 1));
286
+ int equal_base;
287
+ if (lane == 0) {
288
+ equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count);
289
+ }
290
+ equal_base = __shfl_sync(0xffffffff, equal_base, 0);
291
+ if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) {
292
+ bufs[0][equal_base + equal_rank] = elem_idx;
293
+ }
294
+ }
295
+ }
296
+ }
297
+ }
298
+ __syncthreads();
299
+
300
+ int remaining_k = TopK - decode_smem[SBASE + sOUT];
301
+ if (remaining_k <= 0) return;
302
+
303
+ // If all buffered elements fit, output them all (common for short seqs)
304
+ const int raw_buf0 = decode_smem[SBASE + sBUF0];
305
+ if (raw_buf0 <= remaining_k) {
306
+ const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF;
307
+ const int base = decode_smem[SBASE + sOUT];
308
+ for (int i = tx; i < nb; i += kThreadsPerBlock) {
309
+ output_indices[base + i] = bufs[0][i];
310
+ }
311
+ __syncthreads();
312
+ return;
313
+ }
314
+
315
+ // ---- Phase 3: Deferred refinement (rare path) ----
316
+ int* refine[2] = {decode_smem, decode_smem + RHIST};
317
+ const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF;
318
+
319
+ for (int i = tx; i < RHIST; i += kThreadsPerBlock) {
320
+ refine[0][i] = 0;
321
+ }
322
+ __syncthreads();
323
+
324
+ for (int i = tx; i < num_buf0; i += kThreadsPerBlock) {
325
+ const uint32_t fp32 = convert_to_uint32_v2(logits[bufs[0][i]]);
326
+ atomicAdd(&refine[0][(fp32 >> 24) & 0xFF], 1);
327
+ }
328
+ __syncthreads();
329
+
330
+ auto compute_suffix_sum = [&]() {
331
+#pragma unroll 8
332
+ for (int i = 0; i < 8; ++i) {
333
+ if (tx < RADIX) {
334
+ const int stride = 1 << i;
335
+ const int s = i & 1;
336
+ const int d = s ^ 1;
337
+ int value = refine[s][tx];
338
+ if (tx < RADIX - stride) value += refine[s][tx + stride];
339
+ refine[d][tx] = value;
340
+ }
341
+ __syncthreads();
342
+ }
343
+ };
344
+
345
+#pragma unroll 4
346
+ for (int pass = 0; pass < 4; ++pass) {
347
+ const int src = pass & 1;
348
+ const int dst = src ^ 1;
349
+
350
+ const int raw_buf = decode_smem[SBASE + sBUF0 + src];
351
+ const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF;
352
+
353
+ compute_suffix_sum();
354
+
355
+ if (tx < RADIX && refine[0][tx] > remaining_k &&
356
+ refine[0][tx + 1] <= remaining_k) {
357
+ decode_smem[SBASE + sREF] = tx;
358
+ decode_smem[SBASE + sBUF0 + dst] = 0;
359
+ decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1];
360
+ }
361
+ __syncthreads();
362
+
363
+ const int ref_thr = decode_smem[SBASE + sREF];
364
+ remaining_k -= refine[0][ref_thr + 1];
365
+ const int bit_offset = 24 - pass * 8;
366
+
367
+ if (remaining_k == 0) {
368
+ for (int i = tx; i < num_buffered; i += kThreadsPerBlock) {
369
+ const int idx = bufs[src][i];
370
+ const uint32_t fp32 = convert_to_uint32_v2(logits[idx]);
371
+ if (((fp32 >> bit_offset) & 0xFF) > static_cast<uint32_t>(ref_thr)) {
372
+ const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1);
373
+ output_indices[pos] = idx;
374
+ }
375
+ }
376
+ __syncthreads();
377
+ break;
378
+ }
379
+
380
+ __syncthreads();
381
+ if (tx < RADIX + 1) refine[0][tx] = 0;
382
+ __syncthreads();
383
+
384
+ for (int i = tx; i < num_buffered; i += kThreadsPerBlock) {
385
+ const int idx = bufs[src][i];
386
+ const float logit_val = logits[idx];
387
+ const uint32_t fp32 = convert_to_uint32_v2(logit_val);
388
+ const int bin = (fp32 >> bit_offset) & 0xFF;
389
+
390
+ if (bin > ref_thr) {
391
+ const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1);
392
+ output_indices[pos] = idx;
393
+ } else if (bin == ref_thr) {
394
+ if (pass == 3) {
395
+ const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1);
396
+ if (slot > 0) output_indices[TopK - slot] = idx;
397
+ } else {
398
+ const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1);
399
+ if (__builtin_expect(bp < DBUF, 1)) {
400
+ bufs[dst][bp] = idx;
401
+ const int nbo = bit_offset - 8;
402
+ atomicAdd(&refine[0][(fp32 >> nbo) & 0xFF], 1);
403
+ }
404
+ }
405
+ }
406
+ }
407
+ __syncthreads();
408
+ }
409
+}
410
+
411
+// ============================================================================
412
+// Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement
413
+// For sequences 8K < seq_len <= 64K.
414
+// ============================================================================
415
+
416
+// Adapted from:
417
+// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87
418
+// by: DarkSharpness
419
+// which at the same time is an optimized topk kernel copied from tilelang
420
+// kernel
421
+__device__ __noinline__ void histogram_256_topk(
422
+ const float* __restrict__ logits, int* __restrict__ output_indices,
423
+ int logits_offset, int seq_len) {
424
+ // All shared state lives in dynamic shared memory to avoid static
425
+ extern __shared__ char medium_smem[];
426
+
427
+ int (*shared_histogram)[RADIX + 128] =
428
+ reinterpret_cast<int (*)[RADIX + 128]>(medium_smem);
429
+ int* medium_scalars = reinterpret_cast<int*>(medium_smem + kMediumHistBytes);
430
+ int& shared_output_count = medium_scalars[0];
431
+ int& shared_threshold_bin = medium_scalars[1];
432
+ int* shared_buffered_count = &medium_scalars[2];
433
+ int& shared_final_k = medium_scalars[4];
434
+ int (*buffered_indices)[MAX_BUFFERED_ITEMS] =
435
+ reinterpret_cast<int (*)[MAX_BUFFERED_ITEMS]>(medium_smem +
436
+ kMediumHeaderSize);
437
+
438
+ const int thread_id = threadIdx.x;
439
+ int remaining_k = TopK;
440
+
441
+ if (thread_id < RADIX + 1) {
442
+ shared_histogram[0][thread_id] = 0;
443
+ }
444
+ __syncthreads();
445
+
446
+ for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
447
+ const auto bin = convert_to_uint8(logits[idx + logits_offset]);
448
+ atomicAdd(&shared_histogram[0][bin], 1);
449
+ }
450
+ __syncthreads();
451
+
452
+ auto compute_cumulative_sum = [&]() {
453
+#pragma unroll 8
454
+ for (int i = 0; i < 8; ++i) {
455
+ if (__builtin_expect(thread_id < RADIX, 1)) {
456
+ const int stride = 1 << i;
457
+ const int src_buffer = i & 1;
458
+ const int dst_buffer = src_buffer ^ 1;
459
+ int value = shared_histogram[src_buffer][thread_id];
460
+ if (thread_id < RADIX - stride) {
461
+ value += shared_histogram[src_buffer][thread_id + stride];
462
+ }
463
+ shared_histogram[dst_buffer][thread_id] = value;
464
+ }
465
+ __syncthreads();
466
+ }
467
+ };
468
+
469
+ compute_cumulative_sum();
470
+
471
+ if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
472
+ shared_histogram[0][thread_id + 1] <= remaining_k) {
473
+ shared_threshold_bin = thread_id;
474
+ shared_buffered_count[0] = 0;
475
+ shared_output_count = 0;
476
+ }
477
+ __syncthreads();
478
+
479
+ const int threshold_bin = shared_threshold_bin;
480
+ remaining_k -= shared_histogram[0][threshold_bin + 1];
481
+
482
+ if (remaining_k == 0) {
483
+ for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
484
+ const int bin = convert_to_uint8(logits[idx + logits_offset]);
485
+ if (bin > threshold_bin) {
486
+ const int output_pos = atomicAdd(&shared_output_count, 1);
487
+ output_indices[output_pos] = idx;
488
+ }
489
+ }
490
+ __syncthreads();
491
+ return;
492
+ }
493
+
494
+ __syncthreads();
495
+ if (thread_id < RADIX + 1) {
496
+ shared_histogram[0][thread_id] = 0;
497
+ }
498
+ __syncthreads();
499
+
500
+ for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
501
+ const float logit_value = logits[idx + logits_offset];
502
+ const int bin = convert_to_uint8(logit_value);
503
+ if (bin > threshold_bin) {
504
+ const int output_pos = atomicAdd(&shared_output_count, 1);
505
+ output_indices[output_pos] = idx;
506
+ } else if (bin == threshold_bin) {
507
+ const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1);
508
+ if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) {
509
+ buffered_indices[0][buffer_pos] = idx;
510
+ const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
511
+ const int next_bin = (fp32_bits >> 24) & 0xFF;
512
+ atomicAdd(&shared_histogram[0][next_bin], 1);
513
+ }
514
+ }
515
+ }
516
+ __syncthreads();
517
+
518
+#pragma unroll 4
519
+ for (int pass = 0; pass < 4; ++pass) {
520
+ const int src_buffer = pass % 2;
521
+ const int dst_buffer = src_buffer ^ 1;
522
+ const int raw_buffered = shared_buffered_count[src_buffer];
523
+ const int num_buffered =
524
+ (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS;
525
+
526
+ compute_cumulative_sum();
527
+
528
+ if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
529
+ shared_histogram[0][thread_id + 1] <= remaining_k) {
530
+ shared_threshold_bin = thread_id;
531
+ shared_buffered_count[dst_buffer] = 0;
532
+ shared_final_k = remaining_k - shared_histogram[0][thread_id + 1];
533
+ }
534
+ __syncthreads();
535
+
536
+ const int threshold_bin = shared_threshold_bin;
537
+ remaining_k -= shared_histogram[0][threshold_bin + 1];
538
+ const int bit_offset = 24 - pass * 8;
539
+
540
+ if (remaining_k == 0) {
541
+ for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
542
+ const int idx = buffered_indices[src_buffer][i];
543
+ const uint32_t fp32_bits =
544
+ convert_to_uint32_v2(logits[idx + logits_offset]);
545
+ const int bin = (fp32_bits >> bit_offset) & 0xFF;
546
+ if (bin > threshold_bin) {
547
+ const int output_pos = atomicAdd(&shared_output_count, 1);
548
+ output_indices[output_pos] = idx;
549
+ }
550
+ }
551
+ __syncthreads();
552
+ break;
553
+ }
554
+
555
+ __syncthreads();
556
+ if (thread_id < RADIX + 1) {
557
+ shared_histogram[0][thread_id] = 0;
558
+ }
559
+ __syncthreads();
560
+
561
+ for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
562
+ const int idx = buffered_indices[src_buffer][i];
563
+ const float logit_value = logits[idx + logits_offset];
564
+ const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
565
+ const int bin = (fp32_bits >> bit_offset) & 0xFF;
566
+ if (bin > threshold_bin) {
567
+ const int output_pos = atomicAdd(&shared_output_count, 1);
568
+ output_indices[output_pos] = idx;
569
+ } else if (bin == threshold_bin) {
570
+ if (pass == 3) {
571
+ const int slot = atomicAdd(&shared_final_k, -1);
572
+ if (slot > 0) {
573
+ output_indices[TopK - slot] = idx;
574
+ }
575
+ } else {
576
+ const int buffer_pos =
577
+ atomicAdd(&shared_buffered_count[dst_buffer], 1);
578
+ if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) {
579
+ buffered_indices[dst_buffer][buffer_pos] = idx;
580
+ const int next_bit_offset = bit_offset - 8;
581
+ const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF;
582
+ atomicAdd(&shared_histogram[0][next_bin], 1);
583
+ }
584
+ }
585
+ }
586
+ }
587
+ __syncthreads();
588
+ }
589
+}
590
+
591
+// ============================================================================
592
+// Inter-CTA sync primitives
593
+// ============================================================================
594
+
595
+__device__ __forceinline__ int ld_acquire(int* ptr) {
596
+ int state = 0;
597
+#if (__CUDA_ARCH__ >= 700)
598
+ asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
599
+ : "=r"(state)
600
+ : "l"(ptr));
601
+#else
602
+ asm volatile("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr));
603
+#endif
604
+ return state;
605
+}
606
+
607
+__device__ __forceinline__ void red_release(int* ptr, int val) {
608
+#if (__CUDA_ARCH__ >= 700)
609
+ asm volatile("fence.acq_rel.gpu;\n");
610
+ asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
611
+ :
612
+ : "l"(ptr), "r"(val));
613
+#else
614
+ __threadfence();
615
+ atomicAdd(ptr, val);
616
+#endif
617
+}
618
+
619
+__device__ __forceinline__ void st_release(int* ptr, int val) {
620
+#if (__CUDA_ARCH__ >= 700)
621
+ asm volatile("fence.acq_rel.gpu;\n");
622
+ asm volatile("st.release.gpu.global.b32 [%0], %1;\n" : : "l"(ptr), "r"(val));
623
+#else
624
+ __threadfence();
625
+ atomicExch(ptr, val);
626
+#endif
627
+}
628
+
629
+__device__ __forceinline__ void wait_ge(int* ptr, int target_val,
630
+ int thread_idx) {
631
+ if (thread_idx == 0) {
632
+#pragma unroll 1
633
+ while (ld_acquire(ptr) < target_val) {
634
+ }
635
+ }
636
+ __syncthreads();
637
+}
638
+
639
+// ============================================================================
640
+// Large path: multi-CTA radix select for sequences > 64K
641
+//
642
+// Each row is processed by a group of CTAs. Each CTA loads its chunk into
643
+// shared memory as ordered uint32, then participates in 4 rounds of
644
+// coordinated radix select via global-memory histograms and barriers.
645
+// ============================================================================
646
+
647
+// ============================================================================
648
+// Multi-CTA cooperative RadixTopK for a single large row.
649
+// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
650
+// ============================================================================
651
+
652
+template <uint32_t VEC_SIZE>
653
+__device__ void radix_topk(const float* __restrict__ row_input,
654
+ int32_t* __restrict__ row_output, uint32_t seq_len,
655
+ uint32_t my_chunk_start, uint32_t chunk_size,
656
+ uint32_t* local_histogram, uint32_t* suffix_sum,
657
+ uint32_t* shared_scalars, uint32_t* shared_ordered,
658
+ RadixRowState* state, uint32_t cta_in_group,
659
+ uint32_t ctas_per_group, int& barrier_phase,
660
+ uint32_t iter, uint32_t tx) {
661
+ const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
662
+ ? my_chunk_start + chunk_size
663
+ : seq_len;
664
+ const uint32_t actual_chunk_size =
665
+ (my_chunk_start < seq_len) ? (my_chunk_end - my_chunk_start) : 0;
666
+
667
+ // -- Stage 1: Load chunk to shared memory as ordered uint32 --
668
+ {
669
+ const uint32_t aligned_size = (actual_chunk_size / VEC_SIZE) * VEC_SIZE;
670
+
671
+ for (uint32_t i = tx * VEC_SIZE; i < aligned_size;
672
+ i += kThreadsPerBlock * VEC_SIZE) {
673
+ const float* src = row_input + my_chunk_start + i;
674
+ if constexpr (VEC_SIZE == 4) {
675
+ float4 v = *reinterpret_cast<const float4*>(src);
676
+ shared_ordered[i] = convert_to_uint32_v2(v.x);
677
+ shared_ordered[i + 1] = convert_to_uint32_v2(v.y);
678
+ shared_ordered[i + 2] = convert_to_uint32_v2(v.z);
679
+ shared_ordered[i + 3] = convert_to_uint32_v2(v.w);
680
+ } else if constexpr (VEC_SIZE == 2) {
681
+ float2 v = *reinterpret_cast<const float2*>(src);
682
+ shared_ordered[i] = convert_to_uint32_v2(v.x);
683
+ shared_ordered[i + 1] = convert_to_uint32_v2(v.y);
684
+ } else {
685
+ shared_ordered[i] = convert_to_uint32_v2(*src);
686
+ }
687
+ }
688
+ for (uint32_t i = aligned_size + tx; i < actual_chunk_size;
689
+ i += kThreadsPerBlock) {
690
+ shared_ordered[i] = convert_to_uint32_v2(row_input[my_chunk_start + i]);
691
+ }
692
+ }
693
+ __syncthreads();
694
+
695
+ // -- Init radix select state --
696
+ if (tx == 0) {
697
+ shared_scalars[0] = 0; // prefix
698
+ shared_scalars[1] = TopK; // remaining_k
699
+ }
700
+ __syncthreads();
701
+
702
+ // -- Initial barrier --
703
+ if (tx == 0) {
704
+ red_release(&state->arrival_counter, 1);
705
+ }
706
+ wait_ge(&state->arrival_counter,
707
+ (barrier_phase + 1) * static_cast<int>(ctas_per_group), tx);
708
+ barrier_phase++;
709
+ __syncthreads();
710
+
711
+ if (cta_in_group == 0 && tx == 0) {
712
+ st_release(&state->output_counter, 0);
713
+ }
714
+
715
+ // -- Stage 2: 4 rounds of radix select --
716
+ for (uint32_t round = 0; round < 4; round++) {
717
+ const uint32_t global_round = iter * 4 + round;
718
+ const uint32_t shift = 24 - round * 8;
719
+ const uint32_t prefix = shared_scalars[0];
720
+ const uint32_t remaining_k = shared_scalars[1];
721
+
722
+ uint32_t* current_hist = state->histogram[global_round % 3];
723
+ uint32_t* next_hist = state->histogram[(global_round + 1) % 3];
724
+
725
+ for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
726
+ local_histogram[i] = 0;
727
+ }
728
+ __syncthreads();
729
+
730
+ for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) {
731
+ uint32_t ordered = shared_ordered[i];
732
+ uint32_t mask = (round == 0) ? 0u : (~0u << (32 - round * 8));
733
+ if ((ordered & mask) == prefix) {
734
+ uint32_t bucket = (ordered >> shift) & 0xFF;
735
+ atomicAdd(&local_histogram[bucket], 1);
736
+ }
737
+ }
738
+ __syncthreads();
739
+
740
+ for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
741
+ if (local_histogram[i] > 0) {
742
+ atomicAdd(¤t_hist[i], local_histogram[i]);
743
+ }
744
+ }
745
+
746
+ if (cta_in_group == 0) {
747
+ for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
748
+ next_hist[i] = 0;
749
+ }
750
+ }
751
+
752
+ if (tx == 0) {
753
+ red_release(&state->arrival_counter, 1);
754
+ }
755
+ wait_ge(&state->arrival_counter,
756
+ (barrier_phase + 1) * static_cast<int>(ctas_per_group), tx);
757
+ barrier_phase++;
758
+ __syncthreads();
759
+
760
+ for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
761
+ suffix_sum[i] = current_hist[i];
762
+ }
763
+ __syncthreads();
764
+
765
+ for (uint32_t stride = 1; stride < RADIX; stride *= 2) {
766
+ uint32_t val = 0;
767
+ if (tx < RADIX) {
768
+ val = suffix_sum[tx];
769
+ if (tx + stride < RADIX) val += suffix_sum[tx + stride];
770
+ }
771
+ __syncthreads();
772
+ if (tx < RADIX) suffix_sum[tx] = val;
773
+ __syncthreads();
774
+ }
775
+
776
+ if (tx == 0) {
777
+ shared_scalars[2] = 0;
778
+ shared_scalars[3] = remaining_k;
779
+ }
780
+ __syncthreads();
781
+
782
+ if (tx < RADIX) {
783
+ uint32_t count_ge = suffix_sum[tx];
784
+ uint32_t count_gt = (tx + 1 < RADIX) ? suffix_sum[tx + 1] : 0;
785
+ if (count_ge >= remaining_k && count_gt < remaining_k) {
786
+ shared_scalars[2] = tx;
787
+ shared_scalars[3] = remaining_k - count_gt;
788
+ }
789
+ }
790
+ __syncthreads();
791
+
792
+ if (tx == 0) {
793
+ shared_scalars[0] = prefix | (shared_scalars[2] << shift);
794
+ shared_scalars[1] = shared_scalars[3];
795
+ }
796
+ __syncthreads();
797
+ } // end 4 radix rounds
798
+
799
+ // -- Count local > pivot elements --
800
+ const uint32_t ordered_pivot = shared_scalars[0];
801
+
802
+ if (tx == 0) suffix_sum[0] = 0;
803
+ __syncthreads();
804
+
805
+ uint32_t my_gt_count = 0;
806
+ for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) {
807
+ if (shared_ordered[i] > ordered_pivot) my_gt_count++;
808
+ }
809
+ for (int offset = 16; offset > 0; offset /= 2) {
810
+ my_gt_count += __shfl_down_sync(0xffffffff, my_gt_count, offset);
811
+ }
812
+ if (tx % 32 == 0 && my_gt_count > 0) {
813
+ atomicAdd(&suffix_sum[0], my_gt_count);
814
+ }
815
+ __syncthreads();
816
+ const uint32_t local_gt_count = suffix_sum[0];
817
+
818
+ // -- Stage 3: Collect top-k indices --
819
+ if (tx == 0) {
820
+ local_histogram[0] = 0;
821
+ if (local_gt_count > 0) {
822
+ local_histogram[1] =
823
+ atomicAdd(&state->output_counter, static_cast<int>(local_gt_count));
824
+ }
825
+ }
826
+ __syncthreads();
827
+
828
+ for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) {
829
+ if (shared_ordered[i] > ordered_pivot) {
830
+ uint32_t local_pos = atomicAdd(&local_histogram[0], 1);
831
+ int pos = static_cast<int>(local_histogram[1]) + local_pos;
832
+ row_output[pos] = static_cast<int32_t>(my_chunk_start + i);
833
+ }
834
+ }
835
+
836
+ if (tx == 0) {
837
+ red_release(&state->arrival_counter, 1);
838
+ }
839
+ wait_ge(&state->arrival_counter,
840
+ (barrier_phase + 1) * static_cast<int>(ctas_per_group), tx);
841
+ barrier_phase++;
842
+ __syncthreads();
843
+
844
+ for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) {
845
+ if (shared_ordered[i] == ordered_pivot) {
846
+ int pos = atomicAdd(&state->output_counter, 1);
847
+ if (pos < TopK) {
848
+ row_output[pos] = static_cast<int32_t>(my_chunk_start + i);
849
+ }
850
+ }
851
+ }
852
+}
853
+
854
+// ============================================================================
855
+// Persistent kernel — BS≤32, decode/medium/large paths with RadixTopK
856
+// BS>32 uses standalone histogram_256_buffered_topk (separate kernel,
857
+// see filtered_topk.cuh)
858
+// ============================================================================
859
+
860
+template <uint32_t VEC_SIZE = 1>
861
+__global__ void __launch_bounds__(kThreadsPerBlock, 2)
862
+ persistent_topk_kernel(PersistentTopKParams params) {
863
+ const uint32_t tx = threadIdx.x;
864
+ extern __shared__ uint8_t smem_raw[];
865
+
866
+ // ========================================================================
867
+ // Group mode: multi-CTA groups with static round-robin row assignment.
868
+ // Non-large rows: CTA-0 handles trivial/decode/medium.
869
+ // Large rows: all CTAs in the group cooperate via RadixTopK.
870
+ // ========================================================================
871
+ const uint32_t ctas_per_group = params.ctas_per_group;
872
+ const uint32_t group_id = blockIdx.x / ctas_per_group;
873
+ const uint32_t cta_in_group = blockIdx.x % ctas_per_group;
874
+ const uint32_t num_groups = gridDim.x / ctas_per_group;
875
+ const uint32_t chunk_size = params.chunk_size;
876
+
877
+ if (blockIdx.x >= num_groups * ctas_per_group) return;
878
+
879
+ // Early exit: non-CTA-0 threads are never needed if no large rows exist
880
+ if (cta_in_group != 0 && params.max_seq_len <= RADIX_THRESHOLD) return;
881
+
882
+ uint32_t* local_histogram = reinterpret_cast<uint32_t*>(smem_raw);
883
+ uint32_t* suffix_sum = local_histogram + RADIX;
884
+ uint32_t* shared_scalars = suffix_sum + RADIX;
885
+ uint32_t* shared_ordered =
886
+ reinterpret_cast<uint32_t*>(smem_raw + kFixedSmemLarge);
887
+
888
+ // RadixRowState for multi-CTA cooperative radix
889
+ RadixRowState* state = ¶ms.row_states[group_id];
890
+
891
+ // -- Initialize RadixRowState (only needed if large rows exist) --
892
+ if (params.max_seq_len > RADIX_THRESHOLD) {
893
+ if (cta_in_group == 0) {
894
+ for (uint32_t buf = 0; buf < 3; buf++) {
895
+ for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
896
+ state->histogram[buf][i] = 0;
897
+ }
898
+ }
899
+ if (tx == 0) {
900
+ state->remaining_k = 0;
901
+ state->prefix = 0;
902
+ state->arrival_counter = 0;
903
+ state->output_counter = 0;
904
+ }
905
+ }
906
+ __syncthreads();
907
+ }
908
+
909
+ int barrier_phase = 0;
910
+ const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;
911
+
912
+ for (uint32_t iter = 0; iter < total_iters; iter++) {
913
+ // Static round-robin: all CTAs in the group implicitly agree on the row
914
+ uint32_t row_idx = group_id + iter * num_groups;
915
+ if (row_idx >= params.num_rows) break;
916
+
917
+ const uint32_t seq_len = params.lengths[row_idx];
918
+ int32_t* row_output = params.output + row_idx * TopK;
919
+ const float* row_input = params.input + row_idx * params.stride;
920
+
921
+ if (seq_len <= RADIX_THRESHOLD) {
922
+ if (cta_in_group == 0) {
923
+ if (seq_len <= static_cast<uint32_t>(TopK)) {
924
+ // Trivial case: seq_len <= TopK
925
+ for (uint32_t i = tx; i < static_cast<uint32_t>(TopK);
926
+ i += kThreadsPerBlock) {
927
+ row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
928
+ }
929
+ } else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
930
+ histogram_2048_topk(row_input, row_output, seq_len);
931
+ } else {
932
+ histogram_256_topk(row_input, row_output, 0, seq_len);
933
+ }
934
+ }
935
+ continue;
936
+ }
937
+
938
+ const uint32_t my_chunk_start = cta_in_group * chunk_size;
939
+ radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
940
+ chunk_size, local_histogram, suffix_sum,
941
+ shared_scalars, shared_ordered, state, cta_in_group,
942
+ ctas_per_group, barrier_phase, iter, tx);
943
+ }
944
+}
945
+
946
+} // namespace persistent
947
+
948
+// ============================================================================
949
+// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only.
950
+// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent).
951
+// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
952
+// ============================================================================
953
+
954
+#define FLASHINFER_CUDA_CALL(func, ...) \
955
+ { \
956
+ cudaError_t e = (func); \
957
+ if (e != cudaSuccess) { \
958
+ return e; \
959
+ } \
960
+ }
961
+
962
+#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__
963
+
964
+template <typename T, size_t N>
965
+struct vec_t {
966
+ T data[N];
967
+
968
+ FLASHINFER_INLINE T& operator[](size_t i) { return data[i]; }
969
+ FLASHINFER_INLINE const T& operator[](size_t i) const { return data[i]; }
970
+
971
+ FLASHINFER_INLINE void cast_load(const T* ptr) {
972
+#pragma unroll
973
+ for (size_t i = 0; i < N; ++i) {
974
+ data[i] = ptr[i];
975
+ }
976
+ }
977
+
978
+ FLASHINFER_INLINE void cast_store(T* ptr) const {
979
+#pragma unroll
980
+ for (size_t i = 0; i < N; ++i) {
981
+ ptr[i] = data[i];
982
+ }
983
+ }
984
+};
985
+#undef FLASHINFER_INLINE
986
+
987
+// FilteredTopK traits for different data types
988
+template <typename DType>
989
+struct FilteredTopKTraits;
990
+
991
+// Specialization for float (32-bit): coarse histogram uses FP16 high 8 bits, 4
992
+// refinement rounds
993
+template <>
994
+struct FilteredTopKTraits<float> {
995
+ using OrderedType = uint32_t;
996
+ static constexpr int NUM_REFINE_ROUNDS = 4;
997
+ static constexpr int FIRST_REFINE_SHIFT = 24;
998
+
999
+ __device__ __forceinline__ static uint8_t ToCoarseKey(float x) {
1000
+ // Convert to FP16 representation and extract high 8 bits
1001
+ __half h = __float2half_rn(x);
1002
+ uint16_t bits = __half_as_ushort(h);
1003
+ uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
1004
+ : static_cast<uint16_t>(bits | 0x8000);
1005
+ return static_cast<uint8_t>(key >> 8);
1006
+ }
1007
+
1008
+ __device__ __forceinline__ static OrderedType ToOrdered(float x) {
1009
+ uint32_t bits = __float_as_uint(x);
1010
+ return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
1011
+ }
1012
+};
1013
+
1014
+constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
1015
+constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
1016
+constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
1017
+ 16 * 1024; // 16K indices per buffer
1018
+constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
1019
+ sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB
1020
+
1021
+/*!
1022
+ * \brief Filtered Top-K kernel for ragged sequences.
1023
+ *
1024
+ * \tparam DType Data type (float, half, nv_bfloat16)
1025
+ * \tparam IdType Index type (int32_t)
1026
+ * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
1027
+ */
1028
+template <typename DType, typename IdType, int VEC_SIZE>
1029
+__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1030
+ FilteredTopKUnifiedKernel(const DType* __restrict__ input,
1031
+ IdType* __restrict__ output,
1032
+ const IdType* __restrict__ lengths,
1033
+ uint32_t num_rows, uint32_t top_k,
1034
+ uint32_t max_len) {
1035
+ constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS;
1036
+ constexpr int RADIX = 256;
1037
+ constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE;
1038
+
1039
+ const uint32_t bid = blockIdx.x;
1040
+ const int tx = threadIdx.x;
1041
+
1042
+ if (bid >= num_rows) return;
1043
+
1044
+ const int length =
1045
+ (lengths != nullptr) ? lengths[bid] : static_cast<int>(max_len);
1046
+ const DType* score = input + bid * max_len;
1047
+ IdType* dst = output + bid * top_k;
1048
+
1049
+ // Trivial case: length <= top_k
1050
+ if (length <= static_cast<int>(top_k)) {
1051
+ for (int i = tx; i < static_cast<int>(top_k); i += BLOCK_SIZE) {
1052
+ dst[i] = (i < length) ? static_cast<IdType>(i) : static_cast<IdType>(-1);
1053
+ }
1054
+ return;
1055
+ }
1056
+
1057
+ // Static shared memory
1058
+ alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
1059
+ alignas(128) __shared__ int s_counter;
1060
+ alignas(128) __shared__ int s_threshold_bin_id;
1061
+ alignas(128) __shared__ int s_num_input[2];
1062
+ alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
1063
+
1064
+ auto& s_histogram = s_histogram_buf[0];
1065
+
1066
+ // Dynamic shared memory for input double buffer
1067
+ extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
1068
+
1069
+ using Traits = FilteredTopKTraits<DType>;
1070
+ int topk = top_k;
1071
+
1072
+ // Stage 1: 8-bit coarse histogram with vectorized loads
1073
+ if (tx < RADIX + 1) s_histogram[tx] = 0;
1074
+ __syncthreads();
1075
+
1076
+ vec_t<DType, VEC_SIZE> score_vec;
1077
+
1078
+ const int aligned_length = (length / VEC_SIZE) * VEC_SIZE;
1079
+#pragma unroll 2
1080
+ for (int base = tx * VEC_SIZE; base < aligned_length;
1081
+ base += BLOCK_SIZE * VEC_SIZE) {
1082
+ score_vec.cast_load(&score[base]);
1083
+#pragma unroll
1084
+ for (int j = 0; j < VEC_SIZE; ++j) {
1085
+ const auto bin = Traits::ToCoarseKey(score_vec[j]);
1086
+ atomicAdd(&s_histogram[bin], 1);
1087
+ }
1088
+ }
1089
+ // Handle tail
1090
+ for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) {
1091
+ const auto bin = Traits::ToCoarseKey(score[i]);
1092
+ atomicAdd(&s_histogram[bin], 1);
1093
+ }
1094
+ __syncthreads();
1095
+
1096
+ // Suffix sum
1097
+ const auto run_cumsum = [&]() {
1098
+#pragma unroll 8
1099
+ for (int i = 0; i < 8; ++i) {
1100
+ if (tx < RADIX) {
1101
+ const auto j = 1 << i;
1102
+ const auto k = i & 1;
1103
+ auto value = s_histogram_buf[k][tx];
1104
+ if (tx < RADIX - j) {
1105
+ value += s_histogram_buf[k][tx + j];
1106
+ }
1107
+ s_histogram_buf[k ^ 1][tx] = value;
1108
+ }
1109
+ __syncthreads();
1110
+ }
1111
+ };
1112
+
1113
+ run_cumsum();
1114
+ if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
1115
+ s_threshold_bin_id = tx;
1116
+ s_num_input[0] = 0;
1117
+ s_counter = 0;
1118
+ }
1119
+ __syncthreads();
1120
+
1121
+ const auto threshold_bin = s_threshold_bin_id;
1122
+ topk -= s_histogram[threshold_bin + 1];
1123
+
1124
+ constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS;
1125
+ constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT;
1126
+
1127
+ if (topk == 0) {
1128
+ // Collect indices where bin > threshold
1129
+#pragma unroll 2
1130
+ for (int base = tx * VEC_SIZE; base < aligned_length;
1131
+ base += BLOCK_SIZE * VEC_SIZE) {
1132
+ score_vec.cast_load(&score[base]);
1133
+#pragma unroll
1134
+ for (int j = 0; j < VEC_SIZE; ++j) {
1135
+ const auto bin = static_cast<int>(Traits::ToCoarseKey(score_vec[j]));
1136
+ if (bin > threshold_bin) {
1137
+ const auto pos = atomicAdd(&s_counter, 1);
1138
+ s_indices[pos] = base + j;
1139
+ }
1140
+ }
1141
+ }
1142
+ // Handle tail
1143
+ for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) {
1144
+ const auto bin = static_cast<int>(Traits::ToCoarseKey(score[i]));
1145
+ if (bin > threshold_bin) {
1146
+ const auto pos = atomicAdd(&s_counter, 1);
1147
+ s_indices[pos] = i;
1148
+ }
1149
+ }
1150
+ __syncthreads();
1151
+ } else {
1152
+ __syncthreads();
1153
+ if (tx < RADIX + 1) s_histogram[tx] = 0;
1154
+ __syncthreads();
1155
+
1156
+ // Filter + histogram for refinement
1157
+ auto filter_and_add_to_histogram = [&](auto raw_input, int index) {
1158
+ const auto bin = static_cast<int>(Traits::ToCoarseKey(raw_input));
1159
+ if (bin > threshold_bin) {
1160
+ const auto pos = atomicAdd(&s_counter, 1);
1161
+ s_indices[pos] = index;
1162
+ } else if (bin == threshold_bin) {
1163
+ const auto pos = atomicAdd(&s_num_input[0], 1);
1164
+ if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) {
1165
+ s_input_idx[0][pos] = index;
1166
+ const auto ordered = Traits::ToOrdered(raw_input);
1167
+ const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF;
1168
+ atomicAdd(&s_histogram[sub_bin], 1);
1169
+ }
1170
+ }
1171
+ };
1172
+#pragma unroll 2
1173
+ for (int base = tx * VEC_SIZE; base < aligned_length;
1174
+ base += BLOCK_SIZE * VEC_SIZE) {
1175
+ score_vec.cast_load(&score[base]);
1176
+#pragma unroll
1177
+ for (int j = 0; j < VEC_SIZE; ++j) {
1178
+ filter_and_add_to_histogram(score_vec[j], base + j);
1179
+ }
1180
+ }
1181
+ // Handle tail
1182
+ for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) {
1183
+ filter_and_add_to_histogram(score[i], i);
1184
+ }
1185
+ __syncthreads();
1186
+
1187
+ // Stage 2: refine with 8bit radix passes
1188
+#pragma unroll
1189
+ for (int round = 0; round < NUM_ROUNDS; ++round) {
1190
+ __shared__ int s_last_remain;
1191
+ const auto r_idx = round % 2;
1192
+
1193
+ const auto _raw_num_input = s_num_input[r_idx];
1194
+ const auto num_input =
1195
+ (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE;
1196
+
1197
+ run_cumsum();
1198
+ if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
1199
+ s_threshold_bin_id = tx;
1200
+ s_num_input[r_idx ^ 1] = 0;
1201
+ s_last_remain = topk - s_histogram[tx + 1];
1202
+ }
1203
+ __syncthreads();
1204
+
1205
+ const auto threshold = s_threshold_bin_id;
1206
+ topk -= s_histogram[threshold + 1];
1207
+
1208
+ const int offset = FIRST_SHIFT - round * 8;
1209
+ const bool is_last_round = (round == NUM_ROUNDS - 1);
1210
+
1211
+ if (topk == 0) {
1212
+ for (int i = tx; i < num_input; i += BLOCK_SIZE) {
1213
+ const auto idx = s_input_idx[r_idx][i];
1214
+ const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF;
1215
+ if (static_cast<int>(bin) > threshold) {
1216
+ const auto pos = atomicAdd(&s_counter, 1);
1217
+ s_indices[pos] = idx;
1218
+ }
1219
+ }
1220
+ __syncthreads();
1221
+ break;
1222
+ } else {
1223
+ __syncthreads();
1224
+ if (tx < RADIX + 1) s_histogram[tx] = 0;
1225
+ __syncthreads();
1226
+ for (int i = tx; i < num_input; i += BLOCK_SIZE) {
1227
+ const auto idx = s_input_idx[r_idx][i];
1228
+ const auto raw_input = score[idx];
1229
+ const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF;
1230
+ if (static_cast<int>(bin) > threshold) {
1231
+ const auto pos = atomicAdd(&s_counter, 1);
1232
+ s_indices[pos] = idx;
1233
+ } else if (static_cast<int>(bin) == threshold) {
1234
+ if (is_last_round) {
1235
+ const auto pos = atomicAdd(&s_last_remain, -1);
1236
+ if (pos > 0) {
1237
+ s_indices[top_k - pos] = idx;
1238
+ }
1239
+ } else {
1240
+ const auto pos = atomicAdd(&s_num_input[r_idx ^ 1], 1);
1241
+ if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) {
1242
+ s_input_idx[r_idx ^ 1][pos] = idx;
1243
+ const auto bin32 = Traits::ToOrdered(raw_input);
1244
+ const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF;
1245
+ atomicAdd(&s_histogram[sub_bin], 1);
1246
+ }
1247
+ }
1248
+ }
1249
+ }
1250
+ __syncthreads();
1251
+ }
1252
+ }
1253
+ }
1254
+
1255
+ // Output phase - mode-specific
1256
+#pragma unroll 2
1257
+ for (int base = tx; base < static_cast<int>(top_k); base += BLOCK_SIZE) {
1258
+ const int idx = s_indices[base];
1259
+ dst[base] = static_cast<IdType>(idx);
1260
+ }
1261
+}
1262
+
1263
+// Helper to compute GCD for VEC_SIZE selection
1264
+constexpr uint32_t gcd(uint32_t a, uint32_t b) {
1265
+ while (b != 0) {
1266
+ uint32_t t = b;
1267
+ b = a % b;
1268
+ a = t;
1269
+ }
1270
+ return a;
1271
+}
1272
+
1273
+// Compute optimal VEC_SIZE based on max_len and dtype
1274
+// Returns 1, 2, 4, or 8
1275
+template <typename DType>
1276
+constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
1277
+ constexpr int MAX_VEC = 16 / sizeof(DType); // 4 for float32, 8 for fp16/bf16
1278
+ // Use GCD to find largest power-of-2 divisor
1279
+ const uint32_t g = gcd(max_len, static_cast<uint32_t>(MAX_VEC));
1280
+ return static_cast<int>(g);
1281
+}
1282
+
1283
+template <typename DType, typename IdType>
1284
+cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
1285
+ IdType* lengths, uint32_t num_rows,
1286
+ uint32_t top_k_val, uint32_t max_len,
1287
+ cudaStream_t stream = 0) {
1288
+ constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC;
1289
+ constexpr int MAX_VEC = 16 / sizeof(DType);
1290
+
1291
+ dim3 grid(num_rows);
1292
+ dim3 block(FILTERED_TOPK_BLOCK_THREADS);
1293
+ void* args[] = {&input, &output_indices, &lengths,
1294
+ &num_rows, &top_k_val, &max_len};
1295
+
1296
+ const int vec_size = ComputeFilteredTopKVecSize<DType>(max_len);
1297
+
1298
+#define DISPATCH_VEC_SIZE(VS) \
1299
+ if (vec_size == VS) { \
1300
+ auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>; \
1301
+ FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
1302
+ kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
1303
+ FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
1304
+ smem_size, stream)); \
1305
+ return cudaSuccess; \
1306
+ }
1307
+
1308
+ DISPATCH_VEC_SIZE(1)
1309
+ DISPATCH_VEC_SIZE(2)
1310
+ DISPATCH_VEC_SIZE(4)
1311
+ if constexpr (MAX_VEC >= 8) {
1312
+ DISPATCH_VEC_SIZE(8)
1313
+ }
1314
+#undef DISPATCH_VEC_SIZE
1315
+
1316
+ return cudaSuccess;
1317
+}
1318
+
1319
+} // namespace vllm
1320
+
1321
+#endif // PERSISTENT_TOPK_CUH_
csrc/topk.cu
+139
-358
@@ -1,373 +1,154 @@
1
-// Portions of this file are adapted from SGLang PR:
2
-// https://github.com/sgl-project/sglang/pull/11194
3
-// and
4
-// https://github.com/sgl-project/sglang/pull/17747
1
+// Persistent TopK kernel for DeepSeek V3 sparse attention indexer.
2
+// See persistent_topk.cuh for kernel implementation.
3
6
-#include "cuda_compat.h"
7
-#include "dispatch_utils.h"
8
-
9
-#include <torch/cuda.h>
10
-#include <c10/cuda/CUDAGuard.h>
4
+#include <torch/all.h>
5
+#include <ATen/cuda/CUDAContext.h>
6
+#include <cuda_runtime.h>
7
+#include <algorithm>
8
9
#ifndef USE_ROCM
13
- #include <cub/cub.cuh>
14
-#else
15
- #include <hipcub/hipcub.hpp>
10
+ #include "persistent_topk.cuh"
11
#endif
12
18
-namespace vllm {
19
-
20
-constexpr int TopK = 2048; // DeepSeek V3 sparse attention top-k
21
-constexpr int kThreadsPerBlock = 1024; // Threads per block
22
-
23
-// Shared memory budget
24
-#if defined(USE_ROCM)
25
-constexpr size_t kSmem = 48 * 1024; // ROCm default: 48KB
26
-#else
27
-// Reduced from 128KB to 32KB to improve occupancy.
28
-// Each radix pass needs at most ~TopK candidates in the threshold bin,
29
-// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
30
-constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t); // 32KB (bytes)
31
-#endif
32
-
33
-struct FastTopKParams {
34
- const float* __restrict__ input; // [batch, seq_len] Logits
35
- const int32_t* __restrict__ row_starts; // [batch] Offset into each row
36
- // (optional)
37
- int32_t* __restrict__ indices; // [batch, TopK] Output top-k indices
38
- int32_t* __restrict__ lengths; // [batch] Sequence lengths per row
39
- int64_t input_stride; // Stride between rows
40
-};
41
-
42
-__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
43
- uint32_t bits = __float_as_uint(x);
44
- return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
45
-}
46
-
47
-__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
48
- __half h = __float2half_rn(x);
49
- uint16_t bits = __half_as_ushort(h);
50
- uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
51
- : static_cast<uint16_t>(bits | 0x8000);
52
- return static_cast<uint8_t>(key >> 8);
53
-}
54
-
55
-__device__ void naive_topk_cuda(const float* __restrict__ logits,
56
- int32_t* __restrict__ output_indices,
57
- int32_t seq_len) {
58
- const int thread_id = threadIdx.x;
59
- for (int i = thread_id; i < TopK; i += kThreadsPerBlock) {
60
- output_indices[i] = (i < seq_len) ? i : -1;
61
- }
62
-}
63
-
64
-// Adapted from:
65
-// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87
66
-// by: DarkSharpness
67
-// which at the same time is an optimized topk kernel copied from tilelang
68
-// kernel
69
-__device__ void fast_topk_cuda_tl(
70
- const float* __restrict__ logits, // Input logits [seq_len]
71
- int* __restrict__ output_indices, // Output top-k indices [TopK]
72
- int logits_offset, // Starting offset in logits array
73
- int seq_len) // Number of valid logits to process
74
-{
75
- constexpr int RADIX = 256;
76
- constexpr int MAX_BUFFERED_ITEMS = kSmem / (2 * sizeof(int));
77
-
78
- alignas(128) __shared__ int shared_histogram[2][RADIX + 128];
79
- alignas(128) __shared__ int shared_output_count;
80
- alignas(128) __shared__ int shared_threshold_bin;
81
- alignas(128) __shared__ int shared_buffered_count[2];
82
-
83
- extern __shared__ int buffered_indices[][MAX_BUFFERED_ITEMS];
84
-
85
- const int thread_id = threadIdx.x;
86
- int remaining_k = TopK;
87
-
88
- // Pass 0: Build coarse 8-bit histogram using FP16 high bits
89
- if (thread_id < RADIX + 1) {
90
- shared_histogram[0][thread_id] = 0;
91
- }
92
- __syncthreads();
93
-
94
- for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
95
- const auto bin = convert_to_uint8(logits[idx + logits_offset]);
96
- ::atomicAdd(&shared_histogram[0][bin], 1);
97
- }
98
- __syncthreads();
99
-
100
- // Helper: Compute cumulative sum (suffix sum) over histogram using ping-pong
101
- // buffers
102
- auto compute_cumulative_sum = [&]() {
103
- static_assert(1 << 8 == RADIX,
104
- "Radix must be 256 for 8 unrolled iterations");
105
-#pragma unroll 8
106
- for (int i = 0; i < 8; ++i) {
107
- if (C10_LIKELY(thread_id < RADIX)) {
108
- const int stride = 1 << i;
109
- const int src_buffer = i & 1;
110
- const int dst_buffer = src_buffer ^ 1;
111
-
112
- int value = shared_histogram[src_buffer][thread_id];
113
- if (thread_id < RADIX - stride) {
114
- value += shared_histogram[src_buffer][thread_id + stride];
115
- }
116
- shared_histogram[dst_buffer][thread_id] = value;
117
- }
118
- __syncthreads();
119
- }
120
- };
121
-
122
- compute_cumulative_sum();
123
-
124
- // Find threshold bin where cumsum crosses remaining_k
125
- if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
126
- shared_histogram[0][thread_id + 1] <= remaining_k) {
127
- shared_threshold_bin = thread_id;
128
- shared_buffered_count[0] = 0;
129
- shared_output_count = 0;
130
- }
131
- __syncthreads();
132
-
133
- const int threshold_bin = shared_threshold_bin;
134
- remaining_k -= shared_histogram[0][threshold_bin + 1];
135
-
136
- // Early exit if threshold bin perfectly matches remaining_k
137
- if (remaining_k == 0) {
138
- for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
139
- const int bin = convert_to_uint8(logits[idx + logits_offset]);
140
- if (bin > threshold_bin) {
141
- const int output_pos = ::atomicAdd(&shared_output_count, 1);
142
- output_indices[output_pos] = idx;
143
- }
144
- }
145
- __syncthreads();
146
- return;
147
- }
148
-
149
- // Prepare for refinement passes: Process threshold bin
150
- __syncthreads();
151
- if (thread_id < RADIX + 1) {
152
- shared_histogram[0][thread_id] = 0;
153
- }
154
- __syncthreads();
155
-
156
- // Scan all elements and:
157
- // 1. Write indices > threshold_bin to output
158
- // 2. Buffer indices == threshold_bin for refinement
159
- // 3. Build histogram for next refinement pass (fused optimization)
160
- for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
161
- const float logit_value = logits[idx + logits_offset];
162
- const int bin = convert_to_uint8(logit_value);
163
-
164
- if (bin > threshold_bin) {
165
- // in top-k, write to output
166
- const int output_pos = ::atomicAdd(&shared_output_count, 1);
167
- output_indices[output_pos] = idx;
168
- } else if (bin == threshold_bin) {
169
- // Candidate for top-k, needs refinement
170
- const int buffer_pos = ::atomicAdd(&shared_buffered_count[0], 1);
171
- if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
172
- buffered_indices[0][buffer_pos] = idx;
173
- // Fused: Build histogram for next pass
174
- const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
175
- const int next_bin = (fp32_bits >> 24) & 0xFF;
176
- ::atomicAdd(&shared_histogram[0][next_bin], 1);
177
- }
178
- }
13
+void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
14
+ torch::Tensor& output, torch::Tensor& workspace, int64_t k,
15
+ int64_t max_seq_len) {
16
+#ifndef USE_ROCM
17
+ TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
18
+ TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
19
+ TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
20
+ TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
21
+ TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
22
+ TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
23
+ TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
24
+ TORCH_CHECK(lengths.dim() == 1, "lengths must be 1D");
25
+ TORCH_CHECK(output.dim() == 2, "output must be 2D");
26
+
27
+ const int64_t num_rows = logits.size(0);
28
+ const int64_t stride = logits.size(1);
29
+
30
+ TORCH_CHECK(lengths.size(0) == num_rows, "lengths size mismatch");
31
+ TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
32
+ "output size mismatch");
33
+ namespace P = vllm::persistent;
34
+
35
+ TORCH_CHECK(k == P::TopK, "k must be 2048");
36
+ TORCH_CHECK(k <= stride, "k out of range");
37
+
38
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream();
39
+
40
+ static int num_sms = 0;
41
+ static int max_smem_per_block = 0;
42
+ if (num_sms == 0) {
43
+ int device;
44
+ cudaGetDevice(&device);
45
+ cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device);
46
+ cudaDeviceGetAttribute(&max_smem_per_block,
47
+ cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
48
}
180
- __syncthreads();
181
-
182
- // ============================================================================
183
- // Passes 1-4: Refine using 8-bit passes over FP32 bits
184
- // ============================================================================
185
- // FP32 bits [31:0] split into 4 bytes processed MSB-first:
186
- // Pass 1: bits [31:24], Pass 2: bits [23:16], Pass 3: bits [15:8], Pass 4:
187
- // bits [7:0]
188
-#pragma unroll 4
189
- for (int pass = 0; pass < 4; ++pass) {
190
- __shared__ int shared_final_k; // For final pass: remaining slots to fill
191
- const int src_buffer = pass % 2;
192
- const int dst_buffer = src_buffer ^ 1;
193
-
194
- // Clamp buffered count to prevent overflow
195
- const int raw_buffered = shared_buffered_count[src_buffer];
196
- const int num_buffered =
197
- (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS;
198
-
199
- compute_cumulative_sum();
200
-
201
- // Find threshold bin for this pass
202
- if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
203
- shared_histogram[0][thread_id + 1] <= remaining_k) {
204
- shared_threshold_bin = thread_id;
205
- shared_buffered_count[dst_buffer] = 0;
206
- shared_final_k = remaining_k - shared_histogram[0][thread_id + 1];
207
- }
208
- __syncthreads();
209
-
210
- const int threshold_bin = shared_threshold_bin;
211
- remaining_k -= shared_histogram[0][threshold_bin + 1];
212
-
213
- // Bit offset for this pass: 24, 16, 8, 0
214
- const int bit_offset = 24 - pass * 8;
49
216
- // Early exit if threshold bin perfectly matches
217
- if (remaining_k == 0) {
218
- for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
219
- const int idx = buffered_indices[src_buffer][i];
220
- const uint32_t fp32_bits =
221
- convert_to_uint32_v2(logits[idx + logits_offset]);
222
- const int bin = (fp32_bits >> bit_offset) & 0xFF;
223
- if (bin > threshold_bin) {
224
- const int output_pos = ::atomicAdd(&shared_output_count, 1);
225
- output_indices[output_pos] = idx;
226
- }
227
- }
228
- __syncthreads();
229
- break;
50
+ if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
51
+ cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
52
+ logits.data_ptr<float>(), output.data_ptr<int32_t>(),
53
+ lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
54
+ static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
55
+ TORCH_CHECK(status == cudaSuccess,
56
+ "FilteredTopK failed: ", cudaGetErrorString(status));
57
+ } else {
58
+ TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
59
+ TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
60
+
61
+ // Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
62
+ // large path. Empirically tuned.
63
+ int effective_max_smem;
64
+ if (num_rows <= 4) {
65
+ effective_max_smem =
66
+ std::min(max_smem_per_block, static_cast<int>(P::kSmemMedium));
67
+ } else if (num_rows <= 8) {
68
+ constexpr int kSmemCapMedium = 48 * 1024;
69
+ effective_max_smem = std::min(max_smem_per_block, kSmemCapMedium);
70
+ } else {
71
+ effective_max_smem = max_smem_per_block;
72
}
73
232
- // Continue refinement
233
- __syncthreads();
234
- if (thread_id < RADIX + 1) {
235
- shared_histogram[0][thread_id] = 0;
74
+ size_t available_for_ordered =
75
+ static_cast<size_t>(effective_max_smem) - P::kFixedSmemLarge;
76
+ uint32_t max_chunk_elements =
77
+ static_cast<uint32_t>(available_for_ordered / sizeof(uint32_t));
78
+
79
+ uint32_t vec_size = 1;
80
+ if (stride % 4 == 0)
81
+ vec_size = 4;
82
+ else if (stride % 2 == 0)
83
+ vec_size = 2;
84
+
85
+ max_chunk_elements = (max_chunk_elements / vec_size) * vec_size;
86
+ uint32_t min_chunk = vec_size * P::kThreadsPerBlock;
87
+ if (max_chunk_elements < min_chunk) max_chunk_elements = min_chunk;
88
+
89
+ uint32_t ctas_per_group =
90
+ (static_cast<uint32_t>(stride) + max_chunk_elements - 1) /
91
+ max_chunk_elements;
92
+ uint32_t chunk_size =
93
+ (static_cast<uint32_t>(stride) + ctas_per_group - 1) / ctas_per_group;
94
+ chunk_size = ((chunk_size + vec_size - 1) / vec_size) * vec_size;
95
+ if (chunk_size > max_chunk_elements) chunk_size = max_chunk_elements;
96
+
97
+ size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
98
+ if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
99
+
100
+ int occupancy = 1;
101
+ cudaOccupancyMaxActiveBlocksPerMultiprocessor(
102
+ &occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
103
+ smem_size);
104
+ if (occupancy < 1) occupancy = 1;
105
+
106
+ uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
107
+ uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
108
+ static_cast<uint32_t>(num_rows));
109
+ if (num_groups == 0) num_groups = 1;
110
+ uint32_t total_ctas = num_groups * ctas_per_group;
111
+
112
+ size_t state_bytes = num_groups * sizeof(P::RadixRowState);
113
+ TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
114
+ "workspace too small, need ", state_bytes, " bytes");
115
+
116
+ P::PersistentTopKParams params;
117
+ params.input = logits.data_ptr<float>();
118
+ params.output = output.data_ptr<int32_t>();
119
+ params.lengths = lengths.data_ptr<int32_t>();
120
+ params.num_rows = static_cast<uint32_t>(num_rows);
121
+ params.stride = static_cast<uint32_t>(stride);
122
+ params.chunk_size = chunk_size;
123
+ params.row_states =
124
+ reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
125
+ params.ctas_per_group = ctas_per_group;
126
+ params.max_seq_len = static_cast<uint32_t>(max_seq_len);
127
+
128
+ #define LAUNCH_PERSISTENT(VS) \
129
+ do { \
130
+ auto kernel = &P::persistent_topk_kernel<VS>; \
131
+ cudaError_t err = cudaFuncSetAttribute( \
132
+ kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
133
+ TORCH_CHECK(err == cudaSuccess, \
134
+ "Failed to set smem: ", cudaGetErrorString(err)); \
135
+ kernel<<<total_ctas, P::kThreadsPerBlock, smem_size, stream>>>(params); \
136
+ } while (0)
137
+
138
+ if (vec_size == 4) {
139
+ LAUNCH_PERSISTENT(4);
140
+ } else if (vec_size == 2) {
141
+ LAUNCH_PERSISTENT(2);
142
+ } else {
143
+ LAUNCH_PERSISTENT(1);
144
}
237
- __syncthreads();
238
-
239
- for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
240
- const int idx = buffered_indices[src_buffer][i];
241
- const float logit_value = logits[idx + logits_offset];
242
- const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
243
- const int bin = (fp32_bits >> bit_offset) & 0xFF;
244
-
245
- if (bin > threshold_bin) {
246
- // Definitely in top-k
247
- const int output_pos = ::atomicAdd(&shared_output_count, 1);
248
- output_indices[output_pos] = idx;
249
- } else if (bin == threshold_bin) {
250
- if (pass == 3) {
251
- // Final pass (bits [7:0]): No more refinement possible
252
- // Fill remaining slots in reverse order to maintain descending order
253
- const int slot = ::atomicAdd(&shared_final_k, -1);
254
- if (slot > 0) {
255
- output_indices[TopK - slot] = idx;
256
- }
257
- } else {
258
- // Buffer for next pass and build next histogram
259
- const int buffer_pos =
260
- ::atomicAdd(&shared_buffered_count[dst_buffer], 1);
261
- if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
262
- buffered_indices[dst_buffer][buffer_pos] = idx;
263
- // Fused: Build histogram for next pass
264
- const int next_bit_offset = bit_offset - 8;
265
- const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF;
266
- ::atomicAdd(&shared_histogram[0][next_bin], 1);
267
- }
268
- }
269
- }
270
- }
271
- __syncthreads();
145
+ #undef LAUNCH_PERSISTENT
146
}
273
-}
147
275
-__global__ __launch_bounds__(kThreadsPerBlock) void topk_kernel(
276
- const FastTopKParams params) {
277
- const auto& [input, row_starts, indices, lengths, input_stride] = params;
278
- const uint64_t batch_idx = blockIdx.x;
279
- const int logits_offset = row_starts == nullptr ? 0 : row_starts[batch_idx];
280
- const int seq_len = lengths[batch_idx];
281
- int* output_indices = indices + batch_idx * TopK;
282
- const float* logits = input + batch_idx * input_stride;
283
-
284
- if (seq_len <= TopK) {
285
- // Shortcut: All elements are in top-k
286
- return naive_topk_cuda(logits, output_indices, seq_len);
287
- } else {
288
- return fast_topk_cuda_tl(logits, output_indices, logits_offset, seq_len);
289
- }
290
-}
291
-
292
-FastTopKParams get_params(
293
- const at::Tensor& score, const at::Tensor& lengths,
294
- std::optional<at::Tensor> row_starts_opt = std::nullopt,
295
- std::optional<at::Tensor> indices_opt = std::nullopt) {
296
- const int64_t batch_size = score.size(0);
297
-
298
- TORCH_CHECK(score.dim() == 2 && score.stride(1) == 1,
299
- "score must be 2D with contiguous rows");
300
- TORCH_CHECK(lengths.dim() == 1 && lengths.is_contiguous() &&
301
- lengths.size(0) == batch_size,
302
- "lengths must be 1D contiguous with size matching batch");
303
-
304
- const int32_t* row_starts_ptr = nullptr;
305
- if (row_starts_opt.has_value()) {
306
- const auto& row_starts = *row_starts_opt;
307
- TORCH_CHECK(row_starts.dim() == 1 && row_starts.size(0) == batch_size,
308
- "row_starts must be 1D with size matching batch");
309
- row_starts_ptr = row_starts.data_ptr<int32_t>();
310
- }
311
-
312
- int32_t* indices_ptr = nullptr;
313
- if (indices_opt.has_value()) {
314
- const auto& indices = *indices_opt;
315
- TORCH_CHECK(indices.dim() == 2 && indices.is_contiguous() &&
316
- indices.size(0) == batch_size && indices.size(1) == TopK,
317
- "indices must be 2D contiguous [batch, TopK]");
318
- indices_ptr = indices.data_ptr<int32_t>();
319
- }
320
-
321
- return FastTopKParams{
322
- .input = score.data_ptr<float>(),
323
- .row_starts = row_starts_ptr,
324
- .indices = indices_ptr,
325
- .lengths = lengths.data_ptr<int32_t>(),
326
- .input_stride = score.stride(0),
327
- };
328
-}
329
-
330
-template <auto* kernel_func, size_t smem_bytes>
331
-void setup_kernel_smem_once() {
332
- static const cudaError_t result = []() -> cudaError_t {
333
-#ifdef USE_ROCM
334
- auto func_ptr = reinterpret_cast<const void*>(kernel_func);
148
+ cudaError_t err = cudaGetLastError();
149
+ TORCH_CHECK(err == cudaSuccess,
150
+ "persistent_topk failed: ", cudaGetErrorString(err));
151
#else
336
- auto func_ptr = kernel_func;
152
+ TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
153
#endif
338
- return cudaFuncSetAttribute(
339
- func_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
340
- }();
341
-
342
- TORCH_CHECK(
343
- result == cudaSuccess,
344
- "Failed to set kernel shared memory limit: ", cudaGetErrorString(result));
154
}
346
-
347
-} // namespace vllm
348
-
349
-void large_context_topk(
350
- const torch::Tensor& logits, torch::Tensor& indices,
351
- const torch::Tensor& seq_lens,
352
- std::optional<torch::Tensor> row_starts = std::nullopt) {
353
- TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
354
- TORCH_CHECK(indices.is_cuda(), "indices must be a CUDA tensor");
355
- TORCH_CHECK(seq_lens.is_cuda(), "seq_lens must be a CUDA tensor");
356
- if (row_starts.has_value()) {
357
- TORCH_CHECK(row_starts->is_cuda(), "row_starts must be a CUDA tensor");
358
- }
359
-
360
- const auto params = vllm::get_params(logits, seq_lens, row_starts, indices);
361
- const int64_t batch_size = logits.size(0);
362
-
363
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
364
- const dim3 grid(static_cast<uint32_t>(batch_size));
365
- const dim3 block(vllm::kThreadsPerBlock);
366
-
367
- vllm::setup_kernel_smem_once<vllm::topk_kernel, vllm::kSmem>();
368
- vllm::topk_kernel<<<grid, block, vllm::kSmem, stream>>>(params);
369
-
370
- const cudaError_t result = cudaGetLastError();
371
- TORCH_CHECK(result == cudaSuccess,
372
- "large_context_topk kernel failed: ", cudaGetErrorString(result));
373
-}
\ No newline at end of file
csrc/torch_bindings.cpp
+3
-4
@@ -197,10 +197,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
197
ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
198
199
ops.def(
200
- "large_context_topk(Tensor score, Tensor indices, Tensor lengths, "
201
- "Tensor? "
202
- "row_starts_opt) -> ()");
203
- ops.impl("large_context_topk", torch::kCUDA, &large_context_topk);
200
+ "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
201
+ "Tensor workspace, int k, int max_seq_len) -> ()");
202
+ ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
203
204
// Layernorm-quant
205
// Apply Root Mean Square (RMS) Normalization to the input tensor.
tests/kernels/test_top_k_per_row.py
+540
-78
@@ -122,6 +122,39 @@ def compare_top_k_results(
122
return True
123
124
125
+def validate_topk_against_reference(
126
+ logits: torch.Tensor,
127
+ cuda_indices: torch.Tensor,
128
+ row_starts: torch.Tensor,
129
+ row_ends: torch.Tensor,
130
+ top_k: int,
131
+ kernel_name: str,
132
+) -> None:
133
+ """
134
+ Validate CUDA top-k results against PyTorch reference implementation.
135
+
136
+ Args:
137
+ logits: Input logits tensor
138
+ cuda_indices: CUDA kernel output indices
139
+ row_starts: Row start positions
140
+ row_ends: Row end positions
141
+ top_k: Number of top elements to select
142
+ kernel_name: Name of the kernel being tested (for error messages)
143
+ """
144
+ num_rows = cuda_indices.shape[0]
145
+ torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
146
+
147
+ for i in range(num_rows):
148
+ row_end = int(row_ends[i])
149
+ k_i = min(top_k, row_end)
150
+ idx = logits[i, :row_end].topk(k_i, dim=-1)[1]
151
+ torch_indices[i, :k_i] = idx
152
+
153
+ assert compare_top_k_results(
154
+ logits, cuda_indices, torch_indices, row_starts, row_ends, top_k
155
+ ), f"{kernel_name} results don't match torch.topk"
156
+
157
+
158
@pytest.mark.parametrize("num_rows", NUM_ROWS)
159
@pytest.mark.parametrize("top_k", TOP_K_VALUES)
160
@pytest.mark.parametrize("clean_logits", [True, False])
@@ -278,111 +311,540 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
311
312
313
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
314
+@pytest.mark.parametrize(
315
+ "seq_len_range,test_id",
316
+ [
317
+ pytest.param((4000, 8000), "short_sequences", id="short"),
318
+ pytest.param((8000, 32000), "medium_sequences", id="medium"),
319
+ pytest.param((32000, 163840), "long_sequences", id="long"),
320
+ ],
321
+)
322
@pytest.mark.parametrize("clean_logits", [True, False])
323
+@pytest.mark.parametrize("top_k", [2048])
324
+@pytest.mark.parametrize("next_n", [1, 4])
325
@torch.inference_mode()
283
-def test_deepseek_hybrid_topk(clean_logits: bool) -> None:
326
+def test_deepseek_persistent_topk(
327
+ seq_len_range: tuple[int, int],
328
+ test_id: str,
329
+ clean_logits: bool,
330
+ top_k: int,
331
+ next_n: int,
332
+) -> None:
333
+ """
334
+ Test persistent_topk with varying sequence lengths and speculative decoding.
335
+ Supports speculative decoding with next_n > 1.
336
+ """
337
+ set_random_seed(42 if test_id == "short_sequences" else 43)
338
torch.set_default_device("cuda:0")
339
286
- top_k = 2048
287
-
288
- # Test case 1: Short sequences (< 8192)
289
- batch_size_short = 4
290
- next_n = 1
291
- num_rows_short = batch_size_short * next_n
340
+ batch_size = 4
341
+ num_rows = batch_size * next_n
342
293
- # Create sequences with max length < 8192
294
- seq_lens_short = torch.randint(
295
- 4000, 8000, (batch_size_short,), dtype=torch.int32, device="cuda"
343
+ seq_lens = torch.randint(
344
+ seq_len_range[0],
345
+ seq_len_range[1],
346
+ (batch_size,),
347
+ dtype=torch.int32,
348
+ device="cuda",
349
)
350
298
- row_starts_short = torch.zeros(num_rows_short, dtype=torch.int32, device="cuda")
299
- row_indices_short = torch.arange(num_rows_short, device="cuda") // next_n
300
- next_n_offset_short = torch.arange(num_rows_short, device="cuda") % next_n
301
- row_ends_short = (
302
- seq_lens_short[row_indices_short] - next_n + next_n_offset_short + 1
351
+ # Compute row boundaries for speculative decoding
352
+ row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
353
+ row_indices = torch.arange(num_rows, device="cuda") // next_n
354
+ next_n_offset = torch.arange(num_rows, device="cuda") % next_n
355
+ row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1
356
+
357
+ logits = create_random_logits(
358
+ row_starts, row_ends, torch.float32, 42, clean_logits, "random"
359
)
360
305
- logits_short = create_random_logits(
306
- row_starts_short, row_ends_short, torch.float32, 42, clean_logits, "random"
361
+ indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
362
+
363
+ if next_n == 1:
364
+ lengths = seq_lens
365
+ else:
366
+ offsets = torch.arange(next_n, device=logits.device, dtype=torch.int32)
367
+ lengths = (seq_lens.unsqueeze(1) - next_n + 1 + offsets).flatten()
368
+
369
+ workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
370
+ max_seq_len = int(seq_lens.max().item())
371
+ torch.ops._C.persistent_topk(
372
+ logits, lengths, indices, workspace, top_k, max_seq_len
373
)
374
309
- indices_vllm = torch.empty(
310
- (num_rows_short, top_k), dtype=torch.int32, device="cuda"
375
+ validate_topk_against_reference(
376
+ logits, indices, row_starts, row_ends, top_k, f"persistent_topk ({test_id})"
377
)
378
313
- # Use vllm's kernel for short sequences
314
- torch.ops._C.top_k_per_row_decode(
315
- logits_short,
316
- next_n,
317
- seq_lens_short,
318
- indices_vllm,
319
- num_rows_short,
320
- logits_short.stride(0),
321
- logits_short.stride(1),
322
- top_k,
379
+
380
+def run_large_context_topk_test(
381
+ batch_size: int,
382
+ seq_lens: list[int],
383
+ top_k: int,
384
+ data_type: str = "random",
385
+ seed: int = 42,
386
+) -> None:
387
+ """
388
+ Helper to run persistent_topk kernel test with given parameters.
389
+
390
+ Args:
391
+ batch_size: Number of rows/sequences
392
+ seq_lens: List of sequence lengths (one per row)
393
+ top_k: Number of top elements to select
394
+ data_type: Type of test data to generate
395
+ seed: Random seed for reproducibility
396
+ """
397
+ torch.set_default_device("cuda:0")
398
+ set_random_seed(seed)
399
+
400
+ # Create test data
401
+ num_rows = batch_size
402
+ max_len = max(seq_lens)
403
+ lengths = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
404
+
405
+ if data_type == "random":
406
+ logits = torch.randn(num_rows, max_len, dtype=torch.float32, device="cuda")
407
+ elif data_type == "sorted_asc":
408
+ # Each row gets its own ascending sequence based on its length
409
+ logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda")
410
+ for i, length in enumerate(seq_lens):
411
+ logits[i, :length] = torch.arange(
412
+ length, dtype=torch.float32, device="cuda"
413
+ )
414
+ if length < max_len:
415
+ logits[i, length:] = float("-inf")
416
+ elif data_type == "sorted_desc":
417
+ # Each row gets its own descending sequence based on its length
418
+ logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda")
419
+ for i, length in enumerate(seq_lens):
420
+ logits[i, :length] = torch.arange(
421
+ length, 0, -1, dtype=torch.float32, device="cuda"
422
+ )
423
+ if length < max_len:
424
+ logits[i, length:] = float("-inf")
425
+ elif data_type == "all_same":
426
+ logits = torch.ones(num_rows, max_len, dtype=torch.float32, device="cuda")
427
+ for i, length in enumerate(seq_lens):
428
+ if length < max_len:
429
+ logits[i, length:] = float("-inf")
430
+ elif data_type == "many_ties":
431
+ # Only 10 unique values, many duplicates
432
+ logits = torch.randint(0, 10, (num_rows, max_len), device="cuda").float() / 10.0
433
+ for i, length in enumerate(seq_lens):
434
+ if length < max_len:
435
+ logits[i, length:] = float("-inf")
436
+ elif data_type == "small_differences":
437
+ # Very small differences to test float precision
438
+ base = torch.randn(num_rows, max_len, dtype=torch.float32, device="cuda")
439
+ noise = (
440
+ torch.randn(num_rows, max_len, dtype=torch.float32, device="cuda") * 1e-6
441
+ )
442
+ logits = base + noise
443
+ for i, length in enumerate(seq_lens):
444
+ if length < max_len:
445
+ logits[i, length:] = float("-inf")
446
+ else:
447
+ raise ValueError(f"Unknown data_type: {data_type}")
448
+
449
+ # Create output tensor
450
+ indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
451
+
452
+ workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
453
+ max_seq_len = max(seq_lens)
454
+ torch.ops._C.persistent_topk(
455
+ logits, lengths, indices, workspace, top_k, max_seq_len
456
)
457
325
- # Test case 2: Long sequences (>= 8192) - should use large_context_topk kernel
326
- batch_size_long = 4
327
- num_rows_long = batch_size_long * next_n
458
+ torch.accelerator.synchronize()
459
+
460
+ torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
461
+ for i in range(num_rows):
462
+ length = seq_lens[i]
463
+ k_i = min(top_k, length)
464
+ if k_i > 0:
465
+ idx = logits[i, :length].topk(k_i, dim=-1)[1]
466
+ torch_indices[i, :k_i] = idx
467
+ if k_i < top_k:
468
+ torch_indices[i, k_i:] = -1
469
+ else:
470
+ torch_indices[i, :] = -1
471
+
472
+ # Compare results
473
+ for i in range(num_rows):
474
+ length = seq_lens[i]
475
+ k_i = min(top_k, length)
476
+
477
+ if k_i == 0:
478
+ continue
479
+
480
+ cuda_row = indices[i, :k_i].cpu()
481
+ torch_row = torch_indices[i, :k_i].cpu()
482
+
483
+ # Filter out -1 padding values from cuda_row
484
+ valid_mask = cuda_row >= 0
485
+ cuda_row = cuda_row[valid_mask]
486
+
487
+ # Compare sets (order may differ for ties)
488
+ cuda_set = set(cuda_row.tolist())
489
+ torch_set = set(torch_row.tolist())
490
+
491
+ if cuda_set == torch_set:
492
+ continue
493
+
494
+ # If sets differ, check if it's due to equal values (ties)
495
+ cuda_vals = logits[i, cuda_row].cpu()
496
+ torch_vals = logits[i, torch_row].cpu()
497
+
498
+ # Check that min CUDA value >= max of values NOT in top-k
499
+ if k_i < length:
500
+ non_topk_indices = torch.tensor(
501
+ list(set(range(length)) - cuda_set), dtype=torch.int32
502
+ )
503
+ if len(non_topk_indices) > 0:
504
+ non_topk_vals = logits[i, non_topk_indices].cpu()
505
+ min_cuda_val = cuda_vals.min()
506
+ max_non_topk = non_topk_vals.max()
507
+
508
+ # Allow small tolerance for floating point errors
509
+ assert min_cuda_val >= max_non_topk - 1e-4, (
510
+ f"Row {i}: CUDA top-k contains values smaller than non-top-k. "
511
+ f"Min CUDA: {min_cuda_val}, Max non-top-k: {max_non_topk}, "
512
+ f"Length: {length}, k: {k_i}, CUDA indices: {sorted(cuda_set)[:10]}..., " # noqa: E501
513
+ f"Expected indices: {sorted(torch_set)[:10]}..."
514
+ )
515
+
516
+ # For ties, verify the values are close
517
+ assert torch.allclose(
518
+ cuda_vals.sort(descending=True)[0],
519
+ torch_vals.sort(descending=True)[0],
520
+ rtol=1e-4,
521
+ atol=1e-4,
522
+ ), f"""Row {i}: Top-k values don't match.
523
+ CUDA: {cuda_vals.sort(descending=True)[0][:10]},
524
+ Torch: {torch_vals.sort(descending=True)[0][:10]}"""
525
329
- # Create sequences with max length >= 8192
330
- seq_lens_long = torch.randint(
331
- 8192, 16384, (batch_size_long,), dtype=torch.int32, device="cuda"
526
+
527
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
528
+@pytest.mark.parametrize(
529
+ "test_config",
530
+ [
531
+ # ==================== CATEGORY: Sequence Length Edge Cases ====================
532
+ pytest.param(
533
+ {"seq_lens": [1, 10, 100, 2048], "top_k": 2048, "data_type": "random"},
534
+ id="seq_len_edge_very_small_to_medium",
535
+ ),
536
+ pytest.param(
537
+ {
538
+ "seq_lens": [2049, 2100, 2500, 3000],
539
+ "top_k": 2048,
540
+ "data_type": "random",
541
+ },
542
+ id="seq_len_edge_above_k",
543
+ ),
544
+ pytest.param(
545
+ {"seq_lens": [8000, 16384, 20000], "top_k": 2048, "data_type": "random"},
546
+ id="algo_transition_filtered_radix",
547
+ ),
548
+ # ==================== CATEGORY: Data Distributions ====================
549
+ pytest.param(
550
+ {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "sorted_asc"},
551
+ id="data_sorted_ascending",
552
+ ),
553
+ pytest.param(
554
+ {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "sorted_desc"},
555
+ id="data_sorted_descending",
556
+ ),
557
+ pytest.param(
558
+ {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "all_same"},
559
+ id="data_all_same",
560
+ ),
561
+ pytest.param(
562
+ {"seq_lens": [5000, 10000], "top_k": 2048, "data_type": "many_ties"},
563
+ id="data_many_ties",
564
+ ),
565
+ pytest.param(
566
+ {
567
+ "seq_lens": [5000, 10000],
568
+ "top_k": 2048,
569
+ "data_type": "small_differences",
570
+ },
571
+ id="data_float_precision",
572
+ ),
573
+ # ==================== CATEGORY: Alignment / Vectorization ====================
574
+ pytest.param(
575
+ {
576
+ "seq_lens": [2055, 2056, 2057, 2063],
577
+ "top_k": 2048,
578
+ "data_type": "random",
579
+ },
580
+ id="align_vec_boundaries_low",
581
+ ),
582
+ pytest.param(
583
+ {
584
+ "seq_lens": [4095, 4096, 4097, 4102],
585
+ "top_k": 2048,
586
+ "data_type": "random",
587
+ },
588
+ id="align_4k_boundary",
589
+ ),
590
+ pytest.param(
591
+ {
592
+ "seq_lens": [8191, 8192, 8193, 8198],
593
+ "top_k": 2048,
594
+ "data_type": "random",
595
+ },
596
+ id="align_8k_boundary",
597
+ ),
598
+ pytest.param(
599
+ {
600
+ "seq_lens": [16383, 16384, 16385, 16390],
601
+ "top_k": 2048,
602
+ "data_type": "random",
603
+ },
604
+ id="align_16k_boundary",
605
+ ),
606
+ ],
607
+)
608
+@torch.inference_mode()
609
+def test_persistent_topk_correctness(test_config: dict) -> None:
610
+ """
611
+ Comprehensive correctness tests covering:
612
+ - Sequence length edge cases (trivial, boundary, varied)
613
+ - Very small sequences (< 100 elements)
614
+ - Mixed sequence lengths in same batch
615
+ - Data distributions (sorted, ties, precision)
616
+ - Memory alignment / vectorization boundaries
617
+ """
618
+ run_large_context_topk_test(
619
+ batch_size=len(test_config["seq_lens"]),
620
+ seq_lens=test_config["seq_lens"],
621
+ top_k=test_config["top_k"],
622
+ data_type=test_config.get("data_type", "random"),
623
)
624
334
- row_starts_long = torch.zeros(num_rows_long, dtype=torch.int32, device="cuda")
335
- row_indices_long = torch.arange(num_rows_long, device="cuda") // next_n
336
- next_n_offset_long = torch.arange(num_rows_long, device="cuda") % next_n
337
- row_ends_long = seq_lens_long[row_indices_long] - next_n + next_n_offset_long + 1
625
339
- logits_long = create_random_logits(
340
- row_starts_long, row_ends_long, torch.float32, 43, clean_logits, "random"
626
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
627
+@pytest.mark.parametrize(
628
+ "test_config",
629
+ [
630
+ # ==================== CATEGORY: Batch Size Scalability ====================
631
+ pytest.param(
632
+ {"batch_size": 1, "seq_len": 5000, "top_k": 2048},
633
+ id="batch_1",
634
+ ),
635
+ pytest.param(
636
+ {"batch_size": 4, "seq_len": 5000, "top_k": 2048},
637
+ id="batch_4",
638
+ ),
639
+ pytest.param(
640
+ {"batch_size": 32, "seq_len": 5000, "top_k": 2048},
641
+ id="batch_32",
642
+ ),
643
+ pytest.param(
644
+ {"batch_size": 256, "seq_len": 5000, "top_k": 2048},
645
+ id="batch_256",
646
+ ),
647
+ # ==================== CATEGORY: Single-CTA vs Multi-CTA ====================
648
+ pytest.param(
649
+ {"batch_size": 2, "seq_len": 4096, "top_k": 2048},
650
+ id="single_cta_4k",
651
+ ),
652
+ pytest.param(
653
+ {"batch_size": 2, "seq_len": 8192, "top_k": 2048},
654
+ id="single_cta_8k",
655
+ ),
656
+ pytest.param(
657
+ {"batch_size": 2, "seq_len": 163840, "top_k": 2048},
658
+ id="multi_cta_163840_dsv3_max",
659
+ ),
660
+ # ==================== CATEGORY: Extreme Cases ====================
661
+ pytest.param(
662
+ {"batch_size": 512, "seq_len": 5000, "top_k": 2048},
663
+ id="extreme_large_batch",
664
+ ),
665
+ pytest.param(
666
+ {"batch_size": 2, "seq_len": 163840, "top_k": 2048},
667
+ id="extreme_dsv3_max_context",
668
+ ),
669
+ ],
670
+)
671
+@torch.inference_mode()
672
+def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
673
+ """
674
+ Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2):
675
+ - Batch size scalability (1, 4, 32, 256)
676
+ - Single-CTA vs Multi-CTA execution
677
+ - Extreme configurations (large batch, max context length)
678
+ """
679
+ run_large_context_topk_test(
680
+ batch_size=test_config["batch_size"],
681
+ seq_lens=[test_config["seq_len"]] * test_config["batch_size"],
682
+ top_k=test_config["top_k"],
683
)
684
343
- indices = torch.empty((num_rows_long, top_k), dtype=torch.int32, device="cuda")
685
345
- # Use large_context_topk kernel for long sequences
346
- if next_n == 1:
347
- lengths = seq_lens_long
348
- else:
349
- offsets = torch.arange(next_n, device=logits_long.device, dtype=torch.int32)
350
- lengths = (seq_lens_long.unsqueeze(1) - next_n + 1 + offsets).flatten()
686
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
687
+@torch.inference_mode()
688
+def test_persistent_topk_stress() -> None:
689
+ """
690
+ Stress test with random configurations to catch edge cases.
691
+ Capped at 163840 (DeepSeek V3.2 max context) for realistic testing.
692
+ """
693
+ torch.set_default_device("cuda:0")
694
+ top_k = 2048
695
352
- torch.ops._C.large_context_topk(
353
- logits_long,
354
- indices,
355
- lengths,
356
- None,
696
+ for seed in range(3):
697
+ set_random_seed(seed)
698
+
699
+ # Random batch size (limited for speed)
700
+ batch_size = torch.randint(1, 32, (1,)).item()
701
+
702
+ # Random sequence lengths capped at DeepSeek V3.2 max context
703
+ seq_lens = torch.randint(100, 163840, (batch_size,)).tolist()
704
+
705
+ run_large_context_topk_test(
706
+ batch_size=batch_size,
707
+ seq_lens=seq_lens,
708
+ top_k=top_k,
709
+ seed=seed,
710
+ )
711
+
712
+
713
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
714
+@pytest.mark.parametrize(
715
+ "test_config",
716
+ [
717
+ # Mixed batch: rows spanning all four paths (trivial, decode, medium, large)
718
+ pytest.param(
719
+ {
720
+ "seq_lens": [2000, 6000, 30000, 80000],
721
+ "top_k": 2048,
722
+ "data_type": "random",
723
+ },
724
+ id="mixed_all_paths",
725
+ ),
726
+ # All decode/medium rows (typical decode scenario)
727
+ pytest.param(
728
+ {
729
+ "seq_lens": [2048, 4096, 8192, 16000],
730
+ "top_k": 2048,
731
+ "data_type": "random",
732
+ },
733
+ id="all_decode_medium",
734
+ ),
735
+ # All large rows
736
+ pytest.param(
737
+ {
738
+ "seq_lens": [70000, 100000, 163840],
739
+ "top_k": 2048,
740
+ "data_type": "random",
741
+ },
742
+ id="all_large",
743
+ ),
744
+ # Boundary around LARGE_THRESHOLD (32K)
745
+ pytest.param(
746
+ {
747
+ "seq_lens": [32767, 32768, 32769, 32772],
748
+ "top_k": 2048,
749
+ "data_type": "random",
750
+ },
751
+ id="large_threshold_boundary",
752
+ ),
753
+ # Single row medium
754
+ pytest.param(
755
+ {
756
+ "seq_lens": [5000],
757
+ "top_k": 2048,
758
+ "data_type": "random",
759
+ },
760
+ id="single_row_medium",
761
+ ),
762
+ # Single row large
763
+ pytest.param(
764
+ {
765
+ "seq_lens": [100000],
766
+ "top_k": 2048,
767
+ "data_type": "random",
768
+ },
769
+ id="single_row_large",
770
+ ),
771
+ # Trivial rows mixed with medium and large
772
+ pytest.param(
773
+ {
774
+ "seq_lens": [100, 2048, 10000, 80000],
775
+ "top_k": 2048,
776
+ "data_type": "random",
777
+ },
778
+ id="trivial_medium_large_mix",
779
+ ),
780
+ ],
781
+)
782
+@torch.inference_mode()
783
+def test_persistent_topk(test_config: dict) -> None:
784
+ """
785
+ Tests specific to the persistent_topk kernel:
786
+ - Mixed medium/large rows in the same batch (dynamic per-row dispatch)
787
+ - Boundary around LARGE_THRESHOLD (32K)
788
+ - Trivial + medium + large rows in a single batch
789
+ """
790
+ run_large_context_topk_test(
791
+ batch_size=len(test_config["seq_lens"]),
792
+ seq_lens=test_config["seq_lens"],
793
+ top_k=test_config["top_k"],
794
+ data_type=test_config.get("data_type", "random"),
795
)
796
359
- torch_indices_short = torch.empty(
360
- (num_rows_short, top_k), dtype=torch.int32, device="cuda"
797
+
798
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
799
+@torch.inference_mode()
800
+def test_persistent_topk_padded_stride() -> None:
801
+ """
802
+ Test persistent_topk with padded logits (large stride, small seq_len)
803
+ to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
804
+ returns [B, max_model_len] with max_model_len=163840.
805
+ """
806
+ set_random_seed(42)
807
+ torch.set_default_device("cuda:0")
808
+
809
+ top_k = 2048
810
+ batch_size = 4
811
+ padded_stride = 163840 # DeepSeek-V3.2 max_model_len
812
+ actual_seq_lens = [3000, 5000, 8000, 12000]
813
+
814
+ # Create padded logits tensor (like fp8_paged_mqa_logits output)
815
+ logits = torch.full(
816
+ (batch_size, padded_stride),
817
+ float("-inf"),
818
+ dtype=torch.float32,
819
+ device="cuda",
820
)
362
- for i in range(num_rows_short):
363
- row_end = int(row_ends_short[i])
364
- k_i = min(top_k, row_end)
365
- idx = logits_short[i, :row_end].topk(k_i, dim=-1)[1]
366
- torch_indices_short[i, :k_i] = idx
821
+ for i, sl in enumerate(actual_seq_lens):
822
+ logits[i, :sl] = torch.randn(sl, dtype=torch.float32, device="cuda")
823
368
- assert compare_top_k_results(
369
- logits_short,
370
- indices_vllm,
371
- torch_indices_short,
372
- row_starts_short,
373
- row_ends_short,
374
- top_k,
375
- ), "top_k_per_row_decode kernel (short sequences) doesn't match torch.topk"
824
+ lengths = torch.tensor(actual_seq_lens, dtype=torch.int32, device="cuda")
825
+ indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
826
+ workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
827
377
- torch_indices_long = torch.empty(
378
- (num_rows_long, top_k), dtype=torch.int32, device="cuda"
828
+ torch.ops._C.persistent_topk(
829
+ logits, lengths, indices, workspace, top_k, max(actual_seq_lens)
830
)
380
- for i in range(num_rows_long):
381
- row_end = int(row_ends_long[i])
382
- k_i = min(top_k, row_end)
383
- idx = logits_long[i, :row_end].topk(k_i, dim=-1)[1]
384
- torch_indices_long[i, :k_i] = idx
831
+ torch.accelerator.synchronize()
832
386
- assert compare_top_k_results(
387
- logits_long, indices, torch_indices_long, row_starts_long, row_ends_long, top_k
388
- ), "large_context_topk kernel (long sequences) doesn't match torch.topk"
833
+ # Validate against torch.topk
834
+ for i in range(batch_size):
835
+ sl = actual_seq_lens[i]
836
+ k_i = min(top_k, sl)
837
+ expected = logits[i, :sl].topk(k_i, dim=-1)[1].cpu()
838
+ actual = indices[i, :k_i].cpu()
839
+
840
+ expected_set = set(expected.tolist())
841
+ actual_set = set(actual.tolist())
842
+
843
+ if expected_set != actual_set:
844
+ # Allow ties
845
+ expected_vals = logits[i, expected].cpu().sort(descending=True)[0]
846
+ actual_vals = logits[i, actual].cpu().sort(descending=True)[0]
847
+ assert torch.allclose(expected_vals, actual_vals, rtol=1e-4, atol=1e-4), (
848
+ f"Row {i}: persistent_topk with padded stride doesn't match. "
849
+ f"seq_len={sl}, stride={padded_stride}"
850
+ )
vllm/model_executor/layers/sparse_attn_indexer.py
+24
-24
@@ -25,6 +25,8 @@ elif current_platform.is_xpu():
25
26
logger = init_logger(__name__)
27
28
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
29
+
30
31
def sparse_attn_indexer(
32
hidden_states: torch.Tensor,
@@ -51,6 +53,7 @@ def sparse_attn_indexer(
53
current_workspace_manager().get_simultaneous(
54
((total_seq_lens, head_dim), torch.float8_e4m3fn),
55
((total_seq_lens, 4), torch.uint8),
56
+ ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
57
)
58
59
# Dummy allocation to simulate for peak logits tensor memory during inference.
@@ -157,15 +160,6 @@ def sparse_attn_indexer(
160
topk_tokens,
161
)
162
160
- # Compute lengths from row spans
161
- # lengths = (chunk.cu_seqlen_ke - chunk.cu_seqlen_ks).to(torch.int32)
162
- # torch.ops._C.large_context_topk(
163
- # logits,
164
- # topk_indices,
165
- # lengths,
166
- # chunk.cu_seqlen_ks, # row_starts
167
- # )
168
-
163
if has_decode:
164
decode_metadata = attn_metadata.decode
165
assert decode_metadata is not None
@@ -204,23 +198,29 @@ def sparse_attn_indexer(
198
num_rows = logits.shape[0]
199
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
200
207
- if decode_metadata.use_large_context_topk:
208
- if next_n == 1:
209
- lengths = decode_metadata.seq_lens
210
- else:
211
- # (bs,) -> (bs, 1) + (next_n,) -> (bs, next_n) -> (bs * next_n,)
212
- lengths = (
213
- decode_metadata.seq_lens.unsqueeze(1)
214
- - next_n
215
- + 1
216
- + decode_metadata.offsets
217
- ).flatten()
218
-
219
- torch.ops._C.large_context_topk(
201
+ if next_n == 1:
202
+ lengths = decode_metadata.seq_lens
203
+ else:
204
+ # (bs,) -> (bs, 1) + (next_n,) -> (bs, next_n) -> (bs * next_n,)
205
+ lengths = (
206
+ decode_metadata.seq_lens.unsqueeze(1)
207
+ - next_n
208
+ + 1
209
+ + decode_metadata.offsets
210
+ ).flatten()
211
+
212
+ if current_platform.is_cuda():
213
+ workspace_manager = current_workspace_manager()
214
+ (topk_workspace,) = workspace_manager.get_simultaneous(
215
+ ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
216
+ )
217
+ torch.ops._C.persistent_topk(
218
logits,
221
- topk_indices,
219
lengths,
223
- None,
220
+ topk_indices,
221
+ topk_workspace,
222
+ topk_tokens,
223
+ attn_metadata.max_seq_len,
224
)
225
else:
226
if current_platform.is_xpu():
vllm/model_executor/models/deepseek_v2.py
+6
-2
@@ -67,7 +67,9 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
67
per_token_group_quant_fp8,
68
)
69
from vllm.model_executor.layers.rotary_embedding import get_rope
70
-from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
70
+from vllm.model_executor.layers.sparse_attn_indexer import (
71
+ SparseAttnIndexer,
72
+)
73
from vllm.model_executor.layers.vocab_parallel_embedding import (
74
ParallelLMHead,
75
VocabParallelEmbedding,
@@ -1203,7 +1205,9 @@ class DeepseekV2Model(nn.Module):
1205
self.start_layer, self.end_layer, self.layers = make_layers(
1206
config.num_hidden_layers,
1207
lambda prefix: DeepseekV2DecoderLayer(
1206
- vllm_config, prefix, topk_indices_buffer=topk_indices_buffer
1208
+ vllm_config,
1209
+ prefix,
1210
+ topk_indices_buffer=topk_indices_buffer,
1211
),
1212
prefix=f"{prefix}.layers",
1213
)
vllm/v1/attention/backends/mla/indexer.py
-12
@@ -145,7 +145,6 @@ class DeepSeekV32IndexerDecodeMetadata:
145
decode_lens: torch.Tensor
146
requires_padding: bool
147
schedule_metadata: torch.Tensor
148
- use_large_context_topk: bool
148
offsets: torch.Tensor | None # Precomputed offsets for speculative decoding
149
150
@@ -437,7 +436,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
436
437
if use_native and next_n > 1:
438
offsets = self.offsets_buffer
440
- batch_size = num_decodes
439
elif max_decode_len > 1:
440
# Flatten multi-token decode requests into single-token
441
# batch entries, expanding seq_lens and block tables so
@@ -496,10 +494,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
494
self.decode_lens_buffer[:num_decode_tokens] = 1
495
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
496
offsets = None
499
- batch_size = num_decode_tokens
497
else:
498
offsets = None
502
- batch_size = num_decodes
499
500
# DeepGEMM is required for the paged MQA logits on CUDA devices
501
if current_platform.is_cuda() and has_deep_gemm():
@@ -509,20 +505,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
505
self.num_sms,
506
)
507
512
- # Decide which top-k kernel to use based on batch size and sequence length
513
- # Decision logic based on micro-benchmark results:
514
- # - large_context_topk wins for batch <= 128 and seq_len > 8K
515
- # - top_k_per_row_decode wins for batch > 128 or seq_len <= 8K
516
- _is_large_context = common_attn_metadata.max_seq_len > 8192
517
- use_large_context_topk = batch_size <= 128 and _is_large_context
518
-
508
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
509
block_table=block_table,
510
seq_lens=seq_lens,
511
decode_lens=decode_lens,
512
requires_padding=False,
513
schedule_metadata=self.scheduler_metadata_buffer,
525
- use_large_context_topk=use_large_context_topk,
514
offsets=offsets,
515
)
516