llama-kv-cache.cpp 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718
  1. #include "llama-kv-cache.h"
  2. #include "llama-impl.h"
  3. #include "llama-batch.h"
  4. #include "llama-cparams.h"
  5. #include "llama-model.h"
  6. #include <algorithm>
  7. #include <limits>
  8. #include <map>
  9. static const llama_kv_cache_slot_info llama_kv_cache_slot_info_failed{false};
  10. uint32_t llama_kv_cache_get_padding(const struct llama_cparams & cparams) {
  11. // the FA kernels require padding to avoid extra runtime boundary checks
  12. return cparams.flash_attn ? 256u : 32u;
  13. }
  14. bool llama_kv_cache_init(
  15. struct llama_kv_cache & cache,
  16. const llama_model & model,
  17. const llama_cparams & cparams,
  18. ggml_type type_k,
  19. ggml_type type_v,
  20. uint32_t kv_size,
  21. bool offload) {
  22. const struct llama_hparams & hparams = model.hparams;
  23. const int32_t n_layer = hparams.n_layer;
  24. cache.has_shift = false;
  25. cache.recurrent = llama_model_is_recurrent(&model);
  26. cache.v_trans = !cache.recurrent && !cparams.flash_attn;
  27. cache.can_shift = !cache.recurrent && model.arch != LLM_ARCH_DEEPSEEK2; // not supported due to MLA
  28. LLAMA_LOG_INFO("%s: kv_size = %d, offload = %d, type_k = '%s', type_v = '%s', n_layer = %d, can_shift = %d\n",
  29. __func__, kv_size, offload, ggml_type_name(type_k), ggml_type_name(type_v), n_layer, cache.can_shift);
  30. cache.head = 0;
  31. cache.size = kv_size;
  32. cache.used = 0;
  33. cache.type_k = type_k;
  34. cache.type_v = type_v;
  35. cache.cells.clear();
  36. cache.cells.resize(kv_size);
  37. // create a context for each buffer type
  38. std::map<ggml_backend_buffer_type_t, ggml_context *> ctx_map;
  39. auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {
  40. auto it = ctx_map.find(buft);
  41. if (it == ctx_map.end()) {
  42. struct ggml_init_params params = {
  43. /*.mem_size =*/ size_t(2u*n_layer*ggml_tensor_overhead()),
  44. /*.mem_buffer =*/ NULL,
  45. /*.no_alloc =*/ true,
  46. };
  47. ggml_context * ctx = ggml_init(params);
  48. if (!ctx) {
  49. return nullptr;
  50. }
  51. ctx_map[buft] = ctx;
  52. cache.ctxs.emplace_back(ctx);
  53. return ctx;
  54. }
  55. return it->second;
  56. };
  57. cache.k_l.reserve(n_layer);
  58. cache.v_l.reserve(n_layer);
  59. for (int i = 0; i < n_layer; i++) {
  60. const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i) + hparams.n_embd_k_s();
  61. const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i) + hparams.n_embd_v_s();
  62. LLAMA_LOG_DEBUG("%s: layer %d: n_embd_k_gqa = %d, n_embd_v_gqa = %d\n", __func__, i, n_embd_k_gqa, n_embd_v_gqa);
  63. ggml_backend_buffer_type_t buft;
  64. if (offload) {
  65. auto * dev = model.dev_layer(i);
  66. buft = ggml_backend_dev_buffer_type(dev);
  67. } else {
  68. buft = ggml_backend_cpu_buffer_type();
  69. }
  70. ggml_context * ctx = ctx_for_buft(buft);
  71. if (!ctx) {
  72. LLAMA_LOG_ERROR("%s: failed to create ggml context for kv cache\n", __func__);
  73. return false;
  74. }
  75. ggml_tensor * k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size);
  76. ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, n_embd_v_gqa*kv_size);
  77. ggml_format_name(k, "cache_k_l%d", i);
  78. ggml_format_name(v, "cache_v_l%d", i);
  79. cache.k_l.push_back(k);
  80. cache.v_l.push_back(v);
  81. }
  82. // allocate tensors and initialize the buffers to avoid NaNs in the padding
  83. for (auto it : ctx_map) {
  84. auto * buft = it.first;
  85. auto * ctx = it.second;
  86. ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft);
  87. if (!buf) {
  88. LLAMA_LOG_ERROR("%s: failed to allocate buffer for kv cache\n", __func__);
  89. return false;
  90. }
  91. ggml_backend_buffer_clear(buf, 0);
  92. 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);
  93. cache.bufs.emplace_back(buf);
  94. }
  95. return true;
  96. }
  97. struct llama_kv_cache_slot_info llama_kv_cache_find_slot(
  98. struct llama_kv_cache & cache,
  99. const struct llama_ubatch & ubatch) {
  100. const uint32_t n_tokens = ubatch.n_tokens;
  101. const uint32_t n_seqs = ubatch.n_seqs;
  102. const uint32_t n_seq_tokens = ubatch.n_seq_tokens;
  103. if (cache.recurrent) {
  104. // For recurrent state architectures (like Mamba or RWKV),
  105. // each cache cell can store the state for a whole sequence.
  106. // A slot should be always be contiguous.
  107. // can only process batches with an equal number of new tokens in each sequence
  108. GGML_ASSERT(ubatch.equal_seqs);
  109. int32_t min = cache.size - 1;
  110. int32_t max = 0;
  111. // everything should fit if all seq_ids are smaller than the max
  112. for (uint32_t s = 0; s < n_seqs; ++s) {
  113. const uint32_t n_seq_id = ubatch.n_seq_id[s];
  114. for (uint32_t j = 0; j < n_seq_id; ++j) {
  115. const llama_seq_id seq_id = ubatch.seq_id[s][j];
  116. if (seq_id < 0 || (uint32_t) seq_id >= cache.size) {
  117. // too big seq_id
  118. // TODO: would it be possible to resize the cache instead?
  119. LLAMA_LOG_ERROR("%s: seq_id=%d >= n_seq_max=%d Try using a bigger --parallel value\n", __func__, seq_id, cache.size);
  120. return llama_kv_cache_slot_info_failed;
  121. }
  122. if (j > 0) {
  123. llama_kv_cell & seq = cache.cells[seq_id];
  124. if (seq.tail >= 0) {
  125. llama_kv_cell & cell = cache.cells[seq.tail];
  126. // clear cells from seq_ids that become shared
  127. // (should not normally happen, but let's handle it anyway)
  128. cell.seq_id.erase(seq_id);
  129. seq.tail = -1;
  130. if (cell.seq_id.empty()) {
  131. cell.pos = -1;
  132. cell.src = -1;
  133. cache.used -= 1;
  134. }
  135. }
  136. }
  137. }
  138. }
  139. #ifndef NDEBUG
  140. {
  141. std::vector<int32_t> tails_verif;
  142. tails_verif.assign(cache.size, -1);
  143. for (uint32_t i = 0; i < cache.size; ++i) {
  144. llama_kv_cell & cell = cache.cells[i];
  145. for (llama_seq_id seq_id : cell.seq_id) {
  146. if (tails_verif[seq_id] != -1) {
  147. LLAMA_LOG_ERROR("%s: duplicate tail for seq_id %d in cell %d and %d\n", __func__, seq_id, i, tails_verif[seq_id]);
  148. }
  149. tails_verif[seq_id] = i;
  150. }
  151. }
  152. for (uint32_t i = 0; i < cache.size; ++i) {
  153. if (tails_verif[i] != cache.cells[i].tail) {
  154. LLAMA_LOG_ERROR("%s: wrong tail for seq_id %d, (%d instead of %d)\n", __func__, i, cache.cells[i].tail, tails_verif[i]);
  155. }
  156. }
  157. }
  158. #endif
  159. // find next empty cell
  160. uint32_t next_empty_cell = cache.head;
  161. for (uint32_t i = 0; i < cache.size; ++i) {
  162. if (next_empty_cell >= cache.size) { next_empty_cell -= cache.size; }
  163. llama_kv_cell & cell = cache.cells[next_empty_cell];
  164. if (cell.is_empty()) { break; }
  165. next_empty_cell += 1;
  166. }
  167. // find usable cell range
  168. for (uint32_t s = 0; s < n_seqs; ++s) {
  169. const llama_seq_id seq_id = ubatch.seq_id[s][0];
  170. llama_kv_cell & seq_meta = cache.cells[seq_id];
  171. bool has_cell = false;
  172. if (seq_meta.tail >= 0) {
  173. llama_kv_cell & cell = cache.cells[seq_meta.tail];
  174. GGML_ASSERT(cell.has_seq_id(seq_id));
  175. // does this seq_id "own" the cell?
  176. if (cell.seq_id.size() == 1) { has_cell = true; }
  177. }
  178. if (!has_cell) {
  179. llama_kv_cell & empty_cell = cache.cells[next_empty_cell];
  180. GGML_ASSERT(empty_cell.is_empty());
  181. // copy old tail into the empty cell
  182. if (seq_meta.tail >= 0) {
  183. llama_kv_cell & orig_cell = cache.cells[seq_meta.tail];
  184. empty_cell.pos = orig_cell.pos;
  185. empty_cell.src = orig_cell.src;
  186. orig_cell.seq_id.erase(seq_id);
  187. empty_cell.seq_id.insert(seq_id); // will be overwritten
  188. }
  189. seq_meta.tail = next_empty_cell;
  190. // find next empty cell
  191. if (s + 1 < n_seqs) {
  192. next_empty_cell += 1;
  193. for (uint32_t i = 0; i < cache.size; ++i) {
  194. if (next_empty_cell >= cache.size) { next_empty_cell -= cache.size; }
  195. llama_kv_cell & cell = cache.cells[next_empty_cell];
  196. if (cell.is_empty()) { break; }
  197. next_empty_cell += 1;
  198. }
  199. }
  200. }
  201. if (min > seq_meta.tail) { min = seq_meta.tail; }
  202. if (max < seq_meta.tail) { max = seq_meta.tail; }
  203. }
  204. // gather and re-order
  205. for (uint32_t s = 0; s < n_seqs; ++s) {
  206. int32_t dst_id = s + min;
  207. int32_t src_id = cache.cells[ubatch.seq_id[s][0]].tail;
  208. if (dst_id != src_id) {
  209. llama_kv_cell & dst_cell = cache.cells[dst_id];
  210. llama_kv_cell & src_cell = cache.cells[src_id];
  211. std::swap(dst_cell.pos, src_cell.pos);
  212. std::swap(dst_cell.src, src_cell.src);
  213. std::swap(dst_cell.seq_id, src_cell.seq_id);
  214. // swap tails (assuming they NEVER overlap)
  215. for (const llama_seq_id seq_id : src_cell.seq_id) {
  216. cache.cells[seq_id].tail = src_id;
  217. }
  218. for (const llama_seq_id seq_id : dst_cell.seq_id) {
  219. cache.cells[seq_id].tail = dst_id;
  220. }
  221. }
  222. }
  223. // update the pos of the used seqs
  224. for (uint32_t s = 0; s < n_seqs; ++s) {
  225. const llama_pos last_pos = ubatch.pos[n_seq_tokens * s + n_seq_tokens - 1];
  226. int32_t cell_id = s + min;
  227. llama_kv_cell & cell = cache.cells[cell_id];
  228. if (cell.pos >= 0 && last_pos != cell.pos + (llama_pos) n_seq_tokens) {
  229. // What should happen when the pos backtracks or skips a value?
  230. // Clearing the state mid-batch would require special-casing which isn't done.
  231. LLAMA_LOG_WARN("%s: non-consecutive token position %d after %d for sequence %d with %u new tokens\n",
  232. __func__, last_pos, cell.pos, ubatch.seq_id[s][0], n_seq_tokens);
  233. }
  234. cell.pos = last_pos;
  235. cell.seq_id.clear();
  236. for (int32_t j = 0; j < ubatch.n_seq_id[s]; ++j) {
  237. const llama_seq_id seq_id = ubatch.seq_id[s][j];
  238. cell.seq_id.insert(seq_id);
  239. cache.cells[seq_id].tail = cell_id;
  240. }
  241. }
  242. // allow getting the range of used cells, from head to head + n
  243. cache.head = min;
  244. cache.n = max - min + 1;
  245. cache.used = std::count_if(cache.cells.begin(), cache.cells.end(),
  246. [](const llama_kv_cell& cell){ return !cell.is_empty(); });
  247. // sanity check
  248. return llama_kv_cache_slot_info(cache.n >= n_seqs);
  249. }
  250. // otherwise, one cell per token.
  251. if (n_tokens > cache.size) {
  252. LLAMA_LOG_ERROR("%s: n_tokens=%d > cache.size=%d\n", __func__, n_tokens, cache.size);
  253. return llama_kv_cache_slot_info_failed;
  254. }
  255. uint32_t n_tested = 0;
  256. while (true) {
  257. if (cache.head + n_tokens > cache.size) {
  258. n_tested += cache.size - cache.head;
  259. cache.head = 0;
  260. continue;
  261. }
  262. bool found = true;
  263. for (uint32_t i = 0; i < n_tokens; i++) {
  264. if (cache.cells[cache.head + i].pos >= 0) {
  265. found = false;
  266. cache.head += i + 1;
  267. n_tested += i + 1;
  268. break;
  269. }
  270. }
  271. if (found) {
  272. break;
  273. }
  274. if (n_tested >= cache.size) {
  275. //LLAMA_LOG_ERROR("%s: failed to find a slot for %d tokens\n", __func__, n_tokens);
  276. return llama_kv_cache_slot_info_failed;
  277. }
  278. }
  279. for (uint32_t s = 0; s < n_seqs; s++) {
  280. for (uint32_t i = 0; i < n_seq_tokens; ++i) {
  281. uint32_t k = s*n_seq_tokens + i;
  282. cache.cells[cache.head + k].pos = ubatch.pos[k];
  283. for (int32_t j = 0; j < ubatch.n_seq_id[s]; j++) {
  284. cache.cells[cache.head + k].seq_id.insert(ubatch.seq_id[s][j]);
  285. }
  286. }
  287. }
  288. cache.used += n_tokens;
  289. return llama_kv_cache_slot_info(cache.head, cache.head + n_tokens);
  290. }
  291. uint32_t llama_kv_cache_cell_max(const struct llama_kv_cache & cache) {
  292. for (uint32_t i = cache.size; i > 0; --i) {
  293. const llama_kv_cell & cell = cache.cells[i - 1];
  294. if (cell.pos >= 0 && !cell.is_empty()) {
  295. return i;
  296. }
  297. }
  298. return 0;
  299. }
  300. void llama_kv_cache_clear(struct llama_kv_cache & cache) {
  301. for (int32_t i = 0; i < (int32_t) cache.size; ++i) {
  302. cache.cells[i].pos = -1;
  303. cache.cells[i].seq_id.clear();
  304. cache.cells[i].src = -1;
  305. cache.cells[i].tail = -1;
  306. }
  307. cache.head = 0;
  308. cache.used = 0;
  309. for (auto & buf : cache.bufs) {
  310. ggml_backend_buffer_clear(buf.get(), 0);
  311. }
  312. }
  313. bool llama_kv_cache_seq_rm(
  314. struct llama_kv_cache & cache,
  315. llama_seq_id seq_id,
  316. llama_pos p0,
  317. llama_pos p1) {
  318. uint32_t new_head = cache.size;
  319. if (p0 < 0) p0 = 0;
  320. if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
  321. // models like Mamba or RWKV can't have a state partially erased
  322. if (cache.recurrent) {
  323. if (seq_id >= (int64_t) cache.size) {
  324. // could be fatal
  325. return false;
  326. }
  327. if (0 <= seq_id) {
  328. int32_t & tail_id = cache.cells[seq_id].tail;
  329. if (tail_id >= 0) {
  330. const llama_kv_cell & cell = cache.cells[tail_id];
  331. // partial intersection is invalid
  332. if ((0 < p0 && p0 <= cell.pos) || (0 < p1 && p1 <= cell.pos)) {
  333. return false;
  334. }
  335. // invalidate tails which will be cleared
  336. if (p0 <= cell.pos && cell.pos < p1) {
  337. tail_id = -1;
  338. }
  339. }
  340. } else {
  341. // seq_id is negative, then the range should include everything or nothing
  342. if (p0 != p1 && (p0 != 0 || p1 != std::numeric_limits<llama_pos>::max())) {
  343. return false;
  344. }
  345. }
  346. }
  347. for (uint32_t i = 0; i < cache.size; ++i) {
  348. if (cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
  349. if (seq_id < 0) {
  350. cache.cells[i].seq_id.clear();
  351. } else if (cache.cells[i].has_seq_id(seq_id)) {
  352. cache.cells[i].seq_id.erase(seq_id);
  353. } else {
  354. continue;
  355. }
  356. if (cache.cells[i].is_empty()) {
  357. // keep count of the number of used cells
  358. if (cache.cells[i].pos >= 0) cache.used--;
  359. cache.cells[i].pos = -1;
  360. cache.cells[i].src = -1;
  361. if (new_head == cache.size) new_head = i;
  362. }
  363. }
  364. }
  365. // If we freed up a slot, set head to it so searching can start there.
  366. if (new_head != cache.size && new_head < cache.head) cache.head = new_head;
  367. return true;
  368. }
  369. void llama_kv_cache_seq_cp(
  370. struct llama_kv_cache & cache,
  371. llama_seq_id seq_id_src,
  372. llama_seq_id seq_id_dst,
  373. llama_pos p0,
  374. llama_pos p1) {
  375. if (p0 < 0) p0 = 0;
  376. if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
  377. if (cache.recurrent) {
  378. if ((uint32_t) seq_id_dst < cache.size && (uint32_t) seq_id_src < cache.size) {
  379. llama_kv_cell & tail_src = cache.cells[seq_id_src];
  380. llama_kv_cell & tail_dst = cache.cells[seq_id_dst];
  381. if (tail_dst.tail >= 0) {
  382. // clear destination seq_id if it wasn't empty
  383. llama_kv_cell & cell_dst = cache.cells[tail_dst.tail];
  384. cell_dst.seq_id.erase(seq_id_dst);
  385. tail_dst.tail = -1;
  386. if (cell_dst.seq_id.empty()) {
  387. cell_dst.pos = -1;
  388. cell_dst.delta = -1;
  389. cell_dst.src = -1;
  390. cache.used -= 1;
  391. }
  392. }
  393. if (tail_src.tail >= 0) {
  394. llama_kv_cell & cell_src = cache.cells[tail_src.tail];
  395. cell_src.seq_id.insert(seq_id_dst);
  396. tail_dst.tail = tail_src.tail;
  397. }
  398. }
  399. return;
  400. }
  401. // otherwise, this is the KV cache of a Transformer-like model
  402. cache.head = 0;
  403. for (uint32_t i = 0; i < cache.size; ++i) {
  404. if (cache.cells[i].has_seq_id(seq_id_src) && cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
  405. cache.cells[i].seq_id.insert(seq_id_dst);
  406. }
  407. }
  408. }
  409. void llama_kv_cache_seq_keep(struct llama_kv_cache & cache, llama_seq_id seq_id) {
  410. uint32_t new_head = cache.size;
  411. for (uint32_t i = 0; i < cache.size; ++i) {
  412. if (cache.recurrent && (llama_seq_id) i != seq_id) {
  413. cache.cells[i].tail = -1;
  414. }
  415. if (!cache.cells[i].has_seq_id(seq_id)) {
  416. if (cache.cells[i].pos >= 0) cache.used--;
  417. cache.cells[i].pos = -1;
  418. cache.cells[i].src = -1;
  419. cache.cells[i].seq_id.clear();
  420. if (new_head == cache.size) new_head = i;
  421. } else {
  422. cache.cells[i].seq_id.clear();
  423. cache.cells[i].seq_id.insert(seq_id);
  424. }
  425. }
  426. // If we freed up a slot, set head to it so searching can start there.
  427. if (new_head != cache.size && new_head < cache.head) cache.head = new_head;
  428. }
  429. void llama_kv_cache_seq_add(
  430. struct llama_kv_cache & cache,
  431. llama_seq_id seq_id,
  432. llama_pos p0,
  433. llama_pos p1,
  434. llama_pos delta) {
  435. uint32_t new_head = cache.size;
  436. if (p0 < 0) p0 = 0;
  437. if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
  438. // If there is no range then return early to avoid looping over the cache.
  439. if (p0 == p1) return;
  440. if (cache.recurrent) {
  441. // for Mamba-like or RWKV models, only the pos needs to be shifted
  442. if (0 <= seq_id && seq_id < (int64_t) cache.size) {
  443. const int32_t tail_id = cache.cells[seq_id].tail;
  444. if (tail_id >= 0) {
  445. llama_kv_cell & cell = cache.cells[tail_id];
  446. if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {
  447. cell.pos += delta;
  448. }
  449. }
  450. }
  451. return;
  452. }
  453. for (uint32_t i = 0; i < cache.size; ++i) {
  454. if (cache.cells[i].has_seq_id(seq_id) && cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
  455. cache.has_shift = true;
  456. cache.cells[i].pos += delta;
  457. cache.cells[i].delta += delta;
  458. if (cache.cells[i].pos < 0) {
  459. if (!cache.cells[i].is_empty()) {
  460. cache.used--;
  461. }
  462. cache.cells[i].pos = -1;
  463. cache.cells[i].seq_id.clear();
  464. if (new_head == cache.size) {
  465. new_head = i;
  466. }
  467. }
  468. }
  469. }
  470. // If we freed up a slot, set head to it so searching can start there.
  471. // Otherwise we just start the next search from the beginning.
  472. cache.head = new_head != cache.size ? new_head : 0;
  473. }
  474. void llama_kv_cache_seq_div(
  475. struct llama_kv_cache & cache,
  476. llama_seq_id seq_id,
  477. llama_pos p0,
  478. llama_pos p1,
  479. int d) {
  480. if (p0 < 0) p0 = 0;
  481. if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
  482. // If there is no range then return early to avoid looping over the cache.
  483. if (p0 == p1) return;
  484. if (cache.recurrent) {
  485. // for Mamba-like or RWKV models, only the pos needs to be changed
  486. if (0 <= seq_id && seq_id < (int64_t) cache.size) {
  487. const int32_t tail_id = cache.cells[seq_id].tail;
  488. if (tail_id >= 0) {
  489. llama_kv_cell & cell = cache.cells[tail_id];
  490. if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {
  491. cell.pos /= d;
  492. }
  493. }
  494. }
  495. return;
  496. }
  497. for (uint32_t i = 0; i < cache.size; ++i) {
  498. if (cache.cells[i].has_seq_id(seq_id) && cache.cells[i].pos >= p0 && cache.cells[i].pos < p1) {
  499. cache.has_shift = true;
  500. {
  501. llama_pos p_old = cache.cells[i].pos;
  502. cache.cells[i].pos /= d;
  503. cache.cells[i].delta += cache.cells[i].pos - p_old;
  504. }
  505. }
  506. }
  507. }
  508. llama_pos llama_kv_cache_seq_pos_max(struct llama_kv_cache & cache, llama_seq_id seq_id) {
  509. llama_pos result = 0;
  510. for (uint32_t i = 0; i < cache.size; ++i) {
  511. if (cache.cells[i].has_seq_id(seq_id)) {
  512. result = std::max(result, cache.cells[i].pos);
  513. }
  514. }
  515. return result;
  516. }
  517. void llama_kv_cache_defrag(struct llama_kv_cache & cache) {
  518. if (!cache.recurrent) {
  519. cache.do_defrag = true;
  520. }
  521. }
  522. int32_t llama_get_kv_cache_token_count(const struct llama_kv_cache & kv) {
  523. int result = 0;
  524. for (uint32_t i = 0; i < kv.size; i++) {
  525. result += kv.cells[i].seq_id.size();
  526. }
  527. return result;
  528. }
  529. int32_t llama_get_kv_cache_used_cells(const struct llama_kv_cache & kv) {
  530. return kv.used;
  531. }
  532. bool llama_kv_cache_can_shift(const struct llama_kv_cache & kv) {
  533. return kv.can_shift;
  534. }
  535. //
  536. // kv cache view
  537. //
  538. struct llama_kv_cache_view llama_kv_cache_view_init(const struct llama_kv_cache & kv, int32_t n_seq_max) {
  539. struct llama_kv_cache_view result = {
  540. /*.n_cells = */ 0,
  541. /*.n_seq_max = */ n_seq_max,
  542. /*.token_count = */ 0,
  543. /*.used_cells = */ llama_get_kv_cache_used_cells(kv),
  544. /*.max_contiguous = */ 0,
  545. /*.max_contiguous_idx = */ -1,
  546. /*.cells = */ nullptr,
  547. /*.cells_sequences = */ nullptr,
  548. };
  549. return result;
  550. }
  551. void llama_kv_cache_view_free(struct llama_kv_cache_view * view) {
  552. if (view->cells != nullptr) {
  553. free(view->cells);
  554. view->cells = nullptr;
  555. }
  556. if (view->cells_sequences != nullptr) {
  557. free(view->cells_sequences);
  558. view->cells_sequences = nullptr;
  559. }
  560. }
  561. void llama_kv_cache_view_update(struct llama_kv_cache_view * view, const struct llama_kv_cache & kv) {
  562. if (uint32_t(view->n_cells) < kv.size || view->cells == nullptr) {
  563. view->n_cells = int32_t(kv.size);
  564. void * p = realloc(view->cells, sizeof(struct llama_kv_cache_view_cell) * view->n_cells);
  565. GGML_ASSERT(p != nullptr && "Failed to alloc kv_cache_view cells");
  566. view->cells = (struct llama_kv_cache_view_cell *)p;
  567. p = realloc(view->cells_sequences, sizeof(llama_seq_id) * view->n_seq_max * view->n_cells);
  568. GGML_ASSERT(p != nullptr && "Failed to alloc kv_cache_view cells sequences");
  569. view->cells_sequences = (llama_seq_id *)p;
  570. }
  571. const std::vector<llama_kv_cell> & kv_cells = kv.cells;
  572. llama_kv_cache_view_cell * c_curr = view->cells;
  573. llama_seq_id * cs_curr = view->cells_sequences;
  574. int32_t used_cells = 0;
  575. int32_t token_count = 0;
  576. int32_t curr_contig_idx = -1;
  577. uint32_t max_contig = 0;
  578. int32_t max_contig_idx = -1;
  579. for (int32_t i = 0; i < int32_t(kv.size); i++, c_curr++, cs_curr += view->n_seq_max) {
  580. const size_t curr_size = kv_cells[i].seq_id.size();
  581. token_count += curr_size;
  582. c_curr->pos = kv_cells[i].pos + kv_cells[i].delta;
  583. if (curr_size > 0) {
  584. if (curr_contig_idx >= 0 && uint32_t(i - curr_contig_idx) > max_contig) {
  585. max_contig = i - curr_contig_idx;
  586. max_contig_idx = curr_contig_idx;
  587. }
  588. curr_contig_idx = -1;
  589. } else if (curr_contig_idx < 0) {
  590. curr_contig_idx = i;
  591. }
  592. int seq_idx = 0;
  593. for (const llama_seq_id it : kv_cells[i].seq_id) {
  594. if (seq_idx >= view->n_seq_max) {
  595. break;
  596. }
  597. cs_curr[seq_idx] = it;
  598. seq_idx++;
  599. }
  600. if (seq_idx != 0) {
  601. used_cells++;
  602. }
  603. for (; seq_idx < view->n_seq_max; seq_idx++) {
  604. cs_curr[seq_idx] = -1;
  605. }
  606. }
  607. if (curr_contig_idx >= 0 && kv_cells.size() - curr_contig_idx > max_contig) {
  608. max_contig_idx = curr_contig_idx;
  609. max_contig = kv_cells.size() - curr_contig_idx;
  610. }
  611. view->max_contiguous = max_contig;
  612. view->max_contiguous_idx = max_contig_idx;
  613. view->token_count = token_count;
  614. view->used_cells = used_cells;
  615. if (uint32_t(used_cells) != kv.used) {
  616. LLAMA_LOG_ERROR("%s: used cells mismatch. kv_cache says %d but we calculated %d\n",
  617. __func__, kv.used, used_cells);
  618. }
  619. }