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 8, 14:13 UTC
[Perf][Kernel] Add faster topKperRow decode kernel for DeepSeek-V3.2 sparse attention (#33680)
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>
Roberto L. Castro committed
Feb 10, 2026 at 16:29 UTC
afdce12c89555ce7b7bd4f3215b5d844de0a32ed
8 files changed
+554
-12
CMakeLists.txt
+1
@@ -293,6 +293,7 @@ set(VLLM_EXT_SRC
293
"csrc/fused_qknorm_rope_kernel.cu"
294
"csrc/layernorm_quant_kernels.cu"
295
"csrc/sampler.cu"
296
+ "csrc/topk.cu"
297
"csrc/cuda_view.cu"
298
"csrc/quantization/gptq/q_gemm.cu"
299
"csrc/quantization/w8a8/int8/scaled_quant.cu"
csrc/ops.h
+4
@@ -114,6 +114,10 @@ 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);
120
+
121
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
122
torch::Tensor& weight, torch::Tensor& scale,
123
double epsilon);
csrc/sampler.cu
+1
-1
@@ -725,4 +725,4 @@ void top_k_per_row_prefill(const torch::Tensor& logits,
725
static_cast<int>(stride0), static_cast<int>(stride1),
726
static_cast<int>(topK), kSortingAlgorithmThreshold);
727
}
728
-}
728
+}
\ No newline at end of file
csrc/topk.cu
new
+373
@@ -0,0 +1,373 @@
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
5
+
6
+#include "cuda_compat.h"
7
+#include "dispatch_utils.h"
8
+
9
+#include <torch/cuda.h>
10
+#include <c10/cuda/CUDAGuard.h>
11
+
12
+#ifndef USE_ROCM
13
+ #include <cub/cub.cuh>
14
+#else
15
+ #include <hipcub/hipcub.hpp>
16
+#endif
17
+
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
+ }
179
+ }
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;
215
+
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;
230
+ }
231
+
232
+ // Continue refinement
233
+ __syncthreads();
234
+ if (thread_id < RADIX + 1) {
235
+ shared_histogram[0][thread_id] = 0;
236
+ }
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();
272
+ }
273
+}
274
+
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);
335
+#else
336
+ auto func_ptr = kernel_func;
337
+#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));
345
+}
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
+ c10::optional<torch::Tensor> row_starts = c10::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
+6
@@ -190,6 +190,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
190
"int numRows, int stride0, int stride1, int topK) -> ()");
191
ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
192
193
+ ops.def(
194
+ "large_context_topk(Tensor score, Tensor indices, Tensor lengths, "
195
+ "Tensor? "
196
+ "row_starts_opt) -> ()");
197
+ ops.impl("large_context_topk", torch::kCUDA, &large_context_topk);
198
+
199
// Layernorm-quant
200
// Apply Root Mean Square (RMS) Normalization to the input tensor.
201
ops.def(
tests/kernels/test_top_k_per_row.py
+111
@@ -275,3 +275,114 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
275
_run_top_k_per_row_decode_test(
276
top_k, batch_size, next_n, vocab_size, clean_logits, data_generation
277
)
278
+
279
+
280
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
281
+@pytest.mark.parametrize("clean_logits", [True, False])
282
+@torch.inference_mode()
283
+def test_deepseek_hybrid_topk(clean_logits: bool) -> None:
284
+ torch.set_default_device("cuda:0")
285
+
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
292
+
293
+ # Create sequences with max length < 8192
294
+ seq_lens_short = torch.randint(
295
+ 4000, 8000, (batch_size_short,), dtype=torch.int32, device="cuda"
296
+ )
297
+
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
303
+ )
304
+
305
+ logits_short = create_random_logits(
306
+ row_starts_short, row_ends_short, torch.float32, 42, clean_logits, "random"
307
+ )
308
+
309
+ indices_vllm = torch.empty(
310
+ (num_rows_short, top_k), dtype=torch.int32, device="cuda"
311
+ )
312
+
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,
323
+ )
324
+
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
328
+
329
+ # Create sequences with max length >= 8192
330
+ seq_lens_long = torch.randint(
331
+ 8192, 16384, (batch_size_long,), dtype=torch.int32, device="cuda"
332
+ )
333
+
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
338
+
339
+ logits_long = create_random_logits(
340
+ row_starts_long, row_ends_long, torch.float32, 43, clean_logits, "random"
341
+ )
342
+
343
+ indices = torch.empty((num_rows_long, top_k), dtype=torch.int32, device="cuda")
344
+
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()
351
+
352
+ torch.ops._C.large_context_topk(
353
+ logits_long,
354
+ indices,
355
+ lengths,
356
+ None,
357
+ )
358
+
359
+ torch_indices_short = torch.empty(
360
+ (num_rows_short, top_k), dtype=torch.int32, device="cuda"
361
+ )
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
367
+
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"
376
+
377
+ torch_indices_long = torch.empty(
378
+ (num_rows_long, top_k), dtype=torch.int32, device="cuda"
379
+ )
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
385
+
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"
vllm/model_executor/layers/sparse_attn_indexer.py
+39
-11
@@ -126,6 +126,15 @@ def sparse_attn_indexer(
126
topk_tokens,
127
)
128
129
+ # Compute lengths from row spans
130
+ # lengths = (chunk.cu_seqlen_ke - chunk.cu_seqlen_ks).to(torch.int32)
131
+ # torch.ops._C.large_context_topk(
132
+ # logits,
133
+ # topk_indices,
134
+ # lengths,
135
+ # chunk.cu_seqlen_ks, # row_starts
136
+ # )
137
+
138
if has_decode:
139
decode_metadata = attn_metadata.decode
140
# kv_cache size requirement [num_block, block_size, n_head, head_dim],
@@ -162,18 +171,37 @@ def sparse_attn_indexer(
171
)
172
173
num_rows = logits.shape[0]
165
-
174
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
167
- torch.ops._C.top_k_per_row_decode(
168
- logits,
169
- next_n,
170
- decode_metadata.seq_lens,
171
- topk_indices,
172
- num_rows,
173
- logits.stride(0),
174
- logits.stride(1),
175
- topk_tokens,
176
- )
175
+
176
+ if decode_metadata.use_large_context_topk:
177
+ if next_n == 1:
178
+ lengths = decode_metadata.seq_lens
179
+ else:
180
+ # (bs,) -> (bs, 1) + (next_n,) -> (bs, next_n) -> (bs * next_n,)
181
+ lengths = (
182
+ decode_metadata.seq_lens.unsqueeze(1)
183
+ - next_n
184
+ + 1
185
+ + decode_metadata.offsets
186
+ ).flatten()
187
+
188
+ torch.ops._C.large_context_topk(
189
+ logits,
190
+ topk_indices,
191
+ lengths,
192
+ None,
193
+ )
194
+ else:
195
+ torch.ops._C.top_k_per_row_decode(
196
+ logits,
197
+ next_n,
198
+ decode_metadata.seq_lens,
199
+ topk_indices,
200
+ num_rows,
201
+ logits.stride(0),
202
+ logits.stride(1),
203
+ topk_tokens,
204
+ )
205
206
if decode_metadata.requires_padding:
207
# if padded, we need to unpack
vllm/v1/attention/backends/mla/indexer.py
+19
@@ -86,6 +86,8 @@ class DeepSeekV32IndexerDecodeMetadata:
86
decode_lens: torch.Tensor
87
requires_padding: bool
88
schedule_metadata: torch.Tensor
89
+ use_large_context_topk: bool
90
+ offsets: torch.Tensor | None # Precomputed offsets for speculative decoding
91
92
93
@dataclass
@@ -320,6 +322,21 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
322
# Use CPU to avoid GPU sync; breaking async scheduling
323
requires_padding = (decode_lens_cpu.max() > decode_lens_cpu.min()).item()
324
325
+ # Decide which top-k kernel to use based on batch size and sequence length
326
+ batch_size = num_decodes
327
+ _is_large_context = common_attn_metadata.max_seq_len > 8192
328
+
329
+ # Decision logic based on micro-benchmark results:
330
+ # - large_context_topk wins for batch <= 128 and seq_len > 8K
331
+ # - top_k_per_row_decode wins for batch > 128 or seq_len <= 8K
332
+ use_large_context_topk = batch_size <= 128 and _is_large_context
333
+
334
+ next_n = 1 + self.num_speculative_tokens
335
+ if next_n > 1:
336
+ offsets = torch.arange(next_n, device=self.device, dtype=torch.int32)
337
+ else:
338
+ offsets = None
339
+
340
seq_lens = common_attn_metadata.seq_lens[:num_decodes]
341
if is_deep_gemm_supported():
342
self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
@@ -331,6 +348,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
348
decode_lens=decode_lens,
349
requires_padding=requires_padding,
350
schedule_metadata=self.scheduler_metadata_buffer,
351
+ use_large_context_topk=use_large_context_topk,
352
+ offsets=offsets,
353
)
354
355
attn_metadata = DeepseekV32IndexerMetadata(