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

_prepare_probs validates the probability tensor for the top-p/top-k renormalization Triton kernels, which process a batch of rows in parallel and therefore require a 2-D [batch, vocab] layout. 1-D, 3-D, or 0-dim tensors are rejected before any kernel launch.

Source

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

@triton.jit
def _normalize_kernel(
    out_ptr,
    row_sums_ptr,
    numel,
    vocab_size: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    offsets = tl.program_id(0).to(tl.int64) * 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 _prepare_probs(probs: torch.Tensor) -> torch.Tensor:
    if probs.ndim != 2:
        raise ValueError(f"probs must be 2D, got shape={tuple(probs.shape)}")
    if not probs.is_cuda:
        raise ValueError("renorm kernels require a CUDA/HIP tensor")
    return probs.float().contiguous()


def _renorm_from_pivots(probs_fp32: torch.Tensor, pivots: torch.Tensor) -> torch.Tensor:
    batch_size, vocab_size = probs_fp32.shape
    num_chunks = triton.cdiv(vocab_size, _BLOCK_SIZE)
    out = torch.empty_like(probs_fp32)
    partial_sums = torch.empty(
        (batch_size, num_chunks), device=probs_fp32.device, dtype=torch.float32
    )
    _mask_and_partial_sum_kernel[(batch_size, num_chunks)](
        probs_fp32,
        pivots,
        out,
        partial_sums,
        vocab_size=vocab_size,

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape to [batch, vocab]: probs = probs.reshape(-1, probs.shape[-1]) or .unsqueeze(0) for a single row
2. Select the last timestep if you accidentally passed the full seq of logits, then softmax before renorm

Example fix

# before
out = top_p_renorm_probs_triton(probs_1d, 0.9)
# after
out = top_p_renorm_probs_triton(probs_1d.unsqueeze(0), 0.9).squeeze(0)

if probs.ndim != 2:
    probs = probs.reshape(-1, probs.shape[-1])

def is_2d_probs(p): return p.ndim == 2

When it happens

Trigger: Calling top_p_renorm_probs_triton or top_k_renorm_probs_triton with a 1-D vocab vector, a 3-D tensor, or a scalar.

Common situations: Passing a single row without .unsqueeze(0), or feeding logits/hidden states [batch, seq, vocab] from prefill instead of final-step probabilities [batch, vocab].

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/97f7af51c69a65a1. Report an issue: GitHub.