CS348Project/llama.cpp
5
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-11-04 09:32:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

llama : add struct llama_kv_cache (wip) [no ci]

Browse Source
This commit is contained in:
Georgi Gerganov
2025-01-13 14:13:11 +02:00
parent 178a7eb952
commit f78b396ee7
8 changed files with 428 additions and 415 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
common/common.cpp
View File

@@ -952,7 +952,9 @@ struct common_init_result common_init_from_params(common_params & params) {
return iparams;
}
if (params.ctx_shift && !llama_kv_cache_can_shift(lctx)) {
llama_kv_cache * kv = llama_get_kv_cache(lctx);
if (params.ctx_shift && !llama_kv_cache_can_shift(kv)) {
LOG_WRN("%s: KV cache shifting is not supported for this model, disabling KV cache shifting\n", __func__);
params.ctx_shift = false;
}
@@ -1057,7 +1059,7 @@ struct common_init_result common_init_from_params(common_params & params) {
if (llama_model_has_decoder(model)) {
llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch)));
}
llama_kv_cache_clear(lctx);
llama_kv_cache_clear(kv);
llama_synchronize(lctx);
llama_perf_context_reset(lctx);
}

10
common/speculative.cpp
View File

@@ -171,8 +171,10 @@ llama_tokens common_speculative_gen_draft(
llama_tokens result;
result.reserve(params.n_draft);
llama_kv_cache * kv = llama_get_kv_cache(ctx);
if (reuse_n == 0) {
llama_kv_cache_clear(ctx);
llama_kv_cache_clear(kv);
prompt.clear();
} else {
@@ -191,14 +193,14 @@ llama_tokens common_speculative_gen_draft(
}
if (reuse_i > 0) {
llama_kv_cache_seq_rm (ctx, 0, 0, reuse_i);
llama_kv_cache_seq_add(ctx, 0, reuse_i, -1, -reuse_i);
llama_kv_cache_seq_rm (kv, 0, 0, reuse_i);
llama_kv_cache_seq_add(kv, 0, reuse_i, -1, -reuse_i);
prompt.erase(prompt.begin(), prompt.begin() + reuse_i);
}
if (reuse_n < (int) prompt.size()) {
llama_kv_cache_seq_rm (ctx, 0, reuse_n, -1);
llama_kv_cache_seq_rm (kv, 0, reuse_n, -1);
prompt.erase(prompt.begin() + reuse_n, prompt.end());
}

5
examples/embedding/embedding.cpp
View File

@@ -34,10 +34,11 @@ static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & toke
static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd, int embd_norm) {
const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);
const struct llama_model * model = llama_get_model(ctx);
const llama_model * model = llama_get_model(ctx);
llama_kv_cache * kv = llama_get_kv_cache(ctx);
// clear previous kv_cache values (irrelevant for embeddings)
llama_kv_cache_clear(ctx);
llama_kv_cache_clear(kv);
// run model
LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);

79
include/llama.h
View File

