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.cpp
+20 -278
View File
@@ -112,7 +112,7 @@ llama_kv_cache::llama_kv_cache(
auto it = ctx_map.find(buft);
if (it == ctx_map.end()) {
ggml_init_params params = {
/*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), //Reserve tensor metadata for up to 3 tensors per layer (K, V, and optional K_idx), plus one view per tensor per stream.
/*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
@@ -242,25 +242,9 @@ llama_kv_cache::llama_kv_cache(
v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr);
}
const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il);
ggml_tensor * k_idx = n_embd_k_idx > 0
? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream)
: nullptr;
if (k_idx) {
ggml_format_name(k_idx, "cache_k_idx_l%d", il);
msa_strict_slots = (n_stream == n_seq_max);
}
std::vector<ggml_tensor *> k_idx_stream;
for (uint32_t s = 0; s < n_stream; ++s) {
k_idx_stream.push_back(k_idx
? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2])
: nullptr);
}
map_layer_ids[il] = layers.size();
layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream });
layers.push_back({ il, k, v, k_stream, v_stream, });
}
if (reuse) {
@@ -309,24 +293,13 @@ llama_kv_cache::llama_kv_cache(
}
{
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
const size_t memory_size_k_idx = size_k_idx_bytes();
const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx;
const size_t memory_size_k = size_k_bytes();
const size_t memory_size_v = size_v_bytes();
constexpr float mib = 1024.0f * 1024.0f;
const std::string k_log = format(", K (%s): %7.2f MiB", ggml_type_name(type_k), (float) memory_size_k / mib);
const std::string v_log = format(", V (%s): %7.2f MiB", ggml_type_name(type_v), (float) memory_size_v / mib);
std::string k_idx_log;
if (memory_size_k_idx > 0) {
k_idx_log = format(", K_idx (%s): %7.2f MiB", ggml_type_name(GGML_TYPE_F32), (float) memory_size_k_idx / mib);
}
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs)%s%s%s\n", __func__,
(float) memory_size_total / mib, kv_size, (int) layers.size(), n_seq_max, n_stream,
k_log.c_str(), v_log.c_str(), k_idx_log.c_str());
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
}
// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
@@ -419,39 +392,6 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
p1 = std::numeric_limits<llama_pos>::max();
}
// empty range - nothing to remove
if (p0 >= p1) {
return true;
}
// MSA anchors block selection to absolute cache slots (slot == position). Tail trim and full removal preserve this invariant, but removing a prefix
// or middle range would free slots while later cells survive, desynchronizing the indexer cache. Reject such removals before modifying the cache.
if (msa_strict_slots) {
for (llama_seq_id sid = 0; sid < (llama_seq_id) seq_to_stream.size(); ++sid) {
if (seq_id >= 0 && sid != seq_id) {
continue;
}
const auto & cells = v_cells[seq_to_stream[sid]];
const llama_pos pmin = cells.seq_pos_min(sid);
const llama_pos pmax = cells.seq_pos_max(sid);
if (pmin < 0) {
continue; // empty sequence
}
const bool overlaps = p0 <= pmax && p1 > pmin; // the range removes something
const bool leaves_tail = p1 <= pmax; // cells beyond the range survive
if (overlaps && leaves_tail) {
LLAMA_LOG_WARN("%s: MSA: partial (non-suffix) removal [%d, %d) for seq %d is not supported "
"(block selection is anchored to cache slots) - rejected\n", __func__, p0, p1, sid);
return false;
}
}
}
if (seq_id >= 0) {
auto & cells = v_cells[seq_to_stream[seq_id]];
auto & head = v_heads[seq_to_stream[seq_id]];
@@ -906,10 +846,6 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
if (layer.v_stream[ssrc]) {
ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);
}
if (layer.k_idx_stream[ssrc]) {
GGML_ASSERT(layer.k_idx_stream[sdst]);
ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]);
}
}
}
}
@@ -1058,44 +994,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
const auto & cells = v_cells[seq_to_stream[seq_id]];
if (n_tokens > cells.size()) {
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
return { };
}
// MSA block selection assumes slot == logical position (append-only streams).
if (msa_strict_slots) {
for (uint32_t ii = 0; ii < n_tokens; ++ii) {
const llama_pos pos = ubatch.pos[s*n_tokens + ii];
if (pos < 0 || (uint64_t) pos >= cells.size()) {
LLAMA_LOG_WARN("%s: MSA: position %d is outside the cache range [0, %u)\n",
__func__, pos, cells.size());
return { };
}
const uint32_t idx = (uint32_t) pos;
if (!cells.is_empty(idx)) {
LLAMA_LOG_WARN("%s: MSA: required slot %u is already occupied (stream %u)\n",
__func__, idx, seq_to_stream[seq_id]);
return { };
}
// strictly increasing positions, rules out duplicates and, for contiguous requests, is tightened to exact adjacency
if (!res.idxs[s].empty() && (cont ? idx != res.idxs[s].back() + 1
: idx <= res.idxs[s].back())) {
LLAMA_LOG_WARN("%s: MSA: token positions are not %s within the ubatch\n",
__func__, cont ? "contiguous" : "strictly increasing");
return { };
}
res.idxs[s].push_back(idx);
}
continue;
}
uint32_t head_cur = v_heads[seq_to_stream[seq_id]];
// if we have enough unused cells before the current head ->
@@ -1104,6 +1002,11 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
head_cur = 0;
}
if (n_tokens > cells.size()) {
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
return { };
}
uint32_t n_tested = 0;
// for continuous slots, we test that all tokens in the ubatch fit, starting from the current head
@@ -1210,15 +1113,6 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
const auto idx = sinfo.idxs[s][ii];
if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) {
LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: "
"writing pos %d into cell %u (stream %u). The indexer cache "
"would desync and block selection would silently corrupt. "
"This is a bug, please report it with reproduction steps.\n",
__func__, ubatch.pos[i], idx, sinfo.strm[s]);
GGML_ABORT("MSA: slot != pos");
}
if (!cells.is_empty(idx)) {
assert(cells.seq_count(idx) == 1);
@@ -1262,8 +1156,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",
__func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
// under MSA strict slots this path should be unreachable, since strict MSA placement never selects occupied cells
GGML_ASSERT(seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1));
seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
}
}
@@ -1283,12 +1176,6 @@ bool llama_kv_cache::get_can_shift() const {
if (hparams.n_pos_per_embd() > 1) {
return false;
}
// shifting would leave k_idx stale
for (const auto & layer : layers) {
if (layer.k_idx) {
return false;
}
}
return true;
}
@@ -1337,6 +1224,12 @@ ggml_tensor * llama_kv_cache::get_k_storage(int32_t il) const {
return layers[ikv].k;
}
const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const {
GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
return v_cells[seq_to_stream[seq_id]];
}
uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
uint32_t result = 0;
@@ -1405,23 +1298,6 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
}
ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
const int32_t ikv = map_layer_ids.at(il);
auto * k_idx = layers[ikv].k_idx;
GGML_ASSERT(k_idx);
const uint64_t kv_size = get_size();
const int64_t n_idx = k_idx->ne[0]; // 128
const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
return ggml_view_4d(ctx, k_idx,
n_idx, 1, n_kv, ns,
ggml_row_size(k_idx->type, n_idx), // nb1 (single head)
ggml_row_size(k_idx->type, n_idx), // nb2 (per cell)
ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream)
ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0);
}
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
@@ -1523,28 +1399,6 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama
return k_idxs;
}
ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
GGML_UNUSED(sinfo);
const int32_t ikv = map_layer_ids.at(il);
ggml_tensor * k_idx = layers[ikv].k_idx;
GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache");
const int64_t n_embd_head = k_idx_cur->ne[0]; // 128
const int64_t n_head = k_idx_cur->ne[1]; // 1
const int64_t n_tokens = k_idx_cur->ne[2];
const int64_t n_embd_gqa = n_embd_head*n_head; // 128
GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]);
k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0);
const int64_t n_stream = k_idx->ne[2];
if (n_stream > 1) {
const int64_t kv_size = get_size();
k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream);
}
return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store
}
ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
const uint32_t n_tokens = ubatch.n_tokens;
@@ -1979,18 +1833,6 @@ size_t llama_kv_cache::size_v_bytes() const {
return size_v_bytes;
}
size_t llama_kv_cache::size_k_idx_bytes() const {
size_t size_k_idx_bytes = 0;
for (const auto & layer : layers) {
if (layer.k_idx) {
size_k_idx_bytes += ggml_nbytes(layer.k_idx);
}
}
return size_k_idx_bytes;
}
ggml_tensor * llama_kv_cache::build_rope_shift(
const llama_cparams & cparams,
ggml_context * ctx,
@@ -2303,36 +2145,6 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
}
}
if (size_k_idx_bytes() > 0) {
const uint32_t has_k_idx_u32 = 1;
io.write(&has_k_idx_u32, sizeof(has_k_idx_u32));
for (const auto & layer : layers) {
const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0;
io.write(&layer_has_k_idx, sizeof(layer_has_k_idx));
if (!layer_has_k_idx) {
continue;
}
GGML_ASSERT(layer.k_idx_stream[cr.strm]);
const int32_t k_idx_type_i = (int32_t) layer.k_idx->type;
io.write(&k_idx_type_i, sizeof(k_idx_type_i));
const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
io.write(&k_idx_size_row, sizeof(k_idx_size_row));
for (const auto & range : cr.data) {
const size_t range_size = range.second - range.first;
const size_t buf_size = range_size * k_idx_size_row;
const size_t offset = range.first * k_idx_size_row;
io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size);
}
}
}
if (!v_trans) {
for (const auto & layer : layers) {
const uint32_t il = layer.il;
@@ -2581,68 +2393,6 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
}
}
if (size_k_idx_bytes() > 0) {
uint32_t has_k_idx_u32 = 0;
io.read(&has_k_idx_u32, sizeof(has_k_idx_u32));
if (has_k_idx_u32 != 1) {
LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__);
return false;
}
for (const auto & layer : layers) {
uint32_t layer_has_k_idx = 0;
io.read(&layer_has_k_idx, sizeof(layer_has_k_idx));
const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0;
if (layer_has_k_idx != expected_layer_has_k_idx) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx state for layer: got %u, expected %u\n",
__func__, layer_has_k_idx, expected_layer_has_k_idx);
return false;
}
if (!layer_has_k_idx) {
continue;
}
GGML_ASSERT(layer.k_idx_stream[strm]);
int32_t k_idx_type_i = -1;
io.read(&k_idx_type_i, sizeof(k_idx_type_i));
if (k_idx_type_i != (int32_t) layer.k_idx->type) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx type: got %d, expected %d\n",
__func__, k_idx_type_i, (int32_t) layer.k_idx->type);
return false;
}
uint64_t k_idx_size_row = 0;
io.read(&k_idx_size_row, sizeof(k_idx_size_row));
const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
if (k_idx_size_row != expected_k_idx_size_row) {
LLAMA_LOG_ERROR(
"%s: mismatched k_idx row size: got %zu, expected %zu\n",
__func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row);
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row);
} else {
for (uint32_t i = 0; i < cell_count; ++i) {
io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row);
}
}
}
}
}
if (!this->v_trans) {
for (const auto & layer : layers) {
const uint32_t il = layer.il;
@@ -2844,10 +2594,6 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const {
return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
}
@@ -2856,10 +2602,6 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_
return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const {
return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]);
}
ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
return kv->build_input_k_idxs(ctx, ubatch);
}