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, 18:30 UTC
[V1][TPU] Speed up top-k on TPU by using torch.topk (#15242)
Signed-off-by: Hyesoo Yang <hyeygit@gmail.com>
Hyesoo Yang committed
Mar 20, 2025 at 19:19 UTC
47195057e9f2b93b35579e2ffd26e2e15cd49979
3 files changed
+29
-4
tests/v1/tpu/test_sampler.py
+2
-1
@@ -39,7 +39,7 @@ def test_sampler_compilation(model_name: str, monkeypatch):
39
sampling_params = SamplingParams(
40
temperature=0.7,
41
# top_p=0.6, # TODO too slow!
42
- # top_k=10,
42
+ top_k=10,
43
min_p=0.2,
44
max_tokens=16)
45
s = time()
@@ -49,6 +49,7 @@ def test_sampler_compilation(model_name: str, monkeypatch):
49
# Second request with different params, but for which we
50
# compiled for in previous eager iteration.
51
sampling_params = SamplingParams(temperature=0.1,
52
+ top_k=12,
53
min_p=0.8,
54
max_tokens=24)
55
s = time()
vllm/envs.py
+6
@@ -95,6 +95,7 @@ if TYPE_CHECKING:
95
VLLM_DP_MASTER_PORT: int = 0
96
VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
97
VLLM_V0_USE_OUTLINES_CACHE: bool = False
98
+ VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION: bool = False
99
100
101
def get_default_cache_root():
@@ -623,6 +624,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
624
# an environment with potentially malicious users.
625
"VLLM_V0_USE_OUTLINES_CACHE":
626
lambda: os.environ.get("VLLM_V0_USE_OUTLINES_CACHE", "0") == "1",
627
+
628
+ # If set, disables TPU-specific optimization for top-k & top-p sampling
629
+ "VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION":
630
+ lambda: bool(int(os.environ["VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION"]))
631
+ if "VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION" in os.environ else None,
632
}
633
634
# end-env-vars-definition
vllm/v1/sample/ops/topk_topp_sampler.py
+21
-3
@@ -66,7 +66,14 @@ class TopKTopPSampler(nn.Module):
66
"best performance, please install FlashInfer.")
67
self.forward = self.forward_native
68
elif current_platform.is_tpu():
69
- self.forward = self.forward_tpu
69
+ if envs.VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION:
70
+ logger.warning(
71
+ "TPU-specific optimization for top-k & top-p sampling are "
72
+ "disabled, falling back to PyTorch-native implementation "
73
+ "which could be very slow.")
74
+ self.forward = self.forward_native
75
+ else:
76
+ self.forward = self.forward_tpu
77
else:
78
self.forward = self.forward_native
79
@@ -105,8 +112,19 @@ class TopKTopPSampler(nn.Module):
112
k: Optional[torch.Tensor],
113
p: Optional[torch.Tensor],
114
) -> torch.Tensor:
108
- # TODO Placeholder for TPU optimized topk/p kernel
109
- # logits = apply_top_k_top_p(logits, k, p)
115
+ # If only top-k is specified, use pytorch's builtin topk op. This leads
116
+ # to significant speed up on TPU compared to using apply_top_k_top_p.
117
+ if k is not None and p is None:
118
+ topk_values, topk_indices = torch.topk(logits, k, dim=-1)
119
+
120
+ mask = torch.ones_like(logits, dtype=torch.bool)
121
+ mask.scatter_(-1, topk_indices, False)
122
+ logits.masked_fill_(mask, float('-inf'))
123
+ else:
124
+ # TODO Placeholder for TPU optimized topp kernel
125
+ # logits = apply_top_k_top_p(logits, k, p)
126
+ pass
127
+
128
probs = logits.softmax(dim=-1, dtype=torch.float32)
129
return random_sample(probs, generators)
130