[pnb]

Overview

I use logprobs with /completion. It is very useful. Recently, it seemed to stop working, but eventually I realized that was because I finally started using speculative decoding.

There was a feature request about it (#24271) from @drazdra, auto-closed due to inactivity. I also found this comment from @ServeurpersoCom that made it seem probably feasible to add. With the help of Qwen 3.8-27B and llama.cpp itself, it was.

Example:

curl http://localhost:6666/completion -H "Content-Type: application/json" \
  -d '{"prompt": "The capital of France is", "n_predict": 10, "logprobs": 3, "seed": 1}'

On master, this will produce correct output for the first token only:

      "logprob": -4.361362457275391,
      "top_logprobs": [
        {
          "id": 11751,
          "token": " Paris",
          "bytes": [
            32,
            80,
            97,
            114,
            105,
            115
          ],
          "logprob": -0.2336617261171341
        },
        {
          "id": 264,
          "token": " a",
          "bytes": [
            32,
            97
          ],
          "logprob": -3.882844924926758
        },
        {
          "id": 3750,
          "token": " known",
          "bytes": [
            32,
            107,
            110,
            111,
            119,
            110
          ],
          "logprob": -4.30119514465332
        }
      ]

After that first token, logprobs are missing. Like:

    {
      "id": 13,
      "token": ".",
      "bytes": [
        46
      ],
      "logprob": 0.0,
      "top_logprobs": []
    },

Whereas on this PR, they are correctly returned for all tokens. From my tests, the logprobs match the same values given by master branch without speculative decoding.

Additional information

Fixes #24271. Also there was a TODO in the source code about this issue.

Note that it still does not work for post_sampling_probs because that is a whole more complicated beast.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, pi + Qwen 3.8-27B + llama.cpp. Changes are very minimal and I've reviewed (and slightly edited) them myself.

[ngxson]

/bot review

[ggml-gh-bot]

Automated code review

Reviewed tools/server/server-context.cpp (the only changed file) plus the surrounding code paths it depends on: populate_token_probs / get_token_probabilities, common_sampler_sample_and_accept_n, llama_get_logits_ith, the sub-batch decode loop, post_decode, and the response serialization in server-task.cpp.

Verdict

The change is small, correct, and reuses existing infrastructure rather than inventing new machinery. I traced the indexing, the off-by-one risk, the move semantics, and the partial-accept/replay path; all check out. No blocking findings. A couple of non-blocking suggestions below.

Correctness notes (verified, no action needed)

• Indexing is consistent with the existing spec sampler. spec_i_batch[i] is the same index handed to common_sampler_sample inside common_sampler_sample_and_accept_n (which already works on master), and get_token_probabilities -> llama_get_logits_ith resolves that index through the same output_ids mapping. The post_decode sub-batch check at server-context.cpp:3698 already guarantees all spec_i_batch entries lie in the current sub-batch, which for a spec slot resolves to off == 0, so absolute and sub-batch-relative indexing coincide. This matches the non-spec path's tok_idx = slot.i_batch - off.
• No off-by-one. ids[i] (= accepted[i]) is the token sampled from logits row spec_i_batch[i], so populating result.tok = ids[i] from spec_i_batch[i] is the right distribution - including the diverging token on partial acceptance and the final fully-accepted token.
• Bounds are safe. spec_i_batch.size() == n_draft + 1 >= ids.size(), and the loop is bounded by ids.size(), so spec_i_batch[i] is always in range.
• The std::move(slot.spec_i_batch) at the top is equivalent to the old explicit .clear(): the moved-from vector is empty, satisfying the GGML_ASSERT(spec_i_batch.empty()) at server-context.cpp:479 on the next round. common_sampler_sample_and_accept_n receives the local copy, so the move-before-clone ordering is fine.
• First-token behavior is unchanged: the first token after prompt still goes through the non-spec path (spec_draft empty), which already populated logprobs - consistent with the PR description's claim that only tokens 2..N were missing on master.
• The partial-acceptance + checkpoint-restore (use_ckpt_tgt) path returns before the ids loop, so no probs are emitted that round; this matches master behavior (replay re-emits them later). No regression.

Will slow the review (non-blocking)

(point 1) Missing regression test. tools/server/tests/unit/test_speculative.py exists and test_completion.py already has test_n_probs / test_n_probs_stream. A test asserting top_logprobs is populated for every token under spec decoding (ideally comparing logprob values against a no-draft run, mirroring test_with_and_without_draft) would lock in this fix and follow the existing patterns. Without it, a future refactor of the spec path could silently regress this again.

(point 2) post_sampling_probs=true + spec is silently misleading. With the new guard, that combination skips populate_token_probs, leaving result.prob = 1.0 and result.probs empty, which serializes as prob: 1.0 / top_probs: [] (server-task.cpp:284). That reads as genuine 100% confidence rather than "unsupported". The PR description documents the limitation, but server users won't see the description. Consider a one-time SRV_WRN when a slot runs spec decoding with post_sampling_probs set, or a note in tools/server/README.md. This is no worse than master, just an unimproved case.

Nits

(point 3) The two added comments are concise and explain why, which is good. The one at server-context.cpp:3909 could optionally say "logprobs use pre-sampling logits only here" to make the false argument self-documenting, but it is fine as is.

(point 4) Pre-existing TODO at server-context.cpp:3784 ("set it here instead of doing inside populate_token_probs") is untouched by this PR and not its responsibility; leaving it is correct.

This review was generated automatically by pi coding agent using zai-org/GLM-5.2. It may contain mistakes. Maintainers make the final call.

[pnb]

The only thing the bot suggests that seems potentially interesting to me is a regression test. Although it is of course a tradeoff for increasing CI time, since this would probably not be a very fast test to run, so I am slightly inclined not to include it. But happy to if others think so.

[sigjhl]

I ran into the same problem while using gemma 4 e4b with MTP as a transcription model.

I did an independent fix with Fable 5, then discovered your PR. Here is the regression test the bot asked for, cherry-pickable, on this branch: https://github.com/sigjhl/llama.cpp/tree/pr27196-regression-test. The results:

# master (f1357e4), test added:
FAILED unit/test_speculative.py::test_draft_token_probs - AssertionError: assert 0 == 4
# this PR (c22562c71) + test:
7 passed in 14.30s

My independent fix has post_sampling_probs via a sampler callback and happy to open it as a follow-up after this PR merges. It will probably settle the bot's point 2.

[pnb]

@sigjhl The test seems reasonably straightforward so I cherry-picked that.

I also rebased onto master since a merge conflict had come up in the meantime. @ngxson is there anything to be done for this PR or is it possible to consider merging?