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 10, 10:15 UTC
[perf][async] support non cpu sync get logprob tensors for spec (#31336)
Signed-off-by: izhuhaoran <izhuhaoran@qq.com> Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
zhrrr committed
Jan 10, 2026 at 05:24 UTC
97ba96fbe9a53416b973bbe05eb3bfefd6408a6c
3 files changed
+48
-29
vllm/v1/outputs.py
+8
@@ -69,6 +69,14 @@ class LogprobsTensors(NamedTuple):
69
self.selected_token_ranks.to("cpu", non_blocking=True),
70
)
71
72
+ def filter(self, mask: torch.Tensor) -> "LogprobsTensors":
73
+ """Filter the logprobs tensors with the given bool mask."""
74
+ return LogprobsTensors(
75
+ self.logprob_token_ids[mask],
76
+ self.logprobs[mask],
77
+ self.selected_token_ranks[mask],
78
+ )
79
+
80
@staticmethod
81
def empty_cpu(
82
num_positions: int, num_tokens_per_position: int
vllm/v1/sample/rejection_sampler.py
+25
-14
@@ -9,7 +9,7 @@ import torch.nn as nn
9
10
from vllm.logger import init_logger
11
from vllm.triton_utils import tl, triton
12
-from vllm.v1.outputs import LogprobsTensors, SamplerOutput
12
+from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
13
from vllm.v1.sample.metadata import SamplingMetadata
14
from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts
15
from vllm.v1.sample.ops.penalties import apply_all_penalties
@@ -185,13 +185,22 @@ class RejectionSampler(nn.Module):
185
final_logits[target_logits_indices] = target_logits.to(torch.float32)
186
final_logits[bonus_logits_indices] = bonus_logits.to(torch.float32)
187
188
- # Compute accepted token indices.
189
- accepted_mask = sampled_token_ids != PLACEHOLDER_TOKEN_ID
190
- num_accepted_tokens = accepted_mask.sum(dim=-1)
191
- accepted_logit_indices = accepted_mask.nonzero(as_tuple=True)[1]
192
- accepted_logit_indices += cu_num_sampled_tokens.repeat_interleave(
193
- num_accepted_tokens
188
+ # NOTE: To avoid cpu-gpu synchronization, we now simply compute indices for
189
+ # all draft tokens, including the rejected ones. The rejected tokens will
190
+ # be filtered out in the `parse_output`.
191
+ logit_start_indices = cu_num_sampled_tokens
192
+ offsets = torch.arange(
193
+ sampled_token_ids.shape[-1],
194
+ device=logit_start_indices.device,
195
+ dtype=logit_start_indices.dtype,
196
)
197
+ accepted_logit_indices = (
198
+ logit_start_indices.unsqueeze(1) + offsets.unsqueeze(0)
199
+ ).flatten()
200
+ accepted_logit_indices.clamp_(max=final_logits.shape[0] - 1)
201
+ accepted_tokens = sampled_token_ids.clone().flatten()
202
+ # we replace rejected token ids with 0 to avoid gather_logprobs error
203
+ accepted_tokens[accepted_tokens == PLACEHOLDER_TOKEN_ID] = 0
204
205
# Compute logprobs for accepted tokens.
206
accepted_logits = final_logits[accepted_logit_indices]
@@ -200,7 +209,6 @@ class RejectionSampler(nn.Module):
209
if self.is_logits_logprobs_mode
210
else self.sampler.compute_logprobs(accepted_logits)
211
)
203
- accepted_tokens = sampled_token_ids[accepted_mask]
212
return self.sampler.gather_logprobs(
213
accepted_logprobs,
214
max_num_logprobs,
@@ -212,8 +220,8 @@ class RejectionSampler(nn.Module):
220
output_token_ids: torch.Tensor,
221
vocab_size: int,
222
discard_req_indices: Sequence[int] = (),
215
- return_cu_num_tokens: bool = False,
216
- ) -> tuple[list[list[int]], list[int] | None]:
223
+ logprobs_tensors: LogprobsTensors | None = None,
224
+ ) -> tuple[list[list[int]], LogprobsLists | None]:
225
"""Parse the output of the rejection sampler.
226
Args:
227
output_token_ids: The sampled token IDs in shape
@@ -222,7 +230,7 @@ class RejectionSampler(nn.Module):
230
and will be filtered out in this function.
231
vocab_size: The size of the vocabulary.
232
discard_req_indices: Optional row indices to discard tokens in.
225
- return_cu_num_tokens: Whether to also return cumulative token counts.
233
+ logprobs_tensors: Optional logprobs tensors to filter.
234
Returns:
235
A list of lists of token IDs.
236
"""
@@ -231,15 +239,18 @@ class RejectionSampler(nn.Module):
239
valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & (
240
output_token_ids_np < vocab_size
241
)
234
- cu_num_tokens = None
235
- if return_cu_num_tokens:
242
+ output_logprobs = None
243
+ if logprobs_tensors is not None:
244
cu_num_tokens = [0] + valid_mask.sum(axis=1).cumsum().tolist()
245
+ filtered_tensors = logprobs_tensors.filter(valid_mask.flatten())
246
+ output_logprobs = filtered_tensors.tolists(cu_num_tokens)
247
+
248
if len(discard_req_indices) > 0:
249
valid_mask[discard_req_indices] = False
250
outputs = [
251
row[valid_mask[i]].tolist() for i, row in enumerate(output_token_ids_np)
252
]
242
- return outputs, cu_num_tokens
253
+ return outputs, output_logprobs
254
255
def apply_logits_processors(
256
self,
vllm/v1/worker/gpu_model_runner.py
+15
-15
@@ -237,19 +237,20 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
237
valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()
238
for i in self._invalid_req_indices:
239
valid_sampled_token_ids[i].clear()
240
- cu_num_tokens = None
240
+ logprobs_lists = None
241
+ if self._logprobs_tensors_cpu is not None:
242
+ logprobs_lists = self._logprobs_tensors_cpu.tolists()
243
else:
242
- valid_sampled_token_ids, cu_num_tokens = RejectionSampler.parse_output(
244
+ valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output(
245
self.sampled_token_ids_cpu,
246
self.vocab_size,
247
self._invalid_req_indices,
246
- return_cu_num_tokens=self._logprobs_tensors_cpu is not None,
248
+ logprobs_tensors=self._logprobs_tensors_cpu,
249
)
250
251
output = self._model_runner_output
252
output.sampled_token_ids = valid_sampled_token_ids
251
- if self._logprobs_tensors_cpu:
252
- output.logprobs = self._logprobs_tensors_cpu.tolists(cu_num_tokens)
253
+ output.logprobs = logprobs_lists
254
return output
255
256
@@ -395,6 +396,9 @@ class GPUModelRunner(
396
else:
397
self.max_encoder_len = 0
398
399
+ # Async scheduling
400
+ self.use_async_scheduling = self.scheduler_config.async_scheduling
401
+
402
# Sampler
403
self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
404
@@ -504,7 +508,6 @@ class GPUModelRunner(
508
cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
509
)
510
507
- self.use_async_scheduling = self.scheduler_config.async_scheduling
511
# Separate cuda stream for overlapping transfer of sampled token ids from
512
# GPU to CPU when async scheduling is enabled.
513
self.async_output_copy_stream: torch.cuda.Stream | None = None
@@ -2784,7 +2787,7 @@ class GPUModelRunner(
2787
sampled_token_ids = sampler_output.sampled_token_ids
2788
logprobs_tensors = sampler_output.logprobs_tensors
2789
invalid_req_indices = []
2787
- cu_num_tokens: list[int] | None = None
2790
+ logprobs_lists = None
2791
if not self.use_async_scheduling:
2792
# Get the valid generated tokens.
2793
max_gen_len = sampled_token_ids.shape[-1]
@@ -2794,13 +2797,16 @@ class GPUModelRunner(
2797
# Mask out the sampled tokens that should not be sampled.
2798
for i in discard_sampled_tokens_req_indices:
2799
valid_sampled_token_ids[int(i)].clear()
2800
+
2801
+ if logprobs_tensors is not None:
2802
+ logprobs_lists = logprobs_tensors.tolists()
2803
else:
2804
# Includes spec decode tokens.
2799
- valid_sampled_token_ids, cu_num_tokens = RejectionSampler.parse_output(
2805
+ valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output(
2806
sampled_token_ids,
2807
self.input_batch.vocab_size,
2808
discard_sampled_tokens_req_indices,
2803
- return_cu_num_tokens=logprobs_tensors is not None,
2809
+ logprobs_tensors=logprobs_tensors,
2810
)
2811
else:
2812
valid_sampled_token_ids = []
@@ -2853,12 +2859,6 @@ class GPUModelRunner(
2859
req_state = self.requests[req_id]
2860
req_state.output_token_ids.extend(sampled_ids)
2861
2856
- logprobs_lists = (
2857
- logprobs_tensors.tolists(cu_num_tokens)
2858
- if not self.use_async_scheduling and logprobs_tensors is not None
2859
- else None
2860
- )
2861
-
2862
# Compute prompt logprobs if needed.
2863
prompt_logprobs_dict = self._get_prompt_logprobs_dict(
2864
hidden_states[:num_scheduled_tokens],