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