From 86debab54c046232014b108d530a8c25d857e9a3 Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Tue, 1 Jul 2025 00:48:10 -0600 Subject: [PATCH] Fix `numel()` downcast in vllm/csrc/moe/moe_align_sum_kernels.cu +2 (#17082) Co-authored-by: mgoin --- csrc/moe/moe_align_sum_kernels.cu | 2 +- csrc/moe/topk_softmax_kernels.cu | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu index 9335e2333b0d9..462dbd1f8b380 100644 --- a/csrc/moe/moe_align_sum_kernels.cu +++ b/csrc/moe/moe_align_sum_kernels.cu @@ -239,7 +239,7 @@ void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size] torch::Tensor& output) // [num_tokens, hidden_size] { const int hidden_size = input.size(-1); - const int num_tokens = output.numel() / hidden_size; + const auto num_tokens = output.numel() / hidden_size; const int topk = input.size(1); dim3 grid(num_tokens); diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu index dea5b1f21ec27..064b76c9cd427 100644 --- a/csrc/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -492,7 +492,7 @@ void topk_softmax( torch::Tensor& gating_output) // [num_tokens, num_experts] { const int num_experts = gating_output.size(-1); - const int num_tokens = gating_output.numel() / num_experts; + const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);