server: return prompt-token logprobs when echo=true (#27174)

[safiullah3915]

Overview

Fixes #27174.

This PR enables support for echo=true + logprobs=N in the /v1/completions endpoint. Previously, the server ignored the echo setting for logprobs, returning an empty list for prompt tokens and breaking downstream evaluation frameworks like lm-eval.

Additional information

Motivation

Evaluation frameworks like lm-eval grade a model by feeding it a specific correct answer as a prompt and adding up the logprobs for those prompt tokens to calculate a total score. Because prompt scores were missing, lm-eval received blanks and assumed the model was randomly guessing.

Changes

• Extract Prompt Logprobs: Piggybacks on the normal prompt decode pass to extract logprobs for prompt positions without requiring a global KV-cache clear.
• OAI JSON Formatting: Correctly populates token_logprobs and top_logprobs for the echoed prompt tokens.
• Null Edge Case: Ensures the very first prompt token explicitly returns null for its logprob data.
• Regression Tests: Added Python unit tests for both streaming and non-streaming endpoints.

Requirements

• I have read and agree with the contributing guidelines (https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md)
• AI usage disclosure: YES - I used AI to understand the code structure and data flow, but wrote and fully understand the submitted logic myself.

Co-author: @simongonzalezdc

[safiullah3915]

Hi @simongonzalezdc
I have implemented a fix for this! Could you please pull this branch, run it on your machine, and review it? I went for the full prompt token logprobs on echo: true. The server now successfully returns the exact prompt tokens alongside their math scores/probabilities

Let me know if it works as expected on your end.

[ggml-gh-bot]

Hi @safiullah3915, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[simongonzalezdc]

Validation report (gfx1151 / Strix Halo, asked in-thread)

Environment: AMD Ryzen AI Max+ 395 (Radeon 8060S, gfx1151), ROCm 7.2.4 HIP build, commit 43e3005 vs current main (2100e59) as control, identical server flags (Qwen3.8-27B UD-Q4_K_XL, q4_0 KV, 8k ctx, temp 0).

1) Control (main) — bug confirmed, and worse than reported. echo=true, logprobs=5 returns an empty list (token_logprobs_len=0, top_logprobs absent). Additionally, lm-eval (local-completions, loglikelihood tasks) doesn't just "assume random guessing" against main — it crashes with KeyError: 'token_logprobs' on arc_easy. Motivation fully validated.

2) PR build — hard abort on the feature request. The first echo=true&logprobs request on a fresh server kills the process:

/srv/external/llama.cpp-mtp-test/src/llama-context.cpp:2189: GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max) failed
#1  ggml_abort
#2  llama_context::decode (libllama)
#3  llama_decode
#4  server_context_impl::decode
#5  server_context_impl::update_slots

Minimal repro: curl http://127.0.0.1:PORT/v1/completions -d '{"prompt":"The capital of France is","max_tokens":4,"temperature":0,"echo":true,"logprobs":5}' → connection closed without response, server dead.

Likely mechanism: the prompt-logprob pass piggybacks on prompt decode, but the outputs buffer isn't sized for the extra prompt positions (n_outputs_max capacity), so the assert fires on the first eligible batch. Might need a capacity reservation when the task carries echo+logprobs (or splitting the prompt-logprob positions into their own decode).

Normal-path generation on the PR build was not reachable in our sequence (the feature request aborts first); happy to re-run the full paired matrix (including byte-identical generation anchors and the lm-eval arc_easy end-to-end) once a fix lands — the motivation here is strong and we want this upstream.

[simongonzalezdc]

Correction + refinement to my earlier report (my crash claim was config-dependent — precise trigger below; the feature itself works).

At the default context size, the PR works correctly: echo=true, logprobs=5 returns a properly populated token_logprobs (null first token as expected) and top_logprobs — 5/5 requests clean in my retest.

The crash I reported is deterministic and context-size dependent: with any -c above the default (tested -c 8192 and -c 262144, 0/5 vs 5/5 at default), the first echo+logprobs request hits the same assert:

llama-context.cpp:2189: GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max) failed

Repro: launch with -c 8192 (or any larger ctx), then the same request as before. This matters for production configs — long-context serving (we run 262k) is exactly where evaluation frameworks want prompt logprobs. Guess in line with before: the outputs capacity reservation for prompt positions likely needs to account for the larger context/batch geometry, not just defaults.

Thanks for the ping to retest — flagging my own overreach in the first report: the 'first request kills the server' framing was true only for non-default context sizes.

