From 2a53cbcc2c0d8b0c84cedc6912cfd76bc94fefda Mon Sep 17 00:00:00 2001 From: Erik LaBianca Date: Wed, 3 Jun 2026 17:29:25 -0400 Subject: [PATCH 1/2] feat(server): soft-close thinking termination (qwen35 + gemma4) Backend mechanism for soft-close thinking termination via a logit-ratio peek with split probe/inject ids, a min_tokens floor, and max_tokens treated as a response-only budget while thinking is active. - qwen35 originates the mechanism in qwen35_backend.{cpp,h}. - gemma4 ports the same mechanism in gemma4_backend.cpp (plus the probe/inject hooks in gemma4_internal.h owned here). - common/model_backend.h carries the shared hook contract. - test_server_unit.cpp adds 533 lines of coverage for both backends. - docs/specs/thinking-budget.md updated; soft-close design plan under docs/experiments/. server_main CLI flags for this mechanism ship in the thinking-control API PR (#341). Lets a model finish its reasoning gracefully rather than being hard-cut when a thinking budget is hit, which preserves answer quality on long reasoning traces. --- docs/specs/thinking-budget.md | 58 ++- server/src/common/model_backend.h | 122 ++++-- server/src/gemma4/gemma4_backend.cpp | 44 +++ server/src/qwen35/qwen35_backend.cpp | 174 +++++++-- server/src/qwen35/qwen35_backend.h | 6 +- server/test/test_server_unit.cpp | 533 +++++++++++++++++++++++++++ 6 files changed, 859 insertions(+), 78 deletions(-) diff --git a/docs/specs/thinking-budget.md b/docs/specs/thinking-budget.md index 5ebc731be..54d0979c0 100644 --- a/docs/specs/thinking-budget.md +++ b/docs/specs/thinking-budget.md @@ -538,13 +538,44 @@ The current taxonomy is: | Value | Meaning | |---|---| | `natural` | The model emitted `` on its own, either before reaching the phase-1 cap or before Level 2 had to force-close. | -| `hard` | The phase-1 cap was reached without a model-emitted ``. Either Level 2 force-closed the block in-loop (preserving KV) or Level 1 ran the phase-2 reprompt. | +| `soft` | The soft-close logit-ratio peek (Level 2.5) fired before the hard cap — `prob[] / prob[chosen_tok]` cleared the operator-configured `soft_close_min_ratio` threshold, and the AR loop injected `` while the model was already "near" closing. Indicates voluntary cooperation: the model would have closed soon anyway; we just hurried it along to reclaim tokens. Currently Qwen3.5/3.6 only. | +| `hard` | The phase-1 cap was reached without a model-emitted `` and without the soft path triggering. Either Level 2 force-closed the block in-loop (preserving KV) or Level 1 ran the phase-2 reprompt. | + +When both `soft` and `hard` could fire on the same AR step (the +soft threshold cleared at exactly the budget-edge step), `soft` +wins — the soft trigger carries more information (the model agreed +it was time) than the hard trigger (which only reports coercion). +See `docs/experiments/soft-close-thinking-termination-plan.md` §4 + +§12 for the design rationale. + +Soft-close is enabled by the operator via the CLI flag +`--think-soft-close-min-ratio `. Default `0.0` keeps the legacy +two-value taxonomy (`natural` / `hard`); any positive value +activates the third. The dial is a probability ratio in `[0, 1]`: + +| `min_ratio` | Behaviour | +|---|---| +| `0.0` | Disabled. Soft path inert; per-request overrides silently ignored. | +| `0.05`–`0.2` | Conservative — fires only when `` is within 5×–20× of the argmax probability. Recommended starting range. | +| `0.5` | Aggressive — fires when `` has at least half the probability of the chosen token. | +| `1.0` | Strict — fires only when `` IS the most-likely token. Useful as a safety check. | + +Per-request override (Anthropic envelope, see §4.1): + +```jsonc +{ + "thinking": { + "type": "enabled", + "soft_close_min_ratio": 0.1 + } +} +``` -A third value `soft` is reserved for a future voluntary-close -mechanism (logit-biasing the model toward `` as the cap -approaches, before forcing it). Reserved so consumers can switch on -the value without an exhaustive-match warning when a future server -version adds it; not emitted today. +The per-request value clamps to `min(requested, server_default)` — +clients can tighten (lower the threshold, fire more aggressively) +but not loosen (raise it above the operator's ceiling). When the +server has the dial disabled (`0.0`), per-request overrides are +silently ignored — the feature is operator-policy gated. ## 8. Streaming @@ -564,9 +595,18 @@ in the terminal `message_delta` event for Anthropic. server-configured ceiling, never looser. Allowing full override would re-create the silent-truncation footgun of middleboxes that drop unknown fields. -- **Soft close-kind / soft-budget hint.** The mechanism (logit bias - to nudge `` selection before the hard cap) is sketched in - §7 but not specified. +- **Spec-decode soft-close peek.** Soft-close fires inside the AR + loop. When spec-decode is in use, the close still triggers at the + spec-decode → AR tail-off boundary (slightly later than pure-AR + mode); the verify/accept inner loop does not run the comparator. + Gemma 4 and Laguna are pure-AR; this only matters for Qwen3.5/3.6 + with a draft model. +- **Multi-token close joint probability.** When `` tokenizes + to multiple ids, the soft-close comparator peeks only the FIRST + id's logit (the existing multi-token inject machinery drives the + remainder of the sequence on subsequent steps). The joint + `P(t_0, t_1, …)` peek is left to a v2 if false-positive rates + warrant it. - **Per-token close-info metadata.** The upstream reference exposes `(token_index, remaining_budget, rank)` for the close event. The current `finish_details` reports aggregate counts only. diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 139159e2b..36443f7f7 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -10,6 +10,7 @@ #pragma once +#include #include #include #include @@ -60,39 +61,95 @@ struct DaemonIO { // ─── Generate request/result ──────────────────────────────────────────── -// Thinking-budget force-close hook. Mirrors antirez/ds4 ds4_eval.c's -// hard_limit_reply_budget semantics: when the budget remaining (n_gen -// minus tokens committed so far) falls to hard_limit_remaining, the -// next sampled tokens get overridden with close_token_ids in order, -// giving the model the remaining budget to write a visible answer -// after the injected close-tag sequence. -// -// Single vs multi-token close: -// Qwen3.6: is one added_token (id 248069). close_token_ids -// has size 1. One override + budget_close_injected=true. -// DeepSeek/laguna: tokenizes to 3 ordinary tokens -// ([1718, 37947, 32] for DS-V3). close_token_ids has -// size 3. Three consecutive overrides, then resume. -// -// This is "Level 2" of our thinking-budget migration: in-process -// mid-stream force-close, KV-continuous. Beats Level 1's phase-2 -// reprompt because the model never sees a fresh prefill — its KV -// state continues naturally after the injected close. -// -// Current implementation: AR-decode only. When budget_hook is set, -// backends MAY route generation through their AR path (skipping spec -// decode) — the perf trade-off is acceptable since this only kicks in -// for thinking-enabled requests. Spec-decode integration is a follow-up. +// Thinking-budget force-close hook; see docs/specs/thinking-budget.md. +// When (n_gen - committed) == hard_limit_remaining, overrides sampled +// tokens with close_token_ids (AR path only). Empty = disabled. struct BudgetHook { - // Multi-token close sequence injected when `(n_gen - committed)` - // drops to `hard_limit_remaining`. For Qwen3.x this is the - // canonical "Considering the limited time..." summarize-and-stop - // lead-in (tokenized at server startup); for non-qwen arches it's - // a single close-tag token. Empty = hook disabled. + // Inject sequence written when the hard cap fires OR when soft-close + // fires. This is the verbatim tokenization of the model card's + // `thinking_terminator_hint` (e.g. for Qwen3.6 the lead-in + // "Considering the limited time by the user, ... \n\n"). + // May be many tokens long; the first element is what the AR loop + // writes on the firing step, with the rest streamed out on + // subsequent steps. Empty = disabled. std::vector close_token_ids; + // Short PROBE sequence used by the soft-close logit-ratio peek. + // Conceptually this is the tokenization of just the close MARKER + // (e.g. `` — a single token id 248069 on Qwen3.6) rather + // than the full inject directive above. Splitting probe-vs-inject + // matters because the inject sequence for trained-hint models + // starts with a content token like "Considering" whose logit is + // 19-35 nats below the chosen token at every step, masking the + // close-marker's true probability and preventing soft-close from + // ever firing. + // When empty, the soft-close peek falls back to + // `close_token_ids.front()` (legacy behavior — kept so models that + // haven't been updated keep working identically to before the split). + std::vector soft_close_probe_ids; int hard_limit_remaining = 0; + // Soft-close (Level 2 voluntary). When > 0, at each AR step the + // loop compares the probe-token logit against the chosen-token + // logit; if `prob[probe[0]] / prob[chosen] >= soft_close_min_ratio` + // (equivalently `logit[probe[0]] - logit[chosen] >= log(min_ratio)`), + // the inject sequence (close_token_ids) is written BEFORE the hard + // limit is reached. 0.0 = disabled (default); 1.0 = fire only when + // the probe token is already the most-likely token; lower values = + // fire more aggressively. See docs/specs/thinking-budget.md §7 and + // docs/experiments/soft-close-thinking-termination-plan.md. + float soft_close_min_ratio = 0.0f; + // Minimum thinking tokens before soft-close is allowed to fire. + // Soft-close peek runs on every AR step but the fire decision is + // gated by this floor — protects against premature termination on + // prompts where the close-marker logit briefly spikes mid-thought. + // 0 = floor disabled (default). Per empirical trajectory data on + // qwen3.6-27b (5 diverse prompts), only becomes + // argmax-competitive at 66-94% of natural reasoning length — so a + // floor in the 64-256 range is the typical operating point. + int soft_close_min_tokens = 0; + // Diagnostic: when true, emit one stderr line per AR step inside the + // thinking phase with (committed, chosen_tok, logit[probe0], + // logit[chosen], diff). Used to record the close-vs-chosen logit + // trajectory across a full thinking run so a sliding-threshold curve + // can be designed from empirical data rather than guessed. Zero cost + // when off. See server_main.cpp --debug-thinking-logits. + bool debug_thinking_logits = false; + + // Probe token id used by the soft-close peek. Returns the first + // element of soft_close_probe_ids when set, otherwise falls back to + // close_token_ids.front() (legacy behavior). Callers must guard + // against an empty hook before calling this. + int32_t soft_close_probe_token() const { + if (!soft_close_probe_ids.empty()) return soft_close_probe_ids.front(); + return close_token_ids.front(); + } }; +namespace soft_close { + +// Returns true when the soft-close comparator would fire on this AR +// step. Side-effect free; safe to call from unit tests. +// +// Fast path: returns false in O(1) when min_ratio <= 0 (the disabled +// default). When the model has already chosen the close token on its +// own, also returns false — the natural-close path handles that. +// +// Math: `prob[i]/prob[j] = exp(logit[i] - logit[j])`, so +// `prob[close]/prob[chosen] >= min_ratio` ⟺ +// `logit[close] - logit[chosen] >= log(min_ratio)`. We compare on +// logits to avoid `exp()` and full-softmax cost; this is numerically +// stable in fp32 for typical LLM logit ranges (~±20). +inline bool should_fire(const float * logits, + int32_t chosen_tok, + int32_t close0_tok, + float min_ratio) { + if (min_ratio <= 0.0f) return false; + if (chosen_tok == close0_tok) return false; + const float log_ratio = std::log(min_ratio); + return (logits[close0_tok] - logits[chosen_tok]) >= log_ratio; +} + +} // namespace soft_close + struct GenerateRequest { std::vector prompt; int n_gen = 0; @@ -140,6 +197,13 @@ struct GenerateResult { // stream and grepping for "" cannot distinguish the two // (the injected close decodes identically). bool budget_forced_close = false; + // True when the soft-close path (logit-ratio peek) injected the + // close sequence in this generation. Mutually exclusive + // with budget_forced_close: when both could fire on the same step, + // soft wins and budget_forced_close stays false. The server uses + // this to attribute close_kind="soft" (vs "hard"). See + // docs/specs/thinking-budget.md §7. + bool soft_forced_close = false; // True iff the AR decode loop's post-close watchdog detected an n-gram // repetition loop and broke out early. Caller surfaces this so clients // can mark the answer as unreliable rather than treating the @@ -239,6 +303,8 @@ struct ModelBackend { retry.spec_decode_ran = first.spec_decode_ran || retry.spec_decode_ran; retry.budget_forced_close = first.budget_forced_close || retry.budget_forced_close; + retry.soft_forced_close = + first.soft_forced_close || retry.soft_forced_close; retry.degenerate_decode_close = first.degenerate_decode_close || retry.degenerate_decode_close; return retry; diff --git a/server/src/gemma4/gemma4_backend.cpp b/server/src/gemma4/gemma4_backend.cpp index 9e7f131a4..103958343 100644 --- a/server/src/gemma4/gemma4_backend.cpp +++ b/server/src/gemma4/gemma4_backend.cpp @@ -411,6 +411,49 @@ bool Gemma4Backend::do_decode(int committed, int n_gen, } }; + // Soft-close: fires when close-token logit is within soft_close_min_ratio + // of the chosen-token logit, injecting the close sequence early rather + // than waiting for the hard think_max boundary. Mirrors qwen35's + // maybe_soft_close lambda — same BudgetHook fields, same soft_close:: + // should_fire predicate from model_backend.h. + auto maybe_soft_close = [&](int32_t & tok, + const float * logits_row, + int committed_now) { + if (budget_close_started) return; + if (budget_hook.close_token_ids.empty()) return; + if (budget_hook.debug_thinking_logits) { + const int32_t close0 = budget_hook.close_token_ids.front(); + const int generated = committed_now - committed_at_entry; + const float diff = logits_row[close0] - logits_row[tok]; + const float ratio = (diff > 50.0f) ? std::exp(50.0f) : std::exp(diff); + std::fprintf(stderr, + "[soft-trace] gemma4 step=%d committed=%d chosen=%d close0=%d " + "logit_close=%.4f logit_chosen=%.4f diff=%.4f prob_ratio=%.6g\n", + generated, committed_now, tok, close0, + logits_row[close0], logits_row[tok], diff, ratio); + } + if (budget_hook.soft_close_min_ratio <= 0.0f) return; + const int32_t close0 = budget_hook.close_token_ids.front(); + if (!soft_close::should_fire(logits_row, tok, close0, + budget_hook.soft_close_min_ratio)) return; + const int generated = committed_now - committed_at_entry; + const int remaining = n_gen - generated; + std::fprintf(stderr, + "[budget-hook] gemma4 soft-close at committed=%d/%d " + "(remaining=%d, min_ratio=%.4f, logit[close0]=%.3f " + "logit[chosen]=%.3f diff=%.3f): overriding token %d with " + "close[0]=%d (seq len %zu)\n", + committed_now, n_gen, remaining, + budget_hook.soft_close_min_ratio, + logits_row[close0], logits_row[tok], + logits_row[close0] - logits_row[tok], + tok, close0, budget_hook.close_token_ids.size()); + tok = close0; + budget_close_started = true; + close_inject_pos = 1; + if (forced_close_out) *forced_close_out = true; + }; + for (int i = 0; i < n_gen; ++i) { // Seed for this iteration's embed step: // - Normal case: previous iteration just pushed a sampled @@ -453,6 +496,7 @@ bool Gemma4Backend::do_decode(int committed, int n_gen, if (logits[j] > best) { best = logits[j]; next = j; } } } + maybe_soft_close(next, logits.data(), committed); maybe_force_close(next, committed); out_tokens.push_back(next); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 639f10245..6c94f704e 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -766,21 +766,22 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, auto t_prefill_end = std::chrono::steady_clock::now(); result.prefill_s = std::chrono::duration(t_prefill_end - t_prefill_start).count(); - // Decode (speculative) + // Decode (speculative or AR) if (req.n_gen > 0) { auto t_decode_start = std::chrono::steady_clock::now(); - // Pass the budget hook into spec-decode. When token count nears - // the budget edge, do_spec_decode breaks out and tails off via - // AR with the hook still active — force-close fires correctly - // without sacrificing spec-decode throughput for the bulk of - // generation. Most requests never hit the tail because the - // model closes naturally well before the budget edge. + // AR decode fires when the upper layer retried after an empty + // spec-decode (force_ar_decode, PR #314 from origin/main). + // Otherwise spec decode. The budget hook flows into both paths; + // do_spec_decode breaks out and tails off via AR with the hook + // still active so force-close fires correctly without sacrificing + // spec-decode throughput for the bulk of generation. bool decode_ok = false; if (req.force_ar_decode) { decode_ok = do_ar_decode(committed, req.n_gen, result.tokens, out_io, req.budget_hook, &result.budget_forced_close, - &result.degenerate_decode_close); + &result.degenerate_decode_close, + &result.soft_forced_close); out_io.emit(-1); } else { decode_ok = do_spec_decode(committed, req.n_gen, result.tokens, out_io, @@ -791,7 +792,8 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, req.stall_skip_tokens, &req.budget_hook, &result.budget_forced_close, - &result.degenerate_decode_close); + &result.degenerate_decode_close, + &result.soft_forced_close); if (decode_ok) { result.empty_visible_output = qwen35_empty_visible_output(result.tokens, w_); @@ -907,7 +909,8 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, decode_ok = do_ar_decode(committed, req.n_gen, result.tokens, out_io, req.budget_hook, &result.budget_forced_close, - &result.degenerate_decode_close); + &result.degenerate_decode_close, + &result.soft_forced_close); out_io.emit(-1); } else { decode_ok = do_spec_decode(committed, req.n_gen, result.tokens, out_io, @@ -918,7 +921,8 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, req.stall_skip_tokens, &req.budget_hook, &result.budget_forced_close, - &result.degenerate_decode_close); + &result.degenerate_decode_close, + &result.soft_forced_close); if (decode_ok) { result.empty_visible_output = qwen35_empty_visible_output(result.tokens, w_); @@ -1301,27 +1305,14 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, const DaemonIO & io, const BudgetHook & budget_hook, bool * forced_close_out, - bool * degenerate_close_out) { - // Budget hook state. - // - budget_close_started: true once we've begun injecting the close - // sequence. Prevents re-triggering on continued forward generation. - // - close_inject_pos: index into budget_hook.close_token_ids for the - // NEXT token to inject. While < close_token_ids.size(), each - // iteration overrides the sampled token with the corresponding - // close-sequence token (single-token close = 1 override and done; - // multi-token close like DeepSeek/laguna [1718,37947,32] = 3 - // consecutive overrides). Once equal to close_token_ids.size(), - // normal sampling resumes (model writes visible answer). + bool * degenerate_close_out, + bool * soft_forced_close_out) { + // budget_close_started: prevents re-triggering; close_inject_pos: next + // token index to inject from close_token_ids. See docs/specs/thinking-budget.md. bool budget_close_started = false; int close_inject_pos = 0; - // Capture entry KV position so the budget check is in the - // "generated since entry" frame, not the absolute KV frame. - // n_gen is the gen-only count (or the remaining-budget remap done by - // spec-decode tail-off); subtracting committed_now (absolute KV = - // prompt_len + tokens generated this call) directly would treat - // prompt-length tokens as if they were generated output, firing - // force-close prompt_len tokens early on prompted requests and - // potentially going negative after spec-decode tail-off. + // committed_at_entry: anchors budget check to "generated since entry" frame, + // not absolute KV (avoids firing prompt_len tokens early). const int committed_at_entry = committed; auto maybe_force_close = [&](int32_t & tok, int committed_now) { if (budget_hook.close_token_ids.empty()) return; @@ -1383,6 +1374,79 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, if (forced_close_out) *forced_close_out = true; } }; + + // Soft-close (logit-ratio peek). Fires BEFORE the hard-cap check so a + // soft trigger on the same step as a hard trigger is reported as + // close_kind="soft" (the more informative signal — the model agreed it + // was time to close, even if the budget was also about to run out). + // Once this lambda starts the close sequence, the maybe_force_close + // continuation branch handles steps 2..N of a multi-token close. + // Zero-cost-when-disabled invariant: when soft_close_min_ratio == 0 + // the outer guard short-circuits and we do not even read logits_buf. + // See docs/experiments/soft-close-thinking-termination-plan.md §3. + auto maybe_soft_close = [&](int32_t & tok, + const float * logits_row, + int committed_now) { + if (budget_close_started) return; // sequence already in progress + if (budget_hook.close_token_ids.empty()) return; // hook disabled + + // PROBE vs INJECT split: + // - probe0 is the token id we PEEK to decide whether to fire + // (the short close marker, e.g. `` = 248069 on Qwen3.6). + // - inject0 / inject sequence is what we WRITE when it fires + // (the full trained-hint directive). + // Fall back to close_token_ids.front() when no separate probe is + // configured (legacy / single-token-marker models). See + // BudgetHook::soft_close_probe_token(). + const int32_t probe0 = budget_hook.soft_close_probe_token(); + const int32_t inject0 = budget_hook.close_token_ids.front(); + + // Diagnostic trajectory log. Fires every AR step (gated on the + // operator flag) regardless of soft_close_min_ratio, so we can + // record close-vs-chosen logit curves even when the dial is off. + // close0 reports the PROBE token id (what the comparator uses). + if (budget_hook.debug_thinking_logits) { + const int generated = committed_now - committed_at_entry; + const float diff = logits_row[probe0] - logits_row[tok]; + const float ratio = (diff > 50.0f) ? std::exp(50.0f) : std::exp(diff); + std::fprintf(stderr, + "[soft-trace] step=%d committed=%d chosen=%d close0=%d " + "logit_close=%.4f logit_chosen=%.4f diff=%.4f prob_ratio=%.6g\n", + generated, committed_now, tok, probe0, + logits_row[probe0], logits_row[tok], diff, ratio); + } + + if (budget_hook.soft_close_min_ratio <= 0.0f) return; // dial disabled + + // Minimum-thinking-tokens floor: false-positive guard. When set, + // suppress fire until the segment has committed at least this + // many tokens. 0 = floor disabled (default). + const int generated_so_far = committed_now - committed_at_entry; + if (generated_so_far < budget_hook.soft_close_min_tokens) return; + + if (!soft_close::should_fire(logits_row, tok, probe0, + budget_hook.soft_close_min_ratio)) { + return; + } + const int generated = committed_now - committed_at_entry; + const int remaining = n_gen - generated; + std::fprintf(stderr, + "[budget-hook] soft-close at committed=%d/%d (remaining=%d, " + "min_ratio=%.4f, logit[probe0=%d]=%.3f logit[chosen]=%.3f " + "diff=%.3f log_ratio=%.3f): overriding sampled token %d with " + "inject[0]=%d (inject seq len %zu)\n", + committed_now, n_gen, remaining, + budget_hook.soft_close_min_ratio, + probe0, logits_row[probe0], logits_row[tok], + logits_row[probe0] - logits_row[tok], + std::log(budget_hook.soft_close_min_ratio), + tok, inject0, budget_hook.close_token_ids.size()); + tok = inject0; + budget_close_started = true; + close_inject_pos = 1; + if (soft_forced_close_out) *soft_forced_close_out = true; + }; + if (n_gen <= 0) return true; auto t_dec0_ar = std::chrono::steady_clock::now(); @@ -1416,12 +1480,33 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, const int initial_emitted = out_tokens.empty() ? 1 : 0; if (initial_emitted == 1) { int32_t first_tok; - if (sampler_.needs_logit_processing()) { - if (!prefill_last_logits_valid_) return false; - ggml_backend_tensor_get(sg_.logits, logits_buf.data(), prefill_last_logits_offset_, - sizeof(float) * vocab); - first_tok = sample_logits(logits_buf.data(), vocab, sampler_, - out_tokens, sampler_rng_); + // Soft-close needs the logits row for the comparator; greedy + // (argmax-only) path normally skips the logits read. Pull the + // prefill's last logits row to CPU when soft is enabled so the + // first AR step participates in the comparator. Zero-cost when + // disabled: only fetched when soft_close_min_ratio > 0. + const bool need_logits = + sampler_.needs_logit_processing() || + budget_hook.soft_close_min_ratio > 0.0f || + budget_hook.debug_thinking_logits; + if (need_logits) { + if (!prefill_last_logits_valid_) { + if (sampler_.needs_logit_processing()) return false; + // Soft-close wanted logits but prefill didn't keep them. + // Skip soft check on this single token rather than error. + first_tok = cache_.last_tok; + } else { + ggml_backend_tensor_get(sg_.logits, logits_buf.data(), + prefill_last_logits_offset_, + sizeof(float) * vocab); + if (sampler_.needs_logit_processing()) { + first_tok = sample_logits(logits_buf.data(), vocab, sampler_, + out_tokens, sampler_rng_); + } else { + first_tok = cache_.last_tok; + } + maybe_soft_close(first_tok, logits_buf.data(), committed); + } } else { first_tok = cache_.last_tok; } @@ -1531,6 +1616,13 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, } } + // Soft check runs BEFORE hard-cap check. If soft fires, it sets + // budget_close_started=true so maybe_force_close's continuation + // branch handles steps 2..N of a multi-token close (and the + // remaining-check branch is skipped because the sequence is + // already started). If soft does not fire (disabled or threshold + // not met), maybe_force_close proceeds as today. + maybe_soft_close(next_tok, logits_buf.data(), committed); maybe_force_close(next_tok, committed); out_tokens.push_back(next_tok); @@ -1656,7 +1748,8 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, const std::vector * stall_skip_tokens, const BudgetHook * budget_hook, bool * forced_close_out, - bool * degenerate_close_out) { + bool * degenerate_close_out, + bool * soft_forced_close_out) { out_accept_rate = 0.0f; out_spec_ran = false; const int hidden = w_.n_embd; @@ -1689,10 +1782,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (!can_spec) { // AR fallback consumes the final prefill position itself, then advances // one token at a time. Pass the budget hook through so force-close - // still fires when spec-decode is unavailable. + // still fires when spec-decode is unavailable. Soft-close pointer + // also forwards so close_kind="soft" can be attributed correctly + // even on the AR fallback path. bool ok = do_ar_decode(committed, n_gen, out_tokens, io, budget_hook ? *budget_hook : BudgetHook{}, - forced_close_out, degenerate_close_out); + forced_close_out, degenerate_close_out, + soft_forced_close_out); io.emit(-1); return ok; } diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 0df4df036..244304912 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -269,7 +269,8 @@ class Qwen35Backend : public ModelBackend { const std::vector * stall_skip_tokens = nullptr, const BudgetHook * budget_hook = nullptr, bool * forced_close_out = nullptr, - bool * degenerate_close_out = nullptr); + bool * degenerate_close_out = nullptr, + bool * soft_forced_close_out = nullptr); // AR decode fallback (no draft model or sampling mode). // budget_hook (when close_token_ids is non-empty) overrides the next @@ -289,7 +290,8 @@ class Qwen35Backend : public ModelBackend { const DaemonIO & io, const BudgetHook & budget_hook = {}, bool * forced_close_out = nullptr, - bool * degenerate_close_out = nullptr); + bool * degenerate_close_out = nullptr, + bool * soft_forced_close_out = nullptr); bool sync_remote_draft_features(int start_pos, int n_tokens); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 1ddbf2d1f..03c0fa7f1 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -3825,6 +3825,520 @@ static void test_flowkv_T5_inert_guard_token_count() { TEST_ASSERT(1024 >= kFkvInertMinTokens); } +// ─── Soft-close comparator + state machine ───────────────────────────── +// +// Tests the logit-ratio peek that lets the AR loop force early +// once the close-token logit comes within a configured probability ratio +// of the chosen-token logit. Default disabled. See +// docs/experiments/soft-close-thinking-termination-plan.md. +// +// We test two layers: +// 1. The pure comparator (`soft_close::should_fire`) — math-only. +// 2. A small state-machine helper that mimics the AR loop's +// precedence (soft first, then hard) so we can exercise the +// multi-token inject path and the soft/hard tie-break without a +// GPU. + +// Mirror of qwen35_backend.cpp's close-injection state for unit testing. +struct CloseState { + bool started = false; + int inject_pos = 0; + bool soft_fired = false; + bool hard_fired = false; +}; + +// Returns the (possibly overridden) token for this step, advancing +// CloseState. Mirrors the soft-then-hard ordering in the real loop. +// committed_now / committed_at_entry / n_gen track the budget arithmetic +// for the hard check identically to qwen35_backend.cpp:909-944. +static int32_t step_close_state(int32_t chosen_tok, + const float * logits, + const dflash::common::BudgetHook & hook, + int committed_now, + int committed_at_entry, + int n_gen, + CloseState & state) { + // Continue an in-progress close sequence. + if (state.started && + state.inject_pos < (int)hook.close_token_ids.size()) + { + int32_t inj = hook.close_token_ids[state.inject_pos]; + state.inject_pos++; + return inj; + } + if (state.started) return chosen_tok; // sequence already complete + + // Soft check (BEFORE hard, per plan §4). + // Probe-vs-inject split: peek hook.soft_close_probe_token(), write + // hook.close_token_ids.front() on fire. + // Min-tokens floor: suppress fire until generated_so_far >= + // hook.soft_close_min_tokens. Mirrors qwen35_backend.cpp gate. + const int generated_so_far = committed_now - committed_at_entry; + if (!hook.close_token_ids.empty() && + hook.soft_close_min_ratio > 0.0f && + generated_so_far >= hook.soft_close_min_tokens && + dflash::common::soft_close::should_fire( + logits, chosen_tok, + hook.soft_close_probe_token(), + hook.soft_close_min_ratio)) + { + state.started = true; + state.inject_pos = 1; + state.soft_fired = true; + return hook.close_token_ids.front(); // INJECT, not probe + } + + // Hard check: remaining <= hard_limit_remaining. + if (!hook.close_token_ids.empty()) { + const int generated = committed_now - committed_at_entry; + const int remaining = n_gen - generated; + if (remaining <= hook.hard_limit_remaining) { + int32_t close0 = hook.close_token_ids.front(); + if (chosen_tok == close0) { + // Model self-closed at boundary; consume as first of seq. + state.started = true; + state.inject_pos = 1; + return chosen_tok; + } + state.started = true; + state.inject_pos = 1; + state.hard_fired = true; + return close0; + } + } + return chosen_tok; +} + +// Build a logits row where the chosen-token gets logit `l_chosen` and the +// close token gets logit `l_close`; all other vocab tokens are far below. +static std::vector make_logits(int vocab, int chosen_tok, + int close_tok, + float l_chosen, float l_close) { + std::vector row(vocab, -100.0f); + row[chosen_tok] = l_chosen; + if (close_tok != chosen_tok) row[close_tok] = l_close; + return row; +} + +static void test_soft_close_disabled_default() { + // min_ratio=0 → never fires, regardless of logit configuration. + auto logits = make_logits(64, /*chosen=*/3, /*close=*/7, + /*l_chosen=*/2.0f, /*l_close=*/10.0f); + bool fired = dflash::common::soft_close::should_fire( + logits.data(), /*chosen=*/3, /*close0=*/7, /*min_ratio=*/0.0f); + TEST_ASSERT(fired == false); + // Even with close as argmax, disabled means false. + fired = dflash::common::soft_close::should_fire( + logits.data(), /*chosen=*/3, /*close0=*/3, /*min_ratio=*/0.0f); + TEST_ASSERT(fired == false); +} + +static void test_soft_close_strict_ratio_one() { + // min_ratio=1.0 → fires only when logit[close] >= logit[chosen] + // (i.e. close is the argmax or tied). chosen!=close already guarded. + auto eq = make_logits(64, 3, 7, /*l_chosen=*/5.0f, /*l_close=*/5.0f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + eq.data(), 3, 7, 1.0f) == true); + + auto below = make_logits(64, 3, 7, /*l_chosen=*/5.001f, /*l_close=*/5.0f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + below.data(), 3, 7, 1.0f) == false); + + auto above = make_logits(64, 3, 7, /*l_chosen=*/4.0f, /*l_close=*/5.0f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + above.data(), 3, 7, 1.0f) == true); +} + +static void test_soft_close_aggressive_half_prob() { + // min_ratio=0.5 — prob[close]/prob[chosen] >= 0.5 ⟺ + // logit_diff >= log(0.5) ≈ -0.6931. + const float ln_half = std::log(0.5f); + + // Boundary inclusive: diff exactly log(0.5). + auto boundary = make_logits(64, 3, 7, + /*l_chosen=*/5.0f, + /*l_close=*/5.0f + ln_half); + TEST_ASSERT(dflash::common::soft_close::should_fire( + boundary.data(), 3, 7, 0.5f) == true); + + // Just below: diff slightly less than log(0.5) (further negative). + auto below = make_logits(64, 3, 7, + /*l_chosen=*/5.0f, + /*l_close=*/5.0f + ln_half - 0.001f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + below.data(), 3, 7, 0.5f) == false); + + // Way above: close strongly favoured (but not argmax). + auto strong = make_logits(64, 3, 7, + /*l_chosen=*/5.0f, + /*l_close=*/4.9f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + strong.data(), 3, 7, 0.5f) == true); +} + +static void test_soft_close_below_threshold() { + // min_ratio=0.5, prob_ratio≈0.3 (well below) → no fire. + const float ln_03 = std::log(0.3f); + auto row = make_logits(64, 3, 7, + /*l_chosen=*/5.0f, + /*l_close=*/5.0f + ln_03); + TEST_ASSERT(dflash::common::soft_close::should_fire( + row.data(), 3, 7, 0.5f) == false); +} + +static void test_soft_close_chosen_is_close() { + // When the sampler already picks the close token, soft check never + // fires — natural-close path handles it. + auto row = make_logits(64, 7, 7, /*l_chosen=*/10.0f, /*l_close=*/10.0f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + row.data(), /*chosen=*/7, /*close0=*/7, /*min_ratio=*/0.5f) == false); + TEST_ASSERT(dflash::common::soft_close::should_fire( + row.data(), /*chosen=*/7, /*close0=*/7, /*min_ratio=*/1.0f) == false); +} + +static void test_soft_close_tiny_ratio_numerical() { + // min_ratio = 1e-6 ⇒ log_ratio ≈ -13.8155. Verify no NaN, threshold + // triggers when diff >= -13.8. + auto on = make_logits(64, 3, 7, /*l_chosen=*/5.0f, /*l_close=*/-8.5f); + auto off = make_logits(64, 3, 7, /*l_chosen=*/5.0f, /*l_close=*/-9.0f); + TEST_ASSERT(dflash::common::soft_close::should_fire( + on.data(), 3, 7, 1e-6f) == true); + TEST_ASSERT(dflash::common::soft_close::should_fire( + off.data(), 3, 7, 1e-6f) == false); +} + +// ── State-machine integration: soft + hard precedence ───────────────── + +static void test_soft_close_single_token_inject() { + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 248069 }; // Qwen3.6 single-token + hook.hard_limit_remaining = 16; + hook.soft_close_min_ratio = 0.1f; + + CloseState state; + // Step where soft should fire: close logit within 10% of chosen. + // log(0.1) ≈ -2.3026. + auto row = make_logits(/*vocab=*/250000, /*chosen=*/100, /*close=*/248069, + /*l_chosen=*/5.0f, /*l_close=*/3.0f); + int32_t out = step_close_state( + /*chosen=*/100, row.data(), + hook, + /*committed_now=*/100, /*committed_at_entry=*/50, /*n_gen=*/200, + state); + TEST_ASSERT(out == 248069); + TEST_ASSERT(state.started == true); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(state.hard_fired == false); + TEST_ASSERT(state.inject_pos == 1); + + // Next step: sequence is complete (single-token close); returns chosen. + auto row2 = make_logits(/*vocab=*/250000, /*chosen=*/200, /*close=*/248069, + /*l_chosen=*/5.0f, /*l_close=*/-50.0f); + out = step_close_state( + /*chosen=*/200, row2.data(), + hook, + /*committed_now=*/101, /*committed_at_entry=*/50, /*n_gen=*/200, + state); + TEST_ASSERT(out == 200); + TEST_ASSERT(state.soft_fired == true); // sticky +} + +static void test_soft_close_multi_token_inject() { + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 1718, 37947, 32 }; // Laguna-style multi-token + hook.hard_limit_remaining = 16; + hook.soft_close_min_ratio = 0.1f; + + CloseState state; + auto row = make_logits(/*vocab=*/250000, /*chosen=*/100, /*close=*/1718, + /*l_chosen=*/5.0f, /*l_close=*/3.0f); + int32_t out = step_close_state( + /*chosen=*/100, row.data(), + hook, + /*committed_now=*/100, /*committed_at_entry=*/50, /*n_gen=*/200, + state); + TEST_ASSERT(out == 1718); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(state.inject_pos == 1); + + // Step 2: inject 37947 regardless of chosen_tok. + auto row2 = make_logits(/*vocab=*/250000, /*chosen=*/300, /*close=*/1718, + /*l_chosen=*/5.0f, /*l_close=*/-50.0f); + out = step_close_state( + /*chosen=*/300, row2.data(), + hook, + /*committed_now=*/101, 50, 200, state); + TEST_ASSERT(out == 37947); + TEST_ASSERT(state.inject_pos == 2); + + // Step 3: inject 32. + out = step_close_state( + /*chosen=*/400, row2.data(), + hook, + /*committed_now=*/102, 50, 200, state); + TEST_ASSERT(out == 32); + TEST_ASSERT(state.inject_pos == 3); + + // Step 4: sequence complete, returns chosen. + out = step_close_state( + /*chosen=*/500, row2.data(), + hook, + /*committed_now=*/103, 50, 200, state); + TEST_ASSERT(out == 500); +} + +static void test_soft_close_then_hard_would_fire() { + // Soft fires at step 100; hard remaining-check would fire at + // committed_now=190 (remaining=10 <= hard_limit=16). Hard path + // skipped because state.started is already true. Telemetry: + // close_kind="soft". + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 248069 }; + hook.hard_limit_remaining = 16; + hook.soft_close_min_ratio = 0.1f; + + CloseState state; + // Soft trigger at committed_now=100. + auto soft_row = make_logits(/*vocab=*/250000, 100, 248069, 5.0f, 3.0f); + (void)step_close_state(100, soft_row.data(), hook, + /*committed_now=*/100, + /*committed_at_entry=*/50, + /*n_gen=*/200, state); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(state.hard_fired == false); + + // Now jump to committed_now=190 (remaining=10) — hard would have + // fired here but state.started=true so it's skipped. + auto far_row = make_logits(/*vocab=*/250000, 999, 248069, 5.0f, -100.0f); + int32_t out = step_close_state(999, far_row.data(), hook, + /*committed_now=*/190, + /*committed_at_entry=*/50, + /*n_gen=*/200, state); + // Single-token close already complete; returns chosen. + TEST_ASSERT(out == 999); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(state.hard_fired == false); +} + +static void test_soft_close_disabled_hard_still_fires() { + // min_ratio=0 (disabled): hard cap should still fire at the budget + // edge. Existing behavior preserved. + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 248069 }; + hook.hard_limit_remaining = 16; + hook.soft_close_min_ratio = 0.0f; // disabled + + CloseState state; + // Big gap between chosen and close — would fire soft if enabled. + auto row = make_logits(/*vocab=*/250000, 100, 248069, 5.0f, 4.99f); + int32_t out = step_close_state(100, row.data(), hook, + /*committed_now=*/100, 50, 200, state); + // Soft disabled: chosen passes through, not yet at hard boundary. + TEST_ASSERT(out == 100); + TEST_ASSERT(state.started == false); + + // At hard boundary: committed_now-entry=184 → remaining=16 ≤ hard. + // (entry=50, n_gen=200, hard_limit=16 ⇒ trigger at committed_now=234.) + out = step_close_state(100, row.data(), hook, + /*committed_now=*/234, 50, 200, state); + TEST_ASSERT(out == 248069); + TEST_ASSERT(state.hard_fired == true); + TEST_ASSERT(state.soft_fired == false); +} + +static void test_soft_close_natural_at_boundary() { + // Model picks close on its own (chosen == close0). Soft check skips + // (chosen==close0 guard); hard check also skips because the model + // self-emitted close. Neither flag set; close_kind would be + // "natural" downstream. + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 248069 }; + hook.hard_limit_remaining = 16; + hook.soft_close_min_ratio = 0.5f; + + CloseState state; + auto row = make_logits(/*vocab=*/250000, 248069, 248069, 5.0f, 5.0f); + // Far from hard boundary; chosen == close. + int32_t out = step_close_state(248069, row.data(), hook, + /*committed_now=*/100, 50, 200, state); + TEST_ASSERT(out == 248069); + TEST_ASSERT(state.soft_fired == false); + TEST_ASSERT(state.hard_fired == false); +} + +// Probe-vs-inject split. When soft_close_probe_ids is set, the +// comparator MUST peek the probe[0] logit, NOT inject[0]. Otherwise +// trained-hint sidecars (inject[0] = content lead-in token) keep +// the dial pinned at zero. +static void test_soft_close_probe_uses_probe_ids_not_inject_ids() { + using namespace dflash::common; + BudgetHook hook; + // Multi-token inject (mirrors a trained-hint sidecar). + hook.close_token_ids = { 99, 100, 101 }; + // Distinct single-token probe (the close marker). + hook.soft_close_probe_ids = { 42 }; + hook.soft_close_min_ratio = 0.5f; + + std::vector row(250000, -100.0f); + row[300] = 11.0f; // chosen + row[42] = 10.0f; // probe — within ratio 0.5 (exp(10-11)=0.37 < 0.5? no, 0.367) + row[42] = 10.31f; // exp(10.31-11) ≈ 0.502 — JUST fires at 0.5 + row[99] = -50.0f; // inject[0] far below — must not influence fire + + CloseState state; + int32_t out = step_close_state(/*chosen=*/300, row.data(), hook, + /*committed_now=*/200, 50, 500, state); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(out == 99); // wrote inject[0], not probe[0] + TEST_ASSERT(state.inject_pos == 1); +} + +// Empty soft_close_probe_ids ⇒ legacy fallback: peek close_token_ids +// front. Guarantees zero churn for any caller that doesn't set the +// new probe field. +static void test_soft_close_probe_ids_empty_falls_back_to_close_token_ids() { + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 248069 }; + // hook.soft_close_probe_ids left empty (legacy). + hook.soft_close_min_ratio = 0.5f; + + std::vector row(250000, -100.0f); + row[300] = 11.0f; + row[248069] = 10.31f; // close_token_ids[0]'s logit — same as before + + CloseState state; + int32_t out = step_close_state(/*chosen=*/300, row.data(), hook, + /*committed_now=*/200, 50, 500, state); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(out == 248069); + // Sanity: soft_close_probe_token() returns inject[0] when probe is empty. + TEST_ASSERT(hook.soft_close_probe_token() == 248069); +} + +// When soft-close fires, the WRITTEN sequence MUST be close_token_ids +// (the full inject), regardless of what soft_close_probe_ids contains. +// The probe is read-only — never appears in the output stream. +static void test_soft_close_inject_sequence_unchanged_when_fires() { + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 1718, 37947, 32 }; + hook.soft_close_probe_ids = { 42 }; + hook.soft_close_min_ratio = 0.1f; + + std::vector row(250000, -100.0f); + row[300] = 5.0f; + row[42] = 3.0f; // probe within ratio 0.1 + row[1718] = -80.0f; // inject[0] far below — must not matter + + CloseState state; + int32_t out = step_close_state(/*chosen=*/300, row.data(), hook, + /*committed_now=*/100, 50, 200, state); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(out == 1718); + + std::vector row2(250000, -100.0f); + row2[999] = 5.0f; + out = step_close_state(/*chosen=*/999, row2.data(), hook, + /*committed_now=*/101, 50, 200, state); + TEST_ASSERT(out == 37947); + out = step_close_state(/*chosen=*/999, row2.data(), hook, + /*committed_now=*/102, 50, 200, state); + TEST_ASSERT(out == 32); + out = step_close_state(/*chosen=*/999, row2.data(), hook, + /*committed_now=*/103, 50, 200, state); + TEST_ASSERT(out == 999); +} + +// min_thinking_tokens floor: when set, fire is suppressed until +// generated_so_far >= soft_close_min_tokens. +static void test_soft_close_min_tokens_blocks_early_fire() { + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 1718 }; + hook.soft_close_probe_ids = { 42 }; + hook.soft_close_min_ratio = 0.5f; + hook.soft_close_min_tokens = 100; + + std::vector row(250000, -100.0f); + row[300] = 5.0f; + row[42] = 5.0f; // prob_ratio = 1.0 ≫ 0.5 + + // Below floor: generated_so_far = 90 - 50 = 40 < 100 ⇒ no fire. + CloseState state_early; + int32_t out = step_close_state(/*chosen=*/300, row.data(), hook, + /*committed_now=*/90, 50, 500, + state_early); + TEST_ASSERT(state_early.soft_fired == false); + TEST_ASSERT(out == 300); + + // Above floor: generated_so_far = 200 - 50 = 150 >= 100 ⇒ fires. + CloseState state_late; + out = step_close_state(/*chosen=*/300, row.data(), hook, + /*committed_now=*/200, 50, 500, + state_late); + TEST_ASSERT(state_late.soft_fired == true); + TEST_ASSERT(out == 1718); +} + +// Default soft_close_min_tokens=0 ⇒ no floor ⇒ fire as soon as +// qualifying logits show up. Confirms the floor is opt-in. +static void test_soft_close_min_tokens_default_zero_unchanged_behavior() { + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 1718 }; + hook.soft_close_probe_ids = { 42 }; + hook.soft_close_min_ratio = 0.5f; + // soft_close_min_tokens left at default 0. + + std::vector row(250000, -100.0f); + row[300] = 5.0f; + row[42] = 5.0f; + + CloseState state; + int32_t out = step_close_state(/*chosen=*/300, row.data(), hook, + /*committed_now=*/1, 0, 500, state); + TEST_ASSERT(state.soft_fired == true); + TEST_ASSERT(out == 1718); +} + +static void test_soft_close_determinism_when_disabled() { + // Byte-identical generation invariant: with min_ratio=0, the + // override token MUST equal the chosen token for every step, for + // any logit configuration. This is the "zero-cost-when-disabled" + // generation determinism guarantee from plan §3.6. + using namespace dflash::common; + BudgetHook hook; + hook.close_token_ids = { 248069 }; + hook.hard_limit_remaining = 0; // disable hard too + hook.soft_close_min_ratio = 0.0f; // disabled + + CloseState state; + std::mt19937 rng(12345); + for (int step = 0; step < 100; step++) { + int32_t chosen = (int32_t)(rng() % 1000); + float l_chosen = (float)(rng() % 100) / 10.0f - 5.0f; + float l_close = (float)(rng() % 100) / 10.0f - 5.0f; + // vocab=250000 covers close_tok=248069. Pre-existing OOB on the + // 1000-element row was silently passing in Release builds; new + // tests perturbing heap layout could turn it into a crash. + auto row = make_logits(/*vocab=*/250000, chosen, /*close=*/248069, + l_chosen, l_close); + int32_t out = step_close_state(chosen, row.data(), hook, + /*committed_now=*/step, 0, 200, + state); + TEST_ASSERT(out == chosen); + } + TEST_ASSERT(state.soft_fired == false); + TEST_ASSERT(state.hard_fired == false); +} + + int main() { std::fprintf(stderr, "══════════════════════════════════════════\n"); std::fprintf(stderr, " Server Unit Tests\n"); @@ -4076,6 +4590,25 @@ int main() { RUN_TEST(test_flowkv_T1_system_end_boundary_first); RUN_TEST(test_flowkv_T5_inert_guard_token_count); + std::fprintf(stderr, "\n── Soft-close comparator + state machine ──\n"); + RUN_TEST(test_soft_close_disabled_default); + RUN_TEST(test_soft_close_strict_ratio_one); + RUN_TEST(test_soft_close_aggressive_half_prob); + RUN_TEST(test_soft_close_below_threshold); + RUN_TEST(test_soft_close_chosen_is_close); + RUN_TEST(test_soft_close_tiny_ratio_numerical); + RUN_TEST(test_soft_close_single_token_inject); + RUN_TEST(test_soft_close_multi_token_inject); + RUN_TEST(test_soft_close_then_hard_would_fire); + RUN_TEST(test_soft_close_disabled_hard_still_fires); + RUN_TEST(test_soft_close_natural_at_boundary); + RUN_TEST(test_soft_close_probe_uses_probe_ids_not_inject_ids); + RUN_TEST(test_soft_close_probe_ids_empty_falls_back_to_close_token_ids); + RUN_TEST(test_soft_close_inject_sequence_unchanged_when_fires); + RUN_TEST(test_soft_close_min_tokens_blocks_early_fire); + RUN_TEST(test_soft_close_min_tokens_default_zero_unchanged_behavior); + RUN_TEST(test_soft_close_determinism_when_disabled); + std::fprintf(stderr, "\n══════════════════════════════════════════\n"); std::fprintf(stderr, " Results: %d assertions, %d failures\n", test_count, test_failures); From 3c81645f72c1d4f12891523959e1cd54e3ae12c0 Mon Sep 17 00:00:00 2001 From: Erik LaBianca Date: Thu, 4 Jun 2026 19:23:10 -0400 Subject: [PATCH 2/2] fix(server): address cubic PR #339 review (close_inject_pos, probe, GPU-argmax logits) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three P1 fixes from cubic AI review on PR #339: 1) qwen35 + gemma4: soft-close set close_inject_pos=1 after firing, which caused maybe_force_close's same-step continuation branch to overwrite tok=inject[0] with close_token_ids[1] — skipping the first close token in multi-token close sequences. Set it to 0 so the continuation reads close_token_ids[0] (no-op rewrite) and advances to 1 for the next step. 2) gemma4: maybe_soft_close used close_token_ids.front() (the INJECT token) for the logit comparator instead of soft_close_probe_token() (the PROBE marker), and skipped the soft_close_min_tokens floor. Mirror qwen35: peek the probe, write the inject, gate on the floor. 3) qwen35 GPU-argmax path: maybe_soft_close was called with stale logits because the GPU-argmax branch only reads the 4-byte argmax token id and leaves logits_buf untouched. Fetch logits from GPU when soft-close (or its diagnostic) is enabled. Zero-cost when off. --- server/src/gemma4/gemma4_backend.cpp | 45 ++++++++++++++++++++-------- server/src/qwen35/qwen35_backend.cpp | 26 +++++++++++++++- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/server/src/gemma4/gemma4_backend.cpp b/server/src/gemma4/gemma4_backend.cpp index 103958343..455faabd3 100644 --- a/server/src/gemma4/gemma4_backend.cpp +++ b/server/src/gemma4/gemma4_backend.cpp @@ -416,41 +416,60 @@ bool Gemma4Backend::do_decode(int committed, int n_gen, // than waiting for the hard think_max boundary. Mirrors qwen35's // maybe_soft_close lambda — same BudgetHook fields, same soft_close:: // should_fire predicate from model_backend.h. + // + // PROBE vs INJECT split (matches qwen35; see model_backend.h): + // - probe0 is the token id we PEEK for the fire decision (the short + // close marker token, configured via soft_close_probe_ids). + // - inject0 is the first token of the close sequence we WRITE when + // the comparator fires (close_token_ids.front()). + // When soft_close_probe_ids is empty, soft_close_probe_token() falls + // back to close_token_ids.front() (legacy single-marker behavior). auto maybe_soft_close = [&](int32_t & tok, const float * logits_row, int committed_now) { if (budget_close_started) return; if (budget_hook.close_token_ids.empty()) return; + const int32_t probe0 = budget_hook.soft_close_probe_token(); + const int32_t inject0 = budget_hook.close_token_ids.front(); if (budget_hook.debug_thinking_logits) { - const int32_t close0 = budget_hook.close_token_ids.front(); const int generated = committed_now - committed_at_entry; - const float diff = logits_row[close0] - logits_row[tok]; + const float diff = logits_row[probe0] - logits_row[tok]; const float ratio = (diff > 50.0f) ? std::exp(50.0f) : std::exp(diff); std::fprintf(stderr, "[soft-trace] gemma4 step=%d committed=%d chosen=%d close0=%d " "logit_close=%.4f logit_chosen=%.4f diff=%.4f prob_ratio=%.6g\n", - generated, committed_now, tok, close0, - logits_row[close0], logits_row[tok], diff, ratio); + generated, committed_now, tok, probe0, + logits_row[probe0], logits_row[tok], diff, ratio); } if (budget_hook.soft_close_min_ratio <= 0.0f) return; - const int32_t close0 = budget_hook.close_token_ids.front(); - if (!soft_close::should_fire(logits_row, tok, close0, + // Minimum-thinking-tokens floor: false-positive guard. When set, + // suppress fire until the segment has committed at least this + // many tokens. 0 = floor disabled (default). Mirrors qwen35. + const int generated_so_far = committed_now - committed_at_entry; + if (generated_so_far < budget_hook.soft_close_min_tokens) return; + if (!soft_close::should_fire(logits_row, tok, probe0, budget_hook.soft_close_min_ratio)) return; const int generated = committed_now - committed_at_entry; const int remaining = n_gen - generated; std::fprintf(stderr, "[budget-hook] gemma4 soft-close at committed=%d/%d " - "(remaining=%d, min_ratio=%.4f, logit[close0]=%.3f " + "(remaining=%d, min_ratio=%.4f, logit[probe0=%d]=%.3f " "logit[chosen]=%.3f diff=%.3f): overriding token %d with " - "close[0]=%d (seq len %zu)\n", + "inject[0]=%d (inject seq len %zu)\n", committed_now, n_gen, remaining, budget_hook.soft_close_min_ratio, - logits_row[close0], logits_row[tok], - logits_row[close0] - logits_row[tok], - tok, close0, budget_hook.close_token_ids.size()); - tok = close0; + probe0, logits_row[probe0], logits_row[tok], + logits_row[probe0] - logits_row[tok], + tok, inject0, budget_hook.close_token_ids.size()); + tok = inject0; budget_close_started = true; - close_inject_pos = 1; + // close_inject_pos = 0 (NOT 1) so maybe_force_close's continuation + // in the same AR step picks up close_token_ids[0] and emits it + // as-is (then advances to 1 for the next step). Setting it to 1 + // here would cause the continuation to overwrite tok with + // close_token_ids[1], skipping the first close token. See cubic + // PR #339 comment 1 (same bug exists symmetrically in qwen35). + close_inject_pos = 0; if (forced_close_out) *forced_close_out = true; }; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 6c94f704e..21c3ead6e 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1443,7 +1443,14 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, tok, inject0, budget_hook.close_token_ids.size()); tok = inject0; budget_close_started = true; - close_inject_pos = 1; + // close_inject_pos = 0 (NOT 1) so that maybe_force_close's + // continuation branch, which runs immediately after this lambda + // in the same AR step, picks up close_token_ids[0] and emits + // inject[0] as-is (and advances close_inject_pos to 1 for the + // next step). If we set it to 1 here, the continuation would + // overwrite tok=inject[0] with close_token_ids[1], skipping + // the first close token entirely. See cubic PR #339 comment 1. + close_inject_pos = 0; if (soft_forced_close_out) *soft_forced_close_out = true; }; @@ -1622,6 +1629,23 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, // remaining-check branch is skipped because the sequence is // already started). If soft does not fire (disabled or threshold // not met), maybe_force_close proceeds as today. + // + // GPU-argmax path: when soft-close or its diagnostic is enabled, + // pull the logits row to CPU before the comparator runs — the + // GPU-argmax branch above only reads 4 bytes (the argmax token id) + // and leaves logits_buf with stale data from the previous step. + // Without this fetch, maybe_soft_close would compare against stale + // logits and could misfire / emit invalid trajectory traces. + // Zero-cost when soft-close is disabled (dial == 0 and no debug + // flag): we skip the D2H copy entirely. See cubic PR #339 comment 3. + if (!sampler_.needs_logit_processing() && + kGpuArgmaxAR && sg_.argmax_tokens && + (budget_hook.soft_close_min_ratio > 0.0f || + budget_hook.debug_thinking_logits)) + { + ggml_backend_tensor_get(sg_.logits, logits_buf.data(), 0, + sizeof(float) * vocab); + } maybe_soft_close(next_tok, logits_buf.data(), committed); maybe_force_close(next_tok, committed);