feat(server): report logprobs when the request asks for them

src/common/AutoModel/automodel.cpp

@@ -391,7 +391,15 @@ void AutoModel::_shared_validate_capacity(std::size_t rendered_tokens,
     }
 }
 
-bool AutoModel::_shared_insert(chat_meta_info_t& meta_info, std::vector<int>& tokens, std::function<bool()> is_cancelled, void* payload, int first_len_run, std::optional<int> requested_max_new_tokens) {
+bool AutoModel::_shared_insert(chat_meta_info_t& meta_info, std::vector<int>& tokens, std::function<bool()> is_cancelled, void* payload, int first_len_run, std::optional<int> requested_max_new_tokens,
+                        bool record_sampled_token) {
+    // Disarmed here, before anything that can return early: every request goes
+    // through this function, and a request that does not ask for logprobs must
+    // not inherit the arming of an earlier one. Arming happens where the
+    // recording happens, which for Qwen3.5-VL is its own generate().
+    this->sampler->capture_top_logprobs = false;
+    this->sampler->captured = false;
+
     this->_shared_validate_capacity(tokens.size(), requested_max_new_tokens);
 
     this->_begin_turn(meta_info, tokens);
@@ -459,7 +467,7 @@ bool AutoModel::_shared_insert(chat_meta_info_t& meta_info, std::vector<int>& to
         header_print("WARNING", "Max length reached, stopping prefilling...");
     }
     this->profiler_list[SAMPLING_TIME].start();
-    this->last_token = this->_sample(y, meta_info);
+    this->last_token = this->_sample(y, meta_info, record_sampled_token);
     meta_info.insert_sample_pending = true;
     this->profiler_list[SAMPLING_TIME].stop(1);
     return true;
@@ -537,12 +545,12 @@ void AutoModel::_begin_turn(chat_meta_info_t& meta_info, const std::vector<int>&
     }
 }
 
-int AutoModel::_sample(buffer<bf16>& logits, chat_meta_info_t& meta_info) {
+int AutoModel::_sample(buffer<bf16>& logits, chat_meta_info_t& meta_info, bool record_in_history) {
     decode_constraint* constraint = meta_info.constraint.get();
     const bool masked = constraint != nullptr && constraint->fill_mask(this->constraint_mask_);
     // The mask goes to the sampler, which applies it to its own copy: the logits
     // live in the engine's device-mapped buffer and are never written here.
-    int token = this->sampler->sample(logits, masked ? &this->constraint_mask_ : nullptr);
+    int token = this->sampler->sample(logits, masked ? &this->constraint_mask_ : nullptr, record_in_history);
     auto allowed = [this](size_t i) {
         return i / 32 < this->constraint_mask_.size() && ((this->constraint_mask_[i / 32] >> (i % 32)) & 1u);
     };
@@ -633,6 +641,83 @@ std::shared_ptr<decode_constraint> AutoModel::make_constraint(const chat_constra
     return chat_make_constraint(this->chat_format, effective, vocab, engine);
 }
 
+/// True when nlohmann::json can serialise this string.
+///
+/// A token can be the start of a multi-byte sequence: decoded on its own it is
+/// not valid text, and json::dump() throws on it. Asking the same library that
+/// would throw is exact, where a hand-written UTF-8 check is only an
+/// approximation of what it accepts.
+static bool is_serialisable(const std::string& s) {
+    try {
+        (void)json(s).dump();
+        return true;
+    } catch (const std::exception&) {
+        return false;
+    }
+}
+
+void AutoModel::_request_logprobs_capture(chat_meta_info_t& meta_info) {
+    // Once per request: the sampler stays armed for all of it, because every
+    // token that reaches the caller's text needs its own distribution.
+    if (!meta_info.want_logprobs || meta_info.logprobs_armed) {
+        return;
+    }
+    meta_info.logprobs_armed = true;
+    this->sampler->captured = false;
+    this->sampler->capture_count = meta_info.top_logprobs_count;
+    this->sampler->capture_top_logprobs = true;
+}
+
+void AutoModel::_record_top_logprobs(chat_meta_info_t& meta_info, int sampled_token) {
+    if (!meta_info.want_logprobs || !this->sampler || !this->sampler->captured) {
+        return;
+    }
+
+    // The sampler's logits are not touched by the sampling -- only
+    // top_k_logits is -- so the value here and the captured logsumexp come
+    // from the same moment.
+    const float lse = this->sampler->captured_logsumexp;
+    auto entry_for = [&](int id) {
+        top_logprob_t e;
+        e.token_id = id;
+        // True log-probability: logit - logsumexp over the WHOLE vocabulary,
+        // not one renormalised over the entries shown. The two differ whenever
+        // top_k is smaller than the vocabulary, which is always.
+        e.logprob = this->sampler->logits[id] - lse;
+        const std::string raw = this->tokenizer
+            ? this->tokenizer->run_time_decoder(id)
+            : std::string();
+        // Decode here, not in the REST handler: that one has no tokenizer.
+        // As unsigned: `char` is signed on x86-64, so a byte above 0x7F would
+        // land here as a negative number, and OpenAI defines `bytes` as
+        // integers 0-255. A client rebuilding the text from them would fail on
+        // every umlaut, every CJK character, every emoji.
+        for (unsigned char c : raw) {
+            e.bytes.push_back(c);
+        }
+        e.token = is_serialisable(raw) ? raw : std::string();
+        return e;
+    };
+
+    logprobs_entry_t entry;
+    if (sampled_token >= 0 && sampled_token < static_cast<int>(this->sampler->logits.size())) {
+        entry.chosen = entry_for(sampled_token);
+    }
+    const auto& src = this->sampler->captured_top_k;
+    const size_t limit = std::min<size_t>(meta_info.top_logprobs_count, src.size());
+    entry.alternatives.reserve(limit);
+    for (size_t i = 0; i < limit; i++) {
+        // A decoding constraint pushes the tokens it forbids to -inf. They
+        // were never drawable, and -inf has no JSON number; the list is sorted,
+        // so the first one ends it.
+        if (!std::isfinite(src[i].logits)) {
+            break;
+        }
+        entry.alternatives.push_back(entry_for(src[i].token_id));
+    }
+    meta_info.logprobs_entries.push_back(std::move(entry));
+}
+
 std::string AutoModel::_shared_generate(chat_meta_info_t& meta_info, int length_limit, std::ostream& os, std::function<bool()> is_cancelled) {
     std::vector<int> sampled_tokens;
     std::string result;

src/common/AutoModel/modeling_qwen3_5vl.cpp

@@ -290,9 +290,15 @@ bool Qwen3_5VL::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input, s
     size_t n = tokens.size();
     tokens.resize(n - (this->enable_think ? 2 : 4));
 
+    // The draw below is thrown away: generate() force-feeds the think block and
+    // samples again. It must not enter the repetition and frequency windows --
+    // it was never generated text, and for a model that answers without a
+    // thinking block the token predicted here IS the answer letter. Counting it
+    // pushes that letter down by count * freq_penalty + pre_penalty at the position
+    // where it matters, so the model's preferred answer arrives penalised.
     bool success = has_images
-        ? this->_shared_insert(meta_info, tokens, is_cancelled, &image_payload, last_image_token_index)
-        : this->_shared_insert(meta_info, tokens, is_cancelled, nullptr);
+        ? this->_shared_insert(meta_info, tokens, is_cancelled, &image_payload, last_image_token_index, std::nullopt, false)
+        : this->_shared_insert(meta_info, tokens, is_cancelled, nullptr, 0, std::nullopt, false);
 
     checkpoint_his = token_history;
     int checkpoint_idx = this->lm_engine->checkpoint();
@@ -335,6 +341,11 @@ std::string Qwen3_5VL::generate(chat_meta_info_t& meta_info, int length_limit, s
         token_str = this->tokenizer->run_time_decoder(198);
         result += token_str;
         // First token of the reasoning block.
+        // No capture here: this draw is the first token of the reasoning
+        // block, and that lands in reasoning_content, not in content. OpenAI's
+        // logprobs.content describes the tokens of message.content, so a
+        // report from here would name a token the caller cannot find in the
+        // text it was given. Requests with think:true get logprobs: null.
         sampled_token = this->_sample(y, meta_info);
         os << token_str << std::flush;
     }
@@ -363,6 +374,7 @@ std::string Qwen3_5VL::generate(chat_meta_info_t& meta_info, int length_limit, s
         buffer<bf16> y = this->lm_engine->forward(271);
         this->profiler_list[DECODING_TIME].stop(1);
         token_str = this->tokenizer->run_time_decoder(271);
+        this->_request_logprobs_capture(meta_info);
         sampled_token = this->_sample(y, meta_info);
     }
     this->total_tokens++;
@@ -371,6 +383,12 @@ std::string Qwen3_5VL::generate(chat_meta_info_t& meta_info, int length_limit, s
     token_str = this->tokenizer->run_time_decoder(last_sampled_token);
     result += token_str;
     os << token_str << std::flush;
+    // Recorded where the token reaches the caller's text, not where it was
+    // drawn, and under the same filter the loop below uses: a token the engine
+    // would have dropped must not get an entry either.
+    if (this->is_normal_token(last_sampled_token)) {
+        this->_record_top_logprobs(meta_info, last_sampled_token);
+    }
 
     if (this->total_tokens >= this->MAX_L){
         header_print("WARNING", "Max length reached, stopping generation...");
@@ -397,6 +415,7 @@ std::string Qwen3_5VL::generate(chat_meta_info_t& meta_info, int length_limit, s
             std::string token_str = this->tokenizer->run_time_decoder(sampled_token);
             os << token_str << std::flush;
             result += token_str;
+            this->_record_top_logprobs(meta_info, sampled_token);
         }
         this->profiler_list[TKOEN_DECODE_TIME].stop(1);
         this->token_history.push_back(sampled_token);

src/common/AutoModel/modeling_qwen3_6_moe.cpp

@@ -299,9 +299,15 @@ bool Qwen3_6_MOE::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input,
     size_t n = tokens.size();
     tokens.resize(n - (this->enable_think ? 2 : 4));
 
+    // The draw below is thrown away: generate() force-feeds the think block and
+    // samples again. It must not enter the repetition and frequency windows --
+    // it was never generated text, and for a model that answers without a
+    // thinking block the token predicted here IS the answer letter. Counting it
+    // pushes that letter down by count * freq_penalty + pre_penalty at the position
+    // where it matters, so the model's preferred answer arrives penalised.
     bool success = has_images
-        ? this->_shared_insert(meta_info, tokens, is_cancelled, &image_payload, last_image_token_index)
-        : this->_shared_insert(meta_info, tokens, is_cancelled, nullptr);
+        ? this->_shared_insert(meta_info, tokens, is_cancelled, &image_payload, last_image_token_index, std::nullopt, false)
+        : this->_shared_insert(meta_info, tokens, is_cancelled, nullptr, 0, std::nullopt, false);
 
     checkpoint_his = token_history;
     int checkpoint_idx = this->lm_engine->checkpoint();

src/common/modules/sampler.cpp

@@ -8,6 +8,9 @@
 
 #include "modules/sampler.hpp"
 
+#include <array>
+#include <cstring>
+#include <numeric>
 #include <deque>
 #include <algorithm>  // for std::sort
 #include <chrono>
@@ -354,10 +357,112 @@ void Sampler::ring_buffer_update_sparse(int sampled_index) {
     this->total_tokens++;
 }
 
+namespace {
+/// exp() for every value a bf16 can hold.
+///
+/// The logits arrive as bf16 from the engine and the widening to fp32 is
+/// exact, so this->logits can only take 65536 distinct values. Looking the
+/// exponential up is the same number std::exp would return -- precomputed,
+/// not approximated -- and it turns a quarter of a million exp() calls into as
+/// many reads. 512 KB, which fits in L2.
+const std::array<double, 65536>& bf16_exp_table() {
+    static const std::array<double, 65536> table = [] {
+        std::array<double, 65536> t{};
+        for (uint32_t b = 0; b < 65536; b++) {
+            const uint32_t bits = b << 16;
+            float f;
+            std::memcpy(&f, &bits, sizeof(f));
+            t[b] = std::exp(static_cast<double>(f));
+        }
+        return t;
+    }();
+    return table;
+}
+}  // namespace
+
+float Sampler::log_sum_exp_all() const {
+    if (this->logits.empty()) {
+        return 0.0f;
+    }
+    const auto& table = bf16_exp_table();
+    double sum = 0.0;
+    for (int i = 0; i < this->in_features; i++) {
+        uint32_t bits;
+        std::memcpy(&bits, &this->logits[i], sizeof(bits));
+        // The low 16 bits are non-zero for anything the penalties moved off
+        // the bf16 grid; those take the exact call. The branch predicts well
+        // and costs nothing measurable.
+        sum += (bits & 0xFFFFu) != 0u
+            ? std::exp(static_cast<double>(this->logits[i]))
+            : table[bits >> 16];
+    }
+    if (std::isfinite(sum) && sum > 0.0) {
+        return static_cast<float>(std::log(sum));
+    }
+    // Nothing on the bf16 grid produced a usable sum -- reachable only for a
+    // logit beyond about +-700, where exp overflows a double. Fall back to the
+    // shifted form. Note what it does NOT cover: +-inf and NaN are not
+    // meaningful here either way, and a logit of +inf would come back as 0.
+    float max_logit = -std::numeric_limits<float>::infinity();
+    for (float v : this->logits) {
+        if (v > max_logit) {
+            max_logit = v;
+        }
+    }
+    if (!std::isfinite(max_logit)) {
+        return 0.0f;
+    }
+    double shifted = 0.0;
+    for (float v : this->logits) {
+        shifted += std::exp(static_cast<double>(v) - max_logit);
+    }
+    return max_logit + static_cast<float>(std::log(shifted));
+}
+
+void Sampler::capture_if_requested() {
+    if (!this->capture_top_logprobs) {
+        return;
+    }
+    // Own top-N by logit, taken from the full logits rather than from
+    // top_k_logits: the report should not depend on the sampler's top_k, and
+    // the greedy path keeps only one entry there.
+    const int n = std::min(this->capture_count, this->in_features);
+    this->captured_top_k.clear();
+    this->captured_logsumexp = this->log_sum_exp_all();
+    if (n > 0 && static_cast<size_t>(n) <= this->top_k_logits.size()) {
+        // top_k_logits is already sorted by logit with the same tie-break, so
+        // its beginning IS the top-n. No second sort, no vocab-sized index
+        // vector -- which matters because this runs once per token.
+        this->captured_top_k.assign(this->top_k_logits.begin(),
+                                    this->top_k_logits.begin() + n);
+    } else if (n > 0) {
+        std::vector<int> idx(this->in_features);
+        std::iota(idx.begin(), idx.end(), 0);
+        // Same order as sampler_topk_apply: bf16 logits over 100k+ entries
+        // are often equal, partial_sort is not stable, and a tie broken
+        // differently would put another token at the top than the one the
+        // draw would have taken.
+        std::partial_sort(idx.begin(), idx.begin() + n, idx.end(),
+                          [this](int a, int b) {
+                              if (this->logits[a] != this->logits[b]) {
+                                  return this->logits[a] > this->logits[b];
+                              }
+                              return a < b;
+                          });
+        this->captured_top_k.reserve(n);
+        for (int i = 0; i < n; i++) {
+            this->captured_top_k.push_back(logits_t{ this->logits[idx[i]], idx[i], 0.0f });
+        }
+    }
+    // Stays armed: every token that reaches the caller's text needs its own
+    // distribution. `captured` marks that THIS draw has one.
+    this->captured = true;
+}
+
 /// \brief Sample the token
 /// \param x the input buffer
 /// \return the sampled token
-int Sampler::sample_greedy(buffer<bf16>& x, const std::vector<uint32_t>* allowed) {
+int Sampler::sample_greedy(buffer<bf16>& x, const std::vector<uint32_t>* allowed, bool record_in_history) {
     const bool penalties_active =
         (this->repeat_last_n != 0) &&
         (this->rep_penalty != 1.0f || this->freq_penalty != 0.0f || this->pre_penalty != 0.0f);
@@ -447,7 +552,35 @@ int Sampler::sample_greedy(buffer<bf16>& x, const std::vector<uint32_t>* allowed
 
     // callers that inspect the surviving candidates still get a consistent view
     this->top_k_logits.assign(1, logits_t{ best, best_id, 1.0f });
-    ring_buffer_update_sparse(best_id);
+    if (this->capture_top_logprobs) {
+        // log_sum_exp_all() reads this->logits, which the branch above only
+        // fills when penalties or a constraint mask are active -- and then
+        // penalised and masked, as the draw saw them, so it must not be
+        // overwritten. Otherwise fill it here the same way sample() does: the
+        // logits come back in device-mapped memory where element-at-a-time
+        // reads are punishing, so this pulls eight at a time.
+        if (!penalties_active && allowed == nullptr) {
+            #if USEAVX2
+            const int simd_width = 8;
+            int i = 0;
+            for (; i <= in_features - simd_width; i += simd_width) {
+                __m128i bf16_vals_x = _mm_loadu_si128((__m128i*)&x[i]);
+                _mm256_storeu_ps(&this->logits[i], bf16o_fp32(bf16_vals_x));
+            }
+            for (; i < in_features; i++) {
+                this->logits[i] = x[i];
+            }
+            #else
+            for (int i = 0; i < in_features; i++) {
+                this->logits[i] = x[i];
+            }
+            #endif
+        }
+        this->capture_if_requested();
+    }
+    if (record_in_history) {
+        ring_buffer_update_sparse(best_id);
+    }
     return best_id;
 }
 
@@ -471,10 +604,10 @@ void Sampler::sampler_allowed_apply(const std::vector<uint32_t>* allowed) {
     }
 }
 
-int Sampler::sample(buffer<bf16>& x, const std::vector<uint32_t>* allowed) {
+int Sampler::sample(buffer<bf16>& x, const std::vector<uint32_t>* allowed, bool record_in_history) {
     // PRNG is seeded once at construction, no per-call seeding needed
     if (this->top_k == 1) {
-        return sample_greedy(x, allowed);
+        return sample_greedy(x, allowed, record_in_history);
     }
 
     // COPY FROM `x` → `this->logits[]`
@@ -500,6 +633,8 @@ int Sampler::sample(buffer<bf16>& x, const std::vector<uint32_t>* allowed) {
     sampler_penalty_apply_sparse();
     sampler_allowed_apply(allowed);
     sampler_topk_apply(this->top_k);
+
+    this->capture_if_requested();
     if (this->use_optimized_sampling) {
         softmax_with_topp_minp(this->top_p, this->min_p);
         sampler_temp_apply(this->temperature);
@@ -515,7 +650,9 @@ int Sampler::sample(buffer<bf16>& x, const std::vector<uint32_t>* allowed) {
     }
 
     int sampled_index = sample_from_probs();
-    ring_buffer_update_sparse(sampled_index);
+    if (record_in_history) {
+        ring_buffer_update_sparse(sampled_index);
+    }
 
     return sampled_index;
 }