{"record":{"id":"e51bd3cdebbb7bc6","repo":"sgl-project/sglang","slug":"probs-must-be-2d-got-shape-tuple-probs-shape-e51bd3","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/top_p_renorm_triton.py","lineNumber":64,"sourceCode":"    offsets = tl.program_id(0) * 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 top_p_renorm_probs_triton(\n    probs: torch.Tensor, top_p: Union[torch.Tensor, float]\n) -> torch.Tensor:\n    \"\"\"Apply exact top-p thresholding and renormalize each probability row.\n\n    Sorting and prefix sums use PyTorch's device kernels because a vocabulary-sized\n    in-register Triton sort does not scale to 100K+ vocabularies. Triton performs\n    the bandwidth-heavy masking, partial reduction, and normalization.\n    \"\"\"\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(\"top_p_renorm_probs_triton requires a CUDA/HIP tensor\")\n\n    probs_fp32 = probs.float().contiguous()\n    batch_size, vocab_size = probs_fp32.shape\n    if batch_size == 0 or vocab_size == 0:\n        return probs_fp32\n\n    if isinstance(top_p, torch.Tensor):\n        top_ps = top_p.to(device=probs.device, dtype=torch.float32).reshape(-1)\n        if top_ps.numel() == 1:\n            top_ps = top_ps.expand(batch_size)\n        elif top_ps.numel() != batch_size:\n            raise ValueError(\n                f\"top_p must be scalar or have one value per row, got \"\n                f\"{top_ps.numel()} values for {batch_size} rows\"\n            )\n    else:","sourceCodeStart":46,"sourceCodeEnd":82,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/sampling/top_p_renorm_triton.py#L46-L82","documentation":"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.","triggerScenarios":"Calling top_p_renorm_probs_triton with probs.ndim != 2 — a single vocab vector, a scalar, or [batch, seq, vocab].","commonSituations":"Forgetting to unsqueeze a single-row distribution, or passing pre-prefill logits tensors with a sequence dimension instead of last-token probabilities.","solutions":["Reshape to exactly 2 dims: probs.reshape(-1, probs.shape[-1])","Apply softmax to logits first, and take the last timestep if input has a seq dim"],"exampleFix":"# before\nout = top_p_renorm_probs_triton(probs.unsqueeze(-1), 0.9)  # 3-D\n# after\nout = top_p_renorm_probs_triton(probs, 0.9)  # [batch, vocab]","handlingStrategy":"validation","validationCode":"if probs.ndim != 2:\n    probs = probs.reshape(-1, probs.shape[-1]) if probs.ndim > 2 else probs.unsqueeze(0)\nassert probs.is_cuda","typeGuard":"def valid_renorm_input(p): return p.ndim == 2 and p.is_cuda","tryCatchPattern":null,"preventionTips":["Pass last-token softmaxed probabilities, never seq-dim logits","Reshape at the sampler entry point"],"tags":["sampling","top-p","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"}