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 4, 00:00 UTC
[Bugfix] Fix DeepSeek v4 topk numerical issue for unaligned max-model-len (#42169)
Wei Zhao committed
May 9, 2026 at 23:30 UTC
986edc858a13d92aa9bd2731710bd95e9e5c672d
1 file changed
+2
-2
csrc/topk.cu
+2
-2
@@ -20,7 +20,7 @@ void launch_persistent_topk(const torch::Tensor& logits,
20
namespace P = vllm::persistent;
21
22
const int64_t num_rows = logits.size(0);
23
- const int64_t stride = logits.size(1);
23
+ const int64_t stride = logits.stride(0);
24
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
25
26
static int num_sms = 0;
@@ -243,7 +243,7 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
243
TORCH_CHECK(output.dim() == 2, "output must be 2D");
244
245
const int64_t num_rows = logits.size(0);
246
- const int64_t stride = logits.size(1);
246
+ const int64_t stride = logits.stride(0);
247
248
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
249
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,