A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 08:45 UTC
[Bugfix][DSV32] Fix overflow in topk. (#30754)
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Daniel Cámpora committed
Dec 16, 2025 at 23:21 UTC
eaa82a709a963ab744647a701fe267223ed7b02b
1 file changed
+10
-8
csrc/sampler.cu
+10
-8
@@ -550,8 +550,8 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(
550
int rowEnd = rowEnds[rowIdx];
551
552
// Local pointers to this block
553
- outIndices += rowIdx * topK;
554
- logits += rowIdx * stride0;
553
+ outIndices += static_cast<int64_t>(rowIdx) * topK;
554
+ logits += static_cast<int64_t>(rowIdx) * stride0;
555
556
topKPerRowJob<kNumThreadsPerBlock, kNumBins, useRadixSort>(
557
nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK);
@@ -576,19 +576,21 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
576
577
// Local pointers to this block
578
if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
579
- outIndices += rowIdx * topK;
579
+ outIndices += static_cast<int64_t>(rowIdx) * topK;
580
} else if constexpr (multipleBlocksPerRow) {
581
const auto blockSize = rowEnd / gridDim.y; // 16384 / 2 = 8192
582
rowStart = blockSize * blockIdx.y; // 8192 * 1 = 8192
583
rowEnd = gridDim.y == blockIdx.y + 1 ? rowEnd : rowStart + blockSize;
584
- outIndices += rowIdx * gridDim.y * topK + blockIdx.y * topK;
585
- outLogits += rowIdx * gridDim.y * topK + blockIdx.y * topK;
584
+ outIndices +=
585
+ static_cast<int64_t>(rowIdx) * gridDim.y * topK + blockIdx.y * topK;
586
+ outLogits +=
587
+ static_cast<int64_t>(rowIdx) * gridDim.y * topK + blockIdx.y * topK;
588
} else if constexpr (mergeBlocks) {
589
rowEnd = numBlocksToMerge * topK;
588
- indices += rowIdx * numBlocksToMerge * topK;
589
- outIndices += rowIdx * topK;
590
+ indices += static_cast<int64_t>(rowIdx) * numBlocksToMerge * topK;
591
+ outIndices += static_cast<int64_t>(rowIdx) * topK;
592
}
591
- logits += rowIdx * stride0;
593
+ logits += static_cast<int64_t>(rowIdx) * stride0;
594
595
topKPerRowJob<kNumThreadsPerBlock, kNumBins, useRadixSort,
596
multipleBlocksPerRow, mergeBlocks>(