{"record":{"id":"f25430c164616a3e","repo":"sgl-project/sglang","slug":"uniform-samples-for-final-sampling-shape-mismatch","errorCode":null,"errorMessage":"uniform_samples_for_final_sampling shape mismatch. Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}.","messagePattern":"uniform_samples_for_final_sampling shape mismatch\\. Expected (.+?), got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/speculative/dflash_utils.py","lineNumber":946,"sourceCode":"    if uniform_samples is None:\n        uniform_samples = torch.rand(\n            (bs, draft_token_num), dtype=torch.float32, device=device\n        )\n    else:\n        if uniform_samples.shape != (bs, draft_token_num):\n            raise ValueError(\n                \"uniform_samples shape mismatch. \"\n                f\"Expected {(bs, draft_token_num)}, got {tuple(uniform_samples.shape)}.\"\n            )\n        uniform_samples = uniform_samples.to(device=device, dtype=torch.float32)\n\n    if uniform_samples_for_final_sampling is None:\n        uniform_samples_for_final_sampling = torch.rand(\n            (bs,), dtype=torch.float32, device=device\n        )\n    else:\n        if uniform_samples_for_final_sampling.shape != (bs,):\n            raise ValueError(\n                \"uniform_samples_for_final_sampling shape mismatch. \"\n                f\"Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}.\"\n            )\n        uniform_samples_for_final_sampling = uniform_samples_for_final_sampling.to(\n            device=device,\n            dtype=torch.float32,\n        )\n\n    target_probs = build_dflash_verify_target_probs(\n        next_token_logits=next_token_logits,\n        sampling_info=sampling_info,\n        draft_token_num=draft_token_num,\n        bs=bs,\n        max_top_k=max_top_k,\n        uniform_top_k_value=uniform_top_k_value,\n        use_sparse_topk=use_sparse_topk,\n    )\n    draft_probs = torch.zeros_like(target_probs)","sourceCodeStart":928,"sourceCodeEnd":964,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/speculative/dflash_utils.py#L928-L964","documentation":"Raised when a caller supplies uniform_samples_for_final_sampling whose shape is not exactly (bs,). This vector holds one uniform sample per sequence used for the final (post-acceptance) sampling step. The check mirrors the optional-samples validation for the acceptance samples and enforces one sample per batch row.","triggerScenarios":"Passing a scalar, a (bs, 1) tensor, or a (bs, draft_token_num) tensor (reusing the acceptance samples) as uniform_samples_for_final_sampling; caching samples across batch-size changes.","commonSituations":"Replay/seeding harnesses that share one random buffer for both sampling stages; refactors that change decode batch size between sample generation and use.","solutions":["Generate with shape (bs,): torch.rand((bs,), dtype=torch.float32, device=device)","Pass None to let the function sample internally","Ensure the buffer is regenerated whenever bs changes (new requests, retract/replay)"],"exampleFix":"// before\nfinal_samples = torch.rand(bs, draft_token_num)  # wrong shape\n... uniform_samples_for_final_sampling=final_samples)\n\n// after\nfinal_samples = torch.rand(\n    (bs,), dtype=torch.float32, device=candidates.device\n)\n... uniform_samples_for_final_sampling=final_samples)","handlingStrategy":"validation","validationCode":"if uniform_samples_for_final_sampling is not None:\n    assert tuple(uniform_samples_for_final_sampling.shape) == (candidates.shape[0],), (\n        f\"final samples {tuple(uniform_samples_for_final_sampling.shape)} != {(candidates.shape[0],)}\"\n    )","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use a separate (bs,)-shaped buffer for final sampling; don't reuse the (bs, draft_token_num) acceptance samples","Regenerate per-batch-size random buffers on schedule changes"],"tags":["sglang","speculative-decoding","dflash","shape-mismatch","rng"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}