mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-18 02:32:40 +02:00
Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 5234b9d267 | |||
| 087f94d82e | |||
| 533b18257b | |||
| ed1c3a20f5 | |||
| d8df12ebc4 | |||
| b75ecd1971 |
+3
-1
@@ -224,10 +224,12 @@ add_subdirectory(src)
|
||||
# utils, programs, examples and tests
|
||||
#
|
||||
|
||||
# mtmd needs this even when common is not built
|
||||
add_subdirectory(vendor/hash)
|
||||
|
||||
if (LLAMA_BUILD_COMMON)
|
||||
add_subdirectory(common)
|
||||
add_subdirectory(vendor/cpp-httplib)
|
||||
add_subdirectory(vendor/hash)
|
||||
endif()
|
||||
|
||||
if (LLAMA_BUILD_COMMON AND LLAMA_BUILD_TESTS AND NOT CMAKE_JS_VERSION)
|
||||
|
||||
@@ -1710,6 +1710,38 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
params.cache_ram_mib = value;
|
||||
}
|
||||
).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
|
||||
add_opt(common_arg(
|
||||
{"-cdisk", "--cache-disk"}, "PATH",
|
||||
"directory for the disk prompt cache; prompts evicted from the RAM cache are saved here and restored on later requests, including across restarts (default: disabled, requires cache-ram)",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.cache_disk_path = value;
|
||||
if (!fs_is_directory(params.cache_disk_path)) {
|
||||
throw std::invalid_argument("not a directory: " + value);
|
||||
}
|
||||
// if doesn't end with DIRECTORY_SEPARATOR, add it
|
||||
if (params.cache_disk_path[params.cache_disk_path.size() - 1] != DIRECTORY_SEPARATOR) {
|
||||
params.cache_disk_path += DIRECTORY_SEPARATOR;
|
||||
}
|
||||
}
|
||||
).set_env("LLAMA_ARG_CACHE_DISK").set_examples({LLAMA_EXAMPLE_SERVER}));
|
||||
add_opt(common_arg(
|
||||
{"--cache-disk-limit"}, "N",
|
||||
string_format("total size budget of the disk prompt cache directory in MiB; oldest entries are deleted when exceeded (default: %d, -1 - no limit)", params.cache_disk_limit_mib),
|
||||
[](common_params & params, int value) {
|
||||
if (value == 0 || value < -1) {
|
||||
throw std::invalid_argument("cache-disk-limit must be positive or -1 (no limit)");
|
||||
}
|
||||
params.cache_disk_limit_mib = value;
|
||||
}
|
||||
).set_env("LLAMA_ARG_CACHE_DISK_LIMIT").set_examples({LLAMA_EXAMPLE_SERVER}));
|
||||
add_opt(common_arg(
|
||||
{"--cache-disk-write-through"},
|
||||
{"--no-cache-disk-write-through"},
|
||||
"write prompts to the disk cache every time they are saved to the RAM cache, instead of only when evicted from it (default: disabled)",
|
||||
[](common_params & params, bool value) {
|
||||
params.cache_disk_write_through = value;
|
||||
}
|
||||
).set_env("LLAMA_ARG_CACHE_DISK_WRITE_THROUGH").set_examples({LLAMA_EXAMPLE_SERVER}));
|
||||
add_opt(common_arg(
|
||||
{"-kvu", "--kv-unified"},
|
||||
{"-no-kvu", "--no-kv-unified"},
|
||||
|
||||
@@ -614,6 +614,10 @@ struct common_params {
|
||||
int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints
|
||||
int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
|
||||
|
||||
std::string cache_disk_path; // disk prompt cache directory, empty = disabled
|
||||
int32_t cache_disk_limit_mib = -1; // total size budget for the disk prompt cache dir, -1 = no limit
|
||||
bool cache_disk_write_through = false; // also write to disk whenever a prompt is saved to the RAM cache
|
||||
|
||||
std::string hostname = "127.0.0.1";
|
||||
std::string public_path = ""; // NOLINT
|
||||
std::string api_prefix = ""; // NOLINT
|
||||
|
||||
@@ -18,13 +18,16 @@ extern "C" {
|
||||
#endif
|
||||
|
||||
#include "xxhash/xxhash.h"
|
||||
#include "sha1/sha1.h"
|
||||
#include "sha256/sha256.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
// sha1 is compiled as C++ and lives in a namespace, see scripts/sync_vendor.py
|
||||
#include "sha1/sha1.h"
|
||||
using namespace vendor_hash;
|
||||
|
||||
|
||||
// uuid.uuid5(uuid.NAMESPACE_URL, 'en.wikipedia.org/wiki/Llama.cpp')
|
||||
#define UUID_NAMESPACE_LLAMA_CPP "ef001206-dadc-5f6d-a15f-3359e577d4e5"
|
||||
|
||||
@@ -56,6 +56,44 @@ patches = {
|
||||
' && (defined(_MSC_VER) && (_MSC_VER >= 1000) || !defined(_MSC_VER)) /* >= C11 */\n'
|
||||
)],
|
||||
|
||||
# sha1 exports a bare "SHA1" symbol, which clashes with the boringssl one at link time.
|
||||
# we compile it as C++ (see vendor/hash/CMakeLists.txt) and put it in a namespace.
|
||||
"vendor/hash/sha1/sha1.h": [
|
||||
(
|
||||
'#if defined(__cplusplus)\n'
|
||||
'extern "C" {\n'
|
||||
'#endif\n',
|
||||
|
||||
'namespace vendor_hash {\n'
|
||||
),
|
||||
(
|
||||
'#if defined(__cplusplus)\n'
|
||||
'}\n'
|
||||
'#endif\n',
|
||||
|
||||
'} // namespace vendor_hash\n'
|
||||
),
|
||||
],
|
||||
|
||||
"vendor/hash/sha1/sha1.c": [
|
||||
(
|
||||
'#include "sha1.h"\n',
|
||||
|
||||
'#include "sha1.h"\n'
|
||||
'\n'
|
||||
'namespace vendor_hash {\n'
|
||||
),
|
||||
(
|
||||
' SHA1Final((unsigned char *)hash_out, &ctx);\n'
|
||||
'}\n',
|
||||
|
||||
' SHA1Final((unsigned char *)hash_out, &ctx);\n'
|
||||
'}\n'
|
||||
'\n'
|
||||
'} // namespace vendor_hash\n'
|
||||
),
|
||||
],
|
||||
|
||||
# silence a maybe-uninitialized warning
|
||||
"vendor/hash/sha256/sha256.c": [(
|
||||
" uint32_t W[16];\n",
|
||||
|
||||
+17
-4
@@ -2428,17 +2428,24 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx);
|
||||
|
||||
const float * scores = nullptr;
|
||||
const int * iscores = nullptr;
|
||||
const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str());
|
||||
if (score_idx != -1) {
|
||||
if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32) {
|
||||
const gguf_type kv_type = gguf_get_kv_type(ctx, score_idx);
|
||||
const gguf_type arr_type = kv_type == GGUF_TYPE_ARRAY ? gguf_get_arr_type(ctx, score_idx) : GGUF_TYPE_COUNT;
|
||||
if (arr_type != GGUF_TYPE_INT32 &&
|
||||
arr_type != GGUF_TYPE_FLOAT32) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str()));
|
||||
}
|
||||
const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx);
|
||||
if (n_scores < n_tokens) {
|
||||
throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n");
|
||||
}
|
||||
scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
|
||||
if (arr_type == GGUF_TYPE_INT32) {
|
||||
iscores = (const int *) gguf_get_arr_data(ctx, score_idx);
|
||||
} else {
|
||||
scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
|
||||
}
|
||||
}
|
||||
|
||||
const int * toktypes = nullptr;
|
||||
@@ -2469,7 +2476,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
|
||||
auto & token_data = id_to_token[i];
|
||||
token_data.text = std::move(word);
|
||||
token_data.score = scores ? scores[i] : 0.0f;
|
||||
if (scores) {
|
||||
token_data.score = scores[i];
|
||||
} else if (iscores) {
|
||||
token_data.score = static_cast<float>(iscores[i]);
|
||||
} else {
|
||||
token_data.score = 0.0f;
|
||||
}
|
||||
token_data.attr = LLAMA_TOKEN_ATTR_NORMAL;
|
||||
|
||||
if (toktypes) { //TODO: remove, required until per token attributes are available from GGUF file
|
||||
|
||||
@@ -78,7 +78,7 @@ set_target_properties(mtmd PROPERTIES
|
||||
)
|
||||
|
||||
target_link_libraries (mtmd PUBLIC ggml llama)
|
||||
target_link_libraries (mtmd PRIVATE Threads::Threads)
|
||||
target_link_libraries (mtmd PRIVATE Threads::Threads vendor-hash)
|
||||
target_include_directories(mtmd PUBLIC .)
|
||||
target_include_directories(mtmd PRIVATE ../..)
|
||||
target_include_directories(mtmd PRIVATE ../../vendor)
|
||||
|
||||
@@ -12,6 +12,8 @@
|
||||
#include "mtmd-helper-common.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include "hash.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <vector>
|
||||
@@ -356,25 +358,14 @@ static bool decode_audio_from_buf(const unsigned char * buf_in, size_t len, int
|
||||
|
||||
} // namespace audio_helpers
|
||||
|
||||
// Computes FNV-1a hash of the data
|
||||
static std::string fnv_hash(const uint8_t * data, size_t len) {
|
||||
const uint64_t fnv_prime = 0x100000001b3ULL;
|
||||
uint64_t hash = 0xcbf29ce484222325ULL;
|
||||
|
||||
for (size_t i = 0; i < len; ++i) {
|
||||
hash ^= data[i];
|
||||
hash *= fnv_prime;
|
||||
}
|
||||
return std::to_string(hash);
|
||||
}
|
||||
|
||||
mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder) {
|
||||
// calculate the hash if needed
|
||||
std::string id;
|
||||
mtmd_bitmap * result = nullptr;
|
||||
|
||||
if (!placeholder) {
|
||||
id = fnv_hash(buf, len);
|
||||
// use sha256 to prevent cache poisoning
|
||||
id = hash_sha256_hex(buf, len);
|
||||
}
|
||||
|
||||
if (audio_helpers::is_audio_file((const char *)buf, len)) {
|
||||
|
||||
@@ -49,7 +49,7 @@ MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtm
|
||||
// note:
|
||||
// - for now, video input is only supported via C++ helper functions
|
||||
// - audio files will be auto-detected based on magic bytes
|
||||
// - output bitmap will have FNV hash as the ID
|
||||
// - output bitmap will have SHA-256 hash (hex string) as the ID
|
||||
// returns nullptr on failure
|
||||
// this function is thread-safe
|
||||
MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder);
|
||||
|
||||
@@ -999,6 +999,20 @@ mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_mi
|
||||
// mtmd_image_preprocessor_lfm2
|
||||
//
|
||||
|
||||
mtmd_image_preproc_out mtmd_image_preprocessor_lfm2::preprocess(const clip_image_u8 & img) {
|
||||
auto const inst = get_slice_instructions(img.get_size());
|
||||
if (!inst.slices.empty()) {
|
||||
return mtmd_image_preprocessor_llava_uhd::preprocess(img);
|
||||
}
|
||||
|
||||
// single tile: no thumbnail
|
||||
// note: not using output.overview here because it will emit <|img_thumbnail|> token, which we don't want in this case
|
||||
auto sliced = slice_image(img, inst);
|
||||
mtmd_image_preproc_out output;
|
||||
output.append(hparams, sliced.overview, true);
|
||||
return output;
|
||||
}
|
||||
|
||||
mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lfm2::get_slice_instructions(const clip_image_size & original_size) {
|
||||
mtmd_image_preprocessor_llava_uhd::slice_instructions inst;
|
||||
const int align_size = hparams.patch_size * hparams.n_merge;
|
||||
|
||||
@@ -145,6 +145,7 @@ struct mtmd_image_preprocessor_lfm2 : mtmd_image_preprocessor_llava_uhd {
|
||||
static constexpr int tile_size = 512;
|
||||
|
||||
using mtmd_image_preprocessor_llava_uhd::mtmd_image_preprocessor_llava_uhd;
|
||||
mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override;
|
||||
slice_instructions get_slice_instructions(const clip_image_size & original_size) override;
|
||||
|
||||
private:
|
||||
|
||||
+32
-14
@@ -2322,23 +2322,12 @@ void mtmd_input_chunk_free(mtmd_input_chunk * chunk) {
|
||||
}
|
||||
}
|
||||
|
||||
int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len) {
|
||||
// returns 0 on success
|
||||
static int32_t mtmd_input_chunk_save_impl(const mtmd_input_chunk * chunk, std::vector<char> & out_buf) {
|
||||
try {
|
||||
mtmd_serialization ser(MTMD_SERIALIZATION_VERSION);
|
||||
chunk->serialize(ser);
|
||||
|
||||
if (expected_out_len) {
|
||||
*expected_out_len = ser.data.size();
|
||||
}
|
||||
if (!out_buf) {
|
||||
// caller is only querying the required size
|
||||
return 0;
|
||||
}
|
||||
if (out_len < ser.data.size()) {
|
||||
LOG_ERR("%s: out_buf is too small, need %zu bytes, got %zu\n", __func__, ser.data.size(), out_len);
|
||||
return -1;
|
||||
}
|
||||
std::memcpy(out_buf, ser.data.data(), ser.data.size());
|
||||
out_buf = std::move(ser.data);
|
||||
return 0;
|
||||
} catch (const std::exception & e) {
|
||||
LOG_ERR("%s: %s\n", __func__, e.what());
|
||||
@@ -2346,6 +2335,35 @@ int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, si
|
||||
}
|
||||
}
|
||||
|
||||
mtmd_input_chunk * mtmd_input_chunk_get_placeholder(const mtmd_input_chunk * chunk) {
|
||||
// this is hacky, but still faster than copy the whole batch data
|
||||
std::vector<char> buf;
|
||||
if (mtmd_input_chunk_save_impl(chunk, buf) != 0) {
|
||||
return nullptr;
|
||||
}
|
||||
return mtmd_input_chunk_load(buf.data(), buf.size());
|
||||
}
|
||||
|
||||
int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len) {
|
||||
std::vector<char> buf;
|
||||
if (mtmd_input_chunk_save_impl(chunk, buf) != 0) {
|
||||
return -1;
|
||||
}
|
||||
if (expected_out_len) {
|
||||
*expected_out_len = buf.size();
|
||||
}
|
||||
if (!out_buf) {
|
||||
// caller is only querying the required size
|
||||
return 0;
|
||||
}
|
||||
if (out_len < buf.size()) {
|
||||
LOG_ERR("%s: out_buf is too small, need %zu bytes, got %zu\n", __func__, buf.size(), out_len);
|
||||
return -1;
|
||||
}
|
||||
std::memcpy(out_buf, buf.data(), buf.size());
|
||||
return 0;
|
||||
}
|
||||
|
||||
mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len) {
|
||||
try {
|
||||
mtmd_serialization ser(MTMD_SERIALIZATION_VERSION, buf, len);
|
||||
|
||||
@@ -233,6 +233,9 @@ MTMD_API llama_pos mtmd_input_chunk_get_n_pos (const mtmd
|
||||
MTMD_API mtmd_input_chunk * mtmd_input_chunk_copy(const mtmd_input_chunk * chunk);
|
||||
MTMD_API void mtmd_input_chunk_free(mtmd_input_chunk * chunk);
|
||||
|
||||
// similar to mtmd_input_chunk_copy, but returns a placeholder chunk
|
||||
MTMD_API mtmd_input_chunk * mtmd_input_chunk_get_placeholder(const mtmd_input_chunk * chunk);
|
||||
|
||||
// save/load an input chunk to/from a buffer (useful for KV save/load)
|
||||
// important: only chunk's metadata will be saved, the actual image/audio data will not be saved
|
||||
// the loaded chunk will always be a placeholder, cannot be used for mtmd_encode() or mtmd_batch_encode()
|
||||
|
||||
@@ -5,6 +5,8 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_BINARY_DIR})
|
||||
set(TARGET server-context)
|
||||
|
||||
add_library(${TARGET} STATIC
|
||||
server-cache-disk.cpp
|
||||
server-cache-disk.h
|
||||
server-chat.cpp
|
||||
server-chat.h
|
||||
server-task.cpp
|
||||
@@ -31,7 +33,7 @@ endif()
|
||||
|
||||
target_include_directories(${TARGET} PRIVATE ../mtmd)
|
||||
target_include_directories(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR})
|
||||
target_link_libraries(${TARGET} PUBLIC llama-common mtmd ${CMAKE_THREAD_LIBS_INIT})
|
||||
target_link_libraries(${TARGET} PUBLIC llama-common mtmd vendor-hash ${CMAKE_THREAD_LIBS_INIT})
|
||||
|
||||
# llama-server-impl: server logic, reusable by app
|
||||
|
||||
|
||||
@@ -164,6 +164,9 @@ For the full list of features, please refer to [server's changelog](https://gith
|
||||
| `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)<br/>(env: LLAMA_ARG_CTX_CHECKPOINTS) |
|
||||
| `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)<br/>(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) |
|
||||
| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)<br/>(env: LLAMA_ARG_CACHE_RAM) |
|
||||
| `-cdisk, --cache-disk PATH` | directory for the disk prompt cache; prompts evicted from the RAM cache are saved here and restored on later requests, including across restarts (default: disabled, requires cache-ram)<br/>(env: LLAMA_ARG_CACHE_DISK) |
|
||||
| `--cache-disk-limit N` | total size budget of the disk prompt cache directory in MiB; oldest entries are deleted when exceeded (default: -1, -1 - no limit)<br/>(env: LLAMA_ARG_CACHE_DISK_LIMIT) |
|
||||
| `--cache-disk-write-through, --no-cache-disk-write-through` | write prompts to the disk cache every time they are saved to the RAM cache, instead of only when evicted from it (default: disabled)<br/>(env: LLAMA_ARG_CACHE_DISK_WRITE_THROUGH) |
|
||||
| `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)<br/>(env: LLAMA_ARG_KV_UNIFIED) |
|
||||
| `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)<br/>(env: LLAMA_ARG_CACHE_IDLE_SLOTS) |
|
||||
| `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)<br/>(env: LLAMA_ARG_CONTEXT_SHIFT) |
|
||||
@@ -327,6 +330,22 @@ services:
|
||||
LLAMA_ARG_PORT: 8080
|
||||
```
|
||||
|
||||
### Prompt disk cache
|
||||
|
||||
The server keeps recently used prompts (their processed KV cache state) in RAM, controlled by `--cache-ram`. With `--cache-disk PATH`, a disk tier is added below the RAM cache: entries evicted from RAM are written to the given directory, and all RAM entries are flushed there on graceful shutdown. On later requests - including after a server restart - the longest cached prefix of the incoming prompt is restored from disk instead of being re-processed.
|
||||
|
||||
```sh
|
||||
llama-server -m model.gguf --cache-disk /path/to/cache --cache-disk-limit 32768
|
||||
```
|
||||
|
||||
Details:
|
||||
|
||||
- Files are named `{compat_hash}-{n_tokens}-{chain_hash}.kvc`, where the hashes identify the server configuration and the exact token prefix the file contains. Lookup is a single directory scan at startup plus one hash pass per prompt - no database is used.
|
||||
- The cache is invalidated automatically when the model file, mmproj, LoRA adapters, KV cache types, or rope parameters change (stale files are ignored, and deleted once the size budget is exceeded).
|
||||
- `--cache-disk-limit` bounds the total size of the directory in MiB; the oldest files (by modification time) are deleted first, including files left over from other models or configurations. The same directory can be shared by multiple servers.
|
||||
- By default, files are only written when an entry is evicted from the RAM cache (or on shutdown). With `--cache-disk-write-through`, every prompt saved to the RAM cache is also written to disk immediately, which is more crash-resilient at the cost of extra I/O.
|
||||
- Note that KV cache states can be large (potentially multiple GiB per prompt, depending on the model and prompt length), so make sure the disk budget is sized accordingly.
|
||||
|
||||
### Multimodal support
|
||||
|
||||
Multimodal support was added in [#12898](https://github.com/ggml-org/llama.cpp/pull/12898) and is currently an experimental feature.
|
||||
@@ -343,6 +362,58 @@ The server includes a set of built-in tools that enable the LLM to access the lo
|
||||
|
||||
To use this feature, start the server with `--tools all`. You can also enable only specific tools by passing a comma-separated list: `--tools name1,name2,...`. Run `--help` for the full list of available tool names.
|
||||
|
||||
### MCP servers
|
||||
|
||||
Besides the built-in tools, the server can expose tools coming from MCP servers, added in [#26062](https://github.com/ggml-org/llama.cpp/pull/26062). Only the stdio transport is supported: such a server is a child process reading JSON-RPC messages on its stdin and writing replies on its stdout, so nothing has to be started or maintained outside `llama-server`.
|
||||
|
||||
Servers are declared in a Cursor-compatible JSON file:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"example": { "command": "/path/to/server", "args": [] }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
```sh
|
||||
llama-server -m model.gguf --mcp-servers-config mcp.json
|
||||
```
|
||||
|
||||
The same JSON can be passed inline with `--mcp-servers-json`. Each entry under `mcpServers` accepts:
|
||||
|
||||
| Key | Explanation |
|
||||
| --- | ----------- |
|
||||
| `command` | executable to spawn, required, entries without it are skipped |
|
||||
| `args` | array of arguments |
|
||||
| `env` | object merged over the parent environment |
|
||||
| `cwd` | working directory of the child process |
|
||||
| `timeout_ms` | per-tool-call timeout (default: 30000) |
|
||||
|
||||
Every server is spawned once at startup to list its tools, then stopped, and respawned on demand when one of its tools is called. Tools are exposed as `<server>_<tool>` alongside the built-in ones: they show up in the Web UI and in `GET /tools`, and the model calls them like any other tool. A name colliding with an already registered tool is skipped. This is independent of `--tools`, MCP servers can be the only tools available.
|
||||
|
||||
The child process runs with the same privileges as the server, so only declare commands you trust. As with `--tools`, `--cors-origins` then defaults to `localhost`.
|
||||
|
||||
Note: `--ui-mcp-proxy` is unrelated, it only lets the Web UI reach remote MCP servers from the browser.
|
||||
|
||||
Any server written against the [MCP specification](https://modelcontextprotocol.io) works as is, whether it uses an official SDK or not: the transport is one JSON-RPC message per line on stdio, so a script wrapping an existing program is a valid server too.
|
||||
|
||||
### CORS
|
||||
|
||||
By default the server reflects any `Origin` header back with credentials allowed. This matches the old, always-on `*` behavior and is fine as long as the server only exposes stateless, read-only endpoints.
|
||||
|
||||
Enabling `--tools` or `--agent` exposes file read/write over the API, so in that case `--cors-origins` defaults to `localhost` instead: only pages served from localhost can reach the server. Pass `--cors-origins` explicitly to override either default.
|
||||
|
||||
Recommended `--cors-origins` setting, depending on where the server runs:
|
||||
|
||||
| Deployment | Recommendation |
|
||||
| ---------- | --------------- |
|
||||
| Public | set an API key, put the server behind a reverse proxy, `--cors-origins` optional |
|
||||
| Local network | set `--cors-origins` to your frontend's origin |
|
||||
| Same machine | `--cors-origins localhost` (default once `--agent` is set) |
|
||||
|
||||
Related flags: `--cors-origins`, `--cors-methods`, `--cors-headers`, `--cors-credentials` / `--no-cors-credentials`. Background and rationale: [#25655](https://github.com/ggml-org/llama.cpp/pull/25655).
|
||||
|
||||
## Build
|
||||
|
||||
`llama-server` is built alongside everything else from the root of the project
|
||||
|
||||
@@ -0,0 +1,580 @@
|
||||
#include "server-cache-disk.h"
|
||||
|
||||
#include "common.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include "xxhash/xxhash.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <functional>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr uint32_t SERVER_CACHE_DISK_MAGIC = 0x3143564B; // "KVC1"
|
||||
constexpr uint32_t SERVER_CACHE_DISK_VERSION = 1;
|
||||
|
||||
// seed for the chained prefix hash - changing it invalidates all filenames
|
||||
constexpr uint64_t SERVER_CACHE_DISK_CHAIN_SEED = 0x6b7663636861696eULL;
|
||||
|
||||
struct server_cache_disk_file_header {
|
||||
uint32_t magic = SERVER_CACHE_DISK_MAGIC;
|
||||
uint32_t version = SERVER_CACHE_DISK_VERSION;
|
||||
uint64_t compat_hash = 0; // full 64-bit value (the filename only carries the low 32 bits)
|
||||
uint64_t chain_hash = 0;
|
||||
uint32_t n_tokens = 0;
|
||||
uint32_t pad = 0;
|
||||
uint64_t tokens_size = 0; // bytes of the server_tokens::serialize() section
|
||||
uint64_t state_size = 0; // bytes of the llama_state_seq_get_data section
|
||||
};
|
||||
|
||||
static_assert(sizeof(server_cache_disk_file_header) == 48, "unexpected header size");
|
||||
|
||||
std::string make_filename(uint64_t compat_hash, uint32_t n_tokens, uint64_t chain_hash) {
|
||||
char buf[64];
|
||||
snprintf(buf, sizeof(buf), "%08x-%u-%016" PRIx64 ".kvc", (uint32_t) compat_hash, n_tokens, chain_hash);
|
||||
return buf;
|
||||
}
|
||||
|
||||
bool parse_filename(const std::string & name, uint32_t & compat32, uint32_t & n_tokens, uint64_t & chain_hash) {
|
||||
if (sscanf(name.c_str(), "%8x-%u-%16" SCNx64 ".kvc", &compat32, &n_tokens, &chain_hash) != 3) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// reject padding/case/suffix variations by requiring the canonical spelling
|
||||
return name == make_filename(compat32, n_tokens, chain_hash);
|
||||
}
|
||||
|
||||
int64_t file_mtime(const std::filesystem::path & path) {
|
||||
std::error_code ec;
|
||||
const auto t = std::filesystem::last_write_time(path, ec);
|
||||
return ec ? 0 : (int64_t) t.time_since_epoch().count();
|
||||
}
|
||||
|
||||
uint64_t covered_key(uint32_t n_tokens, uint64_t chain_hash) {
|
||||
const uint64_t buf[2] = { n_tokens, chain_hash };
|
||||
return XXH64(buf, sizeof(buf), 0);
|
||||
}
|
||||
|
||||
// walk the chained hash over the token list, invoking cb(n, h) at every valid prefix boundary:
|
||||
// after each text token and after each complete media chunk (never mid-chunk)
|
||||
// returns true if the walk reached n_max
|
||||
bool tokens_chain_hash_walk(const server_tokens & tokens, size_t n_max, const std::function<bool(size_t, uint64_t)> & cb) {
|
||||
uint64_t h = SERVER_CACHE_DISK_CHAIN_SEED;
|
||||
|
||||
size_t i = 0;
|
||||
|
||||
try {
|
||||
while (i < n_max) {
|
||||
const llama_token tok = tokens[i];
|
||||
|
||||
if (tok == LLAMA_TOKEN_NULL) {
|
||||
// media chunk - fold in its content id instead of the placeholder token ids,
|
||||
// otherwise different images would hash identically
|
||||
const auto & chunk = tokens.find_chunk(i);
|
||||
|
||||
const char * id = mtmd_input_chunk_get_id(chunk.get());
|
||||
const size_t n_tok = mtmd_input_chunk_get_n_tokens(chunk.get());
|
||||
|
||||
if (id == nullptr || id[0] == '\0' || n_tok == 0 || i + n_tok > n_max) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> buf;
|
||||
buf.reserve(5 + strlen(id));
|
||||
buf.push_back(0x01);
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
buf.push_back((uint8_t) (n_tok >> (8*b)));
|
||||
}
|
||||
buf.insert(buf.end(), id, id + strlen(id));
|
||||
|
||||
h = XXH64(buf.data(), buf.size(), h);
|
||||
|
||||
i += n_tok;
|
||||
} else {
|
||||
uint8_t buf[5] = { 0x00 };
|
||||
memcpy(buf + 1, &tok, sizeof(tok));
|
||||
|
||||
h = XXH64(buf, sizeof(buf), h);
|
||||
|
||||
i += 1;
|
||||
}
|
||||
|
||||
if (!cb(i, h)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
} catch (const std::exception & e) {
|
||||
SRV_WRN("failed to hash token list: %s\n", e.what());
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
server_prompt_cache_disk::server_prompt_cache_disk(const std::string & dir_, uint64_t compat_hash, bool has_mtmd, int32_t limit_mib, bool write_through) :
|
||||
write_through(write_through),
|
||||
dir(dir_.empty() || dir_.back() == DIRECTORY_SEPARATOR ? dir_ : dir_ + DIRECTORY_SEPARATOR),
|
||||
compat_hash(compat_hash),
|
||||
has_mtmd(has_mtmd),
|
||||
limit_bytes(limit_mib < 0 ? 0 : 1024ull*1024ull*limit_mib) {
|
||||
scan_dir();
|
||||
}
|
||||
|
||||
void server_prompt_cache_disk::scan_dir() {
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
std::error_code ec;
|
||||
|
||||
for (const auto & ent : fs::directory_iterator(dir, ec)) {
|
||||
if (!ent.is_regular_file(ec)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const std::string name = ent.path().filename().string();
|
||||
|
||||
// leftover temporary files from a previous crash
|
||||
if (name.size() > 4 && name.compare(name.size() - 4, 4, ".tmp") == 0 && name[0] == '.') {
|
||||
fs::remove(ent.path(), ec);
|
||||
continue;
|
||||
}
|
||||
|
||||
uint32_t compat32 = 0;
|
||||
uint32_t n_tokens = 0;
|
||||
uint64_t chain = 0;
|
||||
|
||||
if (!parse_filename(name, compat32, n_tokens, chain)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
server_cache_disk_file file;
|
||||
file.name = name;
|
||||
file.chain_hash = chain;
|
||||
file.n_tokens = n_tokens;
|
||||
file.n_bytes = ent.file_size(ec);
|
||||
file.mtime = file_mtime(ent.path());
|
||||
|
||||
total_bytes += file.n_bytes;
|
||||
|
||||
if (compat32 == (uint32_t) compat_hash) {
|
||||
index[n_tokens][chain] = std::move(file);
|
||||
} else {
|
||||
foreign.push_back(std::move(file));
|
||||
}
|
||||
}
|
||||
|
||||
SRV_INF("disk prompt cache '%s': %zu usable entries, %zu from other configurations, %.3f MiB total (budget: %.3f MiB)\n",
|
||||
dir.c_str(), n_files(), foreign.size(), total_bytes / (1024.0 * 1024.0), limit_bytes / (1024.0 * 1024.0));
|
||||
}
|
||||
|
||||
size_t server_prompt_cache_disk::n_files() const {
|
||||
size_t res = 0;
|
||||
|
||||
for (const auto & [n, files] : index) {
|
||||
res += files.size();
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
server_cache_disk_file * server_prompt_cache_disk::find_file(uint32_t n_tokens, uint64_t chain_hash) {
|
||||
const auto it = index.find(n_tokens);
|
||||
if (it == index.end()) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
const auto it_file = it->second.find(chain_hash);
|
||||
|
||||
return it_file == it->second.end() ? nullptr : &it_file->second;
|
||||
}
|
||||
|
||||
const server_cache_disk_file * server_prompt_cache_disk::lookup(const server_tokens & tokens, size_t n_max) const {
|
||||
if (index.empty()) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// no file can be longer than the largest indexed length - cap the walk
|
||||
n_max = std::min<size_t>(n_max, index.rbegin()->first);
|
||||
|
||||
const server_cache_disk_file * best = nullptr;
|
||||
|
||||
tokens_chain_hash_walk(tokens, n_max, [&](size_t n, uint64_t h) {
|
||||
const auto it = index.find((uint32_t) n);
|
||||
if (it != index.end()) {
|
||||
const auto it_file = it->second.find(h);
|
||||
if (it_file != it->second.end()) {
|
||||
best = &it_file->second;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
});
|
||||
|
||||
return best;
|
||||
}
|
||||
|
||||
void server_prompt_cache_disk::touch(const server_cache_disk_file & file) {
|
||||
std::error_code ec;
|
||||
std::filesystem::last_write_time(dir + file.name, std::filesystem::file_time_type::clock::now(), ec);
|
||||
|
||||
if (auto * f = find_file(file.n_tokens, file.chain_hash)) {
|
||||
f->mtime = file_mtime(dir + file.name);
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache_disk::forget(const server_cache_disk_file & file) {
|
||||
// copy the fields first - the reference may point into the index entry being erased
|
||||
const uint32_t n_tokens = file.n_tokens;
|
||||
const uint64_t chain = file.chain_hash;
|
||||
const uint64_t n_bytes = file.n_bytes;
|
||||
|
||||
const auto it = index.find(n_tokens);
|
||||
if (it == index.end()) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (it->second.erase(chain) > 0) {
|
||||
total_bytes -= std::min<size_t>(total_bytes, n_bytes);
|
||||
}
|
||||
|
||||
if (it->second.empty()) {
|
||||
index.erase(it);
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache_disk::remove_file(const server_cache_disk_file & file) {
|
||||
SRV_WRN("disk prompt cache: removing '%s'\n", file.name.c_str());
|
||||
|
||||
std::error_code ec;
|
||||
std::filesystem::remove(dir + file.name, ec);
|
||||
|
||||
forget(file);
|
||||
}
|
||||
|
||||
void server_prompt_cache_disk::enforce_budget(const std::string & name_protected) {
|
||||
if (limit_bytes == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
while (total_bytes > limit_bytes) {
|
||||
// find the oldest file, ours and foreign alike
|
||||
const server_cache_disk_file * oldest = nullptr;
|
||||
bool oldest_foreign = false;
|
||||
|
||||
for (const auto & [n, files] : index) {
|
||||
for (const auto & [h, file] : files) {
|
||||
if (file.name != name_protected && (!oldest || file.mtime < oldest->mtime)) {
|
||||
oldest = &file;
|
||||
oldest_foreign = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const auto & file : foreign) {
|
||||
if (file.name != name_protected && (!oldest || file.mtime < oldest->mtime)) {
|
||||
oldest = &file;
|
||||
oldest_foreign = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!oldest) {
|
||||
break;
|
||||
}
|
||||
|
||||
SRV_INF("disk prompt cache: size %.3f MiB over budget %.3f MiB, evicting oldest entry '%s'\n",
|
||||
total_bytes / (1024.0 * 1024.0), limit_bytes / (1024.0 * 1024.0), oldest->name.c_str());
|
||||
|
||||
if (oldest_foreign) {
|
||||
std::error_code ec;
|
||||
std::filesystem::remove(dir + oldest->name, ec);
|
||||
|
||||
total_bytes -= std::min<size_t>(total_bytes, oldest->n_bytes);
|
||||
|
||||
foreign.erase(foreign.begin() + (oldest - foreign.data()));
|
||||
} else {
|
||||
remove_file(*oldest);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
bool server_prompt_cache_disk::store(const server_tokens & tokens, const std::vector<uint8_t> & state_main) {
|
||||
if (tokens.empty() || state_main.empty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<std::pair<size_t, uint64_t>> bounds;
|
||||
|
||||
if (!tokens_chain_hash_walk(tokens, tokens.size(), [&](size_t n, uint64_t h) { bounds.emplace_back(n, h); return true; }) ||
|
||||
bounds.empty() || bounds.back().first != tokens.size()) {
|
||||
SRV_WRN("%s", "disk prompt cache: token list cannot be hashed, skipping\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
const uint32_t n_tokens = (uint32_t) tokens.size();
|
||||
const uint64_t chain = bounds.back().second;
|
||||
|
||||
if (auto * existing = find_file(n_tokens, chain)) {
|
||||
SRV_TRC("disk prompt cache: '%s' already exists, refreshing\n", existing->name.c_str());
|
||||
touch(*existing);
|
||||
return true;
|
||||
}
|
||||
|
||||
if (covered.count(covered_key(n_tokens, chain)) > 0) {
|
||||
SRV_TRC(" - prompt with %u tokens is a prefix of an already persisted entry, skipping\n", n_tokens);
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<char> tok_data;
|
||||
try {
|
||||
tok_data = tokens.serialize();
|
||||
} catch (const std::exception & e) {
|
||||
SRV_WRN("disk prompt cache: failed to serialize tokens: %s\n", e.what());
|
||||
return false;
|
||||
}
|
||||
|
||||
server_cache_disk_file_header header;
|
||||
header.compat_hash = compat_hash;
|
||||
header.chain_hash = chain;
|
||||
header.n_tokens = n_tokens;
|
||||
header.tokens_size = tok_data.size();
|
||||
header.state_size = state_main.size();
|
||||
|
||||
const std::string name = make_filename(compat_hash, n_tokens, chain);
|
||||
|
||||
char tmp_buf[64];
|
||||
snprintf(tmp_buf, sizeof(tmp_buf), ".%08x-%u.tmp", (uint32_t) (uintptr_t) this, tmp_counter++);
|
||||
|
||||
const std::string path_tmp = dir + tmp_buf;
|
||||
const std::string path = dir + name;
|
||||
|
||||
{
|
||||
std::ofstream out(path_tmp, std::ios::binary | std::ios::trunc);
|
||||
|
||||
out.write((const char *) &header, sizeof(header));
|
||||
out.write(tok_data.data(), tok_data.size());
|
||||
out.write((const char *) state_main.data(), state_main.size());
|
||||
|
||||
if (!out.good()) {
|
||||
SRV_ERR("disk prompt cache: failed to write '%s'\n", path_tmp.c_str());
|
||||
|
||||
out.close();
|
||||
|
||||
std::error_code ec;
|
||||
std::filesystem::remove(path_tmp, ec);
|
||||
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
std::error_code ec;
|
||||
std::filesystem::rename(path_tmp, path, ec);
|
||||
if (ec) {
|
||||
SRV_ERR("disk prompt cache: failed to rename '%s' to '%s': %s\n", path_tmp.c_str(), path.c_str(), ec.message().c_str());
|
||||
|
||||
std::filesystem::remove(path_tmp, ec);
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
server_cache_disk_file file;
|
||||
file.name = name;
|
||||
file.chain_hash = chain;
|
||||
file.n_tokens = n_tokens;
|
||||
file.n_bytes = sizeof(header) + tok_data.size() + state_main.size();
|
||||
file.mtime = file_mtime(path);
|
||||
|
||||
total_bytes += file.n_bytes;
|
||||
|
||||
index[n_tokens][chain] = std::move(file);
|
||||
|
||||
for (const auto & [n, h] : bounds) {
|
||||
covered.insert(covered_key((uint32_t) n, h));
|
||||
}
|
||||
|
||||
SRV_INF("disk prompt cache: saved prompt with %u tokens, %.3f MiB to '%s'\n",
|
||||
n_tokens, (sizeof(header) + tok_data.size() + state_main.size()) / (1024.0 * 1024.0), name.c_str());
|
||||
SRV_DBG("%s", "__TEST_TAG_CACHE_DISK_STORE__\n");
|
||||
|
||||
enforce_budget(name);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
server_prompt_cache_disk::load_status server_prompt_cache_disk::load(
|
||||
server_cache_disk_file file, const server_tokens & tokens_new, llama_context * ctx, int32_t id_slot, server_tokens & tokens_out) {
|
||||
const std::string path = dir + file.name;
|
||||
|
||||
std::error_code ec;
|
||||
const uint64_t n_bytes = std::filesystem::file_size(path, ec);
|
||||
|
||||
if (ec) {
|
||||
// deleted by another process - not an error, just a miss
|
||||
forget(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
std::ifstream in(path, std::ios::binary);
|
||||
if (!in.good()) {
|
||||
forget(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
server_cache_disk_file_header header;
|
||||
in.read((char *) &header, sizeof(header));
|
||||
|
||||
if (!in.good() ||
|
||||
header.magic != SERVER_CACHE_DISK_MAGIC ||
|
||||
header.version != SERVER_CACHE_DISK_VERSION ||
|
||||
header.chain_hash != file.chain_hash ||
|
||||
header.n_tokens != file.n_tokens ||
|
||||
header.tokens_size % sizeof(llama_token) != 0 ||
|
||||
sizeof(header) + header.tokens_size + header.state_size != n_bytes) {
|
||||
SRV_WRN("disk prompt cache: '%s' is corrupt\n", file.name.c_str());
|
||||
remove_file(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
if (header.compat_hash != compat_hash) {
|
||||
// same low 32 bits, different configuration - leave the file for its owner
|
||||
SRV_WRN("disk prompt cache: '%s' belongs to a different configuration, ignoring\n", file.name.c_str());
|
||||
forget(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
llama_tokens packed(header.tokens_size / sizeof(llama_token));
|
||||
in.read((char *) packed.data(), header.tokens_size);
|
||||
|
||||
if (!in.good()) {
|
||||
SRV_WRN("disk prompt cache: '%s' is truncated\n", file.name.c_str());
|
||||
remove_file(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
server_tokens loaded;
|
||||
try {
|
||||
loaded = server_tokens::deserialize(packed, has_mtmd);
|
||||
} catch (const std::exception & e) {
|
||||
SRV_WRN("disk prompt cache: failed to deserialize tokens from '%s': %s\n", file.name.c_str(), e.what());
|
||||
remove_file(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
// the filename hash only proves an exact prefix probabilistically - verify against the actual tokens
|
||||
if (loaded.size() != file.n_tokens ||
|
||||
loaded.get_common_prefix(tokens_new) != file.n_tokens ||
|
||||
!loaded.validate(ctx)) {
|
||||
SRV_WRN("disk prompt cache: token mismatch in '%s' (hash collision?)\n", file.name.c_str());
|
||||
remove_file(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> state;
|
||||
try {
|
||||
state.resize(header.state_size);
|
||||
} catch (const std::bad_alloc &) {
|
||||
SRV_ERR("disk prompt cache: failed to allocate %" PRIu64 " bytes for '%s'\n", header.state_size, file.name.c_str());
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
in.read((char *) state.data(), state.size());
|
||||
|
||||
if (!in.good()) {
|
||||
SRV_WRN("disk prompt cache: '%s' is truncated\n", file.name.c_str());
|
||||
remove_file(file);
|
||||
return LOAD_MISS;
|
||||
}
|
||||
|
||||
const size_t n = llama_state_seq_set_data_ext(ctx, state.data(), state.size(), id_slot, 0);
|
||||
if (n != state.size()) {
|
||||
SRV_WRN("disk prompt cache: failed to restore state from '%s' (%zu / %zu bytes)\n", file.name.c_str(), n, state.size());
|
||||
|
||||
// the sequence may hold a partial state now - clear it and let the caller recover
|
||||
llama_memory_seq_rm(llama_get_memory(ctx), id_slot, -1, -1);
|
||||
|
||||
return LOAD_FAIL_SEQ_DIRTY;
|
||||
}
|
||||
|
||||
tokens_out = std::move(loaded);
|
||||
|
||||
covered.insert(covered_key(file.n_tokens, file.chain_hash));
|
||||
|
||||
touch(file);
|
||||
|
||||
SRV_INF("disk prompt cache: restored prompt with %u tokens, %.3f MiB from '%s'\n",
|
||||
file.n_tokens, state.size() / (1024.0 * 1024.0), file.name.c_str());
|
||||
SRV_DBG("%s", "__TEST_TAG_CACHE_DISK_HIT__\n");
|
||||
|
||||
return LOAD_OK;
|
||||
}
|
||||
|
||||
//
|
||||
// compat hash
|
||||
//
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename T>
|
||||
void hash_pod(std::string & blob, const T & value) {
|
||||
static_assert(std::is_trivially_copyable<T>::value, "hash_pod requires a POD type");
|
||||
blob.append((const char *) &value, sizeof(value));
|
||||
}
|
||||
|
||||
void hash_str(std::string & blob, const std::string & value) {
|
||||
blob += value;
|
||||
blob += '\0';
|
||||
}
|
||||
|
||||
// path + size + mtime: conservative, but never misses a changed file
|
||||
void hash_file_meta(std::string & blob, const std::string & path) {
|
||||
hash_str(blob, path);
|
||||
|
||||
std::error_code ec;
|
||||
|
||||
const uint64_t size = path.empty() ? 0 : (uint64_t) std::filesystem::file_size(path, ec);
|
||||
hash_pod(blob, ec ? (uint64_t) 0 : size);
|
||||
|
||||
hash_pod(blob, path.empty() ? (int64_t) 0 : file_mtime(path));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
uint64_t server_cache_disk_compat_hash(const common_params & params) {
|
||||
std::string blob;
|
||||
|
||||
// format versions
|
||||
hash_pod(blob, (uint32_t) SERVER_CACHE_DISK_VERSION);
|
||||
hash_pod(blob, (uint32_t) LLAMA_STATE_SEQ_VERSION);
|
||||
hash_pod(blob, (uint32_t) server_tokens::SERVER_TOKENS_STATE_VERSION);
|
||||
|
||||
// model identity
|
||||
hash_file_meta(blob, params.model.path);
|
||||
hash_file_meta(blob, params.mmproj.path);
|
||||
|
||||
for (const auto & la : params.lora_adapters) {
|
||||
hash_file_meta(blob, la.path);
|
||||
hash_pod(blob, la.scale);
|
||||
}
|
||||
|
||||
// KV cache layout
|
||||
hash_pod(blob, (int32_t) params.cache_type_k);
|
||||
hash_pod(blob, (int32_t) params.cache_type_v);
|
||||
hash_pod(blob, (uint8_t) params.swa_full);
|
||||
|
||||
// rope params change the KV content for the same tokens
|
||||
hash_pod(blob, params.rope_freq_base);
|
||||
hash_pod(blob, params.rope_freq_scale);
|
||||
hash_pod(blob, (int32_t) params.rope_scaling_type);
|
||||
hash_pod(blob, params.yarn_ext_factor);
|
||||
hash_pod(blob, params.yarn_attn_factor);
|
||||
hash_pod(blob, params.yarn_beta_fast);
|
||||
hash_pod(blob, params.yarn_beta_slow);
|
||||
hash_pod(blob, params.yarn_orig_ctx);
|
||||
|
||||
return XXH64(blob.data(), blob.size(), 0);
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
#pragma once
|
||||
|
||||
#include "server-common.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
struct common_params;
|
||||
struct llama_context;
|
||||
|
||||
// disk-backed prompt cache: a cold tier below the in-RAM server_prompt_cache
|
||||
//
|
||||
// each entry is one file in a flat directory, named after the exact token prefix it contains:
|
||||
//
|
||||
// {compat_hash8}-{n_tokens}-{chain_hash16}.kvc
|
||||
//
|
||||
// - compat_hash: hash of everything that invalidates a KV state (model file, mmproj, loras,
|
||||
// cache types, rope params, ...) - see server_cache_disk_compat_hash()
|
||||
// - chain_hash: chained hash over the first n_tokens tokens, so a filename identifies an exact
|
||||
// prefix and lookup is a single rolling-hash pass over the incoming prompt plus an index probe
|
||||
//
|
||||
// file contents mirror what the RAM cache holds for the target context:
|
||||
//
|
||||
// header | server_tokens::serialize() bytes | llama_state_seq_get_data (FLAGS_NONE) bytes
|
||||
|
||||
struct server_cache_disk_file {
|
||||
std::string name; // filename inside the cache directory
|
||||
|
||||
uint64_t chain_hash = 0;
|
||||
uint32_t n_tokens = 0;
|
||||
uint64_t n_bytes = 0;
|
||||
int64_t mtime = 0; // only used for relative ordering during eviction
|
||||
};
|
||||
|
||||
struct server_prompt_cache_disk {
|
||||
server_prompt_cache_disk(const std::string & dir, uint64_t compat_hash, bool has_mtmd, int32_t limit_mib, bool write_through);
|
||||
|
||||
enum load_status {
|
||||
LOAD_OK, // state restored into the sequence
|
||||
LOAD_MISS, // file unusable (corrupt, collision, ...) - sequence untouched
|
||||
LOAD_FAIL_SEQ_DIRTY, // restore failed mid-way - the sequence was cleared and must be re-filled
|
||||
};
|
||||
|
||||
// largest exact-prefix hit for the first n_max tokens, or nullptr on miss
|
||||
const server_cache_disk_file * lookup(const server_tokens & tokens, size_t n_max) const;
|
||||
|
||||
// restore the state from a file into sequence id_slot of ctx
|
||||
// on LOAD_OK, tokens_out receives the cached token list (an exact prefix of tokens_new)
|
||||
load_status load(server_cache_disk_file file, const server_tokens & tokens_new, llama_context * ctx, int32_t id_slot, server_tokens & tokens_out);
|
||||
|
||||
// write one entry; deduplicates against existing files and enforces the size budget
|
||||
bool store(const server_tokens & tokens, const std::vector<uint8_t> & state_main);
|
||||
|
||||
size_t n_files() const;
|
||||
size_t n_bytes_total() const { return total_bytes; }
|
||||
|
||||
const bool write_through;
|
||||
|
||||
private:
|
||||
void scan_dir();
|
||||
|
||||
server_cache_disk_file * find_file(uint32_t n_tokens, uint64_t chain_hash);
|
||||
|
||||
void touch (const server_cache_disk_file & file); // bump mtime so eviction treats it as fresh
|
||||
void forget(const server_cache_disk_file & file); // drop from the index without touching the filesystem
|
||||
void remove_file(const server_cache_disk_file & file); // delete from disk and drop from the index
|
||||
|
||||
// delete oldest-mtime files (ours and foreign alike) while over the size budget
|
||||
void enforce_budget(const std::string & name_protected);
|
||||
|
||||
const std::string dir;
|
||||
const uint64_t compat_hash;
|
||||
const bool has_mtmd;
|
||||
const size_t limit_bytes; // 0 = no limit
|
||||
|
||||
// n_tokens -> chain_hash -> file, for our compat hash only
|
||||
std::map<uint32_t, std::unordered_map<uint64_t, server_cache_disk_file>> index;
|
||||
|
||||
// .kvc files with a different compat hash prefix - never opened, but counted toward the budget
|
||||
std::vector<server_cache_disk_file> foreign;
|
||||
|
||||
size_t total_bytes = 0; // ours + foreign
|
||||
|
||||
// (n_tokens, chain_hash) prefixes known to be covered by a file written or loaded this
|
||||
// session - lets store() skip prefixes of already-persisted prompts
|
||||
std::unordered_set<uint64_t> covered;
|
||||
|
||||
uint32_t tmp_counter = 0;
|
||||
};
|
||||
|
||||
// hash of everything that invalidates a saved KV state for the current server configuration
|
||||
uint64_t server_cache_disk_compat_hash(const common_params & params);
|
||||
@@ -266,8 +266,6 @@ static inline raw_buffer base64_decode(const std::string & encoded_string) {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr uint32_t SERVER_TOKENS_STATE_VERSION = 1;
|
||||
|
||||
uint32_t server_tokens_state_u32(size_t value) {
|
||||
if (value > std::numeric_limits<uint32_t>::max()) {
|
||||
throw std::runtime_error("Server tokens state is too large");
|
||||
@@ -507,6 +505,23 @@ void server_tokens::push_back(const mtmd_input_chunk * chunk) {
|
||||
}
|
||||
}
|
||||
|
||||
void server_tokens::push_back_placeholder(const mtmd_input_chunk * chunk) {
|
||||
auto type = mtmd_input_chunk_get_type(chunk);
|
||||
if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE || type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||||
GGML_ASSERT(has_mtmd);
|
||||
mtmd::input_chunk_ptr new_chunk(mtmd_input_chunk_get_placeholder(chunk));
|
||||
GGML_ASSERT(new_chunk != nullptr && "failed to create placeholder chunk");
|
||||
const size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk);
|
||||
size_t start_idx = tokens.size();
|
||||
for (size_t i = 0; i < n_tokens; ++i) {
|
||||
tokens.emplace_back(LLAMA_TOKEN_NULL);
|
||||
}
|
||||
map_idx_to_media[start_idx] = std::move(new_chunk);
|
||||
} else {
|
||||
push_back(chunk);
|
||||
}
|
||||
}
|
||||
|
||||
void server_tokens::push_back(server_tokens & tokens) {
|
||||
size_t start_idx = size();
|
||||
for (size_t i = 0; i < tokens.size(); i++) {
|
||||
|
||||
@@ -156,6 +156,9 @@ private: // disallow accessing these members directly, risking out-of-sync
|
||||
// map_idx_to_media will contain: {5, img0}, {8, img1}
|
||||
|
||||
public:
|
||||
// version of the serialize()/deserialize() format below
|
||||
static constexpr uint32_t SERVER_TOKENS_STATE_VERSION = 1;
|
||||
|
||||
server_tokens() = default;
|
||||
~server_tokens() = default;
|
||||
|
||||
@@ -195,6 +198,10 @@ public:
|
||||
// will create a copy of the chunk if it contains non-text data
|
||||
void push_back(const mtmd_input_chunk * chunk);
|
||||
|
||||
// same as push_back, but media chunks are stored as placeholders (no image/audio data)
|
||||
// only use this if the chunk will never be encoded again (e.g. it is already in the KV cache)
|
||||
void push_back_placeholder(const mtmd_input_chunk * chunk);
|
||||
|
||||
// appends server tokens, updates the media map. copies media chunks.
|
||||
void push_back(server_tokens & tokens);
|
||||
|
||||
|
||||
@@ -275,11 +275,13 @@ struct server_slot {
|
||||
llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE);
|
||||
}
|
||||
|
||||
prompt_cache.disk_store_write_through(*cur);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) {
|
||||
bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id);
|
||||
bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id, n_ctx);
|
||||
if (!res) {
|
||||
SLT_WRN(*this, "%s", "failed to load prompt from cache\n");
|
||||
}
|
||||
@@ -1308,7 +1310,26 @@ private:
|
||||
SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n");
|
||||
|
||||
prompt_cache = std::make_unique<server_prompt_cache>(params_base.cache_ram_mib, n_ctx);
|
||||
|
||||
if (!params_base.cache_disk_path.empty()) {
|
||||
const uint64_t compat_hash = server_cache_disk_compat_hash(params_base);
|
||||
|
||||
SRV_INF("disk prompt cache is enabled, dir: '%s', compat hash: %08x\n",
|
||||
params_base.cache_disk_path.c_str(), (uint32_t) compat_hash);
|
||||
|
||||
prompt_cache->disk = std::make_unique<server_prompt_cache_disk>(
|
||||
params_base.cache_disk_path,
|
||||
compat_hash,
|
||||
mctx != nullptr,
|
||||
params_base.cache_disk_limit_mib,
|
||||
params_base.cache_disk_write_through);
|
||||
}
|
||||
} else {
|
||||
if (!params_base.cache_disk_path.empty()) {
|
||||
SRV_ERR("%s", "--cache-disk requires the RAM prompt cache - remove `--cache-ram 0`\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n");
|
||||
}
|
||||
SRV_TRC("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n");
|
||||
@@ -3416,7 +3437,8 @@ private:
|
||||
// add the mtmd chunk to cache
|
||||
{
|
||||
const auto & chunk = input_tokens.find_chunk(cur_token_idx);
|
||||
slot.prompt.tokens.push_back(chunk.get()); // copy
|
||||
// the chunk is already in the KV cache at this point, so we don't need to keep its data around
|
||||
slot.prompt.tokens.push_back_placeholder(chunk.get());
|
||||
}
|
||||
|
||||
has_mtmd = true;
|
||||
@@ -4057,6 +4079,11 @@ bool server_context::load_model(common_params & params) {
|
||||
void server_context::start_loop() {
|
||||
auto & params = impl->params_base;
|
||||
impl->queue_tasks.start_loop(params.sleep_idle_seconds * 1000);
|
||||
|
||||
// on graceful shutdown, give the RAM prompt cache entries a chance to survive the restart
|
||||
if (impl->prompt_cache) {
|
||||
impl->prompt_cache->disk_flush();
|
||||
}
|
||||
}
|
||||
|
||||
void server_context::terminate() {
|
||||
|
||||
@@ -1750,6 +1750,8 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
|
||||
SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n",
|
||||
states.front().size() / (1024.0 * 1024.0));
|
||||
|
||||
spill_front();
|
||||
|
||||
states.pop_front();
|
||||
}
|
||||
}
|
||||
@@ -1787,7 +1789,7 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro
|
||||
return &states.back();
|
||||
}
|
||||
|
||||
bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) {
|
||||
bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, int32_t n_ctx_slot) {
|
||||
const int lcp_best = prompt.tokens.get_common_prefix(tokens_new);
|
||||
|
||||
float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins
|
||||
@@ -1797,6 +1799,8 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
|
||||
|
||||
auto it_best = states.end();
|
||||
|
||||
int lcp_it_best = 0;
|
||||
|
||||
// find the most similar cached prompt, that would also preserve the most context
|
||||
for (auto it = states.begin(); it != states.end(); ++it) {
|
||||
const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new);
|
||||
@@ -1815,7 +1819,41 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
|
||||
f_keep_best = f_keep_cur;
|
||||
f_sim_best = f_sim_cur;
|
||||
|
||||
it_best = it;
|
||||
it_best = it;
|
||||
lcp_it_best = lcp_cur;
|
||||
}
|
||||
}
|
||||
|
||||
// check the disk tier for an exact-prefix match longer than what RAM (or the slot itself) offers
|
||||
if (disk) {
|
||||
const int lcp_sel = std::max(lcp_best, lcp_it_best);
|
||||
|
||||
const size_t n_max = std::min<size_t>(tokens_new.size(), std::max(0, n_ctx_slot));
|
||||
|
||||
const auto * file = disk->lookup(tokens_new, n_max);
|
||||
|
||||
if (file && (int64_t) file->n_tokens > (int64_t) lcp_sel) {
|
||||
server_tokens tokens_disk;
|
||||
|
||||
const auto status = disk->load(*file, tokens_new, ctx_tgt, id_slot, tokens_disk);
|
||||
|
||||
if (status == server_prompt_cache_disk::LOAD_OK) {
|
||||
// disk entries carry no draft state - clear the draft sequence so it re-prefills
|
||||
if (ctx_dft) {
|
||||
llama_memory_seq_rm(llama_get_memory(ctx_dft), id_slot, -1, -1);
|
||||
}
|
||||
|
||||
prompt.tokens = std::move(tokens_disk);
|
||||
prompt.checkpoints.clear();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
if (status == server_prompt_cache_disk::LOAD_FAIL_SEQ_DIRTY && it_best == states.end()) {
|
||||
// the slot's sequence was cleared during the failed restore and there is no RAM
|
||||
// candidate to restore over it - the caller has to clear the slot
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1869,6 +1907,8 @@ void server_prompt_cache::update() {
|
||||
while (!states.empty() && size() > limit_size) {
|
||||
SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
|
||||
|
||||
spill_front();
|
||||
|
||||
states.pop_front();
|
||||
}
|
||||
}
|
||||
@@ -1884,6 +1924,8 @@ void server_prompt_cache::update() {
|
||||
SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
|
||||
limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));
|
||||
|
||||
spill_front();
|
||||
|
||||
states.pop_front();
|
||||
}
|
||||
}
|
||||
@@ -1896,3 +1938,39 @@ void server_prompt_cache::update() {
|
||||
(const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache::disk_store(const server_prompt_cache_state & state) const {
|
||||
if (!disk || state.data.main.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
disk->store(state.prompt.tokens, state.data.main);
|
||||
}
|
||||
|
||||
void server_prompt_cache::disk_store_write_through(const server_prompt_cache_state & state) const {
|
||||
if (!disk || !disk->write_through) {
|
||||
return;
|
||||
}
|
||||
|
||||
disk_store(state);
|
||||
}
|
||||
|
||||
void server_prompt_cache::disk_flush() const {
|
||||
if (!disk) {
|
||||
return;
|
||||
}
|
||||
|
||||
SRV_INF("flushing %zu prompt cache entries to disk\n", states.size());
|
||||
|
||||
for (const auto & state : states) {
|
||||
disk_store(state);
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache::spill_front() const {
|
||||
if (!disk || states.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
disk_store(states.front());
|
||||
}
|
||||
|
||||
@@ -7,8 +7,10 @@
|
||||
#include <unordered_set>
|
||||
#include <list>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
|
||||
// TODO: prevent including the whole server-common.h as we only use server_tokens
|
||||
#include "server-cache-disk.h"
|
||||
#include "server-common.h"
|
||||
|
||||
using json = nlohmann::ordered_json;
|
||||
@@ -612,6 +614,10 @@ struct server_prompt_cache {
|
||||
|
||||
std::list<server_prompt_cache_state> states;
|
||||
|
||||
// optional cold tier - entries evicted from RAM are spilled here and can be restored later,
|
||||
// including across server restarts
|
||||
std::unique_ptr<server_prompt_cache_disk> disk;
|
||||
|
||||
// in bytes, 0 = no limit
|
||||
size_t limit_size = 0;
|
||||
|
||||
@@ -624,9 +630,22 @@ struct server_prompt_cache {
|
||||
|
||||
server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
|
||||
|
||||
bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot);
|
||||
bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot, int32_t n_ctx_slot);
|
||||
|
||||
void update();
|
||||
|
||||
// write one RAM cache entry to the disk tier (no-op when the disk tier is disabled)
|
||||
void disk_store(const server_prompt_cache_state & state) const;
|
||||
|
||||
// disk_store, but only when write-through mode is enabled
|
||||
void disk_store_write_through(const server_prompt_cache_state & state) const;
|
||||
|
||||
// spill all RAM entries to the disk tier (e.g. on graceful shutdown)
|
||||
void disk_flush() const;
|
||||
|
||||
private:
|
||||
// spill the entry that is about to be evicted
|
||||
void spill_front() const;
|
||||
};
|
||||
|
||||
// used exclusively by router mode
|
||||
|
||||
@@ -0,0 +1,268 @@
|
||||
import base64
|
||||
import glob
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from utils import *
|
||||
|
||||
server = ServerPreset.tinyllama2()
|
||||
|
||||
cache_dir: str = ""
|
||||
|
||||
|
||||
class LogReader:
|
||||
def __init__(self, path):
|
||||
self.path = path
|
||||
self.pos = 0
|
||||
def drain(self):
|
||||
with open(self.path) as f:
|
||||
f.seek(self.pos)
|
||||
content = f.read()
|
||||
self.pos = f.tell()
|
||||
return content
|
||||
def wait_for(self, tag, timeout=10) -> bool:
|
||||
# the server log is pumped to the file asynchronously - poll for the tag
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if tag in self.drain():
|
||||
return True
|
||||
time.sleep(0.25)
|
||||
return False
|
||||
|
||||
|
||||
def kvc_files() -> list[str]:
|
||||
return sorted(glob.glob(os.path.join(cache_dir, "*.kvc")))
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def create_server():
|
||||
global server, cache_dir
|
||||
cache_dir = tempfile.mkdtemp(prefix="llama_cache_disk_")
|
||||
server = ServerPreset.tinyllama2()
|
||||
server.n_slots = 1
|
||||
server.temperature = 0.0
|
||||
server.debug = True
|
||||
server.cache_disk = cache_dir
|
||||
fd, server.log_path = tempfile.mkstemp(suffix='.log')
|
||||
os.close(fd)
|
||||
yield
|
||||
shutil.rmtree(cache_dir, ignore_errors=True)
|
||||
|
||||
|
||||
PROMPT_A = (
|
||||
"Once upon a time in a land far away, there lived a brave knight "
|
||||
"who traveled across mountains and rivers to find the legendary "
|
||||
"golden sword hidden deep within the enchanted forest of whispers."
|
||||
)
|
||||
|
||||
PROMPT_B = "The quick brown fox jumps over the lazy dog."
|
||||
|
||||
|
||||
def make_prompt_request(prompt, n_predict=0):
|
||||
global server
|
||||
res = server.make_request("POST", "/completion", data={
|
||||
"prompt": prompt,
|
||||
"n_predict": n_predict, # 0 = evaluate the prompt into the KV cache only
|
||||
"cache_prompt": True,
|
||||
})
|
||||
assert res.status_code == 200
|
||||
return res
|
||||
|
||||
|
||||
def test_write_through_and_restart_hit():
|
||||
global server
|
||||
server.cache_disk_write_through = True
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
res = make_prompt_request(PROMPT_A)
|
||||
prompt_n_full = res.body["timings"]["prompt_n"]
|
||||
assert prompt_n_full > 0
|
||||
|
||||
# nothing is written while the prompt is still live in the slot
|
||||
assert len(kvc_files()) == 0
|
||||
|
||||
# a different prompt takes over the only slot - the previous one is saved
|
||||
# to the RAM cache and, in write-through mode, to disk immediately
|
||||
make_prompt_request(PROMPT_B)
|
||||
assert log.wait_for("__TEST_TAG_CACHE_DISK_STORE__")
|
||||
assert len(kvc_files()) == 1
|
||||
|
||||
# the state must survive a full server restart
|
||||
server.stop()
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
res = make_prompt_request(PROMPT_A)
|
||||
assert log.wait_for("__TEST_TAG_CACHE_DISK_HIT__")
|
||||
assert res.body["timings"]["prompt_n"] == 1 # only the last token is re-evaluated
|
||||
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
|
||||
|
||||
|
||||
def test_spill_on_shutdown_flush():
|
||||
global server
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
make_prompt_request(PROMPT_A)
|
||||
make_prompt_request(PROMPT_B) # forces PROMPT_A into the RAM cache
|
||||
|
||||
# without write-through, nothing reaches the disk while running
|
||||
time.sleep(0.5)
|
||||
assert "__TEST_TAG_CACHE_DISK_STORE__" not in log.drain()
|
||||
assert len(kvc_files()) == 0
|
||||
|
||||
# a graceful shutdown flushes the RAM cache entries to disk
|
||||
server.stop()
|
||||
assert len(kvc_files()) == 1
|
||||
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
res = make_prompt_request(PROMPT_A)
|
||||
assert log.wait_for("__TEST_TAG_CACHE_DISK_HIT__")
|
||||
assert res.body["timings"]["prompt_n"] == 1
|
||||
|
||||
|
||||
def test_ram_cache_hit_takes_priority():
|
||||
global server
|
||||
server.cache_disk_write_through = True
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
make_prompt_request(PROMPT_A)
|
||||
make_prompt_request(PROMPT_B)
|
||||
assert len(kvc_files()) == 1
|
||||
|
||||
# PROMPT_A is in both the RAM cache and on disk - the RAM copy must win
|
||||
# (the disk entry is never longer than the RAM one here)
|
||||
res = make_prompt_request(PROMPT_A)
|
||||
time.sleep(0.5)
|
||||
assert "__TEST_TAG_CACHE_DISK_HIT__" not in log.drain()
|
||||
assert res.body["timings"]["cache_n"] > 0
|
||||
|
||||
|
||||
def test_budget_eviction():
|
||||
global server
|
||||
server.n_ctx = 2048
|
||||
server.n_batch = 512
|
||||
server.cache_disk_write_through = True
|
||||
server.cache_disk_limit = 1 # MiB
|
||||
server.start()
|
||||
|
||||
# three long, distinct token-array prompts; each state is close to 1 MiB
|
||||
n_len = 1500
|
||||
for i in range(3):
|
||||
make_prompt_request([100 + i] * n_len)
|
||||
|
||||
# one final small prompt to force the last long prompt out of the slot
|
||||
make_prompt_request(PROMPT_B)
|
||||
|
||||
files = kvc_files()
|
||||
assert len(files) >= 1
|
||||
assert len(files) < 3 # the oldest entries were evicted
|
||||
|
||||
# the budget is respected (a single over-budget file is allowed to remain)
|
||||
if len(files) > 1:
|
||||
assert sum(os.path.getsize(f) for f in files) <= 1024 * 1024
|
||||
|
||||
|
||||
def test_corrupt_file_is_removed():
|
||||
global server
|
||||
server.cache_disk_write_through = True
|
||||
server.start()
|
||||
|
||||
make_prompt_request(PROMPT_A)
|
||||
make_prompt_request(PROMPT_B)
|
||||
files = kvc_files()
|
||||
assert len(files) == 1
|
||||
|
||||
server.stop()
|
||||
|
||||
# corrupt the serialized token section (starts right after the 48-byte header)
|
||||
with open(files[0], "r+b") as f:
|
||||
f.seek(48 + 4)
|
||||
f.write(b"\xff\xff\xff\xff")
|
||||
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
# the request must still succeed, with the prompt fully re-processed
|
||||
res = make_prompt_request(PROMPT_A)
|
||||
time.sleep(0.5)
|
||||
assert "__TEST_TAG_CACHE_DISK_HIT__" not in log.drain()
|
||||
assert res.body["timings"]["prompt_n"] > 1
|
||||
|
||||
# the corrupt file was deleted
|
||||
assert len(kvc_files()) == 0
|
||||
|
||||
|
||||
IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png"
|
||||
|
||||
|
||||
def _get_img_base64(url: str) -> str:
|
||||
response = requests.get(url)
|
||||
response.raise_for_status()
|
||||
return base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mmproj_server():
|
||||
global cache_dir
|
||||
os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>'
|
||||
mm_server = ServerPreset.tinygemma3()
|
||||
mm_server.n_slots = 1
|
||||
mm_server.temperature = 0.0
|
||||
mm_server.debug = True
|
||||
# use the full SWA cache so the restored image prefix can be reused
|
||||
mm_server.swa_full = True
|
||||
mm_server.cache_disk = cache_dir
|
||||
mm_server.cache_disk_write_through = True
|
||||
fd, mm_server.log_path = tempfile.mkstemp(suffix='.log')
|
||||
os.close(fd)
|
||||
return mm_server
|
||||
|
||||
|
||||
def test_image_prompt_across_restart(mmproj_server):
|
||||
server = mmproj_server
|
||||
server.start()
|
||||
|
||||
prompt_cat = {
|
||||
"prompt_string": "What is this: <__media__>\n",
|
||||
"multimodal_data": [_get_img_base64(IMG_URL_CAT)],
|
||||
}
|
||||
|
||||
res = server.make_request("POST", "/completions", data={
|
||||
"n_predict": 0,
|
||||
"cache_prompt": True,
|
||||
"prompt": prompt_cat,
|
||||
})
|
||||
assert res.status_code == 200
|
||||
prompt_n_full = res.body["timings"]["prompt_n"]
|
||||
|
||||
res = server.make_request("POST", "/completions", data={
|
||||
"n_predict": 0,
|
||||
"cache_prompt": True,
|
||||
"prompt": "The quick brown fox",
|
||||
})
|
||||
assert res.status_code == 200
|
||||
assert len(kvc_files()) == 1
|
||||
|
||||
server.stop()
|
||||
server.start()
|
||||
log = LogReader(server.log_path)
|
||||
|
||||
# the image KV must be restored from disk in the new process
|
||||
res = server.make_request("POST", "/completions", data={
|
||||
"n_predict": 0,
|
||||
"cache_prompt": True,
|
||||
"prompt": prompt_cat,
|
||||
})
|
||||
assert res.status_code == 200
|
||||
assert log.wait_for("__TEST_TAG_CACHE_DISK_HIT__")
|
||||
assert res.body["timings"]["prompt_n"] == 1
|
||||
assert res.body["timings"]["cache_n"] == prompt_n_full - 1
|
||||
@@ -111,6 +111,9 @@ class ServerProcess:
|
||||
media_path: str | None = None
|
||||
sleep_idle_seconds: int | None = None
|
||||
cache_ram: int | None = None
|
||||
cache_disk: str | None = None
|
||||
cache_disk_limit: int | None = None
|
||||
cache_disk_write_through: bool = False
|
||||
no_cache_idle_slots: bool = False
|
||||
log_path: str | None = None
|
||||
ui_mcp_proxy: bool = False
|
||||
@@ -271,6 +274,12 @@ class ServerProcess:
|
||||
server_args.extend(["--sleep-idle-seconds", self.sleep_idle_seconds])
|
||||
if self.cache_ram is not None:
|
||||
server_args.extend(["--cache-ram", self.cache_ram])
|
||||
if self.cache_disk is not None:
|
||||
server_args.extend(["--cache-disk", self.cache_disk])
|
||||
if self.cache_disk_limit is not None:
|
||||
server_args.extend(["--cache-disk-limit", self.cache_disk_limit])
|
||||
if self.cache_disk_write_through:
|
||||
server_args.append("--cache-disk-write-through")
|
||||
if self.no_cache_idle_slots:
|
||||
server_args.append("--no-cache-idle-slots")
|
||||
if self.ui_mcp_proxy:
|
||||
|
||||
@@ -259,6 +259,8 @@ int main(int argc, char ** argv) {
|
||||
}
|
||||
cpp += fmt("static const unsigned char asset_%d_data[] = {", i);
|
||||
append_bytes_hex(cpp, bytes);
|
||||
|
||||
// note: this is a simple hash for cache busting, not a cryptographic hash; fnv is enough here
|
||||
const auto hash = fnv_hash(bytes.data(), bytes.size());
|
||||
|
||||
cpp += fmt("};\nstatic const std::size_t asset_%d_size = %zu;\n",
|
||||
|
||||
Vendored
+16
-4
@@ -4,18 +4,30 @@ llama_add_compile_flags()
|
||||
|
||||
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
||||
|
||||
add_library(${TARGET} STATIC
|
||||
set(VENDOR_SRCS
|
||||
xxhash/xxhash.c
|
||||
sha1/sha1.c
|
||||
sha256/sha256.c
|
||||
)
|
||||
|
||||
# disable warnings in 3rd party code
|
||||
add_library(${TARGET} STATIC
|
||||
hash.cpp
|
||||
hash.h
|
||||
${VENDOR_SRCS}
|
||||
)
|
||||
|
||||
target_compile_features(${TARGET} PRIVATE cxx_std_17)
|
||||
|
||||
# disable warnings in 3rd party code, but keep them for hash.cpp
|
||||
if (CMAKE_C_COMPILER_ID STREQUAL "MSVC")
|
||||
target_compile_options(${TARGET} PRIVATE /w)
|
||||
set(NO_WARN_FLAG /w)
|
||||
else()
|
||||
target_compile_options(${TARGET} PRIVATE -w)
|
||||
set(NO_WARN_FLAG -w)
|
||||
endif()
|
||||
set_source_files_properties(${VENDOR_SRCS} PROPERTIES COMPILE_OPTIONS ${NO_WARN_FLAG})
|
||||
|
||||
# sha1 lives in a namespace to avoid a clash with boringssl, see scripts/sync_vendor.py
|
||||
set_source_files_properties(sha1/sha1.c PROPERTIES LANGUAGE CXX)
|
||||
|
||||
# sha256.c includes "rotate-bits/rotate-bits.h", so consumers get this dir too
|
||||
target_include_directories(${TARGET} PUBLIC .)
|
||||
|
||||
Vendored
+23
@@ -0,0 +1,23 @@
|
||||
#include "hash.h"
|
||||
|
||||
extern "C" {
|
||||
#include "sha256/sha256.h"
|
||||
}
|
||||
|
||||
static std::string to_hex(const unsigned char * digest, size_t len) {
|
||||
static const char hex[] = "0123456789abcdef";
|
||||
|
||||
std::string out;
|
||||
out.reserve(2*len);
|
||||
for (size_t i = 0; i < len; ++i) {
|
||||
out += hex[digest[i] >> 4];
|
||||
out += hex[digest[i] & 0xf];
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::string hash_sha256_hex(const void * data, size_t len) {
|
||||
unsigned char digest[SHA256_DIGEST_SIZE];
|
||||
sha256_hash(digest, (const unsigned char *) data, len);
|
||||
return to_hex(digest, SHA256_DIGEST_SIZE);
|
||||
}
|
||||
Vendored
+9
@@ -0,0 +1,9 @@
|
||||
#pragma once
|
||||
|
||||
// C++ wrapper for the vendored hash functions
|
||||
|
||||
#include <cstddef>
|
||||
#include <string>
|
||||
|
||||
// returns the SHA-256 digest as a lowercase hex string
|
||||
std::string hash_sha256_hex(const void * data, size_t len);
|
||||
Vendored
+4
@@ -25,6 +25,8 @@ A million repetitions of "a"
|
||||
|
||||
#include "sha1.h"
|
||||
|
||||
namespace vendor_hash {
|
||||
|
||||
|
||||
#define rol(value, bits) (((value) << (bits)) | ((value) >> (32 - (bits))))
|
||||
|
||||
@@ -293,3 +295,5 @@ void SHA1(
|
||||
SHA1Final((unsigned char *)hash_out, &ctx);
|
||||
}
|
||||
|
||||
} // namespace vendor_hash
|
||||
|
||||
|
||||
Vendored
+2
-6
@@ -9,9 +9,7 @@
|
||||
|
||||
#include "stdint.h"
|
||||
|
||||
#if defined(__cplusplus)
|
||||
extern "C" {
|
||||
#endif
|
||||
namespace vendor_hash {
|
||||
|
||||
typedef struct
|
||||
{
|
||||
@@ -45,8 +43,6 @@ void SHA1(
|
||||
const char *str,
|
||||
uint32_t len);
|
||||
|
||||
#if defined(__cplusplus)
|
||||
}
|
||||
#endif
|
||||
} // namespace vendor_hash
|
||||
|
||||
#endif /* SHA1_H */
|
||||
|
||||
Reference in New Issue
Block a user