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 14, 05:30 UTC
Add topk logits torch op for DS3.2. (#25945)
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Cámpora <961215+dcampora@users.noreply.github.com> Co-authored-by: youkaichao <youkaichao@gmail.com>
Daniel Cámpora committed
Oct 7, 2025 at 12:07 UTC
e1098ced95146d98a4ed46c81ee709013d54fb1f
5 files changed
+446
-25
csrc/ops.h
+5
@@ -100,6 +100,11 @@ void apply_repetition_penalties_(torch::Tensor& logits,
100
const torch::Tensor& output_mask,
101
const torch::Tensor& repetition_penalties);
102
103
+void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
104
+ const torch::Tensor& rowEnds, torch::Tensor& indices,
105
+ torch::Tensor& values, int64_t numRows, int64_t stride0,
106
+ int64_t stride1);
107
+
108
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
109
torch::Tensor& weight, torch::Tensor& scale,
110
double epsilon);
csrc/sampler.cu
+256
-1
@@ -44,6 +44,245 @@ __global__ void apply_repetition_penalties_kernel(
44
}
45
}
46
47
+static inline __device__ uint16_t extractBinIdx(float x) {
48
+ union {
49
+ __half h;
50
+ uint16_t u16;
51
+ } tmp;
52
+ tmp.h = __float2half_rn(x);
53
+ tmp.u16 = (x < 0.f) ? (~tmp.u16 & 0xffff) : (tmp.u16 | 0x8000);
54
+ return 511 - (tmp.u16 >> 7);
55
+}
56
+
57
+template <int kNumThreadsPerBlock = 512>
58
+static __global__ void topKPerRow(const float* logits, const int* rowStarts,
59
+ const int* rowEnds, int* outIndices,
60
+ float* outLogits, int stride0, int stride1) {
61
+ // The number of bins in the histogram.
62
+ static constexpr int kNumBins = 512;
63
+
64
+ // The top-k width.
65
+ static constexpr int kTopK = 2048;
66
+ // The number of elements per thread for the final top-k sort.
67
+ static constexpr int kNumTopKItemsPerThread = kTopK / kNumThreadsPerBlock;
68
+ // The class to sort the elements during the final top-k sort.
69
+ using TopKSort = cub::BlockRadixSort<float, kNumThreadsPerBlock,
70
+ kNumTopKItemsPerThread, int>;
71
+
72
+ // The number of slots for the final pass.
73
+ static constexpr int kNumFinalItems = 3072;
74
+ // The number of elements per thread for the final sort.
75
+ static constexpr int kNumFinalItemsPerThread =
76
+ kNumFinalItems / kNumThreadsPerBlock;
77
+ // The class to sort the elements during the final pass.
78
+ using FinalSort = cub::BlockRadixSort<float, kNumThreadsPerBlock,
79
+ kNumFinalItemsPerThread, int>;
80
+
81
+ // The class to compute the inclusive prefix-sum over the histogram.
82
+ using Scan = cub::BlockScan<int, kNumThreadsPerBlock>;
83
+
84
+ // Shared memory to compute the block scan.
85
+ __shared__ typename Scan::TempStorage smemScan;
86
+
87
+ // The structure to store the final items (for the final pass).
88
+ struct FinalItems {
89
+ // Shared memory to store the indices for the final pass.
90
+ int indices[kNumFinalItems];
91
+ // Shared memory to store the logits for the final pass.
92
+ float logits[kNumFinalItems];
93
+ };
94
+
95
+ // Shared memory to compute the block sort.
96
+ __shared__ union {
97
+ FinalItems items;
98
+ typename FinalSort::TempStorage finalSort;
99
+ typename TopKSort::TempStorage topKSort;
100
+ } smemFinal;
101
+
102
+ // Shared memory to store the histogram.
103
+ __shared__ int smemHistogram[kNumBins];
104
+ // Shared memory to store the selected indices.
105
+ __shared__ int smemIndices[kTopK];
106
+ // Shared memory to store the selected logits.
107
+ __shared__ float smemLogits[kTopK];
108
+ // Shared memory to store the threshold bin.
109
+ __shared__ int smemThresholdBinIdx[1];
110
+ // Shared memory counter to register the candidates for the final phase.
111
+ __shared__ int smemFinalDstIdx[1];
112
+
113
+ // The row computed by this block.
114
+ int rowIdx = blockIdx.x;
115
+ // The range of logits within the row.
116
+ int rowStart = rowStarts[rowIdx], rowEnd = rowEnds[rowIdx];
117
+ // The length of the row.
118
+ int rowLen = rowEnd - rowStart;
119
+
120
+ // Shortcut if the length of the row is smaller than Top-K. Indices are not
121
+ // sorted by their corresponding logit.
122
+ if (rowLen <= kTopK) {
123
+ for (int rowIt = threadIdx.x; rowIt < rowLen;
124
+ rowIt += kNumThreadsPerBlock) {
125
+ int idx = rowStart + rowIt;
126
+ outIndices[rowIdx * kTopK + rowIt] = idx - rowStart;
127
+ outLogits[rowIdx * kTopK + rowIt] =
128
+ logits[rowIdx * stride0 + idx * stride1];
129
+ }
130
+ for (int rowIt = rowLen + threadIdx.x; rowIt < kTopK;
131
+ rowIt += kNumThreadsPerBlock) {
132
+ outIndices[rowIdx * kTopK + rowIt] = -1;
133
+ outLogits[rowIdx * kTopK + rowIt] = -FLT_MAX;
134
+ }
135
+ return;
136
+ }
137
+
138
+ // Clear the histogram.
139
+ if (threadIdx.x < kNumBins) {
140
+ smemHistogram[threadIdx.x] = 0;
141
+ }
142
+
143
+ // Make sure the histogram is ready.
144
+ __syncthreads();
145
+
146
+ // Fetch elements one-by-one.
147
+ for (int rowIt = rowStart + threadIdx.x; rowIt < rowEnd;
148
+ rowIt += kNumThreadsPerBlock) {
149
+ uint16_t idx = extractBinIdx(logits[rowIdx * stride0 + rowIt * stride1]);
150
+ atomicAdd(&smemHistogram[idx], 1);
151
+ }
152
+
153
+ // Make sure the histogram is ready.
154
+ __syncthreads();
155
+
156
+ // Read the values from SMEM.
157
+ int binCount{0};
158
+ if (threadIdx.x < kNumBins) {
159
+ binCount = smemHistogram[threadIdx.x];
160
+ }
161
+
162
+ // Make sure each thread has read its value.
163
+ __syncthreads();
164
+
165
+ // Compute the prefix sum.
166
+ int prefixSum{0}, totalSum{0};
167
+ Scan(smemScan).ExclusiveSum(binCount, prefixSum, totalSum);
168
+
169
+ // Update the histogram with the prefix sums.
170
+ if (threadIdx.x < kNumBins) {
171
+ smemHistogram[threadIdx.x] = prefixSum;
172
+ }
173
+
174
+ // Make sure the data is in shared memory.
175
+ __syncthreads();
176
+
177
+ // Find the last valid bin.
178
+ if (threadIdx.x < kNumBins) {
179
+ int nextPrefixSum =
180
+ threadIdx.x == kNumBins - 1 ? totalSum : smemHistogram[threadIdx.x + 1];
181
+ if (prefixSum < kTopK && nextPrefixSum >= kTopK) {
182
+ smemThresholdBinIdx[0] = threadIdx.x;
183
+ }
184
+ }
185
+
186
+ // Clear the counter to store the items for the final phase.
187
+ if (threadIdx.x == 0) {
188
+ smemFinalDstIdx[0] = 0;
189
+ }
190
+
191
+ // Make sure the data is in shared memory.
192
+ __syncthreads();
193
+
194
+ // The threshold bin.
195
+ int thresholdBinIdx = smemThresholdBinIdx[0];
196
+
197
+ // Fetch elements one-by-one and populate the shared memory buffers.
198
+ for (int rowIt = rowStart + threadIdx.x; rowIt < rowEnd;
199
+ rowIt += kNumThreadsPerBlock) {
200
+ float logit = logits[rowIdx * stride0 + rowIt * stride1];
201
+ uint16_t idx = extractBinIdx(logit);
202
+ if (idx < thresholdBinIdx) {
203
+ int dstIdx = atomicAdd(&smemHistogram[idx], 1);
204
+ smemLogits[dstIdx] = logit;
205
+ smemIndices[dstIdx] = rowIt;
206
+ } else if (idx == thresholdBinIdx) {
207
+ int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
208
+ if (dstIdx < kNumFinalItems) {
209
+ smemFinal.items.logits[dstIdx] = logit;
210
+ smemFinal.items.indices[dstIdx] = rowIt;
211
+ }
212
+ }
213
+ }
214
+
215
+ // Make sure the elements are in shared memory.
216
+ __syncthreads();
217
+
218
+ // The logits of the elements to be sorted in the final pass.
219
+ float finalLogits[kNumFinalItemsPerThread];
220
+ // The indices of the elements to be sorted in the final pass.
221
+ int finalIndices[kNumFinalItemsPerThread];
222
+
223
+// Init.
224
+#pragma unroll
225
+ for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
226
+ finalLogits[ii] = -FLT_MAX;
227
+ }
228
+
229
+// Read the elements from SMEM.
230
+#pragma unroll
231
+ for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
232
+ int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
233
+ if (srcIdx < smemFinalDstIdx[0]) {
234
+ finalLogits[ii] = smemFinal.items.logits[srcIdx];
235
+ finalIndices[ii] = smemFinal.items.indices[srcIdx];
236
+ }
237
+ }
238
+
239
+ // Make sure the shared memory has been read.
240
+ __syncthreads();
241
+
242
+ // Sort the elements.
243
+ FinalSort(smemFinal.finalSort)
244
+ .SortDescendingBlockedToStriped(finalLogits, finalIndices);
245
+
246
+ // Copy the data back to the shared memory storage.
247
+ int baseIdx = thresholdBinIdx > 0 ? smemHistogram[thresholdBinIdx - 1] : 0;
248
+#pragma unroll
249
+ for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
250
+ int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
251
+ int dstIdx = baseIdx + srcIdx;
252
+ if (dstIdx < kTopK) {
253
+ smemLogits[dstIdx] = finalLogits[ii];
254
+ smemIndices[dstIdx] = finalIndices[ii];
255
+ }
256
+ }
257
+
258
+ // Make sure the data is in shared memory.
259
+ __syncthreads();
260
+
261
+ // The topK logits.
262
+ float topKLogits[kNumTopKItemsPerThread];
263
+ // The topK indices.
264
+ int topKIndices[kNumTopKItemsPerThread];
265
+
266
+// Load from shared memory.
267
+#pragma unroll
268
+ for (int ii = 0; ii < kNumTopKItemsPerThread; ++ii) {
269
+ topKLogits[ii] = smemLogits[ii * kNumThreadsPerBlock + threadIdx.x];
270
+ topKIndices[ii] = smemIndices[ii * kNumThreadsPerBlock + threadIdx.x];
271
+ }
272
+
273
+ // Sort the elements.
274
+ TopKSort(smemFinal.topKSort)
275
+ .SortDescendingBlockedToStriped(topKLogits, topKIndices);
276
+
277
+// Store to global memory.
278
+#pragma unroll
279
+ for (int ii = 0; ii < kNumTopKItemsPerThread; ++ii) {
280
+ int offset = rowIdx * kTopK + ii * kNumThreadsPerBlock + threadIdx.x;
281
+ outIndices[offset] = topKIndices[ii] - rowStart;
282
+ outLogits[offset] = topKLogits[ii];
283
+ }
284
+}
285
+
286
} // namespace vllm
287
288
void apply_repetition_penalties_(
@@ -85,4 +324,20 @@ void apply_repetition_penalties_(
324
repetition_penalties.data_ptr<scalar_t>(), num_seqs, vocab_size,
325
tile_size);
326
});
88
-}
\ No newline at end of file
327
+}
328
+
329
+void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
330
+ const torch::Tensor& rowEnds, torch::Tensor& indices,
331
+ torch::Tensor& values, int64_t numRows, int64_t stride0,
332
+ int64_t stride1) {
333
+ // Compute the results on the device.
334
+ constexpr int kNumThreadsPerBlock = 512;
335
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
336
+
337
+ vllm::topKPerRow<kNumThreadsPerBlock>
338
+ <<<numRows, kNumThreadsPerBlock, 0, stream>>>(
339
+ logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
340
+ rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
341
+ values.data_ptr<float>(), static_cast<int>(stride0),
342
+ static_cast<int>(stride1));
343
+}
csrc/torch_bindings.cpp
+7
@@ -188,6 +188,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
188
ops.impl("apply_repetition_penalties_", torch::kCUDA,
189
&apply_repetition_penalties_);
190
191
+ // Optimized top-k per row operation
192
+ ops.def(
193
+ "top_k_per_row(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
194
+ "Tensor! indices, Tensor! values, int numRows, int stride0, "
195
+ "int stride1) -> ()");
196
+ ops.impl("top_k_per_row", torch::kCUDA, &top_k_per_row);
197
+
198
// Layernorm-quant
199
// Apply Root Mean Square (RMS) Normalization to the input tensor.
200
ops.def(
tests/kernels/test_top_k_per_row.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+
8
+from vllm.platforms import current_platform
9
+
10
+# Test parameters
11
+NUM_ROWS = [1, 32, 2050]
12
+TOP_K_VALUES = [2048]
13
+
14
+
15
+def create_random_logits(
16
+ row_starts: torch.Tensor,
17
+ row_ends: torch.Tensor,
18
+ vocab_size: int,
19
+ dtype: torch.dtype,
20
+ seed: int,
21
+) -> torch.Tensor:
22
+ """Create random logits tensor for testing."""
23
+ torch.manual_seed(seed)
24
+ np.random.seed(seed)
25
+ # Generate logits with some structure to make testing more meaningful
26
+ logits = torch.randn(row_starts.shape[0], max(row_ends), dtype=dtype, device="cuda")
27
+ for i, end in enumerate(row_ends):
28
+ logits[i, end:] = float("-inf")
29
+ return logits
30
+
31
+
32
+def create_row_boundaries(
33
+ seq_len: int, vocab_size: int
34
+) -> tuple[torch.Tensor, torch.Tensor]:
35
+ """Create row start and end indices for testing."""
36
+ row_starts = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
37
+ row_ends = torch.arange(1, seq_len + 1, device="cuda", dtype=torch.int32)
38
+ return row_starts, row_ends
39
+
40
+
41
+def compare_top_k_results(
42
+ cuda_indices: torch.Tensor,
43
+ cuda_values: torch.Tensor,
44
+ torch_indices: torch.Tensor,
45
+ torch_values: torch.Tensor,
46
+ row_starts: torch.Tensor,
47
+ row_ends: torch.Tensor,
48
+ top_k: int,
49
+ tolerance: float = 1e-5,
50
+) -> bool:
51
+ """
52
+ Compare results from CUDA top_k_per_row with torch.topk.
53
+ Both results should be sorted and contain the same top-k elements.
54
+ """
55
+ num_rows = cuda_indices.shape[0]
56
+
57
+ for row_idx in range(num_rows):
58
+ # Get valid elements using row boundaries
59
+ row_start = row_starts[row_idx].item()
60
+ row_end = row_ends[row_idx].item()
61
+ row_length = row_end - row_start
62
+ num_valid = min(top_k, row_length)
63
+ cuda_row_indices = cuda_indices[row_idx][:num_valid].cpu()
64
+ torch_row_indices = torch_indices[row_idx][:num_valid].cpu()
65
+
66
+ # Compare the sets of indices first
67
+ cuda_set = set(cuda_row_indices.tolist())
68
+ torch_set = set(torch_row_indices.tolist())
69
+ if cuda_set == torch_set:
70
+ continue
71
+
72
+ # Any difference in elements, compare the values
73
+ cuda_row_values = cuda_values[row_idx][:num_valid].cpu()
74
+ torch_row_values = torch_values[row_idx][:num_valid].cpu()
75
+
76
+ cuda_only_values, torch_only_values = [], []
77
+ for idx in cuda_set - torch_set:
78
+ cuda_pos = (cuda_row_indices == idx).nonzero(as_tuple=True)[0]
79
+ cuda_only_values.append(cuda_row_values[cuda_pos[0]])
80
+
81
+ for idx in torch_set - cuda_set:
82
+ torch_pos = (torch_row_indices == idx).nonzero(as_tuple=True)[0]
83
+ torch_only_values.append(torch_row_values[torch_pos[0]])
84
+
85
+ if len(cuda_only_values) != len(torch_only_values):
86
+ return False
87
+ if not torch.allclose(
88
+ torch.tensor(cuda_only_values),
89
+ torch.tensor(torch_only_values),
90
+ rtol=tolerance,
91
+ atol=tolerance,
92
+ ):
93
+ return False
94
+
95
+ return True
96
+
97
+
98
+@pytest.mark.parametrize("num_rows", NUM_ROWS)
99
+@pytest.mark.parametrize("top_k", TOP_K_VALUES)
100
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
101
+@torch.inference_mode()
102
+def test_top_k_per_row(
103
+ num_rows: int,
104
+ top_k: int,
105
+) -> None:
106
+ """
107
+ Test top_k_per_row.
108
+ """
109
+ torch.set_default_device("cuda:0")
110
+
111
+ # Create test data
112
+ vocab_size = 20000
113
+ row_starts, row_ends = create_row_boundaries(num_rows, vocab_size)
114
+ logits = create_random_logits(row_starts, row_ends, vocab_size, torch.float32, 42)
115
+
116
+ # Create output tensors
117
+ indices = torch.empty((num_rows, 2048), dtype=torch.int32, device="cuda")
118
+ values = torch.empty((num_rows, 2048), dtype=torch.float32, device="cuda")
119
+
120
+ # Run CUDA implementation
121
+ torch.ops._C.top_k_per_row(
122
+ logits,
123
+ row_starts,
124
+ row_ends,
125
+ indices,
126
+ values,
127
+ num_rows,
128
+ top_k,
129
+ logits.stride(0),
130
+ logits.stride(1),
131
+ )
132
+
133
+ # Run reference implementation
134
+ torch_values, torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)
135
+ mask_lo = torch_indices >= 0
136
+ mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0
137
+ mask = mask_lo & mask_hi
138
+ torch_indices = torch_indices.masked_fill(~mask, -1)
139
+
140
+ # Compare results
141
+ assert compare_top_k_results(
142
+ indices, values, torch_indices, torch_values, row_starts, row_ends, top_k
143
+ ), "CUDA top_k_per_row results don't match torch.topk"
vllm/model_executor/models/deepseek_v2.py
+35
-24
@@ -643,17 +643,24 @@ def sparse_attn_indexer(
643
chunk.cu_seqlen_ks,
644
chunk.cu_seqlen_ke,
645
)
646
- topk_indices = logits.topk(min(topk_tokens, logits.shape[-1]), dim=-1)[1]
647
- topk_indices -= chunk.cu_seqlen_ks[:, None]
648
- mask_lo = topk_indices >= 0
649
- mask_hi = (
650
- topk_indices - (chunk.cu_seqlen_ke - chunk.cu_seqlen_ks)[:, None] < 0
646
+ num_rows = logits.shape[0]
647
+ assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
648
+ topk_indices = torch.empty(
649
+ num_rows, topk_tokens, dtype=torch.int32, device=logits.device
650
)
652
- mask = torch.full_like(
653
- topk_indices, False, dtype=torch.bool, device=topk_indices.device
651
+ topk_values = torch.empty(
652
+ num_rows, topk_tokens, dtype=logits.dtype, device=logits.device
653
+ )
654
+ torch.ops._C.top_k_per_row(
655
+ logits,
656
+ chunk.cu_seqlen_ks,
657
+ chunk.cu_seqlen_ke,
658
+ topk_indices,
659
+ topk_values,
660
+ num_rows,
661
+ logits.stride(0),
662
+ logits.stride(1),
663
)
655
- mask = mask_lo & mask_hi
656
- topk_indices = topk_indices.masked_fill(~mask, -1)
664
topk_indices_buffer[
665
chunk.token_start : chunk.token_end, : topk_indices.shape[-1]
666
] = topk_indices.to(dtype=torch.int32)
@@ -693,28 +700,32 @@ def sparse_attn_indexer(
700
# padded query len
701
current_device = padded_q_fp8_decode_tokens.device
702
padded_num_tokens = batch_size * next_n
696
- positions = (
697
- torch.arange(max_model_len, device=current_device)
698
- .unsqueeze(0)
699
- .expand(batch_size * next_n, -1)
700
- )
703
row_indices = torch.arange(padded_num_tokens, device=current_device) // next_n
704
next_n_offset = (
705
torch.arange(padded_num_tokens, device=padded_q_fp8_decode_tokens.device)
706
% next_n
707
)
708
index_end_pos = (
707
- decode_metadata.seq_lens[row_indices] - next_n + next_n_offset
709
+ decode_metadata.seq_lens[row_indices] - next_n + next_n_offset + 1
710
).unsqueeze(1)
709
- # index_end_pos: [B * N, 1]
710
- mask = positions <= index_end_pos
711
- # mask: [B * N, L]
712
- logits = logits.masked_fill(~mask, float("-inf"))
713
- topk_indices = logits.topk(topk_tokens, dim=-1)[1].to(torch.int32) # [B * N, K]
714
- # ensure we don't set indices for the top k
715
- # that is out of range(masked already)
716
- # this will happen if context length is shorter than K
717
- topk_indices[topk_indices > index_end_pos] = -1
711
+ num_rows = logits.shape[0]
712
+ assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
713
+ topk_indices = torch.empty(
714
+ num_rows, topk_tokens, dtype=torch.int32, device=logits.device
715
+ )
716
+ topk_values = torch.empty(
717
+ num_rows, topk_tokens, dtype=logits.dtype, device=logits.device
718
+ )
719
+ torch.ops._C.top_k_per_row(
720
+ logits,
721
+ torch.zeros(num_rows, dtype=torch.int32, device=logits.device),
722
+ index_end_pos.to(dtype=torch.int32, device=logits.device),
723
+ topk_indices,
724
+ topk_values,
725
+ num_rows,
726
+ logits.stride(0),
727
+ logits.stride(1),
728
+ )
729
if decode_metadata.requires_padding:
730
# if padded, we need to unpack
731
# the topk indices removing padded tokens