Felipe97/llama-cpp-compiled
01.1k
1#include "llama-kv-cache.h"2 3#include "llama-impl.h"4#include "llama-io.h"5#include "llama-model.h"6#include "llama-context.h"7 8#include <algorithm>9#include <cassert>10#include <cmath>11#include <cstring>12#include <limits>13#include <map>14#include <stdexcept>15#include <unordered_map>16 17static bool ggml_is_power_of_2(int n) {18 return (n & (n - 1)) == 0;19}20 21// orthonormal Walsh-Hadamard rotation matrix22// note: res^2 == I23static void ggml_gen_hadamard(ggml_tensor * tensor) {24 assert(tensor->type == GGML_TYPE_F32);25 26 const int n = tensor->ne[0];27 28 assert(ggml_is_power_of_2(n));29 assert(tensor->ne[1] == n);30 assert(tensor->ne[2] == 1);31 assert(tensor->ne[3] == 1);32 33 std::vector<float> data_f32;34 35 float * data = (float *) tensor->data;36 37 if (tensor->type != GGML_TYPE_F32) {38 data_f32.resize(n*n);39 data = data_f32.data();40 }41 42 data[0*n + 0] = 1.0 / sqrtf(n);43 44 for (int s = 1; s < n; s *= 2) {45 for (int i = 0; i < s; i++) {46 for (int j = 0; j < s; j++) {47 const float val = data[i*n + j];48 49 data[(i + s)*n + (j )] = val;50 data[(i )*n + (j + s)] = val;51 data[(i + s)*n + (j + s)] = -val;52 }53 }54 }55 56 if (tensor->type != GGML_TYPE_F32) {57 ggml_quantize_chunk(tensor->type, data, tensor->data, 0, 1, n*n, nullptr);58 }59}60 61//62// llama_kv_cache63//64 65llama_kv_cache::llama_kv_cache(66 const llama_model & model,67 const llama_hparams & hparams,68 ggml_type type_k,69 ggml_type type_v,70 bool v_trans,71 bool offload,72 bool unified,73 uint32_t kv_size,74 uint32_t n_seq_max,75 uint32_t n_pad,76 uint32_t n_swa,77 llama_swa_type swa_type,78 llama_memory_t mem_other,79 const layer_filter_cb & filter,80 const layer_reuse_cb & reuse,81 const layer_share_cb & share,82 const char * name_tag) :83 model(model), hparams(hparams), v_trans(v_trans),84 n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type),85 other(static_cast<llama_kv_cache *>(mem_other)),86 v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),87 v_cells(*v_cells_impl) {88 89 // shared cells view the source cache's K/V tensors, so the cell count90 // follows the source allocation: a fitted target can be smaller than the91 // draft default and oversized views would overflow the source tensors92 if (other) {93 const uint32_t size_other = other->get_size();94 if (kv_size != size_other) {95 LLAMA_LOG_WARN("%s: kv_size = %u overridden to %u to match the shared source cache\n", __func__, kv_size, size_other);96 kv_size = size_other;97 }98 }99 100 GGML_ASSERT(kv_size % n_pad == 0);101 102 const uint32_t n_layer = hparams.n_layer_all;103 104 // define a comparator for the buft -> ctx map to ensure that the order is well-defined:105 struct ggml_backend_buft_comparator {106 bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {107 return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;108 }109 };110 std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> ctx_map;111 112 // create a context for each buffer type113 auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {114 auto it = ctx_map.find(buft);115 if (it == ctx_map.end()) {116 ggml_init_params params = {117 /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),118 /*.mem_buffer =*/ NULL,119 /*.no_alloc =*/ true,120 };121 122 ggml_context * ctx = ggml_init(params);123 if (!ctx) {124 return nullptr;125 }126 127 ctx_map.emplace(buft, ctx);128 129 return ctx;130 }131 132 return it->second.get();133 };134 135 GGML_ASSERT(n_stream == 1 || n_stream == n_seq_max);136 137 v_heads.resize(n_stream);138 for (uint32_t s = 0; s < n_stream; ++s) {139 v_heads[s] = 0;140 }141 142 v_cells.resize(n_stream);143 for (uint32_t s = 0; s < n_stream; ++s) {144 v_cells[s].resize(kv_size);145 }146 147 // by default, all sequence ids are mapped to the 0th stream148 seq_to_stream.resize(LLAMA_MAX_SEQ, 0);149 150 if (n_stream > 1) {151 seq_to_stream.resize(n_stream, 0);152 for (uint32_t s = 0; s < n_stream; ++s) {153 seq_to_stream[s] = s;154 }155 }156 157 // [TAG_V_CACHE_VARIABLE]158 if (v_trans && hparams.is_n_embd_v_gqa_variable()) {159 LLAMA_LOG_WARN("%s: the V embeddings have different sizes across layers and FA is not enabled - padding V cache to %d\n",160 __func__, hparams.n_embd_v_gqa_max());161 }162 163 const bool is_mla = hparams.is_mla();164 165 for (uint32_t il = 0; il < n_layer; il++) {166 if (!hparams.has_kv(il)) {167 LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il);168 continue;169 }170 171 if (filter && !filter(il)) {172 LLAMA_LOG_DEBUG("%s: layer %3d: filtered\n", __func__, il);173 continue;174 }175 176 if (share && other) {177 const int32_t il_share = share(il);178 179 if (il_share >= 0) {180 const auto & layer_share = other->layers[other->map_layer_ids[il_share]];181 182 LLAMA_LOG_WARN("%s: layer %3d: sharing with layer %d. k = %p, v = %p\n", __func__, il, il_share,183 layer_share.k->data, layer_share.v->data);184 185 map_layer_ids[il] = layers.size();186 187 layers.push_back(layer_share);188 layers.back().il = il;189 190 continue;191 }192 }193 194 if (n_embd_head_k_all == 0) {195 n_embd_head_k_all = (int32_t) hparams.n_embd_head_k(il);196 } else if (n_embd_head_k_all > 0 && n_embd_head_k_all != (int32_t) hparams.n_embd_head_k(il)) {197 n_embd_head_k_all = -1;198 }199 200 if (!is_mla) {201 if (n_embd_head_v_all == 0) {202 n_embd_head_v_all = (int32_t) hparams.n_embd_head_v(il);203 } else if (n_embd_head_v_all > 0 && n_embd_head_v_all != (int32_t) hparams.n_embd_head_v(il)) {204 n_embd_head_v_all = -1;205 }206 }207 208 // [TAG_V_CACHE_VARIABLE]209 const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);210 const uint32_t n_embd_v_gqa = !v_trans ? hparams.n_embd_v_gqa(il) : hparams.n_embd_v_gqa_max();211 212 const char * dev_name = "CPU";213 214 ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();215 216 if (offload) {217 auto * dev = model.dev_layer(il);218 buft = ggml_backend_dev_buffer_type(dev);219 220 dev_name = ggml_backend_dev_name(dev);221 }222 223 LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);224 225 ggml_context * ctx = ctx_for_buft(buft);226 if (!ctx) {227 throw std::runtime_error("failed to create ggml context for kv cache");228 }229 230 const bool has_k = true;231 const bool has_v = !is_mla;232 233 ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr;234 ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr;235 236 has_k && ggml_format_name(k, "cache_%sk_l%d", name_tag, il);237 has_v && ggml_format_name(v, "cache_%sv_l%d", name_tag, il);238 239 std::vector<ggml_tensor *> k_stream;240 std::vector<ggml_tensor *> v_stream;241 242 for (uint32_t s = 0; s < n_stream; ++s) {243 k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gqa, kv_size, k->nb[1], s*k->nb[2]) : nullptr);244 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);245 }246 247 map_layer_ids[il] = layers.size();248 249 layers.push_back({ il, k, v, k_stream, v_stream, });250 }251 252 if (reuse) {253 LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__);254 255 for (uint32_t il = 0; il < n_layer; il++) {256 const int32_t il_reuse = reuse(il);257 258 if (il_reuse < 0) {259 LLAMA_LOG_DEBUG("%s: - layer %3d: no reuse\n", __func__, il);260 continue;261 }262 263 if (filter && !filter(il)) {264 LLAMA_LOG_DEBUG("%s: - layer %3d: filtered\n", __func__, il);265 continue;266 }267 268 GGML_ASSERT(map_layer_ids.find(il_reuse) != map_layer_ids.end());269 270 map_layer_ids[il] = map_layer_ids[il_reuse];271 272 LLAMA_LOG_DEBUG("%s: - layer %3d: reuse layer %d, is_swa = %d\n", __func__, il, il_reuse, hparams.is_swa(il));273 }274 }275 276 // allocate tensors and initialize the buffers to avoid NaNs in the padding277 for (auto & [buft, ctx] : ctx_map) {278 ggml_backend_buffer_t buf;279 if (hparams.no_alloc) {280 buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer281 for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {282 t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it283 }284 } else {285 buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer286 }287 if (!buf) {288 throw std::runtime_error("failed to allocate buffer for kv cache");289 }290 291 LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);292 293 ggml_backend_buffer_clear(buf, 0);294 ctxs_bufs.emplace_back(std::move(ctx), buf);295 }296 297 {298 const size_t memory_size_k = size_k_bytes();299 const size_t memory_size_v = size_v_bytes();300 301 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__,302 (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,303 ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),304 ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));305 }306 307 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]308 if (other) {309 n_embd_head_k_all = other->n_embd_head_k_all;310 n_embd_head_v_all = other->n_embd_head_v_all;311 312 attn_rot_k = other->attn_rot_k;313 attn_rot_v = other->attn_rot_v;314 } else {315 const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");316 const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;317 if (attn_rot_disable) {318 LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);319 }320 321 attn_rot_k =322 !attn_rot_disable &&323 n_embd_head_k_all > 0 &&324 ggml_is_quantized(type_k) &&325 hparams.n_embd_head_k() % 64 == 0;326 327 // always create Hadamard rotation tensors for DeepSeek lightning indexers328 if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 ||329 model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_DOTS3NOTE) &&330 hparams.n_embd_head_k_full == hparams.indexer_head_size) {331 attn_rot_k = true;332 }333 334 attn_rot_v =335 !attn_rot_disable &&336 n_embd_head_v_all > 0 &&337 ggml_is_quantized(type_v) &&338 hparams.n_embd_head_v() % 64 == 0;339 }340 341 LLAMA_LOG_INFO("%s: attn_rot_k = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_k, n_embd_head_k_all);342 LLAMA_LOG_INFO("%s: attn_rot_v = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_v, n_embd_head_v_all);343 344 // pre-compute the haramard matrices and keep them in host memory345 // TODO: in the future, we can make copies in the backend buffers to avoid host -> device transfers346 if (attn_rot_k || attn_rot_v) {347 for (int64_t n = 64; n <= std::max(n_embd_head_k_all, n_embd_head_v_all); n *= 2) {348 attn_rot_hadamard[n] = std::vector<float>(n*n);349 350 ggml_init_params params = {351 /* .mem_size = */ 1*ggml_tensor_overhead(),352 /* .mem_buffer = */ nullptr,353 /* .no_alloc = */ true,354 };355 356 ggml_context_ptr ctx { ggml_init(params) };357 358 ggml_tensor * tmp = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, n, n);359 tmp->data = attn_rot_hadamard[n].data();360 361 ggml_gen_hadamard(tmp);362 }363 }364 365 const char * LLAMA_KV_CACHE_DEBUG = getenv("LLAMA_KV_CACHE_DEBUG");366 debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0;367}368 369void llama_kv_cache::clear(bool data) {370 for (uint32_t s = 0; s < n_stream; ++s) {371 v_cells[s].reset();372 v_heads[s] = 0;373 }374 375 if (data) {376 for (auto & [_, buf] : ctxs_bufs) {377 ggml_backend_buffer_clear(buf.get(), 0);378 }379 }380}381 382bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {383 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]384 if (other) {385 return true;386 }387 388 // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]389 GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));390 391 if (p0 < 0) {392 p0 = 0;393 }394 395 if (p1 < 0) {396 p1 = std::numeric_limits<llama_pos>::max();397 }398 399 if (seq_id >= 0) {400 auto & cells = v_cells[seq_to_stream[seq_id]];401 auto & head = v_heads[seq_to_stream[seq_id]];402 403 uint32_t new_head = cells.size();404 405 for (uint32_t i = 0; i < cells.size(); ++i) {406 if (!cells.pos_in(i, p0, p1)) {407 continue;408 }409 410 if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {411 if (new_head == cells.size()) {412 new_head = i;413 }414 }415 }416 417 // If we freed up a slot, set head to it so searching can start there.418 if (new_head != cells.size() && new_head < head) {419 head = new_head;420 }421 } else {422 // match any sequence423 for (uint32_t s = 0; s < n_stream; ++s) {424 auto & cells = v_cells[s];425 auto & head = v_heads[s];426 427 uint32_t new_head = cells.size();428 429 for (uint32_t i = 0; i < cells.size(); ++i) {430 if (!cells.pos_in(i, p0, p1)) {431 continue;432 }433 434 cells.rm(i);435 436 if (new_head == cells.size()) {437 new_head = i;438 }439 }440 441 // If we freed up a slot, set head to it so searching can start there.442 if (new_head != cells.size() && new_head < head) {443 head = new_head;444 }445 }446 }447 448 return true;449}450 451void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {452 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]453 if (other) {454 return;455 }456 457 GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());458 GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());459 460 const auto s0 = seq_to_stream[seq_id_src];461 const auto s1 = seq_to_stream[seq_id_dst];462 463 if (s0 == s1) {464 // since both sequences are in the same stream, no data copy is necessary465 // we just have to update the cells meta data466 467 auto & cells = v_cells[s0];468 469 if (seq_id_src == seq_id_dst) {470 return;471 }472 473 if (p0 < 0) {474 p0 = 0;475 }476 477 if (p1 < 0) {478 p1 = std::numeric_limits<llama_pos>::max();479 }480 481 for (uint32_t i = 0; i < cells.size(); ++i) {482 if (!cells.pos_in(i, p0, p1)) {483 continue;484 }485 486 if (cells.seq_has(i, seq_id_src)) {487 cells.seq_add(i, seq_id_dst);488 }489 }490 491 return;492 }493 494 // cross-stream sequence copies require to copy the actual buffer data495 496 bool is_full = true;497 498 if (p0 > 0 && p0 + 1 < (int) get_size()) {499 is_full = false;500 }501 502 if (p1 > 0 && p1 + 1 < (int) get_size()) {503 is_full = false;504 }505 506 GGML_ASSERT(is_full && "seq_cp() is only supported for full KV buffers");507 508 // enqueue the copy operation - the buffer copy will be performed during the next update509 sc_info.ssrc.push_back(s0);510 sc_info.sdst.push_back(s1);511 512 v_cells[s1].reset();513 for (uint32_t i = 0; i < v_cells[s0].size(); ++i) {514 if (v_cells[s0].seq_has(i, seq_id_src)) {515 llama_pos pos = v_cells[s0].pos_get(i);516 llama_pos shift = v_cells[s0].get_shift(i);517 518 llama_kv_cell_ext ext = v_cells[s0].ext_get(i);519 520 if (shift != 0) {521 pos -= shift;522 assert(pos >= 0);523 }524 525 v_cells[s1].pos_set(i, pos);526 v_cells[s1].seq_add(i, seq_id_dst);527 528 if (shift != 0) {529 v_cells[s1].pos_add(i, shift);530 }531 532 v_cells[s1].ext_set(i, ext);533 }534 }535 536 v_heads[s1] = v_heads[s0];537 538 //for (uint32_t s = 0; s < n_stream; ++s) {539 // LLAMA_LOG_WARN("%s: seq %d: min = %d, max = %d\n", __func__, s, v_cells[s].seq_pos_min(s), v_cells[s].seq_pos_max(s));540 //}541}542 543void llama_kv_cache::seq_keep(llama_seq_id seq_id) {544 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]545 if (other) {546 return;547 }548 549 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());550 551 auto & cells = v_cells[seq_to_stream[seq_id]];552 auto & head = v_heads[seq_to_stream[seq_id]];553 554 uint32_t new_head = cells.size();555 556 for (uint32_t i = 0; i < cells.size(); ++i) {557 if (cells.seq_keep(i, seq_id)) {558 if (new_head == cells.size()) {559 new_head = i;560 }561 }562 }563 564 // If we freed up a slot, set head to it so searching can start there.565 if (new_head != cells.size() && new_head < head) {566 head = new_head;567 }568}569 570void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {571 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]572 if (other) {573 return;574 }575 576 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());577 GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1");578 579 auto & cells = v_cells[seq_to_stream[seq_id]];580 auto & head = v_heads[seq_to_stream[seq_id]];581 582 if (shift == 0) {583 return;584 }585 586 uint32_t new_head = cells.size();587 588 if (p0 < 0) {589 p0 = 0;590 }591 592 if (p1 < 0) {593 p1 = std::numeric_limits<llama_pos>::max();594 }595 596 // If there is no range then return early to avoid looping over all cells.597 if (p0 == p1) {598 return;599 }600 601 for (uint32_t i = 0; i < cells.size(); ++i) {602 if (!cells.pos_in(i, p0, p1)) {603 continue;604 }605 606 if (cells.seq_has(i, seq_id)) {607 if (cells.pos_add(i, shift)) {608 if (new_head == cells.size()) {609 new_head = i;610 }611 }612 }613 }614 615 // If we freed up a slot, set head to it so searching can start there.616 // Otherwise we just start the next search from the beginning.617 head = new_head != cells.size() ? new_head : 0;618}619 620void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {621 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]622 if (other) {623 return;624 }625 626 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());627 GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1");628 629 auto & cells = v_cells[seq_to_stream[seq_id]];630 631 if (d == 1) {632 return;633 }634 635 if (p0 < 0) {636 p0 = 0;637 }638 639 if (p1 < 0) {640 p1 = std::numeric_limits<llama_pos>::max();641 }642 643 // If there is no range then return early to avoid looping over the cache.644 if (p0 == p1) {645 return;646 }647 648 for (uint32_t i = 0; i < cells.size(); ++i) {649 if (!cells.pos_in(i, p0, p1)) {650 continue;651 }652 653 if (cells.seq_has(i, seq_id)) {654 cells.pos_div(i, d);655 }656 }657}658 659llama_pos llama_kv_cache::seq_pos_min(llama_seq_id seq_id) const {660 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]661 if (other) {662 return other->seq_pos_min(seq_id);663 }664 665 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());666 667 const auto & cells = v_cells[seq_to_stream[seq_id]];668 669 return cells.seq_pos_min(seq_id);670}671 672llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {673 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]674 if (other) {675 return other->seq_pos_max(seq_id);676 }677 678 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());679 680 const auto & cells = v_cells[seq_to_stream[seq_id]];681 682 return cells.seq_pos_max(seq_id);683}684 685std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {686 std::map<ggml_backend_buffer_type_t, size_t> ret;687 for (const auto & [ctx, buf] : ctxs_bufs) {688 ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());689 690 if (hparams.no_alloc) {691 GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);692 ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);693 } else {694 // GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base695 ret[buft] += ggml_backend_buffer_get_size(buf.get());696 }697 }698 699 return ret;700}701 702llama_memory_context_ptr llama_kv_cache::init_batch(703 llama_batch_allocr & balloc,704 uint32_t n_ubatch,705 bool embd_all) {706 GGML_UNUSED(embd_all);707 708 do {709 balloc.split_reset();710 711 std::vector<llama_ubatch> ubatches;712 while (true) {713 auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);714 715 if (ubatch.n_tokens == 0) {716 break;717 }718 719 ubatches.push_back(std::move(ubatch)); // NOLINT720 }721 722 if (balloc.get_n_used() < balloc.get_n_tokens()) {723 // failed to find a suitable split724 break;725 }726 727 auto sinfos = prepare(ubatches);728 if (sinfos.empty()) {729 break;730 }731 732 return std::make_unique<llama_kv_cache_context>(733 this, std::move(sinfos), std::move(ubatches));734 } while (false);735 736 return std::make_unique<llama_kv_cache_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);737}738 739llama_memory_context_ptr llama_kv_cache::init_full() {740 return std::make_unique<llama_kv_cache_context>(this);741}742 743llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool optimize) {744 GGML_UNUSED(optimize);745 746 bool do_shift = get_has_shift();747 748 return std::make_unique<llama_kv_cache_context>(this, lctx, do_shift, std::move(sc_info));749}750 751llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_ubatch> & ubatches) {752 llama_kv_cache::slot_info_vec_t res;753 754 struct state_t {755 slot_info sinfo; // slot info for the ubatch756 757 std::vector<uint32_t> v_heads_old; // old positions of the heads, before placing the ubatch758 759 std::vector<llama_kv_cells> v_cells; // copy of the old cells, before placing the ubatch760 };761 762 // remember the old state of the cells so we can restore it in the end763 std::vector<state_t> states;764 765 bool success = true;766 767 for (const auto & ubatch : ubatches) {768 // only find a suitable slot for the ubatch. don't modify the cells yet769 const auto sinfo_new = find_slot(ubatch, false);770 if (sinfo_new.empty()) {771 success = false;772 break;773 }774 775 // remember the position that we found776 res.push_back(sinfo_new);777 778 // store the old state of the cells in the recovery stack779 {780 state_t state = { sinfo_new, v_heads, {} };781 782 for (uint32_t s = 0; s < sinfo_new.n_stream(); ++s) {783 auto & cells = v_cells[sinfo_new.strm[s]];784 785 state.v_cells.push_back(cells.cp(sinfo_new.idxs[s]));786 }787 788 states.push_back(std::move(state));789 }790 791 // now emplace the ubatch792 apply_ubatch(sinfo_new, ubatch);793 }794 795 GGML_ASSERT(!states.empty() || !success);796 797 // iterate backwards and restore the cells to their original state798 for (auto it = states.rbegin(); it != states.rend(); ++it) {799 const auto & sinfo = it->sinfo;800 801 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {802 auto & cells = v_cells[sinfo.strm[s]];803 auto & head = v_heads[sinfo.strm[s]];804 805 cells.set(sinfo.idxs[s], it->v_cells[s]);806 head = it->v_heads_old[s];807 }808 }809 810 if (!success) {811 return {};812 }813 814 return res;815}816 817bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info) {818 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]819 if (other) {820 return true;821 }822 823 bool updated = false;824 825 auto * sched = lctx->get_sched();826 827 if (!sc_info.empty()) {828 assert(n_stream > 1 && "stream copy should never happen with a single stream");829 830 llama_synchronize(lctx);831 832 const size_t n_copy = sc_info.ssrc.size();833 834 for (size_t i = 0; i < n_copy; ++i) {835 const auto ssrc = sc_info.ssrc[i];836 const auto sdst = sc_info.sdst[i];837 838 assert(ssrc < n_stream);839 assert(sdst < n_stream);840 841 LLAMA_LOG_DEBUG("%s: copying KV buffer: stream %d to stream %d\n", __func__, ssrc, sdst);842 843 assert(ssrc != sdst);844 845 for (uint32_t il = 0; il < layers.size(); ++il) {846 const auto & layer = layers[il];847 848 ggml_backend_tensor_copy(layer.k_stream[ssrc], layer.k_stream[sdst]);849 850 if (layer.v_stream[ssrc]) {851 ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);852 }853 }854 }855 }856 857 if (do_shift) {858 if (!get_can_shift()) {859 GGML_ABORT("The current KV cache / model configuration does not support K-shift");860 }861 862 LLAMA_LOG_DEBUG("%s: applying K-shift\n", __func__);863 864 // apply K-shift if needed865 if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) {866 ggml_backend_sched_reset(sched);867 868 auto * res = lctx->get_gf_res_reserve();869 870 res->reset();871 872 auto * gf = build_graph_shift(res, lctx);873 if (!ggml_backend_sched_alloc_graph(sched, gf)) {874 LLAMA_LOG_ERROR("%s: failed to allocate compute graph for K-shift\n", __func__);875 return updated;876 }877 878 res->set_inputs(nullptr);879 880 if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) {881 LLAMA_LOG_ERROR("%s: failed to compute K-shift\n", __func__);882 return updated;883 }884 885 updated = true;886 }887 888 for (uint32_t s = 0; s < n_stream; ++s) {889 auto & cells = v_cells[s];890 891 cells.reset_shift();892 }893 }894 895 return updated;896}897 898llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, bool cont) const {899 900 if (debug > 0) {901 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {902 const auto seq_id = ubatch.seq_id_unq[s];903 const auto stream_id = seq_to_stream[seq_id];904 const auto & cells = v_cells[stream_id];905 const uint32_t head_cur = v_heads[stream_id];906 907 LLAMA_LOG_DEBUG("%s: stream[%d], n = %5d, used = %5d, head = %5d, size = %5d, n_swa = %5d\n",908 __func__, stream_id, cells.used_max_p1(), cells.get_used(), head_cur, get_size(), n_swa);909 910 if ((debug == 2 && n_swa > 0) || debug > 2) {911 std::string ss;912 for (uint32_t i = 0; i < cells.size(); ++i) {913 if (cells.is_empty(i)) {914 ss += '.';915 } else {916 assert(cells.seq_count(i) >= 1);917 918 if (cells.seq_count(i) == 1) {919 ss += std::to_string(cells.seq_get(i));920 } else {921 ss += 'M';922 }923 }924 if (i%256 == 255) {925 ss += " *";926 ss += '\n';927 }928 }929 LLAMA_LOG_DEBUG("\n%s\n", ss.c_str());930 }931 932 if ((debug == 2 && n_swa > 0) || debug > 2) {933 std::string ss;934 for (uint32_t i = 0; i < cells.size(); ++i) {935 std::string cur;936 if (cells.is_empty(i)) {937 cur = '.';938 } else {939 cur = std::to_string(cells.pos_get(i));940 }941 const int n = cur.size();942 for (int j = 0; j < 5 - n; ++j) {943 cur += ' ';944 }945 ss += cur;946 if (i%256 == 255) {947 ss += " *";948 }949 if (i%64 == 63) {950 ss += '\n';951 }952 }953 LLAMA_LOG_DEBUG("\n%s\n", ss.c_str());954 }955 956 for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {957 if (cells.seq_pos_min(s) < 0) {958 continue;959 }960 961 LLAMA_LOG_DEBUG("%s: stream[%d] min[%d] = %5d, max[%d] = %5d\n", __func__, stream_id, s, cells.seq_pos_min(s), s, cells.seq_pos_max(s));962 }963 }964 }965 966 uint32_t n_tokens = ubatch.n_tokens;967 uint32_t n_seqs = 1;968 969 if (n_stream > 1) {970 GGML_ASSERT(n_tokens % ubatch.n_seqs_unq == 0);971 972 n_seqs = ubatch.n_seqs_unq;973 n_tokens = n_tokens / n_seqs;974 }975 976 slot_info res = {977 /*.s0 =*/ LLAMA_MAX_SEQ,978 /*.s1 =*/ 0,979 /*.strm =*/ { },980 /*.idxs =*/ { },981 };982 983 res.resize(n_seqs);984 985 for (uint32_t s = 0; s < n_seqs; ++s) {986 const auto seq_id = ubatch.seq_id_unq[s];987 988 if (n_stream > 1) {989 GGML_ASSERT(ubatch.n_seq_id[s*n_tokens] == 1);990 GGML_ASSERT(ubatch.seq_id [s*n_tokens][0] == seq_id);991 }992 993 res.s0 = std::min<uint32_t>(res.s0, seq_to_stream[seq_id]);994 res.s1 = std::max<uint32_t>(res.s1, seq_to_stream[seq_id]);995 996 res.strm[s] = seq_to_stream[seq_id];997 res.idxs[s].reserve(n_tokens);998 999 const auto & cells = v_cells[seq_to_stream[seq_id]];1000 1001 uint32_t head_cur = v_heads[seq_to_stream[seq_id]];1002 1003 // if we have enough unused cells before the current head ->1004 // better to start searching from the beginning of the cache, hoping to fill it1005 if (head_cur > cells.get_used() + 2*n_tokens) {1006 head_cur = 0;1007 }1008 1009 if (n_tokens > cells.size()) {1010 LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());1011 return { };1012 }1013 1014 uint32_t n_tested = 0;1015 1016 // for continuous slots, we test that all tokens in the ubatch fit, starting from the current head1017 // for non-continuous slots, we test the tokens one by one1018 const uint32_t n_test = cont ? n_tokens : 1;1019 1020 while (true) {1021 if (head_cur + n_test > cells.size()) {1022 n_tested += cells.size() - head_cur;1023 head_cur = 0;1024 continue;1025 }1026 1027 for (uint32_t i = 0; i < n_test; i++) {1028 const auto idx = head_cur;1029 1030 head_cur++;1031 n_tested++;1032 1033 //const llama_pos pos = ubatch.pos[i];1034 //const llama_seq_id seq_id = ubatch.seq_id[i][0];1035 1036 // can we use this cell? either:1037 // - the cell is empty1038 // - the cell is occupied only by one sequence:1039 // - (disabled) mask causally, if the sequence is the same as the one we are inserting1040 // - mask SWA, using current max pos for that sequence in the cache1041 // always insert in the cell with minimum pos1042 bool can_use = cells.is_empty(idx);1043 1044 if (!can_use && cells.seq_count(idx) == 1) {1045 const llama_pos pos_cell = cells.pos_get(idx);1046 1047 // (disabled) causal mask1048 // note: it's better to purge any "future" tokens beforehand1049 //if (cells.seq_has(idx, seq_id)) {1050 // can_use = pos_cell >= pos;1051 //}1052 1053 if (!can_use) {1054 const llama_seq_id seq_id_cell = cells.seq_get(idx);1055 1056 // SWA mask1057 if (llama_hparams::is_masked_swa(n_swa, swa_type, pos_cell, cells.seq_pos_max(seq_id_cell) + 1)) {1058 can_use = true;1059 }1060 }1061 }1062 1063 if (can_use) {1064 res.idxs[s].push_back(idx);1065 } else {1066 if (cont) {1067 break;1068 }1069 }1070 }1071 1072 if (res.idxs[s].size() == n_tokens) {1073 break;1074 }1075 1076 if (cont) {1077 res.idxs[s].clear();1078 }1079 1080 if (n_tested >= cells.size()) {1081 //LLAMA_LOG_ERROR("%s: failed to find a slot for %d tokens\n", __func__, n_tokens);1082 return { };1083 }1084 }1085 1086 // we didn't find a suitable slot - return empty result1087 if (res.idxs[s].size() < n_tokens) {1088 return { };1089 }1090 }1091 1092 assert(res.s1 >= res.s0);1093 1094 return res;1095}1096 1097void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) {1098 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]1099 if (other) {1100 return;1101 }1102 1103 // keep track of the max sequence position that we would overwrite with this ubatch1104 // for non-SWA cache, this would be always empty1105 llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ];1106 for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {1107 seq_pos_max_rm[s] = -1;1108 }1109 1110 assert(ubatch.n_tokens == sinfo.n_stream()*sinfo.size());1111 1112 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {1113 for (uint32_t ii = 0; ii < sinfo.size(); ++ii) {1114 const uint32_t i = s*sinfo.size() + ii;1115 1116 auto & cells = v_cells[sinfo.strm[s]];1117 1118 const auto idx = sinfo.idxs[s][ii];1119 1120 if (!cells.is_empty(idx)) {1121 assert(cells.seq_count(idx) == 1);1122 1123 const llama_seq_id seq_id = cells.seq_get(idx);1124 const llama_pos pos = cells.pos_get(idx);1125 1126 seq_pos_max_rm[seq_id] = std::max(seq_pos_max_rm[seq_id], pos);1127 1128 cells.rm(idx);1129 }1130 1131 cells.pos_set(idx, ubatch.pos[i]);1132 1133 if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0) {1134 llama_kv_cell_ext ext;1135 1136 if (ubatch.is_pos_2d()) {1137 ext.x = ubatch.pos[i + ubatch.n_tokens*2];1138 ext.y = ubatch.pos[i + ubatch.n_tokens];1139 }1140 1141 if (ubatch.token) {1142 ext.tok = ubatch.token[i];1143 } else if (hparams.ple_n_heads > 0) {1144 // embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers1145 // TODO @ngxson : check if we can do the same as gemma 3n / gemma 41146 ext.tok = hparams.ple_image_token_id != 01147 ? (llama_token) hparams.ple_image_token_id1148 : (llama_token) hparams.ple_eos_token_id;1149 }1150 1151 cells.ext_set(idx, ext);1152 }1153 1154 for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) {1155 cells.seq_add(idx, ubatch.seq_id[i][s]);1156 }1157 }1158 }1159 1160 // note: we want to preserve the invariant that all positions between [pos_min, pos_max] for each sequence1161 // will be present in the cache. so we have to purge any position which is less than those we would overwrite1162 // ref: https://github.com/ggml-org/llama.cpp/pull/13746#issuecomment-29160570921163 for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {1164 if (seq_pos_max_rm[s] == -1) {1165 continue;1166 }1167 1168 GGML_ASSERT(s < seq_to_stream.size());1169 1170 auto & cells = v_cells[seq_to_stream[s]];1171 1172 if (cells.seq_pos_min(s) <= seq_pos_max_rm[s]) {1173 LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",1174 __func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);1175 1176 seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);1177 }1178 }1179 1180 // move the head at the end of the slot1181 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {1182 auto & head = v_heads[sinfo.strm[s]];1183 1184 head = sinfo.idxs[s].back() + 1;1185 }1186}1187 1188bool llama_kv_cache::get_can_shift() const {1189 // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot.1190 if (model.arch == LLM_ARCH_STEP35) {1191 return false;1192 }1193 if (hparams.n_pos_per_embd() > 1) {1194 return false;1195 }1196 return true;1197}1198 1199uint32_t llama_kv_cache::get_size() const {1200 const auto & cells = v_cells[seq_to_stream[0]];