sgl-project/sglang · error · ValueError

candidates and next_token_logits must be on the same device,

Error message

candidates and next_token_logits must be on the same device, got {candidates.device} and {next_token_logits.device}.

What it means

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.

Source

Thrown at python/sglang/srt/speculative/dflash_utils.py:910

        raise ValueError(f"candidates must be 2D, got shape={tuple(candidates.shape)}")
    if next_token_logits.ndim != 2:
        raise ValueError(
            "next_token_logits must be 2D, "
            f"got shape={tuple(next_token_logits.shape)}."
        )

    bs, draft_token_num = candidates.shape
    if bs <= 0:
        raise ValueError(f"batch size must be positive, got {bs}.")
    if draft_token_num <= 0:
        raise ValueError(f"draft_token_num must be positive, got {draft_token_num}.")
    if next_token_logits.shape[0] != bs * draft_token_num:
        raise ValueError(
            "next_token_logits row count mismatch. "
            f"Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}."
        )
    if candidates.device != next_token_logits.device:
        raise ValueError(
            "candidates and next_token_logits must be on the same device, "
            f"got {candidates.device} and {next_token_logits.device}."
        )

    if threshold_single is None:
        from sglang.srt.runtime_context import get_spec

        threshold_single = get_spec().speculative_accept_threshold_single
    if threshold_acc is None:
        from sglang.srt.runtime_context import get_spec

        threshold_acc = get_spec().speculative_accept_threshold_acc
    threshold_single = float(threshold_single)
    threshold_acc = max(float(threshold_acc), 1e-9)

    device = next_token_logits.device

    if uniform_samples is None:

View on GitHub (pinned to 0132848349)

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

Example fix

// before
result = compute_dflash_sampling_correct_drafts_and_bonus(
    candidates, next_token_logits, ...
)

// after
next_token_logits = next_token_logits.to(candidates.device)
result = compute_dflash_sampling_correct_drafts_and_bonus(
    candidates, next_token_logits, ...
)

Defensive patterns

Strategy: validation

Validate before calling

if candidates.device != next_token_logits.device:
    next_token_logits = next_token_logits.to(candidates.device)

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/b1e287b81d38265a. Report an issue: GitHub.