sgl-project/sglang · error · ValueError

next_token_logits row count mismatch for DFlash verify adjus

Error message

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

What it means

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.

Source

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

    """Apply sampling-time logit adjustments for DFlash verify in place.

    This keeps v1 and v2 verify semantics aligned while letting overlap scheduling
    use the cheaper precomputed `acc_linear_penalties` path instead of allocating a
    repeated `[bs * draft_token_num, vocab]` penalty tensor every step.
    """
    if sampling_info is None:
        return
    if next_token_logits.ndim != 2:
        raise ValueError(
            "next_token_logits must be 2D, "
            f"got shape={tuple(next_token_logits.shape)}."
        )
    if draft_token_num <= 0:
        raise ValueError(f"draft_token_num must be positive, got {draft_token_num}.")

    bs = len(sampling_info)
    if next_token_logits.shape[0] != bs * draft_token_num:
        raise ValueError(
            "next_token_logits row count mismatch for DFlash verify adjustments. "
            f"Expected {bs * draft_token_num}, got {next_token_logits.shape[0]}."
        )

    if sampling_info.has_custom_logit_processor:
        apply_custom_logit_processor(
            next_token_logits,
            sampling_info,
            num_tokens_in_batch=draft_token_num,
        )

    acc_linear_penalties = getattr(sampling_info, "acc_linear_penalties", None)
    penalizer = getattr(sampling_info, "penalizer_orchestrator", None)
    grammar_mask = getattr(sampling_info, "grammar_mask", None)
    logit_bias = getattr(sampling_info, "logit_bias", None)

    logits_3d: Optional[torch.Tensor] = None

View on GitHub (pinned to 0132848349)

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.

Example fix

# before
apply_dflash_verify_logits_adjustments(next_token_logits, sampling_info, draft_token_num=1)
# after
assert next_token_logits.shape[0] == len(sampling_info) * draft_token_num
apply_dflash_verify_logits_adjustments(next_token_logits, sampling_info, draft_token_num=draft_token_num)

Defensive patterns

Strategy: validation

Validate before calling

expected = len(sampling_info) * draft_token_num
assert next_token_logits.ndim == 2 and next_token_logits.shape[0] == expected, (next_token_logits.shape, expected)

Try / catch

try:
    apply_dflash_verify_logits_adjustments(...)
except ValueError as e:
    if 'row count mismatch' in str(e):
        raise RuntimeError(f'batch shape diverged: logits={next_token_logits.shape}') from e
    raise

Prevention

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.

When it happens

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

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

Related errors

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