diff --git a/README.md b/README.md index 57436327eccc..3bd577c1f9c7 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,208 @@ +# llama-mindcontrol + +Mindcontrol for Llama.cpp + +A `llama.cpp` fork that extends the reasoning-budget sampler with staged, in-context budget signaling. + +## Problem +### "But, wait..." + +Reasoning models generate an unbounded `` block prior to their final answer, and the length and quality of that block are sensitive to sampling parameters. At the temperature and top-p/top-k settings needed to avoid degenerate, low-entropy output elsewhere in generation, the reasoning block is prone to failure modes that are distinct from ordinary sampling artifacts: + +- **Repetition loops** — the sampler re-enters a previously visited distribution over reasoning tokens, producing near-identical spans of text with no new information content. +- **Non-convergent revision** — the model repeatedly re-opens a conclusion it has already reached (recurring "wait, actually..." / "but hold on..." transitions) without a stopping condition ever becoming more probable than continuing. +- **Unbounded length** — absent an explicit stopping signal, nothing in the token distribution guarantees termination; the reasoning block can consume arbitrary context budget before (or without) closing. + +Overly-verbose reasoning chain + + +The standard/naive mitigation is a hard token-count cutoff enforced by the sampler: once N tokens have been generated inside the reasoning block, `` is forced regardless of position in the sequence. This bounds worst-case length but does not address any of the above — the cutoff has no dependency on the model's actual generation state and truncates at whatever token index it happens to hit, including mid-token-sequence for a partial word, mid-clause, or mid-computation. It suppresses the symptom (unbounded length) without altering the sampling behavior that produces the loop or the non-convergence in the first place. + +## Mechanism + +This fork adds a state machine around the existing hard-cutoff sampler, with two additional stages that inject fixed text into the reasoning stream at defined points: + +1. **Intro stage** — on entry to the reasoning block, a fixed message stating the token budget is inserted (templated via a `{budget}` placeholder), e.g. `"I'm allowed to think for 512 tokens, so my reasoning should be concise. Let me start by"`. This gives the model an explicit, in-context reference for its own generation length before reasoning begins. +2. **Soft-warning stage** — at a configurable fraction of the budget (default 0.5), the sampler waits for the next newline boundary and inserts a fixed message indicating the budget is half-consumed, e.g. `"I've used up half of my thinking budget, let me start working towards a conclusion"`. +3. **Hard-stop stage with grace period** — once the budget is exhausted, the sampler enters a pending state and waits up to a configurable number of grace tokens for a paragraph boundary (two consecutive newlines) before inserting a fixed closing message and terminating the reasoning block. If no paragraph boundary occurs within the grace period, the cutoff is forced immediately. Total output length remains bounded by `budget + grace_tokens` in all cases. + +Mindcontrol in Action + +Injected text is only inserted at newline or paragraph boundaries, not mid-token or mid-sentence. If the model emits its own `` before a forced stage would trigger, the natural close takes precedence. + +Each stage is opt-in and independently configurable via `LLAMA_ARG_THINK_BUDGET_*` environment variables at server startup, or as per-request overrides in the API call itself: + +| Environment variable | Purpose | +| --- | --- | +| `LLAMA_ARG_THINK_BUDGET` | Token budget for the reasoning block (existing upstream variable) | +| `LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE` | Templated intro message, supports a `{budget}` placeholder | +| `LLAMA_ARG_THINK_BUDGET_SOFT_RATIO` | Fraction of budget at which the soft warning fires (e.g. `0.7`) | +| `LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE` | Templated soft-warning message | +| `LLAMA_ARG_THINK_BUDGET_MESSAGE` | Templated hard-stop closing message | +| `LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS` | How long to wait for a paragraph break before forcing the hard stop | + +Default values preserve upstream's existing hard-cutoff behavior; the new stages are disabled unless configured. + +Planned follow-up work: generalize this mechanism into a configurable reasoning template/grammar, rather than a fixed set of budget-based checkpoints. + +## Benchmarking & Findings + +### Setup + +All results below use Qwen3.6-27B, `UD-Q4_K_XL` quantization, with MTP speculative decoding (3 draft tokens). A separate pass without speculative decoding produced consistent results and is omitted here for brevity. + +Four configurations are compared at several reasoning budgets, each adding one more piece of the mechanism on top of the last: + +- **Naive** — llama.cpp's existing default: the moment the budget is reached, `` is force-injected immediately, with no grace period and no in-context signaling of any kind. This is the behavior described in the Problem section above, and the baseline this fork is trying to improve on. +- **Hard-limit only** — this fork's hard-stop-with-grace-period stage used on its own (no soft warning, no intro message): instead of an immediate cutoff, the sampler waits up to `grace_tokens` for a paragraph boundary before closing the block. +- **Soft + hard** — the grace-period hard stop plus the soft-warning stage (fired at a configurable fraction of the budget), without the intro stage. +- **Intro + soft + hard** — the full three-stage mechanism: intro message, soft warning, hard stop with grace period. + +Two further reference points appear in the charts and tables: **Baseline (unlimited)**, where the reasoning block runs to its own natural ``, and — for LiveCodeBench only — **No reasoning**, where the reasoning block is disabled entirely. + +Benchmarks: HumanEval+ (n=164) and LiveCodeBench (`release_v6`, n=200). Reported token counts are average total completion tokens per test (reasoning block plus final answer), matching the chart axes. + +### Results: HumanEval+ + +HumanEval+: token consumption by configuration + +HumanEval+: accuracy vs. token cost + + +| Budget (tokens) | Naive: tok / pass@1 | Hard-limit only: tok / pass@1 | Soft + hard: tok / pass@1 | Intro + soft + hard: tok / pass@1 | +|---|---|---|---|---| +| 300 | 499 / 92.7% | 489 / 92.1% | 422 / 92.1% | 391 / 92.7% | +| 500 | 749 / 91.5% | 672 / 93.3% | 592 / 93.9% | 569 / 93.3% | +| 750 | 963 / 93.3% | 906 / 93.9% | 809 / 93.9% | 863 / 91.5% | +| 1250 | 1363 / 92.7% | 1348 / 92.7% | 1221 / 93.9% | 1360 / 95.7% | +| Unlimited (baseline) | 2776 / 92.7% | — | — | — | + +Two results stand out here: + +1. **Token consumption drops monotonically as guidance is added, at every budget.** Naive uses the most completion tokens of the four configurations at all four budgets, hard-limit-only is next, and soft + hard / intro + soft + hard are consistently the lowest — e.g. at budget 500: 749 (naive) → 672 (hard-limit only) → 592 (soft + hard) → 569 (intro + soft + hard). Soft + hard and intro + soft + hard are close to each other throughout, and which of the two is marginally lower varies by budget on this benchmark (soft + hard is lowest at 750 and 1250; intro + soft + hard is lowest at 300 and 500) — with n=164 this is likely within run-to-run noise rather than a real ordering between the two. +2. **Most configurations meet or beat the unlimited baseline (92.7%).** Of the 16 budget/configuration combinations, 12 are at or above 92.7%, and 8 exceed it outright — including the best result in the table, intro + soft + hard at budget 1250 (95.7%, using 1360 tokens against the baseline's 2776). The most plausible explanation is that constraining and guiding the reasoning block suppresses the repetition loops and non-convergent revision described in the Problem section above — on a benchmark like HumanEval+, where the model can typically reach a correct answer well within a modest token budget, an unconstrained reasoning block gives the model more opportunity to talk itself into a worse answer, not a better one. + +### Results: LiveCodeBench + +LiveCodeBench: token usage by reasoning budget and style + +LiveCodeBench: accuracy vs. token cost + +| Budget (tokens) | Naive: tok / pass@1 | Hard-limit only: tok / pass@1 | Soft + hard: tok / pass@1 | Intro + soft + hard: tok / pass@1 | +|---|---|---|---|---| +| 500 | 6955 / 61.0% | 4827 / 58.5% | 3894 / 61.5% | 2930 / 56.5% | +| 1000 | 7624 / 64.0% | 5820 / 65.5% | 4582 / 64.5% | 3334 / 60.0% | +| 1750 | 7779 / 62.0% | 6556 / 62.5% | 4864 / 68.5% | 4862 / 66.0% | +| 4000 | 16324 / 70.5% | 11251 / 65.5% | 10277 / 68.5% | 7693 / 69.5% | +| Unlimited (baseline) | 36293 / 72.0% | — | — | — | +| No reasoning | — / 57.0% | — | — | — | + +LiveCodeBench is far more reasoning-intensive at baseline (36293 tokens/task on average, versus 2776 for HumanEval+), and the ordering seen above holds even more cleanly here: **naive > hard-limit only > soft + hard > intro + soft + hard in total token count, at every single budget tested, with no exceptions.** At the 4000-token budget, naive uses 16324 tokens for 70.5% pass@1, while intro + soft + hard uses 7693 tokens — 47% of naive's token count — for 69.5%, a 1-point difference well within what n=200 sampling noise would produce. + +Accuracy differences between configurations at a fixed budget are generally small (a few points, consistent with n=200 noise) and don't show a systematic penalty for the more guided configurations — in most cases they hold accuracy roughly level with naive while using a fraction of the tokens. + +A separate effect shows up in how each configuration's accuracy responds to *increasing* the budget. For soft + hard and intro + soft + hard, pass@1 rises monotonically as budget increases from 500 to 4000, with no reversals. Naive and hard-limit-only do not show this: naive drops from 64.0% (1000 tokens) to 62.0% (1750 tokens) before jumping to 70.5% (4000 tokens), and hard-limit-only drops from 65.5% (1000 tokens) to 62.5% (1750 tokens). The guided configurations turn additional budget into a predictable accuracy gain; naive and hard-limit-only do not — this is the clearest "reduced noise" effect in this data. + +None of the four budget-constrained configurations fully recovers the unlimited baseline's 72.0% at any tested budget. + +### By difficulty (LiveCodeBench) + +| Difficulty | Baseline (unlimited) | 500-token budget (range across 4 configs) | 4000-token budget (range across 4 configs) | +|---|---|---|---| +| Easy (n=53) | 85% | 96–98% | 94–96% | +| Medium (n=61) | 72% | 61–72% | 72–80% | +| Hard (n=86) | 64% | 26–36% | 42–50% | + +This breakdown clarifies where the token savings come from, and echoes the HumanEval+ result above. On easy problems, every budget-constrained configuration at every tested budget scores at or above the unlimited baseline (96–98% vs. 85%) — again consistent with a capped, guided reasoning block reducing the chance the model overthinks its way into a wrong answer on a problem it could already solve. Medium problems are roughly flat to slightly improved at the higher budget. Hard problems are the exception: accuracy stays well below the unlimited baseline at both the smallest (26–36% vs. 64%) and largest (42–50% vs. 64%) budgets tested, for every configuration including intro + soft + hard. Budget-based control, however it's implemented, does not close this gap — the hardest problems still lose accuracy when reasoning length is capped. + +### Summary + +- Naive (llama.cpp's existing immediate-cutoff behavior) uses the most completion tokens of the four configurations at every budget tested, on both benchmarks — this is the behavior the mechanism is designed to improve on. +- Each additional stage of budget-aware guidance (grace period → soft warning → intro message) reduces token consumption further. On LiveCodeBench this ordering is exact at every budget: naive > hard-limit only > soft + hard > intro + soft + hard. +- Aggregate pass@1 does not show a systematic drop from budget constraints. On HumanEval+, 12 of 16 tested combinations meet or exceed the 92.7% unlimited baseline, and the single best result in either benchmark (95.7%) comes from the most heavily guided, budget-constrained configuration. +- Soft + hard and intro + soft + hard produce a monotonic, predictable accuracy/budget relationship on LiveCodeBench; naive and hard-limit-only do not. +- The gains are not evenly distributed across problem difficulty: easy-problem accuracy improves under constrained, guided budgets (consistent with reduced overthinking), while hard-problem accuracy remains below the unlimited baseline at every budget tested, for every configuration. + + +## Quick start + +Configuration is set via `LLAMA_ARG_THINK_BUDGET_*` environment variables, and can be overridden per-request in the API call — see [server API docs](tools/server/README.md) for the request-level parameters. + +### Apple Silicon + +Docker on macOS cannot pass the GPU through to a container, so there is no Metal-accelerated Docker image. Build natively instead, following upstream's [build guide](docs/build.md) (Metal is enabled by default on Apple Silicon): + +```sh +git clone https://github.com/laurencehardman/llama-mindcontrol +cd llama-mindcontrol +cmake -B build +cmake --build build --config Release -j + +LLAMA_ARG_THINK_BUDGET="350" \ +LLAMA_ARG_THINK_BUDGET_SOFT_RATIO="0.7" \ +LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS="64" \ +./build/bin/llama-server -m /path/to/your-model.gguf +``` + +### AMD64 + NVIDIA CUDA + +A pre-built Docker image is provided. Example `docker-compose.yml`: + +```yaml +services: + llama-server: + image: ghcr.io/laurencehardman/llama-mindcontrol:cuda + gpus: all + ports: + - "8080:8080" + volumes: + - ${MODEL_DIR:-./models}:/models:ro + environment: + LLAMA_ARG_THINK_BUDGET: "350" + LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE: " I have {budget} tokens to reason through this - that's enough room to work through it carefully, so I'll think it through step by step rather than rushing to a conclusion." + LLAMA_ARG_THINK_BUDGET_MESSAGE: " [!!NOTE TO SELF] I've used all of my thinking budget, I am now going to wrap up and provide the user their answer." + LLAMA_ARG_THINK_BUDGET_SOFT_RATIO: "0.7" + LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE: " [!NOTE TO SELF] I'm partway through my budget - I should start consolidating toward an answer, but I still have room to finish the important points." + command: + - "--model" + - "/models/your-model.gguf" +``` + +```sh +MODEL_DIR=/path/to/models docker compose up +``` + +Requires the [nvidia-container-toolkit](https://github.com/NVIDIA/nvidia-container-toolkit) on the host. + +`llama-server` exposes the standard OpenAI-compatible API at `http://localhost:8080`. See the upstream documentation below for other configuration options. + +Test it out: + +```sh +curl http://localhost:8080/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "local-model", + "stream": true, + "messages": [ + {"role": "user", "content": "Explain how Flash Attention works."} + ], + + "reasoning_budget_tokens": 350, + "reasoning_budget_message": "I have reached my reasoning budget - I have enough here to answer now.", + + "reasoning_budget_soft_ratio": 0.7, + "reasoning_budget_soft_message": "I am partway through my budget - I should start consolidating toward an answer, but I still have room to finish the important points.", + + "reasoning_budget_intro_message": "I have {budget} tokens to reason through this - that is enough room to work through it carefully, so I will think it through step by step rather than rushing to a conclusion.", + + "reasoning_budget_grace_tokens": 50, + + "reasoning_control": true + }' +``` + +--- # llama.cpp ![llama](https://raw.githubusercontent.com/ggml-org/llama.brand/refs/heads/master/cover/llama-cpp/cover-llama-cpp-dark.svg) diff --git a/common/arg.cpp b/common/arg.cpp index 79480e06f9d2..cffbc495ae70 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3555,6 +3555,36 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING")); + + // unescape \n, \t, \r, \\, \" in message strings passed via CLI/env + static auto unescape = [](const std::string & s) { + std::string r; + r.reserve(s.size()); + for (size_t i = 0; i < s.size(); i++) { + if (s[i] == '\\' && i + 1 < s.size()) { + switch (s[i + 1]) { + case 'n': r += '\n'; i++; break; + case 't': r += '\t'; i++; break; + case 'r': r += '\r'; i++; break; + case '\\': r += '\\'; i++; break; + case '"': r += '"'; i++; break; + default: r += s[i]; break; + } + } else { + r += s[i]; + } + } + return r; + }; + + add_opt(common_arg( + {"--reasoning-budget-enable"}, + {"--no-reasoning-budget-enable"}, + "master switch for the reasoning budget mechanism (hard cutoff, soft warning, intro message, grace period, and the runtime reasoning-control endpoint); if disabled, none of the other --reasoning-budget-* settings take effect regardless of their values (default: disabled)", + [](common_params & params, bool value) { + params.sampling.reasoning_budget_enabled = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_ENABLE")); add_opt(common_arg( {"--reasoning-budget"}, "N", "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)", @@ -3565,11 +3595,40 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET")); add_opt(common_arg( {"--reasoning-budget-message"}, "MESSAGE", - "message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)", + "message forced when the reasoning budget is exhausted; should include the model's own closing tag (e.g. ) as it is not appended automatically, since the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone so the block still always closes (default: none)", [](common_params & params, const std::string & value) { - params.sampling.reasoning_budget_message = value; + params.sampling.reasoning_budget_message = unescape(value); } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE")); + add_opt(common_arg( + {"--reasoning-budget-soft-ratio"}, "N", + "fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff: <= 0 disables, (0,1] enables (default: -1)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_soft_ratio = std::stof(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_SOFT_RATIO")); + add_opt(common_arg( + {"--reasoning-budget-soft-message"}, "MESSAGE", + "message injected at the soft reasoning budget threshold, before the hard cutoff (default: none)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_soft_message = unescape(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE")); + add_opt(common_arg( + {"--reasoning-budget-intro-message"}, "MESSAGE", + string_format("message forced immediately when the reasoning block starts, announcing the token budget; use {budget} as a placeholder for the configured reasoning budget (default: '%s')", + params.sampling.reasoning_budget_intro_message.c_str()), + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_intro_message = unescape(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE")); + add_opt(common_arg( + {"--reasoning-budget-grace-tokens"}, "N", + "once the reasoning budget is exhausted, wait up to N tokens for a paragraph break before forcing the cutoff, instead of forcing immediately (default: 0, disabled)", + [](common_params & params, int value) { + params.sampling.reasoning_budget_grace_tokens = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS")); add_opt(common_arg( {"--reasoning-preserve"}, {"--no-reasoning-preserve"}, diff --git a/common/common.h b/common/common.h index 919c0ea103a4..9fbf61f660dd 100644 --- a/common/common.h +++ b/common/common.h @@ -285,6 +285,7 @@ struct common_params_sampling { // reasoning budget sampler parameters // these are populated by the server/CLI based on chat template params + bool reasoning_budget_enabled = false; // master switch: must be true for the budget/soft/intro/grace mechanism to run at all int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget std::vector reasoning_budget_start; // start tag token sequence std::vector reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence @@ -292,6 +293,16 @@ struct common_params_sampling { std::string reasoning_budget_message; // message injected before end tag when budget exhausted bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime + float reasoning_budget_soft_ratio = -1.0f; // <= 0 = disabled, (0,1] = fraction of budget at which to warn + std::vector reasoning_budget_soft_forced; // tokenized soft warning message (no end tag) + std::string reasoning_budget_soft_message; // soft warning message injected at the soft threshold + + std::vector reasoning_budget_intro_forced; // tokenized intro message forced when the block starts (empty = disabled) + std::string reasoning_budget_intro_message = // intro message announcing the budget; supports a {budget} placeholder + "I'll keep this reasoning under {budget} tokens, so I'll stay focused and efficient. "; + + int32_t reasoning_budget_grace_tokens = 0; // <= 0 = force immediately, N>0 = wait up to N tokens for a paragraph break + bool backend_sampling = false; // print the parameters into a string diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 1fe242d062d1..dd6fcb99e108 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -65,12 +65,88 @@ struct common_reasoning_budget_ctx { size_t force_pos; // next position in forced_tokens to force int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none + + // soft warning + llama_tokens soft_forced_tokens; + bool soft_enabled; // soft_ratio > 0 and soft_forced_tokens non-empty + int32_t soft_threshold; // trigger soft warning once remaining <= this + bool soft_triggered; // soft warning already fired for this reasoning block + size_t soft_force_pos; // next position in soft_forced_tokens to force + + // intro announcement + llama_tokens intro_forced_tokens; + size_t intro_force_pos; // next position in intro_forced_tokens to force + + // graceful hard stop + int32_t grace_tokens; // max tokens to wait for a paragraph boundary once exhausted (<= 0 = disabled) + int32_t grace_remaining; // tokens left in the current grace wait + bool hard_pending_prev_nl; // whether the previous token in HARD_PENDING ended with a newline }; static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { return "reasoning-budget"; } +static bool token_utf8_complete(const common_reasoning_budget_ctx * ctx, llama_token token) { + if (ctx->vocab == nullptr) { + return true; + } + const std::string piece = common_token_to_piece(ctx->vocab, token, false); + return common_utf8_is_complete(piece); +} + +// Transitions into FORCING/WAITING_UTF8 depending on whether this token completes +// a UTF-8 sequence. Shared by every path that decides "start forcing the hard +// cutoff sequence right now". +static void common_reasoning_budget_begin_forcing(common_reasoning_budget_ctx * ctx, llama_token token) { + ctx->end_matcher.reset(); + if (token_utf8_complete(ctx, token)) { + ctx->state = REASONING_BUDGET_FORCING; + ctx->force_pos = 0; + } else { + ctx->state = REASONING_BUDGET_WAITING_UTF8; + } +} + +// Called when the budget hits zero (from COUNTING or SOFT_PENDING): either waits +// (bounded by grace_tokens) for a paragraph boundary, or forces immediately if no +// grace period is configured. +static void common_reasoning_budget_enter_hard_exhausted(common_reasoning_budget_ctx * ctx, llama_token token) { + if (ctx->grace_tokens > 0) { + ctx->state = REASONING_BUDGET_HARD_PENDING; + ctx->grace_remaining = ctx->grace_tokens; + ctx->hard_pending_prev_nl = false; + ctx->end_matcher.reset(); + COM_TRC("budget exhausted, waiting up to %d tokens for a paragraph break\n", ctx->grace_tokens); + return; + } + + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "budget exhausted, forcing end sequence\n"); +} + +// Called whenever a start sequence is (re-)matched, to (re-)activate the reasoning +// block: resets the budget countdown, then routes to the intro message (if +// configured), straight to the hard cutoff (budget <= 0), or normal counting. +static void common_reasoning_budget_activate(common_reasoning_budget_ctx * ctx) { + ctx->remaining = ctx->budget; + ctx->soft_triggered = false; + ctx->end_match = -1; + + if (!ctx->intro_forced_tokens.empty()) { + ctx->state = REASONING_BUDGET_INTRO_FORCING; + ctx->intro_force_pos = 0; + COM_TRC("activated, budget=%d tokens, forcing intro message\n", ctx->budget); + } else if (ctx->remaining <= 0) { + ctx->state = REASONING_BUDGET_FORCING; + ctx->force_pos = 0; + COM_TRC("%s", "budget=0, forcing immediately\n"); + } else { + ctx->state = REASONING_BUDGET_COUNTING; + COM_TRC("activated, budget=%d tokens\n", ctx->budget); + } +} + static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_token token) { auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; @@ -78,20 +154,24 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_IDLE: { if (ctx->start_matcher.advance(token) >= 0) { - ctx->state = REASONING_BUDGET_COUNTING; - ctx->remaining = ctx->budget; - COM_TRC("activated, budget=%d tokens\n", ctx->budget); - + common_reasoning_budget_activate(ctx); + } + break; + } + case REASONING_BUDGET_INTRO_FORCING: + ctx->intro_force_pos++; + if (ctx->intro_force_pos >= ctx->intro_forced_tokens.size()) { if (ctx->remaining <= 0) { ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; - COM_TRC("%s", "budget=0, forcing immediately\n"); + COM_TRC("%s", "intro complete, budget=0, forcing immediately\n"); + } else { + ctx->state = REASONING_BUDGET_COUNTING; + COM_TRC("%s", "intro complete, resuming countdown\n"); } } break; - } case REASONING_BUDGET_COUNTING: - case REASONING_BUDGET_WAITING_UTF8: { const int32_t match = ctx->end_matcher.advance(token); if (match >= 0) { @@ -101,33 +181,93 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } - bool utf8_complete = true; + ctx->remaining--; + if (ctx->remaining <= 0) { + common_reasoning_budget_enter_hard_exhausted(ctx, token); + break; + } + + if (ctx->soft_enabled && !ctx->soft_triggered && ctx->remaining <= ctx->soft_threshold) { + ctx->state = REASONING_BUDGET_SOFT_PENDING; + COM_TRC("soft threshold reached, remaining=%d, waiting for newline\n", ctx->remaining); + } + break; + } + case REASONING_BUDGET_SOFT_PENDING: + { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { + ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; + COM_TRC("%s", "deactivated (natural end)\n"); + break; + } + + ctx->remaining--; + if (ctx->remaining <= 0) { + // hard budget wins: abandon the soft warning, no newline is forced + COM_TRC("%s", "budget exhausted before newline, soft warning skipped\n"); + common_reasoning_budget_enter_hard_exhausted(ctx, token); + break; + } + if (ctx->vocab != nullptr) { const std::string piece = common_token_to_piece(ctx->vocab, token, false); - utf8_complete = common_utf8_is_complete(piece); + if (piece.find('\n') != std::string::npos) { + ctx->state = REASONING_BUDGET_SOFT_FORCING; + ctx->soft_force_pos = 0; + ctx->soft_triggered = true; + COM_TRC("%s", "newline boundary found, forcing soft warning\n"); + } + } + break; + } + case REASONING_BUDGET_SOFT_FORCING: + ctx->soft_force_pos++; + if (ctx->soft_force_pos >= ctx->soft_forced_tokens.size()) { + ctx->state = REASONING_BUDGET_COUNTING; + COM_TRC("%s", "soft warning complete, resuming countdown\n"); + } + break; + case REASONING_BUDGET_HARD_PENDING: + { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { + ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; + COM_TRC("%s", "deactivated (natural end)\n"); + break; } - if (ctx->state == REASONING_BUDGET_WAITING_UTF8) { - if (utf8_complete) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - ctx->end_matcher.reset(); - COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n"); - } - } else if (ctx->state == REASONING_BUDGET_COUNTING) { - ctx->remaining--; - if (ctx->remaining <= 0) { - if (utf8_complete) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - ctx->end_matcher.reset(); - COM_TRC("%s", "budget exhausted, forcing end sequence\n"); - } else { - ctx->state = REASONING_BUDGET_WAITING_UTF8; - ctx->end_matcher.reset(); - COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n"); - } - } + ctx->grace_remaining--; + + const std::string piece = ctx->vocab != nullptr ? common_token_to_piece(ctx->vocab, token, false) : std::string(); + const bool paragraph_boundary = piece.find("\n\n") != std::string::npos || + (ctx->hard_pending_prev_nl && !piece.empty() && piece[0] == '\n'); + ctx->hard_pending_prev_nl = !piece.empty() && piece.back() == '\n'; + + if (paragraph_boundary) { + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "paragraph boundary found, forcing end sequence\n"); + } else if (ctx->grace_remaining <= 0) { + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "grace period expired, forcing end sequence\n"); + } + break; + } + case REASONING_BUDGET_WAITING_UTF8: + { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { + ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; + COM_TRC("%s", "deactivated (natural end)\n"); + break; + } + + if (token_utf8_complete(ctx, token)) { + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n"); } break; } @@ -144,20 +284,12 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } case REASONING_BUDGET_DONE: - // Re-arm on a new start tag: some models emit multiple blocks - // per response, and each should get a fresh budget window. + // Re-arm on a new start sequence: some models emit multiple blocks + // per response, and each should get a fresh budget window (including + // its own intro message, if configured). if (ctx->start_matcher.advance(token) >= 0) { - ctx->state = REASONING_BUDGET_COUNTING; - ctx->remaining = ctx->budget; ctx->end_matcher.reset(); - ctx->end_match = -1; - COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); - - if (ctx->remaining <= 0) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - COM_TRC("%s", "budget=0, forcing immediately\n"); - } + common_reasoning_budget_activate(ctx); } break; } @@ -166,17 +298,28 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) { auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - if (ctx->state != REASONING_BUDGET_FORCING) { - // passthrough — don't modify logits - return; - } + llama_token forced; - if (ctx->force_pos >= ctx->forced_tokens.size()) { + if (ctx->state == REASONING_BUDGET_FORCING) { + if (ctx->force_pos >= ctx->forced_tokens.size()) { + return; + } + forced = ctx->forced_tokens[ctx->force_pos]; + } else if (ctx->state == REASONING_BUDGET_SOFT_FORCING) { + if (ctx->soft_force_pos >= ctx->soft_forced_tokens.size()) { + return; + } + forced = ctx->soft_forced_tokens[ctx->soft_force_pos]; + } else if (ctx->state == REASONING_BUDGET_INTRO_FORCING) { + if (ctx->intro_force_pos >= ctx->intro_forced_tokens.size()) { + return; + } + forced = ctx->intro_forced_tokens[ctx->intro_force_pos]; + } else { + // passthrough — don't modify logits return; } - const llama_token forced = ctx->forced_tokens[ctx->force_pos]; - // set all logits to -inf except the forced token for (size_t i = 0; i < cur_p->size; i++) { if (cur_p->data[i].id != forced) { @@ -193,12 +336,18 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { ctx->end_matcher.reset(); ctx->force_pos = 0; ctx->end_match = -1; + ctx->soft_triggered = false; + ctx->soft_force_pos = 0; + ctx->intro_force_pos = 0; + ctx->grace_remaining = ctx->grace_tokens; + ctx->hard_pending_prev_nl = false; } static struct llama_sampler * common_reasoning_budget_init_state( const struct llama_vocab * vocab, const std::vector & start_seqs, const std::vector & end_seqs, const llama_tokens & forced_tokens, - int32_t budget, common_reasoning_budget_state initial_state); + const llama_tokens & soft_forced_tokens, const llama_tokens & intro_forced_tokens, + int32_t budget, float soft_ratio, int32_t grace_tokens, common_reasoning_budget_state initial_state); static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl); @@ -233,25 +382,46 @@ static struct llama_sampler * common_reasoning_budget_init_state( const std::vector & start_seqs, const std::vector & end_seqs, const llama_tokens & forced_tokens, + const llama_tokens & soft_forced_tokens, + const llama_tokens & intro_forced_tokens, int32_t budget, + float soft_ratio, + int32_t grace_tokens, common_reasoning_budget_state initial_state) { // promote COUNTING with budget <= 0 to FORCING if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { initial_state = REASONING_BUDGET_FORCING; } + const bool soft_enabled = soft_ratio > 0.0f && !soft_forced_tokens.empty(); + int32_t soft_threshold = 0; + if (soft_enabled) { + const float ratio = std::min(soft_ratio, 1.0f); + soft_threshold = std::max(0, budget - (int32_t) std::ceil(budget * ratio)); + } + return llama_sampler_init( /* .iface = */ &common_reasoning_budget_i, /* .ctx = */ new common_reasoning_budget_ctx { - /* .vocab = */ vocab, - /* .start_matcher = */ token_matcher(start_seqs), - /* .end_matcher = */ token_matcher(end_seqs), - /* .forced_tokens = */ forced_tokens, - /* .budget = */ budget, - /* .remaining = */ budget, - /* .state = */ initial_state, - /* .force_pos = */ 0, - /* .end_match = */ -1, + /* .vocab = */ vocab, + /* .start_matcher = */ token_matcher(start_seqs), + /* .end_matcher = */ token_matcher(end_seqs), + /* .forced_tokens = */ forced_tokens, + /* .budget = */ budget, + /* .remaining = */ budget, + /* .state = */ initial_state, + /* .force_pos = */ 0, + /* .end_match = */ -1, + /* .soft_forced_tokens = */ soft_forced_tokens, + /* .soft_enabled = */ soft_enabled, + /* .soft_threshold = */ soft_threshold, + /* .soft_triggered = */ false, + /* .soft_force_pos = */ 0, + /* .intro_forced_tokens = */ intro_forced_tokens, + /* .intro_force_pos = */ 0, + /* .grace_tokens = */ grace_tokens, + /* .grace_remaining = */ grace_tokens, + /* .hard_pending_prev_nl = */ false, } ); } @@ -261,9 +431,13 @@ struct llama_sampler * common_reasoning_budget_init( const std::vector & start_seqs, const std::vector & end_seqs, const llama_tokens & forced_tokens, + const llama_tokens & soft_forced_tokens, + const llama_tokens & intro_forced_tokens, int32_t budget, + float soft_ratio, + int32_t grace_tokens, common_reasoning_budget_state initial_state) { - return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state); + return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, soft_forced_tokens, intro_forced_tokens, budget, soft_ratio, grace_tokens, initial_state); } common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { @@ -293,9 +467,15 @@ bool common_reasoning_budget_force(struct llama_sampler * smpl) { auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - // only a sampler that is actively counting down the budget may be forced; - // any other state (idle, already forcing/waiting, or done) is left untouched - if (ctx->state != REASONING_BUDGET_COUNTING) { + // only a sampler that is actively counting down the budget (or emitting the + // intro/soft messages, or waiting out the post-exhaustion grace period) may + // be forced; any other state (idle, already hard-forcing/waiting, or done) + // is left untouched + if (ctx->state != REASONING_BUDGET_COUNTING && + ctx->state != REASONING_BUDGET_INTRO_FORCING && + ctx->state != REASONING_BUDGET_SOFT_PENDING && + ctx->state != REASONING_BUDGET_SOFT_FORCING && + ctx->state != REASONING_BUDGET_HARD_PENDING) { return false; } diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 1b89a04c42e8..5d3a26e0c3ed 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -8,37 +8,67 @@ #include enum common_reasoning_budget_state { - REASONING_BUDGET_IDLE, // waiting for start sequence - REASONING_BUDGET_COUNTING, // counting down tokens - REASONING_BUDGET_FORCING, // forcing budget message + end sequence - REASONING_BUDGET_WAITING_UTF8, // budget exhausted, waiting for UTF-8 completion - REASONING_BUDGET_DONE, // passthrough forever + REASONING_BUDGET_IDLE, // waiting for start sequence + REASONING_BUDGET_INTRO_FORCING, // forcing the intro/announcement message + REASONING_BUDGET_COUNTING, // counting down tokens + REASONING_BUDGET_SOFT_PENDING, // soft threshold crossed, waiting for a newline boundary + REASONING_BUDGET_SOFT_FORCING, // forcing the soft warning message + REASONING_BUDGET_HARD_PENDING, // budget exhausted, waiting (bounded) for a paragraph boundary + REASONING_BUDGET_FORCING, // forcing budget message + end sequence + REASONING_BUDGET_WAITING_UTF8, // budget exhausted, waiting for UTF-8 completion + REASONING_BUDGET_DONE, // passthrough forever }; // Creates a reasoning budget sampler that limits token generation inside a // reasoning block (e.g. between and ). // -// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE -// IDLE: passthrough, watching for a start sequence -// COUNTING: counting down remaining tokens, watching for a natural end sequence -// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence -// FORCING: forces forced_tokens token-by-token (all other logits -> -inf) -// DONE: passthrough forever +// State machine: IDLE -> INTRO_FORCING -> COUNTING -> SOFT_PENDING -> SOFT_FORCING -> COUNTING -> HARD_PENDING -> WAITING_UTF8 -> FORCING -> DONE +// IDLE: passthrough, watching for a start sequence +// INTRO_FORCING: forces intro_forced_tokens token-by-token right as the block starts, then proceeds to COUNTING (or straight to FORCING if budget <= 0) +// COUNTING: counting down remaining tokens, watching for a natural end sequence +// SOFT_PENDING: soft threshold crossed, waiting for a newline token before warning +// SOFT_FORCING: forces soft_forced_tokens token-by-token, then returns to COUNTING +// HARD_PENDING: budget exhausted, waiting (up to grace_tokens) for a paragraph boundary before forcing +// WAITING_UTF8: waiting to force, allowing tokens to complete a UTF-8 sequence +// FORCING: forces forced_tokens token-by-token (all other logits -> -inf) +// DONE: passthrough forever +// +// The hard cutoff always takes priority over the soft warning: if the budget is +// exhausted before a newline boundary is found in SOFT_PENDING, the soft warning +// is abandoned and the hard-cutoff path proceeds as normal (including any grace +// period below). +// +// Intro tokens (like soft/forced tokens) do not count against the budget: the +// countdown only starts once INTRO_FORCING completes. +// +// When the budget is exhausted (from COUNTING or SOFT_PENDING), if grace_tokens > 0 +// the sampler enters HARD_PENDING instead of forcing immediately: it waits for a +// paragraph boundary (two adjacent newlines) so the cutoff lands at a clean break, +// but forces anyway once grace_tokens more tokens have passed without one - the +// hard guarantee on total length is bounded by budget + grace_tokens, never open-ended. // // Parameters: -// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) -// start_seqs - token sequences, any of which activates counting -// end_seqs - token sequences, any of which naturally deactivates -// forced_tokens - token sequence forced when budget expires -// budget - max tokens allowed in the reasoning block -// initial_state - initial state +// vocab - vocabulary (used for UTF-8/paragraph boundary detection; can be nullptr) +// start_seqs - token sequences, any of which activates counting +// end_seqs - token sequences, any of which naturally deactivates +// forced_tokens - token sequence forced when budget expires +// soft_forced_tokens - token sequence forced once at the soft threshold (empty = disabled) +// intro_forced_tokens - token sequence forced once right as the block starts (empty = disabled) +// budget - max tokens allowed in the reasoning block +// soft_ratio - fraction of budget consumed at which to trigger the soft warning (<= 0 disables it) +// grace_tokens - max tokens to wait for a paragraph boundary after the budget expires (<= 0 = force immediately, no wait) +// initial_state - initial state // struct llama_sampler * common_reasoning_budget_init( const struct llama_vocab * vocab, const std::vector & start_seqs, const std::vector & end_seqs, const llama_tokens & forced_tokens, + const llama_tokens & soft_forced_tokens, + const llama_tokens & intro_forced_tokens, int32_t budget, + float soft_ratio = -1.0f, + int32_t grace_tokens = 0, common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); diff --git a/common/sampling.cpp b/common/sampling.cpp index 256ac161e20f..70d531cd5890 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -295,18 +295,40 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st } } - // reasoning budget sampler (skip when budget is unlimited unless a lazy grammar is active, which needs rbudget for thinking-block suppression) - if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) { + // reasoning budget sampler. The master switch (reasoning_budget_enabled) gates the + // budget/soft/intro/grace mechanism and the manual reasoning_control endpoint - both are + // "thought control" features. It does not gate grammar_lazy's own need for this sampler, + // which is an unrelated grammar/tool-calling concern (suppressing grammar constraints while + // inside a thinking block), not part of the budget-forcing mechanism itself. + const bool reasoning_budget_active = params.reasoning_budget_enabled && (params.reasoning_budget_tokens >= 0 || params.reasoning_control); + if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || reasoning_budget_active)) { rbudget = common_reasoning_budget_init( vocab, {params.reasoning_budget_start}, params.reasoning_budget_end, params.reasoning_budget_forced, - params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); + params.reasoning_budget_soft_forced, + params.reasoning_budget_intro_forced, + params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens, + params.reasoning_budget_soft_ratio, + params.reasoning_budget_grace_tokens); for (const auto & token : prefill_tokens) { llama_sampler_accept(rbudget, token); LOG_DBG("%s: reasoning-budget accepted prefill token (%d)\n", __func__, token); + + // Some chat templates bake the start tag (and trailing text, e.g. a + // newline right after "") directly into the prefill. If that + // activates a forcing state here, any further prefill tokens are + // already-fixed prompt text, not live model output - feeding them + // in would be misread as the model having already emitted the start + // of the forced sequence, silently skipping ahead in it. + const auto state = common_reasoning_budget_get_state(rbudget); + if (state == REASONING_BUDGET_INTRO_FORCING || + state == REASONING_BUDGET_SOFT_FORCING || + state == REASONING_BUDGET_FORCING) { + break; + } } } diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index 3bcc77e1733c..eb9dadd516fe 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -47,7 +47,11 @@ static void test_reasoning_budget( start_seqs, end_seqs, forced_tokens, + {}, // soft_forced_tokens - soft warning not exercised by this helper + {}, // intro_forced_tokens - intro message not exercised by this helper budget, + -1.0f, // soft_ratio - disabled + 0, // grace_tokens - graceful hard stop not exercised by this helper initial_state ); @@ -156,7 +160,7 @@ static void test_reasoning_budget_clone_mid_counting() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 2, -1.0f, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 llama_sampler_accept(sampler, 50); // COUNTING, remaining=1 @@ -175,7 +179,7 @@ static void test_reasoning_budget_clone_mid_forcing() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 0, -1.0f, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token @@ -195,7 +199,7 @@ static void test_reasoning_budget_force_manual() { // if COUNTING, force() succeeds and begins forcing the end sequence from the start { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 5, -1.0f, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=5 llama_sampler_accept(sampler, 50); // COUNTING, remaining=4 @@ -216,7 +220,7 @@ static void test_reasoning_budget_force_manual() { // if IDLE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 5, -1.0f, 0, REASONING_BUDGET_IDLE); GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE"); GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE); @@ -226,7 +230,7 @@ static void test_reasoning_budget_force_manual() { // if DONE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 5, -1.0f, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 101); // natural end -> DONE @@ -240,7 +244,7 @@ static void test_reasoning_budget_force_manual() { // if FORCING, force() is a no-op and must not rewind the force position { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 0, -1.0f, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1) @@ -258,13 +262,340 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } +// Soft warning: crossing the soft threshold moves COUNTING -> SOFT_PENDING, and +// (with a null vocab, so no newline is ever detected) the hard cutoff exhausting +// first correctly abandons the soft warning and forces the hard sequence instead. +static void test_reasoning_budget_soft_warning_skipped_before_hard_cutoff() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector soft_forced = {200, 201}; + + // budget=10, soft_ratio=0.5 -> soft_threshold = 10 - ceil(10*0.5) = 5 + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, soft_forced, {}, 10, 0.5f, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING, remaining=10 + for (llama_token t : {50, 51, 52, 53}) { + llama_sampler_accept(sampler, t); // remaining -> 9,8,7,6 + } + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + llama_sampler_accept(sampler, 54); // remaining=5 <= soft_threshold(5) -> SOFT_PENDING + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + + // no vocab -> no newline is ever found, so the budget clock keeps running + // in SOFT_PENDING until it hits zero, at which point the soft warning must + // be abandoned and the hard cutoff must fire instead + for (llama_token t : {55, 56, 57, 58}) { + llama_sampler_accept(sampler, t); // remaining -> 4,3,2,1 + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + } + llama_sampler_accept(sampler, 59); // remaining=0 -> hard cutoff, soft skipped + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102 && "hard message must fire, not the soft warning"); + + llama_sampler_accept(sampler, 102); + GGML_ASSERT(get_forced_token(sampler, 202) == 101); + llama_sampler_accept(sampler, 101); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'soft warning skipped before hard cutoff' passed\n"); +} + +// SOFT_FORCING forces soft_forced_tokens token-by-token, then resumes COUNTING +// (unlike FORCING, which ends the block by transitioning to DONE). +static void test_reasoning_budget_soft_forcing_resumes_counting() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector soft_forced = {200, 201}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, soft_forced, {}, 5, 0.5f, 0, REASONING_BUDGET_SOFT_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 201) == 200); + llama_sampler_accept(sampler, 200); // advance to the second soft token + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 201) == 201); + llama_sampler_accept(sampler, 201); // soft sequence complete + + // resumes COUNTING (not DONE) - the reasoning block is not over + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'soft forcing resumes counting' passed\n"); +} + +// A manual force() call must always win, abandoning any in-flight soft warning +// and jumping straight to the hard FORCING sequence from force_pos=0. +static void test_reasoning_budget_force_manual_from_soft_states() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector soft_forced = {200, 201}; + + // from SOFT_PENDING + { + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, soft_forced, {}, 10, 0.5f, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING, remaining=10 + for (llama_token t : {50, 51, 52, 53, 54}) { + llama_sampler_accept(sampler, t); // remaining -> 9..5, crosses threshold at 5 + } + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + + GGML_ASSERT(common_reasoning_budget_force(sampler) && "force() should succeed from SOFT_PENDING"); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102 && "force() must jump to the hard sequence, not the soft one"); + + llama_sampler_free(sampler); + } + + // from SOFT_FORCING + { + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, soft_forced, {}, 5, 0.5f, 0, REASONING_BUDGET_SOFT_FORCING); + + llama_sampler_accept(sampler, 200); // advance into the soft sequence (soft_force_pos=1) + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_FORCING); + + GGML_ASSERT(common_reasoning_budget_force(sampler) && "force() should succeed from SOFT_FORCING"); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102 && "force() must restart the hard sequence from force_pos=0"); + + llama_sampler_free(sampler); + } + + fprintf(stderr, " Test 'manual force transition from soft states' passed\n"); +} + +// The intro message fires immediately when the start tag is matched, before any +// budget counting, and does not itself count against the budget. +static void test_reasoning_budget_intro_forcing_then_counting() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector intro_forced = {300, 301}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, intro_forced, 3, -1.0f, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // start tag matched -> straight to INTRO_FORCING (not COUNTING) + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 301) == 300); + llama_sampler_accept(sampler, 300); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 301) == 301); + llama_sampler_accept(sampler, 301); // intro sequence complete -> COUNTING, remaining still full budget + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + // the intro tokens must not have consumed any of the budget: exactly 3 more + // generic tokens are needed to exhaust it + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 51); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + llama_sampler_accept(sampler, 52); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 102); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'intro forcing then counting' passed\n"); +} + +// If the budget is 0, the intro message still fires first (explaining why the +// hard cutoff follows immediately), and only then does the hard FORCING begin. +static void test_reasoning_budget_intro_forcing_budget_zero() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector intro_forced = {300, 301}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, intro_forced, 0, -1.0f, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + llama_sampler_accept(sampler, 300); + llama_sampler_accept(sampler, 301); // intro complete, budget<=0 -> straight to hard FORCING + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 102); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'intro forcing with budget=0' passed\n"); +} + +// A manual force() call must also win from INTRO_FORCING, abandoning the +// partial intro message and jumping straight to the hard sequence. +static void test_reasoning_budget_force_manual_from_intro() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector intro_forced = {300, 301}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, intro_forced, 5, -1.0f, 0, REASONING_BUDGET_INTRO_FORCING); + + llama_sampler_accept(sampler, 300); // advance into the intro sequence (intro_force_pos=1) + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + GGML_ASSERT(common_reasoning_budget_force(sampler) && "force() should succeed from INTRO_FORCING"); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 102 && "force() must jump to the hard sequence, not the intro one"); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'manual force transition from intro' passed\n"); +} + +// Each new block (re-armed after DONE) gets its own intro message too. +static void test_reasoning_budget_intro_rearms_on_multiblock() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector intro_forced = {300, 301}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, intro_forced, 5, -1.0f, 0, REASONING_BUDGET_IDLE); + + // first block: intro, then a natural end before the budget is touched + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + llama_sampler_accept(sampler, 300); + llama_sampler_accept(sampler, 301); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + llama_sampler_accept(sampler, 101); // natural end + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + // second block: re-arm must go through INTRO_FORCING again, from the start + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 300 && "second block must restart the intro sequence from position 0"); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'intro re-arms on multi-block' passed\n"); +} + +// When the budget is exhausted and a grace period is configured, the sampler +// waits in HARD_PENDING rather than forcing immediately. With a null vocab, no +// paragraph boundary can ever be detected (safe fallback, same as the UTF-8 and +// soft-newline checks), so this exercises the "grace period expires" path. +static void test_reasoning_budget_hard_pending_grace_expires() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 2, -1.0f, 3, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 + llama_sampler_accept(sampler, 50); // remaining=1 + llama_sampler_accept(sampler, 51); // remaining=0 -> HARD_PENDING, grace_remaining=3 + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + llama_sampler_accept(sampler, 52); // grace_remaining=2 + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + llama_sampler_accept(sampler, 53); // grace_remaining=1 + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + llama_sampler_accept(sampler, 54); // grace_remaining=0 -> grace expired, force now + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 102) == 102); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'hard pending grace expires' passed\n"); +} + +// A natural end tag seen while waiting out the grace period still wins, same as +// in SOFT_PENDING/COUNTING. +static void test_reasoning_budget_hard_pending_natural_end() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 2, -1.0f, 5, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 + llama_sampler_accept(sampler, 50); // remaining=1 + llama_sampler_accept(sampler, 51); // remaining=0 -> HARD_PENDING + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + llama_sampler_accept(sampler, 101); // natural end tag while pending + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'hard pending natural end' passed\n"); +} + +// Manual force() must also win from HARD_PENDING, skipping the rest of the grace period. +static void test_reasoning_budget_force_manual_from_hard_pending() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, {}, {}, 2, -1.0f, 10, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 51); // remaining=0 -> HARD_PENDING, grace_remaining=10 + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + GGML_ASSERT(common_reasoning_budget_force(sampler) && "force() should succeed from HARD_PENDING"); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 102) == 102); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'manual force transition from hard pending' passed\n"); +} + +// Exhaustion reached via SOFT_PENDING (soft warning abandoned) must also route +// through the grace period when one is configured, not skip straight to FORCING. +static void test_reasoning_budget_soft_pending_exhaustion_uses_grace() { + const std::vector start = {100}; + const std::vector end = {101}; + const std::vector forced = {102, 101}; + const std::vector soft_forced = {200, 201}; + + // budget=10, soft_ratio=0.5 -> soft_threshold=5; grace_tokens=2 + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, soft_forced, {}, 10, 0.5f, 2, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING, remaining=10 + for (llama_token t : {50, 51, 52, 53, 54}) { + llama_sampler_accept(sampler, t); // remaining -> 9..5, crosses soft threshold at 5 + } + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + + // no vocab -> no newline ever found, budget keeps running down in SOFT_PENDING + for (llama_token t : {55, 56, 57, 58}) { + llama_sampler_accept(sampler, t); // remaining -> 4,3,2,1 + } + llama_sampler_accept(sampler, 59); // remaining=0, grace_tokens=2 -> HARD_PENDING, not immediate FORCING + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + llama_sampler_accept(sampler, 60); // grace_remaining=1 + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + llama_sampler_accept(sampler, 61); // grace_remaining=0 -> force + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102 && "hard message must fire, not the soft one"); + + llama_sampler_free(sampler); + + fprintf(stderr, " Test 'soft pending exhaustion uses grace period' passed\n"); +} + +// Upstream multi-pattern matcher: end_match records which end sequence closed the +// block (natural or forced), and is cleared on re-arm. static void test_reasoning_budget_end_match() { const std::vector start = {{100}}; const std::vector end = {{101}, {103, 104}}; // natural end records the sequence that matched; re-arming clears it { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, {}, {}, 5, -1.0f, 0, REASONING_BUDGET_IDLE); GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); @@ -287,7 +618,7 @@ static void test_reasoning_budget_end_match() { { const std::vector end_overlap = {{104}, {103, 104}}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, {}, {}, 5, -1.0f, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 103); @@ -302,7 +633,7 @@ static void test_reasoning_budget_end_match() { // forcing records the end sequence terminating forced_tokens { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, {}, {}, 0, -1.0f, 0, REASONING_BUDGET_FORCING); llama_sampler_accept(sampler, 102); llama_sampler_accept(sampler, 103); @@ -318,7 +649,7 @@ static void test_reasoning_budget_end_match() { // forced_tokens not ending with a known end sequence records nothing { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, {}, {}, 0, -1.0f, 0, REASONING_BUDGET_FORCING); llama_sampler_accept(sampler, 102); // forced sequence complete, DONE GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); @@ -493,9 +824,20 @@ int main(void) { test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); + test_reasoning_budget_soft_warning_skipped_before_hard_cutoff(); + test_reasoning_budget_soft_forcing_resumes_counting(); + test_reasoning_budget_force_manual_from_soft_states(); + test_reasoning_budget_intro_forcing_then_counting(); + test_reasoning_budget_intro_forcing_budget_zero(); + test_reasoning_budget_force_manual_from_intro(); + test_reasoning_budget_intro_rearms_on_multiblock(); + test_reasoning_budget_hard_pending_grace_expires(); + test_reasoning_budget_hard_pending_natural_end(); + test_reasoning_budget_force_manual_from_hard_pending(); + test_reasoning_budget_soft_pending_exhaustion_uses_grace(); test_reasoning_budget_end_match(); - printf("OK (12 tests passed)\n"); + printf("OK (23 tests passed)\n"); printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); diff --git a/tools/cli/README.md b/tools/cli/README.md index bcddd05702bb..b6cd28125bf5 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -172,8 +172,13 @@ | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)
(env: LLAMA_ARG_JINJA) | | `--reasoning-format FORMAT` | controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:
- none: leaves thoughts unparsed in `message.content`
- deepseek: puts thoughts in `message.reasoning_content`
- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`
(default: auto)
(env: LLAMA_ARG_THINK) | | `-rea, --reasoning [on\|off\|auto]` | Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))
(env: LLAMA_ARG_REASONING) | +| `--reasoning-budget-enable, --no-reasoning-budget-enable` | master switch for the reasoning budget mechanism (hard cutoff, soft warning, intro message, grace period, and the runtime reasoning-control endpoint); if disabled, none of the other --reasoning-budget-* settings take effect regardless of their values (default: disabled)
(env: LLAMA_ARG_THINK_BUDGET_ENABLE) | | `--reasoning-budget N` | token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)
(env: LLAMA_ARG_THINK_BUDGET) | -| `--reasoning-budget-message MESSAGE` | message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-message MESSAGE` | message forced when the reasoning budget is exhausted; should include the model's own closing tag (e.g. ) as it is not appended automatically, since the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone so the block still always closes (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-soft-ratio N` | fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff: <= 0 disables, (0,1] enables (default: -1)
(env: LLAMA_ARG_THINK_BUDGET_SOFT_RATIO) | +| `--reasoning-budget-soft-message MESSAGE` | message injected at the soft reasoning budget threshold, before the hard cutoff (default: none)
(env: LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE) | +| `--reasoning-budget-intro-message MESSAGE` | message forced immediately when the reasoning block starts, announcing the token budget; use {budget} as a placeholder for the configured reasoning budget (default: 'I'll keep this reasoning under {budget} tokens, so I'll stay focused and efficient. ')
(env: LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE) | +| `--reasoning-budget-grace-tokens N` | once the reasoning budget is exhausted, wait up to N tokens for a paragraph break before forcing the cutoff, instead of forcing immediately (default: 0, disabled)
(env: LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS) | | `--reasoning-preserve, --no-reasoning-preserve` | preserve reasoning trace in the full history, not just the last assistant message (default: template default)
compatible with certain templates having 'supports_preserve_reasoning' capability
example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking
(env: LLAMA_ARG_REASONING_PRESERVE) | | `--chat-template JINJA_TEMPLATE` | set custom jinja chat template (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE) | | `--chat-template-file JINJA_TEMPLATE_FILE` | set custom jinja chat template file (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE_FILE) | diff --git a/tools/completion/README.md b/tools/completion/README.md index bce71d68d949..9506b08d4b3d 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -253,8 +253,13 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: disabled)
(env: LLAMA_ARG_JINJA) | | `--reasoning-format FORMAT` | controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:
- none: leaves thoughts unparsed in `message.content`
- deepseek: puts thoughts in `message.reasoning_content`
- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`
(default: auto)
(env: LLAMA_ARG_THINK) | | `-rea, --reasoning [on\|off\|auto]` | Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))
(env: LLAMA_ARG_REASONING) | +| `--reasoning-budget-enable, --no-reasoning-budget-enable` | master switch for the reasoning budget mechanism (hard cutoff, soft warning, intro message, grace period, and the runtime reasoning-control endpoint); if disabled, none of the other --reasoning-budget-* settings take effect regardless of their values (default: disabled)
(env: LLAMA_ARG_THINK_BUDGET_ENABLE) | | `--reasoning-budget N` | token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)
(env: LLAMA_ARG_THINK_BUDGET) | -| `--reasoning-budget-message MESSAGE` | message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-message MESSAGE` | message forced when the reasoning budget is exhausted; should include the model's own closing tag (e.g. ) as it is not appended automatically, since the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone so the block still always closes (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-soft-ratio N` | fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff: <= 0 disables, (0,1] enables (default: -1)
(env: LLAMA_ARG_THINK_BUDGET_SOFT_RATIO) | +| `--reasoning-budget-soft-message MESSAGE` | message injected at the soft reasoning budget threshold, before the hard cutoff (default: none)
(env: LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE) | +| `--reasoning-budget-intro-message MESSAGE` | message forced immediately when the reasoning block starts, announcing the token budget; use {budget} as a placeholder for the configured reasoning budget (default: 'I'll keep this reasoning under {budget} tokens, so I'll stay focused and efficient. ')
(env: LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE) | +| `--reasoning-budget-grace-tokens N` | once the reasoning budget is exhausted, wait up to N tokens for a paragraph break before forcing the cutoff, instead of forcing immediately (default: 0, disabled)
(env: LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS) | | `--reasoning-preserve, --no-reasoning-preserve` | preserve reasoning trace in the full history, not just the last assistant message (default: template default)
compatible with certain templates having 'supports_preserve_reasoning' capability
example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking
(env: LLAMA_ARG_REASONING_PRESERVE) | | `--chat-template JINJA_TEMPLATE` | set custom jinja chat template (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE) | | `--chat-template-file JINJA_TEMPLATE_FILE` | set custom jinja chat template file (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE_FILE) | diff --git a/tools/server/README.md b/tools/server/README.md index f45c018972d2..7e0d63ed99a3 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -225,8 +225,13 @@ For the full list of features, please refer to [server's changelog](https://gith | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)
(env: LLAMA_ARG_JINJA) | | `--reasoning-format FORMAT` | controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:
- none: leaves thoughts unparsed in `message.content`
- deepseek: puts thoughts in `message.reasoning_content`
- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`
(default: auto)
(env: LLAMA_ARG_THINK) | | `-rea, --reasoning [on\|off\|auto]` | Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))
(env: LLAMA_ARG_REASONING) | +| `--reasoning-budget-enable, --no-reasoning-budget-enable` | master switch for the reasoning budget mechanism (hard cutoff, soft warning, intro message, grace period, and the runtime reasoning-control endpoint); if disabled, none of the other --reasoning-budget-* settings take effect regardless of their values (default: disabled)
(env: LLAMA_ARG_THINK_BUDGET_ENABLE) | | `--reasoning-budget N` | token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)
(env: LLAMA_ARG_THINK_BUDGET) | -| `--reasoning-budget-message MESSAGE` | message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-message MESSAGE` | message forced when the reasoning budget is exhausted; should include the model's own closing tag (e.g. ) as it is not appended automatically, since the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone so the block still always closes (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-soft-ratio N` | fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff: <= 0 disables, (0,1] enables (default: -1)
(env: LLAMA_ARG_THINK_BUDGET_SOFT_RATIO) | +| `--reasoning-budget-soft-message MESSAGE` | message injected at the soft reasoning budget threshold, before the hard cutoff (default: none)
(env: LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE) | +| `--reasoning-budget-intro-message MESSAGE` | message forced immediately when the reasoning block starts, announcing the token budget; use {budget} as a placeholder for the configured reasoning budget (default: 'I'll keep this reasoning under {budget} tokens, so I'll stay focused and efficient. ')
(env: LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE) | +| `--reasoning-budget-grace-tokens N` | once the reasoning budget is exhausted, wait up to N tokens for a paragraph break before forcing the cutoff, instead of forcing immediately (default: 0, disabled)
(env: LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS) | | `--reasoning-preserve, --no-reasoning-preserve` | preserve reasoning trace in the full history, not just the last assistant message (default: template default)
compatible with certain templates having 'supports_preserve_reasoning' capability
example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking
(env: LLAMA_ARG_REASONING_PRESERVE) | | `--chat-template JINJA_TEMPLATE` | set custom jinja chat template (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE) | | `--chat-template-file JINJA_TEMPLATE_FILE` | set custom jinja chat template file (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE_FILE) | diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index c9109fc9626e..984447511d48 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1131,12 +1131,29 @@ json oaicompat_chat_params_parse( reasoning_budget = opt.reasoning_budget; } + const bool reasoning_control = json_value(body, "reasoning_control", false); + if (!chat_params.thinking_end_tags.empty()) { llama_params["reasoning_budget_tokens"] = reasoning_budget; llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag; llama_params["reasoning_budget_end_tags"] = chat_params.thinking_end_tags; - llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); - llama_params["reasoning_control"] = json_value(body, "reasoning_control", false); + llama_params["reasoning_control"] = reasoning_control; + + const bool has_per_request_reasoning = reasoning_budget >= 0 + || reasoning_control + || body.contains("reasoning_budget_message") + || body.contains("reasoning_budget_soft_ratio") + || body.contains("reasoning_budget_soft_message") + || body.contains("reasoning_budget_intro_message") + || body.contains("reasoning_budget_grace_tokens"); + + if (has_per_request_reasoning) { + llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); + llama_params["reasoning_budget_soft_ratio"] = json_value(body, "reasoning_budget_soft_ratio", opt.reasoning_budget_soft_ratio); + llama_params["reasoning_budget_soft_message"] = json_value(body, "reasoning_budget_soft_message", opt.reasoning_budget_soft_message); + llama_params["reasoning_budget_intro_message"] = json_value(body, "reasoning_budget_intro_message", opt.reasoning_budget_intro_message); + llama_params["reasoning_budget_grace_tokens"] = json_value(body, "reasoning_budget_grace_tokens", opt.reasoning_budget_grace_tokens); + } } } diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6ef797ebb473..a30b9a455db7 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -305,6 +305,10 @@ struct server_chat_params { bool enable_thinking = true; int reasoning_budget = -1; std::string reasoning_budget_message; + float reasoning_budget_soft_ratio = -1.0f; + std::string reasoning_budget_soft_message; + std::string reasoning_budget_intro_message; + int reasoning_budget_grace_tokens = 0; std::string media_path; bool force_pure_content = false; }; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 4655b518e21f..7e81d0469fa3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1523,6 +1523,10 @@ struct server_context_impl { /* enable_thinking */ enable_thinking, /* reasoning_budget */ params_base.sampling.reasoning_budget_tokens, /* reasoning_budget_msg */ params_base.sampling.reasoning_budget_message, + /* reasoning_budget_soft_ratio */ params_base.sampling.reasoning_budget_soft_ratio, + /* reasoning_budget_soft_msg */ params_base.sampling.reasoning_budget_soft_message, + /* reasoning_budget_intro_msg */ params_base.sampling.reasoning_budget_intro_message, + /* reasoning_budget_grace_toks */ params_base.sampling.reasoning_budget_grace_tokens, /* media_path */ params_base.media_path, /* force_pure_content */ params_base.force_pure_content_parser }; diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 674d3ba337bc..97b4b97a5d5c 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -384,6 +384,26 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_hard_limits(-1, INT32_MAX) ->set_desc("Number of tokens in the reasoning budget (-1 = disabled)")); + add((new field_num("reasoning_budget_soft_ratio", params.sampling.reasoning_budget_soft_ratio)) + ->set_hard_limits(-1.0f, 1.0f) + ->set_desc("Fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff (<= 0 = disabled)")); + + add((new field_str("reasoning_budget_intro_message")) + ->set_desc("Message forced immediately when the reasoning block starts, announcing the token budget. Use {budget} as a placeholder for the configured reasoning_budget_tokens value") + ->set_handler([&](field_eval_context & ctx, const json & data) { + GGML_ASSERT(ctx.vocab != nullptr); + std::string message = data.at("reasoning_budget_intro_message").get(); + const std::string budget_str = std::to_string(ctx.params.sampling.reasoning_budget_tokens); + for (size_t pos = 0; (pos = message.find("{budget}", pos)) != std::string::npos; pos += budget_str.size()) { + message.replace(pos, 8, budget_str); + } + ctx.params.sampling.reasoning_budget_intro_forced = common_tokenize(ctx.vocab, message, false, true); + })); + + add((new field_num("reasoning_budget_grace_tokens", params.sampling.reasoning_budget_grace_tokens)) + ->set_hard_limits(0, INT32_MAX) + ->set_desc("Once the reasoning budget is exhausted, wait up to this many tokens for a paragraph break before forcing the cutoff (0 = force immediately)")); + add((new field_str("reasoning_budget_start_tag")) ->set_desc("Token string marking the start of the reasoning budget section") ->set_handler([&](field_eval_context & ctx, const json & data) { @@ -413,20 +433,31 @@ std::vector> make_llama_cmpl_schema(const common_params & })); add((new field_str("reasoning_budget_message")) - ->set_desc("Message to prepend to the reasoning budget end tag when forcing it") + ->set_desc("Message forced when the reasoning budget is exhausted. Should include the model's own closing tag (e.g. ) since it is not appended automatically - the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone, so the reasoning block still always closes") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - if (!ctx.params.sampling.reasoning_budget_end.empty()) { - llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front(); - std::string message = json_value(data, "reasoning_budget_message", std::string()); - if (!message.empty()) { - llama_tokens message_tokens = common_tokenize(ctx.vocab, message, false, true); - end_tag.insert(end_tag.begin(), message_tokens.begin(), message_tokens.end()); + std::string message = data.at("reasoning_budget_message").get(); + if (message.empty()) { + // no custom message: fall back to forcing just the first auto-detected + // closing tag alone, so the block still always closes + if (!ctx.params.sampling.reasoning_budget_end.empty()) { + ctx.params.sampling.reasoning_budget_forced = ctx.params.sampling.reasoning_budget_end.front(); } - ctx.params.sampling.reasoning_budget_forced = std::move(end_tag); + } else { + // the message is expected to already include the model's own closing + // tag (see field description) - tokenized as-is, nothing auto-appended + ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message, false, true); } })); + add((new field_str("reasoning_budget_soft_message")) + ->set_desc("Soft-warning message injected partway through the reasoning budget, at the next newline boundary") + ->set_handler([&](field_eval_context & ctx, const json & data) { + GGML_ASSERT(ctx.vocab != nullptr); + std::string message = data.at("reasoning_budget_soft_message").get(); + ctx.params.sampling.reasoning_budget_soft_forced = common_tokenize(ctx.vocab, message, false, true); + })); + add((new field_json("logit_bias")) ->set_desc("Modify the likelihood of specific tokens. Accepts an array of [token, bias] pairs or an object mapping token to bias. Use false as bias to ban a token") ->set_handler([&](field_eval_context & ctx, const json & data) { @@ -566,11 +597,15 @@ task_params eval_llama_cmpl_schema( // debugging { auto budget = params.sampling.reasoning_budget_tokens; - SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n", + SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks, soft_ratio=%.2f, soft_forced=%zu toks, intro_forced=%zu toks, grace_tokens=%d\n", budget, params.sampling.generation_prompt.c_str(), params.sampling.reasoning_budget_start.size(), params.sampling.reasoning_budget_end.size(), - params.sampling.reasoning_budget_forced.size()); + params.sampling.reasoning_budget_forced.size(), + params.sampling.reasoning_budget_soft_ratio, + params.sampling.reasoning_budget_soft_forced.size(), + params.sampling.reasoning_budget_intro_forced.size(), + params.sampling.reasoning_budget_grace_tokens); } return params;