@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 41ea2dd44a

[Bugfix][V1/V2] Fix prompt_logprobs to respect logprobs_mode (#47680)

Signed-off-by: Wojciech Wais <wojciech.wais@gmail.com> Signed-off-by: Federico Kamelhar <209537060+fede-kamel@users.noreply.github.com> Signed-off-by: Allen Shen <aoshen@inferact.ai> Co-authored-by: Wojciech Wais <wojciech.wais@gmail.com> Co-authored-by: Federico Kamelhar <209537060+fede-kamel@users.noreply.github.com>

aoshen02 committed Jul 18, 2026 at 04:58 UTC 41ea2dd44a3a20c46ebeb985de0022c7673fb953

11 files changed +106 -40

tests/v1/sample/test_logprobs.py

+38

@@ -564,6 +564,44 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):

564 cleanup_dist_env_and_memory()

565

566

567 +def test_prompt_logprobs_mode():

568 + """prompt_logprobs must respect logprobs_mode: *_logits and *_logprobs

569 + must return different values. Prompt tokens skip sampling processors,

570 + so processed_* == raw_* on the prompt side."""

571 + from vllm import LLM

572 +

573 + values: dict[str, float] = {}

574 + for mode in get_args(LogprobsMode):

575 + llm = LLM(

576 + "facebook/opt-125m",

577 + enable_prefix_caching=False,

578 + gpu_memory_utilization=0.05,

579 + max_model_len=16,

580 + logprobs_mode=mode,

581 + )

582 + try:

583 + results = llm.generate(

584 + ["Hello world"],

585 + sampling_params=SamplingParams(

586 + max_tokens=1, prompt_logprobs=0, temperature=0

587 + ),

588 + )

589 + assert results[0].prompt_logprobs is not None

590 + assert results[0].prompt_logprobs[1] is not None

591 + tok_id = results[0].prompt_token_ids[1]

592 + values[mode] = results[0].prompt_logprobs[1][tok_id].logprob

593 + finally:

594 + del llm

595 + torch.accelerator.empty_cache()

596 + cleanup_dist_env_and_memory()

597 +

598 + assert values["raw_logprobs"] <= 0

599 + assert values["processed_logprobs"] <= 0

600 + assert values["raw_logits"] != values["raw_logprobs"]

601 + assert values["processed_logits"] == values["raw_logits"]

602 + assert values["processed_logprobs"] == values["raw_logprobs"]

603 +

604 +

605 class TestCorrectDecodedToken:

606 """Unit tests for _correct_decoded_token method in LogprobsProcessor.

607

vllm/config/model.py

@@ -233,6 +233,8 @@ class ModelConfig:

233 Raw means the values before applying any logit processors, like bad words.

234 Processed means the values after applying all processors, including

235 temperature and top_k/top_p.

236 + Note: for prompt_logprobs, processed_* and raw_* yield identical results

237 + because prompt tokens do not go through sampling processors.

238 """

239 use_fp64_gumbel: bool = False

