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 3, 03:00 UTC
[BugFix] Use correct logprobs for `logprob_token_ids` (#43125)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
May 21, 2026 at 15:42 UTC
565b745ec5d28dafd14585f1b695b159ba336a04
1 file changed
+29
-23
vllm/v1/sample/sampler.py
+29
-23
@@ -13,7 +13,6 @@ from vllm.v1.sample.ops.bad_words import apply_bad_words
13
from vllm.v1.sample.ops.logprobs import batched_count_greater_than
14
from vllm.v1.sample.ops.penalties import apply_all_penalties
15
from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
16
-from vllm.v1.worker.gpu.sample.logprob import compute_token_logprobs
16
17
_SAMPLING_EPS = 1e-5
18
@@ -78,7 +77,8 @@ class Sampler(nn.Module):
77
# This is different from the V0 sampler, which uses the logits that
78
# is used for sampling (after penalties and temperature scaling).
79
num_logprobs = sampling_metadata.max_num_logprobs
81
- if num_logprobs is not None:
80
+ raw_logprobs: torch.Tensor | None = None
81
+ if num_logprobs is not None or sampling_metadata.logprob_token_ids:
82
if logprobs_mode == "raw_logprobs":
83
raw_logprobs = self.compute_logprobs(logits)
84
elif logprobs_mode == "raw_logits":
@@ -107,8 +107,9 @@ class Sampler(nn.Module):
107
# This is used by generative_scoring API to get logprobs for specific tokens
108
logprob_token_ids_tensors = None
109
if sampling_metadata.logprob_token_ids:
110
+ assert raw_logprobs is not None
111
logprob_token_ids_tensors = self.gather_specific_token_logprobs(
111
- logits, sampling_metadata.logprob_token_ids, sampled
112
+ raw_logprobs, sampling_metadata.logprob_token_ids, sampled
113
)
114
115
if num_logprobs is None:
@@ -144,22 +145,19 @@ class Sampler(nn.Module):
145
146
def gather_specific_token_logprobs(
147
self,
147
- logits: torch.Tensor,
148
+ logprobs: torch.Tensor,
149
logprob_token_ids: dict[int, list[int]],
150
sampled: torch.Tensor,
151
) -> LogprobsTensors | None:
151
- """Compute logprobs for specific token IDs using Triton kernel.
152
+ """Gather logprobs for specific token IDs requested per request.
153
153
- This method handles heterogeneous token ID lists across requests by
154
- padding shorter lists to max length and using a fused Triton kernel
155
- for efficient log_softmax + gather computation.
156
-
157
- Benchmarks show the Triton kernel approach is ~1.4x faster than sparse
158
- gather for batch sizes > 1 due to the fused kernel reducing memory
159
- bandwidth requirements.
154
+ Used by the generative_scoring API to return logprobs for an explicit
155
+ set of token ids rather than the top-k. Handles heterogeneous token
156
+ id lists across requests by padding shorter lists to the max length.
157
158
Args:
162
- logits: [batch_size, vocab_size] tensor of logits
159
+ logprobs: [batch_size, vocab_size] tensor of (raw) logprobs to
160
+ gather from.
161
logprob_token_ids: dict mapping req_index -> list of token IDs
162
sampled: [batch_size] tensor of sampled token IDs
163
@@ -170,8 +168,8 @@ class Sampler(nn.Module):
168
if not logprob_token_ids:
169
return None
170
173
- batch_size = logits.shape[0]
174
- device = logits.device
171
+ batch_size = logprobs.shape[0]
172
+ device = logprobs.device
173
174
# Find max number of tokens across all requests
175
max_num_tokens = max(len(tids) for tids in logprob_token_ids.values())
@@ -200,19 +198,24 @@ class Sampler(nn.Module):
198
# tensor so we don't need to D2H + re-upload.
199
token_ids_tensor[:, 0] = sampled
200
203
- # Compute logprobs using the fused Triton kernel (log_softmax + gather)
204
- logprobs = compute_token_logprobs(logits, token_ids_tensor)
201
+ # Gather logprobs at the requested token ids.
202
+ gathered_logprobs = logprobs.gather(-1, token_ids_tensor)
203
204
# Mask invalid (padded) positions with -inf
207
- logprobs = logprobs.masked_fill(~valid_mask, float("-inf"))
205
+ gathered_logprobs = gathered_logprobs.masked_fill(~valid_mask, float("-inf"))
206
209
- # Compute ranks for the sampled token
210
- sampled_logits = logits.gather(-1, sampled.unsqueeze(-1))
211
- token_ranks = (logits > sampled_logits).sum(dim=-1)
207
+ # Compute ranks for the sampled token. log_softmax is monotonic w.r.t.
208
+ # the original logits, so ranks computed from logprobs are equivalent.
209
+ sampled_logprobs = logprobs.gather(-1, sampled.unsqueeze(-1))
210
+ # Avoid 0/1 specialization recompile on the batch dimension of the
211
+ # compiled batched_count_greater_than. See gather_logprobs for context.
212
+ torch._dynamo.decorators.mark_unbacked(logprobs, 0)
213
+ torch._dynamo.decorators.mark_unbacked(sampled_logprobs, 0)
214
+ token_ranks = batched_count_greater_than(logprobs, sampled_logprobs)
215
216
return LogprobsTensors(
217
logprob_token_ids=token_ids_tensor.to(torch.int32),
215
- logprobs=logprobs,
218
+ logprobs=gathered_logprobs,
219
selected_token_ranks=token_ranks,
220
)
221
@@ -252,7 +255,10 @@ class Sampler(nn.Module):
255
greedy_sampled = self.greedy_sample(logits)
256
if sampling_metadata.all_greedy:
257
processed_logprobs = None
255
- if sampling_metadata.max_num_logprobs is not None:
258
+ if (
259
+ sampling_metadata.max_num_logprobs is not None
260
+ or sampling_metadata.logprob_token_ids
261
+ ):
262
if logprobs_mode == "processed_logits":
263
processed_logprobs = logits
264
elif logprobs_mode == "processed_logprobs":