xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 04:26:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix numel() downcast in vllm/csrc/moe/moe_align_sum_kernels.cu +2 (#17082)

Browse Source 
Co-authored-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Richard Barnes 2025-07-01 00:48:10 -06:00 committed by GitHub
parent be250bbc67
commit 86debab54c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/moe/moe_align_sum_kernels.cu
View File

@ -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);

2
csrc/moe/topk_softmax_kernels.cu
View File

@ -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);