240 """Whether to use FP64 (instead of FP32) random noise for Gumbel-max and

vllm/config/vllm.py

-7

@@ -2202,13 +2202,6 @@ class VllmConfig:

2202 if model_config is not None and model_config.enable_prompt_embeds:

2203 unsupported.append("prompt embeds")

2204

2205 - if (

2206 - model_config is not None

2207 - and model_config.runner_type == "generate"

2208 - and model_config.logprobs_mode in ("raw_logits", "processed_logits")

2209 - ):

2210 - unsupported.append(f"logprobs mode '{model_config.logprobs_mode}'")

2211 -

2205 if self.cache_config.kv_sharing_fast_prefill:

2206 # Will be added by https://github.com/vllm-project/vllm/pull/35045

2207 unsupported.append("KV sharing fast prefill")

vllm/model_executor/models/diffusion_gemma.py

+4 -2

@@ -53,7 +53,7 @@ from vllm.v1.worker.gpu.attn_utils import build_attn_metadata

53 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu

54 from vllm.v1.worker.gpu.input_batch import InputBatch

55 from vllm.v1.worker.gpu.model_states.interface import ModelState

56 -from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs

56 +from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores

57 from vllm.v1.worker.gpu.sample.output import SamplerOutput

58 from vllm.v1.worker.gpu.sample.penalties import use_penalty

59 from vllm.v1.worker.gpu.states import RequestState

@@ -1064,6 +1064,7 @@ class DiffusionSampler:

1064 ):

1065 self.sampling_states = sampler.sampling_states

1066 self.req_states = sampler.req_states

1067 + self.logits_mode = sampler.logprobs_mode in ("raw_logits", "processed_logits")

1068 # Self-conditioning soft embed = probs @ embed_weight * normalizer,

1069 # computed in the sampler (see _compiled_sample_step). ``embed_weight``

1070 # is the vocab-parallel shard; [sc_vocab_start, sc_vocab_end) is this

@@ -1359,10 +1360,11 @@ class DiffusionSampler:

1360 # positions are never emitted.

1361 k_i = int(valid_canvas_len_np[start_req + li])

1362 pos = li * CL

1362 - self._pending_logprobs[slot.item()] = compute_topk_logprobs(

1363 + self._pending_logprobs[slot.item()] = compute_topk_scores(

1364 flat_logits[pos : pos + k_i],

1365 max_num_logprobs,

1366 argmax_tokens[local_idx][:k_i],

1367 + logits_mode=self.logits_mode,

1368 )

1369

1370 # Commit steps: is_committing was True at entry. Reassemble previously

vllm/v1/sample/rejection_sampler.py

+8 -2

@@ -67,8 +67,14 @@ class RejectionSampler(nn.Module):

67 self.sampler = sampler

68 self.use_fp64_gumbel = getattr(sampler, "use_fp64_gumbel", False)

69 logprobs_mode = self.sampler.logprobs_mode

70 - self.is_processed_logprobs_mode = logprobs_mode.startswith("processed")

71 - self.is_logits_logprobs_mode = logprobs_mode.endswith("logits")

70 + self.is_processed_logprobs_mode = logprobs_mode in (

71 + "processed_logprobs",

72 + "processed_logits",

73 + )

74 + self.is_logits_logprobs_mode = logprobs_mode in (

75 + "raw_logits",

76 + "processed_logits",

77 + )

79 self.synthetic_conditional_rates: torch.Tensor | None = None

80 if (

vllm/v1/worker/gpu/model_runner.py

+4 -1

@@ -341,7 +341,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

341 self.speculative_config,

342 self.device,

343 )

344 - self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)

344 + self.prompt_logprobs_worker = PromptLogprobsWorker(

345 + self.max_num_reqs,

346 + logprobs_mode=self.model_config.logprobs_mode,

347 + )

348 self.structured_outputs_worker = StructuredOutputsWorker(

349 max_num_logits=self.max_num_reqs * self.decode_query_len,

350 vocab_size=self.vocab_size,

vllm/v1/worker/gpu/sample/logprob.py

+12 -5

@@ -106,7 +106,7 @@ def compute_token_logprobs(

106 return logprobs

107

108

109 -def compute_topk_logprobs(

109 +def compute_topk_scores(

110 logits: torch.Tensor,

111 num_logprobs: int,

112 sampled_token_ids: torch.Tensor,

@@ -114,6 +114,7 @@ def compute_topk_logprobs(

114 logprob_token_ids_state: "LogprobTokenIdsState | None" = None,

115 expanded_idx_mapping: torch.Tensor | None = None,

116 max_per_req_token_ids: int = 0,

117 + logits_mode: bool = False,

118 ) -> LogprobsTensors:

119 assert num_logprobs >= 0

120 batch_size, vocab_size = logits.shape

@@ -124,7 +125,10 @@ def compute_topk_logprobs(

125 if num_logprobs > 0:

126 topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices

127 logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1)

127 - logprobs = compute_token_logprobs(logits, logprob_token_ids)

128 + if logits_mode:

129 + scores = logits.gather(-1, logprob_token_ids).to(torch.float32)

130 + else:

131 + scores = compute_token_logprobs(logits, logprob_token_ids)

132 else:

133 # Some requests specified logprob_token_ids. Build the [batch_size,

134 # 1 + max_cols] token_ids matrix and validity mask on the GPU via a

@@ -158,8 +162,11 @@ def compute_topk_logprobs(

162 NUM_TOPK=num_logprobs,

163 PADDED_COLS=triton.next_power_of_2(num_cols),

164 )

161 - logprobs = compute_token_logprobs(logits, logprob_token_ids)

162 - logprobs = logprobs.masked_fill(~valid_mask, float("-inf"))

165 + if logits_mode:

166 + scores = logits.gather(-1, logprob_token_ids).to(torch.float32)

167 + else:

168 + scores = compute_token_logprobs(logits, logprob_token_ids)

169 + scores = scores.masked_fill(~valid_mask, float("-inf"))

170

171 token_ranks = torch.empty(batch_size, dtype=torch.int64, device=logits.device)

172 _ranks_kernel[(batch_size,)](

@@ -172,7 +179,7 @@ def compute_topk_logprobs(

179 )

180 return LogprobsTensors(

181 logprob_token_ids=logprob_token_ids,

175 - logprobs=logprobs,

182 + logprobs=scores,

183 selected_token_ranks=token_ranks,

184 cu_num_generated_tokens=cu_num_logits,

185 )

vllm/v1/worker/gpu/sample/prompt_logprob.py

+17 -11

@@ -5,16 +5,18 @@ from collections.abc import Callable

5 import numpy as np

6 import torch

8 +from vllm.config.model import LogprobsMode

9 from vllm.sampling_params import SamplingParams

10 from vllm.triton_utils import tl, triton

11 from vllm.v1.outputs import LogprobsTensors

12 from vllm.v1.worker.gpu.input_batch import InputBatch

12 -from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs

13 +from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores

16 class PromptLogprobsWorker:

16 - def __init__(self, max_num_reqs: int):

17 + def __init__(self, max_num_reqs: int, logprobs_mode: LogprobsMode = "raw_logprobs"):

18 self.max_num_reqs = max_num_reqs

19 + self.logprobs_mode = logprobs_mode

21 self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)

22 self.num_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=np.int32)

@@ -82,6 +84,7 @@ class PromptLogprobsWorker:

84 hidden_states[: input_batch.num_tokens],

85 logits_fn,

86 max_num_prompt_logprobs,

87 + self.logprobs_mode,

88 )

89 )

@@ -206,33 +209,36 @@ def compute_prompt_logprobs_with_chunking(

209 prompt_hidden_states: torch.Tensor,

210 logits_fn: Callable[[torch.Tensor], torch.Tensor],

211 num_prompt_logprobs: int,

212 + logprobs_mode: LogprobsMode = "raw_logprobs",

213 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

214 # Since materializing the full prompt logits can take too much memory,

215 # we compute it in chunks.

216 CHUNK_SIZE = 1024

217 token_ids = []

214 - logprobs = []

218 + scores = []

219 ranks = []

220 + logits_mode = logprobs_mode in ("raw_logits", "processed_logits")

221 prompt_token_ids = prompt_token_ids.to(torch.int64)

222 for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):

223 end_idx = start_idx + CHUNK_SIZE

224 # NOTE(woosuk): logits_fn can be slow because it involves all-gather.

225 prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])

221 - requested_num_prompt_logprobs = (

226 + requested_num = (

227 prompt_logits.shape[-1]

228 if num_prompt_logprobs == -1

229 else num_prompt_logprobs

230 )

226 - prompt_logprobs = compute_topk_logprobs(

231 + result = compute_topk_scores(

232 prompt_logits,

228 - requested_num_prompt_logprobs,

233 + requested_num,

234 prompt_token_ids[start_idx:end_idx],

235 + logits_mode=logits_mode,

236 )

231 - token_ids.append(prompt_logprobs.logprob_token_ids)

232 - logprobs.append(prompt_logprobs.logprobs)

233 - ranks.append(prompt_logprobs.selected_token_ranks)

237 + token_ids.append(result.logprob_token_ids)

238 + scores.append(result.logprobs)

239 + ranks.append(result.selected_token_ranks)

240

241 token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0]

236 - logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]

242 + scores = torch.cat(scores, dim=0) if len(scores) > 1 else scores[0]

243 ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]

238 - return token_ids, logprobs, ranks

244 + return token_ids, scores, ranks

vllm/v1/worker/gpu/sample/sampler.py

+8 -6

@@ -19,7 +19,7 @@ from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

19 from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState

20 from vllm.v1.worker.gpu.sample.logprob import (

21 LogprobTokenIdsState,

22 - compute_topk_logprobs,

22 + compute_topk_scores,

23 )

24 from vllm.v1.worker.gpu.sample.output import SamplerOutput

25 from vllm.v1.worker.gpu.sample.penalties import PenaltiesState

@@ -38,8 +38,6 @@ class Sampler:

38 num_speculative_tokens: int = 1,

39 use_fp64_gumbel: bool = False,

40 ):

41 - if logprobs_mode not in ("processed_logprobs", "raw_logprobs"):

42 - raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")

41 self.logprobs_mode = logprobs_mode

42 self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.

43 self.use_fp64_gumbel = use_fp64_gumbel

@@ -102,12 +100,12 @@ class Sampler:

100 )

101

102 if return_logprobs:

105 - if self.logprobs_mode == "processed_logprobs":

103 + if self.logprobs_mode in ("processed_logprobs", "processed_logits"):

104 logits = processed_logits

105 expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]

106 cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None

107 num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0

110 - logprobs_tensors = compute_topk_logprobs(

108 + logprobs_tensors = compute_topk_scores(

109 logits,

110 num_logprobs,

111 sampled,

@@ -115,6 +113,7 @@ class Sampler:

113 logprob_token_ids_state=self.logprob_token_ids_state,

114 expanded_idx_mapping=input_batch.expanded_idx_mapping,

115 max_per_req_token_ids=max_per_req_token_ids,

116 + logits_mode=self.logprobs_mode in ("raw_logits", "processed_logits"),

117 )

118 else:

119 logprobs_tensors = None

@@ -222,7 +221,10 @@ class Sampler:

221 # any greedy requests or per-request seeds, or if post-processed

222 # logprobs need to be returned for any requests.

223 (top_k is None and top_p is None)

225 - or (return_logprobs and self.logprobs_mode == "processed_logprobs")

224 + or (

225 + return_logprobs

226 + and self.logprobs_mode in ("processed_logprobs", "processed_logits")

227 + )

228 or self.sampling_states.any_greedy(idx_mapping_np)

229 or self.sampling_states.any_explicit_seed(idx_mapping_np)

230 )

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

+5 -3

@@ -11,7 +11,7 @@ from vllm.v1.worker.gpu.input_batch import (

11 get_num_sampled_and_rejected,

12 )

13 from vllm.v1.worker.gpu.metrics.logits import get_num_nans

14 -from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs

14 +from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores

15 from vllm.v1.worker.gpu.sample.output import SamplerOutput

16 from vllm.v1.worker.gpu.sample.sampler import Sampler

17 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS

@@ -91,11 +91,13 @@ class RejectionSampler:

91 num_warps=1,

92 )

93 expanded_logits = num_logits != input_batch.idx_mapping.shape[0]

94 - return compute_topk_logprobs(

94 + return compute_topk_scores(

95 logits,

96 max_num_logprobs,

97 flat_sampled,

98 input_batch.cu_num_logits_np.tolist() if expanded_logits else None,

99 + logits_mode=self.sampler.logprobs_mode

100 + in ("raw_logits", "processed_logits"),

101 )

102

103 def __call__(

@@ -139,7 +141,7 @@ class RejectionSampler:

141 sampled,

142 num_sampled,

143 processed_logits

142 - if self.sampler.logprobs_mode == "processed_logprobs"

144 + if self.sampler.logprobs_mode in ("processed_logprobs", "processed_logits")

145 else logits,

146 )

147

vllm/v1/worker/gpu_model_runner.py

+8 -3

@@ -5635,10 +5635,15 @@ class GPUModelRunner(

5635 # to gather the logprob for.

5636 tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits]

5637

5638 - # Compute prompt logprobs.

5639 - logprobs = self.sampler.compute_logprobs(logits)

5638 + # Compute prompt scores respecting logprobs_mode.

5639 + # NOTE: prompt tokens skip sampling processors, so

5640 + # processed_* and raw_* yield the same scores here.

5641 + if self.model_config.logprobs_mode in ("raw_logits", "processed_logits"):

5642 + scores = logits.to(torch.float32)

5643 + else:

5644 + scores = self.sampler.compute_logprobs(logits)

5645 token_ids, logprobs, ranks, _ = self.sampler.gather_logprobs(

5641 - logprobs, num_prompt_logprobs, tgt_token_ids

5646 + scores, num_prompt_logprobs, tgt_token_ids

5647 )

5648

5649 # Transfer GPU->CPU async.