mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-27 14:55:56 +02:00
Compare commits
24 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Aaron Teo
|
ad256ded30 | ||
|
Pascal
|
d73c1d6b22 | ||
|
rankaiyx
|
88b47a755c | ||
|
timkhronos
|
3d1c3a8975 | ||
|
Xuan-Son Nguyen
|
0d47ea7427 | ||
|
Piero Evangelista
|
d4d057b6dd | ||
|
Bartowski
|
7657a6c26a | ||
|
Pascal
|
55b7d6c4c7 | ||
|
Xuan-Son Nguyen
|
d2a818231e | ||
|
Eric Hartford
|
af285020e9 | ||
|
timkhronos
|
b1d4c65524 | ||
|
yzyyzyhhh
|
42fc243060 | ||
|
Pascal
|
ff067f76dd | ||
|
Reese Levine
|
7cdd557f76 | ||
|
Nicky Mouha
|
8bb909374d | ||
|
Xuan-Son Nguyen
|
20455a4ad3 | ||
|
Todor Boinovski
|
355303edab | ||
|
hogeheer499-commits
|
c812c543f8 | ||
|
Xuan-Son Nguyen
|
abc348790e | ||
|
Tekin Ertekin
|
2cfc7670ed | ||
|
Adrien Gallouët
|
720d7fa409 | ||
|
Yongmin Yoo 유용민
|
fb92d8f187 | ||
|
Aldehir Rojas
|
910196f6b3 | ||
|
helanfxz
|
d67c0b4107 |
@@ -0,0 +1,90 @@
|
||||
name: CI (wasm)
|
||||
|
||||
on:
|
||||
workflow_dispatch: # allows manual triggering
|
||||
push:
|
||||
branches:
|
||||
- master
|
||||
paths: [
|
||||
'.github/workflows/build-wasm.yml',
|
||||
'**/CMakeLists.txt',
|
||||
'**/.cmake',
|
||||
'**/*.h',
|
||||
'**/*.hpp',
|
||||
'**/*.c',
|
||||
'**/*.cpp',
|
||||
'**/*.wgsl',
|
||||
'**/*.tmpl',
|
||||
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
|
||||
]
|
||||
|
||||
pull_request:
|
||||
types: [opened, synchronize, reopened]
|
||||
paths: [
|
||||
'.github/workflows/build-wasm.yml',
|
||||
'**/CMakeLists.txt',
|
||||
'**/.cmake',
|
||||
'**/*.h',
|
||||
'**/*.hpp',
|
||||
'**/*.c',
|
||||
'**/*.cpp',
|
||||
'**/*.wgsl',
|
||||
'**/*.tmpl',
|
||||
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
|
||||
]
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
|
||||
cancel-in-progress: true
|
||||
|
||||
env:
|
||||
GGML_NLOOP: 3
|
||||
GGML_N_THREADS: 1
|
||||
LLAMA_ARG_LOG_COLORS: 1
|
||||
LLAMA_ARG_LOG_PREFIX: 1
|
||||
LLAMA_ARG_LOG_TIMESTAMPS: 1
|
||||
|
||||
jobs:
|
||||
ubuntu-webgpu:
|
||||
runs-on: ubuntu-24.04-arm
|
||||
|
||||
steps:
|
||||
- name: Clone
|
||||
id: checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: ccache
|
||||
uses: ggml-org/ccache-action@v1.2.21
|
||||
with:
|
||||
key: webgpu-ubuntu-24.04-arm-wasm
|
||||
evict-old-files: 1d
|
||||
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
|
||||
|
||||
- name: Install Emscripten
|
||||
run: |
|
||||
git clone https://github.com/emscripten-core/emsdk.git
|
||||
cd emsdk
|
||||
./emsdk install latest
|
||||
./emsdk activate latest
|
||||
|
||||
- name: Fetch emdawnwebgpu
|
||||
run: |
|
||||
DAWN_TAG="v20260317.182325"
|
||||
EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
|
||||
echo "Downloading ${EMDAWN_PKG}"
|
||||
curl -L -o emdawn.zip \
|
||||
"https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
|
||||
unzip emdawn.zip
|
||||
|
||||
- name: Build WASM WebGPU
|
||||
run: |
|
||||
source emsdk/emsdk_env.sh
|
||||
emcmake cmake -B build-wasm \
|
||||
-G "Ninja" \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DGGML_WEBGPU=ON \
|
||||
-DGGML_OPENMP=OFF \
|
||||
-DLLAMA_OPENSSL=OFF \
|
||||
-DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
|
||||
|
||||
time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
|
||||
@@ -13,7 +13,9 @@ on:
|
||||
'**/*.hpp',
|
||||
'**/*.c',
|
||||
'**/*.cpp',
|
||||
'**/*.wgsl'
|
||||
'**/*.wgsl',
|
||||
'**/*.tmpl',
|
||||
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
|
||||
]
|
||||
|
||||
pull_request:
|
||||
@@ -151,46 +153,3 @@ jobs:
|
||||
# This is using llvmpipe and runs slower than other backends
|
||||
# test-backend-ops is too slow on llvmpipe, skip it
|
||||
ctest -L main -E test-backend-ops --verbose --timeout 900
|
||||
|
||||
ubuntu-wasm:
|
||||
runs-on: ubuntu-24.04-arm
|
||||
|
||||
steps:
|
||||
- name: Clone
|
||||
id: checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: ccache
|
||||
uses: ggml-org/ccache-action@v1.2.21
|
||||
with:
|
||||
key: webgpu-ubuntu-24.04-arm-wasm
|
||||
evict-old-files: 1d
|
||||
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
|
||||
|
||||
- name: Install Emscripten
|
||||
run: |
|
||||
git clone https://github.com/emscripten-core/emsdk.git
|
||||
cd emsdk
|
||||
./emsdk install latest
|
||||
./emsdk activate latest
|
||||
|
||||
- name: Fetch emdawnwebgpu
|
||||
run: |
|
||||
DAWN_TAG="v20260317.182325"
|
||||
EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
|
||||
echo "Downloading ${EMDAWN_PKG}"
|
||||
curl -L -o emdawn.zip \
|
||||
"https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
|
||||
unzip emdawn.zip
|
||||
|
||||
- name: Build WASM WebGPU
|
||||
run: |
|
||||
source emsdk/emsdk_env.sh
|
||||
emcmake cmake -B build-wasm \
|
||||
-G "Ninja" \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DGGML_WEBGPU=ON \
|
||||
-DLLAMA_OPENSSL=OFF \
|
||||
-DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
|
||||
|
||||
time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
|
||||
|
||||
@@ -84,6 +84,14 @@ else()
|
||||
set(LLAMA_TOOLS_INSTALL_DEFAULT ${LLAMA_STANDALONE})
|
||||
endif()
|
||||
|
||||
# subprocess spawning isn't a supported/sandbox-friendly operation on mobile OSes or in WASM
|
||||
if (CMAKE_SYSTEM_NAME STREQUAL "iOS" OR CMAKE_SYSTEM_NAME STREQUAL "Android" OR ANDROID
|
||||
OR CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN)
|
||||
set(LLAMA_SUBPROCESS_DEFAULT OFF)
|
||||
else()
|
||||
set(LLAMA_SUBPROCESS_DEFAULT ON)
|
||||
endif()
|
||||
|
||||
#
|
||||
# option list
|
||||
#
|
||||
@@ -117,6 +125,7 @@ option(LLAMA_TESTS_INSTALL "llama: install tests" ON)
|
||||
|
||||
# 3rd party libs
|
||||
option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON)
|
||||
option(LLAMA_SUBPROCESS "llama-common: use subprocess, required by server tools and server router mode" ${LLAMA_SUBPROCESS_DEFAULT})
|
||||
option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF)
|
||||
|
||||
|
||||
|
||||
@@ -100,6 +100,10 @@ add_library(${TARGET}
|
||||
sampling.h
|
||||
speculative.cpp
|
||||
speculative.h
|
||||
subproc.cpp
|
||||
subproc.h
|
||||
trie.cpp
|
||||
trie.h
|
||||
unicode.cpp
|
||||
unicode.h
|
||||
jinja/lexer.cpp
|
||||
@@ -125,6 +129,10 @@ set_target_properties(${TARGET} PROPERTIES
|
||||
target_include_directories(${TARGET} PUBLIC . ../vendor)
|
||||
target_compile_features (${TARGET} PUBLIC cxx_std_17)
|
||||
|
||||
if (LLAMA_SUBPROCESS)
|
||||
target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)
|
||||
endif()
|
||||
|
||||
if (BUILD_SHARED_LIBS)
|
||||
set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
|
||||
|
||||
+27
-8
@@ -850,8 +850,9 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
|
||||
params.kv_overrides.back().key[0] = 0;
|
||||
}
|
||||
|
||||
if (!params.server_tools.empty() && !params.cors_origins_explicit) {
|
||||
LOG_WRN("server tools are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
|
||||
const bool mcp_enabled = !params.mcp_servers_config.empty() || !params.mcp_servers_json.empty();
|
||||
if ((!params.server_tools.empty() || mcp_enabled) && !params.cors_origins_explicit) {
|
||||
LOG_WRN("server tools or MCP servers are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
|
||||
params.cors_origins = "localhost";
|
||||
}
|
||||
|
||||
@@ -2507,7 +2508,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
}
|
||||
add_opt(common_arg(
|
||||
{"--mlock"},
|
||||
"DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing",
|
||||
"DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing",
|
||||
[](common_params & params) {
|
||||
LOG_WRN("DEPRECATED: --mlock is deprecated. use --load-mode mlock instead\n");
|
||||
params.load_mode = LLAMA_LOAD_MODE_MLOCK;
|
||||
@@ -2536,13 +2537,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
"model loading mode (default: mmap)\n"
|
||||
"- none: no special loading mode\n"
|
||||
"- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n"
|
||||
"- mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
|
||||
"- mlock: force system to keep model in RAM rather than swapping or compressing\n"
|
||||
"- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
|
||||
"- dio: use DirectIO if available\n",
|
||||
[](common_params & params, const std::string & value) {
|
||||
/**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; }
|
||||
else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; }
|
||||
else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; }
|
||||
else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; }
|
||||
/**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; }
|
||||
else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; }
|
||||
else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; }
|
||||
else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; }
|
||||
else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; }
|
||||
else { throw std::invalid_argument("invalid value"); }
|
||||
}
|
||||
).set_env("LLAMA_ARG_LOAD_MODE"));
|
||||
@@ -3261,6 +3264,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
params.server_tools = parse_csv_row(value);
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS"));
|
||||
add_opt(common_arg(
|
||||
{"--mcp-servers-config"}, "PATH",
|
||||
"experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
|
||||
"note: for security reasons, this will limit --cors-origins to localhost by default",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.mcp_servers_config = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_CONFIG"));
|
||||
add_opt(common_arg(
|
||||
{"--mcp-servers-json"}, "JSON",
|
||||
"experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
|
||||
"note: for security reasons, this will limit --cors-origins to localhost by default",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.mcp_servers_json = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_JSON"));
|
||||
add_opt(common_arg(
|
||||
{"-ag", "--agent"},
|
||||
{"-no-ag", "--no-agent"},
|
||||
|
||||
+14
-8
@@ -1024,7 +1024,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_tag = "[/THINK]";
|
||||
data.thinking_end_tags = {"[/THINK]"};
|
||||
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
|
||||
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
@@ -1150,6 +1150,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
|
||||
data.thinking_start_tag = "<|channel|>analysis<|message|>";
|
||||
data.thinking_end_tags = {"<|end|>"};
|
||||
|
||||
// These special tokens are required to parse properly, so we include them
|
||||
// even if parse_tool_calls is false.
|
||||
data.preserved_tokens = {
|
||||
@@ -1294,7 +1297,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = "<|channel>thought";
|
||||
data.thinking_end_tag = "<channel|>";
|
||||
data.thinking_end_tags = {"<channel|>"};
|
||||
|
||||
data.preserved_tokens = {
|
||||
"<|channel>",
|
||||
@@ -1569,7 +1572,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp
|
||||
const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>";
|
||||
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tag = THINK_END;
|
||||
data.thinking_end_tags = {THINK_END};
|
||||
|
||||
if (inputs.has_continuation()) {
|
||||
const auto & msg = inputs.continue_msg;
|
||||
@@ -1703,7 +1706,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat
|
||||
}
|
||||
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tag = THINK_END;
|
||||
data.thinking_end_tags = {THINK_END};
|
||||
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
|
||||
auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
|
||||
@@ -1943,7 +1946,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = "<think>";
|
||||
data.thinking_end_tag = "</think>";
|
||||
data.thinking_end_tags = {"</think>"};
|
||||
data.preserved_tokens = {
|
||||
"｜DSML｜",
|
||||
"<think>",
|
||||
@@ -2160,7 +2163,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tag = THINK_END;
|
||||
data.thinking_end_tags = {THINK_END};
|
||||
data.preserved_tokens = {
|
||||
TURN_START, TURN_END, CHATBOT, USER, SYSTEM,
|
||||
THINK_START, THINK_END,
|
||||
@@ -2510,7 +2513,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
|
||||
};
|
||||
|
||||
data.thinking_start_tag = "<think>";
|
||||
data.thinking_end_tag = "</think>";
|
||||
data.thinking_end_tags = {"</think>"};
|
||||
|
||||
data.message_delimiters = {
|
||||
{ COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" },
|
||||
@@ -2866,7 +2869,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
|
||||
auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE;
|
||||
if (auto_params.supports_thinking) {
|
||||
auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start);
|
||||
auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end);
|
||||
auto end_tag = trim_whitespace(autoparser.reasoning.end);
|
||||
if (!end_tag.empty()) {
|
||||
auto_params.thinking_end_tags = {std::move(end_tag)};
|
||||
}
|
||||
}
|
||||
common_peg_arena arena;
|
||||
arena.load(auto_params.parser);
|
||||
|
||||
+1
-1
@@ -274,7 +274,7 @@ struct common_chat_params {
|
||||
std::string generation_prompt;
|
||||
bool supports_thinking = false;
|
||||
std::string thinking_start_tag; // e.g., "<think>"
|
||||
std::string thinking_end_tag; // e.g., "</think>"
|
||||
std::vector<std::string> thinking_end_tags; // e.g., "</think>"
|
||||
std::vector<common_grammar_trigger> grammar_triggers;
|
||||
std::vector<std::string> preserved_tokens;
|
||||
std::vector<std::string> additional_stops;
|
||||
|
||||
@@ -1249,7 +1249,6 @@ common_init_result::common_init_result(common_params & params, bool model_only)
|
||||
lora.reset(llama_adapter_lora_init(model, la.path.c_str()));
|
||||
if (lora == nullptr) {
|
||||
COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str());
|
||||
pimpl->model.reset(model);
|
||||
return;
|
||||
}
|
||||
|
||||
|
||||
+10
-6
@@ -284,12 +284,12 @@ struct common_params_sampling {
|
||||
|
||||
// reasoning budget sampler parameters
|
||||
// these are populated by the server/CLI based on chat template params
|
||||
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
|
||||
std::vector<llama_token> reasoning_budget_start; // start tag token sequence
|
||||
std::vector<llama_token> reasoning_budget_end; // end tag token sequence
|
||||
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + end tag)
|
||||
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
|
||||
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
|
||||
std::vector<llama_token> reasoning_budget_start; // start tag token sequence
|
||||
std::vector<llama_tokens> reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence
|
||||
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + first end tag)
|
||||
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
|
||||
|
||||
bool backend_sampling = false;
|
||||
|
||||
@@ -668,6 +668,10 @@ struct common_params {
|
||||
// enable built-in tools
|
||||
std::vector<std::string> server_tools;
|
||||
|
||||
// MCP server configs (Cursor-compatible JSON)
|
||||
std::string mcp_servers_config; // path to JSON file with MCP server definitions
|
||||
std::string mcp_servers_json; // inline JSON with MCP server definitions
|
||||
|
||||
// router server configs
|
||||
std::string models_dir = ""; // directory containing models for the router server
|
||||
std::string models_preset = ""; // directory containing model presets for the router server
|
||||
|
||||
+5
-153
@@ -3,10 +3,10 @@
|
||||
#include "common.h"
|
||||
#include "json-schema-to-grammar.h"
|
||||
#include "log.h"
|
||||
#include "trie.h"
|
||||
#include "unicode.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <deque>
|
||||
#include <initializer_list>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
@@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) {
|
||||
return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F');
|
||||
}
|
||||
|
||||
// Trie for matching multiple literals.
|
||||
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
|
||||
struct trie {
|
||||
struct node {
|
||||
std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
|
||||
bool is_word;
|
||||
};
|
||||
|
||||
std::vector<node> nodes;
|
||||
|
||||
trie(const std::vector<std::string> & words) {
|
||||
create_node(); // root node
|
||||
for (const auto & w : words) {
|
||||
insert(w);
|
||||
}
|
||||
}
|
||||
|
||||
enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
|
||||
|
||||
// Check if a delimiter starts at the given position
|
||||
match_result check_at(std::string_view sv, size_t start_pos) const {
|
||||
size_t current = 0; // Start at root
|
||||
size_t pos = start_pos;
|
||||
|
||||
// LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
|
||||
|
||||
while (pos < sv.size()) {
|
||||
auto result = common_parse_utf8_codepoint(sv, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
auto it = nodes[current].children.find(result.codepoint);
|
||||
if (it == nodes[current].children.end()) {
|
||||
// Can't continue matching
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
current = it->second;
|
||||
pos += result.bytes_consumed;
|
||||
|
||||
// Check if we've matched a complete word
|
||||
if (nodes[current].is_word) {
|
||||
return match_result{match_result::COMPLETE_MATCH};
|
||||
}
|
||||
}
|
||||
|
||||
// Reached end of input while still in the trie (not at root)
|
||||
if (current != 0) {
|
||||
// We're in the middle of a potential match
|
||||
return match_result{match_result::PARTIAL_MATCH};
|
||||
}
|
||||
|
||||
// Reached end at root (no match)
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
private:
|
||||
size_t create_node() {
|
||||
size_t index = nodes.size();
|
||||
nodes.emplace_back();
|
||||
return index;
|
||||
}
|
||||
|
||||
void insert(const std::string & word) {
|
||||
size_t current = 0;
|
||||
size_t pos = 0;
|
||||
while (pos < word.length()) {
|
||||
auto result = common_parse_utf8_codepoint(word, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
uint32_t ch = result.codepoint;
|
||||
pos += result.bytes_consumed;
|
||||
|
||||
auto it = nodes[current].children.find(ch);
|
||||
if (it == nodes[current].children.end()) {
|
||||
size_t child = create_node();
|
||||
nodes[current].children[ch] = child;
|
||||
current = child;
|
||||
} else {
|
||||
current = it->second;
|
||||
}
|
||||
}
|
||||
nodes[current].is_word = true;
|
||||
}
|
||||
};
|
||||
|
||||
// Aho-Corasick automaton
|
||||
struct aho_corasick {
|
||||
trie t;
|
||||
std::vector<size_t> fail; // failure links
|
||||
std::vector<size_t> order; // states in BFS order
|
||||
std::vector<bool> terminal; // match states (directly or via a suffix link)
|
||||
std::set<uint32_t> alphabet; // every character with a transition
|
||||
|
||||
aho_corasick(const std::vector<std::string> & strings) : t(strings) {
|
||||
const auto & nodes = t.nodes;
|
||||
const size_t n = nodes.size();
|
||||
|
||||
fail.assign(n, 0);
|
||||
order.reserve(n);
|
||||
|
||||
std::deque<size_t> queue{ 0 };
|
||||
while (!queue.empty()) {
|
||||
size_t u = queue.front();
|
||||
queue.pop_front();
|
||||
order.push_back(u);
|
||||
for (const auto & [ch, v] : nodes[u].children) {
|
||||
if (u != 0) {
|
||||
size_t f = fail[u];
|
||||
while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
|
||||
f = fail[f];
|
||||
}
|
||||
auto it = nodes[f].children.find(ch);
|
||||
fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
|
||||
}
|
||||
queue.push_back(v);
|
||||
}
|
||||
}
|
||||
|
||||
terminal.assign(n, false);
|
||||
for (size_t u : order) {
|
||||
terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]);
|
||||
}
|
||||
|
||||
for (const auto & node : nodes) {
|
||||
for (const auto & [ch, v] : node.children) {
|
||||
alphabet.insert(ch);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
size_t num_states() const { return t.nodes.size(); }
|
||||
bool is_terminal(size_t s) const { return terminal[s]; }
|
||||
|
||||
// follow failure links until a transition on `ch` exists.
|
||||
size_t next(size_t state, uint32_t ch) const {
|
||||
const auto & nodes = t.nodes;
|
||||
while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
|
||||
state = fail[state];
|
||||
}
|
||||
auto it = nodes[state].children.find(ch);
|
||||
return it != nodes[state].children.end() ? it->second : 0;
|
||||
}
|
||||
};
|
||||
|
||||
static std::pair<uint32_t, size_t> parse_hex_escape(const std::string & str, size_t pos, int hex_count) {
|
||||
if (pos + hex_count > str.length()) {
|
||||
return {0, 0};
|
||||
@@ -797,7 +649,7 @@ struct parser_executor {
|
||||
}
|
||||
|
||||
common_peg_parse_result operator()(const common_peg_until_parser & p) const {
|
||||
trie matcher(p.delimiters);
|
||||
common_trie matcher(p.delimiters);
|
||||
|
||||
// Scan input and check for delimiters
|
||||
size_t pos = start_pos;
|
||||
@@ -824,12 +676,12 @@ struct parser_executor {
|
||||
// Check if a delimiter starts at this position
|
||||
auto match = matcher.check_at(ctx.input, pos);
|
||||
|
||||
if (match == trie::COMPLETE_MATCH) {
|
||||
if (match == common_trie::COMPLETE_MATCH) {
|
||||
// Found a complete delimiter, return everything before it
|
||||
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
|
||||
}
|
||||
|
||||
if (match == trie::PARTIAL_MATCH) {
|
||||
if (match == common_trie::PARTIAL_MATCH) {
|
||||
// Found a partial match extending to end of input, return everything before it
|
||||
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
|
||||
}
|
||||
@@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar(
|
||||
const std::map<size_t, std::vector<uint32_t>> &,
|
||||
const std::vector<uint32_t> &,
|
||||
const std::function<std::string(size_t)> &)> & build_rule) {
|
||||
aho_corasick ac(strings);
|
||||
common_aho_corasick ac(strings);
|
||||
|
||||
auto state_name = [&](size_t s) -> std::string {
|
||||
if (s == 0) {
|
||||
|
||||
@@ -330,6 +330,10 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co
|
||||
}
|
||||
}
|
||||
|
||||
if (preset.name == COMMON_PRESET_DEFAULT_NAME && preset.options.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (preset.name == "*") {
|
||||
// handle global preset
|
||||
global = preset;
|
||||
|
||||
+77
-39
@@ -1,39 +1,52 @@
|
||||
#include "reasoning-budget.h"
|
||||
#include "common.h"
|
||||
#include "trie.h"
|
||||
#include "unicode.h"
|
||||
|
||||
#include "log.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
struct token_matcher {
|
||||
std::vector<llama_token> tokens;
|
||||
size_t pos = 0;
|
||||
std::vector<llama_tokens> seqs;
|
||||
common_aho_corasick ac;
|
||||
size_t state = 0;
|
||||
|
||||
bool advance(llama_token token) {
|
||||
if (tokens.empty()) {
|
||||
return false;
|
||||
}
|
||||
token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
|
||||
|
||||
if (token == tokens[pos]) {
|
||||
pos++;
|
||||
if (pos >= tokens.size()) {
|
||||
pos = 0;
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
pos = 0;
|
||||
if (token == tokens[0]) {
|
||||
pos = 1;
|
||||
static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
|
||||
std::vector<llama_tokens> res;
|
||||
for (const auto & seq : seqs) {
|
||||
if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
|
||||
res.push_back(seq);
|
||||
}
|
||||
}
|
||||
return false;
|
||||
return res;
|
||||
}
|
||||
|
||||
void reset() { pos = 0; }
|
||||
static common_trie build_trie(const std::vector<llama_tokens> & seqs) {
|
||||
common_trie t;
|
||||
for (const auto & seq : seqs) {
|
||||
t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));
|
||||
}
|
||||
return t;
|
||||
}
|
||||
|
||||
// returns the index into seqs of the longest sequence ending at this token, or -1
|
||||
int32_t advance(llama_token token) {
|
||||
state = ac.next(state, (uint32_t) token);
|
||||
const int32_t p = ac.match_pattern(state);
|
||||
if (p >= 0) {
|
||||
state = 0;
|
||||
}
|
||||
return p;
|
||||
}
|
||||
|
||||
void reset() { state = 0; }
|
||||
};
|
||||
|
||||
struct common_reasoning_budget_ctx {
|
||||
@@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx {
|
||||
|
||||
token_matcher start_matcher;
|
||||
token_matcher end_matcher;
|
||||
std::vector<llama_token> forced_tokens;
|
||||
llama_tokens forced_tokens;
|
||||
|
||||
int32_t budget; // maximum tokens in reasoning block
|
||||
int32_t remaining; // tokens remaining in budget
|
||||
@@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx {
|
||||
|
||||
// for forcing
|
||||
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
|
||||
};
|
||||
|
||||
static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) {
|
||||
@@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
|
||||
switch (ctx->state) {
|
||||
case REASONING_BUDGET_IDLE:
|
||||
{
|
||||
if (ctx->start_matcher.advance(token)) {
|
||||
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);
|
||||
@@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
|
||||
case REASONING_BUDGET_COUNTING:
|
||||
case REASONING_BUDGET_WAITING_UTF8:
|
||||
{
|
||||
if (ctx->end_matcher.advance(token)) {
|
||||
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;
|
||||
}
|
||||
@@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
|
||||
break;
|
||||
}
|
||||
case REASONING_BUDGET_FORCING:
|
||||
{
|
||||
// track the end sequence within forced_tokens so it is also reported on DONE
|
||||
const int32_t match = ctx->end_matcher.advance(token);
|
||||
ctx->force_pos++;
|
||||
if (ctx->force_pos >= ctx->forced_tokens.size()) {
|
||||
ctx->state = REASONING_BUDGET_DONE;
|
||||
ctx->end_match = match;
|
||||
COM_TRC("%s", "forced sequence complete, done\n");
|
||||
}
|
||||
break;
|
||||
}
|
||||
case REASONING_BUDGET_DONE:
|
||||
// Re-arm on a new start tag: some models emit multiple <think> blocks
|
||||
// per response, and each should get a fresh budget window.
|
||||
if (ctx->start_matcher.advance(token)) {
|
||||
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) {
|
||||
@@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
|
||||
ctx->start_matcher.reset();
|
||||
ctx->end_matcher.reset();
|
||||
ctx->force_pos = 0;
|
||||
ctx->end_match = -1;
|
||||
}
|
||||
|
||||
static struct llama_sampler * common_reasoning_budget_init_state(
|
||||
const struct llama_vocab * vocab, const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens, const std::vector<llama_token> & forced_tokens,
|
||||
const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,
|
||||
int32_t budget, common_reasoning_budget_state initial_state);
|
||||
|
||||
static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);
|
||||
@@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s
|
||||
}
|
||||
|
||||
static struct llama_sampler * common_reasoning_budget_init_state(
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const llama_tokens & forced_tokens,
|
||||
int32_t budget,
|
||||
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;
|
||||
@@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state(
|
||||
/* .iface = */ &common_reasoning_budget_i,
|
||||
/* .ctx = */ new common_reasoning_budget_ctx {
|
||||
/* .vocab = */ vocab,
|
||||
/* .start_matcher = */ { start_tokens, 0 },
|
||||
/* .end_matcher = */ { end_tokens, 0 },
|
||||
/* .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,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
struct llama_sampler * common_reasoning_budget_init(
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state);
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const llama_tokens & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state);
|
||||
}
|
||||
|
||||
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
|
||||
@@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla
|
||||
return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;
|
||||
}
|
||||
|
||||
const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) {
|
||||
if (!smpl) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
|
||||
if (ctx->end_match < 0) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
return &ctx->end_matcher.seqs[ctx->end_match];
|
||||
}
|
||||
|
||||
bool common_reasoning_budget_force(struct llama_sampler * smpl) {
|
||||
if (!smpl) {
|
||||
return false;
|
||||
|
||||
+16
-10
@@ -2,6 +2,8 @@
|
||||
|
||||
#include "llama.h"
|
||||
|
||||
#include "common.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
@@ -17,30 +19,34 @@ enum common_reasoning_budget_state {
|
||||
// reasoning block (e.g. between <think> and </think>).
|
||||
//
|
||||
// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE
|
||||
// IDLE: passthrough, watching for start_tokens sequence
|
||||
// COUNTING: counting down remaining tokens, watching for natural end_tokens
|
||||
// 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
|
||||
//
|
||||
// Parameters:
|
||||
// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr)
|
||||
// start_tokens - token sequence that activates counting
|
||||
// end_tokens - token sequence for natural deactivation
|
||||
// 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
|
||||
//
|
||||
struct llama_sampler * common_reasoning_budget_init(
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const llama_tokens & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
|
||||
|
||||
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl);
|
||||
|
||||
// The end sequence that transitioned the sampler to DONE, or nullptr if none
|
||||
// was recorded. Cleared when a new start sequence re-arms the sampler.
|
||||
const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl);
|
||||
|
||||
// Manually transition the reasoning budget sampler into the FORCING state.
|
||||
// Returns true if the transition occurred.
|
||||
bool common_reasoning_budget_force(struct llama_sampler * smpl);
|
||||
|
||||
+12
-1
@@ -299,7 +299,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st
|
||||
if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) {
|
||||
rbudget = common_reasoning_budget_init(
|
||||
vocab,
|
||||
params.reasoning_budget_start,
|
||||
{params.reasoning_budget_start},
|
||||
params.reasoning_budget_end,
|
||||
params.reasoning_budget_forced,
|
||||
params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens);
|
||||
@@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo
|
||||
|
||||
if (gsmpl->rbudget && is_generated) {
|
||||
llama_sampler_accept(gsmpl->rbudget, token);
|
||||
|
||||
// if done, replay end sequence which may contain a grammar trigger
|
||||
const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE;
|
||||
if (gsmpl->grmr && !accept_grammar && is_done) {
|
||||
const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget);
|
||||
if (end_seq) {
|
||||
for (const llama_token end_token : *end_seq) {
|
||||
llama_sampler_accept(gsmpl->grmr, end_token);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (gsmpl->grmr && accept_grammar) {
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
#include "subproc.h"
|
||||
|
||||
bool common_subproc::is_supported() {
|
||||
#ifdef LLAMA_SUBPROCESS
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef LLAMA_SUBPROCESS
|
||||
|
||||
static std::vector<char *> to_cstr_vec(const std::vector<std::string> & v) {
|
||||
std::vector<char *> r;
|
||||
r.reserve(v.size() + 1);
|
||||
for (const auto & s : v) {
|
||||
r.push_back(const_cast<char *>(s.c_str()));
|
||||
}
|
||||
r.push_back(nullptr);
|
||||
return r;
|
||||
}
|
||||
|
||||
common_subproc::~common_subproc() {
|
||||
if (is_created) {
|
||||
subprocess_destroy(&proc);
|
||||
is_created = false;
|
||||
}
|
||||
}
|
||||
|
||||
bool common_subproc::create(
|
||||
const std::vector<std::string> & args,
|
||||
int options,
|
||||
const std::vector<std::string> & env,
|
||||
const char * cwd) {
|
||||
auto argv = to_cstr_vec(args);
|
||||
|
||||
int result;
|
||||
if (env.empty() && cwd == nullptr) {
|
||||
result = subprocess_create(argv.data(), options, &proc);
|
||||
} else {
|
||||
auto envp = to_cstr_vec(env);
|
||||
result = subprocess_create_ex(argv.data(), options, env.empty() ? nullptr : envp.data(), cwd, &proc);
|
||||
}
|
||||
|
||||
is_created = result == 0;
|
||||
return is_created;
|
||||
}
|
||||
|
||||
bool common_subproc::has_handle() const {
|
||||
if (!is_created) {
|
||||
return false;
|
||||
}
|
||||
#if defined(_WIN32)
|
||||
return proc.hProcess != nullptr;
|
||||
#else
|
||||
return proc.child > 0;
|
||||
#endif
|
||||
}
|
||||
|
||||
bool common_subproc::alive() {
|
||||
return is_created && subprocess_alive(&proc);
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdin_file() {
|
||||
return is_created ? subprocess_stdin(&proc) : nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdout_file() {
|
||||
return is_created ? subprocess_stdout(&proc) : nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stderr_file() {
|
||||
return is_created ? subprocess_stderr(&proc) : nullptr;
|
||||
}
|
||||
|
||||
void common_subproc::close_stdin() {
|
||||
if (is_created && proc.stdin_file) {
|
||||
fclose(proc.stdin_file);
|
||||
proc.stdin_file = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
void common_subproc::terminate() {
|
||||
if (has_handle()) {
|
||||
subprocess_terminate(&proc);
|
||||
}
|
||||
}
|
||||
|
||||
int common_subproc::join() {
|
||||
int exit_code = -1;
|
||||
if (is_created) {
|
||||
subprocess_join(&proc, &exit_code);
|
||||
subprocess_destroy(&proc);
|
||||
is_created = false;
|
||||
}
|
||||
return exit_code;
|
||||
}
|
||||
|
||||
#else // !LLAMA_SUBPROCESS
|
||||
|
||||
common_subproc::~common_subproc() = default;
|
||||
|
||||
bool common_subproc::create(
|
||||
const std::vector<std::string> &,
|
||||
int,
|
||||
const std::vector<std::string> &,
|
||||
const char *) {
|
||||
(void)(proc);
|
||||
(void)(is_created);
|
||||
return false;
|
||||
}
|
||||
|
||||
bool common_subproc::has_handle() const {
|
||||
return false;
|
||||
}
|
||||
|
||||
bool common_subproc::alive() {
|
||||
return false;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdin_file() {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdout_file() {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stderr_file() {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void common_subproc::close_stdin() {
|
||||
}
|
||||
|
||||
void common_subproc::terminate() {
|
||||
}
|
||||
|
||||
int common_subproc::join() {
|
||||
return -1;
|
||||
}
|
||||
|
||||
#endif // LLAMA_SUBPROCESS
|
||||
@@ -0,0 +1,59 @@
|
||||
#pragma once
|
||||
|
||||
#include <atomic>
|
||||
#include <cstdio>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#ifdef LLAMA_SUBPROCESS
|
||||
#include <sheredom/subprocess.h>
|
||||
#else
|
||||
// dummy values to allow compilation when subprocess is disabled
|
||||
struct subprocess_s {};
|
||||
static constexpr int subprocess_option_no_window = 0;
|
||||
static constexpr int subprocess_option_combined_stdout_stderr = 0;
|
||||
static constexpr int subprocess_option_inherit_environment = 0;
|
||||
static constexpr int subprocess_option_search_user_path = 0;
|
||||
#endif
|
||||
|
||||
// RAII-style wrapper around https://github.com/sheredom/subprocess.h,
|
||||
// exposing method calls instead of free functions operating on subprocess_s.
|
||||
struct common_subproc {
|
||||
common_subproc() = default;
|
||||
~common_subproc();
|
||||
|
||||
common_subproc(const common_subproc &) = delete;
|
||||
common_subproc & operator=(const common_subproc &) = delete;
|
||||
|
||||
// spawn a child process; if env is non-empty it replaces the child's environment
|
||||
// (do not combine with subprocess_option_inherit_environment)
|
||||
bool create(
|
||||
const std::vector<std::string> & args,
|
||||
int options,
|
||||
const std::vector<std::string> & env = {},
|
||||
const char * cwd = nullptr);
|
||||
|
||||
bool alive();
|
||||
|
||||
// true if LLAMA_SUBPROCESS was enabled at build time; when false, create() always fails
|
||||
static bool is_supported();
|
||||
|
||||
FILE * stdin_file();
|
||||
FILE * stdout_file();
|
||||
FILE * stderr_file();
|
||||
|
||||
// close stdin and detach it from the process, so a later join()/destroy() won't double-close it;
|
||||
// use this after writing all input to signal EOF to the child while it's still running
|
||||
void close_stdin();
|
||||
|
||||
void terminate();
|
||||
|
||||
// wait for the process to exit, release the underlying handle and return its exit code
|
||||
int join();
|
||||
|
||||
private:
|
||||
subprocess_s proc {};
|
||||
std::atomic<bool> is_created{false};
|
||||
|
||||
bool has_handle() const;
|
||||
};
|
||||
+123
@@ -0,0 +1,123 @@
|
||||
#include "trie.h"
|
||||
|
||||
#include "unicode.h"
|
||||
|
||||
#include <deque>
|
||||
|
||||
common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const {
|
||||
size_t current = 0; // Start at root
|
||||
size_t pos = start_pos;
|
||||
|
||||
// LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
|
||||
|
||||
while (pos < sv.size()) {
|
||||
auto result = common_parse_utf8_codepoint(sv, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
auto it = nodes[current].children.find(result.codepoint);
|
||||
if (it == nodes[current].children.end()) {
|
||||
// Can't continue matching
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
current = it->second;
|
||||
pos += result.bytes_consumed;
|
||||
|
||||
// Check if we've matched a complete word
|
||||
if (nodes[current].pattern >= 0) {
|
||||
return match_result{match_result::COMPLETE_MATCH};
|
||||
}
|
||||
}
|
||||
|
||||
// Reached end of input while still in the trie (not at root)
|
||||
if (current != 0) {
|
||||
// We're in the middle of a potential match
|
||||
return match_result{match_result::PARTIAL_MATCH};
|
||||
}
|
||||
|
||||
// Reached end at root (no match)
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
int32_t common_trie::insert(const std::string & word) {
|
||||
std::vector<uint32_t> symbols;
|
||||
size_t pos = 0;
|
||||
while (pos < word.length()) {
|
||||
auto result = common_parse_utf8_codepoint(word, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
symbols.push_back(result.codepoint);
|
||||
pos += result.bytes_consumed;
|
||||
}
|
||||
return insert(symbols);
|
||||
}
|
||||
|
||||
int32_t common_trie::insert(const std::vector<uint32_t> & symbols) {
|
||||
size_t current = 0;
|
||||
for (uint32_t ch : symbols) {
|
||||
auto it = nodes[current].children.find(ch);
|
||||
if (it == nodes[current].children.end()) {
|
||||
size_t child = create_node();
|
||||
nodes[current].children[ch] = child;
|
||||
current = child;
|
||||
} else {
|
||||
current = it->second;
|
||||
}
|
||||
}
|
||||
if (nodes[current].pattern < 0) {
|
||||
nodes[current].pattern = n_patterns++;
|
||||
}
|
||||
return nodes[current].pattern;
|
||||
}
|
||||
|
||||
common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) {
|
||||
const auto & nodes = t.nodes;
|
||||
const size_t n = nodes.size();
|
||||
|
||||
fail.assign(n, 0);
|
||||
order.reserve(n);
|
||||
|
||||
std::deque<size_t> queue{ 0 };
|
||||
while (!queue.empty()) {
|
||||
size_t u = queue.front();
|
||||
queue.pop_front();
|
||||
order.push_back(u);
|
||||
for (const auto & [ch, v] : nodes[u].children) {
|
||||
if (u != 0) {
|
||||
size_t f = fail[u];
|
||||
while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
|
||||
f = fail[f];
|
||||
}
|
||||
auto it = nodes[f].children.find(ch);
|
||||
fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
|
||||
}
|
||||
queue.push_back(v);
|
||||
}
|
||||
}
|
||||
|
||||
// fail[u] points to a strictly shorter suffix, so the first pattern found on
|
||||
// the fail chain (including u itself) is the longest pattern ending at u
|
||||
match.assign(n, -1);
|
||||
for (size_t u : order) {
|
||||
match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1);
|
||||
}
|
||||
|
||||
for (const auto & node : nodes) {
|
||||
for (const auto & [ch, v] : node.children) {
|
||||
alphabet.insert(ch);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
size_t common_aho_corasick::next(size_t state, uint32_t ch) const {
|
||||
const auto & nodes = t.nodes;
|
||||
while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
|
||||
state = fail[state];
|
||||
}
|
||||
auto it = nodes[state].children.find(ch);
|
||||
return it != nodes[state].children.end() ? it->second : 0;
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
// Trie for matching multiple literals.
|
||||
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
|
||||
struct common_trie {
|
||||
struct node {
|
||||
std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
|
||||
int32_t pattern = -1; // index of the pattern ending at this node, -1 if none
|
||||
};
|
||||
|
||||
std::vector<node> nodes;
|
||||
|
||||
common_trie() {
|
||||
create_node(); // root node
|
||||
}
|
||||
|
||||
common_trie(const std::vector<std::string> & words) : common_trie() {
|
||||
for (const auto & w : words) {
|
||||
insert(w);
|
||||
}
|
||||
}
|
||||
|
||||
enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
|
||||
|
||||
// Check if a delimiter starts at the given position
|
||||
match_result check_at(std::string_view sv, size_t start_pos) const;
|
||||
|
||||
// Insert a word as a sequence of Unicode codepoints, returns its pattern index
|
||||
int32_t insert(const std::string & word);
|
||||
|
||||
// Insert a raw symbol sequence, returns its pattern index (insertion order,
|
||||
// duplicates keep the first index)
|
||||
int32_t insert(const std::vector<uint32_t> & symbols);
|
||||
|
||||
private:
|
||||
int32_t n_patterns = 0;
|
||||
|
||||
size_t create_node() {
|
||||
size_t index = nodes.size();
|
||||
nodes.emplace_back();
|
||||
return index;
|
||||
}
|
||||
};
|
||||
|
||||
// Aho-Corasick automaton
|
||||
struct common_aho_corasick {
|
||||
common_trie t;
|
||||
std::vector<size_t> fail; // failure links
|
||||
std::vector<size_t> order; // states in BFS order
|
||||
std::vector<int32_t> match; // longest pattern ending at each state (directly or via a suffix link), -1 if none
|
||||
std::set<uint32_t> alphabet; // every character with a transition
|
||||
|
||||
common_aho_corasick(common_trie trie);
|
||||
|
||||
common_aho_corasick(const std::vector<std::string> & strings)
|
||||
: common_aho_corasick(common_trie(strings)) {}
|
||||
|
||||
size_t num_states() const { return t.nodes.size(); }
|
||||
bool is_terminal(size_t s) const { return match[s] >= 0; }
|
||||
|
||||
// index of the longest pattern ending at this state, -1 if none
|
||||
int32_t match_pattern(size_t s) const { return match[s]; }
|
||||
|
||||
// follow failure links until a transition on `ch` exists.
|
||||
size_t next(size_t state, uint32_t ch) const;
|
||||
};
|
||||
@@ -158,6 +158,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"MiniCPMForCausalLM": "minicpm",
|
||||
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
|
||||
"MiniMaxM2ForCausalLM": "minimax",
|
||||
"MiniMaxM3SparseForCausalLM": "minimax",
|
||||
"MiniMaxM3SparseForConditionalGeneration": "minimax",
|
||||
"Ministral3ForCausalLM": "mistral3",
|
||||
"Mistral3ForConditionalGeneration": "mistral3",
|
||||
"MistralForCausalLM": "llama",
|
||||
@@ -267,6 +269,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"Gemma4UnifiedForConditionalGeneration": "gemma",
|
||||
"Glm4vForConditionalGeneration": "qwen3vl",
|
||||
"Glm4vMoeForConditionalGeneration": "qwen3vl",
|
||||
"Glm5vForConditionalGeneration": "kimivl",
|
||||
"GlmOcrForConditionalGeneration": "qwen3vl",
|
||||
"GlmasrModel": "ultravox",
|
||||
"Granite4VisionForConditionalGeneration": "granite",
|
||||
@@ -285,6 +288,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"LlavaForConditionalGeneration": "llava",
|
||||
"MERaLiON2ForConditionalGeneration": "ultravox",
|
||||
"MiMoV2ForCausalLM": "mimo",
|
||||
"MiniMaxM3SparseForConditionalGeneration": "minimax",
|
||||
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
|
||||
"Mistral3ForConditionalGeneration": "llava",
|
||||
"NemotronH_Nano_VL_V2": "nemotron",
|
||||
|
||||
+2
-2
@@ -1156,7 +1156,7 @@ class TextModel(ModelBase):
|
||||
or "projector." in name or "pre_mm_projector_norm" in name \
|
||||
or "image_newline" in name or "view_seperator" in name \
|
||||
or "patch_embed" in name or "patch_embedding" in name \
|
||||
or "patch_merger." in name or "model.connector." in name:
|
||||
or "patch_merger." in name or "patch_merge_mlp." in name or "model.connector." in name:
|
||||
return None
|
||||
|
||||
return super().filter_tensors(item)
|
||||
@@ -1203,7 +1203,7 @@ class TextModel(ModelBase):
|
||||
self.gguf_writer.add_embedding_length(n_embd)
|
||||
logger.info(f"gguf: embedding length = {n_embd}")
|
||||
|
||||
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
|
||||
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
|
||||
self.gguf_writer.add_feed_forward_length(n_ff)
|
||||
logger.info(f"gguf: feed forward length = {n_ff}")
|
||||
|
||||
|
||||
@@ -152,3 +152,19 @@ class KimiK25Model(MmprojModel):
|
||||
name = name.replace(".proj.2.", ".proj.linear_2.")
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("Glm5vForConditionalGeneration")
|
||||
class Glm5vModel(KimiK25Model):
|
||||
"""GLM-5.2-Vision MoonViT3d encoder and projector
|
||||
|
||||
Uses the same vision encoder and projector as Kimi-K2.5, so it reuses the
|
||||
kimik25 projector type. The image begin/end tokens differ, but they are
|
||||
resolved at runtime from the text model vocab.
|
||||
"""
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
if name.startswith("mm_projector.linear_"):
|
||||
name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1)
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
+117
-2
@@ -7,7 +7,7 @@ import torch
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import ModelBase, TextModel, gguf
|
||||
from .base import ModelBase, TextModel, MmprojModel, gguf
|
||||
|
||||
|
||||
@ModelBase.register("MiniMaxM2ForCausalLM")
|
||||
@@ -23,7 +23,7 @@ class MiniMaxM2Model(TextModel):
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
|
||||
# merge expert weights
|
||||
if 'experts' in name:
|
||||
if "block_sparse_moe.experts." in name:
|
||||
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
|
||||
assert bid is not None
|
||||
|
||||
@@ -52,3 +52,118 @@ class MiniMaxM2Model(TextModel):
|
||||
return
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration")
|
||||
class MiniMaxM3Model(MiniMaxM2Model):
|
||||
model_arch = gguf.MODEL_ARCH.MINIMAXM3
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
if ".indexer." in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
self.gguf_writer.add_expert_shared_count(self.find_hparam(["n_shared_experts"]))
|
||||
self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"]))
|
||||
self.gguf_writer.add_expert_weights_norm(True)
|
||||
|
||||
sac = self.find_hparam(["sparse_attention_config"])
|
||||
self.gguf_writer.add_indexer_head_count(sac["sparse_num_index_heads"])
|
||||
self.gguf_writer.add_indexer_key_length(sac["sparse_index_dim"])
|
||||
self.gguf_writer.add_indexer_top_k(sac["sparse_topk_blocks"])
|
||||
self.gguf_writer.add_indexer_block_size(sac["sparse_block_size"])
|
||||
self.gguf_writer.add_indexer_local_blocks(sac["sparse_local_block"])
|
||||
|
||||
moe_layer_freq = self.find_hparam(["moe_layer_freq"])
|
||||
n_dense = 0
|
||||
for v in moe_layer_freq:
|
||||
if v == 0:
|
||||
n_dense += 1
|
||||
else:
|
||||
break
|
||||
self.gguf_writer.add_leading_dense_block_count(n_dense)
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
|
||||
# Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm
|
||||
if name.endswith("norm.weight"):
|
||||
data_torch = data_torch + 1.0
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("MiniMaxM3SparseForConditionalGeneration", "MiniMaxM3VLForConditionalGeneration")
|
||||
class MiniMaxM3VisionModel(MmprojModel):
|
||||
@classmethod
|
||||
def filter_tensors(cls, item):
|
||||
name, gen = item
|
||||
# keep only the vision-side tensors; text / mtp / sparse-index are dropped
|
||||
if not name.startswith(("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.")):
|
||||
return None
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
assert self.hparams_vision is not None
|
||||
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MINIMAXM3)
|
||||
self.gguf_writer.add_vision_use_gelu(True)
|
||||
|
||||
# the ViT carries its own LayerNorm eps (text tower uses a different one)
|
||||
self.gguf_writer.add_vision_attention_layernorm_eps(
|
||||
self.hparams_vision.get("layer_norm_eps", 1e-5)
|
||||
)
|
||||
|
||||
comp = self.hparams_vision.get("img_token_compression_config", {})
|
||||
merge_size = comp.get("spatial_merge_size", 2)
|
||||
self.gguf_writer.add_vision_spatial_merge_size(int(merge_size))
|
||||
|
||||
def modify_tensors(self, data_torch, name, bid):
|
||||
assert self.hparams_vision is not None
|
||||
|
||||
# Conv3d patch embed -> Conv2d slices
|
||||
if name == "vision_tower.vision_model.embeddings.patch_embedding.weight":
|
||||
if data_torch.ndim != 5:
|
||||
raise ValueError(f"unexpected patch_embedding rank {data_torch.ndim} for {name}")
|
||||
kt = data_torch.shape[2]
|
||||
base = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH]
|
||||
for t in range(kt):
|
||||
suffix = ".weight" if t == 0 else f".weight.{t}"
|
||||
yield (base + suffix, data_torch[:, :, t, ...])
|
||||
return
|
||||
|
||||
# Permute ViT q/k. HF [Ta Ha Wa | Tb Hb Wb | pad] reorder to [Ta Tb | Ha Hb | Wa Wb | pad].
|
||||
for new_name, tensor in super().modify_tensors(data_torch, name, bid):
|
||||
if ".attn_q." in new_name or ".attn_k." in new_name:
|
||||
tensor = self._permute_vit_qk(tensor, new_name)
|
||||
yield new_name, tensor
|
||||
|
||||
def _permute_vit_qk(self, t: "Tensor", new_name: str) -> "Tensor":
|
||||
assert self.hparams_vision is not None
|
||||
n_head = self.hparams_vision["num_attention_heads"]
|
||||
d_head = t.shape[0] // n_head
|
||||
axis_dim = 2 * ((2 * (d_head // 2) // 3) // 2)
|
||||
ah = axis_dim // 2
|
||||
half = 3 * ah
|
||||
perm = []
|
||||
perm += list(range(0, ah))
|
||||
perm += list(range(half, half + ah))
|
||||
perm += list(range(ah, 2 * ah))
|
||||
perm += list(range(half + ah, half + 2 * ah))
|
||||
perm += list(range(2 * ah, 3 * ah))
|
||||
perm += list(range(half + 2 * ah, half + 3 * ah))
|
||||
perm += list(range(2 * half, d_head))
|
||||
|
||||
assert axis_dim % 2 == 0
|
||||
assert 3 * axis_dim <= d_head
|
||||
assert len(perm) == d_head
|
||||
assert sorted(perm) == list(range(d_head)), "perm is not a bijection of d_head"
|
||||
assert t.shape[0] == n_head * d_head, f"{new_name}: {t.shape[0]} != {n_head}*{d_head}"
|
||||
assert d_head == 80
|
||||
|
||||
idx = torch.tensor(perm, dtype=torch.long)
|
||||
if t.ndim == 2:
|
||||
return t.reshape(n_head, d_head, t.shape[1])[:, idx, :].reshape(t.shape)
|
||||
return t.reshape(n_head, d_head)[:, idx].reshape(t.shape)
|
||||
|
||||
@@ -3281,6 +3281,35 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session *
|
||||
GGML_UNUSED(sess);
|
||||
}
|
||||
|
||||
static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
|
||||
const struct ggml_tensor * src1 = op->src[1];
|
||||
const struct ggml_tensor * dst = op;
|
||||
|
||||
const bool is_2D = ((const int32_t *) op->op_params)[6] == 1;
|
||||
if (!is_2D) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// For now support F32->F32 and F32->F16 only.
|
||||
if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs.
|
||||
const int32_t p0 = ((const int32_t *) op->op_params)[2];
|
||||
const int32_t p1 = ((const int32_t *) op->op_params)[3];
|
||||
if (p0 != 0 || p1 != 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
GGML_UNUSED(sess);
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
|
||||
const struct ggml_tensor * src0 = op->src[0];
|
||||
const struct ggml_tensor * dst = op;
|
||||
@@ -3430,6 +3459,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
|
||||
case GGML_OP_SOLVE_TRI: return HTP_OP_SOLVE_TRI;
|
||||
case GGML_OP_TRI: return HTP_OP_TRI;
|
||||
case GGML_OP_PAD: return HTP_OP_PAD;
|
||||
case GGML_OP_IM2COL: return HTP_OP_IM2COL;
|
||||
|
||||
case GGML_OP_UNARY:
|
||||
switch (ggml_get_unary_op(t)) {
|
||||
@@ -4152,6 +4182,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
|
||||
supp = ggml_hexagon_supported_ssm_conv(sess, op);
|
||||
break;
|
||||
|
||||
case GGML_OP_IM2COL:
|
||||
supp = ggml_hexagon_supported_im2col(sess, op);
|
||||
break;
|
||||
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
supp = ggml_hexagon_supported_gated_delta_net(sess, op);
|
||||
break;
|
||||
|
||||
@@ -42,6 +42,7 @@ add_library(${HTP_LIB} SHARED
|
||||
solve-tri-ops.c
|
||||
pad-ops.c
|
||||
argsort-ops.c
|
||||
im2col-ops.c
|
||||
)
|
||||
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE
|
||||
|
||||
@@ -140,5 +140,6 @@ int op_diag(struct htp_ops_context * octx);
|
||||
int op_solve_tri(struct htp_ops_context * octx);
|
||||
int op_gated_delta_net(struct htp_ops_context * octx);
|
||||
int op_pad(struct htp_ops_context * octx);
|
||||
int op_im2col(struct htp_ops_context * octx);
|
||||
|
||||
#endif /* HTP_CTX_H */
|
||||
|
||||
@@ -98,6 +98,7 @@ enum htp_op_code {
|
||||
HTP_OP_NORM,
|
||||
HTP_OP_CONCAT,
|
||||
HTP_OP_CLAMP,
|
||||
HTP_OP_IM2COL,
|
||||
|
||||
HTP_OP_INVALID
|
||||
};
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
#pragma clang diagnostic ignored "-Wunused-variable"
|
||||
#pragma clang diagnostic ignored "-Wunused-function"
|
||||
#pragma clang diagnostic ignored "-Wunused-but-set-variable"
|
||||
|
||||
#include <HAP_farf.h>
|
||||
#include <HAP_perf.h>
|
||||
#include <hexagon_protos.h>
|
||||
#include <hexagon_types.h>
|
||||
#include <string.h>
|
||||
|
||||
#define GGML_COMMON_DECL_C
|
||||
#include "ggml-common.h"
|
||||
#include "htp-ctx.h"
|
||||
#include "htp-ops.h"
|
||||
#include "hvx-utils.h"
|
||||
#include "hex-dma.h"
|
||||
#include "hex-profile.h"
|
||||
#include "htp-vtcm.h"
|
||||
|
||||
struct htp_im2col_context {
|
||||
struct htp_ops_context * octx;
|
||||
uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel)
|
||||
|
||||
uint32_t pe_rows_per_thread; // N*OH rows per worker
|
||||
uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256
|
||||
uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256
|
||||
|
||||
// Patch-embed DMA path VTCM ping-pong.
|
||||
uint8_t * pe_vtcm_src; // base of the 2x src buffers region
|
||||
uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region
|
||||
uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes
|
||||
uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes
|
||||
};
|
||||
|
||||
// Per-op VTCM layout for the patch-embed DMA path
|
||||
struct htp_im2col_vtcm_layout {
|
||||
size_t off_src;
|
||||
size_t off_dst;
|
||||
size_t src_bytes_per_thread;
|
||||
size_t dst_bytes_per_thread;
|
||||
size_t total_bytes;
|
||||
};
|
||||
|
||||
static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L,
|
||||
size_t src_row_bytes,
|
||||
size_t dst_row_bytes,
|
||||
uint32_t n_threads) {
|
||||
L->src_bytes_per_thread = 2 * src_row_bytes;
|
||||
L->dst_bytes_per_thread = 2 * dst_row_bytes;
|
||||
|
||||
L->off_src = 0;
|
||||
L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads;
|
||||
L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
|
||||
}
|
||||
|
||||
#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
|
||||
static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
|
||||
struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
|
||||
struct htp_ops_context * octx = ictx->octx; \
|
||||
struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
|
||||
const struct htp_tensor * restrict src1 = octx->src[1]; \
|
||||
const struct htp_tensor * restrict dst = octx->dst; \
|
||||
const int32_t s0 = octx->op_params[0]; \
|
||||
const int32_t s1 = octx->op_params[1]; \
|
||||
const int32_t p0 = octx->op_params[2]; \
|
||||
const int32_t p1 = octx->op_params[3]; \
|
||||
const int32_t d0 = octx->op_params[4]; \
|
||||
const int32_t d1 = octx->op_params[5]; \
|
||||
const uint32_t N = src1->ne[3]; \
|
||||
const uint32_t IC = src1->ne[2]; \
|
||||
const uint32_t IH = src1->ne[1]; \
|
||||
const uint32_t IW = src1->ne[0]; \
|
||||
const uint32_t KH = octx->src[0]->ne[1]; \
|
||||
const uint32_t KW = octx->src[0]->ne[0]; \
|
||||
const uint32_t OH = dst->ne[2]; \
|
||||
const uint32_t OW = dst->ne[1]; \
|
||||
const uint32_t patch_stride = IC * KH * KW; \
|
||||
const float * restrict src_data = (const float *) src1->data; \
|
||||
DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
|
||||
const uint32_t npatches = N * OH * OW; \
|
||||
const uint32_t patch_start = ictx->npatches_per_thread * ith; \
|
||||
const uint32_t patch_end = MIN(patch_start + ictx->npatches_per_thread, npatches); \
|
||||
if (patch_start >= patch_end) { \
|
||||
return; \
|
||||
} \
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
|
||||
for (uint32_t p = patch_start; p < patch_end; p++) { \
|
||||
const uint32_t iow = p % OW; \
|
||||
const uint32_t ioh = (p / OW) % OH; \
|
||||
const uint32_t in = p / (OW * OH); \
|
||||
DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \
|
||||
for (uint32_t iic = 0; iic < IC; iic++) { \
|
||||
const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \
|
||||
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
|
||||
const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \
|
||||
DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
|
||||
if (iih < 0 || iih >= (int32_t) IH) { \
|
||||
SPLAT_FN(out_run, 0.0f, KW); \
|
||||
continue; \
|
||||
} \
|
||||
const int32_t iiw0 = (int32_t) iow * s0 - p0; \
|
||||
const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \
|
||||
if (d0 == 1) { \
|
||||
/* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \
|
||||
const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \
|
||||
int32_t hi = (int32_t) IW - iiw0; \
|
||||
if (hi > (int32_t) KW) { \
|
||||
hi = (int32_t) KW; \
|
||||
} \
|
||||
if (hi <= lo) { \
|
||||
SPLAT_FN(out_run, 0.0f, KW); \
|
||||
} else { \
|
||||
if (lo > 0) { \
|
||||
SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \
|
||||
} \
|
||||
COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \
|
||||
(uint32_t) (hi - lo)); \
|
||||
if (hi < (int32_t) KW) { \
|
||||
SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \
|
||||
} \
|
||||
} \
|
||||
continue; \
|
||||
} \
|
||||
for (uint32_t ikw = 0; ikw < KW; ikw++) { \
|
||||
const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \
|
||||
out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \
|
||||
(DST_CTYPE) 0.0f : \
|
||||
(DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
|
||||
}
|
||||
|
||||
IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
|
||||
IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")
|
||||
|
||||
#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
|
||||
static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
|
||||
struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
|
||||
struct htp_ops_context * octx = ictx->octx; \
|
||||
struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
|
||||
const struct htp_tensor * restrict src1 = octx->src[1]; \
|
||||
const struct htp_tensor * restrict dst = octx->dst; \
|
||||
const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \
|
||||
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \
|
||||
const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \
|
||||
const uint32_t patch_stride = IC * KH * KW; \
|
||||
const float * restrict src_data = (const float *) src1->data; \
|
||||
DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
|
||||
dma_queue * dmaq = octx->ctx->dma[ith]; \
|
||||
uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \
|
||||
uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \
|
||||
float * srcb = (float *) src_base; \
|
||||
DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \
|
||||
const uint32_t nrows = N * OH; \
|
||||
const uint32_t per_thread = ictx->pe_rows_per_thread; \
|
||||
const uint32_t row_start = per_thread * ith; \
|
||||
const uint32_t row_end = MIN(row_start + per_thread, nrows); \
|
||||
if (row_start >= row_end) \
|
||||
return; \
|
||||
for (uint32_t r = row_start; r < row_end; r++) { \
|
||||
const uint32_t in = r / OH; \
|
||||
const uint32_t ioh = r % OH; \
|
||||
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
|
||||
int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
|
||||
int ok = (iih >= 0 && iih < (int32_t) IH); \
|
||||
for (uint32_t iic = 0; iic < IC; iic++) { \
|
||||
float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \
|
||||
const float * _vsrc = \
|
||||
ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \
|
||||
dma_queue_push_ddr_to_vtcm( \
|
||||
dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
|
||||
IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \
|
||||
} \
|
||||
} \
|
||||
for (uint32_t i = 0; i < IC * KH; i++) \
|
||||
dma_queue_pop(dmaq); \
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \
|
||||
for (uint32_t iow = 0; iow < OW; iow++) { \
|
||||
DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \
|
||||
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
|
||||
int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
|
||||
for (uint32_t iic = 0; iic < IC; iic++) { \
|
||||
DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
|
||||
if (iih < 0 || iih >= (int32_t) IH) { \
|
||||
SPLAT_FN(out_run, 0.0f, KW); \
|
||||
continue; \
|
||||
} \
|
||||
const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \
|
||||
COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \
|
||||
DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \
|
||||
dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \
|
||||
OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \
|
||||
dma_queue_flush(dmaq); \
|
||||
} \
|
||||
}
|
||||
|
||||
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
|
||||
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32")
|
||||
|
||||
static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
|
||||
const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];
|
||||
const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];
|
||||
const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];
|
||||
const int is_2D = octx->op_params[6] == 1;
|
||||
if (!is_2D) {
|
||||
return false;
|
||||
}
|
||||
if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
|
||||
return false;
|
||||
}
|
||||
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
|
||||
if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
|
||||
return false; // non-overlapping
|
||||
}
|
||||
if (p0 != 0 || p1 != 0) {
|
||||
return false; // no padding
|
||||
}
|
||||
if (d0 != 1 || d1 != 1) {
|
||||
return false; // no dilation
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path.
|
||||
// Returns false if it doesn't fit the VTCM budget (caller falls back).
|
||||
static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx,
|
||||
struct htp_im2col_context * ictx,
|
||||
uint32_t n_threads) {
|
||||
const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
|
||||
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
|
||||
const uint32_t OW = octx->dst->ne[1];
|
||||
const uint32_t patch_stride = IC * KH * KW;
|
||||
|
||||
ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256);
|
||||
const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
|
||||
ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256);
|
||||
|
||||
// 2 src + 2 dst buffers per thread (ping-pong), src region first then dst.
|
||||
struct htp_im2col_vtcm_layout L;
|
||||
htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads);
|
||||
if (L.total_bytes > octx->ctx->vtcm_size) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint8_t * const base = octx->ctx->vtcm_base;
|
||||
ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
|
||||
ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
|
||||
ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
|
||||
ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
|
||||
return true;
|
||||
}
|
||||
|
||||
int op_im2col(struct htp_ops_context * octx) {
|
||||
const struct htp_tensor * src1 = octx->src[1];
|
||||
const struct htp_tensor * dst = octx->dst;
|
||||
|
||||
if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) {
|
||||
FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported");
|
||||
return HTP_STATUS_NO_SUPPORT;
|
||||
}
|
||||
|
||||
const uint32_t N = src1->ne[3];
|
||||
const uint32_t OH = dst->ne[2];
|
||||
const uint32_t OW = dst->ne[1];
|
||||
const uint32_t npatches = N * OH * OW;
|
||||
const uint32_t n_threads = MIN(octx->n_threads, npatches);
|
||||
|
||||
if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
|
||||
struct htp_im2col_context ictx = { 0 };
|
||||
ictx.octx = octx;
|
||||
ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;
|
||||
|
||||
// Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
|
||||
// everything else (padding/dilation/stride edges) -> pure-DDR kernel.
|
||||
if (im2col_use_patchembed_dma(octx)) {
|
||||
const uint32_t nrows = N * OH;
|
||||
const uint32_t pth = MIN(octx->n_threads, nrows);
|
||||
if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
|
||||
ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
|
||||
if (dst->type == HTP_TYPE_F16) {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
|
||||
} else {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
|
||||
}
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
// else: doesn't fit -> fall through to the pure-DDR kernel below.
|
||||
}
|
||||
|
||||
if (dst->type == HTP_TYPE_F16) {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
|
||||
} else {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads);
|
||||
}
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
@@ -781,6 +781,9 @@ static int execute_op(struct htp_ops_context * octx) {
|
||||
case HTP_OP_PAD:
|
||||
return op_pad(octx);
|
||||
|
||||
case HTP_OP_IM2COL:
|
||||
return op_im2col(octx);
|
||||
|
||||
case HTP_OP_CONCAT:
|
||||
return op_concat(octx);
|
||||
|
||||
|
||||
@@ -12772,7 +12772,7 @@ static void ggml_opencl_op_rms_norm_fused(ggml_backend_t backend, ggml_tensor *
|
||||
ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
|
||||
|
||||
cl_ulong offset0 = extra0->offset + src0->view_offs;
|
||||
cl_ulong offset1 = extra1->offset + src0->view_offs;
|
||||
cl_ulong offset1 = extra1->offset + src1->view_offs;
|
||||
cl_ulong offsetd = extrad->offset + dst->view_offs;
|
||||
|
||||
ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
|
||||
|
||||
+1
-1
@@ -1424,7 +1424,7 @@ void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const vo
|
||||
struct gguf_writer_base {
|
||||
size_t written_bytes {0u};
|
||||
|
||||
~gguf_writer_base(void) = default;
|
||||
virtual ~gguf_writer_base(void) = default;
|
||||
|
||||
// we bet on devirtualization
|
||||
virtual void write(int8_t val) = 0;
|
||||
|
||||
@@ -200,6 +200,8 @@ class Keys:
|
||||
HEAD_COUNT = "{arch}.attention.indexer.head_count"
|
||||
KEY_LENGTH = "{arch}.attention.indexer.key_length"
|
||||
TOP_K = "{arch}.attention.indexer.top_k"
|
||||
BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA
|
||||
LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA
|
||||
TYPES = "{arch}.attention.indexer.types"
|
||||
|
||||
class HyperConnection:
|
||||
@@ -528,6 +530,7 @@ class MODEL_ARCH(IntEnum):
|
||||
APERTUS = auto()
|
||||
COGVLM = auto()
|
||||
MINIMAXM2 = auto()
|
||||
MINIMAXM3 = auto()
|
||||
RND1 = auto()
|
||||
PANGU_EMBED = auto()
|
||||
MISTRAL3 = auto()
|
||||
@@ -774,6 +777,9 @@ class MODEL_TENSOR(IntEnum):
|
||||
INDEXER_PROJ = auto()
|
||||
INDEXER_ATTN_K = auto()
|
||||
INDEXER_ATTN_Q_B = auto()
|
||||
INDEXER_Q_PROJ = auto()
|
||||
INDEXER_K_PROJ = auto()
|
||||
INDEXER_Q_NORM = auto()
|
||||
INDEXER_COMPRESSOR_WKV = auto()
|
||||
INDEXER_COMPRESSOR_WGATE = auto()
|
||||
INDEXER_COMPRESSOR_APE = auto()
|
||||
@@ -851,6 +857,8 @@ class MODEL_TENSOR(IntEnum):
|
||||
V_MM_UP = auto() # cogvlm
|
||||
V_MM_DOWN = auto() # cogvlm
|
||||
V_MM_GATE = auto() # cogvlm
|
||||
V_MM_MERGER_FC1 = auto() # minimax-m3 (patch-merge MLP)
|
||||
V_MM_MERGER_FC2 = auto() # minimax-m3 (patch-merge MLP)
|
||||
V_TOK_BOI = auto() # cogvlm
|
||||
V_TOK_EOI = auto() # cogvlm
|
||||
V_TOK_IMG_BEGIN = auto() # hunyuanvl
|
||||
@@ -1110,6 +1118,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.GROVEMOE: "grovemoe",
|
||||
MODEL_ARCH.APERTUS: "apertus",
|
||||
MODEL_ARCH.MINIMAXM2: "minimax-m2",
|
||||
MODEL_ARCH.MINIMAXM3: "minimax-m3",
|
||||
MODEL_ARCH.COGVLM: "cogvlm",
|
||||
MODEL_ARCH.RND1: "rnd1",
|
||||
MODEL_ARCH.PANGU_EMBED: "pangu-embedded",
|
||||
@@ -1355,6 +1364,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj",
|
||||
MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k",
|
||||
MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b",
|
||||
MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj",
|
||||
MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj",
|
||||
MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm",
|
||||
MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv",
|
||||
MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate",
|
||||
MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape",
|
||||
@@ -1431,6 +1443,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.V_MM_UP: "mm.up",
|
||||
MODEL_TENSOR.V_MM_DOWN: "mm.down",
|
||||
MODEL_TENSOR.V_MM_GATE: "mm.gate",
|
||||
MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1",
|
||||
MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2",
|
||||
MODEL_TENSOR.V_TOK_BOI: "v.boi",
|
||||
MODEL_TENSOR.V_TOK_EOI: "v.eoi",
|
||||
MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm",
|
||||
@@ -1627,6 +1641,8 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.V_RESMPL_QUERY,
|
||||
MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK,
|
||||
MODEL_TENSOR.V_MM_PATCH_MERGER,
|
||||
MODEL_TENSOR.V_MM_MERGER_FC1,
|
||||
MODEL_TENSOR.V_MM_MERGER_FC2,
|
||||
MODEL_TENSOR.V_DS_NORM,
|
||||
MODEL_TENSOR.V_DS_FC1,
|
||||
MODEL_TENSOR.V_DS_FC2,
|
||||
@@ -4163,6 +4179,34 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
],
|
||||
MODEL_ARCH.MINIMAXM3: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.ATTN_NORM,
|
||||
MODEL_TENSOR.ATTN_Q,
|
||||
MODEL_TENSOR.ATTN_Q_NORM,
|
||||
MODEL_TENSOR.ATTN_K,
|
||||
MODEL_TENSOR.ATTN_K_NORM,
|
||||
MODEL_TENSOR.ATTN_V,
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE_INP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
MODEL_TENSOR.FFN_GATE_EXP,
|
||||
MODEL_TENSOR.FFN_DOWN_EXP,
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_GATE_SHEXP,
|
||||
MODEL_TENSOR.FFN_DOWN_SHEXP,
|
||||
MODEL_TENSOR.FFN_UP_SHEXP,
|
||||
MODEL_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.INDEXER_Q_PROJ,
|
||||
MODEL_TENSOR.INDEXER_K_PROJ,
|
||||
MODEL_TENSOR.INDEXER_Q_NORM,
|
||||
MODEL_TENSOR.INDEXER_K_NORM,
|
||||
],
|
||||
MODEL_ARCH.COGVLM: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
@@ -4733,6 +4777,7 @@ class VisionProjectorType:
|
||||
YOUTUVL = "youtuvl"
|
||||
NEMOTRON_V2_VL = "nemotron_v2_vl"
|
||||
HUNYUANVL = "hunyuanvl"
|
||||
MINIMAXM3 = "minimax_m3"
|
||||
MINICPMV4_6 = "minicpmv4_6"
|
||||
GRANITE_SPEECH = "granite_speech" # audio
|
||||
MIMOVL = "mimovl"
|
||||
|
||||
@@ -793,6 +793,12 @@ class GGUFWriter:
|
||||
def add_indexer_top_k(self, top_k: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k)
|
||||
|
||||
def add_indexer_block_size(self, block_size: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.BLOCK_SIZE.format(arch=self.arch), block_size)
|
||||
|
||||
def add_indexer_local_blocks(self, local_blocks: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks)
|
||||
|
||||
def add_indexer_types(self, value: Sequence[bool]) -> None:
|
||||
key = Keys.Attention.Indexer.TYPES.format(arch=self.arch)
|
||||
self.add_array(key, value)
|
||||
|
||||
@@ -1264,7 +1264,8 @@ class TensorNameMap:
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_K_NORM: (
|
||||
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
|
||||
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
|
||||
"model.layers.{bid}.self_attn.index_k_norm", # MSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_PROJ: (
|
||||
@@ -1279,6 +1280,18 @@ class TensorNameMap:
|
||||
"model.layers.{bid}.self_attn.indexer.wq_b", # DSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_Q_PROJ: (
|
||||
"model.layers.{bid}.self_attn.index_q_proj", # MSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_K_PROJ: (
|
||||
"model.layers.{bid}.self_attn.index_k_proj", # MSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_Q_NORM: (
|
||||
"model.layers.{bid}.self_attn.index_q_norm", # MSA
|
||||
),
|
||||
|
||||
############################################################################
|
||||
# TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
|
||||
MODEL_TENSOR.ENC_OUTPUT_NORM: (
|
||||
@@ -1825,6 +1838,14 @@ class TensorNameMap:
|
||||
"visual.downsample", # glm4v
|
||||
),
|
||||
|
||||
MODEL_TENSOR.V_MM_MERGER_FC1: (
|
||||
"patch_merge_mlp.linear_1", # minimax-m3
|
||||
),
|
||||
|
||||
MODEL_TENSOR.V_MM_MERGER_FC2: (
|
||||
"patch_merge_mlp.linear_2", # minimax-m3
|
||||
),
|
||||
|
||||
MODEL_TENSOR.V_DS_NORM: (
|
||||
"model.visual.deepstack_merger_list.{bid}.norm", # deepstack in qwen3vl
|
||||
),
|
||||
|
||||
+5
-4
@@ -203,10 +203,11 @@ extern "C" {
|
||||
};
|
||||
|
||||
enum llama_load_mode {
|
||||
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
|
||||
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
|
||||
LLAMA_LOAD_MODE_MLOCK = 2, // mmap + force system to keep model in RAM rather than swapping or compressing
|
||||
LLAMA_LOAD_MODE_DIRECT_IO = 3, // use direct I/O if available
|
||||
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
|
||||
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
|
||||
LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing
|
||||
LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing
|
||||
LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available
|
||||
};
|
||||
|
||||
LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
|
||||
|
||||
@@ -5,7 +5,7 @@ import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
HTTPLIB_VERSION = "refs/tags/v0.50.1"
|
||||
HTTPLIB_VERSION = "refs/tags/v0.51.0"
|
||||
|
||||
vendor = {
|
||||
"https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp",
|
||||
|
||||
@@ -110,6 +110,15 @@ Public API changes carry a higher bar than internal ones (`CONTRIBUTING.md`). Re
|
||||
- Security: don't trust client-supplied headers (e.g. `X-Forwarded-For`) or add footguns; things like IP allowlisting belong at a reverse proxy unless there's a trusted-proxy design.
|
||||
- Wire new behavior into the existing request/response and checkpoint paths correctly; watch for resource leaks across requests.
|
||||
|
||||
## Multimodal (`tools/mtmd/`)
|
||||
|
||||
- Tensor names must be prefixed by `v.`, `a.`, `mm.` or `a.mm.` (legacy naming doesn't follow this convention - this is expected, but new code should follow it).
|
||||
- Do not use explicit sin/cos for RoPE; use `ggml_rope_ext` instead, see `HOWTO-add-model.md`. If it can't express the needed behavior, that's a design discussion, not a PR.
|
||||
- New GGML ops must not be introduced in the same PR, you must push it as a separate PR.
|
||||
- In most cases, `build_vit` should be enough to build the transformer graph for vision models. Do not add a loop to build the transformer graph manually, unless you have a very good reason to do so. If you do, please explain why in the PR description.
|
||||
- If you need a dedicated preprocessor, there is a high chance that it can be a derived class from one of the existing preprocessors. Check carefully before adding a new preprocessor class.
|
||||
- If the model need a new public API in `mtmd.h`, open a discussion first.
|
||||
|
||||
## General (always)
|
||||
|
||||
Enforce the `AGENTS.md` / `CONTRIBUTING.md` coding and naming guidelines on every changed line - this is a distinct pass from checking that the code works, and matters just as much for review speed:
|
||||
|
||||
@@ -127,6 +127,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_GROVEMOE, "grovemoe" },
|
||||
{ LLM_ARCH_APERTUS, "apertus" },
|
||||
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
|
||||
{ LLM_ARCH_MINIMAX_M3, "minimax-m3" },
|
||||
{ LLM_ARCH_COGVLM, "cogvlm" },
|
||||
{ LLM_ARCH_RND1, "rnd1" },
|
||||
{ LLM_ARCH_PANGU_EMBED, "pangu-embedded" },
|
||||
@@ -253,6 +254,8 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" },
|
||||
@@ -597,6 +600,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
|
||||
{ LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" },
|
||||
{ LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" },
|
||||
{ LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" },
|
||||
{ LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" },
|
||||
{ LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" },
|
||||
{ LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" },
|
||||
{ LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" },
|
||||
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
|
||||
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
|
||||
@@ -832,6 +838,9 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
|
||||
{LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
|
||||
{LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
|
||||
@@ -1001,6 +1010,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
|
||||
case LLM_ARCH_LFM2:
|
||||
case LLM_ARCH_LFM2MOE:
|
||||
case LLM_ARCH_MINIMAX_M2:
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
case LLM_ARCH_MISTRAL4:
|
||||
case LLM_ARCH_KIMI_LINEAR:
|
||||
return false;
|
||||
|
||||
@@ -146,6 +146,7 @@ enum llm_arch {
|
||||
LLM_ARCH_TALKIE,
|
||||
LLM_ARCH_MELLUM,
|
||||
LLM_ARCH_EAGLE3,
|
||||
LLM_ARCH_MINIMAX_M3,
|
||||
LLM_ARCH_DFLASH,
|
||||
LLM_ARCH_UNKNOWN,
|
||||
};
|
||||
@@ -258,6 +259,8 @@ enum llm_kv {
|
||||
LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,
|
||||
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
|
||||
LLM_KV_ATTENTION_INDEXER_TOP_K,
|
||||
LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,
|
||||
LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,
|
||||
LLM_KV_ATTENTION_INDEXER_TYPES,
|
||||
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
|
||||
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
|
||||
@@ -597,6 +600,9 @@ enum llm_tensor {
|
||||
LLM_TENSOR_INDEXER_PROJ,
|
||||
LLM_TENSOR_INDEXER_ATTN_K,
|
||||
LLM_TENSOR_INDEXER_ATTN_Q_B,
|
||||
LLM_TENSOR_INDEXER_Q_PROJ,
|
||||
LLM_TENSOR_INDEXER_K_PROJ,
|
||||
LLM_TENSOR_INDEXER_Q_NORM,
|
||||
LLM_TENSOR_INDEXER_COMPRESSOR_WKV,
|
||||
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
|
||||
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
|
||||
|
||||
@@ -2338,7 +2338,8 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
|
||||
model.arch == LLM_ARCH_KIMI_LINEAR ||
|
||||
model.arch == LLM_ARCH_QWEN35 ||
|
||||
model.arch == LLM_ARCH_QWEN35MOE ||
|
||||
model.arch == LLM_ARCH_DEEPSEEK4) {
|
||||
model.arch == LLM_ARCH_DEEPSEEK4 ||
|
||||
model.arch == LLM_ARCH_MINIMAX_M3) {
|
||||
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
|
||||
}
|
||||
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
|
||||
|
||||
+12
-1
@@ -1709,6 +1709,17 @@ ggml_tensor * llm_graph_context::build_ffn(
|
||||
cur = ggml_swiglu(ctx0, cur);
|
||||
cb(cur, "ffn_swiglu", il);
|
||||
} break;
|
||||
case LLM_FFN_SWIGLU_OAI_MOE:
|
||||
if (gate && type_gate == LLM_FFN_PAR) {
|
||||
// same alpha/limit constants as gpt-oss
|
||||
const float alpha = 1.702f;
|
||||
const float limit = 7.0f;
|
||||
cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit);
|
||||
cb(cur, "ffn_swiglu_oai", il);
|
||||
type_gate = LLM_FFN_SEQ;
|
||||
} else {
|
||||
GGML_ABORT("LLM_FFN_SWIGLU_OAI_MOE requires a parallel gate");
|
||||
} break;
|
||||
case LLM_FFN_GEGLU:
|
||||
{
|
||||
cur = ggml_geglu(ctx0, cur);
|
||||
@@ -2668,7 +2679,7 @@ ggml_tensor * llm_graph_context::build_attn(
|
||||
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
|
||||
}
|
||||
|
||||
const auto & kq_mask = inp->get_kq_mask();
|
||||
ggml_tensor * kq_mask = inp->get_kq_mask();
|
||||
|
||||
ggml_tensor * q = q_cur;
|
||||
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
|
||||
|
||||
@@ -180,6 +180,16 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const {
|
||||
return val;
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const {
|
||||
if (!indexer_kv || indexer_head_size == 0) {
|
||||
return 0; // arch without a MSA indexer
|
||||
}
|
||||
if (il < n_layer_dense_lead) {
|
||||
return 0; // leading dense layers carry no indexer
|
||||
}
|
||||
return indexer_head_size; // 128
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_r() const {
|
||||
if (wkv_head_size != 0) {
|
||||
// for RWKV models
|
||||
|
||||
@@ -226,6 +226,11 @@ struct llama_hparams {
|
||||
uint32_t indexer_n_head = 0;
|
||||
uint32_t indexer_head_size = 0;
|
||||
uint32_t indexer_top_k = 0;
|
||||
// MSA
|
||||
uint32_t indexer_block_size = 0;
|
||||
uint32_t indexer_local_blocks = 0;
|
||||
// MSA stores its indexer keys in the main KV cache (k_idx tensors);
|
||||
bool indexer_kv = false;
|
||||
|
||||
// Indexer is "full" (1) or "shared" (0)
|
||||
// Shared indexers reuse top-k from previous full layer
|
||||
@@ -350,6 +355,9 @@ struct llama_hparams {
|
||||
uint32_t n_embd_k_gqa_max() const;
|
||||
uint32_t n_embd_v_gqa_max() const;
|
||||
|
||||
// dimension of the single-head MSA indexer key stream
|
||||
uint32_t n_embd_k_idx(uint32_t il = 0) const;
|
||||
|
||||
// dimension of the rolling state embeddings
|
||||
// corresponds to Mamba's conv_states size or RWKV's token_shift states size
|
||||
uint32_t n_embd_r() const;
|
||||
|
||||
+278
-14
@@ -112,7 +112,7 @@ llama_kv_cache::llama_kv_cache(
|
||||
auto it = ctx_map.find(buft);
|
||||
if (it == ctx_map.end()) {
|
||||
ggml_init_params params = {
|
||||
/*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
|
||||
/*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), //Reserve tensor metadata for up to 3 tensors per layer (K, V, and optional K_idx), plus one view per tensor per stream.
|
||||
/*.mem_buffer =*/ NULL,
|
||||
/*.no_alloc =*/ true,
|
||||
};
|
||||
@@ -242,9 +242,25 @@ llama_kv_cache::llama_kv_cache(
|
||||
v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr);
|
||||
}
|
||||
|
||||
const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il);
|
||||
ggml_tensor * k_idx = n_embd_k_idx > 0
|
||||
? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream)
|
||||
: nullptr;
|
||||
if (k_idx) {
|
||||
ggml_format_name(k_idx, "cache_k_idx_l%d", il);
|
||||
msa_strict_slots = (n_stream == n_seq_max);
|
||||
}
|
||||
|
||||
std::vector<ggml_tensor *> k_idx_stream;
|
||||
for (uint32_t s = 0; s < n_stream; ++s) {
|
||||
k_idx_stream.push_back(k_idx
|
||||
? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2])
|
||||
: nullptr);
|
||||
}
|
||||
|
||||
map_layer_ids[il] = layers.size();
|
||||
|
||||
layers.push_back({ il, k, v, k_stream, v_stream, });
|
||||
layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream });
|
||||
}
|
||||
|
||||
if (reuse) {
|
||||
@@ -293,13 +309,24 @@ llama_kv_cache::llama_kv_cache(
|
||||
}
|
||||
|
||||
{
|
||||
const size_t memory_size_k = size_k_bytes();
|
||||
const size_t memory_size_v = size_v_bytes();
|
||||
const size_t memory_size_k = size_k_bytes();
|
||||
const size_t memory_size_v = size_v_bytes();
|
||||
const size_t memory_size_k_idx = size_k_idx_bytes();
|
||||
const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx;
|
||||
|
||||
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,
|
||||
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,
|
||||
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
|
||||
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
|
||||
constexpr float mib = 1024.0f * 1024.0f;
|
||||
|
||||
const std::string k_log = format(", K (%s): %7.2f MiB", ggml_type_name(type_k), (float) memory_size_k / mib);
|
||||
const std::string v_log = format(", V (%s): %7.2f MiB", ggml_type_name(type_v), (float) memory_size_v / mib);
|
||||
|
||||
std::string k_idx_log;
|
||||
if (memory_size_k_idx > 0) {
|
||||
k_idx_log = format(", K_idx (%s): %7.2f MiB", ggml_type_name(GGML_TYPE_F32), (float) memory_size_k_idx / mib);
|
||||
}
|
||||
|
||||
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs)%s%s%s\n", __func__,
|
||||
(float) memory_size_total / mib, kv_size, (int) layers.size(), n_seq_max, n_stream,
|
||||
k_log.c_str(), v_log.c_str(), k_idx_log.c_str());
|
||||
}
|
||||
|
||||
// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
|
||||
@@ -392,6 +419,39 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
p1 = std::numeric_limits<llama_pos>::max();
|
||||
}
|
||||
|
||||
// empty range - nothing to remove
|
||||
if (p0 >= p1) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// MSA anchors block selection to absolute cache slots (slot == position). Tail trim and full removal preserve this invariant, but removing a prefix
|
||||
// or middle range would free slots while later cells survive, desynchronizing the indexer cache. Reject such removals before modifying the cache.
|
||||
if (msa_strict_slots) {
|
||||
for (llama_seq_id sid = 0; sid < (llama_seq_id) seq_to_stream.size(); ++sid) {
|
||||
if (seq_id >= 0 && sid != seq_id) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const auto & cells = v_cells[seq_to_stream[sid]];
|
||||
|
||||
const llama_pos pmin = cells.seq_pos_min(sid);
|
||||
const llama_pos pmax = cells.seq_pos_max(sid);
|
||||
|
||||
if (pmin < 0) {
|
||||
continue; // empty sequence
|
||||
}
|
||||
|
||||
const bool overlaps = p0 <= pmax && p1 > pmin; // the range removes something
|
||||
const bool leaves_tail = p1 <= pmax; // cells beyond the range survive
|
||||
|
||||
if (overlaps && leaves_tail) {
|
||||
LLAMA_LOG_WARN("%s: MSA: partial (non-suffix) removal [%d, %d) for seq %d is not supported "
|
||||
"(block selection is anchored to cache slots) - rejected\n", __func__, p0, p1, sid);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (seq_id >= 0) {
|
||||
auto & cells = v_cells[seq_to_stream[seq_id]];
|
||||
auto & head = v_heads[seq_to_stream[seq_id]];
|
||||
@@ -846,6 +906,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
|
||||
if (layer.v_stream[ssrc]) {
|
||||
ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);
|
||||
}
|
||||
if (layer.k_idx_stream[ssrc]) {
|
||||
GGML_ASSERT(layer.k_idx_stream[sdst]);
|
||||
ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -994,6 +1058,44 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
|
||||
|
||||
const auto & cells = v_cells[seq_to_stream[seq_id]];
|
||||
|
||||
if (n_tokens > cells.size()) {
|
||||
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
|
||||
return { };
|
||||
}
|
||||
|
||||
// MSA block selection assumes slot == logical position (append-only streams).
|
||||
if (msa_strict_slots) {
|
||||
for (uint32_t ii = 0; ii < n_tokens; ++ii) {
|
||||
const llama_pos pos = ubatch.pos[s*n_tokens + ii];
|
||||
|
||||
if (pos < 0 || (uint64_t) pos >= cells.size()) {
|
||||
LLAMA_LOG_WARN("%s: MSA: position %d is outside the cache range [0, %u)\n",
|
||||
__func__, pos, cells.size());
|
||||
return { };
|
||||
}
|
||||
|
||||
const uint32_t idx = (uint32_t) pos;
|
||||
|
||||
if (!cells.is_empty(idx)) {
|
||||
LLAMA_LOG_WARN("%s: MSA: required slot %u is already occupied (stream %u)\n",
|
||||
__func__, idx, seq_to_stream[seq_id]);
|
||||
return { };
|
||||
}
|
||||
|
||||
// strictly increasing positions, rules out duplicates and, for contiguous requests, is tightened to exact adjacency
|
||||
if (!res.idxs[s].empty() && (cont ? idx != res.idxs[s].back() + 1
|
||||
: idx <= res.idxs[s].back())) {
|
||||
LLAMA_LOG_WARN("%s: MSA: token positions are not %s within the ubatch\n",
|
||||
__func__, cont ? "contiguous" : "strictly increasing");
|
||||
return { };
|
||||
}
|
||||
|
||||
res.idxs[s].push_back(idx);
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
uint32_t head_cur = v_heads[seq_to_stream[seq_id]];
|
||||
|
||||
// if we have enough unused cells before the current head ->
|
||||
@@ -1002,11 +1104,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
|
||||
head_cur = 0;
|
||||
}
|
||||
|
||||
if (n_tokens > cells.size()) {
|
||||
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
|
||||
return { };
|
||||
}
|
||||
|
||||
uint32_t n_tested = 0;
|
||||
|
||||
// for continuous slots, we test that all tokens in the ubatch fit, starting from the current head
|
||||
@@ -1113,6 +1210,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
|
||||
|
||||
const auto idx = sinfo.idxs[s][ii];
|
||||
|
||||
if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) {
|
||||
LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: "
|
||||
"writing pos %d into cell %u (stream %u). The indexer cache "
|
||||
"would desync and block selection would silently corrupt. "
|
||||
"This is a bug, please report it with reproduction steps.\n",
|
||||
__func__, ubatch.pos[i], idx, sinfo.strm[s]);
|
||||
GGML_ABORT("MSA: slot != pos");
|
||||
}
|
||||
|
||||
if (!cells.is_empty(idx)) {
|
||||
assert(cells.seq_count(idx) == 1);
|
||||
|
||||
@@ -1156,7 +1262,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
|
||||
LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",
|
||||
__func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
|
||||
|
||||
seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
|
||||
// under MSA strict slots this path should be unreachable, since strict MSA placement never selects occupied cells
|
||||
GGML_ASSERT(seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1176,6 +1283,12 @@ bool llama_kv_cache::get_can_shift() const {
|
||||
if (hparams.n_pos_per_embd() > 1) {
|
||||
return false;
|
||||
}
|
||||
// shifting would leave k_idx stale
|
||||
for (const auto & layer : layers) {
|
||||
if (layer.k_idx) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -1292,6 +1405,23 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
|
||||
ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
|
||||
const int32_t ikv = map_layer_ids.at(il);
|
||||
auto * k_idx = layers[ikv].k_idx;
|
||||
GGML_ASSERT(k_idx);
|
||||
|
||||
const uint64_t kv_size = get_size();
|
||||
const int64_t n_idx = k_idx->ne[0]; // 128
|
||||
const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
|
||||
|
||||
return ggml_view_4d(ctx, k_idx,
|
||||
n_idx, 1, n_kv, ns,
|
||||
ggml_row_size(k_idx->type, n_idx), // nb1 (single head)
|
||||
ggml_row_size(k_idx->type, n_idx), // nb2 (per cell)
|
||||
ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream)
|
||||
ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
|
||||
GGML_UNUSED(sinfo);
|
||||
|
||||
@@ -1393,6 +1523,28 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama
|
||||
return k_idxs;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
|
||||
GGML_UNUSED(sinfo);
|
||||
const int32_t ikv = map_layer_ids.at(il);
|
||||
ggml_tensor * k_idx = layers[ikv].k_idx;
|
||||
GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache");
|
||||
|
||||
const int64_t n_embd_head = k_idx_cur->ne[0]; // 128
|
||||
const int64_t n_head = k_idx_cur->ne[1]; // 1
|
||||
const int64_t n_tokens = k_idx_cur->ne[2];
|
||||
const int64_t n_embd_gqa = n_embd_head*n_head; // 128
|
||||
|
||||
GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]);
|
||||
k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0);
|
||||
|
||||
const int64_t n_stream = k_idx->ne[2];
|
||||
if (n_stream > 1) {
|
||||
const int64_t kv_size = get_size();
|
||||
k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream);
|
||||
}
|
||||
return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
|
||||
const uint32_t n_tokens = ubatch.n_tokens;
|
||||
|
||||
@@ -1827,6 +1979,18 @@ size_t llama_kv_cache::size_v_bytes() const {
|
||||
return size_v_bytes;
|
||||
}
|
||||
|
||||
size_t llama_kv_cache::size_k_idx_bytes() const {
|
||||
size_t size_k_idx_bytes = 0;
|
||||
|
||||
for (const auto & layer : layers) {
|
||||
if (layer.k_idx) {
|
||||
size_k_idx_bytes += ggml_nbytes(layer.k_idx);
|
||||
}
|
||||
}
|
||||
|
||||
return size_k_idx_bytes;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::build_rope_shift(
|
||||
const llama_cparams & cparams,
|
||||
ggml_context * ctx,
|
||||
@@ -2139,6 +2303,36 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
|
||||
}
|
||||
}
|
||||
|
||||
if (size_k_idx_bytes() > 0) {
|
||||
const uint32_t has_k_idx_u32 = 1;
|
||||
io.write(&has_k_idx_u32, sizeof(has_k_idx_u32));
|
||||
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0;
|
||||
io.write(&layer_has_k_idx, sizeof(layer_has_k_idx));
|
||||
|
||||
if (!layer_has_k_idx) {
|
||||
continue;
|
||||
}
|
||||
|
||||
GGML_ASSERT(layer.k_idx_stream[cr.strm]);
|
||||
|
||||
const int32_t k_idx_type_i = (int32_t) layer.k_idx->type;
|
||||
io.write(&k_idx_type_i, sizeof(k_idx_type_i));
|
||||
|
||||
const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
|
||||
io.write(&k_idx_size_row, sizeof(k_idx_size_row));
|
||||
|
||||
for (const auto & range : cr.data) {
|
||||
const size_t range_size = range.second - range.first;
|
||||
const size_t buf_size = range_size * k_idx_size_row;
|
||||
const size_t offset = range.first * k_idx_size_row;
|
||||
|
||||
io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!v_trans) {
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t il = layer.il;
|
||||
@@ -2387,6 +2581,68 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
|
||||
}
|
||||
}
|
||||
|
||||
if (size_k_idx_bytes() > 0) {
|
||||
uint32_t has_k_idx_u32 = 0;
|
||||
io.read(&has_k_idx_u32, sizeof(has_k_idx_u32));
|
||||
|
||||
if (has_k_idx_u32 != 1) {
|
||||
LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const auto & layer : layers) {
|
||||
uint32_t layer_has_k_idx = 0;
|
||||
io.read(&layer_has_k_idx, sizeof(layer_has_k_idx));
|
||||
|
||||
const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0;
|
||||
|
||||
if (layer_has_k_idx != expected_layer_has_k_idx) {
|
||||
LLAMA_LOG_ERROR(
|
||||
"%s: mismatched k_idx state for layer: got %u, expected %u\n",
|
||||
__func__, layer_has_k_idx, expected_layer_has_k_idx);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!layer_has_k_idx) {
|
||||
continue;
|
||||
}
|
||||
|
||||
GGML_ASSERT(layer.k_idx_stream[strm]);
|
||||
|
||||
int32_t k_idx_type_i = -1;
|
||||
io.read(&k_idx_type_i, sizeof(k_idx_type_i));
|
||||
|
||||
if (k_idx_type_i != (int32_t) layer.k_idx->type) {
|
||||
LLAMA_LOG_ERROR(
|
||||
"%s: mismatched k_idx type: got %d, expected %d\n",
|
||||
__func__, k_idx_type_i, (int32_t) layer.k_idx->type);
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t k_idx_size_row = 0;
|
||||
io.read(&k_idx_size_row, sizeof(k_idx_size_row));
|
||||
|
||||
const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
|
||||
|
||||
if (k_idx_size_row != expected_k_idx_size_row) {
|
||||
LLAMA_LOG_ERROR(
|
||||
"%s: mismatched k_idx row size: got %zu, expected %zu\n",
|
||||
__func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (cell_count) {
|
||||
if (sinfo.is_contiguous()) {
|
||||
io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row);
|
||||
} else {
|
||||
for (uint32_t i = 0; i < cell_count; ++i) {
|
||||
io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!this->v_trans) {
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t il = layer.il;
|
||||
@@ -2588,6 +2844,10 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
|
||||
return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const {
|
||||
return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
|
||||
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
|
||||
}
|
||||
@@ -2596,6 +2856,10 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_
|
||||
return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const {
|
||||
return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
|
||||
return kv->build_input_k_idxs(ctx, ubatch);
|
||||
}
|
||||
|
||||
@@ -173,10 +173,12 @@ public:
|
||||
// get views of the current state of the cache
|
||||
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
|
||||
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
|
||||
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
|
||||
|
||||
// store k_cur and v_cur in the cache based on the provided head location
|
||||
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
|
||||
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
|
||||
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
|
||||
|
||||
//
|
||||
// preparation API
|
||||
@@ -228,9 +230,11 @@ private:
|
||||
|
||||
ggml_tensor * k;
|
||||
ggml_tensor * v;
|
||||
ggml_tensor * k_idx; // MSA single-head indexer keys, F32
|
||||
|
||||
std::vector<ggml_tensor *> k_stream;
|
||||
std::vector<ggml_tensor *> v_stream;
|
||||
std::vector<ggml_tensor *> k_idx_stream;
|
||||
};
|
||||
|
||||
bool v_trans = true; // the value tensor is transposed
|
||||
@@ -259,6 +263,9 @@ private:
|
||||
// env: LLAMA_KV_CACHE_DEBUG
|
||||
int debug = 0;
|
||||
|
||||
// set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq)
|
||||
bool msa_strict_slots = false;
|
||||
|
||||
// this is the SWA type of the cache - not to be confused with the model SWA type
|
||||
const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
|
||||
|
||||
@@ -291,6 +298,7 @@ private:
|
||||
|
||||
size_t size_k_bytes() const;
|
||||
size_t size_v_bytes() const;
|
||||
size_t size_k_idx_bytes() const;
|
||||
|
||||
ggml_tensor * build_rope_shift(
|
||||
const llama_cparams & cparams,
|
||||
@@ -370,6 +378,7 @@ public:
|
||||
// get views of the current state of the cache
|
||||
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
|
||||
ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;
|
||||
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const;
|
||||
|
||||
// store k_cur and v_cur in the cache based on the provided head location
|
||||
// note: the heads in k_cur and v_cur should be laid out contiguously in memory
|
||||
@@ -379,6 +388,7 @@ public:
|
||||
// - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed
|
||||
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
|
||||
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const;
|
||||
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const;
|
||||
|
||||
// create destination indices for each head of the current batch for where it would be written in the KV cache
|
||||
// the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but
|
||||
|
||||
@@ -542,7 +542,7 @@ llama_model_loader::llama_model_loader(
|
||||
|
||||
tensor_buft_overrides = param_tensor_buft_overrides_p;
|
||||
|
||||
this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MLOCK;
|
||||
this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK;
|
||||
this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO;
|
||||
|
||||
if (!fname.empty()) {
|
||||
|
||||
@@ -281,6 +281,8 @@ void llama_model_saver::add_kv_from_model() {
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true);
|
||||
add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true);
|
||||
|
||||
|
||||
+5
-1
@@ -285,6 +285,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
|
||||
return new llama_model_apertus(params);
|
||||
case LLM_ARCH_MINIMAX_M2:
|
||||
return new llama_model_minimax_m2(params);
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
return new llama_model_minimax_m3(params);
|
||||
case LLM_ARCH_COGVLM:
|
||||
return new llama_model_cogvlm(params);
|
||||
case LLM_ARCH_PANGU_EMBED:
|
||||
@@ -818,6 +820,7 @@ const char * llm_type_name(llm_type type) {
|
||||
case LLM_TYPE_122B_A10B: return "122B.A10B";
|
||||
case LLM_TYPE_196B_A11B: return "196B.A11B";
|
||||
case LLM_TYPE_230B_A10B: return "230B.A10B";
|
||||
case LLM_TYPE_428B_A23B: return "428B.A23B";
|
||||
case LLM_TYPE_235B_A22B: return "235B.A22B";
|
||||
case LLM_TYPE_300B_A47B: return "300B.A47B";
|
||||
case LLM_TYPE_310B_A15B: return "310B.A15B";
|
||||
@@ -1246,7 +1249,7 @@ void llama_model_base::load_vocab(llama_model_loader & ml) {
|
||||
|
||||
bool llama_model_base::load_tensors(llama_model_loader & ml) {
|
||||
const auto & split_mode = params.split_mode;
|
||||
const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK;
|
||||
const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK || params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK;
|
||||
const auto & tensor_split = params.tensor_split;
|
||||
|
||||
const int n_layer_all = hparams.n_layer_all;
|
||||
@@ -2550,6 +2553,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
|
||||
case LLM_ARCH_GROVEMOE:
|
||||
case LLM_ARCH_APERTUS:
|
||||
case LLM_ARCH_MINIMAX_M2:
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
case LLM_ARCH_COGVLM:
|
||||
case LLM_ARCH_PANGU_EMBED:
|
||||
case LLM_ARCH_AFMOE:
|
||||
|
||||
@@ -134,6 +134,7 @@ enum llm_type {
|
||||
LLM_TYPE_122B_A10B, // Qwen3.5
|
||||
LLM_TYPE_196B_A11B, // Step3.5-Flash
|
||||
LLM_TYPE_230B_A10B, // Minimax M2
|
||||
LLM_TYPE_428B_A23B, // Minimax M3
|
||||
LLM_TYPE_235B_A22B,
|
||||
LLM_TYPE_300B_A47B, // Ernie MoE big
|
||||
LLM_TYPE_310B_A15B, // /MiMo-V2-Flash
|
||||
@@ -515,6 +516,12 @@ struct llama_layer {
|
||||
struct ggml_tensor * indexer_attn_k = nullptr;
|
||||
struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias
|
||||
|
||||
// MSA
|
||||
struct ggml_tensor * index_q_proj = nullptr;
|
||||
struct ggml_tensor * index_k_proj = nullptr;
|
||||
struct ggml_tensor * index_q_norm = nullptr;
|
||||
struct ggml_tensor * index_k_norm = nullptr;
|
||||
|
||||
// gemma4 layer output scale, reused for talkie embedding skip scale
|
||||
struct ggml_tensor * out_scale = nullptr;
|
||||
|
||||
|
||||
@@ -326,6 +326,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
|
||||
quantize &= name.find("ssm_conv1d") == std::string::npos;
|
||||
quantize &= name.find("shortconv.conv.weight") == std::string::npos;
|
||||
|
||||
// do not quantize MiniMax's indexer projection weights, they are tiny
|
||||
quantize &= name.find("indexer.k_proj.weight") == std::string::npos;
|
||||
quantize &= name.find("indexer.q_proj.weight") == std::string::npos;
|
||||
|
||||
// do not quantize RWKV's small yet 2D weights
|
||||
quantize &= name.find("time_mix_first.weight") == std::string::npos;
|
||||
quantize &= name.find("time_mix_w0.weight") == std::string::npos;
|
||||
|
||||
@@ -2809,6 +2809,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
|| t.first == "<turn|>" // gemma4
|
||||
|| t.first == "<|tool_response>" // gemma4
|
||||
|| t.first == "<｜end▁of▁sentence｜>" // deepseek-ocr
|
||||
|| t.first == "[e~[" // minimax-m2/m3
|
||||
) {
|
||||
special_eog_ids.insert(t.second);
|
||||
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
|
||||
|
||||
+7
-4
@@ -54,6 +54,8 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) {
|
||||
return "mmap";
|
||||
case LLAMA_LOAD_MODE_MLOCK:
|
||||
return "mlock";
|
||||
case LLAMA_LOAD_MODE_MMAP_MLOCK:
|
||||
return "mmap+mlock";
|
||||
case LLAMA_LOAD_MODE_DIRECT_IO:
|
||||
return "dio";
|
||||
}
|
||||
@@ -61,10 +63,11 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) {
|
||||
}
|
||||
|
||||
enum llama_load_mode llama_load_mode_from_str(const char * str) {
|
||||
if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; }
|
||||
if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; }
|
||||
if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; }
|
||||
if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; }
|
||||
if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; }
|
||||
if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; }
|
||||
if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; }
|
||||
if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; }
|
||||
if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; }
|
||||
throw std::invalid_argument(std::string("unknown load mode: ") + str);
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
#include "models.h"
|
||||
#include "llama-kv-cache.h"
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with
|
||||
// DeepSeek-V3 leading-dense + routed/shared experts (sigmoid gating, routed scaling),
|
||||
// swigluoai activation, and MiniMax Sparse Attention (MSA). MTP is not in released model weights.
|
||||
// Notes: Blocks are anchored to absolute KV cache slots.
|
||||
|
||||
void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
||||
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
|
||||
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
|
||||
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
|
||||
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
|
||||
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
|
||||
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
|
||||
msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks };
|
||||
hparams.indexer_kv = true;
|
||||
|
||||
switch (hparams.n_layer()) {
|
||||
case 60: type = LLM_TYPE_428B_A23B; break;
|
||||
default: type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
}
|
||||
|
||||
void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
const int64_t n_expert_shared = hparams.n_expert_shared;
|
||||
const int64_t n_ff_exp = hparams.n_ff_exp;
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
// output
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
for (int i = 0; i < n_layer; ++i) {
|
||||
auto & layer = layers[i];
|
||||
|
||||
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);
|
||||
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
||||
// per-head QK-norm: a single head_dim vector applied to every head
|
||||
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
|
||||
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
|
||||
|
||||
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
||||
|
||||
if (i < (int) hparams.n_layer_dense_lead) {
|
||||
// leading dense layers
|
||||
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
|
||||
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
|
||||
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
|
||||
} else {
|
||||
// routed experts
|
||||
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
|
||||
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
|
||||
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
|
||||
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
|
||||
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
|
||||
|
||||
// shared expert
|
||||
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
|
||||
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, 0);
|
||||
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
|
||||
|
||||
// indexer
|
||||
layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, 0);
|
||||
layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, 0);
|
||||
layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, 0);
|
||||
layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const {
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
// per-query local-force bias for MSA selection
|
||||
// local window always wins a slot
|
||||
class llm_graph_input_msa_local : public llm_graph_input_i {
|
||||
public:
|
||||
llm_graph_input_msa_local(int blk, int local, int64_t nblk) : blk(blk), local(local), nblk(nblk) {}
|
||||
|
||||
void set_input(const llama_ubatch * ubatch) override {
|
||||
if (!bias || !ubatch->pos) {
|
||||
return;
|
||||
}
|
||||
const int64_t n_tokens = ubatch->n_tokens;
|
||||
std::vector<float> data((size_t) nblk * n_tokens, 0.0f);
|
||||
for (int64_t i = 0; i < n_tokens; ++i) {
|
||||
const int64_t L = ubatch->pos[i] / blk;
|
||||
for (int l = 0; l < local && L - l >= 0; ++l) {
|
||||
if (L - l < nblk) {
|
||||
data[(size_t) i * nblk + (L - l)] = 1e30f;
|
||||
}
|
||||
}
|
||||
}
|
||||
ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float));
|
||||
}
|
||||
|
||||
// valid as long as the bias tensor dims still match the new ubatch/cache window
|
||||
bool can_reuse(const llm_graph_params & params) override {
|
||||
const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);
|
||||
|
||||
bool res = true;
|
||||
res &= bias->ne[1] == params.ubatch.n_tokens;
|
||||
res &= bias->ne[0] * blk == (int64_t) mctx->get_n_kv();
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_tensor * bias = nullptr;
|
||||
int blk;
|
||||
int local;
|
||||
int64_t nblk;
|
||||
};
|
||||
|
||||
// pooled score of a block with no visible token: -inf from the mask, or -FLT_MAX from the
|
||||
// max-pool identity when every element of the block is -inf
|
||||
static inline bool msa_score_masked(float x) { return x <= -1e30f; }
|
||||
|
||||
// MSA block selection (batch regime)
|
||||
// CPU custom op, the token-level expansion and the combination with the causal mask happen on the GPU.
|
||||
static void msa_block_mask_op(struct ggml_tensor * dst, int ith, int nth, void * userdata) {
|
||||
const struct ggml_tensor * bs = dst->src[0];
|
||||
const struct ggml_tensor * bias = dst->src[1];
|
||||
const msa_params * p = (const msa_params *) userdata;
|
||||
|
||||
const int nblk = (int) bs->ne[0];
|
||||
const int Hd = (int) bs->ne[1];
|
||||
const int S = (int) bs->ne[2];
|
||||
|
||||
GGML_ASSERT(bs->type == GGML_TYPE_F32 && ggml_is_contiguous(bs));
|
||||
GGML_ASSERT(bias->type == GGML_TYPE_F32 && ggml_is_contiguous(bias));
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F16 && ggml_is_contiguous(dst));
|
||||
GGML_ASSERT(dst->ne[0] == nblk && dst->ne[1] == S && dst->ne[2] == Hd);
|
||||
GGML_ASSERT(bias->ne[0] == nblk && bias->ne[1] == S);
|
||||
|
||||
const int topk = p->topk_blocks < nblk ? p->topk_blocks : nblk;
|
||||
|
||||
const ggml_fp16_t f16_zero = ggml_fp32_to_fp16(0.0f);
|
||||
const ggml_fp16_t f16_ninf = ggml_fp32_to_fp16(-INFINITY);
|
||||
|
||||
std::vector<float> rank(nblk);
|
||||
std::vector<char> valid(nblk);
|
||||
std::vector<int> ord(nblk);
|
||||
|
||||
ggml_fp16_t * out = (ggml_fp16_t *) dst->data;
|
||||
|
||||
for (int i = ith; i < S; i += nth) {
|
||||
const float * bias_col = (const float *) bias->data + (size_t) i * nblk;
|
||||
for (int h = 0; h < Hd; ++h) {
|
||||
const float * bs_col = (const float *) bs->data + ((size_t) i * Hd + h) * nblk;
|
||||
|
||||
for (int bk = 0; bk < nblk; ++bk) {
|
||||
// a block is selectable if it has a visible token or is locally forced
|
||||
valid[bk] = !msa_score_masked(bs_col[bk]) || bias_col[bk] > 0.0f;
|
||||
rank [bk] = bias_col[bk] > 0.0f ? bias_col[bk] : bs_col[bk];
|
||||
ord [bk] = bk;
|
||||
}
|
||||
|
||||
std::partial_sort(ord.begin(), ord.begin() + topk, ord.end(),
|
||||
[&](int a, int b) { return rank[a] > rank[b]; });
|
||||
|
||||
ggml_fp16_t * dst_col = out + ((size_t) h * S + i) * nblk;
|
||||
for (int bk = 0; bk < nblk; ++bk) {
|
||||
dst_col[bk] = f16_ninf;
|
||||
}
|
||||
for (int t = 0; t < topk; ++t) {
|
||||
const int bk = ord[t];
|
||||
if (!valid[bk]) {
|
||||
break; // sorted desc: first invalid -> fewer than topk selectable blocks
|
||||
}
|
||||
dst_col[bk] = f16_zero;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3])
|
||||
ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa(
|
||||
ggml_tensor * q_cur, // [D, HQ, T]
|
||||
ggml_tensor * k, // [D, n_keys, 1, C]
|
||||
ggml_tensor * v, // [D, n_keys, 1, C]
|
||||
ggml_tensor * mask, // [n_keys, R, 1, C] f16, contiguous
|
||||
int64_t Gp, float kq_scale, int il) const {
|
||||
|
||||
const int64_t D = q_cur->ne[0];
|
||||
const int64_t HQ = q_cur->ne[1];
|
||||
const int64_t T = q_cur->ne[2];
|
||||
const int64_t C = k->ne[3];
|
||||
const int64_t R = HQ*T/(Gp*C);
|
||||
GGML_ASSERT(Gp*C*R == HQ*T);
|
||||
GGML_ASSERT(mask->type == GGML_TYPE_F16);
|
||||
|
||||
// [D, HQ, T] -> [D, Gp, C, R] -> [D, R, Gp, C]
|
||||
// batch (C=HKV, R=T): channel = group
|
||||
// decode (C=HKV*ns, R=1): channel = (group, stream), group innermost
|
||||
ggml_tensor * q = ggml_reshape_4d(ctx0, q_cur, D, Gp, C, R);
|
||||
q = ggml_permute(ctx0, q, 0, 2, 3, 1);
|
||||
|
||||
ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale,
|
||||
hparams.f_max_alibi_bias, 0.0f);
|
||||
ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32);
|
||||
cb(o, "msa_fattn", il);
|
||||
|
||||
// [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T]
|
||||
o = ggml_permute(ctx0, o, 0, 1, 3, 2);
|
||||
if (!ggml_is_contiguous(o)) {
|
||||
o = ggml_cont(ctx0, o); // no-op layout at decode (R == 1), copy at batch
|
||||
}
|
||||
return ggml_reshape_2d(ctx0, o, D*HQ, T);
|
||||
}
|
||||
|
||||
llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
|
||||
const int64_t n_embd_head = hparams.n_embd_head_v();
|
||||
const auto & mm = static_cast<const llama_model_minimax_m3 &>(model);
|
||||
|
||||
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
|
||||
// partial rotary: head_dim != n_rot, so don't assert n_embd_head == n_rot
|
||||
|
||||
ggml_tensor * cur;
|
||||
ggml_tensor * inpL;
|
||||
|
||||
inpL = build_inp_embd(model.tok_embd);
|
||||
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
auto inp_attn = build_attn_inp_kv();
|
||||
|
||||
// MSA calls ggml_flash_attn_ext directly and assumes the non-transposed V layout that
|
||||
// llama.cpp only provides when flash attention is enabled. Block selection is anchored
|
||||
// to absolute KV cache slots, which equal positions only for append-only per-stream
|
||||
// caches either a single sequence, or multiple sequences with kv_unified == false (each
|
||||
// stream then has its own slot space). A unified cache with multiple sequences
|
||||
// interleaves slots and would silently break block anchoring so it falls back to dense.
|
||||
const bool fa_on = cparams.flash_attn;
|
||||
const bool streams_ok = cparams.n_seq_max == 1 || !cparams.kv_unified;
|
||||
const bool msa_enabled = fa_on && streams_ok;
|
||||
|
||||
static bool warned_no_fa = false;
|
||||
if (!fa_on && !warned_no_fa) {
|
||||
LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention "
|
||||
"(output may be degraded). Enable flash attention for MSA.\n", __func__);
|
||||
warned_no_fa = true;
|
||||
}
|
||||
static bool warned_unified = false;
|
||||
if (fa_on && !streams_ok && !warned_unified) {
|
||||
LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; MSA needs per-sequence streams "
|
||||
"-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__);
|
||||
warned_unified = true;
|
||||
}
|
||||
|
||||
// hoisted per-graph MSA state (shared by every sparse layer)
|
||||
llm_graph_input_msa_local * msa_loc = nullptr;
|
||||
ggml_tensor * msa_kqm = nullptr;
|
||||
ggml_tensor * msa_mf = nullptr;
|
||||
int64_t n_kv = 0, nblk = 0, ns = 1, n_tps = 0;
|
||||
bool msa_decode = false; // gather (1 token per stream) vs mask
|
||||
const int blk = mm.msa_p.blk;
|
||||
const int64_t Hd = hparams.indexer_n_head; // one indexer head per GQA group
|
||||
|
||||
if (msa_enabled) {
|
||||
msa_kqm = inp_attn->get_kq_mask();
|
||||
n_kv = msa_kqm->ne[0];
|
||||
n_tps = msa_kqm->ne[1]; // tokens per stream
|
||||
ns = msa_kqm->ne[3]; // streams in this ubatch
|
||||
GGML_ASSERT(msa_kqm->type == GGML_TYPE_F16 && "MSA requires the FA (f16) mask");
|
||||
GGML_ASSERT(n_tps*ns == n_tokens);
|
||||
GGML_ASSERT(n_kv % blk == 0 &&
|
||||
"MSA: KV/mask n_kv must be a multiple of indexer.block_size (128); "
|
||||
"the flash-attention KV padding must be a multiple of the block size. "
|
||||
"A non-multiple would silently drop the partial tail block.");
|
||||
nblk = n_kv / blk;
|
||||
msa_decode = n_tps == 1;
|
||||
|
||||
msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32);
|
||||
|
||||
auto loc = std::make_unique<llm_graph_input_msa_local>(blk, mm.msa_p.local, nblk);
|
||||
loc->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens); // stream-grouped tokens
|
||||
ggml_set_input(loc->bias);
|
||||
msa_loc = (llm_graph_input_msa_local *) res->add_input(std::move(loc));
|
||||
}
|
||||
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
ggml_tensor * inpSA = inpL;
|
||||
|
||||
// self-attention
|
||||
{
|
||||
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
|
||||
n_embd_head, n_head, n_head_kv, il);
|
||||
|
||||
// per-head QK RMSNorm (weights already include Gemma's +1)
|
||||
Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(Qcur, "Qcur_normed", il);
|
||||
Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(Kcur, "Kcur_normed", il);
|
||||
|
||||
// partial rotary: only the first n_rot dims are rotated
|
||||
Qcur = ggml_rope_ext(
|
||||
ctx0, Qcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
Kcur = ggml_rope_ext(
|
||||
ctx0, Kcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
|
||||
cb(Qcur, "Qcur", il);
|
||||
cb(Kcur, "Kcur", il);
|
||||
cb(Vcur, "Vcur", il);
|
||||
|
||||
const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead;
|
||||
|
||||
if (!is_sparse) {
|
||||
cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s,
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,
|
||||
1.0f/sqrtf(float(n_embd_head)), il);
|
||||
} else {
|
||||
const int64_t n_idx_dim = hparams.indexer_head_size; // 128
|
||||
|
||||
GGML_ASSERT(!inp_attn->self_k_rot && !inp_attn->self_v_rot && "MSA: attn-rot not supported");
|
||||
|
||||
// Index Branch, project, norm, partial RoPE, cache
|
||||
ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur);
|
||||
ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur);
|
||||
iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, Hd, n_tokens);
|
||||
ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1, n_tokens);
|
||||
iq = build_norm(iq, model.layers[il].index_q_norm, NULL, LLM_NORM_RMS, il); // +1 baked
|
||||
ik = build_norm(ik, model.layers[il].index_k_norm, NULL, LLM_NORM_RMS, il);
|
||||
iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
|
||||
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,
|
||||
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
|
||||
const auto * mctx_cur = inp_attn->mctx;
|
||||
ggml_build_forward_expand(gf, mctx_cur->cpy_k_idx(ctx0, ik, inp_attn->get_k_idxs(), il));
|
||||
ggml_tensor * ik_kv = mctx_cur->get_k_idx(ctx0, il);
|
||||
|
||||
// Main branch: store K/V, take cache views
|
||||
ggml_build_forward_expand(gf, Qcur);
|
||||
ggml_build_forward_expand(gf, Kcur);
|
||||
ggml_build_forward_expand(gf, Vcur);
|
||||
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));
|
||||
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));
|
||||
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
|
||||
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
|
||||
GGML_ASSERT(!(v->nb[1] > v->nb[2]) && "MSA assumes v_trans=false (FA on)");
|
||||
|
||||
const int64_t D = k->ne[0];
|
||||
const int64_t HKV = k->ne[1];
|
||||
const int64_t Gp = n_head/HKV;
|
||||
GGML_ASSERT(HKV == Hd && "MSA: one indexer head per GQA group");
|
||||
GGML_ASSERT(k->ne[3] == ns);
|
||||
const int K = mm.msa_p.topk_blocks < (int) nblk ? mm.msa_p.topk_blocks : (int) nblk;
|
||||
|
||||
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
|
||||
|
||||
if (msa_decode) {
|
||||
// decode: batched over streams top-k + gather, one grouped FA
|
||||
// scores: per-stream batched matmul over the stream dim (ne[3]).
|
||||
// the cache views are not contiguous across streams (stride = kv_size, not n_kv)
|
||||
ggml_tensor * ikv4 = ggml_view_4d(ctx0, ik_kv, n_idx_dim, n_kv, 1, ns,
|
||||
ik_kv->nb[2], ik_kv->nb[3], ik_kv->nb[3], 0);
|
||||
ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns);
|
||||
ggml_tensor * sc = ggml_mul_mat(ctx0, ikv4, iq4);
|
||||
ggml_mul_mat_set_prec(sc, GGML_PREC_F32);
|
||||
sc = ggml_add_inplace(ctx0, sc, msa_mf);
|
||||
ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
|
||||
cb(bs, "msa_bs", il);
|
||||
|
||||
ggml_tensor * bsf = ggml_add(ctx0, bs,
|
||||
ggml_reshape_4d(ctx0, msa_loc->bias, nblk, 1, 1, ns));
|
||||
ggml_tensor * idx = ggml_top_k(ctx0, bsf, K);
|
||||
|
||||
// token idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (for the mask gather)
|
||||
// row idx: tr[t,k,h,s] = tj*HKV + h (for the per-stream K/V gather)
|
||||
ggml_tensor * a = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk);
|
||||
a = ggml_reshape_4d(ctx0, a, 1, K, Hd, ns);
|
||||
ggml_tensor * tj = ggml_add(ctx0,
|
||||
ggml_repeat_4d(ctx0, a, blk, K, Hd, ns),
|
||||
ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) blk, 1.0f), blk, 1, 1));
|
||||
ggml_tensor * tr = ggml_add(ctx0,
|
||||
ggml_scale(ctx0, tj, (float) HKV),
|
||||
ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd));
|
||||
|
||||
ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);
|
||||
ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);
|
||||
|
||||
ggml_tensor * k3 = ggml_view_3d(ctx0, k, D, HKV*n_kv, ns, k->nb[1], k->nb[3], 0);
|
||||
ggml_tensor * v3 = ggml_view_3d(ctx0, v, D, HKV*n_kv, ns, v->nb[1], v->nb[3], 0);
|
||||
ggml_tensor * m3 = ggml_reshape_3d(ctx0, msa_kqm, 1, n_kv, ns);
|
||||
|
||||
ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr);
|
||||
ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr);
|
||||
ggml_tensor * mg = ggml_get_rows(ctx0, m3, tokj);
|
||||
|
||||
// fold (group, stream) onto the FA channel dim
|
||||
const ggml_type kt = ggml_is_quantized(k->type) ? GGML_TYPE_F16 : k->type;
|
||||
const ggml_type vt = ggml_is_quantized(v->type) ? GGML_TYPE_F16 : v->type;
|
||||
ggml_tensor * kfa = ggml_reshape_4d(ctx0, kg, D, (int64_t) blk*K, 1, Hd*ns);
|
||||
ggml_tensor * vfa = ggml_reshape_4d(ctx0, vg, D, (int64_t) blk*K, 1, Hd*ns);
|
||||
if (kfa->type != kt) { kfa = ggml_cast(ctx0, kfa, kt); }
|
||||
if (vfa->type != vt) { vfa = ggml_cast(ctx0, vfa, vt); }
|
||||
// the FA mask must be F16
|
||||
ggml_tensor * mfa = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, (int64_t) blk*K, 1, 1, Hd*ns), GGML_TYPE_F16);
|
||||
|
||||
cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il);
|
||||
} else {
|
||||
// batch: per-stream loop
|
||||
std::vector<ggml_tensor *> outs(ns);
|
||||
for (int64_t st = 0; st < ns; ++st) {
|
||||
ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, n_tps,
|
||||
iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]);
|
||||
ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv,
|
||||
ik_kv->nb[2], st*ik_kv->nb[3]);
|
||||
ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, 1, n_tps,
|
||||
msa_mf->nb[1], msa_mf->nb[1], st*msa_mf->nb[3]);
|
||||
ggml_tensor * km_s = ggml_view_3d(ctx0, msa_kqm, n_kv, n_tps, 1,
|
||||
msa_kqm->nb[1], msa_kqm->nb[3], st*msa_kqm->nb[3]);
|
||||
ggml_tensor * bias_s = ggml_view_2d(ctx0, msa_loc->bias, nblk, n_tps,
|
||||
msa_loc->bias->nb[1], st*n_tps*msa_loc->bias->nb[1]);
|
||||
ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps,
|
||||
Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]);
|
||||
ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1,
|
||||
k->nb[1], k->nb[2], k->nb[3], st*k->nb[3]);
|
||||
ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1,
|
||||
v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]);
|
||||
|
||||
// block scores: bs = maxpool_blk(idx_q * idx_k^T + causal mask)
|
||||
// scores are unscaled, only the top-k ordering matters
|
||||
ggml_tensor * sc = ggml_mul_mat(ctx0, ik_s,
|
||||
ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps));
|
||||
// indexer scores run in F32
|
||||
ggml_mul_mat_set_prec(sc, GGML_PREC_F32);
|
||||
sc = ggml_reshape_3d(ctx0, sc, n_kv, Hd, n_tps);
|
||||
sc = ggml_add_inplace(ctx0, sc, mf_s);
|
||||
ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
|
||||
cb(bs, "msa_bs", il);
|
||||
|
||||
// block-level 0/-inf keep mask on the CPU, tiny transfer
|
||||
ggml_tensor * srcs[2] = { bs, bias_s };
|
||||
ggml_tensor * bm = ggml_custom_4d(ctx0, GGML_TYPE_F16,
|
||||
nblk, n_tps, Hd, 1,
|
||||
srcs, 2, msa_block_mask_op, GGML_N_TASKS_MAX,
|
||||
const_cast<msa_params *>(&mm.msa_p));
|
||||
cb(bm, "msa_block_mask", il);
|
||||
|
||||
// expand block -> token granularity on the GPU (j = bk*blk + t),
|
||||
// then combine with the causal mask in place
|
||||
ggml_tensor * bmx = ggml_repeat_4d(ctx0,
|
||||
ggml_reshape_3d(ctx0, bm, 1, nblk, n_tps*Hd),
|
||||
blk, nblk, n_tps*Hd, 1);
|
||||
bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd);
|
||||
ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, km_s);
|
||||
mask4 = ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd);
|
||||
cb(mask4, "msa_mask4", il);
|
||||
|
||||
// cache views with groups on ne[3];
|
||||
ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2);
|
||||
ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2);
|
||||
|
||||
outs[st] = build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il);
|
||||
}
|
||||
cur = outs[0];
|
||||
for (int64_t st = 1; st < ns; ++st) {
|
||||
cur = ggml_concat(ctx0, cur, outs[st], 1);
|
||||
}
|
||||
}
|
||||
|
||||
cb(cur, "kqv_out", il);
|
||||
if (model.layers[il].wo) {
|
||||
cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (il == n_layer - 1 && inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
}
|
||||
|
||||
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
||||
cb(ffn_inp, "ffn_inp", il);
|
||||
|
||||
cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
if ((uint32_t) il < hparams.n_layer_dense_lead) {
|
||||
// leading dense FFN (swigluoai)
|
||||
cur = build_ffn(cur,
|
||||
model.layers[il].ffn_up, NULL, NULL,
|
||||
model.layers[il].ffn_gate, NULL, NULL,
|
||||
model.layers[il].ffn_down, NULL, NULL,
|
||||
NULL,
|
||||
LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
|
||||
cb(cur, "ffn_out", il);
|
||||
} else {
|
||||
// routed experts (swigluoai MoE)
|
||||
ggml_tensor * moe_out = build_moe_ffn(cur,
|
||||
model.layers[il].ffn_gate_inp,
|
||||
model.layers[il].ffn_up_exps,
|
||||
model.layers[il].ffn_gate_exps,
|
||||
model.layers[il].ffn_down_exps,
|
||||
model.layers[il].ffn_exp_probs_b,
|
||||
n_expert, n_expert_used,
|
||||
LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm,
|
||||
hparams.expert_weights_scale,
|
||||
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
||||
il);
|
||||
cb(moe_out, "ffn_moe_out", il);
|
||||
|
||||
// shared expert (swigluoai)
|
||||
ggml_tensor * ffn_shexp = build_ffn(cur,
|
||||
model.layers[il].ffn_up_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_gate_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_down_shexp, NULL, NULL,
|
||||
NULL,
|
||||
LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "ffn_shexp", il);
|
||||
|
||||
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
||||
cb(cur, "ffn_out", il);
|
||||
}
|
||||
|
||||
cur = ggml_add(ctx0, cur, ffn_inp);
|
||||
|
||||
cur = build_cvec(cur, il);
|
||||
cb(cur, "l_out", il);
|
||||
|
||||
// input for next layer
|
||||
inpL = cur;
|
||||
}
|
||||
|
||||
cur = inpL;
|
||||
|
||||
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
|
||||
cb(cur, "result_norm", -1);
|
||||
res->t_embd = cur;
|
||||
|
||||
// lm_head
|
||||
cur = build_lora_mm(model.output, cur, model.output_s);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
@@ -1902,6 +1902,29 @@ struct llama_model_minimax_m2 : public llama_model_base {
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
struct msa_params {
|
||||
int blk;
|
||||
int topk_blocks;
|
||||
int local;
|
||||
};
|
||||
|
||||
struct llama_model_minimax_m3 : public llama_model_base {
|
||||
llama_model_minimax_m3(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
void load_arch_tensors(llama_model_loader & ml) override;
|
||||
msa_params msa_p;
|
||||
struct graph : public llm_graph_context {
|
||||
graph(const llama_model & model, const llm_graph_params & params);
|
||||
|
||||
ggml_tensor * build_attn_msa_fa(
|
||||
ggml_tensor * q_cur, // [D, HQ, S] f32
|
||||
ggml_tensor * k, // [D, n_keys, 1, C] C = HKV or HKV*n_stream
|
||||
ggml_tensor * v, // [D, n_keys, 1, C]
|
||||
ggml_tensor * mask, // [n_keys, R, 1, C] f16, R = HQ*T/(Gp*C)
|
||||
int64_t Gp, float kq_scale, int il) const;
|
||||
};
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
struct llama_model_cogvlm : public llama_model_base {
|
||||
llama_model_cogvlm(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
|
||||
@@ -143,6 +143,10 @@ static void test(void) {
|
||||
assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
|
||||
assert(params.load_mode == LLAMA_LOAD_MODE_MLOCK);
|
||||
|
||||
argv = {"binary_name", "-lm", "mmap+mlock"};
|
||||
assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
|
||||
assert(params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK);
|
||||
|
||||
argv = {"binary_name", "-lm", "dio"};
|
||||
assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
|
||||
assert(params.load_mode == LLAMA_LOAD_MODE_DIRECT_IO);
|
||||
@@ -187,6 +191,11 @@ static void test(void) {
|
||||
assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
|
||||
assert(params.load_mode == LLAMA_LOAD_MODE_MLOCK);
|
||||
|
||||
setenv("LLAMA_ARG_LOAD_MODE", "mmap+mlock", true);
|
||||
argv = {"binary_name"};
|
||||
assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
|
||||
assert(params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK);
|
||||
|
||||
setenv("LLAMA_ARG_LOAD_MODE", "dio", true);
|
||||
argv = {"binary_name"};
|
||||
assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON));
|
||||
|
||||
+9
-7
@@ -1144,7 +1144,7 @@ static void test_peg_parser(common_chat_templates * tmpls,
|
||||
// budget sampler inhibits grammar application while inside thinking blocks —
|
||||
// triggers inside <think>...</think> are suppressed.
|
||||
bool use_reasoning_budget_path = false;
|
||||
if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tag.empty()) {
|
||||
if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tags.empty()) {
|
||||
use_reasoning_budget_path = true;
|
||||
for (const auto & trigger : parser.params_.grammar_triggers) {
|
||||
if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) {
|
||||
@@ -1162,7 +1162,7 @@ static void test_peg_parser(common_chat_templates * tmpls,
|
||||
// Walk through full_input tracking thinking state; only match triggers
|
||||
// when outside thinking blocks.
|
||||
const auto & think_start = parser.params_.thinking_start_tag;
|
||||
const auto & think_end = parser.params_.thinking_end_tag;
|
||||
const auto & think_ends = parser.params_.thinking_end_tags;
|
||||
|
||||
bool in_thinking = false;
|
||||
for (size_t i = 0; i < full_input.size(); ++i) {
|
||||
@@ -1172,12 +1172,14 @@ static void test_peg_parser(common_chat_templates * tmpls,
|
||||
i += think_start.size() - 1;
|
||||
continue;
|
||||
}
|
||||
if (in_thinking && full_input.compare(i, think_end.size(), think_end) == 0) {
|
||||
in_thinking = false;
|
||||
i += think_end.size() - 1;
|
||||
continue;
|
||||
}
|
||||
if (in_thinking) {
|
||||
for (const auto & think_end : think_ends) {
|
||||
if (full_input.compare(i, think_end.size(), think_end) == 0) {
|
||||
in_thinking = false;
|
||||
i += think_end.size() - 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// Outside thinking — check if any trigger word starts here
|
||||
|
||||
+11
-14
@@ -4,7 +4,7 @@
|
||||
#include <cstdlib>
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <sheredom/subprocess.h>
|
||||
#include "subproc.h"
|
||||
|
||||
#include "jinja/runtime.h"
|
||||
#include "jinja/parser.h"
|
||||
@@ -2135,21 +2135,20 @@ static void test_template_py(testing & t, const std::string & name, const std::s
|
||||
const char * python_executable = "python3";
|
||||
#endif
|
||||
|
||||
const char * command_line[] = {python_executable, "-c", py_script.c_str(), NULL};
|
||||
std::vector<std::string> args = {python_executable, "-c", py_script, };
|
||||
|
||||
struct subprocess_s subprocess;
|
||||
common_subproc subprocess;
|
||||
int options = subprocess_option_combined_stdout_stderr
|
||||
| subprocess_option_no_window
|
||||
| subprocess_option_inherit_environment
|
||||
| subprocess_option_search_user_path;
|
||||
int result = subprocess_create(command_line, options, &subprocess);
|
||||
|
||||
if (result != 0) {
|
||||
t.log("Failed to create subprocess, error code: " + std::to_string(result));
|
||||
if (!subprocess.create(args, options)) {
|
||||
t.log("Failed to create subprocess");
|
||||
t.assert_true("subprocess creation", false);
|
||||
return;
|
||||
}
|
||||
FILE * p_stdin = subprocess_stdin(&subprocess);
|
||||
FILE * p_stdin = subprocess.stdin_file();
|
||||
|
||||
// Write input
|
||||
std::string input = merged.dump();
|
||||
@@ -2157,24 +2156,22 @@ static void test_template_py(testing & t, const std::string & name, const std::s
|
||||
if (written != input.size()) {
|
||||
t.log("Failed to write complete input to subprocess stdin");
|
||||
t.assert_true("subprocess stdin write", false);
|
||||
subprocess_destroy(&subprocess);
|
||||
subprocess.close_stdin();
|
||||
subprocess.join();
|
||||
return;
|
||||
}
|
||||
fflush(p_stdin);
|
||||
fclose(p_stdin); // Close stdin to signal EOF to the Python process
|
||||
subprocess.stdin_file = nullptr;
|
||||
subprocess.close_stdin(); // Close stdin to signal EOF to the Python process
|
||||
|
||||
// Read output
|
||||
std::string output;
|
||||
char buffer[1024];
|
||||
FILE * p_stdout = subprocess_stdout(&subprocess);
|
||||
FILE * p_stdout = subprocess.stdout_file();
|
||||
while (fgets(buffer, sizeof(buffer), p_stdout)) {
|
||||
output += buffer;
|
||||
}
|
||||
|
||||
int process_return;
|
||||
subprocess_join(&subprocess, &process_return);
|
||||
subprocess_destroy(&subprocess);
|
||||
int process_return = subprocess.join();
|
||||
|
||||
if (process_return != 0) {
|
||||
t.log("Python script failed with exit code: " + std::to_string(process_return));
|
||||
|
||||
@@ -168,6 +168,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64));
|
||||
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192));
|
||||
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128));
|
||||
} else if (arch == LLM_ARCH_MINIMAX_M3) {
|
||||
// partial rotary: n_rot must not exceed the indexer key length (64)
|
||||
ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64));
|
||||
}
|
||||
ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV, 1.0f);
|
||||
ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, 1e-5f);
|
||||
@@ -198,9 +201,13 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2));
|
||||
}
|
||||
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, uint32_t(1));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8));
|
||||
// MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the
|
||||
// indexer head count is independent of the main attention head count.
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(1));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4));
|
||||
ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1));
|
||||
ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector<uint32_t>({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4}));
|
||||
ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab");
|
||||
// ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd);
|
||||
@@ -355,6 +362,7 @@ static bool moe_mandatory(const llm_arch arch) {
|
||||
case LLM_ARCH_LLADA_MOE:
|
||||
case LLM_ARCH_GROVEMOE:
|
||||
case LLM_ARCH_MINIMAX_M2:
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
case LLM_ARCH_RND1:
|
||||
case LLM_ARCH_PADDLEOCR:
|
||||
case LLM_ARCH_MIMO2:
|
||||
|
||||
+130
-19
@@ -20,8 +20,8 @@
|
||||
static void test_reasoning_budget(
|
||||
const char * test_name,
|
||||
const std::vector<llama_token> & sequence,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state,
|
||||
@@ -31,8 +31,12 @@ static void test_reasoning_budget(
|
||||
// Find the maximum token ID to ensure our vocab covers all tokens
|
||||
llama_token max_token = 0;
|
||||
for (auto t : sequence) max_token = std::max(max_token, t);
|
||||
for (auto t : start_tokens) max_token = std::max(max_token, t);
|
||||
for (auto t : end_tokens) max_token = std::max(max_token, t);
|
||||
for (const auto & seq : start_seqs) {
|
||||
for (auto t : seq) max_token = std::max(max_token, t);
|
||||
}
|
||||
for (const auto & seq : end_seqs) {
|
||||
for (auto t : seq) max_token = std::max(max_token, t);
|
||||
}
|
||||
for (auto t : forced_tokens) max_token = std::max(max_token, t);
|
||||
|
||||
// Create a minimal sampler with mock vocabulary
|
||||
@@ -40,8 +44,8 @@ static void test_reasoning_budget(
|
||||
// The UTF-8 boundary check will treat all tokens as complete (safe fallback)
|
||||
auto * sampler = common_reasoning_budget_init(
|
||||
nullptr, // vocab - not used for basic state machine tests
|
||||
start_tokens,
|
||||
end_tokens,
|
||||
start_seqs,
|
||||
end_seqs,
|
||||
forced_tokens,
|
||||
budget,
|
||||
initial_state
|
||||
@@ -152,7 +156,7 @@ static void test_reasoning_budget_clone_mid_counting() {
|
||||
const std::vector<llama_token> end = {101};
|
||||
const std::vector<llama_token> 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, REASONING_BUDGET_IDLE);
|
||||
|
||||
llama_sampler_accept(sampler, 100); // COUNTING, remaining=2
|
||||
llama_sampler_accept(sampler, 50); // COUNTING, remaining=1
|
||||
@@ -171,7 +175,7 @@ static void test_reasoning_budget_clone_mid_forcing() {
|
||||
const std::vector<llama_token> end = {101};
|
||||
const std::vector<llama_token> 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, REASONING_BUDGET_FORCING);
|
||||
|
||||
GGML_ASSERT(get_forced_token(sampler, 102) == 102);
|
||||
llama_sampler_accept(sampler, 102); // advance to the second forced token
|
||||
@@ -191,7 +195,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, REASONING_BUDGET_IDLE);
|
||||
|
||||
llama_sampler_accept(sampler, 100); // COUNTING, remaining=5
|
||||
llama_sampler_accept(sampler, 50); // COUNTING, remaining=4
|
||||
@@ -212,7 +216,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, 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);
|
||||
@@ -222,7 +226,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, REASONING_BUDGET_IDLE);
|
||||
|
||||
llama_sampler_accept(sampler, 100); // COUNTING
|
||||
llama_sampler_accept(sampler, 101); // natural end -> DONE
|
||||
@@ -236,7 +240,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, 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)
|
||||
@@ -254,6 +258,81 @@ static void test_reasoning_budget_force_manual() {
|
||||
fprintf(stderr, " Test 'manual force transition' passed\n");
|
||||
}
|
||||
|
||||
static void test_reasoning_budget_end_match() {
|
||||
const std::vector<llama_tokens> start = {{100}};
|
||||
const std::vector<llama_tokens> 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);
|
||||
|
||||
GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
|
||||
|
||||
llama_sampler_accept(sampler, 100); // COUNTING
|
||||
llama_sampler_accept(sampler, 50);
|
||||
llama_sampler_accept(sampler, 103);
|
||||
llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE
|
||||
|
||||
const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
|
||||
GGML_ASSERT(matched != nullptr);
|
||||
GGML_ASSERT(*matched == llama_tokens({103, 104}));
|
||||
|
||||
llama_sampler_accept(sampler, 100); // re-arm, COUNTING
|
||||
GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
|
||||
|
||||
llama_sampler_free(sampler);
|
||||
}
|
||||
|
||||
// overlapping end sequences: the longest one ending at the position wins
|
||||
{
|
||||
const std::vector<llama_tokens> end_overlap = {{104}, {103, 104}};
|
||||
|
||||
auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE);
|
||||
|
||||
llama_sampler_accept(sampler, 100); // COUNTING
|
||||
llama_sampler_accept(sampler, 103);
|
||||
llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here
|
||||
|
||||
const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
|
||||
GGML_ASSERT(matched != nullptr);
|
||||
GGML_ASSERT(*matched == llama_tokens({103, 104}));
|
||||
|
||||
llama_sampler_free(sampler);
|
||||
}
|
||||
|
||||
// forcing records the end sequence terminating forced_tokens
|
||||
{
|
||||
auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING);
|
||||
|
||||
llama_sampler_accept(sampler, 102);
|
||||
llama_sampler_accept(sampler, 103);
|
||||
GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
|
||||
llama_sampler_accept(sampler, 104); // forced sequence complete, DONE
|
||||
|
||||
const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
|
||||
GGML_ASSERT(matched != nullptr);
|
||||
GGML_ASSERT(*matched == llama_tokens({103, 104}));
|
||||
|
||||
llama_sampler_free(sampler);
|
||||
}
|
||||
|
||||
// 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);
|
||||
|
||||
llama_sampler_accept(sampler, 102); // forced sequence complete, DONE
|
||||
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);
|
||||
GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
|
||||
|
||||
llama_sampler_free(sampler);
|
||||
}
|
||||
|
||||
// a null sampler is safely ignored
|
||||
GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr);
|
||||
|
||||
fprintf(stderr, " Test 'matched end sequence' passed\n");
|
||||
}
|
||||
|
||||
// UTF-8 boundary detection unit test
|
||||
// Tests common_utf8_is_complete() from reasoning-budget.h
|
||||
static void test_utf8_boundary_detection() {
|
||||
@@ -290,7 +369,7 @@ int main(void) {
|
||||
const std::vector<llama_token> forced = {102}; // forced token (not used in this test)
|
||||
const std::vector<llama_token> sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more
|
||||
|
||||
test_reasoning_budget("natural end before budget exhausted", sequence, start, end, forced,
|
||||
test_reasoning_budget("natural end before budget exhausted", sequence, {start}, {end}, forced,
|
||||
5, // budget of 5 tokens
|
||||
REASONING_BUDGET_IDLE,
|
||||
SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)
|
||||
@@ -306,7 +385,7 @@ int main(void) {
|
||||
const std::vector<llama_token> forced = {102, 101}; // forced message + end
|
||||
const std::vector<llama_token> sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2)
|
||||
|
||||
test_reasoning_budget("budget exhausted forcing", sequence, start, end, forced,
|
||||
test_reasoning_budget("budget exhausted forcing", sequence, {start}, {end}, forced,
|
||||
2, // budget of 2 tokens
|
||||
REASONING_BUDGET_IDLE,
|
||||
3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces)
|
||||
@@ -321,7 +400,7 @@ int main(void) {
|
||||
const std::vector<llama_token> forced = {102, 101};
|
||||
const std::vector<llama_token> sequence = {100, 50, 51, 52}; // start token first, then 3 tokens
|
||||
|
||||
test_reasoning_budget("activate immediately budget=0", sequence, start, end, forced,
|
||||
test_reasoning_budget("activate immediately budget=0", sequence, {start}, {end}, forced,
|
||||
0, // budget of 0 tokens
|
||||
REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0
|
||||
0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately)
|
||||
@@ -335,7 +414,7 @@ int main(void) {
|
||||
const std::vector<llama_token> forced = {102};
|
||||
const std::vector<llama_token> sequence = {50, 51, 52, 53};
|
||||
|
||||
test_reasoning_budget("no start/end configured", sequence, start, end, forced,
|
||||
test_reasoning_budget("no start/end configured", sequence, {start}, {end}, forced,
|
||||
2, // budget
|
||||
REASONING_BUDGET_IDLE,
|
||||
SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured)
|
||||
@@ -350,7 +429,7 @@ int main(void) {
|
||||
const std::vector<llama_token> forced = {102, 101};
|
||||
const std::vector<llama_token> sequence = {50, 51, 52, 53};
|
||||
|
||||
test_reasoning_budget("activate immediately with budget", sequence, start, end, forced,
|
||||
test_reasoning_budget("activate immediately with budget", sequence, {start}, {end}, forced,
|
||||
2, // budget of 2 tokens
|
||||
REASONING_BUDGET_COUNTING,
|
||||
2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces)
|
||||
@@ -373,18 +452,50 @@ int main(void) {
|
||||
const std::vector<llama_token> forced = {102, 101};
|
||||
const std::vector<llama_token> sequence = {100, 50, 101, 100, 60, 61, 62, 63};
|
||||
|
||||
test_reasoning_budget("multi-block re-arms budget after DONE", sequence, start, end, forced,
|
||||
test_reasoning_budget("multi-block re-arms budget after DONE", sequence, {start}, {end}, forced,
|
||||
2, // budget of 2 tokens (per block)
|
||||
REASONING_BUDGET_IDLE,
|
||||
6, // forcing starts at i=6 (after second block exhausts at i=5)
|
||||
7); // forcing continues through i=7
|
||||
}
|
||||
|
||||
// Test 7: Multiple start sequences - the second sequence activates counting
|
||||
// Flow: i=0 accept(110), i=1 accept(111)->COUNTING rem=2; i=2 accept(50)->rem=1;
|
||||
// i=3 accept(51)->rem=0->FORCING; i=4..5 apply() forces the end sequence
|
||||
{
|
||||
const std::vector<llama_tokens> start = {{100}, {110, 111}};
|
||||
const std::vector<llama_tokens> end = {{101}};
|
||||
const std::vector<llama_token> forced = {102, 101};
|
||||
const std::vector<llama_token> sequence = {110, 111, 50, 51, 52, 53};
|
||||
|
||||
test_reasoning_budget("multiple start sequences", sequence, start, end, forced,
|
||||
2, // budget of 2 tokens
|
||||
REASONING_BUDGET_IDLE,
|
||||
4, // forcing starts at i=4 (accept at i=3 depletes budget)
|
||||
5); // forcing continues through i=5
|
||||
}
|
||||
|
||||
// Test 8: Multiple end sequences - natural end via the second sequence
|
||||
// Flow: i=0 accept(100)->COUNTING rem=5; i=1 accept(50)->rem=4;
|
||||
// i=2 accept(103)->partial end, rem=3; i=3 accept(104)->end matched, DONE
|
||||
{
|
||||
const std::vector<llama_tokens> start = {{100}};
|
||||
const std::vector<llama_tokens> end = {{101}, {103, 104}};
|
||||
const std::vector<llama_token> forced = {102, 101};
|
||||
const std::vector<llama_token> sequence = {100, 50, 103, 104, 52};
|
||||
|
||||
test_reasoning_budget("multiple end sequences", sequence, start, end, forced,
|
||||
5, // budget of 5 tokens
|
||||
REASONING_BUDGET_IDLE,
|
||||
SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)
|
||||
}
|
||||
|
||||
test_reasoning_budget_clone_mid_counting();
|
||||
test_reasoning_budget_clone_mid_forcing();
|
||||
test_reasoning_budget_force_manual();
|
||||
test_reasoning_budget_end_match();
|
||||
|
||||
printf("OK (9 tests passed)\n");
|
||||
printf("OK (12 tests passed)\n");
|
||||
|
||||
printf("Testing UTF-8 boundary detection... ");
|
||||
test_utf8_boundary_detection();
|
||||
|
||||
@@ -44,6 +44,8 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i
|
||||
n_past++;
|
||||
}
|
||||
|
||||
llama_synchronize(ctx);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
+2
-2
@@ -55,10 +55,10 @@
|
||||
| `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
|
||||
| `-np, --parallel N` | number of parallel sequences to decode (default: 1)<br/>(env: LLAMA_ARG_N_PARALLEL) |
|
||||
| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
|
||||
| `--mlock` | DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
|
||||
| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
|
||||
| `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
|
||||
| `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
|
||||
| `-lm, --load-mode MODE` | model loading mode (default: mmap)<br/>- none: no special loading mode<br/>- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>- mlock: mmap + force system to keep model in RAM rather than swapping or compressing<br/>- dio: use DirectIO if available<br/><br/>(env: LLAMA_ARG_LOAD_MODE) |
|
||||
| `-lm, --load-mode MODE` | model loading mode (default: mmap)<br/>- none: no special loading mode<br/>- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>- mlock: force system to keep model in RAM rather than swapping or compressing<br/>- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing<br/>- dio: use DirectIO if available<br/><br/>(env: LLAMA_ARG_LOAD_MODE) |
|
||||
| `--numa TYPE` | attempt optimizations that help on some NUMA systems<br/>- distribute: spread execution evenly over all nodes<br/>- isolate: only spawn threads on CPUs on the node that execution started on<br/>- numactl: use the CPU map provided by numactl<br/>if run without this previously, it is recommended to drop the system page cache before using this<br/>see https://github.com/ggml-org/llama.cpp/issues/1437<br/>(env: LLAMA_ARG_NUMA) |
|
||||
| `-dev, --device <dev1,dev2,..>` | comma-separated list of devices to use for offloading (none = don't offload)<br/>use --list-devices to see a list of available devices<br/>(env: LLAMA_ARG_DEVICE) |
|
||||
| `--list-devices` | print list of available devices and exit |
|
||||
|
||||
@@ -138,10 +138,10 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
|
||||
| `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
|
||||
| `-np, --parallel N` | number of parallel sequences to decode (default: 1)<br/>(env: LLAMA_ARG_N_PARALLEL) |
|
||||
| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
|
||||
| `--mlock` | DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
|
||||
| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
|
||||
| `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
|
||||
| `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
|
||||
| `-lm, --load-mode MODE` | model loading mode (default: mmap)<br/>- none: no special loading mode<br/>- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>- mlock: mmap + force system to keep model in RAM rather than swapping or compressing<br/>- dio: use DirectIO if available<br/><br/>(env: LLAMA_ARG_LOAD_MODE) |
|
||||
| `-lm, --load-mode MODE` | model loading mode (default: mmap)<br/>- none: no special loading mode<br/>- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>- mlock: force system to keep model in RAM rather than swapping or compressing<br/>- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing<br/>- dio: use DirectIO if available<br/><br/>(env: LLAMA_ARG_LOAD_MODE) |
|
||||
| `--numa TYPE` | attempt optimizations that help on some NUMA systems<br/>- distribute: spread execution evenly over all nodes<br/>- isolate: only spawn threads on CPUs on the node that execution started on<br/>- numactl: use the CPU map provided by numactl<br/>if run without this previously, it is recommended to drop the system page cache before using this<br/>see https://github.com/ggml-org/llama.cpp/issues/1437<br/>(env: LLAMA_ARG_NUMA) |
|
||||
| `-dev, --device <dev1,dev2,..>` | comma-separated list of devices to use for offloading (none = don't offload)<br/>use --list-devices to see a list of available devices<br/>(env: LLAMA_ARG_DEVICE) |
|
||||
| `--list-devices` | print list of available devices and exit |
|
||||
|
||||
@@ -429,45 +429,45 @@ static void print_usage(int /* argc */, char ** argv) {
|
||||
}
|
||||
printf("\n");
|
||||
printf("test parameters:\n");
|
||||
printf(" -m, --model <filename> (default: %s)\n", join(cmd_params_defaults.model, ",").c_str());
|
||||
printf(" -hf, -hfr, --hf-repo <user>/<model>[:quant] Hugging Face model repository; quant is optional, case-insensitive\n");
|
||||
printf(" default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n");
|
||||
printf(" example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M\n");
|
||||
printf(" (default: unused)\n");
|
||||
printf(" -hff, --hf-file <file> Hugging Face model file. If specified, it will override the quant in --hf-repo\n");
|
||||
printf(" (default: unused)\n");
|
||||
printf(" -hft, --hf-token <token> Hugging Face access token\n");
|
||||
printf(" (default: value from HF_TOKEN environment variable)\n");
|
||||
printf(" --offline Offline mode: forces use of cache, prevents network access\n");
|
||||
printf(" (default: disabled)\n");
|
||||
printf(" -p, --n-prompt <n> (default: %s)\n", join(cmd_params_defaults.n_prompt, ",").c_str());
|
||||
printf(" -n, --n-gen <n> (default: %s)\n", join(cmd_params_defaults.n_gen, ",").c_str());
|
||||
printf(" -pg <pp,tg> (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str());
|
||||
printf(" -d, --n-depth <n> (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str());
|
||||
printf(" -b, --batch-size <n> (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str());
|
||||
printf(" -ub, --ubatch-size <n> (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str());
|
||||
printf(" -ctk, --cache-type-k <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str());
|
||||
printf(" -ctv, --cache-type-v <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str());
|
||||
printf(" -t, --threads <n> (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str());
|
||||
printf(" -C, --cpu-mask <hex,hex> (default: %s)\n", join(cmd_params_defaults.cpu_mask, ",").c_str());
|
||||
printf(" --cpu-strict <0|1> (default: %s)\n", join(cmd_params_defaults.cpu_strict, ",").c_str());
|
||||
printf(" --poll <0...100> (default: %s)\n", join(cmd_params_defaults.poll, ",").c_str());
|
||||
printf(" -ngl, --n-gpu-layers <n> (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").c_str());
|
||||
printf(" -ncmoe, --n-cpu-moe <n> (default: %s)\n", join(cmd_params_defaults.n_cpu_moe, ",").c_str());
|
||||
printf(" -sm, --split-mode <none|layer|row|tensor> (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str());
|
||||
printf(" -mg, --main-gpu <i> (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str());
|
||||
printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str());
|
||||
printf(" -fa, --flash-attn <on|off|auto> (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str());
|
||||
printf(" -dev, --device <dev0/dev1/...> (default: auto)\n");
|
||||
printf(" -lm, --load-mode <none|mmap|mlock|dio> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
|
||||
printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n");
|
||||
printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n");
|
||||
printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str());
|
||||
printf(" -ts, --tensor-split <ts0/ts1/..> (default: 0)\n");
|
||||
printf(" -m, --model <filename> (default: %s)\n", join(cmd_params_defaults.model, ",").c_str());
|
||||
printf(" -hf, -hfr, --hf-repo <user>/<model>[:quant] Hugging Face model repository; quant is optional, case-insensitive\n");
|
||||
printf(" default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n");
|
||||
printf(" example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M\n");
|
||||
printf(" (default: unused)\n");
|
||||
printf(" -hff, --hf-file <file> Hugging Face model file. If specified, it will override the quant in --hf-repo\n");
|
||||
printf(" (default: unused)\n");
|
||||
printf(" -hft, --hf-token <token> Hugging Face access token\n");
|
||||
printf(" (default: value from HF_TOKEN environment variable)\n");
|
||||
printf(" --offline Offline mode: forces use of cache, prevents network access\n");
|
||||
printf(" (default: disabled)\n");
|
||||
printf(" -p, --n-prompt <n> (default: %s)\n", join(cmd_params_defaults.n_prompt, ",").c_str());
|
||||
printf(" -n, --n-gen <n> (default: %s)\n", join(cmd_params_defaults.n_gen, ",").c_str());
|
||||
printf(" -pg <pp,tg> (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str());
|
||||
printf(" -d, --n-depth <n> (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str());
|
||||
printf(" -b, --batch-size <n> (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str());
|
||||
printf(" -ub, --ubatch-size <n> (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str());
|
||||
printf(" -ctk, --cache-type-k <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str());
|
||||
printf(" -ctv, --cache-type-v <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str());
|
||||
printf(" -t, --threads <n> (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str());
|
||||
printf(" -C, --cpu-mask <hex,hex> (default: %s)\n", join(cmd_params_defaults.cpu_mask, ",").c_str());
|
||||
printf(" --cpu-strict <0|1> (default: %s)\n", join(cmd_params_defaults.cpu_strict, ",").c_str());
|
||||
printf(" --poll <0...100> (default: %s)\n", join(cmd_params_defaults.poll, ",").c_str());
|
||||
printf(" -ngl, --n-gpu-layers <n> (default: %s)\n", join(cmd_params_defaults.n_gpu_layers, ",").c_str());
|
||||
printf(" -ncmoe, --n-cpu-moe <n> (default: %s)\n", join(cmd_params_defaults.n_cpu_moe, ",").c_str());
|
||||
printf(" -sm, --split-mode <none|layer|row|tensor> (default: %s)\n", join(transform_to_str(cmd_params_defaults.split_mode, split_mode_str), ",").c_str());
|
||||
printf(" -mg, --main-gpu <i> (default: %s)\n", join(cmd_params_defaults.main_gpu, ",").c_str());
|
||||
printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str());
|
||||
printf(" -fa, --flash-attn <on|off|auto> (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str());
|
||||
printf(" -dev, --device <dev0/dev1/...> (default: auto)\n");
|
||||
printf(" -lm, --load-mode <none|mmap|mlock|mmap+mlock|dio> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
|
||||
printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n");
|
||||
printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n");
|
||||
printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str());
|
||||
printf(" -ts, --tensor-split <ts0/ts1/..> (default: 0)\n");
|
||||
printf(" -ot --override-tensor <tensor name pattern>=<buffer type>;...\n");
|
||||
printf(" (default: disabled)\n");
|
||||
printf(" -nopo, --no-op-offload <0|1> (default: 0)\n");
|
||||
printf(" --no-host <0|1> (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str());
|
||||
printf(" (default: disabled)\n");
|
||||
printf(" -nopo, --no-op-offload <0|1> (default: 0)\n");
|
||||
printf(" --no-host <0|1> (default: %s)\n", join(cmd_params_defaults.no_host, ",").c_str());
|
||||
printf("\n");
|
||||
printf(
|
||||
"Multiple values can be given for each parameter by separating them with ','\n"
|
||||
@@ -785,6 +785,8 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
|
||||
mode = LLAMA_LOAD_MODE_MMAP;
|
||||
} else if (m == "mlock") {
|
||||
mode = LLAMA_LOAD_MODE_MLOCK;
|
||||
} else if (m == "mmap+mlock") {
|
||||
mode = LLAMA_LOAD_MODE_MMAP_MLOCK;
|
||||
} else if (m == "dio") {
|
||||
mode = LLAMA_LOAD_MODE_DIRECT_IO;
|
||||
} else {
|
||||
|
||||
@@ -1,8 +1,15 @@
|
||||
# mtmd
|
||||
|
||||
set(MTMD_VIDEO ON CACHE BOOL "enable video support in mtmd (requires ffmpeg binary in PATH)")
|
||||
set(MTMD_VIDEO_HELP "enable video support in mtmd (requires ffmpeg binary in PATH)")
|
||||
|
||||
set(MTMD_VIDEO ON CACHE BOOL "${MTMD_VIDEO_HELP}")
|
||||
# TODO: add MTMD_VIDEO_METHOD in the future to select between ffmpeg and other backends
|
||||
|
||||
if (MTMD_VIDEO AND NOT LLAMA_SUBPROCESS)
|
||||
message(STATUS "Disabling MTMD_VIDEO because LLAMA_SUBPROCESS is OFF")
|
||||
set(MTMD_VIDEO OFF CACHE BOOL "${MTMD_VIDEO_HELP}" FORCE)
|
||||
endif()
|
||||
|
||||
find_package(Threads REQUIRED)
|
||||
|
||||
add_library(mtmd
|
||||
@@ -40,6 +47,7 @@ add_library(mtmd
|
||||
models/paddleocr.cpp
|
||||
models/pixtral.cpp
|
||||
models/qwen2vl.cpp
|
||||
models/minimax-m3.cpp
|
||||
models/qwen3vl.cpp
|
||||
models/mimovl.cpp
|
||||
models/qwen3a.cpp
|
||||
|
||||
@@ -131,6 +131,8 @@
|
||||
#define TN_MM_SOFT_EMB_N "mm.soft_emb_norm.weight" // gemma3
|
||||
#define TN_MM_PROJECTOR "mm.model.fc.%s" // idefics3, deepseekocr
|
||||
#define TN_MM_PATCH_MERGER "mm.patch_merger.%s" // mistral small 3.1, glm4v
|
||||
#define TN_MM_MERGER_FC1 "mm.merger.fc1.%s" // minimax-m3 patch-merge MLP
|
||||
#define TN_MM_MERGER_FC2 "mm.merger.fc2.%s"
|
||||
#define TN_TOK_IMG_BREAK "v.token_embd.img_break" // pixtral
|
||||
#define TN_TOK_GLM_BOI "adapter.boi" // glm-edge (these embeddings are not in text model)
|
||||
#define TN_TOK_GLM_EOI "adapter.eoi" // glm-edge (these embeddings are not in text model)
|
||||
@@ -370,6 +372,7 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_MINICPMV4_6,
|
||||
PROJECTOR_TYPE_GRANITE_SPEECH,
|
||||
PROJECTOR_TYPE_MIMOVL,
|
||||
PROJECTOR_TYPE_MINIMAX_M3,
|
||||
PROJECTOR_TYPE_GRANITE4_VISION,
|
||||
PROJECTOR_TYPE_UNKNOWN,
|
||||
};
|
||||
@@ -424,6 +427,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_MINICPMV4_6, "minicpmv4_6"},
|
||||
{ PROJECTOR_TYPE_GRANITE_SPEECH, "granite_speech"},
|
||||
{ PROJECTOR_TYPE_MIMOVL, "mimovl"},
|
||||
{ PROJECTOR_TYPE_MINIMAX_M3, "minimax_m3"},
|
||||
{ PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"},
|
||||
};
|
||||
|
||||
|
||||
@@ -397,6 +397,10 @@ struct clip_model {
|
||||
ggml_tensor * mm_0_b = nullptr;
|
||||
ggml_tensor * mm_2_w = nullptr;
|
||||
ggml_tensor * mm_2_b = nullptr;
|
||||
ggml_tensor * mm_merger_fc1_w = nullptr; // minimax-m3
|
||||
ggml_tensor * mm_merger_fc1_b = nullptr;
|
||||
ggml_tensor * mm_merger_fc2_w = nullptr;
|
||||
ggml_tensor * mm_merger_fc2_b = nullptr;
|
||||
|
||||
ggml_tensor * image_newline = nullptr;
|
||||
ggml_tensor * view_seperator = nullptr;
|
||||
|
||||
@@ -915,6 +915,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
{
|
||||
builder = std::make_unique<clip_graph_mimovl>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_minimax_m3>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_STEP3VL:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_step3vl>(ctx, img);
|
||||
@@ -1469,6 +1473,17 @@ struct clip_model_loader {
|
||||
LOG_WRN("%s: more info: https://github.com/ggml-org/llama.cpp/issues/16842\n\n", __func__);
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
{
|
||||
hparams.n_merge = 2; // spatial_merge_size
|
||||
hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW;
|
||||
hparams.image_resize_pad = PAD_NONE;
|
||||
get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
|
||||
hparams.rope_theta = 10000.0f; // vision_config.rope_theta
|
||||
// MiniMax-M3: max_pixels 451584 (=672^2) -> 576 merged tokens (image_seq_length)
|
||||
hparams.set_limit_image_tokens(8, 576);
|
||||
hparams.set_warmup_n_tokens(16*16);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_MIMOVL:
|
||||
{
|
||||
hparams.n_merge = 2; // spatial_merge_size
|
||||
@@ -2089,6 +2104,19 @@ struct clip_model_loader {
|
||||
model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
|
||||
model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"), false);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
{
|
||||
// per-patch MLP: mm.1 -> gelu -> mm.2
|
||||
model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
|
||||
model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias"));
|
||||
model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
|
||||
model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
|
||||
// 2x2 merge MLP: mm.merge.fc1 -> gelu -> mm.merge.fc2
|
||||
model.mm_merger_fc1_w = get_tensor(string_format(TN_MM_MERGER_FC1, "weight"));
|
||||
model.mm_merger_fc1_b = get_tensor(string_format(TN_MM_MERGER_FC1, "bias"));
|
||||
model.mm_merger_fc2_w = get_tensor(string_format(TN_MM_MERGER_FC2, "weight"));
|
||||
model.mm_merger_fc2_b = get_tensor(string_format(TN_MM_MERGER_FC2, "bias"));
|
||||
} break;
|
||||
case PROJECTOR_TYPE_STEP3VL:
|
||||
{
|
||||
model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight"));
|
||||
@@ -3360,6 +3388,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
case PROJECTOR_TYPE_QWEN3VL:
|
||||
case PROJECTOR_TYPE_EXAONE4_5:
|
||||
case PROJECTOR_TYPE_MIMOVL:
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
case PROJECTOR_TYPE_GLM4V:
|
||||
case PROJECTOR_TYPE_YOUTUVL:
|
||||
{
|
||||
@@ -3866,6 +3895,24 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
|
||||
set_input_i32("positions", positions);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
{
|
||||
const int n_merge = hparams.n_merge;
|
||||
const int gh = image_size_height / patch_size;
|
||||
const int gw = image_size_width / patch_size;
|
||||
std::vector<int32_t> pos_h, pos_w;
|
||||
pos_h.reserve(gh * gw);
|
||||
pos_w.reserve(gh * gw);
|
||||
for (int bh = 0; bh < gh / n_merge; bh++)
|
||||
for (int bw = 0; bw < gw / n_merge; bw++)
|
||||
for (int mh = 0; mh < n_merge; mh++)
|
||||
for (int mw = 0; mw < n_merge; mw++) {
|
||||
pos_h.push_back(bh * n_merge + mh);
|
||||
pos_w.push_back(bw * n_merge + mw);
|
||||
}
|
||||
set_input_i32("minimax_pos_h", pos_h);
|
||||
set_input_i32("minimax_pos_w", pos_w);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DOTS_OCR:
|
||||
{
|
||||
const int pw = image_size_width / patch_size;
|
||||
@@ -4569,6 +4616,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
return ctx->model.mm_ffn_down_w->ne[1];
|
||||
case PROJECTOR_TYPE_GLM_EDGE:
|
||||
return ctx->model.mm_model_mlp_3_w->ne[1];
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
return ctx->model.mm_merger_fc2_b->ne[0];
|
||||
case PROJECTOR_TYPE_QWEN2VL:
|
||||
case PROJECTOR_TYPE_QWEN25VL:
|
||||
case PROJECTOR_TYPE_EXAONE4_5:
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
#include "models.h"
|
||||
|
||||
ggml_tensor * clip_graph_minimax_m3::apply_rope(
|
||||
ggml_tensor * x, ggml_tensor * pos_h, ggml_tensor * pos_w) {
|
||||
const int64_t Hn = x->ne[1];
|
||||
const int64_t P = x->ne[2];
|
||||
const size_t es = ggml_element_size(x);
|
||||
const int dh = (int) x->ne[0];
|
||||
const int axd = 2 * ((2 * (dh / 2) / 3) / 2);
|
||||
|
||||
GGML_ASSERT(x->nb[0] == es);
|
||||
GGML_ASSERT(3 * axd <= dh);
|
||||
|
||||
const float th = hparams.rope_theta;
|
||||
|
||||
// layout of x is [t, h, w, pad]
|
||||
// t is unrotated, h and w are rotated, pad is unrotated
|
||||
// note: everything from n_dims onward untouched, so w and pad are rotated in one call.
|
||||
auto sl = [&](int off, int n) {
|
||||
return ggml_cont(ctx0, ggml_view_3d(ctx0, x, n, Hn, P, x->nb[1], x->nb[2], (size_t) off * es));
|
||||
};
|
||||
ggml_tensor * t = sl(0, axd);
|
||||
ggml_tensor * h = sl(axd, axd);
|
||||
ggml_tensor * w = sl(2 * axd, dh - 2 * axd); // w + pad
|
||||
|
||||
h = ggml_rope_ext(ctx0, h, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
w = ggml_rope_ext(ctx0, w, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
return ggml_concat(ctx0, ggml_concat(ctx0, t, h, 0), w, 0);
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_minimax_m3::build() {
|
||||
GGML_ASSERT(model.patch_bias == nullptr);
|
||||
GGML_ASSERT(model.class_embedding == nullptr);
|
||||
GGML_ASSERT(model.patch_embeddings_0 && model.patch_embeddings_1);
|
||||
GGML_ASSERT(model.mm_1_w && model.mm_2_w);
|
||||
GGML_ASSERT(model.mm_merger_fc1_w && model.mm_merger_fc2_w);
|
||||
|
||||
const int batch_size = 1;
|
||||
const int n_pos = n_patches;
|
||||
const int merge = hparams.n_merge;
|
||||
|
||||
// patch embedding
|
||||
ggml_tensor * inp_raw = build_inp_raw();
|
||||
ggml_tensor * inp = ggml_add(ctx0,
|
||||
ggml_conv_2d(ctx0, model.patch_embeddings_0, inp_raw, patch_size, patch_size, 0, 0, 1, 1),
|
||||
ggml_conv_2d(ctx0, model.patch_embeddings_1, inp_raw, patch_size, patch_size, 0, 0, 1, 1));
|
||||
|
||||
// spatial merge
|
||||
{
|
||||
inp = ggml_permute(ctx0, inp, 1, 2, 0, 3);
|
||||
inp = ggml_cont_4d(ctx0, inp, n_embd * merge, n_patches_x / merge, n_patches_y, batch_size);
|
||||
inp = ggml_reshape_4d(ctx0, inp, n_embd * merge, n_patches_x / merge, merge, batch_size * (n_patches_y / merge));
|
||||
inp = ggml_permute(ctx0, inp, 0, 2, 1, 3);
|
||||
inp = ggml_cont_3d(ctx0, inp, n_embd, n_patches_x * n_patches_y, batch_size);
|
||||
}
|
||||
|
||||
// t (time axis) is always 0 for now, so we leave it unrotated
|
||||
ggml_tensor * pos_h = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos);
|
||||
ggml_set_name(pos_h, "minimax_pos_h"); ggml_set_input(pos_h);
|
||||
ggml_tensor * pos_w = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos);
|
||||
ggml_set_name(pos_w, "minimax_pos_w"); ggml_set_input(pos_w);
|
||||
|
||||
ggml_tensor * inpL = build_vit(
|
||||
inp, n_pos, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr,
|
||||
[&](ggml_tensor * c, const clip_layer &) {
|
||||
return apply_rope(c, pos_h, pos_w);
|
||||
});
|
||||
|
||||
// projector
|
||||
ggml_tensor * emb = inpL;
|
||||
emb = build_ffn(emb, model.mm_1_w, model.mm_1_b,
|
||||
nullptr, nullptr,
|
||||
model.mm_2_w, model.mm_2_b, FFN_GELU_ERF, -1);
|
||||
|
||||
const int64_t proj = emb->ne[0];
|
||||
emb = ggml_reshape_2d(ctx0, emb, proj * merge * merge, n_pos / (merge * merge));
|
||||
|
||||
emb = build_ffn(emb, model.mm_merger_fc1_w, model.mm_merger_fc1_b,
|
||||
nullptr, nullptr,
|
||||
model.mm_merger_fc2_w, model.mm_merger_fc2_b, FFN_GELU_ERF, -1);
|
||||
|
||||
ggml_build_forward_expand(gf, emb);
|
||||
return gf;
|
||||
}
|
||||
@@ -40,6 +40,12 @@ struct clip_graph_qwen3vl : clip_graph_qwen2vl {
|
||||
ggml_cgraph * build() override;
|
||||
};
|
||||
|
||||
struct clip_graph_minimax_m3 : clip_graph {
|
||||
clip_graph_minimax_m3(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
ggml_tensor * apply_rope(ggml_tensor * x, ggml_tensor * pos_h, ggml_tensor * pos_w);
|
||||
};
|
||||
|
||||
struct clip_graph_mimovl : clip_graph {
|
||||
clip_graph_mimovl(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -640,6 +640,7 @@ bool mtmd_helper_support_video(mtmd_context * ctx) {
|
||||
#ifdef MTMD_VIDEO
|
||||
return mtmd_support_vision(ctx);
|
||||
#else
|
||||
GGML_UNUSED(ctx);
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
@@ -1007,6 +1008,9 @@ mtmd_helper_video * mtmd_helper_video_init(
|
||||
|
||||
return ctx;
|
||||
#else
|
||||
GGML_UNUSED(mctx);
|
||||
GGML_UNUSED(path);
|
||||
GGML_UNUSED(params);
|
||||
LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__);
|
||||
return nullptr;
|
||||
#endif
|
||||
@@ -1039,6 +1043,10 @@ mtmd_helper_video * mtmd_helper_video_init_from_buf(
|
||||
|
||||
return ctx;
|
||||
#else
|
||||
GGML_UNUSED(mctx);
|
||||
GGML_UNUSED(buf);
|
||||
GGML_UNUSED(len);
|
||||
GGML_UNUSED(params);
|
||||
LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__);
|
||||
return nullptr;
|
||||
#endif
|
||||
@@ -1050,6 +1058,7 @@ void mtmd_helper_video_free(mtmd_helper_video * ctx) {
|
||||
ctx->stop_ffmpeg();
|
||||
delete ctx;
|
||||
#else
|
||||
GGML_UNUSED(ctx);
|
||||
LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__);
|
||||
#endif
|
||||
}
|
||||
@@ -1058,6 +1067,7 @@ mtmd_helper_video_info mtmd_helper_video_get_info(const mtmd_helper_video * ctx)
|
||||
#ifdef MTMD_VIDEO
|
||||
return ctx->info;
|
||||
#else
|
||||
GGML_UNUSED(ctx);
|
||||
GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)");
|
||||
#endif
|
||||
}
|
||||
@@ -1068,6 +1078,9 @@ int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx,
|
||||
if (!ctx) return -2;
|
||||
return ctx->read_next(out_bitmap, out_text);
|
||||
#else
|
||||
GGML_UNUSED(ctx);
|
||||
GGML_UNUSED(out_bitmap);
|
||||
GGML_UNUSED(out_text);
|
||||
GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)");
|
||||
#endif
|
||||
}
|
||||
|
||||
+18
-3
@@ -463,6 +463,13 @@ struct mtmd_context {
|
||||
img_end = "<|vision_end|>";
|
||||
image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_MINIMAX_M3:
|
||||
{
|
||||
// ]<]start of image[>[ ... (image embeddings) ... ]<]end of image[>[
|
||||
img_beg = "]<]start of image[>[";
|
||||
img_end = "]<]end of image[>[";
|
||||
image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_YOUTUVL:
|
||||
{
|
||||
// <|vision_start|> ... (image embeddings) ... <|vision_end|>
|
||||
@@ -555,9 +562,17 @@ struct mtmd_context {
|
||||
} break;
|
||||
case PROJECTOR_TYPE_KIMIK25:
|
||||
{
|
||||
// <|media_begin|> ... (image embeddings) ... <|media_end|>
|
||||
img_beg = "<|media_begin|>";
|
||||
img_end = "<|media_end|>";
|
||||
// GLM-5.2-V reuses the Kimi-K2.5 vision encoder and projector, but marks
|
||||
// images with its own tokens, so decide based on the text model vocab
|
||||
if (lookup_token("<|begin_of_image|>") != LLAMA_TOKEN_NULL) {
|
||||
// <|begin_of_image|> ... (image embeddings) ... <|end_of_image|>
|
||||
img_beg = "<|begin_of_image|>";
|
||||
img_end = "<|end_of_image|>";
|
||||
} else {
|
||||
// <|media_begin|> ... (image embeddings) ... <|media_end|>
|
||||
img_beg = "<|media_begin|>";
|
||||
img_end = "<|media_end|>";
|
||||
}
|
||||
image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_LIGHTONOCR:
|
||||
|
||||
@@ -19,6 +19,8 @@ add_library(${TARGET} STATIC
|
||||
server-stream.h
|
||||
server-tools.cpp
|
||||
server-tools.h
|
||||
server-mcp.cpp
|
||||
server-mcp.h
|
||||
server-schema.cpp
|
||||
server-schema.h
|
||||
)
|
||||
|
||||
@@ -136,13 +136,13 @@ Producer side: `server_res_generator` extends `server_res_spipe`, which keeps al
|
||||
|
||||
Lifetime safety: the session holds no back reference to the response, so `spipe` is a plain `unique_ptr` touched only by the http worker. `cancel` raises an atomic the producer polls; the producer finalizes the session from its destructor, which also runs `~server_response_reader::stop()` to cancel the generation at the queue level. A `DELETE` stops work by raising the flag and letting the worker unwind.
|
||||
|
||||
Consumer side: `GET /v1/stream/<conv_id>?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400.
|
||||
Consumer side: `GET /v1/stream?conv_id=<id>&from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400.
|
||||
|
||||
Routes:
|
||||
|
||||
- `GET /v1/stream/:conv_id?from=N`: replay or live reattach.
|
||||
- `GET /v1/stream?conv_id=<id>&from=N`: replay or live reattach. The id travels in the query string because it can embed a model name containing slashes.
|
||||
- `POST /v1/streams/lookup` with `{"conversation_ids": [...]}`: returns session status only for ids the caller already owns. There is no listing route, so live sessions cannot be enumerated (an earlier `GET /v1/streams` was removed for exactly this reason).
|
||||
- `DELETE /v1/stream/:conv_id`: explicit Stop, idempotent (`evict_and_cancel`).
|
||||
- `DELETE /v1/stream?conv_id=<id>`: explicit Stop, idempotent (`evict_and_cancel`).
|
||||
|
||||
Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport.
|
||||
|
||||
@@ -166,8 +166,8 @@ graph TD
|
||||
GC[GC thread] -- drop after TTL --> Sess
|
||||
end
|
||||
Sess -- read_from offset --> Cons[stream_pipe_consumer]
|
||||
Cons -- "GET /v1/stream/:id?from=N" --> Client
|
||||
DEL[DELETE /v1/stream/:id] -- evict_and_cancel --> Sess
|
||||
Cons -- "GET /v1/stream?conv_id=id&from=N" --> Client
|
||||
DEL[DELETE /v1/stream?conv_id=id] -- evict_and_cancel --> Sess
|
||||
```
|
||||
|
||||
The diagram shows the buffer touch points. The live wire (chunks streamed to the original client during a normal generation) is the producer's default output, described under "Producer side" above.
|
||||
@@ -189,7 +189,7 @@ This endpoint is intended to be used internally by the Web UI and subject to cha
|
||||
Get a list of tools, each tool has these fields:
|
||||
- `tool` (string): the ID name of the tool, to be used in POST call. Example: `read_file`
|
||||
- `display_name` (string): the name to be displayed on UI. Example: `Read file`
|
||||
- `type` (string): always be `"builtin"` for now
|
||||
- `type` (string): `"builtin"` for a built-in tool, or `"mcp"` for a tool exposed by an MCP server
|
||||
- `permissions` (object): a mapping string --> boolean that indicates the permission required by this tool. This is useful for the UI to ask the user before calling the tool. For now, the only permission supported is `"write"`
|
||||
- `definition` (object): the OAI-compat definition of this tool
|
||||
|
||||
@@ -199,7 +199,7 @@ Invoke a tool call, request body is a JSON object with:
|
||||
- `tool` (string): the name of the tool
|
||||
- `params` (object): a mapping from argument name (string) to argument value
|
||||
|
||||
Returns JSON object. There are two response formats:
|
||||
Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string):
|
||||
|
||||
Format 1: Plain text. The text will be placed into a field called `plain_text_response`, example:
|
||||
|
||||
|
||||
@@ -72,10 +72,10 @@ For the full list of features, please refer to [server's changelog](https://gith
|
||||
| `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
|
||||
| `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
|
||||
| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
|
||||
| `--mlock` | DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
|
||||
| `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
|
||||
| `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
|
||||
| `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
|
||||
| `-lm, --load-mode MODE` | model loading mode (default: mmap)<br/>- none: no special loading mode<br/>- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>- mlock: mmap + force system to keep model in RAM rather than swapping or compressing<br/>- dio: use DirectIO if available<br/><br/>(env: LLAMA_ARG_LOAD_MODE) |
|
||||
| `-lm, --load-mode MODE` | model loading mode (default: mmap)<br/>- none: no special loading mode<br/>- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>- mlock: force system to keep model in RAM rather than swapping or compressing<br/>- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing<br/>- dio: use DirectIO if available<br/><br/>(env: LLAMA_ARG_LOAD_MODE) |
|
||||
| `--numa TYPE` | attempt optimizations that help on some NUMA systems<br/>- distribute: spread execution evenly over all nodes<br/>- isolate: only spawn threads on CPUs on the node that execution started on<br/>- numactl: use the CPU map provided by numactl<br/>if run without this previously, it is recommended to drop the system page cache before using this<br/>see https://github.com/ggml-org/llama.cpp/issues/1437<br/>(env: LLAMA_ARG_NUMA) |
|
||||
| `-dev, --device <dev1,dev2,..>` | comma-separated list of devices to use for offloading (none = don't offload)<br/>use --list-devices to see a list of available devices<br/>(env: LLAMA_ARG_DEVICE) |
|
||||
| `--list-devices` | print list of available devices and exit |
|
||||
|
||||
@@ -1131,10 +1131,10 @@ json oaicompat_chat_params_parse(
|
||||
reasoning_budget = opt.reasoning_budget;
|
||||
}
|
||||
|
||||
if (!chat_params.thinking_end_tag.empty()) {
|
||||
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_tag"] = chat_params.thinking_end_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);
|
||||
}
|
||||
|
||||
@@ -9,9 +9,15 @@
|
||||
#define JSON_ASSERT GGML_ASSERT
|
||||
#include <nlohmann/json.hpp>
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <condition_variable>
|
||||
#include <cinttypes>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <cinttypes>
|
||||
|
||||
using json = nlohmann::ordered_json;
|
||||
|
||||
@@ -376,3 +382,67 @@ server_tokens format_prompt_rerank(
|
||||
mtmd_context * mctx,
|
||||
const std::string & query,
|
||||
const std::string & doc);
|
||||
|
||||
// simple implementation of a pipe
|
||||
// used for streaming data between threads
|
||||
template<typename T>
|
||||
struct server_pipe {
|
||||
std::mutex mutex;
|
||||
std::condition_variable cv;
|
||||
std::queue<T> queue;
|
||||
std::atomic<bool> writer_closed{false};
|
||||
std::atomic<bool> reader_closed{false};
|
||||
|
||||
// 0 = unbounded (default)
|
||||
// > 0, write() drops the oldest item once the queue is full
|
||||
size_t max_size = 0;
|
||||
|
||||
void close_write() {
|
||||
writer_closed.store(true, std::memory_order_relaxed);
|
||||
cv.notify_all();
|
||||
}
|
||||
|
||||
void close_read() {
|
||||
reader_closed.store(true, std::memory_order_relaxed);
|
||||
cv.notify_all();
|
||||
}
|
||||
|
||||
// close_on_stop = true: should_stop means the reader is gone for good, so the writer is told the pipe is broken.
|
||||
// close_on_stop = false: should_stop is a per-read deadline and further reads still come, so the pipe stays usable.
|
||||
bool read(T & output, const std::function<bool()> & should_stop, bool close_on_stop = true) {
|
||||
std::unique_lock<std::mutex> lk(mutex);
|
||||
constexpr auto poll_interval = std::chrono::milliseconds(500);
|
||||
while (true) {
|
||||
if (!queue.empty()) {
|
||||
output = std::move(queue.front());
|
||||
queue.pop();
|
||||
return true;
|
||||
}
|
||||
if (writer_closed.load()) {
|
||||
return false; // clean EOF
|
||||
}
|
||||
if (should_stop && should_stop()) { // a null should_stop means "never stop"
|
||||
if (close_on_stop) {
|
||||
close_read(); // signal broken pipe to writer
|
||||
}
|
||||
return false; // cancelled / deadline reached
|
||||
}
|
||||
cv.wait_for(lk, poll_interval);
|
||||
}
|
||||
}
|
||||
|
||||
bool write(T && data) {
|
||||
std::lock_guard<std::mutex> lk(mutex);
|
||||
if (reader_closed.load()) {
|
||||
return false; // broken pipe
|
||||
}
|
||||
if (max_size > 0) {
|
||||
while (queue.size() >= max_size) {
|
||||
queue.pop(); // drop oldest to stay bounded
|
||||
}
|
||||
}
|
||||
queue.push(std::move(data));
|
||||
cv.notify_one();
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -0,0 +1,820 @@
|
||||
#include "server-mcp.h"
|
||||
|
||||
#include "subproc.h"
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cstdio>
|
||||
#include <fstream>
|
||||
#include <functional>
|
||||
#include <sstream>
|
||||
#include <thread>
|
||||
|
||||
#if defined(_WIN32)
|
||||
# include <io.h>
|
||||
# include <windows.h>
|
||||
#else
|
||||
# include <errno.h>
|
||||
# include <fcntl.h>
|
||||
# include <poll.h>
|
||||
# include <unistd.h>
|
||||
extern char ** environ;
|
||||
#endif
|
||||
|
||||
// read NDJSON lines from a child pipe, calling on_line per line until `running` clears, EOF/error, or on_line returns false.
|
||||
// polled, not blocking: a grandchild can inherit the pipe's write end and hold it open (terminate() kills only the direct child), so a blocking read would hang teardown on an EOF that never comes.
|
||||
static void mcp_pump_ndjson(FILE * f, std::atomic<bool> & running,
|
||||
const std::function<bool(std::string &&)> & on_line) {
|
||||
if (!f) {
|
||||
return;
|
||||
}
|
||||
const int poll_ms = 50;
|
||||
const size_t max_line = 8 * 1024 * 1024; // drop any single NDJSON line larger than this, so a child that never emits '\n' can't grow buf without bound
|
||||
#if defined(_WIN32)
|
||||
HANDLE h = (HANDLE) _get_osfhandle(_fileno(f));
|
||||
#else
|
||||
int fd = fileno(f);
|
||||
int fl = fcntl(fd, F_GETFL, 0);
|
||||
if (fl >= 0) {
|
||||
fcntl(fd, F_SETFL, fl | O_NONBLOCK);
|
||||
}
|
||||
#endif
|
||||
std::string buf;
|
||||
bool skipping = false; // discarding an over-long line until its terminating newline
|
||||
char chunk[4096];
|
||||
while (running.load()) {
|
||||
size_t n = 0;
|
||||
#if defined(_WIN32)
|
||||
DWORD avail = 0;
|
||||
if (!PeekNamedPipe(h, NULL, 0, NULL, &avail, NULL)) {
|
||||
break; // pipe broken / child gone
|
||||
}
|
||||
if (avail == 0) {
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(poll_ms));
|
||||
continue;
|
||||
}
|
||||
DWORD to_read = avail < (DWORD) sizeof(chunk) ? avail : (DWORD) sizeof(chunk);
|
||||
DWORD got = 0;
|
||||
if (!ReadFile(h, chunk, to_read, &got, NULL) || got == 0) {
|
||||
break;
|
||||
}
|
||||
n = (size_t) got;
|
||||
#else
|
||||
struct pollfd pfd;
|
||||
pfd.fd = fd;
|
||||
pfd.events = POLLIN;
|
||||
pfd.revents = 0;
|
||||
int pr = poll(&pfd, 1, poll_ms);
|
||||
if (pr < 0) {
|
||||
if (errno == EINTR) {
|
||||
continue;
|
||||
}
|
||||
break;
|
||||
}
|
||||
if (pr == 0) {
|
||||
continue; // timeout -> re-check running
|
||||
}
|
||||
if (pfd.revents & (POLLERR | POLLNVAL)) {
|
||||
break;
|
||||
}
|
||||
ssize_t r = read(fd, chunk, sizeof(chunk));
|
||||
if (r < 0) {
|
||||
if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) {
|
||||
continue;
|
||||
}
|
||||
break;
|
||||
}
|
||||
if (r == 0) {
|
||||
break; // EOF: child (and any pipe writers) closed the stream
|
||||
}
|
||||
n = (size_t) r;
|
||||
#endif
|
||||
buf.append(chunk, n);
|
||||
|
||||
// resync after an over-long, unterminated line: discard bytes until the next newline
|
||||
if (skipping) {
|
||||
size_t nl = buf.find('\n');
|
||||
if (nl == std::string::npos) {
|
||||
if (buf.size() > max_line) {
|
||||
buf.clear(); // stay bounded while waiting for a terminator
|
||||
}
|
||||
continue;
|
||||
}
|
||||
buf.erase(0, nl + 1);
|
||||
skipping = false;
|
||||
}
|
||||
|
||||
size_t pos;
|
||||
while ((pos = buf.find('\n')) != std::string::npos) {
|
||||
std::string line = buf.substr(0, pos);
|
||||
buf.erase(0, pos + 1);
|
||||
if (!line.empty() && line.back() == '\r') {
|
||||
line.pop_back();
|
||||
}
|
||||
if (line.empty()) {
|
||||
continue;
|
||||
}
|
||||
if (!on_line(std::move(line))) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// a partial line already larger than the cap and still no newline: drop it to avoid unbounded growth
|
||||
if (buf.size() > max_line) {
|
||||
SRV_WRN("MCP: dropping oversized line (> %zu bytes) from child pipe\n", max_line);
|
||||
buf.clear();
|
||||
skipping = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// server_mcp_server_config
|
||||
//
|
||||
|
||||
std::vector<server_mcp_server_config> server_mcp_server_config::parse_from_json(const std::string & json_str) {
|
||||
return parse_cursor_format(json::parse(json_str));
|
||||
}
|
||||
|
||||
std::vector<server_mcp_server_config> server_mcp_server_config::parse_cursor_format(const json & j) {
|
||||
std::vector<server_mcp_server_config> result;
|
||||
|
||||
if (!j.contains("mcpServers") || !j.at("mcpServers").is_object()) {
|
||||
return result;
|
||||
}
|
||||
|
||||
for (const auto & [name, cfg] : j.at("mcpServers").items()) {
|
||||
server_mcp_server_config sc;
|
||||
sc.name = name;
|
||||
sc.command = cfg.value("command", std::string());
|
||||
sc.cwd = cfg.value("cwd", std::string());
|
||||
sc.timeout_ms = cfg.value("timeout_ms", sc.timeout_ms);
|
||||
|
||||
if (cfg.contains("args") && cfg.at("args").is_array()) {
|
||||
for (const auto & a : cfg.at("args")) {
|
||||
sc.args.push_back(a.get<std::string>());
|
||||
}
|
||||
}
|
||||
if (cfg.contains("env") && cfg.at("env").is_object()) {
|
||||
for (const auto & [k, v] : cfg.at("env").items()) {
|
||||
sc.env[k] = v.get<std::string>();
|
||||
}
|
||||
}
|
||||
|
||||
if (sc.command.empty()) {
|
||||
SRV_WRN("MCP server '%s' has no command, skipping\n", name.c_str());
|
||||
continue;
|
||||
}
|
||||
result.push_back(std::move(sc));
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
//
|
||||
// server_mcp_transport
|
||||
//
|
||||
|
||||
static constexpr const char * MCP_PROTOCOL_VERSION = "2024-11-05";
|
||||
|
||||
static std::string rpc_error_message(const json & resp) {
|
||||
if (resp.contains("error")) {
|
||||
const json & e = resp.at("error");
|
||||
if (e.is_object()) {
|
||||
return e.value("message", "unknown error");
|
||||
}
|
||||
if (e.is_string()) {
|
||||
return e.get<std::string>();
|
||||
}
|
||||
}
|
||||
return "unknown error";
|
||||
}
|
||||
|
||||
// normalize an MCP tools/call result to the /tools contract (see README-dev.md):
|
||||
// concat text parts of result.content[], and surface an isError result
|
||||
static json mcp_result_to_response(const json & result) {
|
||||
std::string text;
|
||||
if (result.contains("content") && result.at("content").is_array()) {
|
||||
for (const auto & part : result.at("content")) {
|
||||
if (part.is_object() && part.value("type", "") == "text") {
|
||||
if (!text.empty()) {
|
||||
text += "\n";
|
||||
}
|
||||
text += part.value("text", "");
|
||||
}
|
||||
}
|
||||
}
|
||||
if (result.is_object() && result.value("isError", false)) {
|
||||
return {{"error", text.empty() ? "MCP tool returned an error" : text}};
|
||||
}
|
||||
return {{"plain_text_response", text}};
|
||||
}
|
||||
|
||||
json server_mcp_transport::send_rpc(const json & request, const std::function<bool()> & should_stop) {
|
||||
if (!to_server.write(request.dump())) {
|
||||
return {{"error", {{"code", -32603}, {"message", "transport closed"}}}};
|
||||
}
|
||||
|
||||
const bool has_id = request.contains("id");
|
||||
const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms);
|
||||
auto stop = [&]() {
|
||||
return (should_stop && should_stop()) || std::chrono::steady_clock::now() >= deadline;
|
||||
};
|
||||
|
||||
std::string frame;
|
||||
while (from_server.read(frame, stop, false)) {
|
||||
json reply;
|
||||
try {
|
||||
reply = json::parse(frame);
|
||||
} catch (...) {
|
||||
if (std::chrono::steady_clock::now() >= deadline) {
|
||||
break;
|
||||
}
|
||||
continue; // skip malformed frame
|
||||
}
|
||||
// no id: a notification. mismatched id: a stale reply from a timed-out request (ids are monotonic, never a future one)
|
||||
if (!has_id || (reply.contains("id") && reply.at("id") == request.at("id"))) {
|
||||
return reply;
|
||||
}
|
||||
if (std::chrono::steady_clock::now() >= deadline) {
|
||||
break; // a flood of notifications must not outrun the deadline
|
||||
}
|
||||
}
|
||||
|
||||
if (should_stop && should_stop()) {
|
||||
return {{"error", {{"code", -32603}, {"message", "cancelled"}}}};
|
||||
}
|
||||
if (std::chrono::steady_clock::now() >= deadline) {
|
||||
return {{"error", {{"code", -32603}, {"message", "request timed out"}}}};
|
||||
}
|
||||
return {{"error", {{"code", -32603}, {"message", "transport closed"}}}};
|
||||
}
|
||||
|
||||
bool server_mcp_transport::ensure_init(const std::function<bool()> & should_stop) {
|
||||
if (initialized) {
|
||||
return true;
|
||||
}
|
||||
|
||||
json init_req = {
|
||||
{"jsonrpc", "2.0"},
|
||||
{"id", next_id++},
|
||||
{"method", "initialize"},
|
||||
{"params", {
|
||||
{"protocolVersion", MCP_PROTOCOL_VERSION},
|
||||
{"capabilities", json::object()},
|
||||
{"clientInfo", {{"name", "llama.cpp"}, {"version", "1.0"}}},
|
||||
}},
|
||||
};
|
||||
json resp = send_rpc(init_req, should_stop);
|
||||
if (!resp.contains("result")) {
|
||||
last_error = "initialize failed: " + rpc_error_message(resp);
|
||||
return false;
|
||||
}
|
||||
|
||||
// notifications/initialized: no id, no reply expected
|
||||
json notif = {{"jsonrpc", "2.0"}, {"method", "notifications/initialized"}};
|
||||
to_server.write(notif.dump());
|
||||
|
||||
initialized = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<server_mcp_tool_def> server_mcp_transport::list_tools(const std::function<bool()> & should_stop) {
|
||||
std::lock_guard<std::mutex> lock(rpc_mutex);
|
||||
if (!ensure_init(should_stop)) {
|
||||
return {};
|
||||
}
|
||||
if (!tools.empty()) {
|
||||
return tools;
|
||||
}
|
||||
|
||||
json req = {{"jsonrpc", "2.0"}, {"id", next_id++}, {"method", "tools/list"}};
|
||||
json resp = send_rpc(req, should_stop);
|
||||
if (!resp.contains("result")) {
|
||||
last_error = "tools/list failed: " + rpc_error_message(resp);
|
||||
return {};
|
||||
}
|
||||
|
||||
const json & result = resp.at("result");
|
||||
if (result.contains("tools") && result.at("tools").is_array()) {
|
||||
for (const auto & t : result.at("tools")) {
|
||||
server_mcp_tool_def def;
|
||||
def.server_name = name;
|
||||
def.name = t.value("name", "");
|
||||
def.description = t.value("description", "");
|
||||
if (t.contains("inputSchema")) {
|
||||
def.input_schema = t.at("inputSchema");
|
||||
}
|
||||
tools.push_back(std::move(def));
|
||||
}
|
||||
}
|
||||
return tools;
|
||||
}
|
||||
|
||||
json server_mcp_transport::call_tool(const std::string & tool_name,
|
||||
const json & arguments,
|
||||
const std::function<bool()> & should_stop) {
|
||||
std::lock_guard<std::mutex> lock(rpc_mutex);
|
||||
if (!ensure_init(should_stop)) {
|
||||
return {{"error", last_error}};
|
||||
}
|
||||
|
||||
json req = {
|
||||
{"jsonrpc", "2.0"},
|
||||
{"id", next_id++},
|
||||
{"method", "tools/call"},
|
||||
{"params", {{"name", tool_name}, {"arguments", arguments}}},
|
||||
};
|
||||
json resp = send_rpc(req, should_stop);
|
||||
if (resp.contains("error")) {
|
||||
return {{"error", rpc_error_message(resp)}};
|
||||
}
|
||||
if (resp.contains("result")) {
|
||||
return mcp_result_to_response(resp.at("result"));
|
||||
}
|
||||
return {{"error", "invalid response from MCP server"}};
|
||||
}
|
||||
|
||||
//
|
||||
// server_mcp_stdio
|
||||
//
|
||||
|
||||
struct server_mcp_stdio::process_handle {
|
||||
common_subproc sp;
|
||||
FILE * in = nullptr; // child stdin
|
||||
FILE * out = nullptr; // child stdout
|
||||
FILE * err = nullptr; // child stderr
|
||||
};
|
||||
|
||||
#if defined(_WIN32)
|
||||
// config strings are UTF-8 (from JSON) and subprocess.h converts them with CP_UTF8, so inputs must be UTF-8, not the active code page
|
||||
static std::wstring windows_utf8_to_wide(const std::string & s) {
|
||||
if (s.empty()) {
|
||||
return std::wstring();
|
||||
}
|
||||
int n = MultiByteToWideChar(CP_UTF8, 0, s.data(), (int) s.size(), NULL, 0);
|
||||
if (n <= 0) {
|
||||
return std::wstring();
|
||||
}
|
||||
std::wstring w((size_t) n, L'\0');
|
||||
MultiByteToWideChar(CP_UTF8, 0, s.data(), (int) s.size(), &w[0], n);
|
||||
return w;
|
||||
}
|
||||
|
||||
static std::string windows_wide_to_utf8(const wchar_t * s, int len /* -1 for NUL-terminated */) {
|
||||
int n = WideCharToMultiByte(CP_UTF8, 0, s, len, NULL, 0, NULL, NULL);
|
||||
if (n <= 0) {
|
||||
return std::string();
|
||||
}
|
||||
std::string out((size_t) n, '\0');
|
||||
WideCharToMultiByte(CP_UTF8, 0, s, len, &out[0], n, NULL, NULL);
|
||||
if (len == -1 && !out.empty() && out.back() == '\0') {
|
||||
out.pop_back(); // drop the terminator WideCharToMultiByte counts for -1
|
||||
}
|
||||
return out;
|
||||
}
|
||||
#endif
|
||||
|
||||
static std::string mcp_resolve_command(const std::string & command) {
|
||||
#if defined(_WIN32)
|
||||
// For Windows: make sure we handle ".exe" correctly, as well as UTF-8
|
||||
std::wstring wcmd = windows_utf8_to_wide(command);
|
||||
wchar_t buf[MAX_PATH * 4];
|
||||
const DWORD cap = (DWORD) (sizeof(buf) / sizeof(buf[0]));
|
||||
|
||||
auto search = [&](const wchar_t * ext) -> std::string {
|
||||
DWORD n = SearchPathW(NULL, wcmd.c_str(), ext, cap, buf, NULL);
|
||||
return (n > 0 && n < cap) ? windows_wide_to_utf8(buf, (int) n) : std::string();
|
||||
};
|
||||
|
||||
std::string found = search(NULL); // exact path / already-extensioned / .exe on PATH
|
||||
if (!found.empty()) {
|
||||
return found;
|
||||
}
|
||||
|
||||
std::wstring pathext;
|
||||
DWORD need = GetEnvironmentVariableW(L"PATHEXT", NULL, 0);
|
||||
if (need > 0) {
|
||||
pathext.resize(need);
|
||||
DWORD got = GetEnvironmentVariableW(L"PATHEXT", &pathext[0], need);
|
||||
pathext.resize(got);
|
||||
}
|
||||
if (pathext.empty()) {
|
||||
pathext = L".COM;.EXE;.BAT;.CMD";
|
||||
}
|
||||
for (size_t start = 0; start <= pathext.size();) {
|
||||
size_t sep = pathext.find(L';', start);
|
||||
std::wstring ext = pathext.substr(start, sep == std::wstring::npos ? std::wstring::npos : sep - start);
|
||||
if (!ext.empty()) {
|
||||
found = search(ext.c_str());
|
||||
if (!found.empty()) {
|
||||
return found;
|
||||
}
|
||||
}
|
||||
if (sep == std::wstring::npos) {
|
||||
break;
|
||||
}
|
||||
start = sep + 1;
|
||||
}
|
||||
return command; // give up and let subprocess.h report the spawn error
|
||||
#else
|
||||
return command;
|
||||
#endif // _WIN32
|
||||
}
|
||||
|
||||
static std::vector<std::string> mcp_parent_env() {
|
||||
std::vector<std::string> env;
|
||||
#if defined(_WIN32)
|
||||
LPWCH block = GetEnvironmentStringsW();
|
||||
if (block) {
|
||||
for (LPWCH e = block; *e; e += wcslen(e) + 1) {
|
||||
env.emplace_back(windows_wide_to_utf8(e, -1));
|
||||
}
|
||||
FreeEnvironmentStringsW(block);
|
||||
}
|
||||
#else
|
||||
if (environ) {
|
||||
for (char ** e = environ; *e; ++e) {
|
||||
env.emplace_back(*e);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
return env;
|
||||
}
|
||||
|
||||
// parent env with the config overrides applied, in "KEY=VALUE" form
|
||||
static std::vector<std::string> mcp_build_env(const std::map<std::string, std::string> & overrides) {
|
||||
std::vector<std::string> env;
|
||||
for (auto & e : mcp_parent_env()) {
|
||||
size_t eq = e.find('=');
|
||||
std::string key = eq == std::string::npos ? e : e.substr(0, eq);
|
||||
if (overrides.find(key) == overrides.end()) {
|
||||
env.push_back(e);
|
||||
}
|
||||
}
|
||||
for (auto & [k, v] : overrides) {
|
||||
env.push_back(k + "=" + v);
|
||||
}
|
||||
return env;
|
||||
}
|
||||
|
||||
server_mcp_stdio::server_mcp_stdio(const server_mcp_server_config & config) : config(config) {
|
||||
name = config.name;
|
||||
timeout_ms = config.timeout_ms;
|
||||
// bound the reply queue: send_rpc only drains during a call, so unsolicited notifications would otherwise grow it without limit
|
||||
from_server.max_size = 65536;
|
||||
}
|
||||
|
||||
server_mcp_stdio::~server_mcp_stdio() {
|
||||
join_pumps();
|
||||
}
|
||||
|
||||
bool server_mcp_stdio::start() {
|
||||
std::vector<std::string> argv_s;
|
||||
argv_s.push_back(mcp_resolve_command(config.command));
|
||||
argv_s.insert(argv_s.end(), config.args.begin(), config.args.end());
|
||||
|
||||
int options = subprocess_option_no_window | subprocess_option_search_user_path;
|
||||
std::vector<std::string> envp_s;
|
||||
if (config.env.empty()) {
|
||||
options |= subprocess_option_inherit_environment;
|
||||
} else {
|
||||
envp_s = mcp_build_env(config.env);
|
||||
}
|
||||
|
||||
auto handle = std::make_unique<process_handle>();
|
||||
bool ok = handle->sp.create(argv_s, options, envp_s, config.cwd.empty() ? nullptr : config.cwd.c_str());
|
||||
if (!ok) {
|
||||
SRV_WRN("MCP '%s': failed to spawn '%s'\n", config.name.c_str(), config.command.c_str());
|
||||
return false;
|
||||
}
|
||||
handle->in = handle->sp.stdin_file();
|
||||
handle->out = handle->sp.stdout_file();
|
||||
handle->err = handle->sp.stderr_file();
|
||||
|
||||
proc = std::move(handle);
|
||||
running.store(true);
|
||||
reader = std::thread([this] { reader_loop(); });
|
||||
writer = std::thread([this] { writer_loop(); });
|
||||
errlog = std::thread([this] { errlog_loop(); });
|
||||
return true;
|
||||
}
|
||||
|
||||
void server_mcp_stdio::close() {
|
||||
join_pumps();
|
||||
}
|
||||
|
||||
bool server_mcp_stdio::is_alive() const {
|
||||
return running.load();
|
||||
}
|
||||
|
||||
std::string server_mcp_stdio::diagnostics() {
|
||||
std::string out;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(rpc_mutex); // last_error is written by send_rpc's callers
|
||||
out = last_error;
|
||||
}
|
||||
std::lock_guard<std::mutex> lk(err_mu);
|
||||
if (!err_tail.empty()) {
|
||||
if (!out.empty()) {
|
||||
out += "; ";
|
||||
}
|
||||
out += "last stderr: " + err_tail;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
void server_mcp_stdio::reader_loop() {
|
||||
mcp_pump_ndjson(proc->out, running, [this](std::string && line) {
|
||||
return from_server.write(std::move(line)); // false => consumer gone, stop
|
||||
});
|
||||
running.store(false);
|
||||
to_server.close_write(); // stop the writer
|
||||
from_server.close_write(); // EOF to any waiting caller
|
||||
}
|
||||
|
||||
// write all of `data` to child stdin, non-blocking and polled so teardown never hangs (a grandchild can hold the read end of a full pipe open). returns false on error/close/shutdown.
|
||||
static bool mcp_write_all(FILE * f, const std::string & data, std::atomic<bool> & running) {
|
||||
if (!f) {
|
||||
return false;
|
||||
}
|
||||
size_t total = 0;
|
||||
#if defined(_WIN32)
|
||||
HANDLE h = (HANDLE) _get_osfhandle(_fileno(f));
|
||||
DWORD nowait = PIPE_NOWAIT;
|
||||
SetNamedPipeHandleState(h, &nowait, NULL, NULL);
|
||||
while (total < data.size() && running.load()) {
|
||||
DWORD written = 0;
|
||||
BOOL ok = WriteFile(h, data.data() + total, (DWORD) (data.size() - total), &written, NULL);
|
||||
if (ok && written > 0) {
|
||||
total += written;
|
||||
continue;
|
||||
}
|
||||
if (!ok) {
|
||||
DWORD err = GetLastError();
|
||||
if (err != ERROR_NO_DATA && err != ERROR_PIPE_BUSY) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// backpressure (pipe full) is rare for small JSON-RPC frames; sleep rather than spin.
|
||||
// no writable-wait exists for a PIPE_NOWAIT anonymous pipe, so this polls like the POSIX poll() path.
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(10));
|
||||
}
|
||||
#else
|
||||
int fd = fileno(f);
|
||||
int fl = fcntl(fd, F_GETFL, 0);
|
||||
if (fl >= 0) {
|
||||
fcntl(fd, F_SETFL, fl | O_NONBLOCK);
|
||||
}
|
||||
while (total < data.size() && running.load()) {
|
||||
ssize_t n = write(fd, data.data() + total, data.size() - total);
|
||||
if (n > 0) {
|
||||
total += (size_t) n;
|
||||
continue;
|
||||
}
|
||||
if (n == 0) {
|
||||
return false;
|
||||
}
|
||||
if (errno == EINTR) {
|
||||
continue;
|
||||
}
|
||||
if (errno != EAGAIN && errno != EWOULDBLOCK) {
|
||||
return false;
|
||||
}
|
||||
struct pollfd pfd;
|
||||
pfd.fd = fd;
|
||||
pfd.events = POLLOUT;
|
||||
pfd.revents = 0;
|
||||
int pr = poll(&pfd, 1, 50);
|
||||
if (pr < 0) {
|
||||
if (errno == EINTR) {
|
||||
continue;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
if (pfd.revents & (POLLERR | POLLNVAL | POLLHUP)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
return total == data.size();
|
||||
}
|
||||
|
||||
void server_mcp_stdio::writer_loop() {
|
||||
auto should_stop = [this] { return !running.load(); };
|
||||
std::string msg;
|
||||
while (to_server.read(msg, should_stop)) {
|
||||
msg.push_back('\n');
|
||||
if (!mcp_write_all(proc->in, msg, running)) {
|
||||
break; // child gone or shutting down
|
||||
}
|
||||
}
|
||||
running.store(false);
|
||||
to_server.close_read(); // fail fast on any further send_rpc write
|
||||
from_server.close_write(); // wake any caller waiting for a reply
|
||||
}
|
||||
|
||||
void server_mcp_stdio::errlog_loop() {
|
||||
static constexpr size_t ERR_TAIL_MAX = 4096;
|
||||
// drain stderr (an undrained pipe blocks the child):
|
||||
// log it, and keep a bounded tail for reporting when the server dies
|
||||
mcp_pump_ndjson(proc->err, running, [this](std::string && line) {
|
||||
SRV_DBG("MCP '%s' stderr: %s\n", name.c_str(), line.c_str());
|
||||
std::lock_guard<std::mutex> lk(err_mu);
|
||||
err_tail += line;
|
||||
err_tail += '\n';
|
||||
if (err_tail.size() > ERR_TAIL_MAX) {
|
||||
err_tail.erase(0, err_tail.size() - ERR_TAIL_MAX);
|
||||
}
|
||||
return true;
|
||||
});
|
||||
}
|
||||
|
||||
void server_mcp_stdio::join_pumps() {
|
||||
if (!proc) {
|
||||
return;
|
||||
}
|
||||
running.store(false);
|
||||
to_server.close_write(); // wake the writer if it waits for a message
|
||||
from_server.close_write(); // wake any caller waiting for a reply
|
||||
|
||||
proc->sp.terminate(); // child death unblocks the blocked fread/fwrite
|
||||
|
||||
if (writer.joinable()) writer.join();
|
||||
if (reader.joinable()) reader.join();
|
||||
if (errlog.joinable()) errlog.join();
|
||||
|
||||
proc->sp.join(); // reap the child: never waiting would leave the pid a zombie for the process lifetime
|
||||
proc.reset();
|
||||
}
|
||||
|
||||
|
||||
//
|
||||
// server_mcp
|
||||
//
|
||||
|
||||
static constexpr int MCP_COOLDOWN_SECONDS = 5;
|
||||
static constexpr int MCP_WARMUP_TIMEOUT_SECONDS = 10; // cap per-server tool discovery at startup
|
||||
|
||||
server_mcp::~server_mcp() {
|
||||
shutdown();
|
||||
|
||||
std::vector<std::shared_ptr<server_mcp_transport>> to_close;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
for (auto & [name, t] : transports) {
|
||||
to_close.push_back(std::move(t));
|
||||
}
|
||||
transports.clear();
|
||||
}
|
||||
for (auto & t : to_close) {
|
||||
t->close();
|
||||
}
|
||||
}
|
||||
|
||||
std::shared_ptr<server_mcp_transport> server_mcp::create_transport(const server_mcp_server_config & cfg) {
|
||||
return std::make_shared<server_mcp_stdio>(cfg);
|
||||
}
|
||||
|
||||
void server_mcp::shutdown() {
|
||||
stopping.store(true);
|
||||
}
|
||||
|
||||
const server_mcp_server_config * server_mcp::find_config(const std::string & name) const {
|
||||
for (const auto & c : configs) {
|
||||
if (c.name == name) {
|
||||
return &c;
|
||||
}
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void server_mcp::start(const common_params & params) {
|
||||
auto append = [this](const std::string & json_str) {
|
||||
try {
|
||||
auto parsed = server_mcp_server_config::parse_from_json(json_str);
|
||||
if (parsed.empty()) {
|
||||
SRV_WRN("%s", "MCP config: no servers found in JSON\n");
|
||||
}
|
||||
for (auto & p : parsed) {
|
||||
// names must be unique across both config sources: get_or_create / find_config key on the name
|
||||
if (find_config(p.name)) {
|
||||
SRV_WRN("MCP config: duplicate server name '%s', skipping\n", p.name.c_str());
|
||||
continue;
|
||||
}
|
||||
configs.push_back(std::move(p));
|
||||
}
|
||||
} catch (const std::exception & e) {
|
||||
throw std::runtime_error(std::string("failed to parse MCP config JSON: ") + e.what());
|
||||
}
|
||||
};
|
||||
if (!params.mcp_servers_config.empty()) {
|
||||
std::ifstream f = fs_open_ifstream(params.mcp_servers_config, std::ios::in);
|
||||
if (!f) {
|
||||
throw std::runtime_error("failed to open MCP config file: " + params.mcp_servers_config);
|
||||
}
|
||||
std::stringstream ss;
|
||||
ss << f.rdbuf();
|
||||
append(ss.str());
|
||||
}
|
||||
if (!params.mcp_servers_json.empty()) {
|
||||
append(params.mcp_servers_json);
|
||||
}
|
||||
|
||||
if (configs.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<server_mcp_tool_def> discovered;
|
||||
for (const auto & cfg : configs) {
|
||||
auto t = create_transport(cfg);
|
||||
if (!t->start()) {
|
||||
SRV_WRN("MCP warmup: failed to spawn '%s': %s\n", cfg.name.c_str(), t->diagnostics().c_str());
|
||||
continue;
|
||||
}
|
||||
// bound warmup per server so an unresponsive one can't stall startup for the full per-call timeout
|
||||
const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(MCP_WARMUP_TIMEOUT_SECONDS);
|
||||
auto should_stop = [this, deadline]() {
|
||||
return stopping.load() || std::chrono::steady_clock::now() >= deadline;
|
||||
};
|
||||
auto tools = t->list_tools(should_stop);
|
||||
SRV_INF("MCP warmup: '%s' discovered %zu tools\n", cfg.name.c_str(), tools.size());
|
||||
discovered.insert(discovered.end(), tools.begin(), tools.end());
|
||||
t->close();
|
||||
}
|
||||
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
registry.swap(discovered);
|
||||
}
|
||||
|
||||
std::vector<server_mcp_tool_def> server_mcp::list_tools() const {
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
return registry;
|
||||
}
|
||||
|
||||
json server_mcp::call_tool(const std::string & server_name,
|
||||
const std::string & tool_name,
|
||||
const json & arguments,
|
||||
const std::function<bool()> & should_stop) {
|
||||
auto transport = get_or_create(server_name);
|
||||
if (!transport) {
|
||||
return {{"error", "MCP server unavailable: " + server_name}};
|
||||
}
|
||||
|
||||
auto stop = [this, &should_stop]() {
|
||||
return stopping.load() || (should_stop && should_stop());
|
||||
};
|
||||
return transport->call_tool(tool_name, arguments, stop);
|
||||
}
|
||||
|
||||
std::shared_ptr<server_mcp_transport> server_mcp::get_or_create(const std::string & name) {
|
||||
std::vector<std::shared_ptr<server_mcp_transport>> to_close; // closed after unlock
|
||||
std::shared_ptr<server_mcp_transport> result;
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
if (stopping.load()) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto now = std::chrono::steady_clock::now();
|
||||
auto dead_it = dead_servers.find(name);
|
||||
if (dead_it != dead_servers.end()) {
|
||||
if (now < dead_it->second) {
|
||||
return nullptr;
|
||||
}
|
||||
dead_servers.erase(dead_it);
|
||||
}
|
||||
|
||||
auto it = transports.find(name);
|
||||
if (it != transports.end()) {
|
||||
if (it->second->is_alive()) {
|
||||
return it->second;
|
||||
}
|
||||
SRV_WRN("MCP '%s' is no longer alive: %s\n", name.c_str(), it->second->diagnostics().c_str());
|
||||
to_close.push_back(std::move(it->second));
|
||||
transports.erase(it);
|
||||
}
|
||||
|
||||
const server_mcp_server_config * cfg = find_config(name);
|
||||
if (cfg) {
|
||||
auto fresh = create_transport(*cfg);
|
||||
if (fresh->start() && fresh->is_alive()) {
|
||||
transports[name] = fresh;
|
||||
result = fresh;
|
||||
} else {
|
||||
SRV_WRN("MCP '%s': failed to start: %s\n", name.c_str(), fresh->diagnostics().c_str());
|
||||
to_close.push_back(std::move(fresh));
|
||||
dead_servers[name] = now + std::chrono::seconds(MCP_COOLDOWN_SECONDS);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (auto & t : to_close) {
|
||||
t->close(); // blocking call, no leaks
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
#pragma once
|
||||
|
||||
#include "server-common.h"
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <functional>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
//
|
||||
// Configuration (Cursor-compatible "mcpServers" JSON)
|
||||
//
|
||||
|
||||
struct server_mcp_server_config {
|
||||
std::string name; // config key, e.g. "filesystem"
|
||||
std::string command;
|
||||
std::vector<std::string> args;
|
||||
std::map<std::string, std::string> env; // merged over the parent env
|
||||
std::string cwd;
|
||||
int timeout_ms = 30000; // per-tool-call timeout
|
||||
|
||||
// throw on parse errors; missing "mcpServers" yields an empty list; entries without a "command" are skipped
|
||||
static std::vector<server_mcp_server_config> parse_from_json(const std::string & json_str);
|
||||
static std::vector<server_mcp_server_config> parse_cursor_format(const json & j);
|
||||
};
|
||||
|
||||
// a tool advertised by an MCP server
|
||||
struct server_mcp_tool_def {
|
||||
std::string server_name;
|
||||
std::string name; // bare tool name, no "<server>_" prefix
|
||||
std::string description;
|
||||
json input_schema; // JSON Schema for the arguments, or null
|
||||
};
|
||||
|
||||
//
|
||||
// server_mcp_transport: one MCP server session.
|
||||
//
|
||||
// caller --send_rpc--> to_server --[writer]--> framing --> server
|
||||
// caller <--send_rpc-- from_server <--[reader]-- framing <-- server
|
||||
//
|
||||
// each queue item is one complete serialized JSON message.
|
||||
// subclass owns byte I/O and framing; base owns JSON and the JSON-RPC session (handshake, id correlation).
|
||||
//
|
||||
|
||||
struct server_mcp_transport {
|
||||
std::string name;
|
||||
int timeout_ms = 30000;
|
||||
|
||||
server_pipe<std::string> to_server; // serialized messages we send to the server
|
||||
server_pipe<std::string> from_server; // serialized messages read from the server
|
||||
|
||||
virtual ~server_mcp_transport() = default;
|
||||
|
||||
virtual bool start() = 0;
|
||||
virtual void close() = 0; // blocking and idempotent
|
||||
virtual bool is_alive() const = 0; // never blocks behind an in-flight send_rpc()
|
||||
|
||||
// human-readable diagnostics for logging when the transport fails/dies
|
||||
// (example: last RPC error, plus any transport-specific detail)
|
||||
// may run on a different thread than send_rpc(), so last_error is read under rpc_mutex
|
||||
virtual std::string diagnostics() {
|
||||
std::lock_guard<std::mutex> lock(rpc_mutex);
|
||||
return last_error;
|
||||
}
|
||||
|
||||
std::vector<server_mcp_tool_def> list_tools(const std::function<bool()> & should_stop);
|
||||
|
||||
json call_tool(const std::string & tool_name,
|
||||
const json & arguments,
|
||||
const std::function<bool()> & should_stop);
|
||||
|
||||
protected:
|
||||
// per-transport: send_rpc() holds it across the reply wait, so sharing it would stall every server behind one slow call. guards all members below.
|
||||
std::mutex rpc_mutex;
|
||||
uint64_t next_id = 1; // reset to 1 per (re)spawn
|
||||
bool initialized = false;
|
||||
std::string last_error;
|
||||
std::vector<server_mcp_tool_def> tools;
|
||||
|
||||
// both assume rpc_mutex is already held by the public caller
|
||||
bool ensure_init(const std::function<bool()> & should_stop); // initialize handshake, once
|
||||
json send_rpc(const json & request, const std::function<bool()> & should_stop); // returns the reply or an {"error": ...}
|
||||
};
|
||||
|
||||
//
|
||||
// server_mcp_stdio: child process, NDJSON JSON-RPC over stdio (stderr drained to the debug log)
|
||||
//
|
||||
|
||||
struct server_mcp_stdio : server_mcp_transport {
|
||||
explicit server_mcp_stdio(const server_mcp_server_config & config);
|
||||
~server_mcp_stdio() override;
|
||||
|
||||
bool start() override;
|
||||
void close() override;
|
||||
bool is_alive() const override;
|
||||
std::string diagnostics() override;
|
||||
|
||||
private:
|
||||
server_mcp_server_config config;
|
||||
|
||||
// defined in the .cpp so <windows.h> stays out of this header
|
||||
struct process_handle;
|
||||
std::unique_ptr<process_handle> proc;
|
||||
|
||||
std::thread reader; // child stdout -> NDJSON de-framing -> from_server
|
||||
std::thread writer; // to_server -> NDJSON framing -> child stdin
|
||||
std::thread errlog; // child stderr -> debug log (must be drained or the child blocks)
|
||||
|
||||
// cleared by close() or by the reader on stdout EOF; read without rpc_mutex
|
||||
std::atomic<bool> running{false};
|
||||
|
||||
// bounded tail of the child's stderr, for diagnostics when it dies
|
||||
std::mutex err_mu;
|
||||
std::string err_tail;
|
||||
|
||||
void reader_loop();
|
||||
void writer_loop();
|
||||
void errlog_loop();
|
||||
void join_pumps();
|
||||
};
|
||||
|
||||
//
|
||||
// server_mcp
|
||||
// declare before the HTTP context so it outlives every /tools handler.
|
||||
//
|
||||
|
||||
class server_mcp {
|
||||
public:
|
||||
server_mcp() = default;
|
||||
~server_mcp();
|
||||
|
||||
// parse the MCP config from params (file and/or inline JSON),
|
||||
// then spawn each server once, list its tools, and shut it down
|
||||
// throws on config parse errors; spawn failures are logged.
|
||||
void start(const common_params & params);
|
||||
|
||||
// true until start() has parsed at least one server from the config
|
||||
bool empty() const { return configs.empty(); }
|
||||
|
||||
std::vector<server_mcp_tool_def> list_tools() const;
|
||||
|
||||
// lazily (re)spawns the transport. returns the MCP result or an {"error": ...}. should_stop is OR-ed with the manager's cancel flag.
|
||||
json call_tool(const std::string & server_name,
|
||||
const std::string & tool_name,
|
||||
const json & arguments,
|
||||
const std::function<bool()> & should_stop = nullptr);
|
||||
|
||||
// flip the cancel flag so in-flight calls return; blocking teardown is in the destructor. call before the HTTP server drains.
|
||||
// note: multiple calls are idempotent
|
||||
void shutdown();
|
||||
|
||||
private:
|
||||
std::vector<server_mcp_server_config> configs;
|
||||
|
||||
mutable std::mutex mutex; // guards transports, dead_servers, registry
|
||||
|
||||
// shared_ptr: call_tool() hands a transport to the caller and drops the lock for the blocking RPC, so a concurrent evict/respawn must not destroy it mid-call
|
||||
std::map<std::string, std::shared_ptr<server_mcp_transport>> transports;
|
||||
std::map<std::string, std::chrono::steady_clock::time_point> dead_servers; // spawn-failure cooldown
|
||||
std::vector<server_mcp_tool_def> registry;
|
||||
|
||||
std::atomic<bool> stopping{false};
|
||||
|
||||
const server_mcp_server_config * find_config(const std::string & name) const;
|
||||
|
||||
// the only place that names a concrete transport
|
||||
std::shared_ptr<server_mcp_transport> create_transport(const server_mcp_server_config & cfg);
|
||||
|
||||
// nullptr during cooldown or shutdown
|
||||
std::shared_ptr<server_mcp_transport> get_or_create(const std::string & name);
|
||||
};
|
||||
+65
-117
@@ -8,10 +8,10 @@
|
||||
#include "preset.h"
|
||||
#include "download.h"
|
||||
#include "http.h"
|
||||
#include "subproc.h"
|
||||
|
||||
#include <cpp-httplib/httplib.h> // TODO: remove this once we use HTTP client from download.h
|
||||
#include <optional>
|
||||
#include <sheredom/subprocess.h>
|
||||
|
||||
#include <functional>
|
||||
#include <optional>
|
||||
@@ -49,43 +49,24 @@ extern char **environ;
|
||||
#define CHILD_ADDR "127.0.0.1"
|
||||
|
||||
struct server_subproc {
|
||||
std::optional<subprocess_s> sproc; // empty while in DOWNLOADING state
|
||||
common_subproc sproc; // not yet spawned while in DOWNLOADING state
|
||||
std::atomic<bool> stopped{false}; // set to cancel a download or signal child process exit
|
||||
|
||||
subprocess_s & get() {
|
||||
GGML_ASSERT(sproc.has_value() && "subprocess not initialized");
|
||||
return sproc.value();
|
||||
}
|
||||
|
||||
bool is_alive() {
|
||||
return sproc.has_value() && subprocess_alive(&sproc.value());
|
||||
return sproc.alive();
|
||||
}
|
||||
|
||||
void request_exit() {
|
||||
if (sproc.has_value()) {
|
||||
FILE * stdin_file = subprocess_stdin(&sproc.value());
|
||||
if (stdin_file) {
|
||||
fprintf(stdin_file, "%s\n", CMD_ROUTER_TO_CHILD_EXIT);
|
||||
fflush(stdin_file);
|
||||
}
|
||||
FILE * stdin_file = sproc.stdin_file();
|
||||
if (stdin_file) {
|
||||
fprintf(stdin_file, "%s\n", CMD_ROUTER_TO_CHILD_EXIT);
|
||||
fflush(stdin_file);
|
||||
}
|
||||
stopped.store(true, std::memory_order_relaxed);
|
||||
}
|
||||
|
||||
void terminate() {
|
||||
if (!sproc.has_value()) {
|
||||
return;
|
||||
}
|
||||
#if defined(_WIN32)
|
||||
if (sproc->hProcess == NULL) {
|
||||
return;
|
||||
}
|
||||
#else
|
||||
if (sproc->child <= 0) {
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
subprocess_terminate(&sproc.value());
|
||||
sproc.terminate();
|
||||
}
|
||||
};
|
||||
|
||||
@@ -711,18 +692,6 @@ std::optional<server_model_meta> server_models::get_meta(const std::string & nam
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
// helper to convert vector<string> to char **
|
||||
// pointers are only valid as long as the original vector is valid
|
||||
static std::vector<char *> to_char_ptr_array(const std::vector<std::string> & vec) {
|
||||
std::vector<char *> result;
|
||||
result.reserve(vec.size() + 1);
|
||||
for (const auto & s : vec) {
|
||||
result.push_back(const_cast<char*>(s.c_str()));
|
||||
}
|
||||
result.push_back(nullptr);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<server_model_meta> server_models::get_all_meta() {
|
||||
std::unique_lock<std::mutex> lk(mutex);
|
||||
if (need_reload) {
|
||||
@@ -845,15 +814,10 @@ void server_models::load(const std::string & name, const load_options & opts) {
|
||||
}
|
||||
inst.meta.args = child_args; // save for debugging
|
||||
|
||||
std::vector<char *> argv = to_char_ptr_array(child_args);
|
||||
std::vector<char *> envp = to_char_ptr_array(child_env);
|
||||
|
||||
// TODO @ngxson : maybe separate stdout and stderr in the future
|
||||
// so that we can use stdout for commands and stderr for logging
|
||||
int options = subprocess_option_no_window | subprocess_option_combined_stdout_stderr;
|
||||
inst.subproc->sproc.emplace();
|
||||
int result = subprocess_create_ex(argv.data(), options, envp.data(), nullptr, &inst.subproc->get());
|
||||
if (result != 0) {
|
||||
if (!inst.subproc->sproc.create(child_args, options, child_env)) {
|
||||
throw std::runtime_error("failed to spawn server instance");
|
||||
}
|
||||
}
|
||||
@@ -867,8 +831,8 @@ void server_models::load(const std::string & name, const load_options & opts) {
|
||||
stop_timeout = inst.meta.stop_timeout,
|
||||
child_mode = opts.mode
|
||||
]() {
|
||||
FILE * stdin_file = subprocess_stdin(&child_proc->get());
|
||||
FILE * stdout_file = subprocess_stdout(&child_proc->get()); // combined stdout/stderr
|
||||
FILE * stdin_file = child_proc->sproc.stdin_file();
|
||||
FILE * stdout_file = child_proc->sproc.stdout_file(); // combined stdout/stderr
|
||||
|
||||
std::thread log_thread([&]() {
|
||||
// read stdout/stderr and forward to main server log
|
||||
@@ -942,9 +906,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
|
||||
}
|
||||
|
||||
// get the exit code
|
||||
int exit_code = 0;
|
||||
subprocess_join(&child_proc->get(), &exit_code);
|
||||
subprocess_destroy(&child_proc->get());
|
||||
int exit_code = child_proc->sproc.join();
|
||||
|
||||
// update status and exit code
|
||||
if (child_mode == SERVER_CHILD_MODE_DOWNLOAD) {
|
||||
@@ -1210,7 +1172,7 @@ bool server_models::ensure_model_ready(const std::string & name) {
|
||||
return true;
|
||||
}
|
||||
|
||||
server_http_res_ptr server_models::proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used) {
|
||||
server_http_res_ptr server_models::proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used, bool detached) {
|
||||
auto meta = get_meta(name);
|
||||
if (!meta.has_value()) {
|
||||
throw std::runtime_error("model name=" + name + " is not found");
|
||||
@@ -1236,7 +1198,10 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co
|
||||
req.headers,
|
||||
req.body,
|
||||
req.files,
|
||||
req.should_stop,
|
||||
// a detached request belongs to a replay session that outlives the client socket:
|
||||
// it reaches the child even when the downstream died during the load wait, the
|
||||
// session buffer is the recipient and DELETE remains the stop
|
||||
detached ? std::function<bool()>([]() { return false; }) : req.should_stop,
|
||||
base_params.timeout_read,
|
||||
base_params.timeout_write
|
||||
);
|
||||
@@ -1507,13 +1472,9 @@ static bool router_validate_model(std::string & name, server_models & models, bo
|
||||
}
|
||||
// resolve alias to canonical model name
|
||||
name = meta->name;
|
||||
if (models_autoload) {
|
||||
models.ensure_model_ready(name);
|
||||
} else {
|
||||
if (!meta->is_running()) {
|
||||
res_err(res, format_error_response("model is not loaded", ERROR_TYPE_INVALID_REQUEST));
|
||||
return false;
|
||||
}
|
||||
if (!models_autoload && !meta->is_running()) {
|
||||
res_err(res, format_error_response("model is not loaded", ERROR_TYPE_INVALID_REQUEST));
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -1567,6 +1528,10 @@ static std::optional<server_model_meta> resolve_child_for_conv(
|
||||
}
|
||||
|
||||
void server_models_routes::init_routes() {
|
||||
if (!common_subproc::is_supported()) {
|
||||
throw std::runtime_error("subprocess is not enabled on this build");
|
||||
}
|
||||
|
||||
this->get_router_props = [this](const server_http_req & req) {
|
||||
std::string name = req.get_param("model");
|
||||
if (name.empty()) {
|
||||
@@ -1602,6 +1567,9 @@ void server_models_routes::init_routes() {
|
||||
if (!router_validate_model(name, models, autoload, error_res)) {
|
||||
return error_res;
|
||||
}
|
||||
if (autoload) {
|
||||
models.ensure_model_ready(name);
|
||||
}
|
||||
return models.proxy_request(req, method, name, false);
|
||||
};
|
||||
|
||||
@@ -1615,12 +1583,23 @@ void server_models_routes::init_routes() {
|
||||
return error_res;
|
||||
}
|
||||
// remember which child serves this conversation so the stream routes can route straight
|
||||
// to it without polling, keyed on the exact conv id from the header
|
||||
// to it without polling, keyed on the exact conv id from the header. registered before
|
||||
// the load wait so a stop issued while the model loads can erase the entry and cancel
|
||||
// this request instead of leaving an orphan generation
|
||||
std::string conv_id = server_stream_conv_id_from_headers(req.headers);
|
||||
if (!conv_id.empty()) {
|
||||
models.conv_models.remember(conv_id, name);
|
||||
uint64_t ticket = models.conv_models.remember(conv_id, name);
|
||||
bool waited = autoload && models.ensure_model_ready(name);
|
||||
if (ticket != 0 && !models.conv_models.alive(conv_id, ticket)) {
|
||||
SRV_INF("request for conv_id=%s cancelled while model name=%s was loading\n",
|
||||
conv_id.c_str(), name.c_str());
|
||||
res_err(error_res, format_error_response(
|
||||
"request cancelled by a stop while the model was loading", ERROR_TYPE_INVALID_REQUEST));
|
||||
return error_res;
|
||||
}
|
||||
return models.proxy_request(req, method, name, true); // update last usage for POST request only
|
||||
// a session request that waited for a load detaches from the client socket: the
|
||||
// client may have dropped during the wait (page reload) and the session buffer must
|
||||
// still receive the generation for a later resume
|
||||
return models.proxy_request(req, method, name, true, waited && ticket != 0); // update last usage for POST request only
|
||||
};
|
||||
|
||||
this->post_router_models_load = [this](const server_http_req & req) {
|
||||
@@ -1813,7 +1792,7 @@ void server_models_routes::init_routes() {
|
||||
};
|
||||
|
||||
this->router_stream_get = [this](const server_http_req & req) {
|
||||
// GET /v1/stream/<conv_id>?from=N. resolve the owning child from the conv_id -> model
|
||||
// GET /v1/stream?conv_id=<id>&from=N. resolve the owning child from the conv_id -> model
|
||||
// map, 404 when nothing maps
|
||||
auto res = std::make_unique<server_http_res>();
|
||||
std::string conv_id = req.get_param("conv_id");
|
||||
@@ -1823,13 +1802,24 @@ void server_models_routes::init_routes() {
|
||||
}
|
||||
std::optional<server_model_meta> owner = resolve_child_for_conv(models, conv_id);
|
||||
if (!owner.has_value()) {
|
||||
res_err(res, format_error_response("Stream not found or expired", ERROR_TYPE_NOT_FOUND));
|
||||
// a registered conv whose model is still loading earns a retry: the session appears
|
||||
// once the load ends and the pending request reaches the child
|
||||
auto tracked = models.conv_models.lookup(conv_id);
|
||||
auto meta = tracked.has_value() ? models.get_meta(*tracked) : std::nullopt;
|
||||
bool transient = meta.has_value() && (meta->status == SERVER_MODEL_STATUS_LOADING ||
|
||||
meta->status == SERVER_MODEL_STATUS_DOWNLOADING ||
|
||||
meta->status == SERVER_MODEL_STATUS_DOWNLOADED);
|
||||
if (transient) {
|
||||
res_err(res, format_error_response("Stream owner model is loading, retry later", ERROR_TYPE_UNAVAILABLE));
|
||||
} else {
|
||||
res_err(res, format_error_response("Stream not found or expired", ERROR_TYPE_NOT_FOUND));
|
||||
}
|
||||
return res;
|
||||
}
|
||||
std::string from = req.get_param("from");
|
||||
std::string child_path = "/v1/stream/" + encode_qs(conv_id);
|
||||
std::string child_path = "/v1/stream?conv_id=" + encode_qs(conv_id);
|
||||
if (!from.empty()) {
|
||||
child_path += "?from=" + from;
|
||||
child_path += "&from=" + from;
|
||||
}
|
||||
SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n",
|
||||
owner->name.c_str(), owner->port, child_path.c_str());
|
||||
@@ -1909,7 +1899,7 @@ void server_models_routes::init_routes() {
|
||||
};
|
||||
|
||||
this->router_stream_delete = [this](const server_http_req & req) {
|
||||
// DELETE /v1/stream/<conv_id>. resolve the owning child via the map and forward only to
|
||||
// DELETE /v1/stream?conv_id=<id>. resolve the owning child via the map and forward only to
|
||||
// it, evict_and_cancel is idempotent on the child
|
||||
auto res = std::make_unique<server_http_res>();
|
||||
std::string conv_id = req.get_param("conv_id");
|
||||
@@ -1917,7 +1907,7 @@ void server_models_routes::init_routes() {
|
||||
res_err(res, format_error_response("Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST));
|
||||
return res;
|
||||
}
|
||||
std::string child_path = "/v1/stream/" + encode_qs(conv_id);
|
||||
std::string child_path = "/v1/stream?conv_id=" + encode_qs(conv_id);
|
||||
auto owner = resolve_child_for_conv(models, conv_id);
|
||||
if (owner.has_value()) {
|
||||
httplib::Client cli(CHILD_ADDR, owner->port);
|
||||
@@ -1926,6 +1916,11 @@ void server_models_routes::init_routes() {
|
||||
cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000);
|
||||
auto resp = cli.Delete(child_path.c_str());
|
||||
(void) resp; // the child logs its own miss when the session is unknown there
|
||||
} else if (auto tracked = models.conv_models.lookup(conv_id); tracked.has_value()) {
|
||||
// the entry exists but its model is still loading: the forget below erases it,
|
||||
// which cancels the request parked in proxy_post before the generation starts
|
||||
SRV_INF("router stop for conv_id=%s while model name=%s is loading, cancelling the pending request\n",
|
||||
conv_id.c_str(), tracked->c_str());
|
||||
} else {
|
||||
SRV_WRN("router stop for unknown conv_id=%s, no owning child in the conv map\n",
|
||||
conv_id.c_str());
|
||||
@@ -1944,53 +1939,6 @@ void server_models_routes::init_routes() {
|
||||
// server_http_proxy
|
||||
//
|
||||
|
||||
// simple implementation of a pipe
|
||||
// used for streaming data between threads
|
||||
template<typename T>
|
||||
struct pipe_t {
|
||||
std::mutex mutex;
|
||||
std::condition_variable cv;
|
||||
std::queue<T> queue;
|
||||
std::atomic<bool> writer_closed{false};
|
||||
std::atomic<bool> reader_closed{false};
|
||||
void close_write() {
|
||||
writer_closed.store(true, std::memory_order_relaxed);
|
||||
cv.notify_all();
|
||||
}
|
||||
void close_read() {
|
||||
reader_closed.store(true, std::memory_order_relaxed);
|
||||
cv.notify_all();
|
||||
}
|
||||
bool read(T & output, const std::function<bool()> & should_stop) {
|
||||
std::unique_lock<std::mutex> lk(mutex);
|
||||
constexpr auto poll_interval = std::chrono::milliseconds(500);
|
||||
while (true) {
|
||||
if (!queue.empty()) {
|
||||
output = std::move(queue.front());
|
||||
queue.pop();
|
||||
return true;
|
||||
}
|
||||
if (writer_closed.load()) {
|
||||
return false; // clean EOF
|
||||
}
|
||||
if (should_stop()) {
|
||||
close_read(); // signal broken pipe to writer
|
||||
return false; // cancelled / reader no longer alive
|
||||
}
|
||||
cv.wait_for(lk, poll_interval);
|
||||
}
|
||||
}
|
||||
bool write(T && data) {
|
||||
std::lock_guard<std::mutex> lk(mutex);
|
||||
if (reader_closed.load()) {
|
||||
return false; // broken pipe
|
||||
}
|
||||
queue.push(std::move(data));
|
||||
cv.notify_one();
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
static std::string to_lower_copy(const std::string & value) {
|
||||
std::string lowered(value.size(), '\0');
|
||||
std::transform(value.begin(), value.end(), lowered.begin(), [](unsigned char c) { return std::tolower(c); });
|
||||
@@ -2100,7 +2048,7 @@ server_http_proxy::server_http_proxy(
|
||||
) {
|
||||
// shared between reader and writer threads
|
||||
auto cli = std::make_shared<httplib::ClientImpl>(host, port);
|
||||
auto pipe = std::make_shared<pipe_t<msg_t>>();
|
||||
auto pipe = std::make_shared<server_pipe<msg_t>>();
|
||||
|
||||
if (scheme == "https") {
|
||||
#ifdef CPPHTTPLIB_OPENSSL_SUPPORT
|
||||
|
||||
@@ -134,12 +134,24 @@ private:
|
||||
// proxy_request forwards a POST carrying an X-Conversation-Id. best effort: a stale entry just
|
||||
// makes the child answer not found and the client recovers. owns its lock, one mutex per struct
|
||||
struct conv_model_tracker {
|
||||
void remember(const std::string & conv_id, const std::string & model) {
|
||||
// returns the ticket of this registration, 0 when nothing was registered. erasing or
|
||||
// replacing the entry invalidates the ticket, which is how a stop cancels a request
|
||||
// parked in the model load wait
|
||||
uint64_t remember(const std::string & conv_id, const std::string & model) {
|
||||
if (conv_id.empty() || model.empty()) {
|
||||
return;
|
||||
return 0;
|
||||
}
|
||||
std::lock_guard<std::mutex> lock(mu);
|
||||
map[conv_id] = model;
|
||||
uint64_t ticket = next_ticket++;
|
||||
map[conv_id] = { model, ticket };
|
||||
return ticket;
|
||||
}
|
||||
|
||||
// false means a stop erased the entry or a newer request replaced it
|
||||
bool alive(const std::string & conv_id, uint64_t ticket) {
|
||||
std::lock_guard<std::mutex> lock(mu);
|
||||
auto it = map.find(conv_id);
|
||||
return it != map.end() && it->second.ticket == ticket;
|
||||
}
|
||||
|
||||
std::optional<std::string> lookup(const std::string & conv_id) {
|
||||
@@ -151,7 +163,7 @@ private:
|
||||
if (it == map.end()) {
|
||||
return std::nullopt;
|
||||
}
|
||||
return it->second;
|
||||
return it->second.model;
|
||||
}
|
||||
|
||||
void forget(const std::string & conv_id) {
|
||||
@@ -163,8 +175,13 @@ private:
|
||||
}
|
||||
|
||||
private:
|
||||
std::mutex mu;
|
||||
std::unordered_map<std::string, std::string> map;
|
||||
struct entry_t {
|
||||
std::string model;
|
||||
uint64_t ticket;
|
||||
};
|
||||
std::mutex mu;
|
||||
uint64_t next_ticket = 1;
|
||||
std::unordered_map<std::string, entry_t> map;
|
||||
};
|
||||
|
||||
common_preset_context ctx_preset;
|
||||
@@ -249,7 +266,7 @@ public:
|
||||
bool ensure_model_ready(const std::string & name);
|
||||
|
||||
// proxy an HTTP request to the model instance
|
||||
server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used);
|
||||
server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used, bool detached = false);
|
||||
|
||||
// handle message sent from server_child::notify_to_router()
|
||||
// raw input must starts with CMD_CHILD_TO_ROUTER_STATE, followed by a JSON string
|
||||
|
||||
@@ -390,21 +390,40 @@ std::vector<std::unique_ptr<field>> make_llama_cmpl_schema(const common_params &
|
||||
ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get<std::string>(), false, true);
|
||||
}));
|
||||
|
||||
add((new field_str("reasoning_budget_end_tag"))
|
||||
->set_desc("Token string marking the end of the reasoning budget section")
|
||||
add((new field_json("reasoning_budget_end_tags"))
|
||||
->add_alias("reasoning_budget_end_tag")
|
||||
->set_desc("Token strings marking the end of the reasoning budget section; the first is forced when the budget expires")
|
||||
->set_handler([&](field_eval_context & ctx, const json & data) {
|
||||
GGML_ASSERT(ctx.vocab != nullptr);
|
||||
std::string end_tag = data.at("reasoning_budget_end_tag").get<std::string>();
|
||||
ctx.params.sampling.reasoning_budget_end = common_tokenize(ctx.vocab, end_tag, false, true);
|
||||
ctx.params.sampling.reasoning_budget_end.clear();
|
||||
if (data.contains("reasoning_budget_end_tags")) {
|
||||
for (const auto & t : data.at("reasoning_budget_end_tags")) {
|
||||
std::string tag = t.get<std::string>();
|
||||
if (!tag.empty()) {
|
||||
ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true));
|
||||
}
|
||||
}
|
||||
} else if (data.contains("reasoning_budget_end_tag")) {
|
||||
std::string tag = data.at("reasoning_budget_end_tag").get<std::string>();
|
||||
if (!tag.empty()) {
|
||||
ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true));
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
add((new field_str("reasoning_budget_message"))
|
||||
->set_desc("Message to prepend to the reasoning budget end tag when forcing it")
|
||||
->set_handler([&](field_eval_context & ctx, const json & data) {
|
||||
GGML_ASSERT(ctx.vocab != nullptr);
|
||||
std::string end_tag = json_value(data, "reasoning_budget_end_tag", std::string());
|
||||
std::string message = data.at("reasoning_budget_message").get<std::string>();
|
||||
ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message + end_tag, false, true);
|
||||
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());
|
||||
}
|
||||
ctx.params.sampling.reasoning_budget_forced = std::move(end_tag);
|
||||
}
|
||||
}));
|
||||
|
||||
add((new field_json("logit_bias"))
|
||||
@@ -546,7 +565,7 @@ 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 toks, forced=%zu toks\n",
|
||||
SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n",
|
||||
budget, params.sampling.generation_prompt.c_str(),
|
||||
params.sampling.reasoning_budget_start.size(),
|
||||
params.sampling.reasoning_budget_end.size(),
|
||||
|
||||
@@ -453,7 +453,7 @@ static server_http_res_ptr make_error_response(int status, const std::string & m
|
||||
|
||||
server_http_context::handler_t server_stream_make_get_handler() {
|
||||
return [](const server_http_req & req) -> server_http_res_ptr {
|
||||
// GET /v1/stream/<conv_id>?from=N replays buffered SSE bytes then blocks for live
|
||||
// GET /v1/stream?conv_id=<id>&from=N replays buffered SSE bytes then blocks for live
|
||||
// bytes until the session finalizes, streamed as text/event-stream for EventSource
|
||||
std::string conv_id = req.get_param("conv_id");
|
||||
if (conv_id.empty()) {
|
||||
@@ -560,13 +560,13 @@ server_http_context::handler_t server_stream_make_lookup_handler() {
|
||||
|
||||
server_http_context::handler_t server_stream_make_delete_handler() {
|
||||
return [](const server_http_req & req) -> server_http_res_ptr {
|
||||
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer and evicts
|
||||
// DELETE /v1/stream?conv_id=<id> is the explicit user Stop, cancels the producer and evicts
|
||||
// the buffer. idempotent, returns 204 even if the session was already gone
|
||||
std::string conv_id = req.get_param("conv_id");
|
||||
if (conv_id.empty()) {
|
||||
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
|
||||
}
|
||||
SRV_TRC("DELETE /v1/stream/%s -> evict_and_cancel\n", conv_id.c_str());
|
||||
SRV_TRC("DELETE /v1/stream conv_id=%s -> evict_and_cancel\n", conv_id.c_str());
|
||||
g_stream_sessions.evict_and_cancel(conv_id);
|
||||
auto res = std::make_unique<server_http_res>();
|
||||
res->status = 204;
|
||||
@@ -621,7 +621,7 @@ bool server_res_spipe::conn_alive() {
|
||||
|
||||
bool server_res_spipe::should_stop() {
|
||||
if (spipe) {
|
||||
// note: if DELETE /v1/stream/<conv_id> is called, is_cancelled() will be true
|
||||
// note: if DELETE /v1/stream is called for this conv, is_cancelled() will be true
|
||||
return spipe->is_cancelled();
|
||||
} else {
|
||||
return !conn_alive();
|
||||
|
||||
@@ -45,7 +45,13 @@ void server_stream_session_manager_start();
|
||||
void server_stream_session_manager_stop();
|
||||
|
||||
// route handler factories wired under /v1/stream/* by server.cpp
|
||||
// child-side handlers for the resumable stream routes. the conv id travels in the conv_id
|
||||
// query string because it can embed a model name containing slashes (org/repo), which the
|
||||
// decoded path would split before the param is captured
|
||||
server_http_context::handler_t server_stream_make_get_handler();
|
||||
// POST /v1/streams/lookup with body {"conversation_ids": [...]}: only answers for ids the
|
||||
// caller already owns (the WebUI passes the convs visible in its sidebar), the server never
|
||||
// lists ids it has not been asked about, so a random caller cannot enumerate live sessions
|
||||
server_http_context::handler_t server_stream_make_lookup_handler();
|
||||
server_http_context::handler_t server_stream_make_delete_handler();
|
||||
|
||||
|
||||
@@ -63,6 +63,8 @@ json task_params::to_json(bool only_metrics) const {
|
||||
{"mirostat", sampling.mirostat},
|
||||
{"mirostat_tau", sampling.mirostat_tau},
|
||||
{"mirostat_eta", sampling.mirostat_eta},
|
||||
{"adaptive_target", sampling.adaptive_target},
|
||||
{"adaptive_decay", sampling.adaptive_decay},
|
||||
{"max_tokens", n_predict},
|
||||
{"n_predict", n_predict}, // TODO: deduplicate?
|
||||
{"n_keep", n_keep},
|
||||
@@ -114,6 +116,8 @@ json task_params::to_json(bool only_metrics) const {
|
||||
{"mirostat", sampling.mirostat},
|
||||
{"mirostat_tau", sampling.mirostat_tau},
|
||||
{"mirostat_eta", sampling.mirostat_eta},
|
||||
{"adaptive_target", sampling.adaptive_target},
|
||||
{"adaptive_decay", sampling.adaptive_decay},
|
||||
{"stop", antiprompt},
|
||||
{"max_tokens", n_predict},
|
||||
{"n_predict", n_predict}, // TODO: deduplicate?
|
||||
|
||||
+113
-27
@@ -1,18 +1,19 @@
|
||||
#include "server-tools.h"
|
||||
|
||||
#include <sheredom/subprocess.h>
|
||||
#include "subproc.h"
|
||||
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <regex>
|
||||
#include <thread>
|
||||
#include <chrono>
|
||||
#include <ctime>
|
||||
#include <atomic>
|
||||
#include <cstring>
|
||||
#include <climits>
|
||||
#include <algorithm>
|
||||
#include <unordered_set>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
@@ -24,7 +25,7 @@ json server_tool::to_json() const {
|
||||
return {
|
||||
{"display_name", display_name},
|
||||
{"tool", name},
|
||||
{"type", "builtin"},
|
||||
{"type", type()},
|
||||
{"permissions", json{
|
||||
{"write", permission_write}
|
||||
}},
|
||||
@@ -137,15 +138,14 @@ public:
|
||||
const std::function<bool(const std::string &)> & on_chunk = nullptr) const override {
|
||||
exec_result res;
|
||||
|
||||
subprocess_s proc;
|
||||
auto argv = to_cstr_vec(args);
|
||||
common_subproc proc;
|
||||
|
||||
int options = subprocess_option_no_window
|
||||
| subprocess_option_combined_stdout_stderr
|
||||
| subprocess_option_inherit_environment
|
||||
| subprocess_option_search_user_path;
|
||||
|
||||
if (subprocess_create(argv.data(), options, &proc) != 0) {
|
||||
if (!proc.create(args, options)) {
|
||||
res.output = "failed to spawn process";
|
||||
return res;
|
||||
}
|
||||
@@ -158,14 +158,14 @@ public:
|
||||
while (!done.load()) {
|
||||
if (std::chrono::steady_clock::now() >= deadline) {
|
||||
timed_out.store(true);
|
||||
subprocess_terminate(&proc);
|
||||
proc.terminate();
|
||||
return;
|
||||
}
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||
}
|
||||
});
|
||||
|
||||
FILE * f = subprocess_stdout(&proc);
|
||||
FILE * f = proc.stdout_file();
|
||||
std::string output;
|
||||
bool truncated = false;
|
||||
if (f) {
|
||||
@@ -176,7 +176,7 @@ public:
|
||||
if (output.size() + len <= max_output) {
|
||||
output.append(buf, len);
|
||||
if (on_chunk && !on_chunk(std::string(buf, len))) {
|
||||
subprocess_terminate(&proc);
|
||||
proc.terminate();
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
@@ -194,8 +194,7 @@ public:
|
||||
timeout_thread.join();
|
||||
}
|
||||
|
||||
subprocess_join(&proc, &res.exit_code);
|
||||
subprocess_destroy(&proc);
|
||||
res.exit_code = proc.join();
|
||||
|
||||
res.output = output;
|
||||
res.timed_out = timed_out.load();
|
||||
@@ -206,16 +205,6 @@ public:
|
||||
}
|
||||
|
||||
private:
|
||||
static std::vector<char *> to_cstr_vec(const std::vector<std::string> & v) {
|
||||
std::vector<char *> r;
|
||||
r.reserve(v.size() + 1);
|
||||
for (const auto & s : v) {
|
||||
r.push_back(const_cast<char *>(s.c_str()));
|
||||
}
|
||||
r.push_back(nullptr);
|
||||
return r;
|
||||
}
|
||||
|
||||
static const std::unordered_set<std::string> & junk_dir_names() {
|
||||
static const std::unordered_set<std::string> names = {
|
||||
".git", ".svn", ".hg", "node_modules", "__pycache__",
|
||||
@@ -1048,16 +1037,42 @@ struct server_tool_get_datetime : server_tool {
|
||||
{"type", "function"},
|
||||
{"function", {
|
||||
{"name", name},
|
||||
{"description", "Returns the current date and time"},
|
||||
{"description", "Returns the current date and time in UTC"},
|
||||
{"parameters", {
|
||||
{"type", "object"},
|
||||
{"properties", {
|
||||
{"format", {
|
||||
{"type", "string"},
|
||||
{"description",
|
||||
"strftime()-style format string for the output (default: \"%Y-%m-%dT%H:%M:%SZ\", "
|
||||
"e.g. ISO 8601). Choose your own format if you need something else, "
|
||||
"e.g. \"%A, %B %d %Y\" for a human-readable date."},
|
||||
}},
|
||||
}},
|
||||
}},
|
||||
}},
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json, server_tool::stream *) const override {
|
||||
auto now = std::chrono::system_clock::now();
|
||||
auto time = std::chrono::system_clock::to_time_t(now);
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string format = json_value(params, "format", std::string("%Y-%m-%dT%H:%M:%SZ"));
|
||||
|
||||
return {{"result", std::ctime(&time)}};
|
||||
auto now = std::chrono::system_clock::now();
|
||||
auto time = std::chrono::system_clock::to_time_t(now);
|
||||
std::tm tm_utc;
|
||||
#ifdef _WIN32
|
||||
gmtime_s(&tm_utc, &time);
|
||||
#else
|
||||
gmtime_r(&time, &tm_utc);
|
||||
#endif
|
||||
|
||||
char buf[256];
|
||||
size_t len = std::strftime(buf, sizeof(buf), format.c_str(), &tm_utc);
|
||||
if (len == 0) {
|
||||
return {{"error", "invalid format string"}};
|
||||
}
|
||||
|
||||
return {{"result", std::string(buf, len)}};
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1102,6 +1117,49 @@ struct server_tools_res : server_http_res {
|
||||
}
|
||||
};
|
||||
|
||||
//
|
||||
// server_mcp_tool: exposes one tool from a running MCP server as a server_tool.
|
||||
//
|
||||
struct server_mcp_tool : server_tool {
|
||||
std::string server_name;
|
||||
std::string tool_name;
|
||||
server_mcp_tool_def def;
|
||||
server_mcp & mcp_mgr;
|
||||
|
||||
server_mcp_tool(server_mcp_tool_def d, server_mcp & mgr)
|
||||
: server_name(d.server_name)
|
||||
, tool_name(d.name)
|
||||
, def(std::move(d))
|
||||
, mcp_mgr(mgr)
|
||||
{
|
||||
name = server_name + "_" + tool_name;
|
||||
display_name = name;
|
||||
permission_write = false;
|
||||
support_stream = false;
|
||||
}
|
||||
|
||||
std::string type() const override { return "mcp"; }
|
||||
|
||||
json get_definition() const override {
|
||||
json schema = def.input_schema;
|
||||
if (schema.is_null() || !schema.is_object()) {
|
||||
schema = json::object();
|
||||
}
|
||||
return {
|
||||
{"type", "function"},
|
||||
{"function", {
|
||||
{"name", name},
|
||||
{"description", def.description},
|
||||
{"parameters", schema},
|
||||
}},
|
||||
};
|
||||
}
|
||||
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
return mcp_mgr.call_tool(server_name, tool_name, params);
|
||||
}
|
||||
};
|
||||
|
||||
static server_tool & find_tool(std::vector<std::unique_ptr<server_tool>> & tools, const std::string & name, bool require_stream) {
|
||||
for (auto & t : tools) {
|
||||
if (t->name == name) {
|
||||
@@ -1130,8 +1188,13 @@ static std::vector<std::unique_ptr<server_tool>> build_tools() {
|
||||
return tools;
|
||||
}
|
||||
|
||||
void server_tools::setup(const std::vector<std::string> & enabled_tools) {
|
||||
void server_tools::setup(const std::vector<std::string> & enabled_tools,
|
||||
server_mcp & mcp_mgr) {
|
||||
if (!enabled_tools.empty()) {
|
||||
if (!common_subproc::is_supported()) {
|
||||
throw std::runtime_error("subprocess is not enabled on this build");
|
||||
}
|
||||
|
||||
std::unordered_set<std::string> enabled_set(enabled_tools.begin(), enabled_tools.end());
|
||||
auto all_tools = build_tools();
|
||||
|
||||
@@ -1161,6 +1224,29 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools) {
|
||||
}
|
||||
}
|
||||
|
||||
// append MCP tools, skipping any that collide with a built-in or another MCP tool of the same "<server>_<tool>" name
|
||||
if (!mcp_mgr.empty()) {
|
||||
std::unordered_set<std::string> seen_names;
|
||||
for (auto & t : tools) {
|
||||
seen_names.insert(t->name);
|
||||
}
|
||||
size_t n_added = 0;
|
||||
for (const auto & def : mcp_mgr.list_tools()) {
|
||||
std::string mcp_name = def.server_name + "_" + def.name;
|
||||
if (seen_names.count(mcp_name)) {
|
||||
SRV_WRN("MCP tool \"%s\" from server \"%s\" collides with an existing tool, skipping\n",
|
||||
mcp_name.c_str(), def.server_name.c_str());
|
||||
continue;
|
||||
}
|
||||
seen_names.insert(mcp_name);
|
||||
tools.push_back(std::make_unique<server_mcp_tool>(def, mcp_mgr));
|
||||
n_added++;
|
||||
}
|
||||
if (n_added > 0) {
|
||||
SRV_INF("Added %zu MCP tools\n", n_added);
|
||||
}
|
||||
}
|
||||
|
||||
handle_get = [this](const server_http_req &) -> server_http_res_ptr {
|
||||
auto res = std::make_unique<server_http_res>();
|
||||
try {
|
||||
|
||||
@@ -3,9 +3,11 @@
|
||||
#include "server-common.h"
|
||||
#include "server-http.h"
|
||||
#include "server-queue.h"
|
||||
#include "server-mcp.h"
|
||||
|
||||
#include <atomic>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
|
||||
struct server_tool {
|
||||
std::string name;
|
||||
@@ -15,6 +17,7 @@ struct server_tool {
|
||||
|
||||
virtual ~server_tool() = default;
|
||||
virtual json get_definition() const = 0;
|
||||
virtual std::string type() const { return "builtin"; }
|
||||
|
||||
struct stream {
|
||||
server_response & qr;
|
||||
@@ -34,7 +37,8 @@ struct server_tools {
|
||||
server_response queue_res;
|
||||
std::atomic<int> res_id{0};
|
||||
|
||||
void setup(const std::vector<std::string> & enabled_tools);
|
||||
void setup(const std::vector<std::string> & enabled_tools,
|
||||
server_mcp & mcp_mgr);
|
||||
|
||||
server_http_context::handler_t handle_get;
|
||||
server_http_context::handler_t handle_post;
|
||||
|
||||
+33
-15
@@ -88,6 +88,11 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t
|
||||
int llama_server(int argc, char ** argv) {
|
||||
std::setlocale(LC_NUMERIC, "C");
|
||||
|
||||
#ifndef _WIN32
|
||||
// Ignore SIGPIPE so the server does not crash if an MCP child exits while we are writing to its stdin
|
||||
signal(SIGPIPE, SIG_IGN);
|
||||
#endif
|
||||
|
||||
// own arguments required by this example
|
||||
common_params params;
|
||||
|
||||
@@ -157,6 +162,9 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
params.model_alias.insert(model_name);
|
||||
}
|
||||
|
||||
// note: this is guaranteed to out-live ctx_http and tools
|
||||
server_mcp mcp_mgr;
|
||||
|
||||
// struct that contains llama context and inference
|
||||
server_context ctx_server;
|
||||
|
||||
@@ -264,10 +272,8 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
ctx_http.get ("/slots", ex_wrapper(routes.get_slots));
|
||||
ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots));
|
||||
|
||||
// resumable streaming, the conversation_id is the session identity end to end. router and
|
||||
// child wire different handlers under the same paths: a child binds the local session
|
||||
// factories, the router binds proxies that resolve the owning child through the
|
||||
// conv_id -> model map
|
||||
// resumable streaming: a child binds the local session factories, the router binds
|
||||
// proxies that resolve the owning child, see server-stream.h
|
||||
server_http_context::handler_t stream_get_h;
|
||||
server_http_context::handler_t streams_lookup_h;
|
||||
server_http_context::handler_t stream_delete_h;
|
||||
@@ -280,12 +286,9 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
streams_lookup_h = server_stream_make_lookup_handler();
|
||||
stream_delete_h = server_stream_make_delete_handler();
|
||||
}
|
||||
ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h));
|
||||
// POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids
|
||||
// you already own (the WebUI passes the convs visible in its sidebar). the server never
|
||||
// lists ids it has not been asked about, so a random caller cannot enumerate live sessions
|
||||
ctx_http.get ("/v1/stream", ex_wrapper(stream_get_h));
|
||||
ctx_http.post("/v1/streams/lookup", ex_wrapper(streams_lookup_h));
|
||||
ctx_http.del ("/v1/stream/:conv_id", ex_wrapper(stream_delete_h));
|
||||
ctx_http.del ("/v1/stream", ex_wrapper(stream_delete_h));
|
||||
|
||||
// Google Cloud Platform (Vertex AI) compat
|
||||
ctx_http.register_gcp_compat();
|
||||
@@ -326,17 +329,28 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
ctx_http.post("/cors-proxy", ex_wrapper(res_403));
|
||||
}
|
||||
|
||||
// EXPERIMENTAL built-in tools
|
||||
if (!params.server_tools.empty()) {
|
||||
try {
|
||||
mcp_mgr.start(params);
|
||||
} catch (const std::exception & e) {
|
||||
SRV_ERR("MCP starting failed: %s\n", e.what());
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (!params.server_tools.empty() || !mcp_mgr.empty()) {
|
||||
try {
|
||||
tools.setup(params.server_tools);
|
||||
tools.setup(params.server_tools, mcp_mgr);
|
||||
} catch (const std::exception & e) {
|
||||
SRV_ERR("tools setup failed: %s\n", e.what());
|
||||
return 1;
|
||||
}
|
||||
ctx_http.get ("/tools", ex_wrapper(tools.handle_get));
|
||||
ctx_http.post("/tools", ex_wrapper(tools.handle_post));
|
||||
warn_names.push_back("built-in tools (experimental)");
|
||||
if (!params.server_tools.empty()) {
|
||||
warn_names.push_back("built-in tools (experimental)");
|
||||
}
|
||||
if (!mcp_mgr.empty()) {
|
||||
warn_names.push_back("MCP servers (experimental)");
|
||||
}
|
||||
} else {
|
||||
ctx_http.get ("/tools", ex_wrapper(res_403));
|
||||
ctx_http.post("/tools", ex_wrapper(res_403));
|
||||
@@ -378,7 +392,7 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
if (is_router_server) {
|
||||
SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n");
|
||||
|
||||
clean_up = [&models_routes]() {
|
||||
clean_up = [&models_routes, &mcp_mgr]() {
|
||||
SRV_INF("%s: cleaning up before exit...\n", __func__);
|
||||
// stop the session GC first, it finalizes live sessions and wakes pending readers
|
||||
server_stream_session_manager_stop();
|
||||
@@ -386,6 +400,7 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
models_routes->stopping.store(true); // maybe redundant, but just to be safe
|
||||
models_routes->models.unload_all();
|
||||
}
|
||||
mcp_mgr.shutdown();
|
||||
llama_backend_free();
|
||||
};
|
||||
|
||||
@@ -401,17 +416,19 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
// important to disconnect any SSE clients
|
||||
models_routes->stopping.store(true);
|
||||
}
|
||||
mcp_mgr.shutdown();
|
||||
ctx_http.stop();
|
||||
};
|
||||
|
||||
} else {
|
||||
// setup clean up function, to be called before exit
|
||||
clean_up = [&ctx_http, &ctx_server]() {
|
||||
clean_up = [&ctx_http, &ctx_server, &mcp_mgr]() {
|
||||
SRV_INF("%s: cleaning up before exit...\n", __func__);
|
||||
// stop the session GC first, it finalizes live sessions and wakes pending readers
|
||||
server_stream_session_manager_stop();
|
||||
ctx_http.stop();
|
||||
ctx_server.terminate();
|
||||
mcp_mgr.shutdown();
|
||||
llama_backend_free();
|
||||
};
|
||||
|
||||
@@ -444,6 +461,7 @@ int llama_server(common_params & params, int argc, char ** argv) {
|
||||
SRV_INF("%s", "model loaded\n");
|
||||
|
||||
shutdown_handler = [&](int) {
|
||||
mcp_mgr.shutdown();
|
||||
// this will unblock start_loop()
|
||||
ctx_server.terminate();
|
||||
};
|
||||
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Minimal MCP server that writes notification + response in a single write() with no flush.
|
||||
This reproduces the buffering bug where read_message() can strand the response.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Echo back the input message",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message": {"type": "string"}
|
||||
},
|
||||
"required": ["message"]
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
def handle_initialize(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "burst-test", "version": "1.0"}
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_list(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {"tools": TOOLS}
|
||||
}
|
||||
|
||||
def handle_tools_call(params, req_id):
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if tool_name == "echo":
|
||||
message = arguments.get("message", "")
|
||||
notif = {
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/progress",
|
||||
"params": {"progress": 50, "total": 100}
|
||||
}
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"echo: {message}"}]
|
||||
}
|
||||
}
|
||||
# Single os.write() call: both lines land in one pipe packet atomically.
|
||||
# This is the key difference from mcp_malformed_server.py which flushes between writes.
|
||||
data = (json.dumps(notif) + "\n" + json.dumps(response) + "\n").encode("utf-8")
|
||||
os.write(sys.stdout.fileno(), data)
|
||||
return None # already written
|
||||
else:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
|
||||
}
|
||||
return response
|
||||
|
||||
HANDLERS = {
|
||||
"initialize": handle_initialize,
|
||||
"tools/list": handle_tools_list,
|
||||
"tools/call": handle_tools_call,
|
||||
}
|
||||
|
||||
def main():
|
||||
# Use line-buffered text mode for regular responses, but the burst write
|
||||
# uses os.write() directly to guarantee a single kernel write().
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
params = request.get("params", {})
|
||||
|
||||
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
|
||||
if req_id is None:
|
||||
continue
|
||||
|
||||
handler = HANDLERS.get(method)
|
||||
if handler:
|
||||
response = handler(params, req_id)
|
||||
if response is not None:
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
else:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32601, "message": f"Method not found: {method}"}
|
||||
}
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP server that crashes after receiving a specific tool call.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
def handle_initialize(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "crash-test", "version": "1.0"}
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_list(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"tools": [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Echo back the input message",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message": {"type": "string"}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "crash",
|
||||
"description": "Crash the server",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_call(params, req_id):
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if tool_name == "echo":
|
||||
message = arguments.get("message", "")
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"echo: {message}"}]
|
||||
}
|
||||
}
|
||||
elif tool_name == "crash":
|
||||
# Send a partial response then exit
|
||||
sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": "crashing..."}]}}) + "\n")
|
||||
sys.stdout.flush()
|
||||
os._exit(1)
|
||||
else:
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
|
||||
}
|
||||
|
||||
HANDLERS = {
|
||||
"initialize": handle_initialize,
|
||||
"tools/list": handle_tools_list,
|
||||
"tools/call": handle_tools_call,
|
||||
}
|
||||
|
||||
def main():
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
params = request.get("params", {})
|
||||
|
||||
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
|
||||
if req_id is None:
|
||||
continue
|
||||
|
||||
handler = HANDLERS.get(method)
|
||||
if handler:
|
||||
response = handler(params, req_id)
|
||||
else:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32601, "message": f"Method not found: {method}"}
|
||||
}
|
||||
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Minimal MCP server for testing.
|
||||
Implements JSON-RPC 2.0 over stdio (line-delimited JSON).
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Ensure we use python3 from the current environment
|
||||
if sys.platform == "win32":
|
||||
# On Windows, we need to use the same python interpreter
|
||||
pass
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Echo back the input message",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message": {"type": "string", "description": "Message to echo"}
|
||||
},
|
||||
"required": ["message"]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "add",
|
||||
"description": "Add two numbers",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {"type": "number"},
|
||||
"b": {"type": "number"}
|
||||
},
|
||||
"required": ["a", "b"]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "fail_once",
|
||||
"description": "Fails on first call, succeeds on subsequent calls",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
_state = {"fail_once_called": False}
|
||||
|
||||
def handle_initialize(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "echo-test", "version": "1.0"}
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_list(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {"tools": TOOLS}
|
||||
}
|
||||
|
||||
def handle_tools_call(params, req_id):
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if tool_name == "echo":
|
||||
message = arguments.get("message", "")
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"echo: {message}"}]
|
||||
}
|
||||
}
|
||||
elif tool_name == "add":
|
||||
a = arguments.get("a", 0)
|
||||
b = arguments.get("b", 0)
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": str(a + b)}]
|
||||
}
|
||||
}
|
||||
elif tool_name == "fail_once":
|
||||
if not _state["fail_once_called"]:
|
||||
_state["fail_once_called"] = True
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32000, "message": "transient error"}
|
||||
}
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": "ok"}]
|
||||
}
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
|
||||
}
|
||||
|
||||
def handle_ping(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {}
|
||||
}
|
||||
|
||||
HANDLERS = {
|
||||
"initialize": handle_initialize,
|
||||
"tools/list": handle_tools_list,
|
||||
"tools/call": handle_tools_call,
|
||||
"ping": handle_ping,
|
||||
}
|
||||
|
||||
def main():
|
||||
# Use unbuffered output
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
params = request.get("params", {})
|
||||
|
||||
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
|
||||
if req_id is None:
|
||||
continue
|
||||
|
||||
handler = HANDLERS.get(method)
|
||||
if handler:
|
||||
response = handler(params, req_id)
|
||||
else:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32601, "message": f"Method not found: {method}"}
|
||||
}
|
||||
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,100 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP server (NDJSON JSON-RPC over stdio) that spawns a long-lived grandchild which inherits
|
||||
this process's stdin/stdout/stderr and keeps them open.
|
||||
|
||||
This reproduces the reader-teardown deadlock: killing the direct MCP child (SIGKILL, which is
|
||||
all subprocess_terminate() does) does NOT close the stdout/stderr pipe write ends, because the
|
||||
grandchild still holds them. A server that reads those pipes with a blocking read would then
|
||||
wait forever for an EOF that never arrives, hanging teardown (both warmup shutdown at startup
|
||||
and process shutdown). The polled, running-aware reader must exit regardless.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
# Spawn a grandchild that inherits our std handles (fds 0/1/2 = the MCP pipes) and lives well
|
||||
# past any teardown in the tests. We do NOT redirect its stdio, so it keeps the pipe write ends
|
||||
# open even after this process is killed.
|
||||
subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Echo back the input message",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"message": {"type": "string", "description": "Message to echo"}},
|
||||
"required": ["message"],
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def handle_initialize(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "grandchild-test", "version": "1.0"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def handle_tools_list(params, req_id):
|
||||
return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}
|
||||
|
||||
|
||||
def handle_tools_call(params, req_id):
|
||||
if params.get("name") == "echo":
|
||||
message = params.get("arguments", {}).get("message", "")
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {"content": [{"type": "text", "text": f"echo: {message}"}]},
|
||||
}
|
||||
return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32602, "message": "Unknown tool"}}
|
||||
|
||||
|
||||
HANDLERS = {
|
||||
"initialize": handle_initialize,
|
||||
"tools/list": handle_tools_list,
|
||||
"tools/call": handle_tools_call,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
params = request.get("params", {})
|
||||
|
||||
if req_id is None:
|
||||
continue # notification, no response
|
||||
|
||||
handler = HANDLERS.get(method)
|
||||
if handler:
|
||||
response = handler(params, req_id)
|
||||
else:
|
||||
response = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}}
|
||||
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP server that sends malformed responses and notifications during requests.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
def handle_initialize(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "malformed-test", "version": "1.0"}
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_list(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"tools": [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Echo back the input message",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"message": {"type": "string"}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_call(params, req_id):
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if tool_name == "echo":
|
||||
message = arguments.get("message", "")
|
||||
# Send a notification first (no id field)
|
||||
notif = {
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/progress",
|
||||
"params": {"progress": 50, "total": 100}
|
||||
}
|
||||
sys.stdout.write(json.dumps(notif) + "\n")
|
||||
sys.stdout.flush()
|
||||
# Then send the actual response
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"echo: {message}"}]
|
||||
}
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
|
||||
}
|
||||
|
||||
HANDLERS = {
|
||||
"initialize": handle_initialize,
|
||||
"tools/list": handle_tools_list,
|
||||
"tools/call": handle_tools_call,
|
||||
}
|
||||
|
||||
def main():
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
# Send malformed JSON response
|
||||
sys.stdout.write("THIS IS NOT JSON\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
params = request.get("params", {})
|
||||
|
||||
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
|
||||
if req_id is None:
|
||||
continue
|
||||
|
||||
handler = HANDLERS.get(method)
|
||||
if handler:
|
||||
response = handler(params, req_id)
|
||||
else:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32601, "message": f"Method not found: {method}"}
|
||||
}
|
||||
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+132
@@ -0,0 +1,132 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP server that sleeps before responding, for timeout testing.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
import time
|
||||
import argparse
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"name": "sleep",
|
||||
"description": "Sleep for a given number of seconds",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"seconds": {"type": "number", "description": "Seconds to sleep"}
|
||||
},
|
||||
"required": ["seconds"]
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
def handle_initialize(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "slow-test", "version": "1.0"}
|
||||
}
|
||||
}
|
||||
|
||||
def handle_tools_list(params, req_id):
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {"tools": TOOLS}
|
||||
}
|
||||
|
||||
def handle_tools_call(params, req_id):
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if tool_name == "sleep":
|
||||
seconds = arguments.get("seconds", 1)
|
||||
time.sleep(seconds)
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"slept {seconds}s"}]
|
||||
}
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
|
||||
}
|
||||
|
||||
HANDLERS = {
|
||||
"initialize": handle_initialize,
|
||||
"tools/list": handle_tools_list,
|
||||
"tools/call": handle_tools_call,
|
||||
}
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--delay", type=float, default=5.0, help="Delay in seconds for sleep tool")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Override the sleep duration
|
||||
global handle_tools_call
|
||||
def handle_tools_call(params, req_id):
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if tool_name == "sleep":
|
||||
seconds = arguments.get("seconds", args.delay)
|
||||
time.sleep(seconds)
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"slept {seconds}s"}]
|
||||
}
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32602, "message": f"Unknown tool: {tool_name}"}
|
||||
}
|
||||
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1)
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
params = request.get("params", {})
|
||||
|
||||
# JSON-RPC 2.0: a message without an id is a notification and must not receive a response
|
||||
if req_id is None:
|
||||
continue
|
||||
|
||||
handler = HANDLERS.get(method)
|
||||
if handler:
|
||||
response = handler(params, req_id)
|
||||
else:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"error": {"code": -32601, "message": f"Method not found: {method}"}
|
||||
}
|
||||
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -66,6 +66,8 @@ def test_completion_stream(prompt: str, n_predict: int, re_content: str, n_promp
|
||||
assert server.n_predict is not None
|
||||
assert data["generation_settings"]["n_predict"] == min(n_predict, server.n_predict)
|
||||
assert data["generation_settings"]["seed"] == server.seed
|
||||
assert "adaptive_target" in data["generation_settings"]
|
||||
assert "adaptive_decay" in data["generation_settings"]
|
||||
assert match_regex(re_content, content)
|
||||
else:
|
||||
assert len(data["tokens"]) > 0
|
||||
|
||||
@@ -0,0 +1,718 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Tests for MCP server integration via the /tools endpoint.
|
||||
|
||||
Invariants verified:
|
||||
1. MCP tools appear in /tools listing when configured
|
||||
2. MCP tools use <server>_<tool> naming
|
||||
3. MCP tools can be invoked and return correct results
|
||||
4. Misconfigured MCP servers do not crash the server
|
||||
5. Multiple MCP servers can be configured simultaneously
|
||||
6. Warmup populates the tool list at startup
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from utils import *
|
||||
|
||||
# Path to the test MCP server fixture
|
||||
FIXTURES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "fixtures")
|
||||
MCP_ECHO_SERVER = os.path.join(FIXTURES_DIR, "mcp_echo_server.py")
|
||||
|
||||
server: ServerProcess
|
||||
|
||||
|
||||
def _mcp_config_json(servers: dict) -> str:
|
||||
"""Create a JSON config string for --mcp-servers-json."""
|
||||
return json.dumps({"mcpServers": servers})
|
||||
|
||||
|
||||
def _start_server_with_mcp(mcp_json: str, **kwargs) -> ServerProcess:
|
||||
"""Helper to start a router server with MCP config."""
|
||||
srv = ServerPreset.router()
|
||||
srv.server_tools = "all"
|
||||
srv.no_ui = True
|
||||
srv.server_port = 8085 # avoid conflict with load_all() which uses 8080
|
||||
srv.mcp_servers_json = mcp_json
|
||||
for k, v in kwargs.items():
|
||||
setattr(srv, k, v)
|
||||
srv.start()
|
||||
return srv
|
||||
|
||||
|
||||
def test_mcp_tools_listed_in_tools_endpoint():
|
||||
"""MCP tools should appear in GET /tools with server:tool naming."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
|
||||
tools = res.body
|
||||
assert isinstance(tools, list), f"Expected list, got {type(tools)}"
|
||||
|
||||
# Find MCP tools - name is in "tool" field or definition.function.name
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
|
||||
mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
|
||||
assert len(mcp_tools) >= 2, f"Expected at least 2 echo_ tools, got {len(mcp_tools)}: {mcp_tools}"
|
||||
|
||||
tool_names = {get_tool_name(t) for t in mcp_tools}
|
||||
assert "echo_echo" in tool_names
|
||||
assert "echo_add" in tool_names
|
||||
|
||||
# Verify tool structure
|
||||
echo_tool = next(t for t in mcp_tools if get_tool_name(t) == "echo_echo")
|
||||
assert "description" in echo_tool or "definition" in echo_tool
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_tool_invocation():
|
||||
"""MCP tools should be callable via POST /tools and return correct results."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# Call echo_echo
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_echo",
|
||||
"params": {"message": "hello world"}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
body = res.body
|
||||
assert "error" not in body, body
|
||||
# The result format depends on the tool implementation
|
||||
# For MCP tools, it should contain the tool result
|
||||
assert "plain_text_response" in body or "result" in body or "content" in body, body
|
||||
|
||||
# Call echo_add
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_add",
|
||||
"params": {"a": 3, "b": 5}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
body = res.body
|
||||
assert "error" not in body, body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_bad_command_does_not_crash():
|
||||
"""A misconfigured MCP server should not crash the llama-server."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"nonexistent": {
|
||||
"command": "this_executable_does_not_exist_12345",
|
||||
"args": [],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# Server should still be healthy
|
||||
res = server.make_request("GET", "/health")
|
||||
assert res.status_code == 200, res.body
|
||||
|
||||
# Builtin tools should still work
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
tools = res.body
|
||||
# Should have builtin tools but no MCP tools from the bad server
|
||||
mcp_tools = [t for t in tools if t.get("name", "").startswith("nonexistent_")]
|
||||
assert len(mcp_tools) == 0, f"Expected no nonexistent_ tools, got {mcp_tools}"
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_multiple_servers():
|
||||
"""Multiple MCP servers can be configured simultaneously."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
},
|
||||
"echo2": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
|
||||
tools = res.body
|
||||
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
|
||||
echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
|
||||
echo2_tools = [t for t in tools if get_tool_name(t).startswith("echo2_")]
|
||||
|
||||
assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}"
|
||||
assert len(echo2_tools) >= 2, f"Expected echo2_ tools, got {echo2_tools}"
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_tools_not_listed_when_not_configured():
|
||||
"""Without MCP config, no MCP tools should appear."""
|
||||
global server
|
||||
server = ServerPreset.router()
|
||||
server.server_tools = "all"
|
||||
server.no_ui = True
|
||||
server.server_port = 8085
|
||||
server.start()
|
||||
|
||||
try:
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
|
||||
tools = res.body
|
||||
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
|
||||
# Should only have builtin tools, no server: prefixed tools
|
||||
mcp_tools = [t for t in tools if ":" in get_tool_name(t)]
|
||||
assert len(mcp_tools) == 0, f"Expected no MCP tools, got {mcp_tools}"
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_fail_once_tool_eventual_success():
|
||||
"""Test that a tool that fails once eventually succeeds (tests instance respawn)."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# First call should succeed (warmup already spawned and shut down the instance,
|
||||
# but the first actual tool call will spawn a fresh instance)
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_fail_once",
|
||||
"params": {}
|
||||
})
|
||||
# It might fail on first call if the warmup instance was shut down
|
||||
# and a new instance is spawned. The fail_once state is per-process,
|
||||
# so a fresh process will fail once then succeed.
|
||||
# Actually, warmup spawns, lists, then shuts down. So the first tool call
|
||||
# spawns a new process which will fail once.
|
||||
assert res.status_code in (200, 500), res.body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_tools_via_json_config_file():
|
||||
"""Test that --mcp-servers-config (file) works as well as --mcp-servers-json."""
|
||||
global server
|
||||
config = {
|
||||
"mcpServers": {
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
|
||||
json.dump(config, f)
|
||||
config_path = f.name
|
||||
|
||||
try:
|
||||
server = ServerPreset.router()
|
||||
server.server_tools = "all"
|
||||
server.no_ui = True
|
||||
server.server_port = 8085
|
||||
server.mcp_servers_config = config_path
|
||||
server.start()
|
||||
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
|
||||
tools = res.body
|
||||
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
|
||||
mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
|
||||
assert len(mcp_tools) >= 2, f"Expected echo_ tools, got {mcp_tools}"
|
||||
finally:
|
||||
os.unlink(config_path)
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_tools_slot_independent():
|
||||
"""MCP tools should work without any slot concept; /tools is slot-independent."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# Call /tools without any slot binding - should succeed
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_echo",
|
||||
"params": {"message": "hello"}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
body = res.body
|
||||
assert "error" not in body, body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_concurrent_tool_calls():
|
||||
"""Concurrent POST /tools to same MCP server should all succeed."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
def call_tool():
|
||||
return server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_echo",
|
||||
"params": {"message": "hi"}
|
||||
})
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as executor:
|
||||
futures = [executor.submit(call_tool) for _ in range(10)]
|
||||
results = [f.result() for f in futures]
|
||||
|
||||
for res in results:
|
||||
assert res.status_code == 200, res.body
|
||||
assert "error" not in res.body, res.body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_tool_timeout():
|
||||
"""Tool call should timeout if MCP server is too slow."""
|
||||
global server
|
||||
MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py")
|
||||
mcp_json = _mcp_config_json({
|
||||
"slow": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_SLOW_SERVER, "--delay", "5"],
|
||||
"timeout_ms": 500
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "slow_sleep",
|
||||
"params": {"seconds": 5}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
body = res.body
|
||||
assert "error" in body, body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_warmup_partial_failure():
|
||||
"""Good server's tools should appear even if bad server fails warmup."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"good": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
},
|
||||
"bad": {
|
||||
"command": "nonexistent",
|
||||
"args": []
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
tools = res.body
|
||||
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
|
||||
# good server tools should be present
|
||||
assert any("good_" in get_tool_name(t) for t in tools), f"Expected good: tools in {tools}"
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_notification_during_request():
|
||||
"""Notification during request should not be returned as response."""
|
||||
global server
|
||||
MCP_MALFORMED_SERVER = os.path.join(FIXTURES_DIR, "mcp_malformed_server.py")
|
||||
mcp_json = _mcp_config_json({
|
||||
"notifying": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_MALFORMED_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "notifying_echo",
|
||||
"params": {"message": "hi"}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
body = res.body
|
||||
assert "error" not in body, body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_instance_respawn_after_crash():
|
||||
"""Tool call after process crash should respawn and succeed."""
|
||||
global server
|
||||
MCP_CRASH_SERVER = os.path.join(FIXTURES_DIR, "mcp_crash_server.py")
|
||||
mcp_json = _mcp_config_json({
|
||||
"crash": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_CRASH_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# First call succeeds
|
||||
res1 = server.make_request("POST", "/tools", data={
|
||||
"tool": "crash_echo",
|
||||
"params": {"message": "hi"}
|
||||
})
|
||||
assert res1.status_code == 200, res1.body
|
||||
assert "error" not in res1.body, res1.body
|
||||
|
||||
# Second call should also succeed (respawned instance)
|
||||
res2 = server.make_request("POST", "/tools", data={
|
||||
"tool": "crash_echo",
|
||||
"params": {"message": "hi2"}
|
||||
})
|
||||
assert res2.status_code == 200, res2.body
|
||||
assert "error" not in res2.body, res2.body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
|
||||
|
||||
def test_mcp_fail_once_eventual_success_verified():
|
||||
"""Verify that fail_once tool eventually succeeds after respawn."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# First call may fail (fresh process)
|
||||
res1 = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_fail_once",
|
||||
"params": {}
|
||||
})
|
||||
# Second call should succeed
|
||||
res2 = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_fail_once",
|
||||
"params": {}
|
||||
})
|
||||
assert res2.status_code == 200, res2.body
|
||||
assert "error" not in res2.body, res2.body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_config_file_errors():
|
||||
"""Invalid JSON config and missing file should cause server to fail to start."""
|
||||
# Invalid JSON - server should fail to start
|
||||
server = ServerPreset.router()
|
||||
server.server_tools = "all"
|
||||
server.no_ui = True
|
||||
server.server_port = 8085
|
||||
server.mcp_servers_json = "not valid json"
|
||||
try:
|
||||
server.start()
|
||||
assert False, "Server should not have started with invalid MCP JSON config"
|
||||
except RuntimeError:
|
||||
pass # Expected: server process dies due to bad config
|
||||
|
||||
# Missing file - server should fail to start
|
||||
server = ServerPreset.router()
|
||||
server.server_tools = "all"
|
||||
server.no_ui = True
|
||||
server.server_port = 8085
|
||||
server.mcp_servers_config = "/nonexistent/path.json"
|
||||
try:
|
||||
server.start()
|
||||
assert False, "Server should not have started with missing config file"
|
||||
except RuntimeError:
|
||||
pass # Expected: server process dies due to missing config
|
||||
|
||||
|
||||
def test_mcp_empty_tool_list():
|
||||
"""MCP server reporting zero tools should result in empty tool list."""
|
||||
global server
|
||||
# Create a minimal server that returns empty tools list
|
||||
empty_server = os.path.join(FIXTURES_DIR, "_empty_mcp_server.py")
|
||||
with open(empty_server, "w") as f:
|
||||
f.write('''#!/usr/bin/env python3
|
||||
import json, sys, os
|
||||
def main():
|
||||
sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line: continue
|
||||
try: request = json.loads(line)
|
||||
except: continue
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
if method == "initialize":
|
||||
resp = {"jsonrpc": "2.0", "id": req_id, "result": {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "empty", "version": "1.0"}}}
|
||||
elif method == "tools/list":
|
||||
resp = {"jsonrpc": "2.0", "id": req_id, "result": {"tools": []}}
|
||||
else:
|
||||
resp = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": "Method not found"}}
|
||||
sys.stdout.write(json.dumps(resp) + "\\n")
|
||||
sys.stdout.flush()
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
''')
|
||||
try:
|
||||
mcp_json = _mcp_config_json({
|
||||
"empty": {
|
||||
"command": sys.executable,
|
||||
"args": [empty_server],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
tools = res.body
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
mcp_tools = [t for t in tools if get_tool_name(t).startswith("empty:")]
|
||||
assert len(mcp_tools) == 0, f"Expected no empty: tools, got {mcp_tools}"
|
||||
finally:
|
||||
os.unlink(empty_server)
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_rapid_succession_calls():
|
||||
"""Many rapid calls should increment next_id correctly and correlate responses."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
for i in range(20):
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "echo_echo",
|
||||
"params": {"message": f"msg{i}"}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
assert "error" not in res.body, res.body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_notification_burst():
|
||||
"""Notification + response in a single write() with no flush should not strand the response."""
|
||||
global server
|
||||
MCP_BURST_SERVER = os.path.join(FIXTURES_DIR, "mcp_burst_server.py")
|
||||
mcp_json = _mcp_config_json({
|
||||
"burst": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_BURST_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "burst_echo",
|
||||
"params": {"message": "burst test"}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
body = res.body
|
||||
assert "error" not in body, body
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_tool_definition_shape_via_chat_completions():
|
||||
"""MCP tool definitions returned by GET /tools should have the correct shape for chat/completions."""
|
||||
global server
|
||||
mcp_json = _mcp_config_json({
|
||||
"echo": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_ECHO_SERVER],
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# Get MCP tool definitions
|
||||
res = server.make_request("GET", "/tools")
|
||||
assert res.status_code == 200, res.body
|
||||
tools = res.body
|
||||
|
||||
def get_tool_name(t):
|
||||
return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "")
|
||||
|
||||
echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")]
|
||||
assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}"
|
||||
|
||||
echo_tool = next(t for t in echo_tools if get_tool_name(t) == "echo_echo")
|
||||
definition = echo_tool.get("definition", echo_tool)
|
||||
|
||||
# Verify the definition has the standard function-calling shape
|
||||
assert definition.get("type") == "function", f"Expected type=function, got {definition.get('type')}"
|
||||
func = definition.get("function", {})
|
||||
assert "name" in func, "Missing function.name"
|
||||
assert "description" in func, "Missing function.description"
|
||||
assert "parameters" in func, f"Missing function.parameters, got keys: {list(func.keys())}"
|
||||
params = func["parameters"]
|
||||
assert params.get("type") == "object", f"Expected parameters.type=object, got {params.get('type')}"
|
||||
assert "properties" in params, "Missing parameters.properties"
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_slow_tool_call_slot_release():
|
||||
"""A slow tool call should not stall server shutdown for the full I/O timeout."""
|
||||
global server
|
||||
MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py")
|
||||
mcp_json = _mcp_config_json({
|
||||
"slow": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_SLOW_SERVER, "--delay", "10"],
|
||||
"timeout_ms": 30000
|
||||
}
|
||||
})
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# Start a slow tool call in a background thread
|
||||
def slow_call():
|
||||
return server.make_request("POST", "/tools", data={
|
||||
"tool": "slow_sleep",
|
||||
"params": {"seconds": 10}
|
||||
})
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(slow_call)
|
||||
|
||||
# Wait a moment for the call to start
|
||||
time.sleep(2)
|
||||
|
||||
# Stop the server while the tool call is in progress.
|
||||
# With global MCP instances, close_all() is called explicitly at shutdown
|
||||
# (not from slot release), so shutdown should complete promptly.
|
||||
start_time = time.time()
|
||||
server.stop()
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# The server should stop quickly, not wait for the full 30s I/O timeout.
|
||||
# With the terminating flag, send_rpc() bails out within one select()
|
||||
# slice (~50ms). This threshold MUST stay below the 5s force-kill
|
||||
# fallback in ServerProcess.stop(): without the flag, shutdown stalls
|
||||
# on the instance mutex and only completes when stop() sends SIGKILL
|
||||
# at ~5s -- which any threshold above 5 would still accept.
|
||||
assert elapsed < 3, f"Server stop took {elapsed:.1f}s, expected < 3s"
|
||||
|
||||
# Wait for the future to complete (it will get an error response or timeout)
|
||||
try:
|
||||
res = future.result(timeout=5)
|
||||
# If we got a response, it should be an error since the server stopped
|
||||
if hasattr(res, 'status_code'):
|
||||
assert res.status_code in (200, 500, 502, 503, 504), f"Unexpected status: {res.status_code}"
|
||||
except Exception:
|
||||
# Thread may have raised due to connection error - that's acceptable
|
||||
pass
|
||||
finally:
|
||||
server.stop()
|
||||
|
||||
|
||||
def test_mcp_grandchild_holding_pipes_does_not_deadlock():
|
||||
"""An MCP server that leaves a grandchild inheriting its stdout/stderr must not deadlock
|
||||
teardown.
|
||||
|
||||
subprocess_terminate() only SIGKILLs the direct MCP child, so the inherited pipe write ends
|
||||
stay open and a blocking read on them would never see EOF. That hung both warmup shutdown
|
||||
(the server would never reach "ready") and process shutdown. The polled, running-aware reader
|
||||
must exit regardless, so the server both starts and stops promptly here.
|
||||
"""
|
||||
global server
|
||||
MCP_GRANDCHILD_SERVER = os.path.join(FIXTURES_DIR, "mcp_grandchild_server.py")
|
||||
mcp_json = _mcp_config_json({
|
||||
"gc": {
|
||||
"command": sys.executable,
|
||||
"args": [MCP_GRANDCHILD_SERVER],
|
||||
}
|
||||
})
|
||||
|
||||
# If warmup teardown deadlocked, the server would never become ready and start() would time out.
|
||||
server = _start_server_with_mcp(mcp_json)
|
||||
|
||||
try:
|
||||
# invoking the tool spawns a live transport whose reader thread holds the inherited pipe
|
||||
res = server.make_request("POST", "/tools", data={
|
||||
"tool": "gc_echo",
|
||||
"params": {"message": "hello"}
|
||||
})
|
||||
assert res.status_code == 200, res.body
|
||||
assert "error" not in res.body, res.body
|
||||
|
||||
# shutdown must be prompt: a deadlocked reader-join would stall until the 5s SIGKILL
|
||||
# fallback in ServerProcess.stop(), so the threshold has to stay below that
|
||||
start = time.time()
|
||||
server.stop()
|
||||
elapsed = time.time() - start
|
||||
assert elapsed < 3, f"server shutdown took {elapsed:.1f}s (expected < 3s) — teardown likely deadlocked"
|
||||
finally:
|
||||
server.stop()
|
||||
@@ -0,0 +1,153 @@
|
||||
import json
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from urllib.parse import quote
|
||||
import pytest
|
||||
from utils import *
|
||||
|
||||
server: ServerProcess
|
||||
|
||||
# a model name with slashes exercises the query string routing of the stream routes: the id
|
||||
# cannot travel as a path param because the decoded slash would split it before capture
|
||||
MODEL = "ggml-org/tinygemma3-GGUF:Q8_0"
|
||||
STREAM_ID = f"conv-stream-test::{MODEL}"
|
||||
QS = "conv_id=" + quote(STREAM_ID, safe="")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def create_server():
|
||||
global server
|
||||
server = ServerPreset.router()
|
||||
|
||||
|
||||
def test_stream_resume_and_stop_with_slashed_model_name():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
content = ""
|
||||
for data in server.make_stream_request("POST", "/chat/completions", data={
|
||||
"model": MODEL,
|
||||
"stream": True,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}, headers={"X-Conversation-Id": STREAM_ID}):
|
||||
if data["choices"]:
|
||||
content += data["choices"][0]["delta"].get("content") or ""
|
||||
assert len(content) > 0
|
||||
|
||||
# the finished session replays from the beginning through the router
|
||||
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
|
||||
assert res.status_code == 200
|
||||
assert "data: " in str(res.body)
|
||||
|
||||
# the explicit stop reaches the owning child and evicts the session
|
||||
res = server.make_request("DELETE", f"/v1/stream?{QS}")
|
||||
assert res.status_code == 204
|
||||
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
|
||||
assert res.status_code == 404
|
||||
|
||||
|
||||
def test_stream_stop_during_model_load():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
thread_error: list[ServerError] = []
|
||||
thread_done = threading.Event()
|
||||
|
||||
def fire_post():
|
||||
try:
|
||||
for _ in server.make_stream_request("POST", "/chat/completions", data={
|
||||
"model": MODEL,
|
||||
"stream": True,
|
||||
"max_tokens": 512,
|
||||
"messages": [{"role": "user", "content": "Count from 1 to 1000."}],
|
||||
}, headers={"X-Conversation-Id": STREAM_ID}):
|
||||
pass
|
||||
except ServerError as e:
|
||||
thread_error.append(e)
|
||||
finally:
|
||||
thread_done.set()
|
||||
|
||||
t = threading.Thread(target=fire_post)
|
||||
t.start()
|
||||
|
||||
# catch the autoload window, tiny models load fast so poll aggressively
|
||||
saw_loading = False
|
||||
deadline = time.time() + 5.0
|
||||
while time.time() < deadline and not thread_done.is_set():
|
||||
res = server.make_request("GET", "/models")
|
||||
status = next(m["status"]["value"] for m in res.body["data"] if m["id"] == MODEL)
|
||||
if status == "loading":
|
||||
saw_loading = True
|
||||
break
|
||||
time.sleep(0.002)
|
||||
if not saw_loading:
|
||||
t.join()
|
||||
pytest.skip("load window too short to be observed on this machine") # ty: ignore[too-many-positional-arguments]
|
||||
|
||||
# a stop during the load cancels the parked request instead of leaving an orphan
|
||||
res = server.make_request("DELETE", f"/v1/stream?{QS}")
|
||||
assert res.status_code == 204
|
||||
assert thread_done.wait(timeout=60)
|
||||
t.join()
|
||||
assert len(thread_error) == 1
|
||||
assert thread_error[0].code == 400
|
||||
assert "cancelled" in json.dumps(thread_error[0].body)
|
||||
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
|
||||
assert res.status_code == 404
|
||||
|
||||
|
||||
def test_stream_resumes_after_reload_during_model_load():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
# raw socket client so the connection can be dropped mid load like a page reload
|
||||
body = json.dumps({
|
||||
"model": MODEL,
|
||||
"stream": True,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
})
|
||||
request = (
|
||||
f"POST /v1/chat/completions HTTP/1.1\r\n"
|
||||
f"Host: {server.server_host}:{server.server_port}\r\n"
|
||||
f"Content-Type: application/json\r\n"
|
||||
f"X-Conversation-Id: {STREAM_ID}\r\n"
|
||||
f"Content-Length: {len(body)}\r\n"
|
||||
f"Connection: close\r\n\r\n{body}"
|
||||
)
|
||||
sock = socket.create_connection((server.server_host, server.server_port))
|
||||
sock.sendall(request.encode())
|
||||
|
||||
# drop the client while the model loads, poll aggressively to catch the window
|
||||
saw_loading = False
|
||||
saw_503 = False
|
||||
deadline = time.time() + 5.0
|
||||
while time.time() < deadline:
|
||||
res = server.make_request("GET", "/models")
|
||||
status = next(m["status"]["value"] for m in res.body["data"] if m["id"] == MODEL)
|
||||
if status == "loading":
|
||||
saw_loading = True
|
||||
break
|
||||
if status == "loaded":
|
||||
break
|
||||
time.sleep(0.002)
|
||||
sock.close()
|
||||
if not saw_loading:
|
||||
pytest.skip("load window too short to be observed on this machine") # ty: ignore[too-many-positional-arguments]
|
||||
|
||||
# while the model loads the resume route answers retry later, then the session appears,
|
||||
# receives the whole generation despite the dead client, and replays from the beginning
|
||||
deadline = time.time() + 60.0
|
||||
replay = None
|
||||
while time.time() < deadline:
|
||||
res = server.make_request("GET", f"/v1/stream?{QS}&from=0")
|
||||
if res.status_code == 503:
|
||||
saw_503 = True
|
||||
elif res.status_code == 200 and "data: " in str(res.body):
|
||||
replay = res
|
||||
break
|
||||
time.sleep(0.1)
|
||||
assert saw_503, "resume during the load did not answer 503"
|
||||
assert replay is not None, "session never became resumable after the client disconnect"
|
||||
@@ -115,6 +115,8 @@ class ServerProcess:
|
||||
backend_sampling: bool = False
|
||||
gcp_compat: bool = False
|
||||
server_tools: str | None = None
|
||||
mcp_servers_config: str | None = None
|
||||
mcp_servers_json: str | None = None
|
||||
cors_origins: str | None = None
|
||||
|
||||
# session variables
|
||||
@@ -265,6 +267,10 @@ class ServerProcess:
|
||||
server_args.append("--ui-mcp-proxy")
|
||||
if self.server_tools:
|
||||
server_args.extend(["--tools", self.server_tools])
|
||||
if self.mcp_servers_config:
|
||||
server_args.extend(["--mcp-servers-config", self.mcp_servers_config])
|
||||
if self.mcp_servers_json:
|
||||
server_args.extend(["--mcp-servers-json", self.mcp_servers_json])
|
||||
if self.backend_sampling:
|
||||
server_args.append("--backend_sampling")
|
||||
if self.gcp_compat:
|
||||
|
||||
+8
-1
@@ -13,6 +13,13 @@
|
||||
|
||||
const gauge = useContextGauge();
|
||||
|
||||
// The gauge hook wraps a processing state instance that only follows the
|
||||
// live stream while its own monitoring flag is set, so the card instance
|
||||
// starts monitoring like the dial does.
|
||||
$effect(() => {
|
||||
gauge.startMonitoring();
|
||||
});
|
||||
|
||||
let cardEl = $state<HTMLElement | null>(null);
|
||||
|
||||
// Any press outside the card and outside the dial closes the card.
|
||||
@@ -44,7 +51,7 @@
|
||||
<div
|
||||
role="status"
|
||||
bind:this={cardEl}
|
||||
class="absolute z-50 w-64 -translate-x-1/2 rounded-lg border border-border/50 bg-popover p-3 text-popover-foreground shadow-lg"
|
||||
class="absolute z-50 w-64 -translate-x-1/2 rounded-lg border border-border/50 bg-popover p-3 text-sm text-popover-foreground shadow-lg ring-1 ring-foreground/10"
|
||||
style="left: {gaugePopup.centerX}px; bottom: {gaugePopup.bottom}px"
|
||||
onpointerenter={gaugeCardEnter}
|
||||
onpointerleave={gaugeCardLeave}
|
||||
|
||||
+6
-3
@@ -10,7 +10,7 @@
|
||||
} from '$lib/components/app';
|
||||
import { getMessageEditContext } from '$lib/contexts';
|
||||
import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
|
||||
import { isLoading, isChatStreaming } from '$lib/stores/chat.svelte';
|
||||
import { chatStore, isLoading, isChatStreaming } from '$lib/stores/chat.svelte';
|
||||
import { modelLoadProgressText } from '$lib/utils';
|
||||
import { MessageRole } from '$lib/enums';
|
||||
import { config } from '$lib/stores/settings.svelte';
|
||||
@@ -82,8 +82,11 @@
|
||||
let hasNoContent = $derived(!message?.content?.trim());
|
||||
let isActivelyProcessing = $derived(isCurrentlyLoading || isStreaming);
|
||||
|
||||
// during a router auto-load the message has no model yet, so target the selected one
|
||||
let loadTargetModel = $derived(message.model ?? modelsStore.selectedModelName);
|
||||
// during a router auto-load the message has no model yet: target the model frozen in the
|
||||
// persisted stream state (survives a reload), then fall back to the dropdown selection
|
||||
let loadTargetModel = $derived(
|
||||
message.model ?? chatStore.getResumeModel(message.convId) ?? modelsStore.selectedModelName
|
||||
);
|
||||
let modelLoadProgress = $derived(
|
||||
isRouter && loadTargetModel ? modelsStore.getLoadProgress(loadTargetModel) : null
|
||||
);
|
||||
|
||||
+11
-2
@@ -7,7 +7,7 @@
|
||||
import { getMessageEditContext } from '$lib/contexts';
|
||||
import { KeyboardKey, MessageRole } from '$lib/enums';
|
||||
import { config } from '$lib/stores/settings.svelte';
|
||||
import { isIMEComposing } from '$lib/utils';
|
||||
import { autoResizeTextarea, isIMEComposing } from '$lib/utils';
|
||||
|
||||
interface Props {
|
||||
class?: string;
|
||||
@@ -91,6 +91,11 @@
|
||||
resizeObserver.disconnect();
|
||||
};
|
||||
});
|
||||
$effect(() => {
|
||||
if (editCtx.isEditing && textareaElement) {
|
||||
autoResizeTextarea(textareaElement);
|
||||
}
|
||||
});
|
||||
|
||||
function toggleExpand() {
|
||||
isExpanded = !isExpanded;
|
||||
@@ -105,11 +110,15 @@
|
||||
{#if editCtx.isEditing}
|
||||
<div class="w-full max-w-[80%]">
|
||||
<textarea
|
||||
style="max-height: var(--max-message-height);"
|
||||
bind:this={textareaElement}
|
||||
value={editCtx.editedContent}
|
||||
class="min-h-[60px] w-full resize-none rounded-2xl px-3 py-2 text-sm {INPUT_CLASSES}"
|
||||
onkeydown={handleEditKeydown}
|
||||
oninput={(e) => editCtx.setContent(e.currentTarget.value)}
|
||||
oninput={(e) => {
|
||||
autoResizeTextarea(e.currentTarget);
|
||||
editCtx.setContent(e.currentTarget.value);
|
||||
}}
|
||||
placeholder="Edit system message..."
|
||||
></textarea>
|
||||
|
||||
|
||||
@@ -31,11 +31,11 @@
|
||||
import { config } from '$lib/stores/settings.svelte';
|
||||
import { serverLoading, serverError } from '$lib/stores/server.svelte';
|
||||
import { parseFilesToMessageExtras } from '$lib/utils/browser-only';
|
||||
import { onDestroy, onMount } from 'svelte';
|
||||
import { onDestroy, onMount, tick } from 'svelte';
|
||||
import ChatScreenGreeting from './ChatScreenGreeting.svelte';
|
||||
import ChatScreenActionScrollDown from './ChatScreenActionScrollDown.svelte';
|
||||
import ChatScreenDialogsAndAlerts from './ChatScreenDialogsAndAlerts.svelte';
|
||||
import { ROUTES } from '$lib/constants';
|
||||
import { LANDING_SETTLE_MAX_MS, LANDING_STABLE_FRAMES, ROUTES } from '$lib/constants';
|
||||
|
||||
let { showCenteredEmpty = false } = $props();
|
||||
|
||||
@@ -128,6 +128,41 @@
|
||||
return true;
|
||||
}
|
||||
|
||||
let lastScrolledConversationId: string | null = null;
|
||||
|
||||
// Lands at the bottom of a conversation the first time its messages
|
||||
// render, whether the route comes from another conversation or from a
|
||||
// non-conversation route. The page keeps growing after the first pin
|
||||
// without DOM mutations (content-visibility size realizations, syntax
|
||||
// highlight passes), so the instant pin repeats every frame until the
|
||||
// height settles, bailing out on user scroll or conversation change.
|
||||
async function handleMessagesReady(messageCount: number) {
|
||||
if (messageCount === 0) return;
|
||||
const id = activeConversation()?.id ?? null;
|
||||
if (!id || id === lastScrolledConversationId) return;
|
||||
lastScrolledConversationId = id;
|
||||
await tick();
|
||||
autoScroll.scrollToBottom();
|
||||
|
||||
const container = scroll.chatScrollContainer;
|
||||
if (!container) return;
|
||||
const started = performance.now();
|
||||
let stableFrames = 0;
|
||||
let lastHeight = container.scrollHeight;
|
||||
const settle = () => {
|
||||
if (autoScroll.userScrolledUp) return;
|
||||
if (activeConversation()?.id !== id) return;
|
||||
autoScroll.scrollToBottom();
|
||||
const height = container.scrollHeight;
|
||||
stableFrames = height === lastHeight ? stableFrames + 1 : 0;
|
||||
lastHeight = height;
|
||||
if (stableFrames >= LANDING_STABLE_FRAMES) return;
|
||||
if (performance.now() - started > LANDING_SETTLE_MAX_MS) return;
|
||||
requestAnimationFrame(settle);
|
||||
};
|
||||
requestAnimationFrame(settle);
|
||||
}
|
||||
|
||||
function handleSendLikeScroll() {
|
||||
if (!isMobile.current) {
|
||||
autoScroll.enable();
|
||||
@@ -246,6 +281,7 @@
|
||||
{#if !isEmpty}
|
||||
<ChatMessages
|
||||
messages={activeMessages()}
|
||||
onMessagesReady={handleMessagesReady}
|
||||
onUserAction={() => {
|
||||
handleSendLikeScroll();
|
||||
}}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user