{"record":{"id":"ed5c04923c8398f1","repo":"sgl-project/sglang","slug":"next-token-logits-row-count-mismatch-expected-bs","errorCode":null,"errorMessage":"next_token_logits row count mismatch. Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}.","messagePattern":"next_token_logits row count mismatch\\. Expected (.+?), got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"critical","filePath":"python/sglang/srt/speculative/dflash_utils.py","lineNumber":905,"sourceCode":"    if not _DFLASH_SAMPLING_VERIFY_AVAILABLE:\n        raise RuntimeError(\n            \"DFLASH non-greedy verification is unavailable on this build/device.\"\n        )\n    if candidates.ndim != 2:\n        raise ValueError(f\"candidates must be 2D, got shape={tuple(candidates.shape)}\")\n    if next_token_logits.ndim != 2:\n        raise ValueError(\n            \"next_token_logits must be 2D, \"\n            f\"got shape={tuple(next_token_logits.shape)}.\"\n        )\n\n    bs, draft_token_num = candidates.shape\n    if bs <= 0:\n        raise ValueError(f\"batch size must be positive, got {bs}.\")\n    if draft_token_num <= 0:\n        raise ValueError(f\"draft_token_num must be positive, got {draft_token_num}.\")\n    if next_token_logits.shape[0] != bs * draft_token_num:\n        raise ValueError(\n            \"next_token_logits row count mismatch. \"\n            f\"Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}.\"\n        )\n    if candidates.device != next_token_logits.device:\n        raise ValueError(\n            \"candidates and next_token_logits must be on the same device, \"\n            f\"got {candidates.device} and {next_token_logits.device}.\"\n        )\n\n    if threshold_single is None:\n        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)","sourceCodeStart":887,"sourceCodeEnd":923,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/speculative/dflash_utils.py#L887-L923","documentation":"Raised by compute_dflash_sampling_correct_drafts_and_bonus in sglang's DFLASH speculative decoding when the target model's next_token_logits tensor does not have exactly bs*draft_token_num rows. The function verifies that every draft candidate position has a corresponding logit row before running the sampling/acceptance computation. A mismatch means the target forward produced a different number of logit rows than the draft candidate layout expects.","triggerScenarios":"Calling compute_dflash_sampling_correct_drafts_and_bonus(candidates, next_token_logits, ...) where candidates.shape == (bs, draft_token_num) but next_token_logits.shape[0] != bs*draft_token_num — e.g. the target model was run on a flattened (bs*draft_token_num) batch but bs or draft_token_num was recomputed inconsistently, or logits from a non-speculative forward (one row per sequence) were passed in.","commonSituations":"Mismatched draft_token_num between the DFLASH worker config and the actual target forward output; a target worker that only returns last-token logits; refactors that reshape logits before passing them into the speculative sampling utils; off-by-one in prefill vs decode batch assembly.","solutions":["Check how next_token_logits is produced upstream in forward_batch_generation and ensure the target model runs with num_tokens=bs*draft_token_num and returns one logit row per draft token","Verify candidates.shape and next_token_logits.shape[0] with an assert/log right before the call to identify which side diverges","Confirm draft_token_num used to build candidates matches the speculative algorithm config used to launch the target forward","If the target backend cannot produce per-draft-token logits, use a sampling path compatible with last-token-only logits instead of this function"],"exampleFix":"// before\nbonus, accept = compute_dflash_sampling_correct_drafts_and_bonus(\n    candidates, next_token_logits, ...\n)\n\n// after\nassert next_token_logits.shape[0] == candidates.numel(), (\n    f\"logits rows {next_token_logits.shape[0]} != \"\n    f\"candidates {candidates.shape[0]*candidates.shape[1]}\"\n)\nbonus, accept = compute_dflash_sampling_correct_drafts_and_bonus(\n    candidates, next_token_logits, ...\n)","handlingStrategy":"validation","validationCode":"bs, draft_token_num = candidates.shape\nassert next_token_logits.shape[0] == bs * draft_token_num, (\n    f\"next_token_logits rows={next_token_logits.shape[0]}, \"\n    f\"expected={bs * draft_token_num}\"\n)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always derive bs and draft_token_num from candidates.shape rather than separate variables","Run the target forward with num_tokens=bs*draft_token_num when using per-draft-token sampling","Add a shape assert immediately before speculative sampling calls in debug builds"],"tags":["sglang","speculative-decoding","dflash","shape-mismatch","tensor-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}