Skip to content

vllm.v1.worker.gpu.sample.greedy_argmax

Functions:

  • greedy_argmax –

    Like logits.argmax(dim=-1), but treats NaNs as -inf on every path.

greedy_argmax(logits)

Like logits.argmax(dim=-1), but treats NaNs as -inf on every path. Ties resolve to the first index.

Source code in vllm/v1/worker/gpu/sample/greedy_argmax.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def greedy_argmax(logits: torch.Tensor) -> torch.Tensor:
    """Like `logits.argmax(dim=-1)`, but treats NaNs as -inf on every path.
    Ties resolve to the first index.
    """
    if (
        not logits.is_cuda
        or logits.dtype not in (torch.bfloat16, torch.float32)
        or logits.ndim != 2
        or logits.stride(1) != 1
        or logits.shape[1] < 4096
    ):
        return logits.masked_fill(logits.isnan(), float("-inf")).argmax(dim=-1)
    rows, vocab = logits.shape
    out = torch.empty(rows, dtype=torch.int64, device=logits.device)
    if rows == 0:
        return out
    block = 4096
    splits = triton.cdiv(vocab, block)
    values = torch.empty((rows, splits), dtype=torch.float32, device=logits.device)
    indices = torch.empty((rows, splits), dtype=torch.int32, device=logits.device)
    _argmax_partial[(rows, splits)](
        logits, values, indices, vocab, logits.stride(0), splits, block, num_warps=4
    )
    _argmax_finish[(rows,)](
        values, indices, out, splits, triton.next_power_of_2(splits), num_warps=1
    )
    return out