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 12, 04:45 UTC
[Perf] Avoid pageable HtoD transfer in MinTokensLogitsProcessor (#29826)
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
jthomson04 committed
Dec 2, 2025 at 13:25 UTC
1528e079e2b2cf8a807e4dce86ef05540e16a430
1 file changed
+6
-2
vllm/v1/sample/logits_processor/builtin.py
+6
-2
@@ -110,7 +110,7 @@ class MinPLogitsProcessor(LogitsProcessor):
110
# Identify valid tokens using threshold comparison
111
invalid_token_mask = probability_values < adjusted_min_p
112
# Apply mask using boolean indexing
113
- logits[invalid_token_mask] = -float("inf")
113
+ logits.masked_fill_(invalid_token_mask, -float("inf"))
114
return logits
115
116
@@ -178,6 +178,10 @@ class MinTokensLogitsProcessor(LogitsProcessor):
178
self._device_tensor([], torch.int32),
179
)
180
181
+ self.neg_inf_tensor = torch.tensor(
182
+ -float("inf"), dtype=torch.float32, device=self.device
183
+ )
184
+
185
def is_argmax_invariant(self) -> bool:
186
"""By censoring stop tokens, min-tokens can change the outcome
187
of the argmax operation in greedy sampling."""
@@ -229,7 +233,7 @@ class MinTokensLogitsProcessor(LogitsProcessor):
233
def apply(self, logits: torch.Tensor) -> torch.Tensor:
234
if self.min_toks:
235
# Inhibit EOS token for requests which have not reached min length
232
- logits[self.logits_slice] = -float("inf")
236
+ logits.index_put_(self.logits_slice, self.neg_inf_tensor)
237
return logits
238
239