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
[Misc] Misc code simplifications (#26450)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Oct 9, 2025 at 02:10 UTC
ddcbc2f33468138b1cbbd357e669d7ed7f63f5d6
6 files changed
+79
-90
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -1474,7 +1474,7 @@ class Scheduler(SchedulerInterface):
1474
1475
affected_req_ids.add(request.request_id)
1476
1477
- return (affected_req_ids, total_affected_tokens)
1477
+ return affected_req_ids, total_affected_tokens
1478
1479
def _handle_invalid_blocks(self, invalid_block_ids: set[int]) -> set[str]:
1480
total_requests_to_reschedule = 0
vllm/v1/core/sched/utils.py
+1
-2
@@ -59,8 +59,7 @@ def check_stop(
59
sampling_params = request.sampling_params
60
assert sampling_params is not None
61
62
- min_tokens = sampling_params.min_tokens
63
- if request.num_output_tokens < min_tokens:
62
+ if request.num_output_tokens < sampling_params.min_tokens:
63
return False
64
65
last_token_id = request.output_token_ids[-1]
vllm/v1/sample/rejection_sampler.py
+8
-13
@@ -147,22 +147,20 @@ class RejectionSampler(nn.Module):
147
sampling_metadata: SamplingMetadata,
148
metadata: SpecDecodeMetadata,
149
) -> torch.Tensor:
150
+ has_penalties = not sampling_metadata.no_penalties
151
any_penalties_or_bad_words = (
151
- sampling_metadata.bad_words_token_ids or not sampling_metadata.no_penalties
152
+ sampling_metadata.bad_words_token_ids or has_penalties
153
)
154
155
output_token_ids = sampling_metadata.output_token_ids
156
if any_penalties_or_bad_words:
157
output_token_ids = self._combine_outputs_with_spec_tokens(
157
- sampling_metadata.output_token_ids,
158
+ output_token_ids,
159
sampling_metadata.spec_token_ids,
160
)
161
162
# Calculate indices of target logits.
162
- if (
163
- sampling_metadata.allowed_token_ids_mask is not None
164
- or not sampling_metadata.no_penalties
165
- ):
163
+ if sampling_metadata.allowed_token_ids_mask is not None or has_penalties:
164
num_requests = len(sampling_metadata.output_token_ids)
165
num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu")
166
original_indices = torch.arange(num_requests, device="cpu")
@@ -180,18 +178,15 @@ class RejectionSampler(nn.Module):
178
logits.masked_fill_(token_mask, float("-inf"))
179
180
# Apply bad words exclusion.
183
- if sampling_metadata.bad_words_token_ids:
181
+ if bad_words_token_ids := sampling_metadata.bad_words_token_ids:
182
apply_bad_words_with_drafts(
185
- logits,
186
- sampling_metadata.bad_words_token_ids,
187
- output_token_ids,
188
- metadata.num_draft_tokens,
183
+ logits, bad_words_token_ids, output_token_ids, metadata.num_draft_tokens
184
)
185
186
return logits
187
188
+ @staticmethod
189
def apply_penalties(
194
- self,
190
logits: torch.Tensor,
191
sampling_metadata: SamplingMetadata,
192
metadata: SpecDecodeMetadata,
@@ -218,8 +213,8 @@ class RejectionSampler(nn.Module):
213
)
214
return logits
215
216
+ @staticmethod
217
def _combine_outputs_with_spec_tokens(
222
- self,
218
output_token_ids: list[list[int]],
219
spec_token_ids: Optional[list[list[int]]] = None,
220
) -> list[list[int]]:
vllm/v1/sample/sampler.py
+26
-30
@@ -120,8 +120,8 @@ class Sampler(nn.Module):
120
)
121
return sampler_output
122
123
+ @staticmethod
124
def apply_temperature(
124
- self,
125
logits: torch.Tensor,
126
temp: torch.Tensor,
127
all_random: bool,
@@ -132,7 +132,8 @@ class Sampler(nn.Module):
132
temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
133
return logits.div_(temp.unsqueeze(dim=1))
134
135
- def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
135
+ @staticmethod
136
+ def greedy_sample(logits: torch.Tensor) -> torch.Tensor:
137
return logits.argmax(dim=-1).view(-1)
138
139
def sample(
@@ -191,11 +192,12 @@ class Sampler(nn.Module):
192
)
193
return sampled, processed_logprobs
194
194
- def compute_logprobs(self, logits: torch.Tensor) -> torch.Tensor:
195
+ @staticmethod
196
+ def compute_logprobs(logits: torch.Tensor) -> torch.Tensor:
197
return logits.log_softmax(dim=-1, dtype=torch.float32)
198
199
+ @staticmethod
200
def gather_logprobs(
198
- self,
201
logprobs: torch.Tensor,
202
num_logprobs: int,
203
token_ids: torch.Tensor,
@@ -238,8 +240,8 @@ class Sampler(nn.Module):
240
241
return LogprobsTensors(indices, logprobs, token_ranks)
242
243
+ @staticmethod
244
def _combine_outputs_with_spec_tokens(
242
- self,
245
output_token_ids: list[list[int]],
246
spec_token_ids: Optional[list[list[int]]] = None,
247
) -> list[list[int]]:
@@ -257,8 +259,9 @@ class Sampler(nn.Module):
259
sampling_metadata: SamplingMetadata,
260
predict_bonus_token: bool,
261
) -> torch.Tensor:
262
+ bad_words_token_ids = sampling_metadata.bad_words_token_ids
263
any_penalties_or_bad_words = (
261
- sampling_metadata.bad_words_token_ids or not sampling_metadata.no_penalties
264
+ bool(bad_words_token_ids) or not sampling_metadata.no_penalties
265
)
266
267
output_token_ids = sampling_metadata.output_token_ids
@@ -266,7 +269,7 @@ class Sampler(nn.Module):
269
# Combine base outputs with spec tokens when speculative decoding
270
# is enabled.
271
output_token_ids = self._combine_outputs_with_spec_tokens(
269
- sampling_metadata.output_token_ids,
272
+ output_token_ids,
273
sampling_metadata.spec_token_ids,
274
)
275
@@ -275,14 +278,8 @@ class Sampler(nn.Module):
278
logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf"))
279
280
# Apply bad words exclusion.
278
- if sampling_metadata.bad_words_token_ids:
279
- apply_bad_words(
280
- logits,
281
- sampling_metadata.bad_words_token_ids,
282
- output_token_ids
283
- if output_token_ids is not None
284
- else sampling_metadata.output_token_ids,
285
- )
281
+ if bad_words_token_ids:
282
+ apply_bad_words(logits, bad_words_token_ids, output_token_ids)
283
284
# Apply logits processors which can impact greedy sampling.
285
for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
@@ -292,22 +289,21 @@ class Sampler(nn.Module):
289
logits = self.apply_penalties(logits, sampling_metadata, output_token_ids)
290
return logits
291
292
+ @staticmethod
293
def apply_penalties(
296
- self,
294
logits: torch.Tensor,
295
sampling_metadata: SamplingMetadata,
299
- output_token_ids: Optional[list[list[int]]] = None,
296
+ output_token_ids: list[list[int]],
297
) -> torch.Tensor:
301
- if not sampling_metadata.no_penalties:
302
- assert sampling_metadata.prompt_token_ids is not None
303
- logits = apply_all_penalties(
304
- logits,
305
- sampling_metadata.prompt_token_ids,
306
- sampling_metadata.presence_penalties,
307
- sampling_metadata.frequency_penalties,
308
- sampling_metadata.repetition_penalties,
309
- output_token_ids
310
- if output_token_ids is not None
311
- else sampling_metadata.output_token_ids,
312
- )
313
- return logits
298
+ if sampling_metadata.no_penalties:
299
+ return logits
300
+
301
+ assert sampling_metadata.prompt_token_ids is not None
302
+ return apply_all_penalties(
303
+ logits,
304
+ sampling_metadata.prompt_token_ids,
305
+ sampling_metadata.presence_penalties,
306
+ sampling_metadata.frequency_penalties,
307
+ sampling_metadata.repetition_penalties,
308
+ output_token_ids,
309
+ )
vllm/v1/worker/gpu_input_batch.py
+9
-11
@@ -62,10 +62,9 @@ class CachedRequestState:
62
"provided via prompt_embeds, and its ID is unknown."
63
)
64
return self.prompt_token_ids[idx]
65
- elif idx - self.num_prompt_tokens < len(self.output_token_ids):
65
+ if idx - self.num_prompt_tokens < len(self.output_token_ids):
66
return self.output_token_ids[idx - self.num_prompt_tokens]
67
- else:
68
- return -1
67
+ return -1
68
69
70
class InputBatch:
@@ -770,14 +769,13 @@ class InputBatch:
769
not self.no_penalties
770
or self.logits_processing_needs_token_ids[:num_reqs].any()
771
)
773
- if needs_prompt_token_ids:
774
- # The prompt tokens are used only for applying penalties or
775
- # step pooling during the sampling/pooling process.
776
- # Hence copy these tensors only when there are requests which
777
- # need penalties/step_pooler to be applied.
778
- prompt_token_ids = self._make_prompt_token_ids_tensor()
779
- else:
780
- prompt_token_ids = None
772
+ # The prompt tokens are used only for applying penalties or
773
+ # step pooling during the sampling/pooling process.
774
+ # Hence copy these tensors only when there are requests which
775
+ # need penalties/step_pooler to be applied.
776
+ prompt_token_ids = (
777
+ self._make_prompt_token_ids_tensor() if needs_prompt_token_ids else None
778
+ )
779
780
allowed_token_ids_mask: Optional[torch.Tensor] = None
781
if not self.no_allowed_token_ids:
vllm/v1/worker/gpu_model_runner.py
+34
-33
@@ -1996,7 +1996,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1996
# Should be called after attention metadata creation. This just pads
1997
# the second ubatch slice out to the total number of tokens
1998
# (num_tokens + padding)
1999
- def pad_out_ubatch_slice(self, ubatch_slices: UBatchSlices, num_total_tokens: int):
1999
+ @staticmethod
2000
+ def pad_out_ubatch_slice(ubatch_slices: UBatchSlices, num_total_tokens: int):
2001
padded_second_ubatch_slice = slice(
2002
ubatch_slices[1].token_slice.start, num_total_tokens
2003
)
@@ -2085,12 +2086,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2086
dict[str, Any],
2087
]:
2088
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2089
+ is_first_rank = get_pp_group().is_first_rank
2090
2091
# _prepare_inputs may reorder the batch, so we must gather multi
2092
# modal outputs after that to ensure the correct order
2093
if (
2094
self.supports_mm_inputs
2093
- and get_pp_group().is_first_rank
2095
+ and is_first_rank
2096
and not self.model_config.is_encoder_decoder
2097
):
2098
# Run the multimodal encoder if any.
@@ -2115,7 +2117,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2117
**self._init_model_kwargs(num_scheduled_tokens),
2118
**self._extract_mm_kwargs(scheduler_output),
2119
}
2118
- elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
2120
+ elif self.enable_prompt_embeds and is_first_rank:
2121
# Get the input embeddings for the tokens that are not input embeds,
2122
# then put them into the appropriate positions.
2123
# TODO(qthequartermasterman): Since even when prompt embeds are
@@ -2155,7 +2157,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2157
else:
2158
positions = self.positions.gpu[:num_input_tokens]
2159
2158
- if get_pp_group().is_first_rank:
2160
+ if is_first_rank:
2161
intermediate_tensors = None
2162
else:
2163
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
@@ -2186,38 +2188,37 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2188
# Sample the next token and get logprobs if needed.
2189
sampling_metadata = self.input_batch.sampling_metadata
2190
if spec_decode_metadata is None:
2189
- sampler_output = self.sampler(
2191
+ return self.sampler(
2192
logits=logits,
2193
sampling_metadata=sampling_metadata,
2194
)
2193
- else:
2194
- # When indexing with a tensor (bonus_logits_indices), PyTorch
2195
- # creates a new tensor with separate storage from the original
2196
- # logits tensor. This means any in-place operations on bonus_logits
2197
- # won't affect the original logits tensor.
2198
- assert logits is not None
2199
- bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
2200
- sampler_output = self.sampler(
2201
- logits=bonus_logits,
2202
- sampling_metadata=sampling_metadata,
2203
- predict_bonus_token=True,
2204
- )
2205
- bonus_token_ids = sampler_output.sampled_token_ids
2206
-
2207
- # Just like `bonus_logits`, `target_logits` is a new tensor with
2208
- # separate storage from the original `logits` tensor. Therefore,
2209
- # it is safe to update `target_logits` in place.
2210
- target_logits = logits[spec_decode_metadata.target_logits_indices]
2211
- output_token_ids = self.rejection_sampler(
2212
- spec_decode_metadata,
2213
- None, # draft_probs
2214
- target_logits,
2215
- bonus_token_ids,
2216
- sampling_metadata,
2217
- )
2218
- sampler_output.sampled_token_ids = output_token_ids
2219
- self._update_states_after_model_execute(output_token_ids)
2195
2196
+ # When indexing with a tensor (bonus_logits_indices), PyTorch
2197
+ # creates a new tensor with separate storage from the original
2198
+ # logits tensor. This means any in-place operations on bonus_logits
2199
+ # won't affect the original logits tensor.
2200
+ assert logits is not None
2201
+ bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
2202
+ sampler_output = self.sampler(
2203
+ logits=bonus_logits,
2204
+ sampling_metadata=sampling_metadata,
2205
+ predict_bonus_token=True,
2206
+ )
2207
+ bonus_token_ids = sampler_output.sampled_token_ids
2208
+
2209
+ # Just like `bonus_logits`, `target_logits` is a new tensor with
2210
+ # separate storage from the original `logits` tensor. Therefore,
2211
+ # it is safe to update `target_logits` in place.
2212
+ target_logits = logits[spec_decode_metadata.target_logits_indices]
2213
+ output_token_ids = self.rejection_sampler(
2214
+ spec_decode_metadata,
2215
+ None, # draft_probs
2216
+ target_logits,
2217
+ bonus_token_ids,
2218
+ sampling_metadata,
2219
+ )
2220
+ sampler_output.sampled_token_ids = output_token_ids
2221
+ self._update_states_after_model_execute(output_token_ids)
2222
return sampler_output
2223
2224
def _bookkeeping_sync(
@@ -3741,7 +3742,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3742
decode_cudagraph_batch_sizes = [
3743
x
3744
for x in self.cudagraph_batch_sizes
3744
- if x <= max_num_tokens and x >= self.uniform_decode_query_len
3745
+ if max_num_tokens >= x >= self.uniform_decode_query_len
3746
]
3747
compilation_cases_decode = list(reversed(decode_cudagraph_batch_sizes))
3748
self._capture_cudagraphs(