[Bugfix][V1/V2] Fix prompt_logprobs to respect logprobs_mode

prompt_logprobs ignored the logprobs_mode config — all modes returned
log_softmax results. This fix threads the mode through both V1 and V2
GPU model runners so that *_logits modes return raw logits and
*_logprobs modes return log_softmax values.

V1 runner (gpu_model_runner.py): inline mode branch before gather.
V2 runner (prompt_logprob.py): route to compute_topk_logits or
compute_topk_logprobs based on mode.
Config (vllm.py): remove V2 runner unsupported-feature gate for
raw_logits/processed_logits now that V2 handles them.

Fixes: #35832

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

Author: 3 people

tests/v1/sample/test_prompt_logprobs_mode.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for prompt-side logprobs_mode branching.
+
+The mode selects logits (float32) vs log_softmax(logits) for prompt
+scoring. These tests verify the branch without a GPU model runner.
+"""
+
+import pytest
+import torch
+
+
+@pytest.mark.parametrize(
+    "mode",
+    ["raw_logits", "processed_logits", "raw_logprobs", "processed_logprobs"],
+)
+def test_prompt_score_mode_branch(mode: str):
+    logits = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float16)
+    compute_logprobs_called = False
+
+    def compute_logprobs(x: torch.Tensor) -> torch.Tensor:
+        nonlocal compute_logprobs_called
+        compute_logprobs_called = True
+        return x.log_softmax(dim=-1, dtype=torch.float32)
+
+    if mode.endswith("logits"):
+        scores = logits.to(torch.float32)
+    else:
+        scores = compute_logprobs(logits)
+
+    assert scores.dtype == torch.float32
+
+    if mode.endswith("logits"):
+        torch.testing.assert_close(scores, logits.to(torch.float32))
+        assert not compute_logprobs_called
+    else:
+        expected = logits.log_softmax(dim=-1, dtype=torch.float32)
+        torch.testing.assert_close(scores, expected)
+        assert compute_logprobs_called

tests/v1/sample/test_sampling_params_e2e.py

@@ -182,3 +182,25 @@ def test_seed(llm):
 
     assert out_1[0].outputs[0].text == out_2[0].outputs[0].text
     assert out_1[0].outputs[0].text != out_3[0].outputs[0].text
+
+
+def test_prompt_logprobs_respects_logprobs_mode():
+    """prompt_logprobs values must differ between raw_logits and raw_logprobs."""
+    prompt = "Hello world"
+    values: dict[str, float] = {}
+
+    for mode in ("raw_logits", "raw_logprobs"):
+        engine = LLM(MODEL, enforce_eager=True, logprobs_mode=mode)
+        out = engine.generate(
+            prompt,
+            SamplingParams(max_tokens=1, prompt_logprobs=0, temperature=0),
+        )[0]
+        assert out.prompt_logprobs is not None
+        assert out.prompt_logprobs[1] is not None
+        tok_id = out.prompt_token_ids[1]
+        values[mode] = out.prompt_logprobs[1][tok_id].logprob
+        del engine
+
+    assert values["raw_logits"] != values["raw_logprobs"], (
+        "prompt_logprobs should differ between logits and logprobs modes"
+    )

vllm/config/vllm.py

@@ -2103,13 +2103,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
         if model_config is not None and model_config.enable_prompt_embeds:
             unsupported.append("prompt embeds")
 
-        if (
-            model_config is not None
-            and model_config.runner_type == "generate"
-            and model_config.logprobs_mode in ("raw_logits", "processed_logits")
-        ):
-            unsupported.append(f"logprobs mode '{model_config.logprobs_mode}'")
-
         if self.cache_config.kv_sharing_fast_prefill:
             # Will be added by https://github.com/vllm-project/vllm/pull/35045
             unsupported.append("KV sharing fast prefill")

vllm/v1/worker/gpu/model_runner.py

@@ -341,7 +341,10 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
                     self.speculative_config,
                     self.device,
                 )
-            self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
+            self.prompt_logprobs_worker = PromptLogprobsWorker(
+                self.max_num_reqs,
+                logprobs_mode=self.model_config.logprobs_mode,
+            )
             self.structured_outputs_worker = StructuredOutputsWorker(
                 max_num_logits=self.max_num_reqs * self.decode_query_len,
                 vocab_size=self.vocab_size,

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

@@ -106,6 +106,34 @@ def compute_token_logprobs(
     return logprobs
 
 
+def compute_topk_logits(
+    logits: torch.Tensor,
+    num_logprobs: int,
+    sampled_token_ids: torch.Tensor,
+) -> LogprobsTensors:
+    """Like compute_topk_logprobs but returns raw logits (no log_softmax)."""
+    batch_size, vocab_size = logits.shape
+    logits_f32 = logits.to(torch.float32)
+
+    logprob_token_ids = sampled_token_ids.unsqueeze(-1)
+    if num_logprobs > 0:
+        topk_indices = torch.topk(logits_f32, num_logprobs, dim=-1).indices
+        logprob_token_ids = torch.cat(
+            (logprob_token_ids, topk_indices), dim=1
+        )
+
+    values = logits_f32.gather(-1, logprob_token_ids)
+
+    target_vals = logits_f32.gather(-1, sampled_token_ids.unsqueeze(-1))
+    token_ranks = (logits_f32 >= target_vals).sum(-1)
+
+    return LogprobsTensors(
+        logprob_token_ids=logprob_token_ids,
+        logprobs=values,
+        selected_token_ranks=token_ranks,
+    )
+
+
 def compute_topk_logprobs(
     logits: torch.Tensor,
     num_logprobs: int,

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

@@ -9,12 +9,16 @@
 from vllm.triton_utils import tl, triton
 from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.worker.gpu.input_batch import InputBatch
-from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
+from vllm.v1.worker.gpu.sample.logprob import (
+    compute_topk_logits,
+    compute_topk_logprobs,
+)
 
 
 class PromptLogprobsWorker:
-    def __init__(self, max_num_reqs: int):
+    def __init__(self, max_num_reqs: int, logprobs_mode: str = "raw_logprobs"):
         self.max_num_reqs = max_num_reqs
+        self.logprobs_mode = logprobs_mode
 
         self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
         self.num_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=np.int32)
@@ -82,6 +86,7 @@ def compute_prompt_logprobs(
                 hidden_states[: input_batch.num_tokens],
                 logits_fn,
                 max_num_prompt_logprobs,
+                self.logprobs_mode,
             )
         )
 
@@ -206,33 +211,36 @@ def compute_prompt_logprobs_with_chunking(
     prompt_hidden_states: torch.Tensor,
     logits_fn: Callable[[torch.Tensor], torch.Tensor],
     num_prompt_logprobs: int,
+    logprobs_mode: str = "raw_logprobs",
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     # Since materializing the full prompt logits can take too much memory,
     # we compute it in chunks.
     CHUNK_SIZE = 1024
     token_ids = []
-    logprobs = []
+    scores = []
     ranks = []
     prompt_token_ids = prompt_token_ids.to(torch.int64)
     for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
         end_idx = start_idx + CHUNK_SIZE
-        # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
         prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
-        requested_num_prompt_logprobs = (
+        requested_num = (
             prompt_logits.shape[-1]
             if num_prompt_logprobs == -1
             else num_prompt_logprobs
         )
-        prompt_logprobs = compute_topk_logprobs(
-            prompt_logits,
-            requested_num_prompt_logprobs,
-            prompt_token_ids[start_idx:end_idx],
+        compute_fn = (
+            compute_topk_logits if logprobs_mode.endswith("logits") else compute_topk_logprobs
         )
-        token_ids.append(prompt_logprobs.logprob_token_ids)
-        logprobs.append(prompt_logprobs.logprobs)
-        ranks.append(prompt_logprobs.selected_token_ranks)
+        result = compute_fn(
+            prompt_logits, 
+            requested_num, 
+            prompt_token_ids[start_idx:end_idx]
+        )
+        token_ids.append(result.logprob_token_ids)
+        scores.append(result.logprobs)
+        ranks.append(result.selected_token_ranks)
 
     token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0]
-    logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
+    scores = torch.cat(scores, dim=0) if len(scores) > 1 else scores[0]
     ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
-    return token_ids, logprobs, ranks
+    return token_ids, scores, ranks

vllm/v1/worker/gpu_model_runner.py

@@ -5556,10 +5556,12 @@ def _get_prompt_logprobs_dict(
             # to gather the logprob for.
             tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits]
 
-            # Compute prompt logprobs.
-            logprobs = self.sampler.compute_logprobs(logits)
+            if self.model_config.logprobs_mode.endswith("logits"):
+                scores = logits.to(torch.float32)
+            else:
+                scores = self.sampler.compute_logprobs(logits)
             token_ids, logprobs, ranks, _ = self.sampler.gather_logprobs(
-                logprobs, num_prompt_logprobs, tgt_token_ids
+                scores, num_prompt_logprobs, tgt_token_ids
             )
 
             # Transfer GPU->CPU async.