Skip to content

Commit fefbeb8

Browse files
ggerganovarthw
ggerganov
authored and
arthw
committed
server : reuse cached context chunks (ggml-org#9866)
ggml-ci
1 parent d710541 commit fefbeb8

File tree

5 files changed

+78
-6
lines changed

5 files changed

+78
-6
lines changed

‎common/arg.cpp

+7
Original file line numberDiff line numberDiff line change
@@ -1788,6 +1788,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
17881788
params.n_threads_http = value;
17891789
}
17901790
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_THREADS_HTTP"));
1791+
add_opt(common_arg(
1792+
{"--cache-reuse"}, "N",
1793+
string_format("min chunk size to attempt reusing from the cache via KV shifting (default: %d)", params.n_cache_reuse),
1794+
[](common_params & params, int value) {
1795+
params.n_cache_reuse = value;
1796+
}
1797+
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CACHE_REUSE"));
17911798
add_opt(common_arg(
17921799
{"--metrics"},
17931800
string_format("enable prometheus compatible metrics endpoint (default: %s)", params.endpoint_metrics ? "enabled" : "disabled"),

‎common/common.h

+2-1
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,8 @@ struct common_params {
277277
int32_t port = 8080; // server listens on this network port
278278
int32_t timeout_read = 600; // http read timeout in seconds
279279
int32_t timeout_write = timeout_read; // http write timeout in seconds
280-
int n_threads_http = -1; // number of threads to process HTTP requests (TODO: support threadpool)
280+
int32_t n_threads_http = -1; // number of threads to process HTTP requests (TODO: support threadpool)
281+
int32_t n_cache_reuse = 0; // min chunk size to reuse from the cache via KV shifting
281282

282283
std::string hostname = "127.0.0.1";
283284
std::string public_path = ""; // NOLINT

‎examples/server/README.md

+1
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ The project is under active development, and we are [looking for feedback and co
147147
| `--ssl-cert-file FNAME` | path to file a PEM-encoded SSL certificate<br/>(env: LLAMA_ARG_SSL_CERT_FILE) |
148148
| `-to, --timeout N` | server read/write timeout in seconds (default: 600)<br/>(env: LLAMA_ARG_TIMEOUT) |
149149
| `--threads-http N` | number of threads used to process HTTP requests (default: -1)<br/>(env: LLAMA_ARG_THREADS_HTTP) |
150+
| `--cache-reuse N` | min chunk size to attempt reusing from the cache via KV shifting (default: 0)<br/>(env: LLAMA_ARG_CACHE_REUSE) |
150151
| `--metrics` | enable prometheus compatible metrics endpoint (default: disabled)<br/>(env: LLAMA_ARG_ENDPOINT_METRICS) |
151152
| `--slots` | enable slots monitoring endpoint (default: disabled)<br/>(env: LLAMA_ARG_ENDPOINT_SLOTS) |
152153
| `--props` | enable changing global properties via POST /props (default: disabled)<br/>(env: LLAMA_ARG_ENDPOINT_PROPS) |

‎examples/server/server.cpp

+66-3
Original file line numberDiff line numberDiff line change
@@ -800,7 +800,7 @@ struct server_context {
800800
int slot_prompt_len = slot_prompt.size();
801801

802802
// length of the Longest Common Prefix between the current slot's prompt and the input prompt
803-
int lcp_len = common_part(slot_prompt, prompt);
803+
int lcp_len = longest_common_prefix(slot_prompt, prompt);
804804

805805
// fraction of the common substring length compared to the current slot's prompt length
806806
similarity = static_cast<float>(lcp_len) / slot_prompt_len;
@@ -2012,7 +2012,7 @@ struct server_context {
20122012
}
20132013
slot.params.n_keep = std::min(slot.n_ctx - 4, slot.params.n_keep);
20142014

2015-
// if input prompt is too big, truncate it (if group attention self-extend is disabled)
2015+
// if input prompt is too big, truncate it
20162016
if (slot.n_prompt_tokens >= slot.n_ctx) {
20172017
const int n_left = slot.n_ctx - slot.params.n_keep;
20182018

@@ -2042,12 +2042,74 @@ struct server_context {
20422042

20432043
if (slot.params.cache_prompt) {
20442044
// reuse any previously computed tokens that are common with the new prompt
2045-
slot.n_past = common_part(slot.cache_tokens, prompt_tokens);
2045+
slot.n_past = longest_common_prefix(slot.cache_tokens, prompt_tokens);
20462046

20472047
// push the prompt into the sampling context (do not apply grammar)
20482048
for (int i = 0; i < slot.n_past; ++i) {
20492049
common_sampler_accept(slot.smpl, slot.cache_tokens[i], false);
20502050
}
2051+
2052+
// reuse chunks from the cached prompt by shifting their KV cache in the new position
2053+
if (params.n_cache_reuse > 0) {
2054+
size_t head_c = slot.n_past; // cache
2055+
size_t head_p = slot.n_past; // current prompt
2056+
2057+
SLT_DBG(slot, "trying to reuse chunks with size > %d, slot.n_past = %d\n", params.n_cache_reuse, slot.n_past);
2058+
2059+
while (head_c < slot.cache_tokens.size() &&
2060+
head_p < prompt_tokens.size()) {
2061+
if (llama_token_is_control(model, slot.cache_tokens[head_c])) {
2062+
break;
2063+
}
2064+
2065+
if (llama_token_is_control(model, prompt_tokens[head_p])) {
2066+
break;
2067+
}
2068+
2069+
size_t n_match = 0;
2070+
2071+
while (head_c + n_match < slot.cache_tokens.size() &&
2072+
head_p + n_match < prompt_tokens.size() &&
2073+
slot.cache_tokens[head_c + n_match] == prompt_tokens[head_p + n_match]) {
2074+
if (llama_token_is_control(model, slot.cache_tokens[head_c + n_match])) {
2075+
break;
2076+
}
2077+
2078+
if (llama_token_is_control(model, prompt_tokens[head_p + n_match])) {
2079+
break;
2080+
}
2081+
2082+
n_match++;
2083+
}
2084+
2085+
if (n_match >= (size_t) params.n_cache_reuse) {
2086+
SLT_DBG(slot, "reusing chunk with size %zu, shifting KV cache [%zu, %zu) -> [%zu, %zu)\n", n_match, head_c, head_c + n_match, head_p, head_p + n_match);
2087+
//for (size_t i = head_p; i < head_p + n_match; i++) {
2088+
// SLT_DBG(slot, "cache token %3zu: %6d '%s'\n", i, prompt_tokens[i], common_token_to_piece(ctx, prompt_tokens[i]).c_str());
2089+
//}
2090+
2091+
const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c;
2092+
2093+
llama_kv_cache_seq_rm (ctx, slot.id + 1, head_p, head_c);
2094+
llama_kv_cache_seq_add(ctx, slot.id + 1, head_c, -1, kv_shift);
2095+
2096+
for (size_t i = 0; i < n_match; i++) {
2097+
slot.cache_tokens[head_p + i] = slot.cache_tokens[head_c + i];
2098+
2099+
common_sampler_accept(slot.smpl, slot.cache_tokens[head_p + i], false);
2100+
2101+
slot.n_past++;
2102+
}
2103+
2104+
head_c += n_match;
2105+
head_p += n_match;
2106+
} else {
2107+
head_c += 1;
2108+
}
2109+
}
2110+
2111+
SLT_DBG(slot, "after context reuse, new slot.n_past = %d\n", slot.n_past);
2112+
}
20512113
}
20522114
}
20532115

@@ -3257,6 +3319,7 @@ int main(int argc, char ** argv) {
32573319

32583320
ctx_server.queue_tasks.on_new_task(std::bind(
32593321
&server_context::process_single_task, &ctx_server, std::placeholders::_1));
3322+
32603323
ctx_server.queue_tasks.on_update_slots(std::bind(
32613324
&server_context::update_slots, &ctx_server));
32623325

‎examples/server/utils.hpp

+2-2
Original file line numberDiff line numberDiff line change
@@ -195,14 +195,14 @@ static std::string gen_chatcmplid() {
195195
// other common utils
196196
//
197197

198-
static size_t common_part(const std::vector<llama_token> & a, const std::vector<llama_token> & b) {
198+
static size_t longest_common_prefix(const std::vector<llama_token> & a, const std::vector<llama_token> & b) {
199199
size_t i;
200200
for (i = 0; i < a.size() && i < b.size() && a[i] == b[i]; i++) {}
201201

202202
return i;
203203
}
204204

205-
static size_t common_part(const std::string & a, const std::string & b) {
205+
static size_t longest_common_prefix(const std::string & a, const std::string & b) {
206206
size_t i;
207207
for (i = 0; i < a.size() && i < b.size() && a[i] == b[i]; i++) {}
208208

0 commit comments

Comments
 (0)