Skip to content

GitLab

Unverified Commit 8f44a92d authored by Dipika Sikka's avatar Dipika Sikka Committed by GitHub
Browse files

[BugFix] fix group_topk (#8430)

parent 360ddbd3
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
vllm/model_executor/layers/fused_moe/fused_moe.py
View file @ 8f44a92d
...@@ -410,6 +410,7 @@ def fused_topk( ...@@ -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, ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment