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 9, 07:30 UTC
[ModelRunner V2] Fix spec decoding + logprobs (#33391)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 30, 2026 at 19:33 UTC
876a16f4fb092a7954ce2c8db177943c55adb6a1
7 files changed
+22
-5
tests/v1/engine/test_output_processor.py
+1
@@ -335,6 +335,7 @@ def _validate_logprobs(
335
ref_prompt_logprob_toks,
336
ref_prompt_logprob_vals,
337
ref_prompt_token_ranks,
338
+ _,
339
) = ref_prompt_logprobs
340
for idx, (prompt_token, pos_logprob_dict) in enumerate(
341
zip(prompt_token_ids[1:], prompt_logprobs[1:])
vllm/v1/engine/logprobs.py
+1
-1
@@ -130,7 +130,7 @@ class LogprobsProcessor:
130
assert self.num_prompt_logprobs is not None
131
assert self.prompt_logprobs is not None
132
133
- token_ids, logprobs, ranks = prompt_logprobs_tensors
133
+ token_ids, logprobs, ranks, _ = prompt_logprobs_tensors
134
135
# Recover shapes.
136
num_prompt_tokens, num_logprobs = logprobs.shape
vllm/v1/outputs.py
+9
-1
@@ -51,13 +51,17 @@ class LogprobsTensors(NamedTuple):
51
logprobs: torch.Tensor
52
# [num_reqs x num_generated_tokens]
53
selected_token_ranks: torch.Tensor
54
+ # [num_reqs]
55
+ cu_num_generated_tokens: list[int] | None = None
56
57
def tolists(self, cu_num_generated_tokens: list[int] | None = None):
58
return LogprobsLists(
59
self.logprob_token_ids.cpu().numpy(),
60
self.logprobs.cpu().numpy(),
61
self.selected_token_ranks.cpu().numpy(),
60
- cu_num_generated_tokens,
62
+ cu_num_generated_tokens
63
+ if cu_num_generated_tokens is not None
64
+ else self.cu_num_generated_tokens,
65
)
66
67
def to_cpu_nonblocking(self) -> "LogprobsTensors":
@@ -67,10 +71,14 @@ class LogprobsTensors(NamedTuple):
71
self.logprob_token_ids.to("cpu", non_blocking=True),
72
self.logprobs.to("cpu", non_blocking=True),
73
self.selected_token_ranks.to("cpu", non_blocking=True),
74
+ self.cu_num_generated_tokens,
75
)
76
77
def filter(self, mask: torch.Tensor) -> "LogprobsTensors":
78
"""Filter the logprobs tensors with the given bool mask."""
79
+ assert self.cu_num_generated_tokens is None, (
80
+ "filter can't be used with cu_num_generated_tokens"
81
+ )
82
return LogprobsTensors(
83
self.logprob_token_ids[mask],
84
self.logprobs[mask],
vllm/v1/worker/gpu/model_runner.py
+2
-1
@@ -316,7 +316,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
316
# NOTE(woosuk): During the initial memory profiling, the sampler may skip
317
# top_k, top_p, and logprobs, using less GPU memory than what is possible
318
# during actual execution.
319
- self.sampler(logits, idx_mapping, idx_mapping_np, pos)
319
+ self.sampler(logits, idx_mapping, idx_mapping_np, idx_mapping_np, pos)
320
321
@torch.inference_mode()
322
def profile_run(self) -> None:
@@ -686,6 +686,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
686
logits,
687
input_batch.expanded_idx_mapping,
688
input_batch.idx_mapping_np,
689
+ input_batch.cu_num_logits_np,
690
sample_pos,
691
)
692
vllm/v1/worker/gpu/sample/logprob.py
+2
@@ -103,6 +103,7 @@ def compute_topk_logprobs(
103
logits: torch.Tensor,
104
num_logprobs: int,
105
sampled_token_ids: torch.Tensor,
106
+ cu_num_logits: list[int] | None = None,
107
) -> LogprobsTensors:
108
assert num_logprobs >= 0
109
batch_size, vocab_size = logits.shape
@@ -135,4 +136,5 @@ def compute_topk_logprobs(
136
logprob_token_ids=logprob_token_ids,
137
logprobs=logprobs,
138
selected_token_ranks=token_ranks,
139
+ cu_num_generated_tokens=cu_num_logits,
140
)
vllm/v1/worker/gpu/sample/sampler.py
+6
-1
@@ -62,6 +62,7 @@ class Sampler:
62
logits: torch.Tensor,
63
idx_mapping: torch.Tensor,
64
idx_mapping_np: np.ndarray,
65
+ cu_num_logits_np: np.ndarray,
66
pos: torch.Tensor,
67
) -> SamplerOutput:
68
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
@@ -78,7 +79,11 @@ class Sampler:
79
if self.logprobs_mode == "processed_logprobs"
80
else logits
81
)
81
- logprobs_tensors = compute_topk_logprobs(logits, max_num_logprobs, sampled)
82
+ expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
83
+ cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
84
+ logprobs_tensors = compute_topk_logprobs(
85
+ logits, max_num_logprobs, sampled, cu_num_logits
86
+ )
87
else:
88
logprobs_tensors = None
89
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -4449,7 +4449,7 @@ class GPUModelRunner(
4449
4450
# Compute prompt logprobs.
4451
logprobs = self.sampler.compute_logprobs(logits)
4452
- token_ids, logprobs, ranks = self.sampler.gather_logprobs(
4452
+ token_ids, logprobs, ranks, _ = self.sampler.gather_logprobs(
4453
logprobs, num_prompt_logprobs, tgt_token_ids
4454
)
4455