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 15, 03:30 UTC
Fix `numel()` downcast in vllm/csrc/moe/moe_align_sum_kernels.cu +2 (#17082)
Co-authored-by: mgoin <mgoin64@gmail.com>
Richard Barnes committed
Jul 1, 2025 at 00:48 UTC
86debab54c046232014b108d530a8c25d857e9a3
2 files changed
+2
-2
csrc/moe/moe_align_sum_kernels.cu
+1
-1
@@ -239,7 +239,7 @@ void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size]
239
torch::Tensor& output) // [num_tokens, hidden_size]
240
{
241
const int hidden_size = input.size(-1);
242
- const int num_tokens = output.numel() / hidden_size;
242
+ const auto num_tokens = output.numel() / hidden_size;
243
const int topk = input.size(1);
244
245
dim3 grid(num_tokens);
csrc/moe/topk_softmax_kernels.cu
+1
-1
@@ -492,7 +492,7 @@ void topk_softmax(
492
torch::Tensor& gating_output) // [num_tokens, num_experts]
493
{
494
const int num_experts = gating_output.size(-1);
495
- const int num_tokens = gating_output.numel() / num_experts;
495
+ const auto num_tokens = gating_output.numel() / num_experts;
496
const int topk = topk_weights.size(-1);
497
498
const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);