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