Skip to content

Commit

Permalink
mamba : in comments, properly refer to KV cells instead of slots
Browse files Browse the repository at this point in the history
  • Loading branch information
compilade committed Feb 14, 2024
1 parent 06ead3d commit 322686e
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions llama.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1734,7 +1734,7 @@ struct llama_kv_cell {
// ring-buffer of cached KV data
struct llama_kv_cache {
bool has_shift = false;
// with Mamba, a slot can hold the state for more than one past token
// with Mamba, a cell can hold the state for more than one past token
bool unlimited = false;

// Note: The value of head isn't only used to optimize searching
Expand Down Expand Up @@ -1993,7 +1993,7 @@ static bool llama_kv_cache_init(

cache.has_shift = false;

// for now, only Mamba can hold state for more than one past token per slot
// for now, only Mamba can hold state for more than one past token per cell
cache.unlimited = model.arch == LLM_ARCH_MAMBA;

cache.head = 0;
Expand Down Expand Up @@ -2249,7 +2249,7 @@ static void llama_kv_cache_seq_cp(
cache.cells[seq_id_dst].delta = seq_id_src;
// NOTE: a sequence can't have multiple sources, but can have multiple destinations.
// For compatibility with the other KV cache API functions,
// the seq_id(s) of a slot suggests an intent to "copy to" those id(s),
// the seq_id(s) of a cell suggests an intent to "copy to" those id(s),
// so that when a sequence is copied, it can initially be found from the source cell.
cache.cells[seq_id_src].seq_id.insert(seq_id_dst);
// prevent the destination from getting cleared
Expand Down Expand Up @@ -11726,10 +11726,10 @@ struct llama_context * llama_new_context_with_model(
ggml_type type_k = params.type_k;
ggml_type type_v = params.type_v;

// Mamba only needs a constant number of KV cache slots per sequence
// Mamba only needs a constant number of KV cache cells per sequence
if (model->arch == LLM_ARCH_MAMBA) {
// Mamba needs as many slots as there are distinct sequences processed at the same time
// The extra slot allows dedicating a sequence id to the system prompt
// Mamba needs as many KV cells as there are sequences kept at any time
// The extra cell allows dedicating a sequence id to the system prompt
// TODO: find a better way to get the max number of parallel sequences
kv_size = params.n_parallel + 1;
// it's probably best to keep as much precision as possible for the states
Expand Down

0 comments on commit 322686e

Please sign in to comment.