{"record":{"id":"b1e287b81d38265a","repo":"sgl-project/sglang","slug":"candidates-and-next-token-logits-must-be-on-the-sa","errorCode":null,"errorMessage":"candidates and next_token_logits must be on the same device, got {candidates.device} and {next_token_logits.device}.","messagePattern":"candidates and next_token_logits must be on the same device, got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/speculative/dflash_utils.py","lineNumber":910,"sourceCode":"        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)\n    threshold_acc = max(float(threshold_acc), 1e-9)\n\n    device = next_token_logits.device\n\n    if uniform_samples is None:","sourceCodeStart":892,"sourceCodeEnd":928,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/speculative/dflash_utils.py#L892-L928","documentation":"Raised when candidates and next_token_logits live on different CUDA devices (or one is on CPU while the other is on GPU) inside DFLASH speculative sampling. The function performs elementwise tensor ops between the two, so PyTorch would otherwise fail later with a device mismatch; this check fails fast with a clear message. It indicates tensors were created or moved with inconsistent device arguments.","triggerScenarios":"Passing candidates sampled on cuda:0 but logits computed on cuda:1 (multi-GPU TP where the draft samples on one rank and logits are gathered from another), or passing a CPU tensor for one argument and a GPU tensor for the other.","commonSituations":"Tensor-parallel or pipeline setups where tensors are gathered across ranks; test code that creates torch.rand tensors without device=; refactors that drop a .to(device) call; CUDA_VISIBLE_DEVICES remapping making 'cuda' resolve differently in different processes.","solutions":["Move both tensors to the same device before calling: candidates = candidates.to(next_token_logits.device)","Audit upstream code for where each tensor is allocated and pass an explicit device (the target worker's device) everywhere","In TP setups, ensure logits and sampled candidates are gathered onto the same rank/device before speculative sampling"],"exampleFix":"// before\nresult = compute_dflash_sampling_correct_drafts_and_bonus(\n    candidates, next_token_logits, ...\n)\n\n// after\nnext_token_logits = next_token_logits.to(candidates.device)\nresult = compute_dflash_sampling_correct_drafts_and_bonus(\n    candidates, next_token_logits, ...\n)","handlingStrategy":"validation","validationCode":"if candidates.device != next_token_logits.device:\n    next_token_logits = next_token_logits.to(candidates.device)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Standardize one device variable (e.g. the target worker's device) and use it for every tensor allocation","In TP setups, gather logits and candidates onto the same rank before speculative sampling","Never allocate tensors without an explicit device= argument in serving code"],"tags":["sglang","speculative-decoding","dflash","cuda","device-mismatch"],"backgroundTag":"torch-device-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}