@@ -60,6 +60,7 @@ extern "C" {
struct llama_model;
struct llama_context;
struct llama_sampler;
struct llama_kv_cache;
typedef int32_t llama_pos;
typedef int32_t llama_token;
@@ -467,8 +468,9 @@ extern "C" {
DEPRECATED(LLAMA_API int32_t llama_n_vocab (const struct llama_vocab * vocab), "use llama_vocab_n_tokens instead");
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx);
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx);
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx); // TODO: remove const?
LLAMA_API struct llama_kv_cache * llama_get_kv_cache( struct llama_context * ctx);
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx);
LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
LLAMA_API enum llama_rope_type llama_model_rope_type(const struct llama_model * model);
@@ -584,7 +586,7 @@ extern "C" {
// KV cache
//
// TODO: remove llama_kv_cache_view_* API
// TODO: start using struct llama_kv_cache
// Information associated with an individual cell in the KV cache view.
struct llama_kv_cache_view_cell {
@@ -639,14 +641,20 @@ extern "C" {
// Returns the number of tokens in the KV cache (slow, use only for debug)
// If a KV cell has multiple sequences assigned to it, it will be counted multiple times
LLAMA_API int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx);
LLAMA_API int32_t llama_kv_cache_n_tokens(const struct llama_kv_cache * kv);
DEPRECATED(LLAMA_API int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx),
"use llama_kv_cache_n_tokens instead");
// Returns the number of used KV cells (i.e. have at least one sequence assigned to them)
LLAMA_API int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx);
LLAMA_API int32_t llama_kv_cache_used_cells(const struct llama_kv_cache * kv);
DEPRECATED(LLAMA_API int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx),
"use llama_kv_cache_used_cells instead");
// Clear the KV cache - both cell info is erased and KV data is zeroed
LLAMA_API void llama_kv_cache_clear(
struct llama_context * ctx);
struct llama_kv_cache * kv);
// Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
// Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails
@@ -654,26 +662,26 @@ extern "C" {
// p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf)
LLAMA_API bool llama_kv_cache_seq_rm(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1);
struct llama_kv_cache * kv,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1);
// Copy all tokens that belong to the specified sequence to another sequence
// Note that this does not allocate extra KV cache memory - it simply assigns the tokens to the new sequence
// p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf)
LLAMA_API void llama_kv_cache_seq_cp(
struct llama_context * ctx,
llama_seq_id seq_id_src,
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1);
struct llama_kv_cache * kv,
llama_seq_id seq_id_src,
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1);
// Removes all tokens that do not belong to the specified sequence
LLAMA_API void llama_kv_cache_seq_keep(
struct llama_context * ctx,
llama_seq_id seq_id);
struct llama_kv_cache * kv,
llama_seq_id seq_id);
// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1)
// If the KV cache is RoPEd, the KV data is updated accordingly:
@@ -682,11 +690,11 @@ extern "C" {
// p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf)
LLAMA_API void llama_kv_cache_seq_add(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta);
struct llama_kv_cache * kv,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta);
// Integer division of the positions by factor of `d > 1`
// If the KV cache is RoPEd, the KV data is updated accordingly:
@@ -695,31 +703,28 @@ extern "C" {
// p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf)
LLAMA_API void llama_kv_cache_seq_div(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d);
struct llama_kv_cache * kv,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d);
// Returns the largest position present in the KV cache for the specified sequence
LLAMA_API llama_pos llama_kv_cache_seq_pos_max(
struct llama_context * ctx,
llama_seq_id seq_id);
// TODO: the llama_kv_cache_defrag and llama_kv_cache_update API tightly couples llama_context with llama_kv_cache
// how to avoid this?
struct llama_kv_cache * kv,
llama_seq_id seq_id);
// Defragment the KV cache
// This will be applied:
// - lazily on next llama_decode()
// - explicitly with llama_kv_cache_update()
LLAMA_API void llama_kv_cache_defrag(struct llama_context * ctx);
// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
LLAMA_API void llama_kv_cache_update(struct llama_context * ctx);
LLAMA_API void llama_kv_cache_defrag(struct llama_kv_cache * kv);
// Check if the context supports KV cache shifting
LLAMA_API bool llama_kv_cache_can_shift(struct llama_context * ctx);
LLAMA_API bool llama_kv_cache_can_shift(const struct llama_kv_cache * kv);
// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
LLAMA_API void llama_update_kv_cache(struct llama_context * ctx, struct llama_kv_cache * kv);
//
// State / sessions

16
src/llama-context.cpp
View File

