A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
Optimized topk for topk=1 (Llama-4) (#16512)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Apr 12, 2025 at 00:21 UTC
bd6028d6b0bbc0c569ece0535067081c5e8bdc14
2 files changed
+11
-2
vllm/model_executor/models/llama4.py
+2
-2
@@ -37,7 +37,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
39
from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
40
-from .utils import (AutoWeightsLoader, extract_layer_index,
40
+from .utils import (AutoWeightsLoader, extract_layer_index, fast_topk,
41
is_pp_missing_parameter)
42
43
@@ -50,7 +50,7 @@ class Llama4MoE(nn.Module):
50
topk: int,
51
renormalize: bool,
52
) -> Tuple[torch.Tensor, torch.Tensor]:
53
- router_scores, router_indices = torch.topk(gating_output, topk, dim=-1)
53
+ router_scores, router_indices = fast_topk(gating_output, topk, dim=-1)
54
router_scores = torch.sigmoid(router_scores.float()).to(
55
hidden_states.dtype)
56
return (router_scores, router_indices.to(torch.int32))
vllm/model_executor/models/utils.py
+9
@@ -703,3 +703,12 @@ def cast_overflow_tensors(
703
clamp_value = torch.finfo(tensors.dtype).max - offset
704
tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value)
705
return tensors
706
+
707
+
708
+def fast_topk(values, topk, dim):
709
+ if topk == 1:
710
+ # Use max along the specified dimension to get both value and index
711
+ return torch.max(values, dim=dim, keepdim=True)
712
+ else:
713
+ # Use topk for efficiency with larger k values
714
+ return torch.topk(values, topk, dim=dim)