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

llama : rename many llama_kv_cache_* functions

Browse Source
This commit is contained in:
Francis Couture-Harpin
2024-04-29 10:24:45 -04:00
parent d66849f628
commit a09db95eab
2 changed files with 138 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
llama.cpp
View File

@@ -2032,7 +2032,6 @@ struct llama_rs_seq_meta {
// ring-buffered tree of cached recurrent state data // ring-buffered tree of cached recurrent state data
struct llama_rs_cache { struct llama_rs_cache {
bool do_copy = false;
uint32_t head = 0; // first state used for the last slot uint32_t head = 0; // first state used for the last slot
uint32_t size = 0; uint32_t size = 0;
@@ -2769,7 +2768,7 @@ struct llama_context {
}; };
// //
// kv cache helpers // kv and rs cache helpers
// //
static bool llama_cache_init( static bool llama_cache_init(
@@ -2898,7 +2897,7 @@ static bool llama_cache_init(
// updates the cache head // updates the cache head
// Note: On success, it's important that cache.head points // Note: On success, it's important that cache.head points
// to the first cell of the slot. // to the first cell of the slot.
static bool llama_kv_cache_find_slot( static bool llama_cache_find_slot(
struct llama_cache & cache, struct llama_cache & cache,
const struct llama_batch & batch) { const struct llama_batch & batch) {
const uint32_t kv_size = cache.kv.size; const uint32_t kv_size = cache.kv.size;
@@ -3181,7 +3180,6 @@ static void llama_cache_clear(struct llama_cache & cache) {
rs_cell.tail_rc = 0; rs_cell.tail_rc = 0;
rs_cell.seq_nodes.clear(); rs_cell.seq_nodes.clear();
} }
cache.rs.do_copy = false;
cache.rs.head = 0; cache.rs.head = 0;
cache.rs.used = 0; cache.rs.used = 0;
cache.rs.n_seqs = 0; cache.rs.n_seqs = 0;
@@ -3412,8 +3410,8 @@ static llama_pos llama_cache_seq_add(
llama_pos p1, llama_pos p1,
llama_pos delta) { llama_pos delta) {
if (p0 < 0) p0 = 0; if (p0 < 0) { p0 = 0; }
if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max(); if (p1 < 0) { p1 = std::numeric_limits<llama_pos>::max(); }
llama_pos n_past = p0; llama_pos n_past = p0;
@@ -3535,7 +3533,7 @@ static llama_pos llama_cache_seq_div(
} }
static llama_pos llama_cache_seq_pos_max(struct llama_cache & cache, llama_seq_id seq_id) { static llama_pos llama_cache_seq_pos_max(struct llama_cache & cache, llama_seq_id seq_id) {
llama_pos result = 0; llama_pos result = -1;
if (cache.rs.size > 0) { if (cache.rs.size > 0) {
int32_t cell_id = cache.rs.seq_tails[seq_id].tail; int32_t cell_id = cache.rs.seq_tails[seq_id].tail;
@@ -11174,7 +11172,7 @@ static int llama_decode_internal(
if (hparams.causal_attn) { if (hparams.causal_attn) {
llama_kv_cache_update(&lctx); llama_kv_cache_update(&lctx);
if (!llama_kv_cache_find_slot(lctx.cache, u_batch)) { if (!llama_cache_find_slot(lctx.cache, u_batch)) {
return 1; return 1;
} }
@@ -15790,6 +15788,10 @@ void llama_kv_cache_view_update(const struct llama_context * ctx, struct llama_k
} }
} }
bool llama_rs_cache_rebuild(struct llama_context * ctx, bool debug) {
return ctx->cache.rs.rebuild(debug);
}
int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx) { int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx) {
int result = 0; int result = 0;
@@ -15804,55 +15806,96 @@ int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx) {
return ctx->cache.kv.used; return ctx->cache.kv.used;
} }
void llama_kv_cache_clear(struct llama_context * ctx) { int32_t llama_get_rs_cache_used_cells(const struct llama_context * ctx) {
return ctx->cache.rs.used;
}
void llama_cache_clear(struct llama_context * ctx) {
llama_cache_clear(ctx->cache); llama_cache_clear(ctx->cache);
} }
// deprecated
void llama_kv_cache_clear(struct llama_context * ctx) {
llama_cache_clear(ctx);
}
llama_pos llama_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return 0; }
return llama_cache_seq_rm(ctx->cache, seq_id, p0, p1);
}
// deprecated
bool llama_kv_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) { bool llama_kv_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return false; } llama_pos n_past = llama_cache_seq_rm(ctx, seq_id, p0, p1);
llama_pos n_past = llama_cache_seq_rm(ctx->cache, seq_id, p0, p1);
return n_past >= p0; return n_past >= p0;
} }
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) {
llama_pos llama_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) {
uint32_t n_seq_max = llama_n_seq_max(ctx); uint32_t n_seq_max = llama_n_seq_max(ctx);
if (seq_id_src < 0 || seq_id_dst < 0 || (uint32_t) seq_id_src >= n_seq_max || (uint32_t) seq_id_dst >= n_seq_max) { if (seq_id_src < 0 || seq_id_dst < 0 || (uint32_t) seq_id_src >= n_seq_max || (uint32_t) seq_id_dst >= n_seq_max) {
return; return 0;
} }
if (seq_id_src == seq_id_dst) { if (seq_id_src == seq_id_dst) {
return; return llama_cache_seq_pos_max(ctx->cache, seq_id_dst) + 1;
} }
llama_cache_seq_cp(ctx->cache, seq_id_src, seq_id_dst, p0, p1); return llama_cache_seq_cp(ctx->cache, seq_id_src, seq_id_dst, p0, p1);
} }
void llama_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) { // deprecated
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) {
llama_cache_seq_cp(ctx, seq_id_src, seq_id_dst, p0, p1);
}
void llama_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) {
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return; } if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return; }
llama_cache_seq_keep(ctx->cache, seq_id); llama_cache_seq_keep(ctx->cache, 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) { // deprecated
if (delta == 0) { void llama_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) {
return; llama_cache_seq_keep(ctx, seq_id);
}
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return; }
llama_cache_seq_add(ctx->cache, 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) { llama_pos llama_cache_seq_add(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
if (d == 1) {
return;
}
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return; }
llama_cache_seq_div(ctx->cache, seq_id, p0, p1, d);
}
llama_pos llama_kv_cache_seq_pos_max(struct llama_context * ctx, llama_seq_id seq_id) {
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return 0; } if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return 0; }
if (delta == 0) {
return llama_cache_seq_pos_max(ctx->cache, seq_id) + 1;
}
return llama_cache_seq_add(ctx->cache, seq_id, p0, p1, delta);
}
// deprecated
void llama_kv_cache_seq_add(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
llama_cache_seq_add(ctx, seq_id, p0, p1, delta);
}
llama_pos llama_cache_seq_div(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return 0; }
if (d == 1) {
return llama_cache_seq_pos_max(ctx->cache, seq_id) + 1;
}
return llama_cache_seq_div(ctx->cache, seq_id, p0, p1, d);
}
// deprecated
void llama_kv_cache_seq_div(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
llama_cache_seq_div(ctx, seq_id, p0, p1, d);
}
llama_pos llama_cache_seq_pos_max(struct llama_context * ctx, llama_seq_id seq_id) {
if (seq_id < 0 || (uint32_t) seq_id >= llama_n_seq_max(ctx)) { return -1; }
return llama_cache_seq_pos_max(ctx->cache, seq_id); return llama_cache_seq_pos_max(ctx->cache, seq_id);
} }
// deprecated
llama_pos llama_kv_cache_seq_pos_max(struct llama_context * ctx, llama_seq_id seq_id) {
llama_pos max_pos = llama_cache_seq_pos_max(ctx, seq_id);
return max_pos < 0 ? 0 : max_pos;
}
void llama_kv_cache_defrag(struct llama_context * ctx) { void llama_kv_cache_defrag(struct llama_context * ctx) {
llama_kv_cache_defrag(ctx->cache.kv); llama_kv_cache_defrag(ctx->cache.kv);
} }
@@ -16597,7 +16640,7 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src,
batch.n_seq_id[i] = 1; batch.n_seq_id[i] = 1;
batch.seq_id[i][0] = dest_seq_id; batch.seq_id[i][0] = dest_seq_id;
} }
if (!llama_kv_cache_find_slot(cache, batch)) { if (!llama_cache_find_slot(cache, batch)) {
llama_batch_free(batch); llama_batch_free(batch);
LLAMA_LOG_ERROR("%s: failed to find available cells in kv cache\n", __func__); LLAMA_LOG_ERROR("%s: failed to find available cells in kv cache\n", __func__);
return 0; return 0;

72
llama.h
View File

@@ -515,6 +515,12 @@ extern "C" {
// Update the KV cache view structure with the current state of the KV cache. (use only for debugging purposes) // Update the KV cache view structure with the current state of the KV cache. (use only for debugging purposes)
LLAMA_API void llama_kv_cache_view_update(const struct llama_context * ctx, struct llama_kv_cache_view * view); LLAMA_API void llama_kv_cache_view_update(const struct llama_context * ctx, struct llama_kv_cache_view * view);
// Rebuild and check the validity of the recurrent state cache's tree of sequences.
// (slow, use only for debugging purposes)
// Returns whether or not the rs cache was valid.
// The errors are always corrected, but only logged when debug is true.
LLAMA_API bool llama_rs_cache_rebuild(struct llama_context * ctx, bool debug);
// Returns the number of tokens in the KV cache (slow, use only for debug) // 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 // 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_get_kv_cache_token_count(const struct llama_context * ctx);
@@ -522,36 +528,60 @@ extern "C" {
// Returns the number of used KV cells (i.e. have at least one sequence assigned to them) // 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_get_kv_cache_used_cells(const struct llama_context * ctx);
// Clear the KV cache // Returns the number of used recurrent state cells (i.e. have at least one sequence assigned to them)
LLAMA_API void llama_kv_cache_clear( LLAMA_API int32_t llama_get_rs_cache_used_cells(const struct llama_context * ctx);
// Clear the KV and recurrent state caches
LLAMA_API void llama_cache_clear(
struct llama_context * ctx); struct llama_context * ctx);
LLAMA_API DEPRECATED(void llama_kv_cache_clear(
struct llama_context * ctx),
"use llama_cache_clear instead");
// Removes all tokens that belong to the specified sequence and have positions in [p0, p1) // 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
// seq_id < 0 : match any sequence // seq_id < 0 : match any sequence
// p0 < 0 : [0, p1] // p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf) // p1 < 0 : [p0, inf)
LLAMA_API bool llama_kv_cache_seq_rm( // Returns n_past
LLAMA_API llama_pos llama_cache_seq_rm(
struct llama_context * ctx, struct llama_context * ctx,
llama_seq_id seq_id, llama_seq_id seq_id,
llama_pos p0, llama_pos p0,
llama_pos p1); llama_pos p1);
LLAMA_API DEPRECATED(bool llama_kv_cache_seq_rm(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1),
"use llama_cache_seq_rm instead, and handle its return value for partial removals");
// Copy all tokens that belong to the specified sequence to another sequence // 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 // Note that this does not allocate extra KV or RS cache memory - it simply assigns the tokens to the new sequence
// p0 < 0 : [0, p1] // p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf) // p1 < 0 : [p0, inf)
LLAMA_API void llama_kv_cache_seq_cp( // Returns n_past
LLAMA_API llama_pos llama_cache_seq_cp(
struct llama_context * ctx, struct llama_context * ctx,
llama_seq_id seq_id_src, llama_seq_id seq_id_src,
llama_seq_id seq_id_dst, llama_seq_id seq_id_dst,
llama_pos p0, llama_pos p0,
llama_pos p1); llama_pos p1);
LLAMA_API DEPRECATED(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),
"use llama_cache_seq_cp instead, and handle its return value for partial copies");
// Removes all tokens that do not belong to the specified sequence // Removes all tokens that do not belong to the specified sequence
LLAMA_API void llama_kv_cache_seq_keep( LLAMA_API void llama_cache_seq_keep(
struct llama_context * ctx, struct llama_context * ctx,
llama_seq_id seq_id); llama_seq_id seq_id);
LLAMA_API DEPRECATED(void llama_kv_cache_seq_keep(
struct llama_context * ctx,
llama_seq_id seq_id),
"use llama_cache_seq_keep instead");
// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1) // 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: // If the KV cache is RoPEd, the KV data is updated accordingly:
@@ -559,12 +589,20 @@ extern "C" {
// - explicitly with llama_kv_cache_update() // - explicitly with llama_kv_cache_update()
// p0 < 0 : [0, p1] // p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf) // p1 < 0 : [p0, inf)
LLAMA_API void llama_kv_cache_seq_add( // Returns n_past
LLAMA_API llama_pos llama_cache_seq_add(
struct llama_context * ctx, struct llama_context * ctx,
llama_seq_id seq_id, llama_seq_id seq_id,
llama_pos p0, llama_pos p0,
llama_pos p1, llama_pos p1,
llama_pos delta); llama_pos delta);
LLAMA_API DEPRECATED(void llama_kv_cache_seq_add(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta),
"use llama_cache_seq_add instead");
// Integer division of the positions by factor of `d > 1` // Integer division of the positions by factor of `d > 1`
// If the KV cache is RoPEd, the KV data is updated accordingly: // If the KV cache is RoPEd, the KV data is updated accordingly:
@@ -572,17 +610,29 @@ extern "C" {
// - explicitly with llama_kv_cache_update() // - explicitly with llama_kv_cache_update()
// p0 < 0 : [0, p1] // p0 < 0 : [0, p1]
// p1 < 0 : [p0, inf) // p1 < 0 : [p0, inf)
LLAMA_API void llama_kv_cache_seq_div( // Returns n_past
LLAMA_API llama_pos llama_cache_seq_div(
struct llama_context * ctx, struct llama_context * ctx,
llama_seq_id seq_id, llama_seq_id seq_id,
llama_pos p0, llama_pos p0,
llama_pos p1, llama_pos p1,
int d); int d);
LLAMA_API DEPRECATED(void llama_kv_cache_seq_div(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d),
"use llama_cache_seq_div instead");
// Returns the largest position present in the KV cache for the specified sequence // Returns the largest position present in the KV and/or RS cache for the specified sequence
LLAMA_API llama_pos llama_kv_cache_seq_pos_max( LLAMA_API llama_pos llama_cache_seq_pos_max(
struct llama_context * ctx, struct llama_context * ctx,
llama_seq_id seq_id); llama_seq_id seq_id);
LLAMA_API DEPRECATED(llama_pos llama_kv_cache_seq_pos_max(
struct llama_context * ctx,
llama_seq_id seq_id),
"use llama_cache_seq_pos_max instead, which also now returns -1 instead of 0 when the seq_id has no cells");
// Defragment the KV cache // Defragment the KV cache
// This will be applied: // This will be applied: