sgl-project/sglang · critical · ValueError

next_token_logits row count mismatch. Expected {bs * draft_t

Error message

next_token_logits row count mismatch. Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}.

What it means

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.

Source

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

    if not _DFLASH_SAMPLING_VERIFY_AVAILABLE:
        raise RuntimeError(
            "DFLASH non-greedy verification is unavailable on this build/device."
        )
    if candidates.ndim != 2:
        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)

View on GitHub (pinned to 0132848349)

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

Example fix

// before
bonus, accept = compute_dflash_sampling_correct_drafts_and_bonus(
    candidates, next_token_logits, ...
)

// after
assert next_token_logits.shape[0] == candidates.numel(), (
    f"logits rows {next_token_logits.shape[0]} != "
    f"candidates {candidates.shape[0]*candidates.shape[1]}"
)
bonus, accept = compute_dflash_sampling_correct_drafts_and_bonus(
    candidates, next_token_logits, ...
)

Defensive patterns

Strategy: validation

Validate before calling

bs, draft_token_num = candidates.shape
assert next_token_logits.shape[0] == bs * draft_token_num, (
    f"next_token_logits rows={next_token_logits.shape[0]}, "
    f"expected={bs * draft_token_num}"
)

Prevention

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

When it happens

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

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

Related errors

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