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, 16:00 UTC
[V1][TPU] TPU-optimized top-p implementation (avoids scattering). (#15736)
Signed-off-by: Hyesoo Yang <hyeygit@gmail.com> Co-authored-by: root <root@t1v-n-822696b7-w-0.us-central2-b.c.tpu-prod-env-large-adhoc.internal>
Hyesoo Yang committed
Apr 2, 2025 at 17:18 UTC
1b84eff03adc04bdde9b2378a751d862a28326eb
3 files changed
+174
-15
.buildkite/run-tpu-v1-test.sh
+3
-1
@@ -36,7 +36,9 @@ docker run --privileged --net host --shm-size=16G -it \
36
&& echo TEST_6 \
37
&& pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py \
38
&& echo TEST_7 \
39
- && pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py" \
39
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py \
40
+ && echo TEST_8 \
41
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py" \
42
43
44
# TODO: This test fails because it uses RANDOM_SEED sampling
tests/v1/tpu/test_topk_topp_sampler.py
new
+132
@@ -0,0 +1,132 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import math
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_tpu
9
+
10
+if not current_platform.is_tpu():
11
+ pytest.skip("This test needs a TPU.", allow_module_level=True)
12
+import torch_xla.core.xla_model as xm
13
+
14
+BATCH_SIZE = 1024
15
+VOCAB_SIZE = 128 * 1024
16
+TOLERANCE = 1e-6
17
+
18
+
19
+def test_topp_result_sums_past_p():
20
+ with torch.device(xm.xla_device()):
21
+ xm.set_rng_state(seed=33)
22
+
23
+ logits = torch.rand((BATCH_SIZE, VOCAB_SIZE))
24
+ probs = logits.softmax(dim=-1)
25
+
26
+ # Random top-p values between 0 and 1.
27
+ p = torch.rand((BATCH_SIZE, ))
28
+
29
+ # Set p=1 for ~50% of requests in the batch (top-p disabled).
30
+ p.masked_fill_(torch.randint(0, 2, (BATCH_SIZE, ), dtype=bool), 1)
31
+
32
+ no_op_k = torch.tensor([VOCAB_SIZE])
33
+ logits_masked = apply_top_k_top_p_tpu(logits=logits.clone(),
34
+ k=no_op_k,
35
+ p=p)
36
+
37
+ # Verify that the masked logit's probability sums to at least p.
38
+ probs.masked_fill_(logits_masked.isinf(), 0)
39
+ masked_prob_sum = probs.sum(dim=-1)
40
+
41
+ xm.mark_step()
42
+
43
+ # Perform assertion on CPU.
44
+ assert torch.all(torch.ge(masked_prob_sum.cpu() + TOLERANCE, p.cpu()))
45
+
46
+
47
+def test_topp_basic():
48
+ with torch.device(xm.xla_device()):
49
+ logits = torch.tensor([[math.log(0.2),
50
+ math.log(0.3),
51
+ math.log(0.5)],
52
+ [math.log(0.5),
53
+ math.log(0.1),
54
+ math.log(0.4)]])
55
+
56
+ result = apply_top_k_top_p_tpu(logits=logits.clone(),
57
+ k=torch.tensor([3, 3]),
58
+ p=torch.tensor([0.79, 0.79]))
59
+
60
+ xm.mark_step()
61
+
62
+ # Expect the smallest elements to be dropped.
63
+ expected_result = logits.clone().cpu()
64
+ expected_result[0, 0] = float("-inf")
65
+ expected_result[1, 1] = float("-inf")
66
+ assert torch.allclose(expected_result, result.cpu())
67
+
68
+
69
+def test_topp_select_all():
70
+ with torch.device(xm.xla_device()):
71
+ logits = torch.tensor([[math.log(0.2),
72
+ math.log(0.3),
73
+ math.log(0.5)],
74
+ [math.log(0.5),
75
+ math.log(0.1),
76
+ math.log(0.4)]])
77
+
78
+ result = apply_top_k_top_p_tpu(logits=logits.clone(),
79
+ k=torch.tensor([3, 3]),
80
+ p=torch.tensor([1.0, 1.0]))
81
+
82
+ xm.mark_step()
83
+
84
+ assert torch.allclose(logits.cpu(), result.cpu())
85
+
86
+
87
+def test_topp_with_ties():
88
+ with torch.device(xm.xla_device()):
89
+ # Input has multiple math.log(0.3).
90
+ logits = torch.tensor(
91
+ [[math.log(0.3),
92
+ math.log(0.3),
93
+ math.log(0.3),
94
+ math.log(0.1)]])
95
+
96
+ result = apply_top_k_top_p_tpu(logits=logits.clone(),
97
+ k=torch.tensor([4]),
98
+ p=torch.tensor([0.2]))
99
+
100
+ xm.mark_step()
101
+
102
+ # All tie values are included in the top-p set. Tie breaking is left
103
+ # to be done during final sampling (all tie tokens have equal
104
+ # probability of being chosen).
105
+ expected_result = logits.clone().cpu()
106
+ expected_result[0, 3] = float("-inf")
107
+ assert torch.allclose(expected_result, result.cpu())
108
+
109
+
110
+def test_both_topk_topp():
111
+ with torch.device(xm.xla_device()):
112
+ logits = torch.tensor([[math.log(0.2),
113
+ math.log(0.3),
114
+ math.log(0.5)],
115
+ [math.log(0.5),
116
+ math.log(0.1),
117
+ math.log(0.4)]])
118
+
119
+ # Set k=1 for the first batch.
120
+ result = apply_top_k_top_p_tpu(logits=logits.clone(),
121
+ k=torch.tensor([1, 3]),
122
+ p=torch.tensor([0.79, 0.79]))
123
+
124
+ xm.mark_step()
125
+
126
+ # Since for the first batch k=1, expect only the largest element gets
127
+ # selected.
128
+ expected_result = logits.clone().cpu()
129
+ expected_result[0, 0] = float("-inf")
130
+ expected_result[0, 1] = float("-inf")
131
+ expected_result[1, 1] = float("-inf")
132
+ assert torch.allclose(expected_result, result.cpu())
vllm/v1/sample/ops/topk_topp_sampler.py
+39
-14
@@ -122,23 +122,48 @@ class TopKTopPSampler(nn.Module):
122
k: Optional[torch.Tensor],
123
p: Optional[torch.Tensor],
124
) -> torch.Tensor:
125
- # If only top-k is specified, use pytorch's builtin topk op. This leads
126
- # to significant speed up on TPU compared to using apply_top_k_top_p.
127
- if k is not None and p is None:
128
- topk_values, topk_indices = torch.topk(logits, k, dim=-1)
129
-
130
- mask = torch.ones_like(logits, dtype=torch.bool)
131
- mask.scatter_(-1, topk_indices, False)
132
- logits.masked_fill_(mask, float('-inf'))
133
- else:
134
- # TODO Placeholder for TPU optimized topp kernel
135
- # logits = apply_top_k_top_p(logits, k, p)
136
- pass
137
-
125
+ logits = apply_top_k_top_p_tpu(logits, k, p)
126
probs = logits.softmax(dim=-1, dtype=torch.float32)
127
return random_sample(probs, generators)
128
129
130
+def apply_top_k_top_p_tpu(
131
+ logits: torch.Tensor,
132
+ k: torch.Tensor,
133
+ p: torch.Tensor,
134
+) -> torch.Tensor:
135
+ """
136
+ Apply top-k and top-p optimized for TPU.
137
+
138
+ This algorithm avoids using torch.scatter which is extremely slow on TPU.
139
+ This is achieved by finding a "cut-off" element in the original logit, and
140
+ after thresholding the logit using this cut-off, the remaining elements
141
+ shall constitute the top-p set.
142
+
143
+ Note: in the case of tie (i.e. multipple cut-off elements present in the
144
+ logit), all tie elements are included in the top-p set. In other words,
145
+ this function does not break ties. Instead, these tie tokens have equal
146
+ chance of being chosen during final sampling, so we can consider the tie
147
+ being broken then.
148
+ """
149
+ if k is not None:
150
+ logits = apply_top_k_only(logits, k)
151
+
152
+ if p is not None:
153
+ probs = logits.softmax(dim=-1)
154
+ probs_sort, _ = probs.sort(dim=-1, descending=False)
155
+ cumprob = torch.cumsum(probs_sort, dim=-1)
156
+ top_p_mask = cumprob <= 1 - p.unsqueeze(dim=1)
157
+ top_p_mask[:, -1] = False # at least one
158
+
159
+ top_p_count = top_p_mask.sum(dim=-1).unsqueeze(1)
160
+ top_p_cutoff = probs_sort.gather(-1, top_p_count)
161
+ elements_to_discard = probs < top_p_cutoff
162
+ logits.masked_fill_(elements_to_discard, -float("inf"))
163
+
164
+ return logits
165
+
166
+
167
def apply_top_k_top_p(
168
logits: torch.Tensor,
169
k: Optional[torch.Tensor],
@@ -199,7 +224,7 @@ def apply_top_k_only(
224
max_top_k = k.max()
225
# topk.values tensor has shape [batch_size, max_top_k].
226
# Convert top k to 0-based index in range [0, max_top_k).
202
- k_index = k.sub_(1).unsqueeze(1)
227
+ k_index = k.sub_(1).unsqueeze(1).expand(logits.shape[0], 1)
228
top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index.long())
229
# Handle non-topk rows.
230
top_k_mask.masked_fill_(no_top_k_mask.unsqueeze(1), -float("inf"))