sgl-project/sglang · error · ValueError

probs must be 2D, got shape={tuple(probs.shape)}

Error message

probs must be 2D, got shape={tuple(probs.shape)}

What it means

top_p_renorm_probs_triton (the standalone variant in top_p_renorm_triton.py) requires a 2-D [batch, vocab] probability tensor: it sorts rows and computes per-row prefix sums. Tensors with any other rank (1-D vector, 3-D seq of logits) are rejected up front.

Source

Thrown at python/sglang/kernels/ops/sampling/top_p_renorm_triton.py:64

    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < numel
    row = offsets // vocab_size
    values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    denominator = tl.load(row_sums_ptr + row, mask=mask, other=1.0)
    tl.store(out_ptr + offsets, values / denominator, mask=mask)


def top_p_renorm_probs_triton(
    probs: torch.Tensor, top_p: Union[torch.Tensor, float]
) -> torch.Tensor:
    """Apply exact top-p thresholding and renormalize each probability row.

    Sorting and prefix sums use 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.
    """
    if probs.ndim != 2:
        raise ValueError(f"probs must be 2D, got shape={tuple(probs.shape)}")
    if not probs.is_cuda:
        raise ValueError("top_p_renorm_probs_triton requires a CUDA/HIP tensor")

    probs_fp32 = probs.float().contiguous()
    batch_size, vocab_size = probs_fp32.shape
    if batch_size == 0 or vocab_size == 0:
        return probs_fp32

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

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape to exactly 2 dims: probs.reshape(-1, probs.shape[-1])
2. Apply softmax to logits first, and take the last timestep if input has a seq dim

Example fix

# before
out = top_p_renorm_probs_triton(probs.unsqueeze(-1), 0.9)  # 3-D
# after
out = top_p_renorm_probs_triton(probs, 0.9)  # [batch, vocab]

if probs.ndim != 2:
    probs = probs.reshape(-1, probs.shape[-1]) if probs.ndim > 2 else probs.unsqueeze(0)
assert probs.is_cuda

def valid_renorm_input(p): return p.ndim == 2 and p.is_cuda

When it happens

Trigger: Calling top_p_renorm_probs_triton with probs.ndim != 2 — a single vocab vector, a scalar, or [batch, seq, vocab].

Common situations: Forgetting to unsqueeze a single-row distribution, or passing pre-prefill logits tensors with a sequence dimension instead of last-token probabilities.

Related errors

• probs must be 2D, got shape={tuple(probs.shape)}
• Input probs contains NaN.
• renorm kernels require a CUDA/HIP tensor
• top_p must be scalar or have one value per row, got {top_ps.
• top_p values must be in (0, 1]

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