Skip to content

sampler: greedy paths report log of raw logit instead of log of probability #2329

Description

@dx2102

Symptom

Greedy request with logprobs=5 returns positive numbers:

'were'  logprob=+1.322
'was'   logprob=+1.314

Positive values cannot be natural log probabilities since ln(p) <= 0 for p <= 1. After #2327 lands these become ln(raw_logit) instead of log10(raw_logit); neither is ln(softmax(logits)[token]).

Root cause

Sampler::sample has two None (greedy) arms:

  • non-speculative greedy passes the post-penalty logits into sample_argmax
  • speculative greedy passes the post-penalty logits into sample_speculative_top_kp_min_p

Neither arm applies softmax_last_dim, while both Some(temperature) arms do. Both callees then bind the input to let probs: Vec<f32> = logits.to_vec1()?; and treat that variable as a probability distribution: .ln() at the end (both functions), plus top-p and min-p filters that assume values in [0, 1] (only sample_speculative_top_kp_min_p).

The sampled token is usually still correct because argmax_f32 is invariant under softmax. The wrong logprob value is only surfaced to the caller when return_logprobs=true; otherwise it is filled and never read.

Repro

import requests
r = requests.post("http://127.0.0.1:8766/v1/completions", json={
    "model": "Qwen/Qwen3-0.6B",
    "prompt": "Once upon a time, there",
    "max_tokens": 1,
    "temperature": 0.0, "top_p": 1.0, "top_k": 1,
    "logprobs": 5,
})
for t in r.json()["choices"][0]["logprobs"]["content"][0]["top_logprobs"]:
    print(f"  {t['bytes']!r:12}  logprob={t['logprob']:+.4f}")

HF transformers log_softmax reference on the same model and prompt:

'were'   ln(prob)=-0.809
'was'    ln(prob)=-1.059
'are'    ln(prob)=-1.934

Other issues found in sampler.rs while looking at this

  • sample_speculative_top_kp_min_p has "sample" in its name but its last step is argmax_f32(&probs). Its top-k, top-p, min-p filters therefore never change the sampled token, only the probability distribution returned to the user.
  • Both sample_speculative_top_kp_min_p and sample_top_kp_min_p apply all three sampling filters (top-k, top-p, min-p) to their local probs variable, then pass that filtered variable into get_top_logprobs. The API's top_logprobs list (the user-requested count of alternatives) is therefore drawn from post-filter values rather than from the model's real distribution. On temperature paths that filtered variable is a real probability distribution; on the greedy speculative path (see primary root cause) it is post-penalty logits treated as probs and filtered by comparing logit magnitudes against thresholds in [0, 1], so top_logprobs there is doubly wrong: neither derived from probabilities nor filtered on meaningful values.
  • sample_speculative_top_kp_min_p's min-p block computes threshold = max_p * min_p. After penalty application probs can be negative; when max_p is negative, threshold > max_p, so the check zeros every entry including the max, and the subsequent argmax_f32 falls back to index 0. In this corner case the sampled token itself is wrong, not just the reported logprob.
  • sample_speculative_top_kp_min_p lacks the "top_p / min_p out of range means early return to multinomial" guards that sample_top_kp_min_p has.
  • test_gumbel_speculative has no Gumbel sampling in the tested code path.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions