Unify compute_topk_logits into compute_topk_logprobs via logits_mode

Instead of a separate compute_topk_logits function that misses
cu_num_logits, logprob_token_ids_state, and expanded_idx_mapping,
add a logits_mode bool to compute_topk_logprobs. The only difference
is step 2 (score computation): gather raw float32 vs Triton
log_softmax. Steps 1 (token_ids matrix) and 3 (ranks + packaging)
are shared.

This fixes the chatgpt-codex review findings: logits mode now
correctly handles custom logprob_token_ids and spec decode offsets.

Signed-off-by: Allen Shen <aoshen@inferact.ai>

Author: aoshen02

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

@@ -106,32 +106,6 @@ 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,
@@ -140,6 +114,7 @@ def compute_topk_logprobs(
     logprob_token_ids_state: "LogprobTokenIdsState | None" = None,
     expanded_idx_mapping: torch.Tensor | None = None,
     max_per_req_token_ids: int = 0,
+    logits_mode: bool = False,
 ) -> LogprobsTensors:
     assert num_logprobs >= 0
     batch_size, vocab_size = logits.shape
@@ -150,20 +125,18 @@ def compute_topk_logprobs(
         if num_logprobs > 0:
             topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
             logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1)
-        logprobs = compute_token_logprobs(logits, logprob_token_ids)
+        if logits_mode:
+            scores = logits.to(torch.float32).gather(-1, logprob_token_ids)
+        else:
+            scores = compute_token_logprobs(logits, logprob_token_ids)
     else:
-        # Some requests specified logprob_token_ids. Build the [batch_size,
-        # 1 + max_cols] token_ids matrix and validity mask on the GPU via a
-        # single triton kernel, overriding the topk columns with per-request
-        # tokens where applicable.
         assert logprob_token_ids_state is not None
         assert expanded_idx_mapping is not None
 
         if num_logprobs > 0:
             topk_token_ids = torch.topk(logits, num_logprobs, dim=-1).indices
             topk_token_ids = topk_token_ids.to(torch.int32)
         else:
-            # This tensor just used as an int32 pointer, data not accessed.
             topk_token_ids = logprob_token_ids_state.token_ids.gpu
 
         num_cols = max(num_logprobs, max_per_req_token_ids)
@@ -184,8 +157,11 @@ def compute_topk_logprobs(
             NUM_TOPK=num_logprobs,
             PADDED_COLS=triton.next_power_of_2(num_cols),
         )
-        logprobs = compute_token_logprobs(logits, logprob_token_ids)
-        logprobs = logprobs.masked_fill(~valid_mask, float("-inf"))
+        if logits_mode:
+            scores = logits.to(torch.float32).gather(-1, logprob_token_ids)
+        else:
+            scores = compute_token_logprobs(logits, logprob_token_ids)
+        scores = scores.masked_fill(~valid_mask, float("-inf"))
 
     token_ranks = torch.empty(batch_size, dtype=torch.int64, device=logits.device)
     _ranks_kernel[(batch_size,)](
@@ -198,7 +174,7 @@ def compute_topk_logprobs(
     )
     return LogprobsTensors(
         logprob_token_ids=logprob_token_ids,
-        logprobs=logprobs,
+        logprobs=scores,
         selected_token_ranks=token_ranks,
         cu_num_generated_tokens=cu_num_logits,
     )

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

@@ -9,10 +9,7 @@
 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_logits,
-    compute_topk_logprobs,
-)
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
 
 
 class PromptLogprobsWorker:
@@ -232,8 +229,12 @@ def compute_prompt_logprobs_with_chunking(
         )
         chunk_tgt = prompt_token_ids[start_idx:end_idx]
 
-        compute_fn = compute_topk_logits if logits_mode else compute_topk_logprobs
-        result = compute_fn(prompt_logits, requested_num, chunk_tgt)
+        result = compute_topk_logprobs(
+            prompt_logits,
+            requested_num,
+            chunk_tgt,
+            logits_mode=logits_mode,
+        )
         token_ids.append(result.logprob_token_ids)
         scores.append(result.logprobs)
         ranks.append(result.selected_token_ranks)

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

@@ -19,7 +19,6 @@
 from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
 from vllm.v1.worker.gpu.sample.logprob import (
     LogprobTokenIdsState,
-    compute_topk_logits,
     compute_topk_logprobs,
 )
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -106,22 +105,16 @@ def __call__(
             expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
             cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
             num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
-            if self.logprobs_mode.endswith("logits"):
-                logprobs_tensors = compute_topk_logits(
-                    logits,
-                    num_logprobs,
-                    sampled,
-                )
-            else:
-                logprobs_tensors = compute_topk_logprobs(
-                    logits,
-                    num_logprobs,
-                    sampled,
-                    cu_num_logits,
-                    logprob_token_ids_state=self.logprob_token_ids_state,
-                    expanded_idx_mapping=input_batch.expanded_idx_mapping,
-                    max_per_req_token_ids=max_per_req_token_ids,
-                )
+            logprobs_tensors = compute_topk_logprobs(
+                logits,
+                num_logprobs,
+                sampled,
+                cu_num_logits,
+                logprob_token_ids_state=self.logprob_token_ids_state,
+                expanded_idx_mapping=input_batch.expanded_idx_mapping,
+                max_per_req_token_ids=max_per_req_token_ids,
+                logits_mode=self.logprobs_mode.endswith("logits"),
+            )
         else:
             logprobs_tensors = None
 

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

@@ -11,10 +11,7 @@
     get_num_sampled_and_rejected,
 )
 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
-from vllm.v1.worker.gpu.sample.logprob import (
-    compute_topk_logits,
-    compute_topk_logprobs,
-)
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
@@ -94,13 +91,12 @@ def _get_logprobs_tensors(
             num_warps=1,
         )
         expanded_logits = num_logits != input_batch.idx_mapping.shape[0]
-        if self.sampler.logprobs_mode.endswith("logits"):
-            return compute_topk_logits(logits, max_num_logprobs, flat_sampled)
         return compute_topk_logprobs(
             logits,
             max_num_logprobs,
             flat_sampled,
             input_batch.cu_num_logits_np.tolist() if expanded_logits else None,
+            logits_mode=self.sampler.logprobs_mode.endswith("logits"),
         )
 
     def __call__(