Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model: M3: Move MSA into a new memory implementation (#26338)

Browse Source 
* Move MSA logic from llama-kv-cache into llama-kv-cache-msa

* cont : minor

* cont : ws fix

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
timkhronos
2026-08-03 16:30:08 +03:00
committed by GitHub
co-authored by Georgi Gerganov
parent 563dec81c1
commit 67d5978bb1
11 changed files with 826 additions and 378 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-kv-cache.h
+2 -10
View File
@@ -164,6 +164,8 @@ public:
std::vector<uint32_t> get_layer_ids() const;
ggml_tensor * get_k_storage(int32_t il) const;
const llama_kv_cells & get_cells(llama_seq_id seq_id) const;
//
// graph_build API
//
@@ -173,12 +175,10 @@ public:
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
// store k_cur and v_cur in the cache based on the provided head location
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
//
// preparation API
@@ -230,11 +230,9 @@ private:
ggml_tensor * k;
ggml_tensor * v;
ggml_tensor * k_idx; // MSA single-head indexer keys, F32
std::vector<ggml_tensor *> k_stream;
std::vector<ggml_tensor *> v_stream;
std::vector<ggml_tensor *> k_idx_stream;
};
bool v_trans = true; // the value tensor is transposed
@@ -263,9 +261,6 @@ private:
// env: LLAMA_KV_CACHE_DEBUG
int debug = 0;
// set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq)
bool msa_strict_slots = false;
// this is the SWA type of the cache - not to be confused with the model SWA type
const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
@@ -298,7 +293,6 @@ private:
size_t size_k_bytes() const;
size_t size_v_bytes() const;
size_t size_k_idx_bytes() const;
ggml_tensor * build_rope_shift(
const llama_cparams & cparams,
@@ -378,7 +372,6 @@ public:
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const;
// store k_cur and v_cur in the cache based on the provided head location
// note: the heads in k_cur and v_cur should be laid out contiguously in memory
@@ -388,7 +381,6 @@ public:
// - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const;
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const;
// create destination indices for each head of the current batch for where it would be written in the KV cache
// the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but