sgl-project/sglang · error · ValueError

top_k must be scalar or have one value per row, got {top_ks.

Error message

top_k must be scalar or have one value per row, got {top_ks.numel()} values for {batch_size} rows

What it means

top_k_renorm_probs_triton accepts a scalar top_k or a per-row int tensor whose length must equal the batch size. A tensor whose element count is neither 1 nor batch_size cannot be mapped onto rows and is rejected.

Source

Thrown at python/sglang/kernels/ops/sampling/renorm_triton.py:154

    probs: torch.Tensor, top_k: Union[torch.Tensor, int]
) -> torch.Tensor:
    """Apply exact top-k thresholding and renormalize each probability row.

    Sorting uses PyTorch's device kernels because a vocabulary-sized in-register
    Triton sort does not scale to 100K+ vocabularies. Triton performs the
    bandwidth-heavy masking, partial reduction, and normalization.
    """
    probs_fp32 = _prepare_probs(probs)
    batch_size, vocab_size = probs_fp32.shape
    if batch_size == 0 or vocab_size == 0:
        return probs_fp32

    if isinstance(top_k, torch.Tensor):
        top_ks = top_k.to(device=probs.device, dtype=torch.int64).reshape(-1)
        if top_ks.numel() == 1:
            top_ks = top_ks.expand(batch_size)
        elif top_ks.numel() != batch_size:
            raise ValueError(
                f"top_k must be scalar or have one value per row, got "
                f"{top_ks.numel()} values for {batch_size} rows"
            )
    else:
        top_ks = torch.full(
            (batch_size,), int(top_k), device=probs.device, dtype=torch.int64
        )

    # Match FlashInfer's threshold semantics: sort descending, keep the k highest
    # probabilities, and retain all ties at the pivot.
    sorted_probs = torch.sort(probs_fp32, dim=-1, descending=True).values
    cutoff = (top_ks - 1).clamp_(min=0, max=vocab_size - 1)
    pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous()

    return _renorm_from_pivots(probs_fp32, pivots)


__all__ = ["top_k_renorm_probs_triton", "top_p_renorm_probs_triton"]

View on GitHub (pinned to 0132848349)

Solutions

1. Ensure top_k.numel() == 1 or == probs.shape[0]; reshape(-1) first
2. Rebuild per-row top_k arrays whenever the batch composition changes

Example fix

# before
top_k = torch.full((batch_size + 1,), 50, device='cuda')
out = top_k_renorm_probs_triton(probs, top_k)
# after
top_k = torch.full((batch_size,), 50, device='cuda')
out = top_k_renorm_probs_triton(probs, top_k)

if isinstance(top_k, torch.Tensor):
    assert top_k.numel() in (1, probs.shape[0])

def top_k_shape_ok(tk, batch): return not isinstance(tk, torch.Tensor) or tk.numel() in (1, batch)

When it happens

Trigger: Passing a top_k tensor with the wrong length, e.g. [num_requests] while probs has a different batch size, or a 2-D [batch, something>1] tensor.

Common situations: Batched sampling where the per-request top_k array was built for a subset of the batch, or leftover top_k tensors from a previous larger batch being reused.

Related errors

• top_p must be scalar or have one value per row, got {top_ps.
• kv-canary: RealKvSource.read_bytes must be a positive multip
• probs must be 2D, got shape={tuple(probs.shape)}
• top_p values must be in (0, 1]
• Multi-output conditioning requires prompt text so the prompt

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/dc7260e012b24066. Report an issue: GitHub.