vllm.v1.worker.gpu.sample.greedy_argmax

Functions:

• greedy_argmax –

  logits.argmax(dim=-1) over a large vocabulary, split across CTAs.

greedy_argmax(logits)

logits.argmax(dim=-1) over a large vocabulary, split across CTAs.

Ties resolve to the first index. NaN logits are treated as -inf, so a row with NaN may differ from torch.argmax, which returns the NaN's index.

Source code in vllm/v1/worker/gpu/sample/greedy_argmax.py

def greedy_argmax(logits: torch.Tensor) -> torch.Tensor:
    """``logits.argmax(dim=-1)`` over a large vocabulary, split across CTAs.

    Ties resolve to the first index. NaN logits are treated as -inf, so a row
    with NaN may differ from ``torch.argmax``, which returns the NaN's 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.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