Skip to content

vllm.v1.sample.ops.logprobs

Some utilities for logprobs, including logits.

Functions:

batched_count_greater_than(x, values)

Counts elements in each row of x that are greater than the corresponding value in values. Use torch.compile to generate an optimized kernel for this function. otherwise, it will create additional copies of the input tensors and cause memory issues.

Parameters:

  • x

    (Tensor) –

    A 2D tensor of shape (batch_size, n_elements).

  • values

    (Tensor) –

    A 2D tensor of shape (batch_size, 1).

Returns:

  • Tensor –

    torch.Tensor: A 1D tensor of shape (batch_size,) with the counts.

Source code in vllm/v1/sample/ops/logprobs.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
@torch.compile(backend=current_platform.simple_compile_backend)
def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
    """
    Counts elements in each row of x that are greater than the corresponding
    value in values.  Use torch.compile to generate an optimized kernel for
    this function. otherwise, it will create additional copies of the input
    tensors and cause memory issues.

    Args:
        x (torch.Tensor): A 2D tensor of shape (batch_size, n_elements).
        values (torch.Tensor): A 2D tensor of shape (batch_size, 1).

    Returns:
        torch.Tensor: A 1D tensor of shape (batch_size,) with the counts.
    """
    torch._check(x.shape[0] >= 1)
    torch._check(x.shape[0] == values.shape[0])
    return (x >= values).sum(-1)