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
[TPU] Remove TopKTopPSampler dependency for TPU sampler (#24391)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Sep 7, 2025 at 01:12 UTC
105d3d62efde2d750b985020bbcea888f94a139b
3 files changed
+79
-57
tests/v1/tpu/test_topk_topp_sampler.py
+6
-2
@@ -6,8 +6,12 @@ import pytest
6
import torch
7
8
from vllm.platforms import current_platform
9
-from vllm.v1.sample.ops.topk_topp_sampler import (apply_top_k_top_p,
10
- apply_top_k_top_p_tpu)
9
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
10
+
11
+# isort: off
12
+from vllm.v1.sample.tpu.sampler import (apply_top_k_top_p as
13
+ apply_top_k_top_p_tpu)
14
+# isort: on
15
16
if not current_platform.is_tpu():
17
pytest.skip("This test needs a TPU.", allow_module_level=True)
vllm/v1/sample/ops/topk_topp_sampler.py
+2
-51
@@ -73,10 +73,8 @@ class TopKTopPSampler(nn.Module):
73
self.forward = self.forward_native
74
else:
75
self.forward = self.forward_native
76
- if current_platform.is_tpu():
77
- self.apply_top_k_top_p = apply_top_k_top_p_tpu
78
- else:
79
- self.apply_top_k_top_p = apply_top_k_top_p
76
+
77
+ self.apply_top_k_top_p = apply_top_k_top_p
78
79
def forward_native(
80
self,
@@ -125,53 +123,6 @@ class TopKTopPSampler(nn.Module):
123
return flashinfer_sample(logits.contiguous(), k, p, generators), None
124
125
128
-def apply_top_k_top_p_tpu(
129
- logits: torch.Tensor,
130
- k: torch.Tensor,
131
- p: torch.Tensor,
132
-) -> torch.Tensor:
133
- """
134
- Apply top-k and top-p optimized for TPU.
135
-
136
- This algorithm avoids using torch.scatter which is extremely slow on TPU.
137
- This is achieved by finding a "cut-off" element in the original logit, and
138
- after thresholding the logit using this cut-off, the remaining elements
139
- shall constitute the top-p set.
140
-
141
- Note: in the case of tie (i.e. multipple cut-off elements present in the
142
- logit), all tie elements are included in the top-p set. In other words,
143
- this function does not break ties. Instead, these tie tokens have equal
144
- chance of being chosen during final sampling, so we can consider the tie
145
- being broken then.
146
- """
147
- probs = logits.softmax(dim=-1)
148
- probs_sort, _ = probs.sort(dim=-1, descending=False)
149
-
150
- if k is not None:
151
- top_k_count = probs_sort.size(1) - k.to(torch.long) # shape: (batch, )
152
- top_k_count = top_k_count.unsqueeze(dim=1)
153
- top_k_cutoff = probs_sort.gather(-1, top_k_count)
154
-
155
- # Make sure the no top-k rows are no-op.
156
- no_top_k_mask = (k == logits.shape[1]).unsqueeze(dim=1)
157
- top_k_cutoff.masked_fill_(no_top_k_mask, -float("inf"))
158
-
159
- elements_to_discard = probs < top_k_cutoff
160
- logits.masked_fill_(elements_to_discard, -float("inf"))
161
-
162
- if p is not None:
163
- cumprob = torch.cumsum(probs_sort, dim=-1)
164
- top_p_mask = cumprob <= 1 - p.unsqueeze(dim=1)
165
- top_p_mask[:, -1] = False # at least one
166
-
167
- top_p_count = top_p_mask.sum(dim=-1).unsqueeze(1)
168
- top_p_cutoff = probs_sort.gather(-1, top_p_count)
169
- elements_to_discard = probs < top_p_cutoff
170
- logits.masked_fill_(elements_to_discard, -float("inf"))
171
-
172
- return logits
173
-
174
-
126
def apply_top_k_top_p(
127
logits: torch.Tensor,
128
k: Optional[torch.Tensor],
vllm/v1/sample/tpu/sampler.py
+71
-4
@@ -2,11 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Sampler layer implementing TPU supported operations."""
4
5
+from typing import Optional
6
+
7
import torch
8
import torch.nn as nn
9
10
from vllm.v1.outputs import LogprobsTensors, SamplerOutput
9
-from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
11
from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
12
13
_SAMPLING_EPS = 1e-5
@@ -17,7 +18,6 @@ class Sampler(nn.Module):
18
def __init__(self):
19
# TODO(houseroad): Add support for logprobs_mode.
20
super().__init__()
20
- self.topk_topp_sampler = TopKTopPSampler()
21
22
def forward(
23
self,
@@ -65,13 +65,17 @@ class Sampler(nn.Module):
65
logits = self.apply_min_p(logits, sampling_metadata.min_p)
66
67
# Apply top_k and/or top_p.
68
- random_sampled, _ = self.topk_topp_sampler(
68
+ logits = apply_top_k_top_p(
69
logits,
70
- sampling_metadata.generators,
70
sampling_metadata.top_k,
71
sampling_metadata.top_p,
72
)
73
74
+ # Random sample.
75
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
76
+ random_sampled = self.random_sample(probs,
77
+ sampling_metadata.generators)
78
+
79
sampled = torch.where(sampling_metadata.temperature < _SAMPLING_EPS,
80
greedy_sampled, random_sampled)
81
return sampled
@@ -144,3 +148,66 @@ class Sampler(nn.Module):
148
# Apply mask using boolean indexing (xla friendly)
149
logits.masked_fill_(~valid_token_mask, -float("inf"))
150
return logits
151
+
152
+ def random_sample(
153
+ self,
154
+ probs: torch.Tensor,
155
+ generators: dict[int, torch.Generator],
156
+ ) -> torch.Tensor:
157
+ q = torch.empty_like(probs)
158
+ # NOTE(woosuk): To batch-process the requests without their own seeds,
159
+ # which is the common case, we first assume that every request does
160
+ # not have its own seed. Then, we overwrite the values for the requests
161
+ # that have their own seeds.
162
+ q.exponential_()
163
+ if generators:
164
+ for i, generator in generators.items():
165
+ q[i].exponential_(generator=generator)
166
+ return probs.div_(q).argmax(dim=-1).view(-1)
167
+
168
+
169
+def apply_top_k_top_p(
170
+ logits: torch.Tensor,
171
+ k: Optional[torch.Tensor],
172
+ p: Optional[torch.Tensor],
173
+) -> torch.Tensor:
174
+ """
175
+ Apply top-k and top-p optimized for TPU.
176
+
177
+ This algorithm avoids using torch.scatter which is extremely slow on TPU.
178
+ This is achieved by finding a "cut-off" element in the original logit, and
179
+ after thresholding the logit using this cut-off, the remaining elements
180
+ shall constitute the top-p set.
181
+
182
+ Note: in the case of tie (i.e. multipple cut-off elements present in the
183
+ logit), all tie elements are included in the top-p set. In other words,
184
+ this function does not break ties. Instead, these tie tokens have equal
185
+ chance of being chosen during final sampling, so we can consider the tie
186
+ being broken then.
187
+ """
188
+ probs = logits.softmax(dim=-1)
189
+ probs_sort, _ = probs.sort(dim=-1, descending=False)
190
+
191
+ if k is not None:
192
+ top_k_count = probs_sort.size(1) - k.to(torch.long) # shape: (batch, )
193
+ top_k_count = top_k_count.unsqueeze(dim=1)
194
+ top_k_cutoff = probs_sort.gather(-1, top_k_count)
195
+
196
+ # Make sure the no top-k rows are no-op.
197
+ no_top_k_mask = (k == logits.shape[1]).unsqueeze(dim=1)
198
+ top_k_cutoff.masked_fill_(no_top_k_mask, -float("inf"))
199
+
200
+ elements_to_discard = probs < top_k_cutoff
201
+ logits.masked_fill_(elements_to_discard, -float("inf"))
202
+
203
+ if p is not None:
204
+ cumprob = torch.cumsum(probs_sort, dim=-1)
205
+ top_p_mask = cumprob <= 1 - p.unsqueeze(dim=1)
206
+ top_p_mask[:, -1] = False # at least one
207
+
208
+ top_p_count = top_p_mask.sum(dim=-1).unsqueeze(1)
209
+ top_p_cutoff = probs_sort.gather(-1, top_p_count)
210
+ elements_to_discard = probs < top_p_cutoff
211
+ logits.masked_fill_(elements_to_discard, -float("inf"))
212
+
213
+ return logits