reasoning budget for batching

This commit is contained in:
Concedo 2026-10-07 23:01:16 +08:00
parent d2cd79bb3f
commit ea6477ccaf
3 changed files with 256 additions and 87 deletions

View file

@ -259,18 +259,18 @@ For troubleshooting, include your operating system, hardware, KoboldCpp version,
## Official Resources
The following are all official KoboldCpp resources
[Download KoboldCpp Releases](https://github.com/LostRuins/koboldcpp/releases/latest)
[KoboldAI Discord](https://koboldai.org/discord)
[KoboldCpp Huggingface Repo](https://huggingface.co/koboldcpp)
[KoboldCpp Sourceforge Mirror](https://sourceforge.net/projects/koboldcpp/)
[KoboldAI Subreddit](https://www.reddit.com/r/KoboldAi/)
[KoboldCpp Wiki](https://github.com/LostRuins/koboldcpp/wiki)
[KoboldCpp API Reference](https://lite.koboldai.net/koboldcpp_api)
[KoboldAI Lite Online WebUI](https://lite.koboldai.net)
[KoboldCpp Colab Notebook](https://colab.research.google.com/github/LostRuins/koboldcpp/blob/concedo/colab.ipynb)
[KoboldCpp Public Demo](https://koboldai-koboldcpp-tiefighter.hf.space/)
[KoboldCpp RunPod Cloud Image](https://koboldai.org/runpodcpp)
[KoboldCpp Docker image](https://hub.docker.com/r/koboldai/koboldcpp) (Caution: For experts only)
- [Download KoboldCpp Releases](https://github.com/LostRuins/koboldcpp/releases/latest)
- [KoboldAI Discord](https://koboldai.org/discord)
- [KoboldCpp Huggingface Repo](https://huggingface.co/koboldcpp)
- [KoboldCpp Sourceforge Mirror](https://sourceforge.net/projects/koboldcpp/)
- [KoboldAI Subreddit](https://www.reddit.com/r/KoboldAi/)
- [KoboldCpp Wiki](https://github.com/LostRuins/koboldcpp/wiki)
- [KoboldCpp API Reference](https://lite.koboldai.net/koboldcpp_api)
- [KoboldAI Lite Online WebUI](https://lite.koboldai.net)
- [KoboldCpp Colab Notebook](https://colab.research.google.com/github/LostRuins/koboldcpp/blob/concedo/colab.ipynb)
- [KoboldCpp Public Demo](https://koboldai-koboldcpp-tiefighter.hf.space/)
- [KoboldCpp RunPod Cloud Image](https://koboldai.org/runpodcpp)
- [KoboldCpp Docker image](https://hub.docker.com/r/koboldai/koboldcpp) (Caution: For experts only)
## Third Party Resources
These unofficial community projects may be outdated or unmaintained. Contact their maintainers for support.

View file

@ -54,6 +54,7 @@
#include "common/speculative.h"
#include "common/chat.h"
#include "common/log.h"
#include "common/reasoning-budget.h"
#include "llama-grammar.h"
#include "vendor/stb/stb_image.h"
#include "otherarch/sdcpp/thirdparty/stb_image_resize.h"
@ -2290,6 +2291,72 @@ void sample_guidance(struct llama_context * ctx, struct llama_context * guidance
}
}
struct ReasoningBudgetTokenSequences
{
std::vector<int> start;
std::vector<int> end;
std::vector<int> forced;
};
static ReasoningBudgetTokenSequences get_reasoning_budget_token_sequences()
{
ReasoningBudgetTokenSequences result;
if(file_format != FileFormat::GGUF_GENERIC)
{
return result;
}
std::string start = "<think>";
std::string end = "</think>";
std::string budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n</think>";
size_t expected_start_tokens = 1;
size_t expected_end_tokens = 1;
switch(file_format_meta.model_architecture)
{
case llm_arch::LLM_ARCH_GEMMA4:
start = "<|channel>thought";
end = "<channel|>";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n<channel|>";
expected_start_tokens = 2;
break;
case llm_arch::LLM_ARCH_SEED_OSS:
start = "<seed:think>";
end = "</seed:think>";
budget_exceeded = "\n(Reasoning budget exceeded)\n<seed:cot_budget_reflect>The current thinking budget is 0, so I will directly start answering the question.</seed:cot_budget_reflect>\nTime to respond now.\n</seed:think>";
break;
case llm_arch::LLM_ARCH_COHERE2MOE:
start = "<|START_THINKING|>";
end = "<|END_THINKING|>";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n<|END_THINKING|>";
break;
case llm_arch::LLM_ARCH_MISTRAL3:
start = "[THINK]";
end = "[/THINK]";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n[/THINK]";
break;
case llm_arch::LLM_ARCH_MUSE_GLIMMER:
start = " to=self<|message|>";
end = "<|eom|>";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n<|eom|>";
expected_start_tokens = 3;
break;
default:
break;
}
TokenizeString(start, result.start, file_format, false);
TokenizeString(end, result.end, file_format, false);
TokenizeString(budget_exceeded, result.forced, file_format, false);
if(result.start.size() != expected_start_tokens || result.end.size() != expected_end_tokens || result.forced.empty())
{
result.start.clear();
result.end.clear();
result.forced.clear();
}
return result;
}
static int apply_reasoning_budget(int id, const std::vector<int> & start_think, const std::vector<int> & end_think, std::vector<int> & think_end_phrase_toks, int budget)
{
if(budget<0 || start_think.size()==0 || end_think.size()!=1 || think_end_phrase_toks.size()==0) //start_think can be 1-3 tokens long, end_think is always 1 token
@ -4367,6 +4434,7 @@ struct BatchGenerateRequest
float dry_base = 1.75f;
int dry_allowed_length = 2;
int dry_penalty_last_n = 0;
int reasoning_budget = -1;
bool tool_call_fix = false;
bool allow_eos_token = true;
bool bypass_eos_token = false;
@ -4549,10 +4617,6 @@ static bool batch_inputs_eligible(const generation_inputs & inputs)
}
}
}
if(inputs.reasoning_budget >= 0)
{
return false;
}
return true;
}
@ -4780,6 +4844,125 @@ static llama_sampler * batch_tool_call_init(int32_t n_vocab, bool render_special
});
}
struct BatchReasoningGrammarSampler
{
llama_sampler * reasoning = nullptr;
llama_sampler * grammar = nullptr;
};
static bool batch_reasoning_is_active(const llama_sampler * reasoning)
{
const common_reasoning_budget_state state = common_reasoning_budget_get_state(reasoning);
return state == REASONING_BUDGET_COUNTING ||
state == REASONING_BUDGET_WAITING_UTF8 ||
state == REASONING_BUDGET_FORCING;
}
static const char * batch_reasoning_grammar_name(const llama_sampler * /*smpl*/)
{
return "kcpp-batch-reasoning-grammar";
}
static void batch_reasoning_grammar_accept(llama_sampler * smpl, llama_token token)
{
auto * ctx = (BatchReasoningGrammarSampler *) smpl->ctx;
const bool reasoning_was_active = batch_reasoning_is_active(ctx->reasoning);
llama_sampler_accept(ctx->reasoning, token);
const bool reasoning_is_active = batch_reasoning_is_active(ctx->reasoning);
if(ctx->grammar && !reasoning_was_active && !reasoning_is_active)
{
llama_sampler_accept(ctx->grammar, token);
}
}
static void batch_reasoning_grammar_apply(llama_sampler * smpl, llama_token_data_array * cur_p)
{
auto * ctx = (BatchReasoningGrammarSampler *) smpl->ctx;
llama_sampler_apply(ctx->reasoning, cur_p);
if(ctx->grammar && !batch_reasoning_is_active(ctx->reasoning))
{
llama_sampler_apply(ctx->grammar, cur_p);
}
}
static void batch_reasoning_grammar_reset(llama_sampler * smpl)
{
auto * ctx = (BatchReasoningGrammarSampler *) smpl->ctx;
llama_sampler_reset(ctx->reasoning);
llama_sampler_reset(ctx->grammar);
}
static llama_sampler * batch_reasoning_grammar_clone(const llama_sampler * smpl)
{
const auto * ctx = (const BatchReasoningGrammarSampler *) smpl->ctx;
return llama_sampler_init(smpl->iface, new BatchReasoningGrammarSampler {
llama_sampler_clone(ctx->reasoning),
llama_sampler_clone(ctx->grammar),
});
}
static void batch_reasoning_grammar_free(llama_sampler * smpl)
{
auto * ctx = (BatchReasoningGrammarSampler *) smpl->ctx;
llama_sampler_free(ctx->reasoning);
llama_sampler_free(ctx->grammar);
delete ctx;
}
static llama_sampler_i batch_reasoning_grammar_i = {
/* .name = */ batch_reasoning_grammar_name,
/* .accept = */ batch_reasoning_grammar_accept,
/* .apply = */ batch_reasoning_grammar_apply,
/* .reset = */ batch_reasoning_grammar_reset,
/* .clone = */ batch_reasoning_grammar_clone,
/* .free = */ batch_reasoning_grammar_free,
/* .backend_init = */ nullptr,
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
};
static int batch_find_last_sequence_end(const llama_tokens & tokens, const llama_tokens & sequence)
{
if(sequence.empty() || tokens.size() < sequence.size())
{
return -1;
}
for(size_t offset = tokens.size() - sequence.size() + 1; offset-- > 0;)
{
if(std::equal(sequence.begin(), sequence.end(), tokens.begin() + offset))
{
return (int) (offset + sequence.size() - 1);
}
}
return -1;
}
static void batch_prime_reasoning_sampler(
llama_sampler * reasoning_sampler,
const llama_tokens & prompt_tokens,
const llama_tokens & start_tokens,
const llama_tokens & end_tokens)
{
const int start_end = batch_find_last_sequence_end(prompt_tokens, start_tokens);
const int end_end = batch_find_last_sequence_end(prompt_tokens, end_tokens);
if(start_end < 0 || end_end > start_end)
{
return;
}
const size_t start_pos = (size_t) (start_end + 1 - (int) start_tokens.size());
for(size_t i = start_pos; i < prompt_tokens.size(); ++i)
{
llama_sampler_accept(reasoning_sampler, prompt_tokens[i]);
if(common_reasoning_budget_get_state(reasoning_sampler) == REASONING_BUDGET_FORCING)
{
// Remaining prompt tokens are context, not tokens from the forced sequence.
break;
}
}
}
static llama_sampler * batch_build_sampler(const BatchGenerateRequest & req)
{
llama_sampler_chain_params params = llama_sampler_chain_default_params();
@ -4788,6 +4971,58 @@ static llama_sampler * batch_build_sampler(const BatchGenerateRequest & req)
const llama_vocab * sampler_vocab = llama_model_get_vocab(llama_get_model(llama_ctx_v4));
const int32_t sampler_n_vocab = llama_vocab_n_tokens(sampler_vocab);
llama_sampler * reasoning_sampler = nullptr;
llama_sampler * grammar_sampler = nullptr;
if(req.reasoning_budget >= 0)
{
const ReasoningBudgetTokenSequences reasoning_tokens = get_reasoning_budget_token_sequences();
if(!reasoning_tokens.start.empty() && !reasoning_tokens.end.empty() && !reasoning_tokens.forced.empty())
{
const llama_tokens start_tokens(reasoning_tokens.start.begin(), reasoning_tokens.start.end());
const llama_tokens end_tokens(reasoning_tokens.end.begin(), reasoning_tokens.end.end());
const llama_tokens forced_tokens(reasoning_tokens.forced.begin(), reasoning_tokens.forced.end());
reasoning_sampler = common_reasoning_budget_init(
sampler_vocab,
{start_tokens},
{end_tokens},
forced_tokens,
req.reasoning_budget);
batch_prime_reasoning_sampler(reasoning_sampler, req.prompt_tokens, start_tokens, end_tokens);
}
}
if(!req.grammar.empty())
{
grammar_sampler = llama_sampler_init_grammar(sampler_vocab, req.grammar.c_str(), "root");
if(!grammar_sampler)
{
printf("\nIgnored invalid grammar sampler.");
}
}
if(reasoning_sampler && grammar_sampler)
{
llama_sampler_chain_add(chain, llama_sampler_init(&batch_reasoning_grammar_i, new BatchReasoningGrammarSampler {
reasoning_sampler,
grammar_sampler,
}));
reasoning_sampler = nullptr;
grammar_sampler = nullptr;
}
else
{
if(reasoning_sampler)
{
llama_sampler_chain_add(chain, reasoning_sampler);
reasoning_sampler = nullptr;
}
if(grammar_sampler)
{
llama_sampler_chain_add(chain, grammar_sampler);
grammar_sampler = nullptr;
}
}
if(req.tool_call_fix)
{
// This sampler must not be primed with prompt tokens: its three-token window starts at generation.
@ -4825,19 +5060,6 @@ static llama_sampler * batch_build_sampler(const BatchGenerateRequest & req)
llama_sampler_chain_add(chain, dry_sampler);
}
if(!req.grammar.empty())
{
llama_sampler * grammar_sampler = llama_sampler_init_grammar(sampler_vocab, req.grammar.c_str(), "root");
if(grammar_sampler)
{
llama_sampler_chain_add(chain, grammar_sampler);
}
else
{
printf("\nIgnored invalid grammar sampler.");
}
}
llama_sampler * rep_pen_sampler = batch_rep_pen_init(
req.rep_pen_range,
req.rep_pen,
@ -5229,6 +5451,7 @@ int gpttype_batch_generate_submit(const generation_inputs inputs)
req->dry_base = inputs.dry_base;
req->dry_allowed_length = inputs.dry_allowed_length;
req->dry_penalty_last_n = inputs.dry_penalty_last_n;
req->reasoning_budget = inputs.reasoning_budget;
req->tool_call_fix = inputs.tool_call_fix;
req->allow_eos_token = inputs.allow_eos_token;
req->bypass_eos_token = inputs.bypass_eos_token;
@ -6272,61 +6495,10 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
ApplyPromptFormatAdjustments(addedmemory, kcpp_data->prompt);
//thinking budget handling
std::vector<int> thinking_start_sequence;
std::vector<int> thinking_end_sequence;
std::vector<int> thinking_end_phrase_toksleft;
std::string chat_template = "";
if (file_format == FileFormat::GGUF_GENERIC) {
chat_template = gpttype_get_chat_template();
std::string start = "<think>";
std::string end = "</think>";
std::string budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n</think>";
size_t expected_start_tokens = 1;
size_t expected_end_tokens = 1;
switch (file_format_meta.model_architecture) {
case llm_arch::LLM_ARCH_GEMMA4:
start = "<|channel>thought";
end = "<channel|>";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n<channel|>";
expected_start_tokens = 2;
break;
case llm_arch::LLM_ARCH_SEED_OSS:
start = "<seed:think>";
end = "</seed:think>";
budget_exceeded = "\n(Reasoning budget exceeded)\n<seed:cot_budget_reflect>The current thinking budget is 0, so I will directly start answering the question.</seed:cot_budget_reflect>\nTime to respond now.\n</seed:think>";
break;
case llm_arch::LLM_ARCH_COHERE2MOE:
start = "<|START_THINKING|>";
end = "<|END_THINKING|>";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n<|END_THINKING|>";
break;
case llm_arch::LLM_ARCH_MISTRAL3:
start = "[THINK]";
end = "[/THINK]";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n[/THINK]";
break;
case llm_arch::LLM_ARCH_MUSE_GLIMMER:
start = " to=self<|message|>";
end = "<|eom|>";
budget_exceeded = "\n(Reasoning budget exceeded)\nTime to respond now.\n<|eom|>";
expected_start_tokens = 3;
break;
default:
break;
}
TokenizeString(start, thinking_start_sequence, file_format, false);
TokenizeString(end, thinking_end_sequence, file_format, false);
TokenizeString(budget_exceeded, thinking_end_phrase_toksleft, file_format, false);
if (thinking_start_sequence.size() != expected_start_tokens || thinking_end_sequence.size() != expected_end_tokens)
{
thinking_start_sequence.clear();
thinking_end_sequence.clear();
thinking_end_phrase_toksleft.clear();
}
}
ReasoningBudgetTokenSequences reasoning_tokens = get_reasoning_budget_token_sequences();
std::vector<int> thinking_start_sequence = std::move(reasoning_tokens.start);
std::vector<int> thinking_end_sequence = std::move(reasoning_tokens.end);
std::vector<int> thinking_end_phrase_toksleft = std::move(reasoning_tokens.forced);
bool stream_sse = inputs.stream_sse;

View file

@ -2552,9 +2552,6 @@ def continuous_batching_python_eligible(genparams, api_format):
if genparams.get("sampler_order") and genparams.get("sampler_order") != [6, 0, 1, 3, 4, 2, 5]:
utfprint("Batching disabled due to sampler order",2)
return False
if genparams.get("reasoning_effort"):
utfprint("Batching disabled due to reasoning",2)
return False
return True
def sd_get_info():