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, 07:15 UTC
improve logits bias (#19041)
Yu Guo committed
Jun 6, 2025 at 04:59 UTC
8267f9916f9b4ceace3a8daa908d52824068862b
1 file changed
+16
-1
vllm/v1/sample/sampler.py
+16
-1
@@ -5,6 +5,7 @@
5
import torch
6
import torch.nn as nn
7
8
+from vllm.utils import async_tensor_h2d, is_pin_memory_available
9
from vllm.v1.outputs import LogprobsTensors, SamplerOutput
10
from vllm.v1.sample.metadata import SamplingMetadata
11
from vllm.v1.sample.ops.bad_words import apply_bad_words
@@ -20,6 +21,7 @@ class Sampler(nn.Module):
21
def __init__(self):
22
super().__init__()
23
self.topk_topp_sampler = TopKTopPSampler()
24
+ self.pin_memory = is_pin_memory_available()
25
26
def forward(
27
self,
@@ -232,6 +234,10 @@ class Sampler(nn.Module):
234
# One idea is implement this as a PyTorch C++ op, and we may
235
# even optimize the logit_bias layout.
236
237
+ rows: list[int] = []
238
+ cols: list[int] = []
239
+ vals: list[float] = []
240
+
241
# Get vocabulary size from logits
242
vocab_size = logits.shape[-1]
243
@@ -244,7 +250,16 @@ class Sampler(nn.Module):
250
f"token_id {token_id} in logit_bias contains "
251
f"out-of-vocab token id. Vocabulary size: "
252
f"{vocab_size}")
247
- logits[i, token_id] += bias
253
+ rows.append(i)
254
+ cols.append(token_id)
255
+ vals.append(bias)
256
+
257
+ if rows:
258
+ indices = async_tensor_h2d([rows, cols], torch.int64,
259
+ logits.device, self.pin_memory)
260
+ values = async_tensor_h2d(vals, torch.float, logits.device,
261
+ self.pin_memory)
262
+ logits.index_put_(tuple(indices), values=values, accumulate=True)
263
return logits
264
265
def apply_allowed_token_ids(