Fix logprobs

This commit is mostly a cherry-pick of ggml-org/llama.cpp#10783.

That PR and friends were partially cherry-picked by ikawrakow#723, but wasn't
really in a working state yet.

A couple of additional changes:
* Include timing information in response, which was (unintentionally?)
  done in mainline since ggml-org/llama.cpp#10643.
* Also return the actual logprobs for accepted draft tokens. This is
  still a TODO in mainline [1].

Note that there is a TG performance penalty to return the logprobs.
Here are some numbers I got with Qwen2.5-Coder-32B-Instruct:
* no draft, no logprobs: 12.81 tok/s
* no draft, with logprobs: 12.02 tok/s (6.2% drop)
* with draft, no logprobs: 36.59 tok/s
* with draft, with logprobs: 29.08 tok/s (20.5% drop)

[1] https://github.com/ggml-org/llama.cpp/blob/b6548/tools/server/server.cpp#L4019

Author: sayap

examples/server/server.cpp

@@ -556,6 +556,7 @@ struct slot_params {
     std::vector<std::string> antiprompt;
 
     bool timings_per_token = false;
+    bool post_sampling_probs = false;
     json input_prefix;
     json input_suffix;
 
@@ -1545,6 +1546,8 @@ struct server_context {
         slot.sparams.n_probs           = json_value(data, "n_probs",           default_sparams.n_probs);
         slot.sparams.min_keep          = json_value(data, "min_keep",          default_sparams.min_keep);
 
+        slot.params.post_sampling_probs = json_value(data, "post_sampling_probs", default_params.post_sampling_probs);
+
         // speculative decoding parameters
         slot.params.speculative.n_max = json_value(data, "speculative.n_max", params.n_draft);
         slot.params.speculative.n_min = json_value(data, "speculative.n_min", params.n_draft_min);
@@ -1947,26 +1950,7 @@ struct server_context {
         }
 
         // check if there is incomplete UTF-8 character at the end
-        bool incomplete = false;
-        for (unsigned i = 1; i < 5 && i <= slot.generated_text.size(); ++i) {
-            unsigned char c = slot.generated_text[slot.generated_text.size() - i];
-            if ((c & 0xC0) == 0x80) {
-                // continuation byte: 10xxxxxx
-                continue;
-            }
-            if ((c & 0xE0) == 0xC0) {
-                // 2-byte character: 110xxxxx ...
-                incomplete = i < 2;
-            } else if ((c & 0xF0) == 0xE0) {
-                // 3-byte character: 1110xxxx ...
-                incomplete = i < 3;
-            } else if ((c & 0xF8) == 0xF0) {
-                // 4-byte character: 11110xxx ...
-                incomplete = i < 4;
-            }
-            // else 1-byte character or invalid byte
-            break;
-        }
+        bool incomplete = validate_utf8(slot.generated_text) < slot.generated_text.size();
 
         if (!incomplete) {
             size_t pos = std::min(slot.n_sent_text, slot.generated_text.size());
@@ -2062,6 +2046,56 @@ struct server_context {
         return slot.has_next_token; // continue
     }
 
+    void populate_token_probs(const server_slot & slot, completion_token_output & result, bool post_sampling, bool special, int idx) {
+        size_t n_probs = slot.sparams.n_probs;
+        size_t n_vocab = llama_n_vocab(llama_get_model(ctx));
+
+        if (post_sampling) {
+            const auto * cur_p = llama_sampling_get_candidates(slot.ctx_sampling);
+            const size_t max_probs = cur_p->size;
+
+            // set probability for sampled token
+            for (size_t i = 0; i < max_probs; i++) {
+                if (cur_p->data[i].id == result.tok) {
+                    result.prob = cur_p->data[i].p;
+                    break;
+                }
+            }
+
+            // set probability for top n_probs tokens
+            result.probs.reserve(max_probs);
+            for (size_t i = 0; i < std::min(max_probs, n_probs); i++) {
+                result.probs.push_back({
+                    cur_p->data[i].id,
+                    llama_detokenize(ctx, {cur_p->data[i].id}, special),
+                    cur_p->data[i].p
+                });
+            }
+        } else {
+            // TODO: optimize this with min-p optimization
+            std::vector<llama_token_data> cur = get_token_probabilities(ctx, idx);
+
+            // set probability for sampled token
+            for (size_t i = 0; i < n_vocab; i++) {
+                // set probability for sampled token
+                if (cur[i].id == result.tok) {
+                    result.prob = cur[i].p;
+                    break;
+                }
+            }
+
+            // set probability for top n_probs tokens
+            result.probs.reserve(n_probs);
+            for (size_t i = 0; i < std::min(n_vocab, n_probs); i++) {
+                result.probs.push_back({
+                    cur[i].id,
+                    llama_detokenize(ctx, {cur[i].id}, special),
+                    cur[i].p
+                });
+            }
+        }
+    }
+
     json get_formated_generation(const server_slot & slot) const {
         const auto eos_bias   =             slot.sparams.logit_bias.find(llama_token_eos(model));
         const bool ignore_eos = eos_bias != slot.sparams.logit_bias.end() && eos_bias->second < 0.0f && std::isinf(eos_bias->second);
@@ -2159,38 +2193,19 @@ struct server_context {
         res.stop     = false;
         res.stream = slot.params.stream;
         res.content = tkn.text_to_send;
+        res.post_sampling_probs = slot.params.post_sampling_probs;
         res.oaicompat = slot.params.oaicompat;
         res.oaicompat_model = slot.params.oaicompat_model;
         res.oaicompat_cmpl_id = slot.params.oaicompat_cmpl_id;
         res.n_decoded = slot.n_decoded;
         res.n_prompt_tokens = slot.n_prompt_tokens;
-        res.data     = json {
-            {"content",    tkn.text_to_send},
-            {"stop",       false},
-            {"id_slot",    slot.id},
-            {"multimodal", false}
-        };
         slot.update_chat_msg(res.oaicompat_msg_diffs);
-        if (slot.sparams.n_probs > 0) {
-            const std::vector<llama_token> to_send_toks = llama_tokenize(ctx, tkn.text_to_send, false);
-            const size_t probs_pos = std::min(slot.n_sent_token_probs, slot.generated_token_probs.size());
-            const size_t probs_stop_pos = std::min(slot.n_sent_token_probs + to_send_toks.size(), slot.generated_token_probs.size());
-
-            std::vector<completion_token_output> probs_output;
-            if (probs_pos < probs_stop_pos) {
-                probs_output = std::vector<completion_token_output>(
-                    slot.generated_token_probs.begin() + probs_pos,
-                    slot.generated_token_probs.begin() + probs_stop_pos);
-            }
-            slot.n_sent_token_probs = probs_stop_pos;
 
-            res.data["completion_probabilities"] = probs_vector_to_json(ctx, probs_output);
+        // populate res.probs_output
+        if (slot.sparams.n_probs > 0) {
+            res.probs_output = {tkn}; // copy the token probs
         }
 
-        if (slot.oaicompat) {
-            res.data["oaicompat_token_ctr"] = slot.n_decoded;
-            res.data["model"] = slot.oaicompat_model;
-        }
         // populate timings if this is final response or timings_per_token is enabled
         if (slot.params.timings_per_token) {
             res.timings = slot.get_timings();
@@ -2207,56 +2222,30 @@ struct server_context {
         res.stop = true; // to do: set value
         res.stream = slot.params.stream;
         res.content = slot.generated_text;
+        res.timings = slot.get_timings();
+        res.post_sampling_probs = slot.params.post_sampling_probs;
         res.oaicompat = slot.params.oaicompat;
         res.oaicompat_model = slot.params.oaicompat_model;
         res.oaicompat_cmpl_id = slot.params.oaicompat_cmpl_id;
         res.oaicompat_msg = slot.update_chat_msg(res.oaicompat_msg_diffs);
         res.n_decoded = slot.n_decoded;
         res.n_prompt_tokens = slot.n_prompt_tokens;
         res.oaicompat_model = slot.oaicompat_model;
-        res.data     = json {
-            {"content",             !slot.params.stream ? slot.generated_text : ""},
-            {"generated_text",      slot.generated_text},  // Always include full text for finish_reason logic
-            {"id_slot",             slot.id},
-            {"stop",                true},
-            {"model",               params.model_alias},
-            {"tokens_predicted",    slot.n_decoded},
-            {"tokens_evaluated",    slot.n_prompt_tokens},
-            {"generation_settings", get_formated_generation(slot)},
-            {"prompt",              slot.prompt},
-            {"truncated",           slot.truncated},
-            {"stopped_eos",         slot.stopped_eos},
-            {"stopped_word",        slot.stopped_word},
-            {"stopped_limit",       slot.stopped_limit},
-            {"stopping_word",       slot.stopping_word},
-            {"tokens_cached",       slot.n_past},
-            {"timings",             slot.get_formated_timings()},
-            //{"oaicompat_chat_format",  slot.params.oaicompat_chat_format},
-        };
 
+        // populate res.probs_output
         if (slot.sparams.n_probs > 0) {
-            std::vector<completion_token_output> probs;
             if (!slot.params.stream && slot.stopped_word) {
                 const std::vector<llama_token> stop_word_toks = llama_tokenize(ctx, slot.stopping_word, false);
 
                 size_t safe_offset = std::min(slot.generated_token_probs.size(), stop_word_toks.size());
-                probs = std::vector<completion_token_output>(
+                res.probs_output = std::vector<completion_token_output>(
                         slot.generated_token_probs.begin(),
                         slot.generated_token_probs.end() - safe_offset);
             } else {
-                probs = std::vector<completion_token_output>(
+                res.probs_output = std::vector<completion_token_output>(
                         slot.generated_token_probs.begin(),
                         slot.generated_token_probs.end());
             }
-            //res.generation_params = slot.params;
-            res.data["completion_probabilities"] = probs_vector_to_json(ctx, probs);
-        }
-
-        res.timings = slot.get_timings();
-
-        if (slot.oaicompat) {
-            res.data["oaicompat_token_ctr"] = slot.n_decoded;
-            res.data["model"] = slot.oaicompat_model;
         }
 
         queue_results.send(std::move(res));
@@ -3194,7 +3183,8 @@ struct server_context {
                 }
 
                 completion_token_output result;
-                const llama_token id = llama_sampling_sample(slot.ctx_sampling, ctx, NULL, slot.i_batch - i);
+                const int tok_idx = slot.i_batch - i;
+                const llama_token id = llama_sampling_sample(slot.ctx_sampling, ctx, NULL, tok_idx);
 
                 llama_sampling_accept(slot.ctx_sampling, ctx, id, true);
 
@@ -3210,35 +3200,12 @@ struct server_context {
 
                 slot.t_token_generation = (t_current - slot.t_start_generation) / 1e3;
 
-                llama_token_data_array cur_p = { slot.ctx_sampling->cur.data(), slot.ctx_sampling->cur.size(), false };
                 result.tok = id;
+                result.prob = 1.0f; // TODO: set it here instead of doing inside populate_token_probs
                 result.text_to_send = llama_token_to_piece(ctx, result.tok, accept_special_token(slot, result.tok));
 
-                const size_t n_probs = std::min(cur_p.size, (size_t) slot.sparams.n_probs);
-                if (n_probs > 0) {
-                    const size_t n_valid = slot.ctx_sampling->n_valid;
-
-                    // Make sure at least n_probs top tokens are at the front of the vector:
-                    if (slot.sparams.temp == 0.0f && n_probs > n_valid) {
-                        llama_sample_top_k(ctx, &cur_p, n_probs, 0);
-                    }
-
-                    if (slot.sparams.temp == 0.0f) {
-                        // With greedy sampling the probabilities have possibly not been calculated.
-                        for (size_t i = 0; i < n_probs; ++i) {
-                            result.probs.push_back({
-                                cur_p.data[i].id,llama_detokenize(ctx, {cur_p.data[i].id}, params.special),
-                                i == 0 ? 1.0f : 0.0f
-                            });
-                        }
-                    } else {
-                        for (size_t i = 0; i < n_probs; ++i) {
-                            result.probs.push_back({
-                                cur_p.data[i].id, llama_detokenize(ctx, {cur_p.data[i].id}, params.special),
-                                i >= n_valid ? 0.0f : cur_p.data[i].p // Tokens filtered out due to e.g. top_k have 0 probability.
-                            });
-                        }
-                    }
+                if (slot.sparams.n_probs > 0) {
+                    populate_token_probs(slot, result, slot.params.post_sampling_probs, params.special, tok_idx);
                 }
 
                 if (!process_token(result, slot)) {
@@ -3343,7 +3310,11 @@ struct server_context {
 
                     result.tok = ids[i];
                     result.text_to_send = llama_token_to_piece(ctx, result.tok, params.special);
-                    // result.prob = 1.0f; // set later
+                    result.prob         = 1.0f; // set later
+
+                    if (slot.sparams.n_probs > 0) {
+                        populate_token_probs(slot, result, slot.params.post_sampling_probs, params.special, i);
+                    }
 
                     if (!process_token(result, slot)) {
                         // release slot because of stop condition

examples/server/utils.hpp

@@ -381,31 +381,6 @@ struct completion_token_output {
     }
 };
 
-// convert a vector of completion_token_output to json
-static json probs_vector_to_json(const llama_context * ctx, const std::vector<completion_token_output> & probs) {
-    json out = json::array();
-
-    for (const auto & prob : probs) {
-        json probs_for_token = json::array();
-
-        for (const auto & p : prob.probs) {
-            const std::string tok_str = tokens_to_output_formatted_string(ctx, p.tok);
-            probs_for_token.push_back(json {
-                {"tok_str", tok_str},
-                {"prob",    p.prob},
-            });
-        }
-
-        const std::string tok_str = tokens_to_output_formatted_string(ctx, prob.tok);
-        out.push_back(json {
-            {"content", tok_str},
-            {"probs",   probs_for_token},
-        });
-    }
-
-    return out;
-}
-
 
 //
 // OAI utils
@@ -616,13 +591,12 @@ static json oaicompat_chat_params_parse(
 
     // Handle "logprobs" field
     // TODO: The response format of this option is not yet OAI-compatible, but seems like no one really using it; We may need to fix it in the future
-    if (body.contains("logprobs")) {
+    if (json_value(body, "logprobs", false)) {
         if (has_tools && stream) {
             throw std::runtime_error("logprobs is not supported with tools + stream");
         }
         llama_params["n_probs"] = json_value(body, "top_logprobs", 20);
-    }
-    else if (body.contains("top_logprobs")) {
+    } else if (body.contains("top_logprobs") && !body.at("top_logprobs").is_null()) {
         throw std::runtime_error("top_logprobs requires logprobs to be set to true");
     }
 
@@ -715,3 +689,33 @@ static json format_error_response(const std::string & message, const enum error_
         {"type", type_str},
     };
 }
+
+static std::vector<llama_token_data> get_token_probabilities(llama_context * ctx, int idx) {
+    std::vector<llama_token_data> cur;
+    const auto * logits = llama_get_logits_ith(ctx, idx);
+    const int n_vocab = llama_n_vocab(llama_get_model(ctx));
+
+    cur.resize(n_vocab);
+    for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
+        cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f};
+    }
+
+    // sort tokens by logits
+    std::sort(cur.begin(), cur.end(), [](const llama_token_data & a, const llama_token_data & b) {
+        return a.logit > b.logit;
+    });
+
+    // apply softmax
+    float max_l = cur[0].logit;
+    float cum_sum = 0.0f;
+    for (size_t i = 0; i < cur.size(); ++i) {
+        float p = expf(cur[i].logit - max_l);
+        cur[i].p = p;
+        cum_sum += p;
+    }
+    for (size_t i = 0; i < cur.size(); ++i) {
+        cur[i].p /= cum_sum;
+    }
+
+    return cur;
+}