sgl-project/sglang · error · ValueError

uniform_samples_for_final_sampling shape mismatch. Expected

Error message

uniform_samples_for_final_sampling shape mismatch. Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}.

What it means

Raised when a caller supplies uniform_samples_for_final_sampling whose shape is not exactly (bs,). This vector holds one uniform sample per sequence used for the final (post-acceptance) sampling step. The check mirrors the optional-samples validation for the acceptance samples and enforces one sample per batch row.

Source

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

    if uniform_samples is None:
        uniform_samples = torch.rand(
            (bs, draft_token_num), dtype=torch.float32, device=device
        )
    else:
        if uniform_samples.shape != (bs, draft_token_num):
            raise ValueError(
                "uniform_samples shape mismatch. "
                f"Expected {(bs, draft_token_num)}, got {tuple(uniform_samples.shape)}."
            )
        uniform_samples = uniform_samples.to(device=device, dtype=torch.float32)

    if uniform_samples_for_final_sampling is None:
        uniform_samples_for_final_sampling = torch.rand(
            (bs,), dtype=torch.float32, device=device
        )
    else:
        if uniform_samples_for_final_sampling.shape != (bs,):
            raise ValueError(
                "uniform_samples_for_final_sampling shape mismatch. "
                f"Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}."
            )
        uniform_samples_for_final_sampling = uniform_samples_for_final_sampling.to(
            device=device,
            dtype=torch.float32,
        )

    target_probs = build_dflash_verify_target_probs(
        next_token_logits=next_token_logits,
        sampling_info=sampling_info,
        draft_token_num=draft_token_num,
        bs=bs,
        max_top_k=max_top_k,
        uniform_top_k_value=uniform_top_k_value,
        use_sparse_topk=use_sparse_topk,
    )
    draft_probs = torch.zeros_like(target_probs)

View on GitHub (pinned to 0132848349)

Solutions

Generate with shape (bs,): torch.rand((bs,), dtype=torch.float32, device=device)
Pass None to let the function sample internally
Ensure the buffer is regenerated whenever bs changes (new requests, retract/replay)

Example fix

// before
final_samples = torch.rand(bs, draft_token_num)  # wrong shape
... uniform_samples_for_final_sampling=final_samples)

// after
final_samples = torch.rand(
    (bs,), dtype=torch.float32, device=candidates.device
)
... uniform_samples_for_final_sampling=final_samples)

Defensive patterns

Strategy: validation

Validate before calling

if uniform_samples_for_final_sampling is not None:
    assert tuple(uniform_samples_for_final_sampling.shape) == (candidates.shape[0],), (
        f"final samples {tuple(uniform_samples_for_final_sampling.shape)} != {(candidates.shape[0],)}"
    )

Prevention

Use a separate (bs,)-shaped buffer for final sampling; don't reuse the (bs, draft_token_num) acceptance samples
Regenerate per-batch-size random buffers on schedule changes

When it happens

Trigger: Passing a scalar, a (bs, 1) tensor, or a (bs, draft_token_num) tensor (reusing the acceptance samples) as uniform_samples_for_final_sampling; caching samples across batch-size changes.

Common situations: Replay/seeding harnesses that share one random buffer for both sampling stages; refactors that change decode batch size between sample generation and use.

Related errors

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