[safiullah3915]

Validation report (gfx1151 / Strix Halo, asked in-thread)

Environment: AMD Ryzen AI Max+ 395 (Radeon 8060S, gfx1151), ROCm 7.2.4 HIP build, commit 43e3005 vs current main (2100e59) as control, identical server flags (Qwen3.8-27B UD-Q4_K_XL, q4_0 KV, 8k ctx, temp 0).

1) Control (main) — bug confirmed, and worse than reported. echo=true, logprobs=5 returns an empty list (token_logprobs_len=0, top_logprobs absent). Additionally, lm-eval (local-completions, loglikelihood tasks) doesn't just "assume random guessing" against main — it crashes with KeyError: 'token_logprobs' on arc_easy. Motivation fully validated.

2) PR build — hard abort on the feature request. The first echo=true&logprobs request on a fresh server kills the process:

/srv/external/llama.cpp-mtp-test/src/llama-context.cpp:2189: GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max) failed
#1  ggml_abort
#2  llama_context::decode (libllama)
#3  llama_decode
#4  server_context_impl::decode
#5  server_context_impl::update_slots

Minimal repro: curl http://127.0.0.1:PORT/v1/completions -d '{"prompt":"The capital of France is","max_tokens":4,"temperature":0,"echo":true,"logprobs":5}' → connection closed without response, server dead.

Likely mechanism: the prompt-logprob pass piggybacks on prompt decode, but the outputs buffer isn't sized for the extra prompt positions (n_outputs_max capacity), so the assert fires on the first eligible batch. Might need a capacity reservation when the task carries echo+logprobs (or splitting the prompt-logprob positions into their own decode).

Normal-path generation on the PR build was not reachable in our sequence (the feature request aborts first); happy to re-run the full paired matrix (including byte-identical generation anchors and the lm-eval arc_easy end-to-end) once a fix lands — the motivation here is strong and we want this upstream.

@simongonzalezdc Thanks for detailed testing! I ran into the exact same n_outputs_max assertion error during my local tests. I temporarily bypassed it by starting the server with the --embedding flag, which forces it to allocate enough capacity.
I agree the server needs to dynamically reserve this output buffer capacity whenever echo=true and logprobs are used. Do you have any advice on the best place in the code to implement this memory reservation?

[simongonzalezdc]

Great that you reproduced the same assert — and your --embedding workaround is actually the key diagnostic: it flips server_output_limits() into the {n_batch, 1} branch (tools/server/server-context.cpp, the params.embedding || pooling != NONE case), which is exactly the output capacity this feature needs. That confirms the mechanism.

Where the mismatch comes from (traced through your diff): the prompt-logprob path marks prompt tokens as batch outputs during the main prompt decode (/* output = */ need_logits with need_prompt_logits()), so a single decode can carry O(ubatch) outputs. But the context's output capacity is reserved once at startup (params_base.n_outputs_max = output_limits.total at server-context.cpp:967), sized by common_speculative_get_output_limits() for generation+speculation only. Prompt-scale outputs exceed that reservation and hit the assert in llama_context::decode (llama-context.cpp:2189).

Two concrete fix options:

A. Reserve at startup — in server_output_limits(), raise the total when the server must support echo+logprobs:

result.total = std::max<int32_t>(result.total, params.n_batch);

This is exactly what the embedding branch already returns, so it's precedent-consistent. Cost to be aware of: the logits buffer scales as n_vocab * n_outputs_max floats (llama-context.cpp:2057) — on the order of a GB of device memory at n_batch 2048 for Qwen-class vocab, paid by every server instance. If that's a concern, gate it behind a server flag (--prompt-logprobs) or enable when --logprobs-style capacity is requested at startup.

B. Chunk within existing capacity (no reservation, no memory cost): don't mark all prompt positions as outputs in the main batch. Instead, after prompt processing, feed the echo positions back through decode in chunks bounded by the already-reserved capacity (query llama_n_outputs_max(ctx) and slice i_batch_prompt accordingly). More code in the task path, but it works on every configuration without startup-time decisions — and long-context users (we serve 262k) are exactly the ones short on spare memory.

If it's useful context from our repro: default context (4096) survived 5/5 requests while -c 8192 and -c 262144 crashed deterministically on the first request — I haven't traced why the default size passes, so treat that as observation, not mechanism.

Happy to run either patch through the same paired harness (gfx1151, 262k production config) — the setup is standing by.