@@ -602,11 +602,15 @@ uint32_t llama_n_seq_max(const struct llama_context * ctx) {
return ctx->kv_self.size;
}
const struct llama_model * llama_get_model(const struct llama_context * ctx) {
const llama_model * llama_get_model(const llama_context * ctx) {
return &ctx->model;
}
enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx) {
llama_kv_cache * llama_get_kv_cache(llama_context * ctx) {
return &ctx->kv_self;
}
enum llama_pooling_type llama_pooling_type(const llama_context * ctx) {
return ctx->cparams.pooling_type;
}
@@ -1142,7 +1146,7 @@ struct llama_data_read {
if (dest_seq_id != -1) {
// single sequence
llama_kv_cache_seq_rm(kv_self, dest_seq_id, -1, -1);
kv_self.seq_rm(dest_seq_id, -1, -1);
llama_ubatch batch = ctx->sbatch.reserve_ubatch(cell_count, /* has_embd */ false);
batch.n_tokens = cell_count;
@@ -1185,7 +1189,7 @@ struct llama_data_read {
return false;
}
llama_kv_cache_clear(kv_self);
kv_self.clear();
for (uint32_t i = 0; i < cell_count; ++i) {
llama_kv_cell & cell = kv_self.cells[i];
@@ -1362,9 +1366,9 @@ struct llama_data_read {
if (!res) {
if (seq_id == -1) {
llama_kv_cache_clear(ctx);
ctx->kv_self.clear();
} else {
llama_kv_cache_seq_rm(ctx, seq_id, -1, -1);
ctx->kv_self.seq_rm(seq_id, -1, -1);
}
throw std::runtime_error("failed to restore kv cache");
}

286
src/llama-kv-cache.cpp
View File

@@ -350,277 +350,67 @@ uint32_t llama_kv_cache_cell_max(const struct llama_kv_cache & cache) {
return 0;
}
void llama_kv_cache_clear(struct llama_kv_cache & cache) {
for (int32_t i = 0; i < (int32_t) cache.size; ++i) {
cache.cells[i].pos = -1;
cache.cells[i].seq_id.clear();
cache.cells[i].src = -1;
cache.cells[i].tail = -1;
}
cache.head = 0;
cache.used = 0;
for (auto & buf : cache.bufs) {
ggml_backend_buffer_clear(buf.get(), 0);
}
void llama_kv_cache_clear(llama_kv_cache * kv) {
kv->clear();
}
bool llama_kv_cache_seq_rm(
struct llama_kv_cache & cache,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1) {
uint32_t new_head = cache.size;
if (p0 < 0) p0 = 0;
if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
// models like Mamba or RWKV can't have a state partially erased
if (cache.recurrent) {
if (seq_id >= (int64_t) cache.size) {
// could be fatal
return false;
}
if (0 <= seq_id) {
int32_t & tail_id = cache.cells[seq_id].tail;
if (tail_id >= 0) {
const llama_kv_cell & cell = cache.cells[tail_id];
// partial intersection is invalid
if ((0 < p0 && p0 <= cell.pos) || (0 < p1 && p1 <= cell.pos)) {
return false;
}
// invalidate tails which will be cleared
if (p0 <= cell.pos && cell.pos < p1) {
tail_id = -1;
}
}
} else {
// seq_id is negative, then the range should include everything or nothing
if (p0 != p1 && (p0 != 0 || p1 != std::numeric_limits<llama_pos>::max())) {
return false;
}
}
}
for (uint32_t i = 0; i < cache.size; ++i) {
if (cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
if (seq_id < 0) {
cache.cells[i].seq_id.clear();
} else if (cache.cells[i].has_seq_id(seq_id)) {
cache.cells[i].seq_id.erase(seq_id);
} else {
continue;
}
if (cache.cells[i].is_empty()) {
// keep count of the number of used cells
if (cache.cells[i].pos >= 0) cache.used--;
cache.cells[i].pos = -1;
cache.cells[i].src = -1;
if (new_head == cache.size) new_head = i;
}
}
}
// If we freed up a slot, set head to it so searching can start there.
if (new_head != cache.size && new_head < cache.head) cache.head = new_head;
return true;
llama_kv_cache * kv,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1) {
return kv->seq_rm(seq_id, p0, p1);
}
void llama_kv_cache_seq_cp(
struct llama_kv_cache & cache,
llama_seq_id seq_id_src,
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1) {
if (p0 < 0) p0 = 0;
if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
if (cache.recurrent) {
if ((uint32_t) seq_id_dst < cache.size && (uint32_t) seq_id_src < cache.size) {
llama_kv_cell & tail_src = cache.cells[seq_id_src];
llama_kv_cell & tail_dst = cache.cells[seq_id_dst];
if (tail_dst.tail >= 0) {
// clear destination seq_id if it wasn't empty
llama_kv_cell & cell_dst = cache.cells[tail_dst.tail];
cell_dst.seq_id.erase(seq_id_dst);
tail_dst.tail = -1;
if (cell_dst.seq_id.empty()) {
cell_dst.pos = -1;
cell_dst.delta = -1;
cell_dst.src = -1;
cache.used -= 1;
}
}
if (tail_src.tail >= 0) {
llama_kv_cell & cell_src = cache.cells[tail_src.tail];
cell_src.seq_id.insert(seq_id_dst);
tail_dst.tail = tail_src.tail;
}
}
return;
}
// otherwise, this is the KV cache of a Transformer-like model
cache.head = 0;
for (uint32_t i = 0; i < cache.size; ++i) {
if (cache.cells[i].has_seq_id(seq_id_src) && cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
cache.cells[i].seq_id.insert(seq_id_dst);
}
}
llama_kv_cache * kv,
llama_seq_id seq_id_src,
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1) {
kv->seq_cp(seq_id_src, seq_id_dst, p0, p1);
}
void llama_kv_cache_seq_keep(struct llama_kv_cache & cache, llama_seq_id seq_id) {
uint32_t new_head = cache.size;
for (uint32_t i = 0; i < cache.size; ++i) {
if (cache.recurrent && (llama_seq_id) i != seq_id) {
cache.cells[i].tail = -1;
}
if (!cache.cells[i].has_seq_id(seq_id)) {
if (cache.cells[i].pos >= 0) cache.used--;
cache.cells[i].pos = -1;
cache.cells[i].src = -1;
cache.cells[i].seq_id.clear();
if (new_head == cache.size) new_head = i;
} else {
cache.cells[i].seq_id.clear();
cache.cells[i].seq_id.insert(seq_id);
}
}
// If we freed up a slot, set head to it so searching can start there.
if (new_head != cache.size && new_head < cache.head) cache.head = new_head;
void llama_kv_cache_seq_keep(llama_kv_cache * kv, llama_seq_id seq_id) {
kv->seq_keep(seq_id);
}
void llama_kv_cache_seq_add(
struct llama_kv_cache & cache,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta) {
uint32_t new_head = cache.size;
if (p0 < 0) p0 = 0;
if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
// If there is no range then return early to avoid looping over the cache.
if (p0 == p1) return;
if (cache.recurrent) {
// for Mamba-like or RWKV models, only the pos needs to be shifted
if (0 <= seq_id && seq_id < (int64_t) cache.size) {
const int32_t tail_id = cache.cells[seq_id].tail;
if (tail_id >= 0) {
llama_kv_cell & cell = cache.cells[tail_id];
if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {
cell.pos += delta;
}
}
}
return;
}
for (uint32_t i = 0; i < cache.size; ++i) {
if (cache.cells[i].has_seq_id(seq_id) && cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
cache.has_shift = true;
cache.cells[i].pos += delta;
cache.cells[i].delta += delta;
if (cache.cells[i].pos < 0) {
if (!cache.cells[i].is_empty()) {
cache.used--;
}
cache.cells[i].pos = -1;
cache.cells[i].seq_id.clear();
if (new_head == cache.size) {
new_head = i;
}
}
}
}
// If we freed up a slot, set head to it so searching can start there.
// Otherwise we just start the next search from the beginning.
cache.head = new_head != cache.size ? new_head : 0;
llama_kv_cache * kv,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta) {
kv->seq_add(seq_id, p0, p1, delta);
}
void llama_kv_cache_seq_div(
struct llama_kv_cache & cache,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d) {
if (p0 < 0) p0 = 0;
if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
// If there is no range then return early to avoid looping over the cache.
if (p0 == p1) return;
if (cache.recurrent) {
// for Mamba-like or RWKV models, only the pos needs to be changed
if (0 <= seq_id && seq_id < (int64_t) cache.size) {
const int32_t tail_id = cache.cells[seq_id].tail;
if (tail_id >= 0) {
llama_kv_cell & cell = cache.cells[tail_id];
if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {
cell.pos /= d;
}
}
}
return;
}
for (uint32_t i = 0; i < cache.size; ++i) {
if (cache.cells[i].has_seq_id(seq_id) && cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
cache.has_shift = true;
{
llama_pos p_old = cache.cells[i].pos;
cache.cells[i].pos /= d;
cache.cells[i].delta += cache.cells[i].pos - p_old;
}
}
}
llama_kv_cache * kv,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d) {
kv->seq_div(seq_id, p0, p1, d);
}
llama_pos llama_kv_cache_seq_pos_max(struct llama_kv_cache & cache, llama_seq_id seq_id) {
llama_pos result = 0;
for (uint32_t i = 0; i < cache.size; ++i) {
if (cache.cells[i].has_seq_id(seq_id)) {
result = std::max(result, cache.cells[i].pos);
}
}
return result;
llama_pos llama_kv_cache_seq_pos_max(llama_kv_cache * kv, llama_seq_id seq_id) {
return kv->seq_pos_max(seq_id);
}
void llama_kv_cache_defrag(struct llama_kv_cache & cache) {
if (!cache.recurrent) {
cache.do_defrag = true;
}
void llama_kv_cache_defrag(llama_kv_cache * kv) {
kv->defrag();
}
int32_t llama_get_kv_cache_token_count(const struct llama_kv_cache & kv) {
int result = 0;
for (uint32_t i = 0; i < kv.size; i++) {
result += kv.cells[i].seq_id.size();
}
return result;
int32_t llama_kv_cache_n_tokens(const llama_kv_cache * kv) {
return kv->n_tokens();
}
int32_t llama_get_kv_cache_used_cells(const struct llama_kv_cache & kv) {
return kv.used;
int32_t llama_kv_cache_used_cells(const llama_kv_cache * kv) {
return kv->used;
}
bool llama_kv_cache_can_shift(const struct llama_kv_cache & kv) {
return kv.can_shift;
bool llama_kv_cache_can_shift(const llama_kv_cache * kv) {
return kv->can_shift;
}
//
@@ -632,7 +422,7 @@ struct llama_kv_cache_view llama_kv_cache_view_init(const struct llama_kv_cache
/*.n_cells = */ 0,
/*.n_seq_max = */ n_seq_max,
/*.token_count = */ 0,
/*.used_cells = */ llama_get_kv_cache_used_cells(kv),
/*.used_cells = */ llama_kv_cache_used_cells(&kv),
/*.max_contiguous = */ 0,
/*.max_contiguous_idx = */ -1,
/*.cells = */ nullptr,

350
src/llama-kv-cache.h
View File

@@ -57,6 +57,16 @@ struct llama_kv_cache {
std::vector<ggml_context_ptr> ctxs;
std::vector<ggml_backend_buffer_ptr> bufs;
int32_t n_tokens() const {
int32_t result = 0;
for (uint32_t i = 0; i < size; i++) {
result += cells[i].seq_id.size();
}
return result;
}
size_t total_size() const {
size_t size = 0;
for (const auto & buf : bufs) {
@@ -75,6 +85,297 @@ struct llama_kv_cache {
return max_pos;
}
void clear() {
for (int32_t i = 0; i < (int32_t) size; ++i) {
cells[i].pos = -1;
cells[i].seq_id.clear();
cells[i].src = -1;
cells[i].tail = -1;
}
head = 0;
used = 0;
for (auto & buf : bufs) {
ggml_backend_buffer_clear(buf.get(), 0);
}
}
bool seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
uint32_t new_head = size;
if (p0 < 0) {
p0 = 0;
}
if (p1 < 0) {
p1 = std::numeric_limits<llama_pos>::max();
}
// models like Mamba or RWKV can't have a state partially erased
if (recurrent) {
if (seq_id >= (int64_t) size) {
// could be fatal
return false;
}
if (0 <= seq_id) {
int32_t & tail_id = cells[seq_id].tail;
if (tail_id >= 0) {
const llama_kv_cell & cell = cells[tail_id];
// partial intersection is invalid
if ((0 < p0 && p0 <= cell.pos) || (0 < p1 && p1 <= cell.pos)) {
return false;
}
// invalidate tails which will be cleared
if (p0 <= cell.pos && cell.pos < p1) {
tail_id = -1;
}
}
} else {
// seq_id is negative, then the range should include everything or nothing
if (p0 != p1 && (p0 != 0 || p1 != std::numeric_limits<llama_pos>::max())) {
return false;
}
}
}
for (uint32_t i = 0; i < size; ++i) {
if (cells[i].pos >= p0 && cells[i].pos < p1) {
if (seq_id < 0) {
cells[i].seq_id.clear();
} else if (cells[i].has_seq_id(seq_id)) {
cells[i].seq_id.erase(seq_id);
} else {
continue;
}
if (cells[i].is_empty()) {
// keep count of the number of used cells
if (cells[i].pos >= 0) {
used--;
}
cells[i].pos = -1;
cells[i].src = -1;
if (new_head == size) {
new_head = i;
}
}
}
}
// If we freed up a slot, set head to it so searching can start there.
if (new_head != size && new_head < head) {
head = new_head;
}
return true;
}
void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
if (seq_id_src == seq_id_dst) {
return;
}
if (p0 < 0) {
p0 = 0;
}
if (p1 < 0) {
p1 = std::numeric_limits<llama_pos>::max();
}
if (recurrent) {
if ((uint32_t) seq_id_dst < size && (uint32_t) seq_id_src < size) {
llama_kv_cell & tail_src = cells[seq_id_src];
llama_kv_cell & tail_dst = cells[seq_id_dst];
if (tail_dst.tail >= 0) {
// clear destination seq_id if it wasn't empty
llama_kv_cell & cell_dst = cells[tail_dst.tail];
cell_dst.seq_id.erase(seq_id_dst);
tail_dst.tail = -1;
if (cell_dst.seq_id.empty()) {
cell_dst.pos = -1;
cell_dst.delta = -1;
cell_dst.src = -1;
used -= 1;
}
}
if (tail_src.tail >= 0) {
llama_kv_cell & cell_src = cells[tail_src.tail];
cell_src.seq_id.insert(seq_id_dst);
tail_dst.tail = tail_src.tail;
}
}
return;
}
// otherwise, this is the KV of a Transformer-like model
head = 0;
for (uint32_t i = 0; i < size; ++i) {
if (cells[i].has_seq_id(seq_id_src) && cells[i].pos >= p0 && cells[i].pos < p1) {
cells[i].seq_id.insert(seq_id_dst);
}
}
}
void seq_keep(llama_seq_id seq_id) {
uint32_t new_head = size;
for (uint32_t i = 0; i < size; ++i) {
if (recurrent && (llama_seq_id) i != seq_id) {
cells[i].tail = -1;
}
if (!cells[i].has_seq_id(seq_id)) {
if (cells[i].pos >= 0) {
used--;
}
cells[i].pos = -1;
cells[i].src = -1;
cells[i].seq_id.clear();
if (new_head == size){
new_head = i;
}
} else {
cells[i].seq_id.clear();
cells[i].seq_id.insert(seq_id);
}
}
// If we freed up a slot, set head to it so searching can start there.
if (new_head != size && new_head < head) {
head = new_head;
}
}
void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
if (delta == 0) {
return;
}
uint32_t new_head = size;
if (p0 < 0) {
p0 = 0;
}
if (p1 < 0) {
p1 = std::numeric_limits<llama_pos>::max();
}
// If there is no range then return early to avoid looping over the
if (p0 == p1) {
return;
}
if (recurrent) {
// for Mamba-like or RWKV models, only the pos needs to be shifted
if (0 <= seq_id && seq_id < (int64_t) size) {
const int32_t tail_id = cells[seq_id].tail;
if (tail_id >= 0) {
llama_kv_cell & cell = cells[tail_id];
if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {
cell.pos += delta;
}
}
}
return;
}
for (uint32_t i = 0; i < size; ++i) {
if (cells[i].has_seq_id(seq_id) && cells[i].pos >= p0 && cells[i].pos < p1) {
has_shift = true;
cells[i].pos += delta;
cells[i].delta += delta;
if (cells[i].pos < 0) {
if (!cells[i].is_empty()) {
used--;
}
cells[i].pos = -1;
cells[i].seq_id.clear();
if (new_head == size) {
new_head = i;
}
}
}
}
// If we freed up a slot, set head to it so searching can start there.
// Otherwise we just start the next search from the beginning.
head = new_head != size ? new_head : 0;
}
void seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
if (d == 1) {
return;
}
if (p0 < 0) {
p0 = 0;
}
if (p1 < 0) {
p1 = std::numeric_limits<llama_pos>::max();
}
// If there is no range then return early to avoid looping over the cache.
if (p0 == p1) {
return;
}
if (recurrent) {
// for Mamba-like or RWKV models, only the pos needs to be changed
if (0 <= seq_id && seq_id < (int64_t) size) {
const int32_t tail_id = cells[seq_id].tail;
if (tail_id >= 0) {
llama_kv_cell & cell = cells[tail_id];
if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {
cell.pos /= d;
}
}
}
return;
}
for (uint32_t i = 0; i < size; ++i) {
if (cells[i].has_seq_id(seq_id) && cells[i].pos >= p0 && cells[i].pos < p1) {
has_shift = true;
{
llama_pos p_old = cells[i].pos;
cells[i].pos /= d;
cells[i].delta += cells[i].pos - p_old;
}
}
}
}
llama_pos seq_pos_max(llama_seq_id seq_id) {
llama_pos result = 0;
for (uint32_t i = 0; i < size; ++i) {
if (cells[i].has_seq_id(seq_id)) {
result = std::max(result, cells[i].pos);
}
}
return result;
}
void defrag() {
if (!recurrent) {
do_defrag = true;
}
}
};
// a structure holds information about the slot found in llama_kv_cache_find_slot
@@ -112,51 +413,6 @@ struct llama_kv_cache_slot_info llama_kv_cache_find_slot(
// find how many cells are currently in use
uint32_t llama_kv_cache_cell_max(const struct llama_kv_cache & cache);
void llama_kv_cache_clear(struct llama_kv_cache & cache);
bool llama_kv_cache_seq_rm(
struct llama_kv_cache & cache,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1);
void llama_kv_cache_seq_cp(
struct llama_kv_cache & cache,
llama_seq_id seq_id_src,
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1);
void llama_kv_cache_seq_keep(
struct llama_kv_cache & cache,
llama_seq_id seq_id);
void llama_kv_cache_seq_add(
struct llama_kv_cache & cache,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta);
void llama_kv_cache_seq_div(
struct llama_kv_cache & cache,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d);
llama_pos llama_kv_cache_seq_pos_max(
struct llama_kv_cache & cache,
llama_seq_id seq_id);
void llama_kv_cache_defrag(struct llama_kv_cache & cache);
int32_t llama_get_kv_cache_token_count(const struct llama_kv_cache & kv);
int32_t llama_get_kv_cache_used_cells(const struct llama_kv_cache & kv);
bool llama_kv_cache_can_shift(const struct llama_kv_cache & kv);
//
// kv cache view
//
@@ -206,10 +462,10 @@ struct llama_kv_slot_restorer {
cache.n = old_state.n;
if (cache.recurrent) { // recurrent models like Mamba or RWKV can't have a state partially erased
llama_kv_cache_seq_rm(cache, -1, -1, -1);
cache.seq_rm(-1, -1, -1);
} else {
for (auto & slot : slot_boundaries) {
llama_kv_cache_seq_rm(cache, -1, slot.first, slot.second);
cache.seq_rm(-1, slot.first, slot.second);
}
}
}

91
src/llama.cpp
View File

@@ -8564,7 +8564,7 @@ static int llama_decode_impl(
// non-causal masks do not use the KV cache
if (hparams.causal_attn) {
llama_kv_cache_update(&lctx);
llama_update_kv_cache(&lctx, &lctx.kv_self); // TODO: lctx->update_kv_cache()
// if we have enough unused cells before the current head ->
// better to start searching from the beginning of the cache, hoping to fill it
@@ -8760,7 +8760,7 @@ static int llama_decode_impl(
if (fragmentation > cparams.defrag_thold) {
//LLAMA_LOG_INFO("fragmentation: %.2f\n", fragmentation);
llama_kv_cache_defrag(kv_self);
kv_self.defrag();
}
}
@@ -9182,11 +9182,11 @@ static void llama_kv_cache_defrag_impl(struct llama_context & lctx) {
//LLAMA_LOG_INFO("(tmp log) KV defrag time: %.3f ms\n", (t_end - t_start)/1000.0);
}
static void llama_kv_cache_update_impl(struct llama_context & lctx) {
static void llama_update_kv_cache_impl(llama_context & lctx, llama_kv_cache & kv) {
bool need_reserve = false;
if (lctx.kv_self.has_shift) {
if (!llama_kv_cache_can_shift(&lctx)) {
if (kv.has_shift) {
if (!kv.can_shift) {
GGML_ABORT("The current context does not support K-shift");
}
@@ -9206,23 +9206,21 @@ static void llama_kv_cache_update_impl(struct llama_context & lctx) {
}
{
auto & kv_self = lctx.kv_self;
kv.has_shift = false;
kv_self.has_shift = false;
for (uint32_t i = 0; i < kv_self.size; ++i) {
kv_self.cells[i].delta = 0;
for (uint32_t i = 0; i < kv.size; ++i) {
kv.cells[i].delta = 0;
}
}
}
// defragment the KV cache if needed
if (lctx.kv_self.do_defrag) {
if (kv.do_defrag) {
llama_kv_cache_defrag_impl(lctx);
need_reserve = true;
lctx.kv_self.do_defrag = false;
kv.do_defrag = false;
}
// reserve a worst case graph again
@@ -9845,6 +9843,7 @@ struct llama_context * llama_init_from_model(
return ctx;
}
// deprecated
struct llama_context * llama_new_context_with_model(
struct llama_model * model,
struct llama_context_params params) {
@@ -9855,73 +9854,27 @@ struct llama_context * llama_new_context_with_model(
// kv cache
//
// TODO: tmp bridges below until `struct llama_kv_cache` is exposed through the public API
struct llama_kv_cache_view llama_kv_cache_view_init(const struct llama_context * ctx, int32_t n_seq_max) {
struct llama_kv_cache_view llama_kv_cache_view_init(const llama_context * ctx, int32_t n_seq_max) {
return llama_kv_cache_view_init(ctx->kv_self, n_seq_max);
}
void llama_kv_cache_view_update(const struct llama_context * ctx, struct llama_kv_cache_view * view) {
void llama_kv_cache_view_update(const llama_context * ctx, llama_kv_cache_view * view) {
llama_kv_cache_view_update(view, ctx->kv_self);
}
int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx) {
return llama_get_kv_cache_token_count(ctx->kv_self);
// deprecated
int32_t llama_get_kv_cache_token_count(const llama_context * ctx) {
return llama_kv_cache_n_tokens(&ctx->kv_self);
}
int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx) {
return llama_get_kv_cache_used_cells(ctx->kv_self);
// deprecated
int32_t llama_get_kv_cache_used_cells(const llama_context * ctx) {
return llama_kv_cache_used_cells(&ctx->kv_self);
}
void llama_kv_cache_clear(struct llama_context * ctx) {
llama_kv_cache_clear(ctx->kv_self);
}
bool llama_kv_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
return llama_kv_cache_seq_rm(ctx->kv_self, seq_id, p0, p1);
}
void llama_kv_cache_seq_cp(struct llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
if (seq_id_src == seq_id_dst) {
return;
}
llama_kv_cache_seq_cp(ctx->kv_self, seq_id_src, seq_id_dst, p0, p1);
}
void llama_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) {
llama_kv_cache_seq_keep(ctx->kv_self, seq_id);
}
void llama_kv_cache_seq_add(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
if (delta == 0) {
return;
}
llama_kv_cache_seq_add(ctx->kv_self, seq_id, p0, p1, delta);
}
void llama_kv_cache_seq_div(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
if (d == 1) {
return;
}
llama_kv_cache_seq_div(ctx->kv_self, seq_id, p0, p1, d);
}
llama_pos llama_kv_cache_seq_pos_max(struct llama_context * ctx, llama_seq_id seq_id) {
return llama_kv_cache_seq_pos_max(ctx->kv_self, seq_id);
}
void llama_kv_cache_defrag(struct llama_context * ctx) {
llama_kv_cache_defrag(ctx->kv_self);
}
void llama_kv_cache_update(struct llama_context * ctx) {
llama_kv_cache_update_impl(*ctx);
}
bool llama_kv_cache_can_shift(struct llama_context * ctx) {
return llama_kv_cache_can_shift(ctx->kv_self);
// TODO: move to llama-context
void llama_update_kv_cache(llama_context * ctx, llama_kv_cache * kv) {
llama_update_kv_cache_impl(*ctx, *kv);
}
///