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
[Deepseek v3.2] Optimize top_k_per_row (#26763)
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Daniel Cámpora committed
Oct 21, 2025 at 10:30 UTC
80e94529845d78284047e52f6253889c534bce13
5 files changed
+13
-49
csrc/ops.h
+1
-2
@@ -99,8 +99,7 @@ void apply_repetition_penalties_(torch::Tensor& logits,
99
100
void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
101
const torch::Tensor& rowEnds, torch::Tensor& indices,
102
- torch::Tensor& values, int64_t numRows, int64_t stride0,
103
- int64_t stride1);
102
+ int64_t numRows, int64_t stride0, int64_t stride1);
103
104
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
105
torch::Tensor& weight, torch::Tensor& scale,
csrc/sampler.cu
+5
-30
@@ -57,7 +57,7 @@ static inline __device__ uint16_t extractBinIdx(float x) {
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) {
60
+ int stride0, int stride1) {
61
// The number of bins in the histogram.
62
static constexpr int kNumBins = 512;
63
@@ -103,8 +103,6 @@ static __global__ void topKPerRow(const float* logits, const int* rowStarts,
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];
106
// Shared memory to store the threshold bin.
107
__shared__ int smemThresholdBinIdx[1];
108
// Shared memory counter to register the candidates for the final phase.
@@ -124,13 +122,10 @@ static __global__ void topKPerRow(const float* logits, const int* rowStarts,
122
rowIt += kNumThreadsPerBlock) {
123
int idx = rowStart + rowIt;
124
outIndices[rowIdx * kTopK + rowIt] = idx - rowStart;
127
- outLogits[rowIdx * kTopK + rowIt] =
128
- logits[rowIdx * stride0 + idx * stride1];
125
}
126
for (int rowIt = rowLen + threadIdx.x; rowIt < kTopK;
127
rowIt += kNumThreadsPerBlock) {
128
outIndices[rowIdx * kTopK + rowIt] = -1;
133
- outLogits[rowIdx * kTopK + rowIt] = -FLT_MAX;
129
}
130
return;
131
}
@@ -201,7 +196,6 @@ static __global__ void topKPerRow(const float* logits, const int* rowStarts,
196
uint16_t idx = extractBinIdx(logit);
197
if (idx < thresholdBinIdx) {
198
int dstIdx = atomicAdd(&smemHistogram[idx], 1);
204
- smemLogits[dstIdx] = logit;
199
smemIndices[dstIdx] = rowIt;
200
} else if (idx == thresholdBinIdx) {
201
int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
@@ -250,7 +244,6 @@ static __global__ void topKPerRow(const float* logits, const int* rowStarts,
244
int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
245
int dstIdx = baseIdx + srcIdx;
246
if (dstIdx < kTopK) {
253
- smemLogits[dstIdx] = finalLogits[ii];
247
smemIndices[dstIdx] = finalIndices[ii];
248
}
249
}
@@ -258,28 +251,12 @@ static __global__ void topKPerRow(const float* logits, const int* rowStarts,
251
// Make sure the data is in shared memory.
252
__syncthreads();
253
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
-
254
// Store to global memory.
255
#pragma unroll
256
for (int ii = 0; ii < kNumTopKItemsPerThread; ++ii) {
257
int offset = rowIdx * kTopK + ii * kNumThreadsPerBlock + threadIdx.x;
281
- outIndices[offset] = topKIndices[ii] - rowStart;
282
- outLogits[offset] = topKLogits[ii];
258
+ outIndices[offset] =
259
+ smemIndices[ii * kNumThreadsPerBlock + threadIdx.x] - rowStart;
260
}
261
}
262
@@ -328,8 +305,7 @@ void apply_repetition_penalties_(
305
306
void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
307
const torch::Tensor& rowEnds, torch::Tensor& indices,
331
- torch::Tensor& values, int64_t numRows, int64_t stride0,
332
- int64_t stride1) {
308
+ int64_t numRows, int64_t stride0, int64_t stride1) {
309
// Compute the results on the device.
310
constexpr int kNumThreadsPerBlock = 512;
311
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
@@ -338,6 +314,5 @@ void top_k_per_row(const torch::Tensor& logits, const torch::Tensor& rowStarts,
314
<<<numRows, kNumThreadsPerBlock, 0, stream>>>(
315
logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
316
rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
341
- values.data_ptr<float>(), static_cast<int>(stride0),
342
- static_cast<int>(stride1));
317
+ static_cast<int>(stride0), static_cast<int>(stride1));
318
}
csrc/torch_bindings.cpp
+1
-1
@@ -185,7 +185,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
185
// Optimized top-k per row operation
186
ops.def(
187
"top_k_per_row(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
188
- "Tensor! indices, Tensor! values, int numRows, int stride0, "
188
+ "Tensor! indices, int numRows, int stride0, "
189
"int stride1) -> ()");
190
ops.impl("top_k_per_row", torch::kCUDA, &top_k_per_row);
191
tests/kernels/test_top_k_per_row.py
+6
-8
@@ -39,10 +39,9 @@ def create_row_boundaries(
39
40
41
def compare_top_k_results(
42
+ logits: torch.Tensor,
43
cuda_indices: torch.Tensor,
43
- cuda_values: torch.Tensor,
44
torch_indices: torch.Tensor,
45
- torch_values: torch.Tensor,
45
row_starts: torch.Tensor,
46
row_ends: torch.Tensor,
47
top_k: int,
@@ -70,8 +69,9 @@ def compare_top_k_results(
69
continue
70
71
# 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()
72
+ logits_row = logits[row_idx]
73
+ cuda_row_values = [logits_row[i] for i in cuda_row_indices]
74
+ torch_row_values = [logits_row[i] for i in torch_row_indices]
75
76
cuda_only_values, torch_only_values = [], []
77
for idx in cuda_set - torch_set:
@@ -115,7 +115,6 @@ def test_top_k_per_row(
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")
118
119
# Run CUDA implementation
120
torch.ops._C.top_k_per_row(
@@ -123,14 +122,13 @@ def test_top_k_per_row(
122
row_starts,
123
row_ends,
124
indices,
126
- values,
125
num_rows,
126
logits.stride(0),
127
logits.stride(1),
128
)
129
130
# Run reference implementation
133
- torch_values, torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)
131
+ torch_indices = logits.topk(min(top_k, max(row_ends)), dim=-1)[1]
132
mask_lo = torch_indices >= 0
133
mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0
134
mask = mask_lo & mask_hi
@@ -138,5 +136,5 @@ def test_top_k_per_row(
136
137
# Compare results
138
assert compare_top_k_results(
141
- indices, values, torch_indices, torch_values, row_starts, row_ends, top_k
139
+ logits, indices, torch_indices, row_starts, row_ends, top_k
140
), "CUDA top_k_per_row results don't match torch.topk"
vllm/model_executor/models/deepseek_v2.py
-8
@@ -577,15 +577,11 @@ def sparse_attn_indexer(
577
topk_indices = torch.empty(
578
num_rows, topk_tokens, dtype=torch.int32, device=logits.device
579
)
580
- topk_values = torch.empty(
581
- num_rows, topk_tokens, dtype=logits.dtype, device=logits.device
582
- )
580
torch.ops._C.top_k_per_row(
581
logits,
582
chunk.cu_seqlen_ks,
583
chunk.cu_seqlen_ke,
584
topk_indices,
588
- topk_values,
585
num_rows,
586
logits.stride(0),
587
logits.stride(1),
@@ -642,15 +638,11 @@ def sparse_attn_indexer(
638
topk_indices = torch.empty(
639
num_rows, topk_tokens, dtype=torch.int32, device=logits.device
640
)
645
- topk_values = torch.empty(
646
- num_rows, topk_tokens, dtype=logits.dtype, device=logits.device
647
- )
641
torch.ops._C.top_k_per_row(
642
logits,
643
torch.zeros(num_rows, dtype=torch.int32, device=logits.device),
644
index_end_pos.to(dtype=torch.int32, device=logits.device),
645
topk_indices,
653
- topk_values,
646
num_rows,
647
logits.stride(0),
648
logits.stride(1),