biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Optimized topk for topk=1 (Llama-4) (#16512)

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin
2025-04-12 00:21:08 -06:00
committed by GitHub
parent 802329dee9
commit bd6028d6b0
2 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/models/llama4.py
View File

@@ -37,7 +37,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
from .utils import (AutoWeightsLoader, extract_layer_index,
from .utils import (AutoWeightsLoader, extract_layer_index, fast_topk,
is_pp_missing_parameter)
@@ -50,7 +50,7 @@ class Llama4MoE(nn.Module):
topk: int,
renormalize: bool,
) -> Tuple[torch.Tensor, torch.Tensor]:
router_scores, router_indices = torch.topk(gating_output, topk, dim=-1)
router_scores, router_indices = fast_topk(gating_output, topk, dim=-1)
router_scores = torch.sigmoid(router_scores.float()).to(
hidden_states.dtype)
return (router_scores, router_indices.to(torch.int32))