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 4, 08:45 UTC
[Model Runner V2] Add `logprob_token_ids` support (#40559)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Wentao Ye committed
May 1, 2026 at 13:02 UTC
51295793a2eed0eefc7505cb9a7d5f96effd7773
5 files changed
+179
-14
vllm/sampling_params.py
+25
@@ -25,6 +25,10 @@ logger = init_logger(__name__)
25
_SAMPLING_EPS = 1e-5
26
_MAX_TEMP = 1e-2
27
28
+MAX_LOGPROB_TOKEN_IDS = 128
29
+"""Upper bound on `SamplingParams.logprob_token_ids` list length. Must match
30
+the per-request row width allocated by the sampler's `LogprobTokenIdsState`."""
31
+
32
33
class SamplingType(IntEnum):
34
GREEDY = 0
@@ -628,6 +632,16 @@ class SamplingParams(
632
# For internal use only. Backward compatibility not guaranteed
633
return self._bad_words_token_ids
634
635
+ @property
636
+ def num_logprobs(self) -> int | None:
637
+ """Number of sample logprobs to return per output token, or `None` if
638
+ no sample logprobs were requested. Takes `logprob_token_ids` into
639
+ account: when `logprobs` is unset but `logprob_token_ids` is set,
640
+ returns `len(logprob_token_ids)`."""
641
+ if self.logprobs is not None:
642
+ return self.logprobs
643
+ return len(self.logprob_token_ids) if self.logprob_token_ids else None
644
+
645
def clone(self) -> "SamplingParams":
646
"""If skip_clone is True, uses shallow copy instead of deep copy."""
647
if self.skip_clone:
@@ -666,6 +680,17 @@ class SamplingParams(
680
value=num_logprobs,
681
)
682
683
+ # Validate logprob_token_ids.
684
+ if self.logprob_token_ids is not None:
685
+ n = len(self.logprob_token_ids)
686
+ if n > MAX_LOGPROB_TOKEN_IDS:
687
+ raise VLLMValidationError(
688
+ f"Requested logprob_token_ids of length {n}, "
689
+ f"which is greater than max allowed: {MAX_LOGPROB_TOKEN_IDS}",
690
+ parameter="logprob_token_ids",
691
+ value=n,
692
+ )
693
+
694
# Validate prompt logprobs.
695
if num_prompt_logprobs := self.prompt_logprobs:
696
if num_prompt_logprobs == -1:
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -1435,7 +1435,7 @@ class Scheduler(SchedulerInterface):
1435
# Extract sample logprobs if needed.
1436
if (
1437
request.sampling_params is not None
1438
- and request.sampling_params.logprobs is not None
1438
+ and request.sampling_params.num_logprobs is not None
1439
and logprobs
1440
):
1441
new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
vllm/v1/engine/logprobs.py
+1
-1
@@ -47,7 +47,7 @@ class LogprobsProcessor:
47
) -> "LogprobsProcessor":
48
sampling_params = request.sampling_params
49
assert sampling_params is not None
50
- num_logprobs = sampling_params.logprobs
50
+ num_logprobs = sampling_params.num_logprobs
51
num_prompt_logprobs = sampling_params.prompt_logprobs
52
return cls(
53
tokenizer=tokenizer,
vllm/v1/worker/gpu/sample/logprob.py
+133
-9
@@ -1,10 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import numpy as np
5
import torch
6
7
+from vllm.sampling_params import MAX_LOGPROB_TOKEN_IDS, SamplingParams
8
from vllm.triton_utils import tl, triton
9
from vllm.v1.outputs import LogprobsTensors
10
+from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
11
12
13
@triton.jit
@@ -75,6 +78,9 @@ def _ranks_kernel(
78
def compute_token_logprobs(
79
logits: torch.Tensor, token_ids: torch.Tensor
80
) -> torch.Tensor:
81
+ # NOTE(woosuk): To save GPU memory, we do not materialize the full
82
+ # [batch_size, vocab_size] logprobs tensor. The kernel computes
83
+ # max + logsumexp per row and only emits logprobs at `token_ids`.
84
batch_size, vocab_size = logits.shape
85
token_ids = token_ids.to(torch.int64)
86
num_logprobs = token_ids.shape[1]
@@ -97,18 +103,52 @@ def compute_topk_logprobs(
103
num_logprobs: int,
104
sampled_token_ids: torch.Tensor,
105
cu_num_logits: list[int] | None = None,
106
+ logprob_token_ids_state: "LogprobTokenIdsState | None" = None,
107
+ expanded_idx_mapping: torch.Tensor | None = None,
108
+ max_per_req_token_ids: int = 0,
109
) -> LogprobsTensors:
110
assert num_logprobs >= 0
111
batch_size, vocab_size = logits.shape
103
- logprob_token_ids = sampled_token_ids.unsqueeze(-1)
104
- if num_logprobs > 0:
105
- topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
106
- logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1)
107
-
108
- # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full
109
- # logprobs tensor. Instead, we only compute and return the logprobs of
110
- # the topk + 1 tokens.
111
- logprobs = compute_token_logprobs(logits, logprob_token_ids)
112
+
113
+ if max_per_req_token_ids == 0:
114
+ # Fast path: no request asked for custom logprob_token_ids.
115
+ logprob_token_ids = sampled_token_ids.unsqueeze(-1)
116
+ if num_logprobs > 0:
117
+ topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
118
+ logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1)
119
+ logprobs = compute_token_logprobs(logits, logprob_token_ids)
120
+ else:
121
+ # Some requests specified logprob_token_ids. Build the [batch_size,
122
+ # 1 + max_cols] token_ids matrix and validity mask on the GPU via a
123
+ # single triton kernel, overriding the topk columns with per-request
124
+ # tokens where applicable.
125
+ assert logprob_token_ids_state is not None
126
+ assert expanded_idx_mapping is not None
127
+ topk_indices = None
128
+ if num_logprobs > 0:
129
+ topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
130
+
131
+ num_cols = max(num_logprobs, max_per_req_token_ids)
132
+ logprob_token_ids = sampled_token_ids.new_zeros((batch_size, 1 + num_cols))
133
+ valid_mask = torch.zeros_like(logprob_token_ids, dtype=torch.bool)
134
+ _fill_logprob_token_ids_kernel[(batch_size,)](
135
+ logprob_token_ids,
136
+ logprob_token_ids.stride(0),
137
+ valid_mask,
138
+ valid_mask.stride(0),
139
+ sampled_token_ids,
140
+ topk_indices if topk_indices is not None else logprob_token_ids,
141
+ topk_indices.stride(0) if topk_indices is not None else 0,
142
+ expanded_idx_mapping,
143
+ logprob_token_ids_state.num_token_ids.gpu,
144
+ logprob_token_ids_state.token_ids.gpu,
145
+ logprob_token_ids_state.token_ids.gpu.stride(0),
146
+ NUM_TOPK=num_logprobs,
147
+ PADDED_COLS=triton.next_power_of_2(num_cols),
148
+ )
149
+ logprobs = compute_token_logprobs(logits, logprob_token_ids)
150
+ logprobs = logprobs.masked_fill(~valid_mask, float("-inf"))
151
+
152
token_ranks = torch.empty(batch_size, dtype=torch.int64, device=logits.device)
153
_ranks_kernel[(batch_size,)](
154
token_ranks,
@@ -124,3 +164,87 @@ def compute_topk_logprobs(
164
selected_token_ranks=token_ranks,
165
cu_num_generated_tokens=cu_num_logits,
166
)
167
+
168
+
169
+@triton.jit
170
+def _fill_logprob_token_ids_kernel(
171
+ # [batch_size, 1 + num_cols]
172
+ out_token_ids_ptr,
173
+ out_token_ids_stride,
174
+ # [batch_size, 1 + num_cols]
175
+ out_valid_mask_ptr,
176
+ out_valid_mask_stride,
177
+ sampled_token_ids_ptr, # [batch_size]
178
+ topk_indices_ptr, # [batch_size, NUM_TOPK] (unused when NUM_TOPK == 0)
179
+ topk_indices_stride,
180
+ expanded_idx_mapping_ptr, # [batch_size] -> req_state_idx
181
+ num_per_req_token_ids_ptr, # [max_num_reqs]
182
+ per_req_token_ids_ptr, # [max_num_reqs, MAX_LOGPROB_TOKEN_IDS]
183
+ per_req_token_ids_stride,
184
+ NUM_TOPK: tl.constexpr,
185
+ PADDED_COLS: tl.constexpr,
186
+):
187
+ batch_idx = tl.program_id(0)
188
+
189
+ # Column 0: always the sampled token, always valid.
190
+ sampled = tl.load(sampled_token_ids_ptr + batch_idx)
191
+ tl.store(out_token_ids_ptr + batch_idx * out_token_ids_stride, sampled)
192
+ tl.store(out_valid_mask_ptr + batch_idx * out_valid_mask_stride, 1)
193
+
194
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + batch_idx)
195
+ num_custom = tl.load(num_per_req_token_ids_ptr + req_state_idx)
196
+
197
+ col = tl.arange(0, PADDED_COLS)
198
+ tid_base = out_token_ids_ptr + batch_idx * out_token_ids_stride + 1
199
+ mask_base = out_valid_mask_ptr + batch_idx * out_valid_mask_stride + 1
200
+
201
+ if num_custom > 0:
202
+ # Override topk with per-request custom tokens.
203
+ src = per_req_token_ids_ptr + req_state_idx * per_req_token_ids_stride
204
+ valid = col < num_custom
205
+ # per_req_token_ids is int32; output is int64.
206
+ tokens = tl.load(src + col, mask=valid, other=0).to(tl.int64)
207
+ else:
208
+ # Fill with topk indices (no-op when NUM_TOPK == 0).
209
+ src = topk_indices_ptr + batch_idx * topk_indices_stride
210
+ valid = col < NUM_TOPK
211
+ tokens = tl.load(src + col, mask=valid, other=0)
212
+
213
+ tl.store(tid_base + col, tokens, mask=valid)
214
+ tl.store(mask_base + col, tl.full([PADDED_COLS], 1, tl.int1), mask=valid)
215
+
216
+
217
+class LogprobTokenIdsState:
218
+ """Per-request override of which token ids' logprobs to return.
219
+
220
+ See `SamplingParams.logprob_token_ids`.
221
+ """
222
+
223
+ def __init__(self, max_num_reqs: int, device: torch.device):
224
+ self.max_num_reqs = max_num_reqs
225
+ self.num_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
226
+ self.token_ids = StagedWriteTensor(
227
+ (max_num_reqs, MAX_LOGPROB_TOKEN_IDS),
228
+ dtype=torch.int32,
229
+ device=device,
230
+ )
231
+
232
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
233
+ token_ids = sampling_params.logprob_token_ids
234
+ if not token_ids:
235
+ self.num_token_ids.np[req_idx] = 0
236
+ return
237
+ n = len(token_ids)
238
+ if n > MAX_LOGPROB_TOKEN_IDS:
239
+ raise ValueError(
240
+ f"Too many logprob_token_ids: {n}. The max is {MAX_LOGPROB_TOKEN_IDS}."
241
+ )
242
+ self.num_token_ids.np[req_idx] = n
243
+ self.token_ids.stage_write(req_idx, 0, token_ids)
244
+
245
+ def apply_staged_writes(self) -> None:
246
+ self.num_token_ids.copy_to_uva()
247
+ self.token_ids.apply_write()
248
+
249
+ def max_num_token_ids(self, idx_mapping_np: np.ndarray) -> int:
250
+ return int(self.num_token_ids.np[idx_mapping_np].max(initial=0))
vllm/v1/worker/gpu/sample/sampler.py
+19
-3
@@ -12,7 +12,10 @@ from vllm.v1.worker.gpu.metrics.logits import get_num_nans
12
from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
13
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
14
from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
15
-from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
15
+from vllm.v1.worker.gpu.sample.logprob import (
16
+ LogprobTokenIdsState,
17
+ compute_topk_logprobs,
18
+)
19
from vllm.v1.worker.gpu.sample.output import SamplerOutput
20
from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
21
from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
@@ -38,6 +41,7 @@ class Sampler:
41
self.penalties_state = PenaltiesState(req_states)
42
self.logit_bias_state = LogitBiasState(max_num_reqs, device)
43
self.bad_words_state = BadWordsState(req_states)
44
+ self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
45
self.num_speculative_tokens = num_speculative_tokens
46
47
def add_request(
@@ -47,12 +51,14 @@ class Sampler:
51
self.penalties_state.add_request(req_idx, sampling_params)
52
self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)
53
self.bad_words_state.add_request(req_idx, sampling_params)
54
+ self.logprob_token_ids_state.add_request(req_idx, sampling_params)
55
56
def apply_staged_writes(self) -> None:
57
self.sampling_states.apply_staged_writes()
58
self.penalties_state.apply_staged_writes()
59
self.logit_bias_state.apply_staged_writes()
60
self.bad_words_state.apply_staged_writes()
61
+ self.logprob_token_ids_state.apply_staged_writes()
62
63
def __call__(
64
self,
@@ -79,13 +85,23 @@ class Sampler:
85
)
86
87
max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
82
- if max_num_logprobs != NO_LOGPROBS:
88
+ max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(
89
+ idx_mapping_np
90
+ )
91
+ if max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0:
92
if self.logprobs_mode == "processed_logprobs":
93
logits = processed_logits
94
expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
95
cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
96
+ num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
97
logprobs_tensors = compute_topk_logprobs(
88
- logits, max_num_logprobs, sampled, cu_num_logits
98
+ logits,
99
+ num_logprobs,
100
+ sampled,
101
+ cu_num_logits,
102
+ logprob_token_ids_state=self.logprob_token_ids_state,
103
+ expanded_idx_mapping=input_batch.expanded_idx_mapping,
104
+ max_per_req_token_ids=max_per_req_token_ids,
105
)
106
else:
107
logprobs_tensors = None