{"record":{"id":"ae3f74b6e8e871b3","repo":"sgl-project/sglang","slug":"next-token-logits-row-count-mismatch-for-dflash-ve","errorCode":null,"errorMessage":"next_token_logits row count mismatch for DFlash verify adjustments. Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}.","messagePattern":"next_token_logits row count mismatch for DFlash verify adjustments\\. Expected (.+?), got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/speculative/dflash_utils.py","lineNumber":231,"sourceCode":"    \"\"\"Apply sampling-time logit adjustments for DFlash verify in place.\n\n    This keeps v1 and v2 verify semantics aligned while letting overlap scheduling\n    use the cheaper precomputed `acc_linear_penalties` path instead of allocating a\n    repeated `[bs * draft_token_num, vocab]` penalty tensor every step.\n    \"\"\"\n    if sampling_info is None:\n        return\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    if draft_token_num <= 0:\n        raise ValueError(f\"draft_token_num must be positive, got {draft_token_num}.\")\n\n    bs = len(sampling_info)\n    if next_token_logits.shape[0] != bs * draft_token_num:\n        raise ValueError(\n            \"next_token_logits row count mismatch for DFlash verify adjustments. \"\n            f\"Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}.\"\n        )\n\n    if sampling_info.has_custom_logit_processor:\n        apply_custom_logit_processor(\n            next_token_logits,\n            sampling_info,\n            num_tokens_in_batch=draft_token_num,\n        )\n\n    acc_linear_penalties = getattr(sampling_info, \"acc_linear_penalties\", None)\n    penalizer = getattr(sampling_info, \"penalizer_orchestrator\", None)\n    grammar_mask = getattr(sampling_info, \"grammar_mask\", None)\n    logit_bias = getattr(sampling_info, \"logit_bias\", None)\n\n    logits_3d: Optional[torch.Tensor] = None\n","sourceCodeStart":213,"sourceCodeEnd":249,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/speculative/dflash_utils.py#L213-L249","documentation":"Raised by apply_dflash_verify_logits_adjustments when next_token_logits.shape[0] != len(sampling_info) * draft_token_num. During DFlash verify, each sequence contributes exactly draft_token_num logit rows, so the row count must equal batch size times draft tokens. A mismatch means the logits tensor and the sampling metadata describe different batch shapes.","triggerScenarios":"Passing a logits tensor whose batch dimension doesn't cover bs * draft_token_num rows, e.g. draft_token_num passed as 1 while the draft produced 8 tokens per sequence, or a partially sliced logits tensor from a custom model runner.","commonSituations":"Mismatch between the draft token count used to run the draft model and the count passed at verify time; off-by-one slicing of next_token_logits in a custom verify path; chunked/mixed batches where the logits were reshaped incorrectly.","solutions":["Ensure the draft_token_num argument matches the value used when generating the draft tokens.","Verify the logits tensor was produced for the full bs * draft_token_num rows (check .shape[0] before calling).","Print/inspect len(sampling_info) and next_token_logits.shape to find which side diverges, then fix batch assembly upstream."],"exampleFix":"# before\napply_dflash_verify_logits_adjustments(next_token_logits, sampling_info, draft_token_num=1)\n# after\nassert next_token_logits.shape[0] == len(sampling_info) * draft_token_num\napply_dflash_verify_logits_adjustments(next_token_logits, sampling_info, draft_token_num=draft_token_num)","handlingStrategy":"validation","validationCode":"expected = len(sampling_info) * draft_token_num\nassert next_token_logits.ndim == 2 and next_token_logits.shape[0] == expected, (next_token_logits.shape, expected)","typeGuard":null,"tryCatchPattern":"try:\n    apply_dflash_verify_logits_adjustments(...)\nexcept ValueError as e:\n    if 'row count mismatch' in str(e):\n        raise RuntimeError(f'batch shape diverged: logits={next_token_logits.shape}') from e\n    raise","preventionTips":["Carry draft_token_num through the whole draft->verify pipeline as one source of truth.","Log logits shape and batch size at verify entry in debug builds."],"tags":["speculative-decoding","dflash","shape-mismatch"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}