sgl-project/sglang · error · ValueError

uniform_samples shape mismatch. Expected {(bs, draft_token_n

Error message

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

What it means

Raised when a caller explicitly supplies uniform_samples to compute_dflash_sampling_correct_drafts_and_bonus but its shape is not exactly (bs, draft_token_num). These samples drive deterministic/replayable acceptance decisions, so the shape must match the candidate layout exactly. When the argument is None the function samples correctly-shaped randoms itself.

Source

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

        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:
        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,

View on GitHub (pinned to 0132848349)

Solutions

Generate samples with the exact shape: torch.rand((bs, draft_token_num), dtype=torch.float32, device=device) mirroring the internal default
Omit the argument (pass None) to let the function sample internally
If replaying captured samples, slice/resize them to the current (bs, draft_token_num) before passing

Example fix

// before
uniform_samples = torch.rand(bs * draft_token_num)  # wrong: 1D
compute_dflash_sampling_correct_drafts_and_bonus(..., uniform_samples=uniform_samples)

// after
uniform_samples = torch.rand(
    (bs, draft_token_num), dtype=torch.float32, device=candidates.device
)
compute_dflash_sampling_correct_drafts_and_bonus(..., uniform_samples=uniform_samples)

Defensive patterns

Strategy: validation

Validate before calling

if uniform_samples is not None:
    bs, draft_token_num = candidates.shape
    assert tuple(uniform_samples.shape) == (bs, draft_token_num), (
        f"uniform_samples {tuple(uniform_samples.shape)} != {(bs, draft_token_num)}"
    )

Prevention

Prefer passing None and letting the function sample internally unless you need determinism
Regenerate seeded samples whenever batch size or draft_token_num changes
Shape-check captured replay tensors against the current candidate shape before use

When it happens

Trigger: Passing uniform_samples with an extra batch dimension (e.g. (1, bs, draft_token_num) from a leftover model forward), reusing samples cached from a previous batch size, or generating samples with a hardcoded shape that doesn't track bs/draft_token_num.

Common situations: Seeded-reproducibility test harnesses that pre-generate random tensors; cached RNG state replay across differently-sized batches; refactoring that changes draft_token_num without updating test fixtures.

Related errors

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