{"record":{"id":"97f7af51c69a65a1","repo":"sgl-project/sglang","slug":"probs-must-be-2d-got-shape-tuple-probs-shape","errorCode":null,"errorMessage":"probs must be 2D, got shape={tuple(probs.shape)}","messagePattern":"probs must be 2D, got shape=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/sampling/renorm_triton.py","lineNumber":56,"sourceCode":"@triton.jit\ndef _normalize_kernel(\n    out_ptr,\n    row_sums_ptr,\n    numel,\n    vocab_size: tl.constexpr,\n    BLOCK_SIZE: tl.constexpr,\n):\n    offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    mask = offsets < numel\n    row = offsets // vocab_size\n    values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32)\n    denominator = tl.load(row_sums_ptr + row, mask=mask, other=1.0)\n    tl.store(out_ptr + offsets, values / denominator, mask=mask)\n\n\ndef _prepare_probs(probs: torch.Tensor) -> torch.Tensor:\n    if probs.ndim != 2:\n        raise ValueError(f\"probs must be 2D, got shape={tuple(probs.shape)}\")\n    if not probs.is_cuda:\n        raise ValueError(\"renorm kernels require a CUDA/HIP tensor\")\n    return probs.float().contiguous()\n\n\ndef _renorm_from_pivots(probs_fp32: torch.Tensor, pivots: torch.Tensor) -> torch.Tensor:\n    batch_size, vocab_size = probs_fp32.shape\n    num_chunks = triton.cdiv(vocab_size, _BLOCK_SIZE)\n    out = torch.empty_like(probs_fp32)\n    partial_sums = torch.empty(\n        (batch_size, num_chunks), device=probs_fp32.device, dtype=torch.float32\n    )\n    _mask_and_partial_sum_kernel[(batch_size, num_chunks)](\n        probs_fp32,\n        pivots,\n        out,\n        partial_sums,\n        vocab_size=vocab_size,","sourceCodeStart":38,"sourceCodeEnd":74,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/sampling/renorm_triton.py#L38-L74","documentation":"_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.","triggerScenarios":"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.","commonSituations":"Passing a single row without .unsqueeze(0), or feeding logits/hidden states [batch, seq, vocab] from prefill instead of final-step probabilities [batch, vocab].","solutions":["Reshape to [batch, vocab]: probs = probs.reshape(-1, probs.shape[-1]) or .unsqueeze(0) for a single row","Select the last timestep if you accidentally passed the full seq of logits, then softmax before renorm"],"exampleFix":"# before\nout = top_p_renorm_probs_triton(probs_1d, 0.9)\n# after\nout = top_p_renorm_probs_triton(probs_1d.unsqueeze(0), 0.9).squeeze(0)","handlingStrategy":"validation","validationCode":"if probs.ndim != 2:\n    probs = probs.reshape(-1, probs.shape[-1])","typeGuard":"def is_2d_probs(p): return p.ndim == 2","tryCatchPattern":null,"preventionTips":["Standardize sampling inputs as [batch, vocab] at the sampler boundary"],"tags":["sampling","top-p","top-k","renorm","shape-validation"],"backgroundTag":"tensor-shape-validation","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}