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 15, 01:00 UTC
[v1][sampler] Inplace logprobs comparison to get the token rank (#21283)
Signed-off-by: Lu Fang <lufang@fb.com>
Lu Fang committed
Jul 21, 2025 at 13:47 UTC
8d0a01a5f2b53794e4bc6b734d7b63cb8a9b7d7d
2 files changed
+26
-1
vllm/v1/sample/ops/logprobs.py
new
+24
@@ -0,0 +1,24 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Some utilities for logprobs, including logits."""
4
+
5
+import torch
6
+
7
+
8
+@torch.compile(dynamic=True)
9
+def batched_count_greater_than(x: torch.Tensor,
10
+ values: torch.Tensor) -> torch.Tensor:
11
+ """
12
+ Counts elements in each row of x that are greater than the corresponding
13
+ value in values. Use torch.compile to generate an optimized kernel for
14
+ this function. otherwise, it will create additional copies of the input
15
+ tensors and cause memory issues.
16
+
17
+ Args:
18
+ x (torch.Tensor): A 2D tensor of shape (batch_size, n_elements).
19
+ values (torch.Tensor): A 2D tensor of shape (batch_size, 1).
20
+
21
+ Returns:
22
+ torch.Tensor: A 1D tensor of shape (batch_size,) with the counts.
23
+ """
24
+ return (x >= values).sum(-1)
vllm/v1/sample/sampler.py
+2
-1
@@ -9,6 +9,7 @@ from vllm.utils import is_pin_memory_available
9
from vllm.v1.outputs import LogprobsTensors, SamplerOutput
10
from vllm.v1.sample.metadata import SamplingMetadata
11
from vllm.v1.sample.ops.bad_words import apply_bad_words
12
+from vllm.v1.sample.ops.logprobs import batched_count_greater_than
13
from vllm.v1.sample.ops.penalties import apply_all_penalties
14
from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
15
@@ -174,7 +175,7 @@ class Sampler(nn.Module):
175
token_logprobs = logprobs.gather(-1, token_ids)
176
177
# Compute the ranks of the actual token.
177
- token_ranks = (logprobs >= token_logprobs).sum(-1)
178
+ token_ranks = batched_count_greater_than(logprobs, token_logprobs)
179
180
# Concatenate together with the topk.
181
indices = torch.cat((token_ids, topk_indices), dim=1)