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 14, 05:30 UTC
Improve fast_topk function with type hints and documentation (#22530)
Signed-off-by: zitian.zhao <zitian.zhao@tencentmusic.com>
ZiTian Zhao committed
Aug 10, 2025 at 11:25 UTC
534c45b9620d4d97cf2ea2cdee77e8461844a243
1 file changed
+17
-1
vllm/model_executor/models/utils.py
+17
-1
@@ -736,7 +736,23 @@ def cast_overflow_tensors(
736
return tensors
737
738
739
-def fast_topk(values, topk, dim):
739
+def fast_topk(values: torch.Tensor, topk: int,
740
+ dim: int) -> tuple[torch.Tensor, torch.Tensor]:
741
+ """
742
+ Optimized topk implementation that uses torch.max for k=1 case.
743
+
744
+ This function provides better performance for the common case of k=1
745
+ by using torch.max instead of the more general torch.topk.
746
+
747
+ Args:
748
+ values: Input tensor to find top-k values from
749
+ topk: Number of top values to return (k). Must be > 0.
750
+ dim: Dimension along which to compute topk
751
+
752
+ Returns:
753
+ Tuple of (values, indices) where values are the top-k values
754
+ and indices are their corresponding indices in the input tensor
755
+ """
756
if topk == 1:
757
# Use max along the specified dimension to get both value and index
758
return torch.max(values, dim=dim, keepdim=True)