| 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 | from vllm.platforms import current_platform |
| 8 | |
| 9 | |
| 10 | @torch.compile(backend=current_platform.simple_compile_backend) |
| 11 | def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.Tensor: |
| 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 | """ |
| 25 | torch._check(x.shape[0] >= 1) |
| 26 | torch._check(x.shape[0] == values.shape[0]) |
| 27 | return (x >= values).sum(-1) |