From 9427726d94d9053de3169ef498d90c25d47a1cd4 Mon Sep 17 00:00:00 2001 From: Marsh Poulson Date: Mon, 7 Sep 2026 12:48:12 -0400 Subject: [PATCH] glm5.2: preserve input marking in prompt tokenization (fix output stalls on stray message tags) Port of llama.cpp commit f72dba3ad ('glm5.2 fixed stray message tags in input'). - Thread jinja string parts (with is_input metadata) through chat template application; strip BOS/EOS from parts consistently - Add common_tokenize_parts(): user-provided content (is_input) is tokenized with parse_special=false so stray message/special tags in read files cannot be injected as real tokens; template parts keep parse_special=true - OAI-compat path serializes prompt_parts; handle_completions_impl consumes them with input-marking-aware tokenization when present - server-context.cpp hunk adapted to beellama's init_opt signatures - common/arg.cpp.patch included as in the original commit (unapplied reference patch for --list-devices device-id display) --- common/arg.cpp.patch | 16 ++++ common/chat-auto-parser-generator.cpp | 3 +- common/chat.cpp | 124 ++++++++++++++++++++++---- common/chat.h | 14 ++- tools/server/server-common.cpp | 9 ++ tools/server/server-context.cpp | 9 ++ 6 files changed, 155 insertions(+), 20 deletions(-) create mode 100644 common/arg.cpp.patch diff --git a/common/arg.cpp.patch b/common/arg.cpp.patch new file mode 100644 index 000000000000..95ac9cea5dae --- /dev/null +++ b/common/arg.cpp.patch @@ -0,0 +1,16 @@ +--- a/common/arg.cpp ++++ b/common/arg.cpp +@@ -2330,7 +2330,10 @@ + for (auto * dev : devices) { + size_t free, total; + ggml_backend_dev_memory(dev, &free, &total); +- printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); ++ struct ggml_backend_dev_props props; ++ ggml_backend_dev_get_props(dev, &props); ++ printf(" %s: %s [%s] (%zu MiB, %zu MiB free)\n", ++ ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), ++ props.device_id ? props.device_id : "N/A", ++ total / 1024 / 1024, free / 1024 / 1024); + } + exit(0); + } diff --git a/common/chat-auto-parser-generator.cpp b/common/chat-auto-parser-generator.cpp index d7e117e4d98b..dcacf12c4755 100644 --- a/common/chat-auto-parser-generator.cpp +++ b/common/chat-auto-parser-generator.cpp @@ -42,7 +42,7 @@ common_chat_params peg_generator::generate_parser(const common_chat_template & const autoparser & autoparser) { // Create the result structure common_chat_params data; - data.prompt = common_chat_template_direct_apply(tmpl, inputs); + data.prompt = common_chat_template_direct_apply(tmpl, inputs, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.preserved_tokens = autoparser.preserved_tokens; @@ -68,6 +68,7 @@ common_chat_params peg_generator::generate_parser(const common_chat_template & } data.prompt += data.generation_prompt; + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = autoparser.build_parser(inputs, parser_generation_prompt); diff --git a/common/chat.cpp b/common/chat.cpp index 743ecde0a77e..3f294db97f5d 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -933,7 +933,8 @@ static std::string common_chat_template_direct_apply_impl( const autoparser::generation_params & inputs, const std::optional & messages_override = std::nullopt, const std::optional & tools_override = std::nullopt, - const std::optional & additional_context = std::nullopt) { + const std::optional & additional_context = std::nullopt, + std::vector * out_parts = nullptr) { jinja::context ctx(tmpl.source()); // messages_override is already built for this template, do not touch its content parts @@ -979,22 +980,79 @@ static std::string common_chat_template_direct_apply_impl( const jinja::value results = runtime.execute(tmpl.prog); auto parts = jinja::runtime::gather_string_parts(results); + // Preserve the jinja::string parts (with is_input metadata) for the caller + if (out_parts) { + *out_parts = parts->as_string().parts; + } + std::string result = parts->as_string().str(); // TODO: improve this later if (inputs.add_bos && string_starts_with(result, tmpl.bos_token())) { result = result.substr(tmpl.bos_token().size()); + // Also strip BOS from the parts + if (out_parts && !out_parts->empty()) { + auto & first_part = out_parts->front(); + if (first_part.val.find(tmpl.bos_token()) == 0) { + first_part.val = first_part.val.substr(tmpl.bos_token().size()); + if (first_part.val.empty()) { + out_parts->erase(out_parts->begin()); + } + } + } } if (inputs.add_eos && string_ends_with(result, tmpl.eos_token())) { result = result.substr(0, result.size() - tmpl.eos_token().size()); + // Also strip EOS from the parts + if (out_parts && !out_parts->empty()) { + auto & last_part = out_parts->back(); + size_t pos = last_part.val.rfind(tmpl.eos_token()); + if (pos != std::string::npos && pos + tmpl.eos_token().size() == last_part.val.size()) { + last_part.val = last_part.val.substr(0, pos); + if (last_part.val.empty()) { + out_parts->pop_back(); + } + } + } } return result; } std::string common_chat_template_direct_apply( const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - return common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt); + const autoparser::generation_params & inputs, + std::vector * out_parts) { + return common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, out_parts); +} + +// Tokenize prompt parts with input marking awareness. +// Parts marked as is_input (from user content) are tokenized with parse_special=false +// to prevent special token injection. Parts from the template are tokenized with +// parse_special=true so that legitimate special tokens like <|im_start|>, <|im_end|>, etc. +// are properly recognized. +std::vector common_tokenize_parts( + const struct llama_vocab * vocab, + const std::vector & parts, + bool add_special) { + std::vector result; + + bool first = true; + for (const auto & part : parts) { + bool parse_special = !part.is_input; + + if (part.val.empty()) { + continue; + } + + // Only add special (BOS) on the very first non-empty part + bool part_add_special = add_special && first; + first = false; + + auto tokens = common_tokenize(vocab, part.val, part_add_special, parse_special); + result.insert(result.end(), tokens.begin(), tokens.end()); + } + + return result; } static std::string common_chat_template_generation_prompt_impl( @@ -1076,7 +1134,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ data.supports_thinking = true; data.thinking_start_tag = "[THINK]"; data.thinking_end_tags = {"[/THINK]"}; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.preserved_tokens = { @@ -1095,6 +1153,8 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -1165,7 +1225,7 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ const std::string GEN_PREFIX = "<|im_start|>assistant\n"; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -1215,6 +1275,8 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } std::vector tool_call_starts = { "" }; @@ -1355,7 +1417,7 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp adjusted_messages.push_back(msg); } - auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages); + auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages, std::nullopt, std::nullopt, &data.prompt_parts); // Check if we need to replace the return token with end token during // inference and without generation prompt. For more details see: @@ -1400,6 +1462,8 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); @@ -1509,7 +1573,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); if (inputs.add_generation_prompt && string_ends_with(data.prompt, "\n")) { @@ -1518,6 +1582,8 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ // from emitting its proper reasoning token sequence. data.generation_prompt = "<|turn>model\n"; data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } data.message_delimiters = { @@ -1548,6 +1614,8 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); @@ -1673,7 +1741,7 @@ static common_chat_params common_chat_params_init_functionary_v3_2(const common_ const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.preserved_tokens = { @@ -1687,6 +1755,8 @@ static common_chat_params common_chat_params_init_functionary_v3_2(const common_ const auto & msg = inputs.continue_msg; data.generation_prompt = "<|start_header_id|>assistant<|end_header_id|>\n\n>>>all\n" + msg.render_content(); data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -1774,7 +1844,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -1814,6 +1884,8 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -1926,7 +1998,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat adjusted_messages.push_back(msg); } - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -1955,6 +2027,8 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -2017,7 +2091,7 @@ static common_chat_params common_chat_params_init_gigachat_v3( common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = false; @@ -2030,6 +2104,8 @@ static common_chat_params common_chat_params_init_gigachat_v3( const auto & msg = inputs.continue_msg; data.generation_prompt = "assistant<|role_sep|>\n" + msg.render_content(); data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); @@ -2196,7 +2272,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha const std::string TC_SEPARATOR = "\n\n"; data.prompt = common_chat_template_direct_apply_impl( - tmpl, inputs, adjusted_messages, std::nullopt, additional_context); + tmpl, inputs, adjusted_messages, std::nullopt, additional_context, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl( tmpl, inputs, adjusted_messages, std::nullopt, additional_context); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -2225,6 +2301,8 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } bool require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; @@ -2386,7 +2464,7 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -2441,6 +2519,8 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -2587,7 +2667,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t // Stable prefix of the generation prompt that precedes the (forced) <|START_THINKING|> marker. const std::string GEN_PREFIX = TURN_START + CHATBOT; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -2625,6 +2705,8 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -2703,7 +2785,7 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_MINIMAX_M3; data.supports_thinking = true; @@ -2753,6 +2835,8 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -3186,7 +3270,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -3223,6 +3307,8 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -3333,7 +3419,7 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = "<|start|>assistant"; data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -3361,6 +3447,8 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat } data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; @@ -3700,7 +3788,7 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ common_chat_params data; auto params_copy = params; params_copy.reasoning_format = COMMON_REASONING_FORMAT_NONE; - data.prompt = common_chat_template_direct_apply_impl(tmpl, params_copy); + data.prompt = common_chat_template_direct_apply_impl(tmpl, params_copy, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, params); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; auto parser = build_chat_peg_parser([&data](common_chat_peg_builder &p) { diff --git a/common/chat.h b/common/chat.h index cb39e3458f44..d53558451858 100644 --- a/common/chat.h +++ b/common/chat.h @@ -7,6 +7,7 @@ #include "jinja/parser.h" #include "jinja/runtime.h" #include "jinja/caps.h" +#include "jinja/string.h" #include "json.h" @@ -269,6 +270,7 @@ struct common_chat_templates_inputs { struct common_chat_params { common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY; std::string prompt; + std::vector prompt_parts; // preserves is_input metadata for safe tokenization std::string grammar; bool grammar_lazy = false; std::string generation_prompt; @@ -324,6 +326,15 @@ std::string common_chat_templates_source(const struct common_chat_templates * tm struct common_chat_params common_chat_templates_apply(const struct common_chat_templates * tmpls, const struct common_chat_templates_inputs & inputs); +// Tokenize prompt parts with input marking awareness. +// Parts marked as is_input are tokenized with parse_special=false to prevent +// special token injection from user-provided content. +// Parts not marked as is_input (from the template) are tokenized with parse_special=true. +std::vector common_tokenize_parts( + const struct llama_vocab * vocab, + const std::vector & parts, + bool add_special); + // Format single message, while taking into account the position of that message in chat history std::string common_chat_format_single(const struct common_chat_templates * tmpls, const std::vector & past_msg, @@ -365,7 +376,8 @@ std::map common_chat_templates_get_caps(const common_chat_tem std::string common_chat_template_direct_apply( const common_chat_template & tmpl, - const autoparser::generation_params & inputs); + const autoparser::generation_params & inputs, + std::vector * out_parts = nullptr); std::string common_chat_template_generation_prompt( const common_chat_template & tmpl, diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 43b4e1d510d5..07529239d266 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1343,6 +1343,15 @@ json oaicompat_chat_params_parse( llama_params["chat_format"] = static_cast(chat_params.format); llama_params["prompt"] = chat_params.prompt; + // Store prompt parts with is_input metadata for safe tokenization + // (prevents special token injection from user content) + if (!chat_params.prompt_parts.empty()) { + json parts_arr = json::array(); + for (const auto & part : chat_params.prompt_parts) { + parts_arr.push_back({{"is_input", part.is_input}, {"text", part.val}}); + } + llama_params["prompt_parts"] = parts_arr; + } if (!chat_params.grammar.empty()) { llama_params["grammar"] = chat_params.grammar; llama_params["grammar_type"] = std::string("tool_calls"); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 99a9f3a381e1..2f392f3a2fcc 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4977,6 +4977,15 @@ std::unique_ptr server_routes::handle_completions_impl( if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files, ctx_server.init_opt)); + } else if (data.contains("prompt_parts") && !data.at("prompt_parts").empty()) { + // Use input-marking-aware tokenization when prompt_parts is available + // (from chat template application with jinja input marking enabled) + std::vector parts; + for (const auto & p : data.at("prompt_parts")) { + parts.push_back({p.at("is_input").get(), p.at("text").get()}); + } + llama_tokens tokens = common_tokenize_parts(ctx_server.vocab, parts, true); + inputs.push_back(server_tokens(tokens, false)); } else { // Everything else, including multimodal completions. inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt);