{"record":{"id":"6b9aac25e5fa23d9","repo":"sgl-project/sglang","slug":"uniform-samples-shape-mismatch-expected-bs-dra","errorCode":null,"errorMessage":"uniform_samples shape mismatch. Expected {(bs, draft_token_num)}, got {tuple(uniform_samples.shape)}.","messagePattern":"uniform_samples shape mismatch\\. Expected (.+?), got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/speculative/dflash_utils.py","lineNumber":934,"sourceCode":"        from sglang.srt.runtime_context import get_spec\n\n        threshold_single = get_spec().speculative_accept_threshold_single\n    if threshold_acc is None:\n        from sglang.srt.runtime_context import get_spec\n\n        threshold_acc = get_spec().speculative_accept_threshold_acc\n    threshold_single = float(threshold_single)\n    threshold_acc = max(float(threshold_acc), 1e-9)\n\n    device = next_token_logits.device\n\n    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,","sourceCodeStart":916,"sourceCodeEnd":952,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/speculative/dflash_utils.py#L916-L952","documentation":"Raised when a caller explicitly supplies uniform_samples to compute_dflash_sampling_correct_drafts_and_bonus but its shape is not exactly (bs, draft_token_num). These samples drive deterministic/replayable acceptance decisions, so the shape must match the candidate layout exactly. When the argument is None the function samples correctly-shaped randoms itself.","triggerScenarios":"Passing uniform_samples with an extra batch dimension (e.g. (1, bs, draft_token_num) from a leftover model forward), reusing samples cached from a previous batch size, or generating samples with a hardcoded shape that doesn't track bs/draft_token_num.","commonSituations":"Seeded-reproducibility test harnesses that pre-generate random tensors; cached RNG state replay across differently-sized batches; refactoring that changes draft_token_num without updating test fixtures.","solutions":["Generate samples with the exact shape: torch.rand((bs, draft_token_num), dtype=torch.float32, device=device) mirroring the internal default","Omit the argument (pass None) to let the function sample internally","If replaying captured samples, slice/resize them to the current (bs, draft_token_num) before passing"],"exampleFix":"// before\nuniform_samples = torch.rand(bs * draft_token_num)  # wrong: 1D\ncompute_dflash_sampling_correct_drafts_and_bonus(..., uniform_samples=uniform_samples)\n\n// after\nuniform_samples = torch.rand(\n    (bs, draft_token_num), dtype=torch.float32, device=candidates.device\n)\ncompute_dflash_sampling_correct_drafts_and_bonus(..., uniform_samples=uniform_samples)","handlingStrategy":"validation","validationCode":"if uniform_samples is not None:\n    bs, draft_token_num = candidates.shape\n    assert tuple(uniform_samples.shape) == (bs, draft_token_num), (\n        f\"uniform_samples {tuple(uniform_samples.shape)} != {(bs, draft_token_num)}\"\n    )","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer passing None and letting the function sample internally unless you need determinism","Regenerate seeded samples whenever batch size or draft_token_num changes","Shape-check captured replay tensors against the current candidate shape before use"],"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"}