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 9, 18:15 UTC
[Refactor] Remove unused tpu files (#32610)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 22, 2026 at 15:35 UTC
f74481018412e4bc63d0fc396ec675ca4bf9bf18
4 files changed
-353
vllm/v1/sample/tpu/__init__.py
vllm/v1/sample/tpu/metadata.py
deleted
-120
@@ -1,120 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass, field
4
-
5
-import torch
6
-
7
-from vllm.v1.worker.tpu_input_batch import InputBatch
8
-
9
-DEFAULT_SAMPLING_PARAMS = dict(
10
- temperature=-1.0,
11
- min_p=0.0,
12
- # strictly disabled for now
13
- top_k=0,
14
- top_p=1.0,
15
- # frequency_penalties=0.0,
16
- # presence_penalties=0.0,
17
- # repetition_penalties=0.0,
18
-)
19
-
20
-
21
-@dataclass
22
-class TPUSupportedSamplingMetadata:
23
- # This class exposes a more xla-friendly interface than SamplingMetadata
24
- # on TPU, in particular all arguments should be traceable and no optionals
25
- # are allowed, to avoid graph recompilation on Nones.
26
- temperature: torch.Tensor = None
27
-
28
- min_p: torch.Tensor = None
29
- top_k: torch.Tensor = None
30
- top_p: torch.Tensor = None
31
-
32
- all_greedy: bool = True
33
- all_random: bool = False
34
-
35
- # Whether logprobs are to be gathered in this batch of request. To balance
36
- # out compile time and runtime, a fixed `max_number_logprobs` value is used
37
- # when gathering logprobs, regardless of the values specified in the batch.
38
- logprobs: bool = False
39
-
40
- # TODO No penalties for now
41
- no_penalties: bool = True
42
- prompt_token_ids = None
43
- frequency_penalties = None
44
- presence_penalties = None
45
- repetition_penalties = None
46
- # should use tensor
47
- output_token_ids: list[list[int]] = field(default_factory=lambda: list())
48
-
49
- min_tokens = None # impl is not vectorized
50
-
51
- logit_bias: list[dict[int, float] | None] = field(default_factory=lambda: list())
52
-
53
- allowed_token_ids_mask = None
54
- bad_words_token_ids = None
55
-
56
- # Generator not supported by xla
57
- _generators: dict[int, torch.Generator] = field(default_factory=lambda: dict())
58
-
59
- @property
60
- def generators(self) -> dict[int, torch.Generator]:
61
- # Generator not supported by torch/xla. This field must be immutable.
62
- return self._generators
63
-
64
- @classmethod
65
- def from_input_batch(
66
- cls,
67
- input_batch: InputBatch,
68
- padded_num_reqs: int,
69
- xla_device: torch.device,
70
- generate_params_if_all_greedy: bool = False,
71
- ) -> "TPUSupportedSamplingMetadata":
72
- """
73
- Copy sampling tensors slices from `input_batch` to on device tensors.
74
-
75
- `InputBatch._make_sampling_metadata` causes recompilation on XLA as it
76
- slices dynamic shapes on device tensors. This impl moves the dynamic
77
- ops to CPU and produces tensors of fixed `padded_num_reqs` size.
78
-
79
- Args:
80
- input_batch: The input batch containing sampling parameters.
81
- padded_num_reqs: The padded number of requests.
82
- xla_device: The XLA device.
83
- generate_params_if_all_greedy: If True, generate sampling parameters
84
- even if all requests are greedy. this is useful for cases where
85
- we want to pre-compile a graph with sampling parameters, even if
86
- they are not strictly needed for greedy decoding.
87
- """
88
- needs_logprobs = (
89
- input_batch.max_num_logprobs > 0 if input_batch.max_num_logprobs else False
90
- )
91
- # Early return to avoid unnecessary cpu to tpu copy
92
- if input_batch.all_greedy is True and generate_params_if_all_greedy is False:
93
- return cls(all_greedy=True, logprobs=needs_logprobs)
94
-
95
- num_reqs = input_batch.num_reqs
96
-
97
- def fill_slice(cpu_tensor: torch.Tensor, fill_val) -> torch.Tensor:
98
- # Pad value is the default one.
99
- cpu_tensor[num_reqs:padded_num_reqs] = fill_val
100
-
101
- fill_slice(
102
- input_batch.temperature_cpu_tensor, DEFAULT_SAMPLING_PARAMS["temperature"]
103
- )
104
- fill_slice(input_batch.min_p_cpu_tensor, DEFAULT_SAMPLING_PARAMS["min_p"])
105
- fill_slice(input_batch.top_k_cpu_tensor, DEFAULT_SAMPLING_PARAMS["top_k"])
106
- fill_slice(input_batch.top_p_cpu_tensor, DEFAULT_SAMPLING_PARAMS["top_p"])
107
-
108
- # Slice persistent device tensors to a fixed pre-compiled padded shape.
109
- return cls(
110
- temperature=input_batch.temperature_cpu_tensor[:padded_num_reqs].to(
111
- xla_device
112
- ),
113
- all_greedy=input_batch.all_greedy,
114
- all_random=input_batch.all_random,
115
- # TODO enable more and avoid returning None values
116
- top_p=input_batch.top_p_cpu_tensor[:padded_num_reqs].to(xla_device),
117
- top_k=input_batch.top_k_cpu_tensor[:padded_num_reqs].to(xla_device),
118
- min_p=input_batch.min_p_cpu_tensor[:padded_num_reqs].to(xla_device),
119
- logprobs=needs_logprobs,
120
- )
vllm/v1/sample/tpu/sampler.py
deleted
-215
@@ -1,215 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Sampler layer implementing TPU supported operations."""
4
-
5
-import torch
6
-import torch.nn as nn
7
-
8
-from vllm.v1.outputs import LogprobsTensors, SamplerOutput
9
-from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
10
-
11
-_SAMPLING_EPS = 1e-5
12
-
13
-
14
-class Sampler(nn.Module):
15
- def __init__(self):
16
- # TODO(houseroad): Add support for logprobs_mode.
17
- super().__init__()
18
-
19
- def forward(
20
- self,
21
- logits: torch.Tensor,
22
- sampling_metadata: TPUSupportedSamplingMetadata,
23
- ) -> SamplerOutput:
24
- # Use float32 for the logits.
25
- logits = logits.to(torch.float32)
26
- # Sample the next token.
27
- sampled = self.sample(logits, sampling_metadata)
28
-
29
- # These are TPU tensors.
30
- sampler_output = SamplerOutput(
31
- # The sampled tokens are expanded to 2D tensor with shape
32
- # [num_requests, 1], where each row represents one generated
33
- # token per request.
34
- sampled_token_ids=sampled.unsqueeze(-1),
35
- logprobs_tensors=None,
36
- )
37
- return sampler_output
38
-
39
- def apply_temperature(
40
- self,
41
- logits: torch.Tensor,
42
- temp: torch.Tensor,
43
- all_random: bool = False,
44
- ) -> torch.Tensor:
45
- # Avoid division by zero for greedy sampling (temperature ~ 0.0).
46
- if not all_random:
47
- temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
48
- return logits.div_(temp.unsqueeze(dim=1))
49
-
50
- def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
51
- return logits.argmax(dim=-1).view(-1)
52
-
53
- def sample(
54
- self,
55
- logits: torch.Tensor,
56
- sampling_metadata: TPUSupportedSamplingMetadata,
57
- ) -> torch.Tensor:
58
- greedy_sampled = self.greedy_sample(logits)
59
-
60
- assert sampling_metadata.temperature is not None
61
-
62
- # Apply temperature.
63
- logits = self.apply_temperature(
64
- logits, sampling_metadata.temperature, sampling_metadata.all_random
65
- )
66
-
67
- # Apply min_p.
68
- if sampling_metadata.min_p is not None:
69
- logits = self.apply_min_p(logits, sampling_metadata.min_p)
70
-
71
- # Apply top_k and/or top_p.
72
- logits = apply_top_k_top_p(
73
- logits,
74
- sampling_metadata.top_k,
75
- sampling_metadata.top_p,
76
- )
77
-
78
- # Random sample.
79
- probs = logits.softmax(dim=-1, dtype=torch.float32)
80
- random_sampled = self.random_sample(probs, sampling_metadata.generators)
81
-
82
- sampled = torch.where(
83
- sampling_metadata.temperature < _SAMPLING_EPS,
84
- greedy_sampled,
85
- random_sampled,
86
- )
87
- return sampled
88
-
89
- def compute_logprobs(self, logits: torch.Tensor) -> torch.Tensor:
90
- return logits.log_softmax(dim=-1, dtype=torch.float32)
91
-
92
- def gather_logprobs(
93
- self,
94
- logprobs: torch.Tensor,
95
- num_logprobs: int,
96
- token_ids: torch.Tensor,
97
- ) -> LogprobsTensors:
98
- """
99
- Gather logprobs for topk and sampled/prompt token.
100
-
101
- Args:
102
- logprobs: (num tokens) x (vocab) tensor
103
- num_logprobs: minimum number of logprobs to
104
- retain per token
105
- token_ids: prompt tokens (if prompt logprobs)
106
- or sampled tokens (if sampled
107
- logprobs); 1D token ID tensor
108
- with (num tokens) elements
109
-
110
- Returns:
111
- Top-k int indices tensor, (num tokens) x (num_logprobs + 1)
112
- Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1)
113
- Sampled token rank tensor, (num tokens)
114
- """
115
- # Find the topK values.
116
- topk_logprobs, topk_indices = torch.topk(logprobs, num_logprobs, dim=-1)
117
-
118
- # Get with the logprob of the prompt or sampled token.
119
- token_ids = token_ids.unsqueeze(-1)
120
- token_logprobs = logprobs.gather(-1, token_ids)
121
-
122
- # Compute the ranks of the actual token.
123
- token_ranks = (logprobs >= token_logprobs).sum(-1)
124
-
125
- # Concatenate together with the topk.
126
- indices = torch.cat((token_ids, topk_indices), dim=1)
127
- logprobs = torch.cat((token_logprobs, topk_logprobs), dim=1)
128
-
129
- # Use int32 to reduce the tensor size.
130
- indices = indices.to(torch.int32)
131
-
132
- return LogprobsTensors(indices, logprobs, token_ranks)
133
-
134
- def apply_min_p(
135
- self,
136
- logits: torch.Tensor,
137
- min_p: torch.Tensor,
138
- ) -> torch.Tensor:
139
- """
140
- Filters logits using adaptive probability thresholding.
141
- """
142
- # Convert logits to probability distribution
143
- probability_values = torch.nn.functional.softmax(logits, dim=-1)
144
- # Calculate maximum probabilities per sequence
145
- max_probabilities = torch.amax(probability_values, dim=-1, keepdim=True)
146
- # Reshape min_p for broadcasting
147
- adjusted_min_p = min_p.unsqueeze(1) * max_probabilities
148
- # Identify valid tokens using threshold comparison
149
- valid_token_mask = probability_values >= adjusted_min_p
150
- # Apply mask using boolean indexing (xla friendly)
151
- logits.masked_fill_(~valid_token_mask, -float("inf"))
152
- return logits
153
-
154
- def random_sample(
155
- self,
156
- probs: torch.Tensor,
157
- generators: dict[int, torch.Generator],
158
- ) -> torch.Tensor:
159
- q = torch.empty_like(probs)
160
- # NOTE(woosuk): To batch-process the requests without their own seeds,
161
- # which is the common case, we first assume that every request does
162
- # not have its own seed. Then, we overwrite the values for the requests
163
- # that have their own seeds.
164
- q.exponential_()
165
- if generators:
166
- for i, generator in generators.items():
167
- q[i].exponential_(generator=generator)
168
- return probs.div_(q).argmax(dim=-1).view(-1)
169
-
170
-
171
-def apply_top_k_top_p(
172
- logits: torch.Tensor,
173
- k: torch.Tensor | None,
174
- p: torch.Tensor | None,
175
-) -> torch.Tensor:
176
- """
177
- Apply top-k and top-p optimized for TPU.
178
-
179
- This algorithm avoids using torch.scatter which is extremely slow on TPU.
180
- This is achieved by finding a "cut-off" element in the original logit, and
181
- after thresholding the logit using this cut-off, the remaining elements
182
- shall constitute the top-p set.
183
-
184
- Note: in the case of tie (i.e. multiple cut-off elements present in the
185
- logit), all tie elements are included in the top-p set. In other words,
186
- this function does not break ties. Instead, these tie tokens have equal
187
- chance of being chosen during final sampling, so we can consider the tie
188
- being broken then.
189
- """
190
- probs = logits.softmax(dim=-1)
191
- probs_sort, _ = probs.sort(dim=-1, descending=False)
192
-
193
- if k is not None:
194
- top_k_count = probs_sort.size(1) - k.to(torch.long) # shape: (batch, )
195
- top_k_count = top_k_count.unsqueeze(dim=1)
196
- top_k_cutoff = probs_sort.gather(-1, top_k_count)
197
-
198
- # Make sure the no top-k rows are no-op.
199
- no_top_k_mask = (k == logits.shape[1]).unsqueeze(dim=1)
200
- top_k_cutoff.masked_fill_(no_top_k_mask, -float("inf"))
201
-
202
- elements_to_discard = probs < top_k_cutoff
203
- logits.masked_fill_(elements_to_discard, -float("inf"))
204
-
205
- if p is not None:
206
- cumprob = torch.cumsum(probs_sort, dim=-1)
207
- top_p_mask = cumprob <= 1 - p.unsqueeze(dim=1)
208
- top_p_mask[:, -1] = False # at least one
209
-
210
- top_p_count = top_p_mask.sum(dim=-1).unsqueeze(1)
211
- top_p_cutoff = probs_sort.gather(-1, top_p_count)
212
- elements_to_discard = probs < top_p_cutoff
213
- logits.masked_fill_(elements_to_discard, -float("inf"))
214
-
215
- return logits
vllm/v1/worker/tpu_worker.py
deleted
-18
@@ -1,18 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""A TPU worker class."""
4
-
5
-from typing import TypeVar
6
-
7
-from vllm.logger import init_logger
8
-from vllm.platforms.tpu import USE_TPU_INFERENCE
9
-
10
-logger = init_logger(__name__)
11
-
12
-_R = TypeVar("_R")
13
-
14
-# TODO(weiyulin) Remove this file after adding an official way to use hardware plugin
15
-if USE_TPU_INFERENCE:
16
- from tpu_inference.worker.tpu_worker import TPUWorker as TpuInferenceWorker
17
-
18
- TPUWorker = TpuInferenceWorker # type: ignore