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

[BugFix] fix group_topk (#8430)

Browse Source
This commit is contained in:
Dipika Sikka 2024-09-12 21:23:42 -04:00 committed by GitHub
parent 360ddbd37e
commit 8f44a92d85
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/layers/fused_moe/fused_moe.py
View File

@ -410,6 +410,7 @@ def fused_topk(
if renormalize: if renormalize:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
return topk_weights, topk_ids return topk_weights, topk_ids
@ -443,7 +444,8 @@ def grouped_topk(hidden_states: torch.Tensor,
if renormalize: if renormalize:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
return topk_weights, topk_ids
return topk_weights, topk_ids.to(torch.int32)
def get_config_dtype_str(dtype: torch.dtype, def get_config_dtype_str(dtype: torch.dtype,