fix(qwen3.5): capture logprobs on the prefill draw when the think block is prefilled

src/common/AutoModel/automodel.cpp

@@ -232,7 +232,15 @@ void AutoModel::_shared_load_model(std::string model_path, json model_info, int
 }
 
 void AutoModel::_shared_initialize_model_state(
-    std::string model_path, json, int context_length) {
+    std::string model_path, json model_info, int context_length) {
+    // The catalog marks a model single-turn with a label, and some classes set
+    // the flag themselves. Only ever switched on here, so an entry without the
+    // label does not override a class that knows it is single-turn.
+    if (model_info.contains("label") &&
+        std::find(model_info["label"].begin(), model_info["label"].end(),
+                  "single-turn") != model_info["label"].end()) {
+        this->single_turn = true;
+    }
     this->model_path = std::move(model_path);
     header_print("FLM", "Loading model: " << this->model_path);
     this->lm_config = std::make_unique<LM_Config>();
@@ -391,7 +399,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, bool capture_logprobs) {
+    // 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);
@@ -420,7 +436,13 @@ bool AutoModel::_shared_insert(chat_meta_info_t& meta_info, std::vector<int>& to
             break;
         }
     }
-    if (skip_count != idx) {
+    // A request that matches the whole history leaves nothing to prefill, and
+    // then there are no logits to draw from: the sampler would read an empty
+    // buffer. It happens when a request repeats the previous one exactly and
+    // that answer added nothing to the history -- a single-turn Qwen3.5 model
+    // with the think block prefilled and max_tokens=1. Recurrent state cannot
+    // be stepped back by one token, so treat it like a divergence.
+    if (skip_count != idx || skip_count == tokens.size()) {
         if (is_server_mode) {
             header_print("FLM", "Conversation context diverged from cache, clearing context...");
         }
@@ -459,7 +481,12 @@ 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);
+    // A caller whose first token of the answer is this draw arms the capture
+    // here, after the disarming above -- a thrown-away draw never is.
+    if (capture_logprobs && record_sampled_token) {
+        this->_request_logprobs_capture(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 +564,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 +660,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

@@ -287,12 +287,32 @@ bool Qwen3_5VL::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input, s
         this->token_history = checkpoint_his; // restore the token history to be consistent with the restored KV cache, which is crucial for correct functioning of _shared_insert's prefix-matching logic
     }
 
-    size_t n = tokens.size();
-    tokens.resize(n - (this->enable_think ? 2 : 4));
+    // The template ends with the opening of the assistant turn and, without
+    // thinking, an empty think block. Normally the prefill stops before that
+    // block, so the checkpoint below sits where the next turn of the
+    // conversation will diverge, and generate() feeds the block one forward()
+    // at a time. A single-turn model never comes back to that checkpoint, so
+    // the four extra decode steps buy nothing: prefill the block with the
+    // prompt, and the token drawn from the prefill logits is already the first
+    // token of the answer.
+    this->think_block_prefilled = this->single_turn && !this->enable_think;
+    if (!this->think_block_prefilled) {
+        size_t n = tokens.size();
+        tokens.resize(n - (this->enable_think ? 2 : 4));
+    }
 
+    // Unless the think block was prefilled, the draw below is thrown away:
+    // generate() force-feeds the block and samples again. Such a draw 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.
+    // With the block prefilled the draw is the answer's first token, so it is
+    // counted, and it is the draw a logprobs request has to capture.
+    const bool real_draw = this->think_block_prefilled;
     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, real_draw, real_draw)
+        : this->_shared_insert(meta_info, tokens, is_cancelled, nullptr, 0, std::nullopt, real_draw, real_draw);
 
     checkpoint_his = token_history;
     int checkpoint_idx = this->lm_engine->checkpoint();
@@ -316,7 +336,16 @@ std::string Qwen3_5VL::generate(chat_meta_info_t& meta_info, int length_limit, s
     this->profiler_list[TKOEN_DECODE_TIME].reset();
     std::string token_str;
     int sampled_token;
-    if(this->enable_think) {
+    if (this->think_block_prefilled) {
+        // insert() prefilled the empty think block; _shared_insert drew the
+        // answer's first token from those logits, under the constraint. That
+        // draw is kept, so nothing is rolled back: the constraint started
+        // behind the closed block (the prompt tail ends with it) and has
+        // already accepted the token.
+        sampled_token = this->last_token;
+        meta_info.insert_sample_pending = false;
+    }
+    else if(this->enable_think) {
         this->token_history.push_back(think_start_id);
         this->profiler_list[DECODING_TIME].start();
         this->_force_token(meta_info, think_start_id);
@@ -335,6 +364,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 +397,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,13 +406,29 @@ 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...");
         reason = MAX_LENGTH_REACHED;
         return result;
     }
-    while (this->total_tokens < this->MAX_L){
+    while (true){
+        // The first token is drawn before the loop, so the limit can already be
+        // reached here -- with max_tokens=1 it always is. Checked only after a
+        // draw, the loop would draw and emit one token past the limit.
+        if ((length_limit > 0) && (meta_info.generated_tokens >= length_limit)){
+            reason = MAX_LENGTH_REACHED;
+            break;
+        }
+        if (this->total_tokens >= this->MAX_L){
+            break;
+        }
         if (is_cancelled()) {
             reason = CANCEL_DETECTED;
             break;
@@ -397,6 +448,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);
@@ -406,10 +458,6 @@ std::string Qwen3_5VL::generate(chat_meta_info_t& meta_info, int length_limit, s
             break;
         }
         meta_info.generated_tokens++;
-        if ((length_limit > 0) && (meta_info.generated_tokens >= length_limit)){
-            reason = MAX_LENGTH_REACHED;
-            break;
-        }
     }
     meta_info.decoding_duration = (uint64_t)(time_utils::cast_to_us(this->profiler_list[DECODING_TIME].get_total_time()).first) * 1e3;
     meta_info.stop_reason = reason;

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();