Bug Summary

File:root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp
Warning:line 981, column 9
Division by zero

Annotated Source Code

Press '?' to see keyboard shortcuts

clang -cc1 -cc1 -triple x86_64-pc-linux-gnu -O2 -analyze -disable-free -clear-ast-before-backend -disable-llvm-verifier -discard-value-names -main-file-name llama-kv-cache.cpp -analyzer-checker=core -analyzer-checker=apiModeling -analyzer-checker=unix -analyzer-checker=deadcode -analyzer-checker=cplusplus -analyzer-checker=security.insecureAPI.UncheckedReturn -analyzer-checker=security.insecureAPI.getpw -analyzer-checker=security.insecureAPI.gets -analyzer-checker=security.insecureAPI.mktemp -analyzer-checker=security.insecureAPI.mkstemp -analyzer-checker=security.insecureAPI.vfork -analyzer-checker=nullability.NullPassedToNonnull -analyzer-checker=nullability.NullReturnedFromNonnull -analyzer-output plist -w -setup-static-analyzer -analyzer-config-compatibility-mode=true -mrelocation-model pic -pic-level 2 -fhalf-no-semantic-interposition -mframe-pointer=all -relaxed-aliasing -ffp-contract=off -fno-rounding-math -mconstructor-aliases -funwind-tables=2 -target-cpu x86-64 -target-feature +avx -target-feature +avx2 -target-feature +bmi2 -target-feature +f16c -target-feature +fma -target-feature +sse4.2 -tune-cpu generic -debugger-tuning=gdb -fdebug-compilation-dir=/root/firefox-clang/obj-x86_64-pc-linux-gnu/third_party/llama.cpp -fcoverage-compilation-dir=/root/firefox-clang/obj-x86_64-pc-linux-gnu/third_party/llama.cpp -resource-dir /usr/lib/llvm-23/lib/clang/23 -include /root/firefox-clang/config/gcc_hidden.h -include /root/firefox-clang/obj-x86_64-pc-linux-gnu/mozilla-config.h -I /root/firefox-clang/obj-x86_64-pc-linux-gnu/dist/stl_wrappers -D _GLIBCXX_ASSERTIONS=1 -I /root/firefox-clang/obj-x86_64-pc-linux-gnu/dist/system_wrappers -U _FORTIFY_SOURCE -D _FORTIFY_SOURCE=2 -D DEBUG=1 -D _GNU_SOURCE=1 -D GGML_USE_CPU=1 -D GGML_VERSION="GGML_VERSION" -D GGML_COMMIT="GGML_COMMIT" -D GGML_SHARED=1 -D LLAMA_SHARED=1 -D GGML_BUILD=1 -D LLAMA_BUILD=1 -D GGML_BACKEND_SHARED=1 -D GGML_BACKEND_BUILD=1 -D MOZ_HAS_MOZGLUE -I /root/firefox-clang/third_party/llama.cpp -I /root/firefox-clang/obj-x86_64-pc-linux-gnu/third_party/llama.cpp -I /root/firefox-clang/third_party/llama.cpp/ggml -I /root/firefox-clang/third_party/llama.cpp/ggml/include -I /root/firefox-clang/third_party/llama.cpp/ggml/src -I /root/firefox-clang/third_party/llama.cpp/ggml/src/ggml-cpu -I /root/firefox-clang/third_party/llama.cpp/include -I /root/firefox-clang/third_party/llama.cpp/src -I /root/firefox-clang/obj-x86_64-pc-linux-gnu/dist/include -I /root/firefox-clang/obj-x86_64-pc-linux-gnu/dist/include/nspr -I /root/firefox-clang/obj-x86_64-pc-linux-gnu/dist/include/nss -D MOZILLA_CLIENT -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/16/../../../../include/c++/16 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/16/../../../../include/x86_64-linux-gnu/c++/16 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/16/../../../../include/c++/16/backward -internal-isystem /usr/lib/llvm-23/lib/clang/23/include -internal-isystem /usr/local/include -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/16/../../../../x86_64-linux-gnu/include -internal-externc-isystem /usr/include/x86_64-linux-gnu -internal-externc-isystem /include -internal-externc-isystem /usr/include -Wno-error=pessimizing-move -Wno-error=large-by-value-copy=128 -Wno-error=implicit-int-float-conversion -Wno-error=thread-safety-analysis -Wno-error=tautological-type-limit-compare -Wno-invalid-offsetof -Wno-range-loop-analysis -Wno-deprecated-anon-enum-enum-conversion -Wno-deprecated-enum-enum-conversion -Wno-inline-new-delete -Wno-error=deprecated-declarations -Wno-error=array-bounds -Wno-error=free-nonheap-object -Wno-error=atomic-alignment -Wno-error=deprecated-builtins -Wno-psabi -Wno-error=builtin-macro-redefined -Wno-vla-cxx-extension -Wno-unknown-warning-option -Wno-character-conversion -Wno-sign-compare -Wno-unused-function -Wno-tautological-unsigned-enum-zero-compare -Wno-implicit-fallthrough -Wno-unreachable-code -std=gnu++20 -fdeprecated-macro -ferror-limit 19 -fstrict-flex-arrays=1 -stack-protector 2 -fstack-clash-protection -ftrivial-auto-var-init=pattern -fno-rtti -fgnuc-version=4.2.1 -fno-implicit-modules -fskip-odr-check-in-gmf -fno-sized-deallocation -fno-aligned-allocation -fdiagnostics-absolute-paths -vectorize-loops -vectorize-slp -analyzer-checker optin.performance.Padding -analyzer-output=html -analyzer-config stable-report-filename=true -mllvm -dwarf-linkage-names=Abstract -faddrsig -fdwarf2-cfi-asm -o /tmp/scan-build-2026-09-01-224014-2642839-1 -x c++ /root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp
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
16#include "moz-overrides.h"
17
18static bool ggml_is_power_of_2(int n) {
19 return (n & (n - 1)) == 0;
20}
21
22// orthonormal Walsh-Hadamard rotation matrix
23// note: res^2 == I
24static void ggml_gen_hadamard(ggml_tensor * tensor) {
25 assert(tensor->type == GGML_TYPE_F32)(static_cast <bool> (tensor->type == GGML_TYPE_F32) ?
void (0) : __assert_fail ("tensor->type == GGML_TYPE_F32"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
26
27 const int n = tensor->ne[0];
28
29 assert(ggml_is_power_of_2(n))(static_cast <bool> (ggml_is_power_of_2(n)) ? void (0) :
__assert_fail ("ggml_is_power_of_2(n)", __builtin_FILE (), __builtin_LINE
(), __extension__ __PRETTY_FUNCTION__));
30 assert(tensor->ne[1] == n)(static_cast <bool> (tensor->ne[1] == n) ? void (0) :
__assert_fail ("tensor->ne[1] == n", __builtin_FILE (), __builtin_LINE
(), __extension__ __PRETTY_FUNCTION__));
31 assert(tensor->ne[2] == 1)(static_cast <bool> (tensor->ne[2] == 1) ? void (0) :
__assert_fail ("tensor->ne[2] == 1", __builtin_FILE (), __builtin_LINE
(), __extension__ __PRETTY_FUNCTION__));
32 assert(tensor->ne[3] == 1)(static_cast <bool> (tensor->ne[3] == 1) ? void (0) :
__assert_fail ("tensor->ne[3] == 1", __builtin_FILE (), __builtin_LINE
(), __extension__ __PRETTY_FUNCTION__));
33
34 std::vector<float> data_f32;
35
36 float * data = (float *) tensor->data;
37
38 if (tensor->type != GGML_TYPE_F32) {
39 data_f32.resize(n*n);
40 data = data_f32.data();
41 }
42
43 data[0*n + 0] = 1.0 / sqrtf(n);
44
45 for (int s = 1; s < n; s *= 2) {
46 for (int i = 0; i < s; i++) {
47 for (int j = 0; j < s; j++) {
48 const float val = data[i*n + j];
49
50 data[(i + s)*n + (j )] = val;
51 data[(i )*n + (j + s)] = val;
52 data[(i + s)*n + (j + s)] = -val;
53 }
54 }
55 }
56
57 if (tensor->type != GGML_TYPE_F32) {
58 ggml_quantize_chunk(tensor->type, data, tensor->data, 0, 1, n*n, nullptr);
59 }
60}
61
62static ggml_tensor * ggml_mul_mat_aux(
63 ggml_context * ctx,
64 ggml_tensor * cur,
65 ggml_tensor * rot) {
66 const auto n = rot->ne[0];
67
68 ggml_tensor * res;
69
70 res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n);
71 res = ggml_mul_mat (ctx, rot, res);
72 ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD);
73 res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]);
74
75 return res;
76}
77
78//
79// llama_kv_cache
80//
81
82llama_kv_cache::llama_kv_cache(
83 const llama_model & model,
84 const llama_hparams & hparams,
85 ggml_type type_k,
86 ggml_type type_v,
87 bool v_trans,
88 bool offload,
89 bool unified,
90 uint32_t kv_size,
91 uint32_t n_seq_max,
92 uint32_t n_pad,
93 uint32_t n_swa,
94 llama_swa_type swa_type,
95 llama_memory_t mem_other,
96 const layer_filter_cb & filter,
97 const layer_reuse_cb & reuse,
98 const layer_share_cb & share) :
99 model(model), hparams(hparams), v_trans(v_trans),
100 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),
101 other(static_cast<llama_kv_cache *>(mem_other)),
102 v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
103 v_cells(*v_cells_impl) {
104
105 // shared cells view the source cache's K/V tensors, so the cell count
106 // follows the source allocation: a fitted target can be smaller than the
107 // draft default and oversized views would overflow the source tensors
108 if (other) {
109 const uint32_t size_other = other->get_size();
110 if (kv_size != size_other) {
111 LLAMA_LOG_WARN("%s: kv_size = %u overridden to %u to match the shared source cache\n", __func__, kv_size, size_other)llama_log_internal(GGML_LOG_LEVEL_WARN , "%s: kv_size = %u overridden to %u to match the shared source cache\n"
, __func__, kv_size, size_other);
112 kv_size = size_other;
113 }
114 }
115
116 GGML_ASSERT(kv_size % n_pad == 0)if (!(kv_size % n_pad == 0)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 116, "GGML_ASSERT(%s) failed", "kv_size % n_pad == 0");
117
118 const uint32_t n_layer = hparams.n_layer_all;
119
120 // define a comparator for the buft -> ctx map to ensure that the order is well-defined:
121 struct ggml_backend_buft_comparator {
122 bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {
123 return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;
124 }
125 };
126 std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> ctx_map;
127
128 // create a context for each buffer type
129 auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {
130 auto it = ctx_map.find(buft);
131 if (it == ctx_map.end()) {
132 ggml_init_params params = {
133 /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
134 /*.mem_buffer =*/ NULL__null,
135 /*.no_alloc =*/ true,
136 };
137
138 ggml_context * ctx = ggml_init(params);
139 if (!ctx) {
140 return nullptr;
141 }
142
143 ctx_map.emplace(buft, ctx);
144
145 return ctx;
146 }
147
148 return it->second.get();
149 };
150
151 GGML_ASSERT(n_stream == 1 || n_stream == n_seq_max)if (!(n_stream == 1 || n_stream == n_seq_max)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 151, "GGML_ASSERT(%s) failed", "n_stream == 1 || n_stream == n_seq_max"
);
152
153 v_heads.resize(n_stream);
154 for (uint32_t s = 0; s < n_stream; ++s) {
155 v_heads[s] = 0;
156 }
157
158 v_cells.resize(n_stream);
159 for (uint32_t s = 0; s < n_stream; ++s) {
160 v_cells[s].resize(kv_size);
161 }
162
163 // by default, all sequence ids are mapped to the 0th stream
164 seq_to_stream.resize(LLAMA_MAX_SEQ256, 0);
165
166 if (n_stream > 1) {
167 seq_to_stream.resize(n_stream, 0);
168 for (uint32_t s = 0; s < n_stream; ++s) {
169 seq_to_stream[s] = s;
170 }
171 }
172
173 // [TAG_V_CACHE_VARIABLE]
174 if (v_trans && hparams.is_n_embd_v_gqa_variable()) {
175 LLAMA_LOG_WARN("%s: the V embeddings have different sizes across layers and FA is not enabled - padding V cache to %d\n",llama_log_internal(GGML_LOG_LEVEL_WARN , "%s: the V embeddings have different sizes across layers and FA is not enabled - padding V cache to %d\n"
, __func__, hparams.n_embd_v_gqa_max())
176 __func__, hparams.n_embd_v_gqa_max())llama_log_internal(GGML_LOG_LEVEL_WARN , "%s: the V embeddings have different sizes across layers and FA is not enabled - padding V cache to %d\n"
, __func__, hparams.n_embd_v_gqa_max());
177 }
178
179 const bool is_mla = hparams.is_mla();
180
181 for (uint32_t il = 0; il < n_layer; il++) {
182 if (!hparams.has_kv(il)) {
183 LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: layer %3d: does not have KV cache\n"
, __func__, il);
184 continue;
185 }
186
187 if (filter && !filter(il)) {
188 LLAMA_LOG_DEBUG("%s: layer %3d: filtered\n", __func__, il)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: layer %3d: filtered\n"
, __func__, il);
189 continue;
190 }
191
192 if (share && other) {
193 const int32_t il_share = share(il);
194
195 if (il_share >= 0) {
196 const auto & layer_share = other->layers[other->map_layer_ids[il_share]];
197
198 LLAMA_LOG_WARN("%s: layer %3d: sharing with layer %d. k = %p, v = %p\n", __func__, il, il_share,llama_log_internal(GGML_LOG_LEVEL_WARN , "%s: layer %3d: sharing with layer %d. k = %p, v = %p\n"
, __func__, il, il_share, layer_share.k->data, layer_share
.v->data)
199 layer_share.k->data, layer_share.v->data)llama_log_internal(GGML_LOG_LEVEL_WARN , "%s: layer %3d: sharing with layer %d. k = %p, v = %p\n"
, __func__, il, il_share, layer_share.k->data, layer_share
.v->data);
200
201 map_layer_ids[il] = layers.size();
202
203 layers.push_back(layer_share);
204 layers.back().il = il;
205
206 continue;
207 }
208 }
209
210 if (n_embd_head_k_all == 0) {
211 n_embd_head_k_all = (int32_t) hparams.n_embd_head_k(il);
212 } else if (n_embd_head_k_all > 0 && n_embd_head_k_all != (int32_t) hparams.n_embd_head_k(il)) {
213 n_embd_head_k_all = -1;
214 }
215
216 if (n_embd_head_v_all == 0) {
217 n_embd_head_v_all = (int32_t) hparams.n_embd_head_v(il);
218 } else if (n_embd_head_v_all > 0 && n_embd_head_v_all != (int32_t) hparams.n_embd_head_v(il)) {
219 n_embd_head_v_all = -1;
220 }
221
222 // [TAG_V_CACHE_VARIABLE]
223 const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
224 const uint32_t n_embd_v_gqa = !v_trans ? hparams.n_embd_v_gqa(il) : hparams.n_embd_v_gqa_max();
225
226 const char * dev_name = "CPU";
227
228 ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();
229
230 if (offload) {
231 auto * dev = model.dev_layer(il);
232 buft = ggml_backend_dev_buffer_type(dev);
233
234 dev_name = ggml_backend_dev_name(dev);
235 }
236
237 LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: layer %3d: dev = %s\n"
, __func__, il, dev_name);
238
239 ggml_context * ctx = ctx_for_buft(buft);
240 if (!ctx) {
241 throwabort_with_suppression(); if (false) std::runtime_error("failed to create ggml context for kv cache");
242 }
243
244 const bool has_k = true;
245 const bool has_v = !is_mla;
246
247 ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr;
248 ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr;
249
250 has_k && ggml_format_name(k, "cache_k_l%d", il);
251 has_v && ggml_format_name(v, "cache_v_l%d", il);
252
253 std::vector<ggml_tensor *> k_stream;
254 std::vector<ggml_tensor *> v_stream;
255
256 for (uint32_t s = 0; s < n_stream; ++s) {
257 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);
258 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);
259 }
260
261 map_layer_ids[il] = layers.size();
262
263 layers.push_back({ il, k, v, k_stream, v_stream, });
264 }
265
266 if (reuse) {
267 LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: reusing layers:\n"
, __func__);
268
269 for (uint32_t il = 0; il < n_layer; il++) {
270 const int32_t il_reuse = reuse(il);
271
272 if (il_reuse < 0) {
273 LLAMA_LOG_DEBUG("%s: - layer %3d: no reuse\n", __func__, il)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: - layer %3d: no reuse\n"
, __func__, il);
274 continue;
275 }
276
277 if (filter && !filter(il)) {
278 LLAMA_LOG_DEBUG("%s: - layer %3d: filtered\n", __func__, il)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: - layer %3d: filtered\n"
, __func__, il);
279 continue;
280 }
281
282 GGML_ASSERT(map_layer_ids.find(il_reuse) != map_layer_ids.end())if (!(map_layer_ids.find(il_reuse) != map_layer_ids.end())) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 282, "GGML_ASSERT(%s) failed", "map_layer_ids.find(il_reuse) != map_layer_ids.end()"
);
283
284 map_layer_ids[il] = map_layer_ids[il_reuse];
285
286 LLAMA_LOG_DEBUG("%s: - layer %3d: reuse layer %d, is_swa = %d\n", __func__, il, il_reuse, hparams.is_swa(il))llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: - layer %3d: reuse layer %d, is_swa = %d\n"
, __func__, il, il_reuse, hparams.is_swa(il));
287 }
288 }
289
290 // allocate tensors and initialize the buffers to avoid NaNs in the padding
291 for (auto & [buft, ctx] : ctx_map) {
292 ggml_backend_buffer_t buf;
293 if (hparams.no_alloc) {
294 buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer
295 for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
296 t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it
297 }
298 } else {
299 buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer
300 }
301 if (!buf) {
302 throwabort_with_suppression(); if (false) std::runtime_error("failed to allocate buffer for kv cache");
303 }
304
305 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)llama_log_internal(GGML_LOG_LEVEL_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);
306
307 ggml_backend_buffer_clear(buf, 0);
308 ctxs_bufs.emplace_back(std::move(ctx), buf);
309 }
310
311 {
312 const size_t memory_size_k = size_k_bytes();
313 const size_t memory_size_v = size_v_bytes();
314
315 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__,llama_log_internal(GGML_LOG_LEVEL_INFO , "%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n"
, __func__, (float)(memory_size_k + memory_size_v) / (1024.0f
* 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream
, ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f
), ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f
))
316 (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,llama_log_internal(GGML_LOG_LEVEL_INFO , "%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n"
, __func__, (float)(memory_size_k + memory_size_v) / (1024.0f
* 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream
, ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f
), ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f
))
317 ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),llama_log_internal(GGML_LOG_LEVEL_INFO , "%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n"
, __func__, (float)(memory_size_k + memory_size_v) / (1024.0f
* 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream
, ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f
), ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f
))
318 ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f))llama_log_internal(GGML_LOG_LEVEL_INFO , "%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n"
, __func__, (float)(memory_size_k + memory_size_v) / (1024.0f
* 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream
, ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f
), ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f
));
319 }
320
321 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
322 if (other) {
323 n_embd_head_k_all = other->n_embd_head_k_all;
324 n_embd_head_v_all = other->n_embd_head_v_all;
325
326 attn_rot_k = other->attn_rot_k;
327 attn_rot_v = other->attn_rot_v;
328 } else {
329 const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
330 const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
331 if (attn_rot_disable) {
332 LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__)llama_log_internal(GGML_LOG_LEVEL_WARN , "%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n"
, __func__);
333 }
334
335 attn_rot_k =
336 !attn_rot_disable &&
337 n_embd_head_k_all > 0 &&
338 ggml_is_quantized(type_k) &&
339 hparams.n_embd_head_k() % 64 == 0;
340
341 // always create Hadamard rotation tensors for DeepSeek V3.2 DSA lightning indexer
342 if (model.arch == LLM_ARCH_DEEPSEEK32 && hparams.n_embd_head_k_full == hparams.indexer_head_size) {
343 attn_rot_k = true;
344 }
345
346 attn_rot_v =
347 !attn_rot_disable &&
348 n_embd_head_v_all > 0 &&
349 ggml_is_quantized(type_v) &&
350 hparams.n_embd_head_v() % 64 == 0;
351 }
352
353 LLAMA_LOG_INFO("%s: attn_rot_k = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_k, n_embd_head_k_all)llama_log_internal(GGML_LOG_LEVEL_INFO , "%s: attn_rot_k = %d, n_embd_head_k_all = %d\n"
, __func__, attn_rot_k, n_embd_head_k_all);
354 LLAMA_LOG_INFO("%s: attn_rot_v = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_v, n_embd_head_v_all)llama_log_internal(GGML_LOG_LEVEL_INFO , "%s: attn_rot_v = %d, n_embd_head_k_all = %d\n"
, __func__, attn_rot_v, n_embd_head_v_all);
355
356 // pre-compute the haramard matrices and keep them in host memory
357 // TODO: in the future, we can make copies in the backend buffers to avoid host -> device transfers
358 if (attn_rot_k || attn_rot_v) {
359 for (int64_t n = 64; n <= std::max(n_embd_head_k_all, n_embd_head_v_all); n *= 2) {
360 attn_rot_hadamard[n] = std::vector<float>(n*n);
361
362 ggml_init_params params = {
363 /* .mem_size = */ 1*ggml_tensor_overhead(),
364 /* .mem_buffer = */ nullptr,
365 /* .no_alloc = */ true,
366 };
367
368 ggml_context_ptr ctx { ggml_init(params) };
369
370 ggml_tensor * tmp = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, n, n);
371 tmp->data = attn_rot_hadamard[n].data();
372
373 ggml_gen_hadamard(tmp);
374 }
375 }
376
377 const char * LLAMA_KV_CACHE_DEBUG = getenv("LLAMA_KV_CACHE_DEBUG");
378 debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0;
379}
380
381void llama_kv_cache::clear(bool data) {
382 for (uint32_t s = 0; s < n_stream; ++s) {
383 v_cells[s].reset();
384 v_heads[s] = 0;
385 }
386
387 if (data) {
388 for (auto & [_, buf] : ctxs_bufs) {
389 ggml_backend_buffer_clear(buf.get(), 0);
390 }
391 }
392}
393
394bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
395 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
396 if (other) {
397 return true;
398 }
399
400 GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()))if (!(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id
< seq_to_stream.size()))) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 400, "GGML_ASSERT(%s) failed", "seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())"
);
401
402 if (p0 < 0) {
403 p0 = 0;
404 }
405
406 if (p1 < 0) {
407 p1 = std::numeric_limits<llama_pos>::max();
408 }
409
410 if (seq_id >= 0) {
411 auto & cells = v_cells[seq_to_stream[seq_id]];
412 auto & head = v_heads[seq_to_stream[seq_id]];
413
414 uint32_t new_head = cells.size();
415
416 for (uint32_t i = 0; i < cells.size(); ++i) {
417 if (!cells.pos_in(i, p0, p1)) {
418 continue;
419 }
420
421 if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {
422 if (new_head == cells.size()) {
423 new_head = i;
424 }
425 }
426 }
427
428 // If we freed up a slot, set head to it so searching can start there.
429 if (new_head != cells.size() && new_head < head) {
430 head = new_head;
431 }
432 } else {
433 // match any sequence
434 for (uint32_t s = 0; s < n_stream; ++s) {
435 auto & cells = v_cells[s];
436 auto & head = v_heads[s];
437
438 uint32_t new_head = cells.size();
439
440 for (uint32_t i = 0; i < cells.size(); ++i) {
441 if (!cells.pos_in(i, p0, p1)) {
442 continue;
443 }
444
445 cells.rm(i);
446
447 if (new_head == cells.size()) {
448 new_head = i;
449 }
450 }
451
452 // If we freed up a slot, set head to it so searching can start there.
453 if (new_head != cells.size() && new_head < head) {
454 head = new_head;
455 }
456 }
457 }
458
459 return true;
460}
461
462void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
463 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
464 if (other) {
465 return;
466 }
467
468 GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size())if (!(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 468, "GGML_ASSERT(%s) failed", "seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size()"
);
469 GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size())if (!(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 469, "GGML_ASSERT(%s) failed", "seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size()"
);
470
471 const auto s0 = seq_to_stream[seq_id_src];
472 const auto s1 = seq_to_stream[seq_id_dst];
473
474 if (s0 == s1) {
475 // since both sequences are in the same stream, no data copy is necessary
476 // we just have to update the cells meta data
477
478 auto & cells = v_cells[s0];
479
480 if (seq_id_src == seq_id_dst) {
481 return;
482 }
483
484 if (p0 < 0) {
485 p0 = 0;
486 }
487
488 if (p1 < 0) {
489 p1 = std::numeric_limits<llama_pos>::max();
490 }
491
492 for (uint32_t i = 0; i < cells.size(); ++i) {
493 if (!cells.pos_in(i, p0, p1)) {
494 continue;
495 }
496
497 if (cells.seq_has(i, seq_id_src)) {
498 cells.seq_add(i, seq_id_dst);
499 }
500 }
501
502 return;
503 }
504
505 // cross-stream sequence copies require to copy the actual buffer data
506
507 bool is_full = true;
508
509 if (p0 > 0 && p0 + 1 < (int) get_size()) {
510 is_full = false;
511 }
512
513 if (p1 > 0 && p1 + 1 < (int) get_size()) {
514 is_full = false;
515 }
516
517 GGML_ASSERT(is_full && "seq_cp() is only supported for full KV buffers")if (!(is_full && "seq_cp() is only supported for full KV buffers"
)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 517, "GGML_ASSERT(%s) failed", "is_full && \"seq_cp() is only supported for full KV buffers\""
);
518
519 // enqueue the copy operation - the buffer copy will be performed during the next update
520 sc_info.ssrc.push_back(s0);
521 sc_info.sdst.push_back(s1);
522
523 v_cells[s1].reset();
524 for (uint32_t i = 0; i < v_cells[s0].size(); ++i) {
525 if (v_cells[s0].seq_has(i, seq_id_src)) {
526 llama_pos pos = v_cells[s0].pos_get(i);
527 llama_pos shift = v_cells[s0].get_shift(i);
528
529 llama_kv_cell_ext ext = v_cells[s0].ext_get(i);
530
531 if (shift != 0) {
532 pos -= shift;
533 assert(pos >= 0)(static_cast <bool> (pos >= 0) ? void (0) : __assert_fail
("pos >= 0", __builtin_FILE (), __builtin_LINE (), __extension__
__PRETTY_FUNCTION__));
534 }
535
536 v_cells[s1].pos_set(i, pos);
537 v_cells[s1].seq_add(i, seq_id_dst);
538
539 if (shift != 0) {
540 v_cells[s1].pos_add(i, shift);
541 }
542
543 v_cells[s1].ext_set(i, ext);
544 }
545 }
546
547 v_heads[s1] = v_heads[s0];
548
549 //for (uint32_t s = 0; s < n_stream; ++s) {
550 // 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));
551 //}
552}
553
554void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
555 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
556 if (other) {
557 return;
558 }
559
560 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())if (!(seq_id >= 0 && (size_t) seq_id < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 560, "GGML_ASSERT(%s) failed", "seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()"
);
561
562 auto & cells = v_cells[seq_to_stream[seq_id]];
563 auto & head = v_heads[seq_to_stream[seq_id]];
564
565 uint32_t new_head = cells.size();
566
567 for (uint32_t i = 0; i < cells.size(); ++i) {
568 if (cells.seq_keep(i, seq_id)) {
569 if (new_head == cells.size()) {
570 new_head = i;
571 }
572 }
573 }
574
575 // If we freed up a slot, set head to it so searching can start there.
576 if (new_head != cells.size() && new_head < head) {
577 head = new_head;
578 }
579}
580
581void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
582 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
583 if (other) {
584 return;
585 }
586
587 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())if (!(seq_id >= 0 && (size_t) seq_id < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 587, "GGML_ASSERT(%s) failed", "seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()"
);
588 GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1")if (!(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1"
)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 588, "GGML_ASSERT(%s) failed", "hparams.n_pos_per_embd() == 1 && \"seq_add() is only supported for n_pos_per_embd() == 1\""
);
589
590 auto & cells = v_cells[seq_to_stream[seq_id]];
591 auto & head = v_heads[seq_to_stream[seq_id]];
592
593 if (shift == 0) {
594 return;
595 }
596
597 uint32_t new_head = cells.size();
598
599 if (p0 < 0) {
600 p0 = 0;
601 }
602
603 if (p1 < 0) {
604 p1 = std::numeric_limits<llama_pos>::max();
605 }
606
607 // If there is no range then return early to avoid looping over all cells.
608 if (p0 == p1) {
609 return;
610 }
611
612 for (uint32_t i = 0; i < cells.size(); ++i) {
613 if (!cells.pos_in(i, p0, p1)) {
614 continue;
615 }
616
617 if (cells.seq_has(i, seq_id)) {
618 if (cells.pos_add(i, shift)) {
619 if (new_head == cells.size()) {
620 new_head = i;
621 }
622 }
623 }
624 }
625
626 // If we freed up a slot, set head to it so searching can start there.
627 // Otherwise we just start the next search from the beginning.
628 head = new_head != cells.size() ? new_head : 0;
629}
630
631void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
632 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
633 if (other) {
634 return;
635 }
636
637 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())if (!(seq_id >= 0 && (size_t) seq_id < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 637, "GGML_ASSERT(%s) failed", "seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()"
);
638 GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1")if (!(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1"
)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 638, "GGML_ASSERT(%s) failed", "hparams.n_pos_per_embd() == 1 && \"seq_div() is only supported for n_pos_per_embd() == 1\""
);
639
640 auto & cells = v_cells[seq_to_stream[seq_id]];
641
642 if (d == 1) {
643 return;
644 }
645
646 if (p0 < 0) {
647 p0 = 0;
648 }
649
650 if (p1 < 0) {
651 p1 = std::numeric_limits<llama_pos>::max();
652 }
653
654 // If there is no range then return early to avoid looping over the cache.
655 if (p0 == p1) {
656 return;
657 }
658
659 for (uint32_t i = 0; i < cells.size(); ++i) {
660 if (!cells.pos_in(i, p0, p1)) {
661 continue;
662 }
663
664 if (cells.seq_has(i, seq_id)) {
665 cells.pos_div(i, d);
666 }
667 }
668}
669
670llama_pos llama_kv_cache::seq_pos_min(llama_seq_id seq_id) const {
671 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
672 if (other) {
673 return other->seq_pos_min(seq_id);
674 }
675
676 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())if (!(seq_id >= 0 && (size_t) seq_id < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 676, "GGML_ASSERT(%s) failed", "seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()"
);
677
678 const auto & cells = v_cells[seq_to_stream[seq_id]];
679
680 return cells.seq_pos_min(seq_id);
681}
682
683llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
684 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
685 if (other) {
686 return other->seq_pos_max(seq_id);
687 }
688
689 GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())if (!(seq_id >= 0 && (size_t) seq_id < seq_to_stream
.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 689, "GGML_ASSERT(%s) failed", "seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()"
);
690
691 const auto & cells = v_cells[seq_to_stream[seq_id]];
692
693 return cells.seq_pos_max(seq_id);
694}
695
696std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
697 std::map<ggml_backend_buffer_type_t, size_t> ret;
698 for (const auto & [ctx, buf] : ctxs_bufs) {
699 ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
700
701 if (hparams.no_alloc) {
702 GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr)if (!(ggml_backend_buffer_get_base(buf.get()) == nullptr)) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 702, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_get_base(buf.get()) == nullptr"
);
703 ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
704 } else {
705 // GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base
706 ret[buft] += ggml_backend_buffer_get_size(buf.get());
707 }
708 }
709
710 return ret;
711}
712
713llama_memory_context_ptr llama_kv_cache::init_batch(
714 llama_batch_allocr & balloc,
715 uint32_t n_ubatch,
716 bool embd_all) {
717 GGML_UNUSED(embd_all)(void)(embd_all);
718
719 do {
720 balloc.split_reset();
721
722 std::vector<llama_ubatch> ubatches;
723 while (true) {
724 auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true);
1
Loop condition is true. Entering loop body
→
2
←
Assuming field 'n_stream' is not equal to 1
→
3
←
'?' condition is false
→
725
726 if (ubatch.n_tokens == 0) {
4
←
Assuming field 'n_tokens' is equal to 0
→
5
←
Taking true branch
→
727 break;
728 }
729
730 ubatches.push_back(std::move(ubatch)); // NOLINT
731 }
732
733 if (balloc.get_n_used() < balloc.get_n_tokens()) {
6
←
 Execution continues on line 733
→
7
←
Assuming the condition is false
→
8
←
Taking false branch
→
734 // failed to find a suitable split
735 break;
736 }
737
738 auto sinfos = prepare(ubatches);
9
←
Calling 'llama_kv_cache::prepare'
→
739 if (sinfos.empty()) {
740 break;
741 }
742
743 return std::make_unique<llama_kv_cache_context>(
744 this, std::move(sinfos), std::move(ubatches));
745 } while (false);
746
747 return std::make_unique<llama_kv_cache_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
748}
749
750llama_memory_context_ptr llama_kv_cache::init_full() {
751 return std::make_unique<llama_kv_cache_context>(this);
752}
753
754llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool optimize) {
755 GGML_UNUSED(optimize)(void)(optimize);
756
757 bool do_shift = get_has_shift();
758
759 return std::make_unique<llama_kv_cache_context>(this, lctx, do_shift, std::move(sc_info));
760}
761
762llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_ubatch> & ubatches) {
763 llama_kv_cache::slot_info_vec_t res;
764
765 struct state_t {
766 slot_info sinfo; // slot info for the ubatch
767
768 std::vector<uint32_t> v_heads_old; // old positions of the heads, before placing the ubatch
769
770 std::vector<llama_kv_cells> v_cells; // copy of the old cells, before placing the ubatch
771 };
772
773 // remember the old state of the cells so we can restore it in the end
774 std::vector<state_t> states;
775
776 bool success = true;
777
778 for (const auto & ubatch : ubatches) {
779 // only find a suitable slot for the ubatch. don't modify the cells yet
780 const auto sinfo_new = find_slot(ubatch, false);
10
←
Calling 'llama_kv_cache::find_slot'
→
781 if (sinfo_new.empty()) {
782 success = false;
783 break;
784 }
785
786 // remember the position that we found
787 res.push_back(sinfo_new);
788
789 // store the old state of the cells in the recovery stack
790 {
791 state_t state = { sinfo_new, v_heads, {} };
792
793 for (uint32_t s = 0; s < sinfo_new.n_stream(); ++s) {
794 auto & cells = v_cells[sinfo_new.strm[s]];
795
796 state.v_cells.push_back(cells.cp(sinfo_new.idxs[s]));
797 }
798
799 states.push_back(std::move(state));
800 }
801
802 // now emplace the ubatch
803 apply_ubatch(sinfo_new, ubatch);
804 }
805
806 GGML_ASSERT(!states.empty() || !success)if (!(!states.empty() || !success)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 806, "GGML_ASSERT(%s) failed", "!states.empty() || !success"
);
807
808 // iterate backwards and restore the cells to their original state
809 for (auto it = states.rbegin(); it != states.rend(); ++it) {
810 const auto & sinfo = it->sinfo;
811
812 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
813 auto & cells = v_cells[sinfo.strm[s]];
814 auto & head = v_heads[sinfo.strm[s]];
815
816 cells.set(sinfo.idxs[s], it->v_cells[s]);
817 head = it->v_heads_old[s];
818 }
819 }
820
821 if (!success) {
822 return {};
823 }
824
825 return res;
826}
827
828bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info) {
829 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
830 if (other) {
831 return true;
832 }
833
834 bool updated = false;
835
836 auto * sched = lctx->get_sched();
837
838 if (!sc_info.empty()) {
839 assert(n_stream > 1 && "stream copy should never happen with a single stream")(static_cast <bool> (n_stream > 1 && "stream copy should never happen with a single stream"
) ? void (0) : __assert_fail ("n_stream > 1 && \"stream copy should never happen with a single stream\""
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
840
841 llama_synchronize(lctx);
842
843 const size_t n_copy = sc_info.ssrc.size();
844
845 for (size_t i = 0; i < n_copy; ++i) {
846 const auto ssrc = sc_info.ssrc[i];
847 const auto sdst = sc_info.sdst[i];
848
849 assert(ssrc < n_stream)(static_cast <bool> (ssrc < n_stream) ? void (0) : __assert_fail
("ssrc < n_stream", __builtin_FILE (), __builtin_LINE (),
__extension__ __PRETTY_FUNCTION__));
850 assert(sdst < n_stream)(static_cast <bool> (sdst < n_stream) ? void (0) : __assert_fail
("sdst < n_stream", __builtin_FILE (), __builtin_LINE (),
__extension__ __PRETTY_FUNCTION__));
851
852 LLAMA_LOG_DEBUG("%s: copying KV buffer: stream %d to stream %d\n", __func__, ssrc, sdst)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: copying KV buffer: stream %d to stream %d\n"
, __func__, ssrc, sdst);
853
854 assert(ssrc != sdst)(static_cast <bool> (ssrc != sdst) ? void (0) : __assert_fail
("ssrc != sdst", __builtin_FILE (), __builtin_LINE (), __extension__
__PRETTY_FUNCTION__));
855
856 for (uint32_t il = 0; il < layers.size(); ++il) {
857 const auto & layer = layers[il];
858
859 ggml_backend_tensor_copy(layer.k_stream[ssrc], layer.k_stream[sdst]);
860
861 if (layer.v_stream[ssrc]) {
862 ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);
863 }
864 }
865 }
866 }
867
868 if (do_shift) {
869 if (!get_can_shift()) {
870 GGML_ABORT("The current KV cache / model configuration does not support K-shift")ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 870, "The current KV cache / model configuration does not support K-shift"
);
871 }
872
873 LLAMA_LOG_DEBUG("%s: applying K-shift\n", __func__)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: applying K-shift\n"
, __func__);
874
875 // apply K-shift if needed
876 if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) {
877 ggml_backend_sched_reset(sched);
878
879 auto * res = lctx->get_gf_res_reserve();
880
881 res->reset();
882
883 auto * gf = build_graph_shift(res, lctx);
884 if (!ggml_backend_sched_alloc_graph(sched, gf)) {
885 LLAMA_LOG_ERROR("%s: failed to allocate compute graph for K-shift\n", __func__)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: failed to allocate compute graph for K-shift\n"
, __func__);
886 return updated;
887 }
888
889 res->set_inputs(nullptr);
890
891 if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) {
892 LLAMA_LOG_ERROR("%s: failed to compute K-shift\n", __func__)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: failed to compute K-shift\n"
, __func__);
893 return updated;
894 }
895
896 updated = true;
897 }
898
899 for (uint32_t s = 0; s < n_stream; ++s) {
900 auto & cells = v_cells[s];
901
902 cells.reset_shift();
903 }
904 }
905
906 return updated;
907}
908
909llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, bool cont) const {
910
911 if (debug > 0) {
11
←
Assuming field 'debug' is > 0
→
12
←
Taking true branch
→
912 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
13
←
Assuming 's' is >= field 'n_seqs_unq'
→
14
←
Loop condition is false. Execution continues on line 977
→
913 const auto seq_id = ubatch.seq_id_unq[s];
914 const auto stream_id = seq_to_stream[seq_id];
915 const auto & cells = v_cells[stream_id];
916 const uint32_t head_cur = v_heads[stream_id];
917
918 LLAMA_LOG_DEBUG("%s: stream[%d], n = %5d, used = %5d, head = %5d, size = %5d, n_swa = %5d\n",llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: stream[%d], n = %5d, used = %5d, head = %5d, size = %5d, n_swa = %5d\n"
, __func__, stream_id, cells.used_max_p1(), cells.get_used(),
head_cur, get_size(), n_swa)
919 __func__, stream_id, cells.used_max_p1(), cells.get_used(), head_cur, get_size(), n_swa)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: stream[%d], n = %5d, used = %5d, head = %5d, size = %5d, n_swa = %5d\n"
, __func__, stream_id, cells.used_max_p1(), cells.get_used(),
head_cur, get_size(), n_swa);
920
921 if ((debug == 2 && n_swa > 0) || debug > 2) {
922 std::string ss;
923 for (uint32_t i = 0; i < cells.size(); ++i) {
924 if (cells.is_empty(i)) {
925 ss += '.';
926 } else {
927 assert(cells.seq_count(i) >= 1)(static_cast <bool> (cells.seq_count(i) >= 1) ? void
(0) : __assert_fail ("cells.seq_count(i) >= 1", __builtin_FILE
(), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__));
928
929 if (cells.seq_count(i) == 1) {
930 ss += std::to_string(cells.seq_get(i));
931 } else {
932 ss += 'M';
933 }
934 }
935 if (i%256 == 255) {
936 ss += " *";
937 ss += '\n';
938 }
939 }
940 LLAMA_LOG_DEBUG("\n%s\n", ss.c_str())llama_log_internal(GGML_LOG_LEVEL_DEBUG, "\n%s\n", ss.c_str()
);
941 }
942
943 if ((debug == 2 && n_swa > 0) || debug > 2) {
944 std::string ss;
945 for (uint32_t i = 0; i < cells.size(); ++i) {
946 std::string cur;
947 if (cells.is_empty(i)) {
948 cur = '.';
949 } else {
950 cur = std::to_string(cells.pos_get(i));
951 }
952 const int n = cur.size();
953 for (int j = 0; j < 5 - n; ++j) {
954 cur += ' ';
955 }
956 ss += cur;
957 if (i%256 == 255) {
958 ss += " *";
959 }
960 if (i%64 == 63) {
961 ss += '\n';
962 }
963 }
964 LLAMA_LOG_DEBUG("\n%s\n", ss.c_str())llama_log_internal(GGML_LOG_LEVEL_DEBUG, "\n%s\n", ss.c_str()
);
965 }
966
967 for (int s = 0; s < LLAMA_MAX_SEQ256; ++s) {
968 if (cells.seq_pos_min(s) < 0) {
969 continue;
970 }
971
972 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))llama_log_internal(GGML_LOG_LEVEL_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));
973 }
974 }
975 }
976
977 uint32_t n_tokens = ubatch.n_tokens;
978 uint32_t n_seqs = 1;
979
980 if (n_stream > 1) {
15
←
Assuming field 'n_stream' is > 1
→
16
←
Taking true branch
→
981 GGML_ASSERT(n_tokens % ubatch.n_seqs_unq == 0)if (!(n_tokens % ubatch.n_seqs_unq == 0)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 981, "GGML_ASSERT(%s) failed", "n_tokens % ubatch.n_seqs_unq == 0"
);
17
←
Division by zero
982
983 n_seqs = ubatch.n_seqs_unq;
984 n_tokens = n_tokens / n_seqs;
985 }
986
987 slot_info res = {
988 /*.s0 =*/ LLAMA_MAX_SEQ256,
989 /*.s1 =*/ 0,
990 /*.strm =*/ { },
991 /*.idxs =*/ { },
992 };
993
994 res.resize(n_seqs);
995
996 for (uint32_t s = 0; s < n_seqs; ++s) {
997 const auto seq_id = ubatch.seq_id_unq[s];
998
999 if (n_stream > 1) {
1000 GGML_ASSERT(ubatch.n_seq_id[s*n_tokens] == 1)if (!(ubatch.n_seq_id[s*n_tokens] == 1)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1000, "GGML_ASSERT(%s) failed", "ubatch.n_seq_id[s*n_tokens] == 1"
);
1001 GGML_ASSERT(ubatch.seq_id [s*n_tokens][0] == seq_id)if (!(ubatch.seq_id [s*n_tokens][0] == seq_id)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1001, "GGML_ASSERT(%s) failed", "ubatch.seq_id [s*n_tokens][0] == seq_id"
);
1002 }
1003
1004 res.s0 = std::min<uint32_t>(res.s0, seq_to_stream[seq_id]);
1005 res.s1 = std::max<uint32_t>(res.s1, seq_to_stream[seq_id]);
1006
1007 res.strm[s] = seq_to_stream[seq_id];
1008 res.idxs[s].reserve(n_tokens);
1009
1010 const auto & cells = v_cells[seq_to_stream[seq_id]];
1011
1012 uint32_t head_cur = v_heads[seq_to_stream[seq_id]];
1013
1014 // if we have enough unused cells before the current head ->
1015 // better to start searching from the beginning of the cache, hoping to fill it
1016 if (head_cur > cells.get_used() + 2*n_tokens) {
1017 head_cur = 0;
1018 }
1019
1020 if (n_tokens > cells.size()) {
1021 LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size())llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: n_tokens = %d > size = %u\n"
, __func__, n_tokens, cells.size());
1022 return { };
1023 }
1024
1025 uint32_t n_tested = 0;
1026
1027 // for continuous slots, we test that all tokens in the ubatch fit, starting from the current head
1028 // for non-continuous slots, we test the tokens one by one
1029 const uint32_t n_test = cont ? n_tokens : 1;
1030
1031 while (true) {
1032 if (head_cur + n_test > cells.size()) {
1033 n_tested += cells.size() - head_cur;
1034 head_cur = 0;
1035 continue;
1036 }
1037
1038 for (uint32_t i = 0; i < n_test; i++) {
1039 const auto idx = head_cur;
1040
1041 head_cur++;
1042 n_tested++;
1043
1044 //const llama_pos pos = ubatch.pos[i];
1045 //const llama_seq_id seq_id = ubatch.seq_id[i][0];
1046
1047 // can we use this cell? either:
1048 // - the cell is empty
1049 // - the cell is occupied only by one sequence:
1050 // - (disabled) mask causally, if the sequence is the same as the one we are inserting
1051 // - mask SWA, using current max pos for that sequence in the cache
1052 // always insert in the cell with minimum pos
1053 bool can_use = cells.is_empty(idx);
1054
1055 if (!can_use && cells.seq_count(idx) == 1) {
1056 const llama_pos pos_cell = cells.pos_get(idx);
1057
1058 // (disabled) causal mask
1059 // note: it's better to purge any "future" tokens beforehand
1060 //if (cells.seq_has(idx, seq_id)) {
1061 // can_use = pos_cell >= pos;
1062 //}
1063
1064 if (!can_use) {
1065 const llama_seq_id seq_id_cell = cells.seq_get(idx);
1066
1067 // SWA mask
1068 if (llama_hparams::is_masked_swa(n_swa, swa_type, pos_cell, cells.seq_pos_max(seq_id_cell) + 1)) {
1069 can_use = true;
1070 }
1071 }
1072 }
1073
1074 if (can_use) {
1075 res.idxs[s].push_back(idx);
1076 } else {
1077 if (cont) {
1078 break;
1079 }
1080 }
1081 }
1082
1083 if (res.idxs[s].size() == n_tokens) {
1084 break;
1085 }
1086
1087 if (cont) {
1088 res.idxs[s].clear();
1089 }
1090
1091 if (n_tested >= cells.size()) {
1092 //LLAMA_LOG_ERROR("%s: failed to find a slot for %d tokens\n", __func__, n_tokens);
1093 return { };
1094 }
1095 }
1096
1097 // we didn't find a suitable slot - return empty result
1098 if (res.idxs[s].size() < n_tokens) {
1099 return { };
1100 }
1101 }
1102
1103 assert(res.s1 >= res.s0)(static_cast <bool> (res.s1 >= res.s0) ? void (0) : __assert_fail
("res.s1 >= res.s0", __builtin_FILE (), __builtin_LINE ()
, __extension__ __PRETTY_FUNCTION__));
1104
1105 return res;
1106}
1107
1108void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) {
1109 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
1110 if (other) {
1111 return;
1112 }
1113
1114 // keep track of the max sequence position that we would overwrite with this ubatch
1115 // for non-SWA cache, this would be always empty
1116 llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ256];
1117 for (uint32_t s = 0; s < LLAMA_MAX_SEQ256; ++s) {
1118 seq_pos_max_rm[s] = -1;
1119 }
1120
1121 assert(ubatch.n_tokens == sinfo.n_stream()*sinfo.size())(static_cast <bool> (ubatch.n_tokens == sinfo.n_stream(
)*sinfo.size()) ? void (0) : __assert_fail ("ubatch.n_tokens == sinfo.n_stream()*sinfo.size()"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
1122
1123 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
1124 for (uint32_t ii = 0; ii < sinfo.size(); ++ii) {
1125 const uint32_t i = s*sinfo.size() + ii;
1126
1127 auto & cells = v_cells[sinfo.strm[s]];
1128
1129 const auto idx = sinfo.idxs[s][ii];
1130
1131 if (!cells.is_empty(idx)) {
1132 assert(cells.seq_count(idx) == 1)(static_cast <bool> (cells.seq_count(idx) == 1) ? void (
0) : __assert_fail ("cells.seq_count(idx) == 1", __builtin_FILE
(), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__));
1133
1134 const llama_seq_id seq_id = cells.seq_get(idx);
1135 const llama_pos pos = cells.pos_get(idx);
1136
1137 seq_pos_max_rm[seq_id] = std::max(seq_pos_max_rm[seq_id], pos);
1138
1139 cells.rm(idx);
1140 }
1141
1142 cells.pos_set(idx, ubatch.pos[i]);
1143
1144 if (ubatch.is_pos_2d()) {
1145 llama_kv_cell_ext ext {
1146 /*.x =*/ ubatch.pos[i + ubatch.n_tokens*2],
1147 /*.y =*/ ubatch.pos[i + ubatch.n_tokens],
1148 };
1149 cells.ext_set(idx, ext);
1150 }
1151
1152 for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) {
1153 cells.seq_add(idx, ubatch.seq_id[i][s]);
1154 }
1155 }
1156 }
1157
1158 // note: we want to preserve the invariant that all positions between [pos_min, pos_max] for each sequence
1159 // will be present in the cache. so we have to purge any position which is less than those we would overwrite
1160 // ref: https://github.com/ggml-org/llama.cpp/pull/13746#issuecomment-2916057092
1161 for (uint32_t s = 0; s < LLAMA_MAX_SEQ256; ++s) {
1162 if (seq_pos_max_rm[s] == -1) {
1163 continue;
1164 }
1165
1166 GGML_ASSERT(s < seq_to_stream.size())if (!(s < seq_to_stream.size())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1166, "GGML_ASSERT(%s) failed", "s < seq_to_stream.size()"
);
1167
1168 auto & cells = v_cells[seq_to_stream[s]];
1169
1170 if (cells.seq_pos_min(s) <= seq_pos_max_rm[s]) {
1171 LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: purging positions [%d, %d] of sequence %d from KV cache\n"
, __func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s)
1172 __func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: purging positions [%d, %d] of sequence %d from KV cache\n"
, __func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
1173
1174 seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
1175 }
1176 }
1177
1178 // move the head at the end of the slot
1179 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
1180 auto & head = v_heads[sinfo.strm[s]];
1181
1182 head = sinfo.idxs[s].back() + 1;
1183 }
1184}
1185
1186bool llama_kv_cache::get_can_shift() const {
1187 // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot.
1188 if (model.arch == LLM_ARCH_STEP35) {
1189 return false;
1190 }
1191 if (hparams.n_pos_per_embd() > 1) {
1192 return false;
1193 }
1194 return true;
1195}
1196
1197uint32_t llama_kv_cache::get_size() const {
1198 const auto & cells = v_cells[seq_to_stream[0]];
1199
1200 return cells.size();
1201}
1202
1203uint32_t llama_kv_cache::get_n_stream() const {
1204 return n_stream;
1205}
1206
1207bool llama_kv_cache::get_has_shift() const {
1208 bool result = false;
1209
1210 for (uint32_t s = 0; s < n_stream; ++s) {
1211 result |= v_cells[s].get_has_shift();
1212 }
1213
1214 return result;
1215}
1216
1217ggml_type llama_kv_cache::type_k() const {
1218 return layers[0].k->type;
1219}
1220
1221ggml_type llama_kv_cache::type_v() const {
1222 return layers[0].v->type;
1223}
1224
1225uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
1226 uint32_t result = 0;
1227
1228 // pad the n_kv value so that the graph remains constant across batches and can be reused
1229 // note: this also helps some backends with performance (f.ex https://github.com/ggml-org/llama.cpp/pull/16812#issuecomment-3455112220)
1230 const uint32_t n_pad_cur = std::max(n_pad, 256u);
1231
1232 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
1233 const auto & cells = v_cells[sinfo.strm[s]];
1234
1235 result = std::max(std::min(cells.size(), std::max(n_pad_cur, GGML_PAD(cells.used_max_p1(), n_pad_cur)(((cells.used_max_p1()) + (n_pad_cur) - 1) & ~((n_pad_cur
) - 1)))), result);
1236 }
1237
1238 return result;
1239}
1240
1241ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
1242 const int32_t ikv = map_layer_ids.at(il);
1243
1244 auto * k = layers[ikv].k;
1245
1246 const uint64_t kv_size = get_size();
1247 const uint64_t n_embd_k_gqa = k->ne[0];
1248
1249 assert(n_embd_k_gqa == hparams.n_embd_k_gqa(il))(static_cast <bool> (n_embd_k_gqa == hparams.n_embd_k_gqa
(il)) ? void (0) : __assert_fail ("n_embd_k_gqa == hparams.n_embd_k_gqa(il)"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
1250
1251 const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
1252
1253 return ggml_view_4d(ctx, k,
1254 hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv, ns,
1255 ggml_row_size(k->type, hparams.n_embd_head_k(il)),
1256 ggml_row_size(k->type, n_embd_k_gqa),
1257 ggml_row_size(k->type, n_embd_k_gqa*kv_size),
1258 ggml_row_size(k->type, n_embd_k_gqa*kv_size)*sinfo.s0);
1259}
1260
1261ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
1262 const int32_t ikv = map_layer_ids.at(il);
1263
1264 auto * v = layers[ikv].v;
1265
1266 const uint64_t kv_size = get_size();
1267 const uint64_t n_embd_v_gqa = v->ne[0];
1268
1269 // [TAG_V_CACHE_VARIABLE]
1270 assert(n_embd_v_gqa >= hparams.n_embd_v_gqa(il))(static_cast <bool> (n_embd_v_gqa >= hparams.n_embd_v_gqa
(il)) ? void (0) : __assert_fail ("n_embd_v_gqa >= hparams.n_embd_v_gqa(il)"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
1271
1272 const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
1273
1274 if (!v_trans) {
1275 // note: v->nb[1] <= v->nb[2]
1276 return ggml_view_4d(ctx, v,
1277 hparams.n_embd_head_v(il), hparams.n_head_kv(il), n_kv, ns,
1278 ggml_row_size(v->type, hparams.n_embd_head_v(il)), // v->nb[1]
1279 ggml_row_size(v->type, n_embd_v_gqa), // v->nb[2]
1280 ggml_row_size(v->type, n_embd_v_gqa*kv_size), // v->nb[3]
1281 ggml_row_size(v->type, n_embd_v_gqa*kv_size)*sinfo.s0);
1282 }
1283
1284 // note: v->nb[1] > v->nb[2]
1285 return ggml_view_4d(ctx, v,
1286 n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), ns,
1287 ggml_row_size(v->type, kv_size*hparams.n_embd_head_v(il)), // v->nb[1]
1288 ggml_row_size(v->type, kv_size), // v->nb[2]
1289 ggml_row_size(v->type, kv_size*n_embd_v_gqa), // v->nb[3]
1290 ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
1291}
1292
1293ggml_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 {
1294 GGML_UNUSED(sinfo)(void)(sinfo);
1295
1296 const int32_t ikv = map_layer_ids.at(il);
1297
1298 ggml_tensor * k = layers[ikv].k;
1299
1300 const int64_t n_embd_head = k_cur->ne[0];
1301 const int64_t n_head = k_cur->ne[1];
1302 const int64_t n_tokens = k_cur->ne[2];
1303
1304 const int64_t n_embd_gqa = n_embd_head*n_head;
1305
1306 // we can merge dims 0 and 1
1307 // TODO: add ggml helper function for this?
1308 GGML_ASSERT(ggml_row_size(k_cur->type, n_embd_head) == k_cur->nb[1])if (!(ggml_row_size(k_cur->type, n_embd_head) == k_cur->
nb[1])) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1308, "GGML_ASSERT(%s) failed", "ggml_row_size(k_cur->type, n_embd_head) == k_cur->nb[1]"
);
1309
1310 k_cur = ggml_view_2d(ctx, k_cur, n_embd_gqa, n_tokens, k_cur->nb[2], 0);
1311
1312 const int64_t n_stream = k->ne[2];
1313
1314 if (n_stream > 1) {
1315 const int64_t kv_size = get_size();
1316
1317 assert(n_embd_gqa == k->ne[0])(static_cast <bool> (n_embd_gqa == k->ne[0]) ? void (
0) : __assert_fail ("n_embd_gqa == k->ne[0]", __builtin_FILE
(), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__));
1318 assert(kv_size == k->ne[1])(static_cast <bool> (kv_size == k->ne[1]) ? void (0)
: __assert_fail ("kv_size == k->ne[1]", __builtin_FILE ()
, __builtin_LINE (), __extension__ __PRETTY_FUNCTION__));
1319
1320 // merge the buffer across all streams because the idxs are global
1321 k = ggml_reshape_2d(ctx, k, n_embd_gqa, kv_size*n_stream);
1322 }
1323
1324 // store the current K values into the cache
1325 return ggml_set_rows(ctx, k, k_cur, k_idxs);
1326}
1327
1328ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const {
1329 GGML_UNUSED(sinfo)(void)(sinfo);
1330
1331 const int32_t ikv = map_layer_ids.at(il);
1332
1333 auto * v = layers[ikv].v;
1334
1335 const int64_t n_embd_head = v_cur->ne[0];
1336 const int64_t n_head = v_cur->ne[1];
1337 const int64_t n_tokens = v_cur->ne[2];
1338
1339 const int64_t n_embd_gqa = n_embd_head*n_head;
1340
1341 // we can merge dims 0 and 1
1342 GGML_ASSERT(ggml_row_size(v_cur->type, n_embd_head) == v_cur->nb[1])if (!(ggml_row_size(v_cur->type, n_embd_head) == v_cur->
nb[1])) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1342, "GGML_ASSERT(%s) failed", "ggml_row_size(v_cur->type, n_embd_head) == v_cur->nb[1]"
);
1343
1344 const int64_t n_stream = v->ne[2];
1345
1346 // take this branch when FA is enabled (the V cache is not transposed)
1347 if (!v_trans) {
1348 v_cur = ggml_view_2d(ctx, v_cur, n_embd_gqa, n_tokens, v_cur->nb[2], 0);
1349
1350 if (n_stream > 1) {
1351 const int64_t kv_size = get_size();
1352
1353 assert(n_embd_gqa == v->ne[0])(static_cast <bool> (n_embd_gqa == v->ne[0]) ? void (
0) : __assert_fail ("n_embd_gqa == v->ne[0]", __builtin_FILE
(), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__));
1354 assert(kv_size == v->ne[1])(static_cast <bool> (kv_size == v->ne[1]) ? void (0)
: __assert_fail ("kv_size == v->ne[1]", __builtin_FILE ()
, __builtin_LINE (), __extension__ __PRETTY_FUNCTION__));
1355
1356 // merge the buffer across all streams because the idxs are global
1357 v = ggml_reshape_2d(ctx, v, n_embd_gqa, kv_size*n_stream);
1358 }
1359
1360 return ggml_set_rows(ctx, v, v_cur, v_idxs);
1361 }
1362
1363 if (ggml_row_size(v_cur->type, n_embd_gqa) == v_cur->nb[2]) {
1364 // we can merge dims 0, 1 and 2
1365 v_cur = ggml_reshape_2d(ctx, v_cur, n_embd_gqa, n_tokens);
1366 } else {
1367 // otherwise -> make a copy to get contiguous data
1368 v_cur = ggml_cont_2d (ctx, v_cur, n_embd_gqa, n_tokens);
1369 }
1370
1371 // [TAG_V_CACHE_VARIABLE]
1372 if (n_embd_gqa < v->ne[0]) {
1373 v_cur = ggml_pad(ctx, v_cur, v->ne[0] - n_embd_gqa, 0, 0, 0);
1374 }
1375
1376 // in this branch the v_idxs are constructed in such a way that each row is a single head element
1377 ggml_tensor * v_view = ggml_reshape_2d(ctx, v, 1, ggml_nelements(v));
1378
1379 v_cur = ggml_reshape_2d(ctx, v_cur, 1, ggml_nelements(v_cur));
1380
1381 return ggml_set_rows(ctx, v_view, v_cur, v_idxs);
1382}
1383
1384ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
1385 const uint32_t n_tokens = ubatch.n_tokens;
1386
1387 ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens);
1388
1389 ggml_set_input(k_idxs);
1390
1391 return k_idxs;
1392}
1393
1394ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
1395 const uint32_t n_tokens = ubatch.n_tokens;
1396
1397 ggml_tensor * v_idxs;
1398
1399 if (!v_trans) {
1400 v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens);
1401 } else {
1402 v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max());
1403 }
1404
1405 ggml_set_input(v_idxs);
1406
1407 return v_idxs;
1408}
1409
1410ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const {
1411 ggml_tensor * res = nullptr;
1412
1413 if (attn_rot_k) {
1414 int nrot = 64;
1415
1416 // TODO: investigate if using the smallest rotation matrix is beneficial also for K (similar as for V)
1417 // ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4141323088
1418 do {
1419 nrot *= 2;
1420 } while (n_embd_head_k_all % nrot == 0);
1421 nrot /= 2;
1422
1423 res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot);
1424 ggml_set_input(res);
1425 ggml_set_name(res, "attn_inp_k_rot");
1426 }
1427
1428 return res;
1429}
1430
1431ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const {
1432 ggml_tensor * res = nullptr;
1433
1434 if (attn_rot_v) {
1435 int nrot = 64;
1436 // using smaller rotation matrices for V seems beneficial
1437 // ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4146397570
1438 //do {
1439 // nrot *= 2;
1440 //} while (hparams.n_embd_head_v() % nrot == 0);
1441 //nrot /= 2;
1442
1443 res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot);
1444 ggml_set_input(res);
1445 ggml_set_name(res, "attn_inp_v_rot");
1446 }
1447
1448 return res;
1449}
1450
1451void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const {
1452 const uint32_t n_tokens = ubatch->n_tokens;
1453 GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream())if (!(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream())) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1453, "GGML_ASSERT(%s) failed", "n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()"
);
1454
1455 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1455, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1456 int64_t * data = (int64_t *) dst->data;
1457
1458 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
1459 const int64_t offs = sinfo.strm[s]*get_size();
1460
1461 for (uint32_t i = 0; i < sinfo.size(); ++i) {
1462 data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i];
1463 }
1464 }
1465}
1466
1467void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const {
1468 const uint32_t n_tokens = ubatch->n_tokens;
1469 GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream())if (!(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream())) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1469, "GGML_ASSERT(%s) failed", "n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()"
);
1470
1471 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1471, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1472 int64_t * data = (int64_t *) dst->data;
1473
1474 if (!v_trans) {
1475 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
1476 const int64_t offs = sinfo.strm[s]*get_size();
1477
1478 for (uint32_t i = 0; i < sinfo.size(); ++i) {
1479 data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i];
1480 }
1481 }
1482 } else {
1483 // note: the V cache is transposed when not using flash attention
1484 const int64_t kv_size = get_size();
1485
1486 const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa_max();
1487
1488 for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {
1489 const int64_t offs = sinfo.strm[s]*kv_size*n_embd_v_gqa;
1490
1491 for (uint32_t i = 0; i < sinfo.size(); ++i) {
1492 for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
1493 data[s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j] = offs + j*kv_size + sinfo.idxs[s][i];
1494 }
1495 }
1496 }
1497 }
1498}
1499
1500void llama_kv_cache::set_input_k_shift(ggml_tensor * dst) const {
1501 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1501, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1502
1503 int32_t * data = (int32_t *) dst->data;
1504
1505 for (uint32_t s = 0; s < n_stream; ++s) {
1506 const auto & cells = v_cells[s];
1507
1508 for (uint32_t i = 0; i < cells.size(); ++i) {
1509 data[s*cells.size() + i] = cells.is_empty(i) ? 0 : cells.get_shift(i);
1510 }
1511 }
1512}
1513
1514struct args_set_input_kq_mask {
1515 const llama_hparams & hparams;
1516 const llama_ubatch * ubatch;
1517
1518 const std::vector<llama_kv_cells> & v_cells;
1519 const std::vector<uint32_t> & seq_to_stream;
1520
1521 uint32_t n_swa;
1522 llama_swa_type swa_type;
1523
1524 int64_t n_kv;
1525 int64_t n_stream;
1526 int64_t n_tps;
1527};
1528
1529template<typename T, bool causal, bool swa, bool is_2d, bool alibi>
1530static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
1531 //const auto & hparams = args.hparams;
1532 const auto & ubatch = args.ubatch;
1533
1534 const auto & v_cells = args.v_cells;
1535 const auto & seq_to_stream = args.seq_to_stream;
1536
1537 const uint32_t n_swa = args.n_swa;
1538 const llama_swa_type swa_type = args.swa_type;
1539
1540 const int64_t n_kv = args.n_kv;
1541 const int64_t n_stream = args.n_stream;
1542 const int64_t n_tps = args.n_tps;
1543
1544 const T mask_keep = llama_cast<T>(0.0f);
1545 const T mask_drop = llama_cast<T>(-INFINITY(__builtin_inff ()));
1546
1547 // the min position in the batch for each sequence
1548 llama_pos seq_pos_min[LLAMA_MAX_SEQ256];
1549 std::fill(seq_pos_min, seq_pos_min + LLAMA_MAX_SEQ256, INT32_MAX(2147483647));
1550
1551 for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
1552 const llama_seq_id seq_id = ubatch->seq_id[i][0];
1553
1554 seq_pos_min[seq_id] = std::min(seq_pos_min[seq_id], ubatch->pos[i]);
1555 }
1556
1557 for (uint32_t s = 0; s < n_stream; ++s) {
1558 // bookkeeping of the KQ mask cells that could change for other tokens of the same sequence
1559 std::unordered_map<llama_seq_id, uint32_t> seq_srct;
1560 std::unordered_map<llama_seq_id, std::vector<uint32_t>> seq_idxs;
1561
1562 for (uint32_t ii = 0; ii < n_tps; ++ii) {
1563 const uint32_t i = s*n_tps + ii;
1564
1565 const llama_seq_id seq_id = ubatch->seq_id[i][0];
1566
1567 const auto & cells = v_cells.at(seq_to_stream[seq_id]);
1568
1569 llama_pos p0 = -1;
1570 const llama_pos p1 = ubatch->pos[i];
1571
1572 // for M-RoPE
1573 const llama_pos p1_x = is_2d ? ubatch->pos[i + ubatch->n_tokens*2] : 0;
1574 const llama_pos p1_y = is_2d ? ubatch->pos[i + ubatch->n_tokens] : 0;
1575
1576 const uint64_t idst = n_kv*i;
1577
1578 // for tokens of the same sequence, the mask is mostly the same, so we can reuse it
1579 // the only cells that could change are the ones that are with similar positions as the
1580 // ones in the batch (i.e. due to causal masking, SWA, etc.)
1581 // keep track of those cells and shortcut the loop to save time
1582 // note: this optimization is not compatible with Alibi position encoding
1583 // ref: https://github.com/ggml-org/llama.cpp/pull/18842
1584 bool prev = false;
1585
1586 auto & idxs = seq_idxs[seq_id];
1587
1588 if (!alibi) {
1589 if (seq_srct.find(seq_id) != seq_srct.end()) {
1590 const uint32_t srct = seq_srct[seq_id];
1591
1592 const uint64_t idst_prev = n_kv*srct;
1593
1594 std::copy(data + idst_prev, data + idst_prev + n_kv, data + idst);
1595
1596 prev = true;
1597 } else {
1598 idxs.clear();
1599 idxs.reserve(ubatch->n_tokens + n_swa + 32);
1600
1601 seq_srct[seq_id] = i;
1602 }
1603 }
1604
1605 for (uint32_t jj = 0; jj < n_kv; ++jj) {
1606 uint32_t j = jj;
1607
1608 // we have an exiting mask for this sequence -> update just seq_idxs
1609 if (!alibi) {
1610 if (prev) {
1611 if (jj >= idxs.size()) {
1612 break;
1613 }
1614
1615 j = idxs[jj];
1616 }
1617 }
1618
1619 if (cells.is_empty(j)) {
1620 goto skip;
1621 }
1622
1623 // mask the token if not the same sequence
1624 if (!cells.seq_has(j, seq_id)) {
1625 goto skip;
1626 }
1627
1628 p0 = cells.pos_get(j);
1629
1630 if (!alibi) {
1631 if (!prev) {
1632 // record all cells for which: p0 >= seq_pos_min[seq_id] - n_swa - 32
1633 if (p0 + (int32_t) (n_swa + 32) >= seq_pos_min[seq_id]) {
1634 idxs.push_back(j);
1635 }
1636 }
1637 }
1638
1639 if (causal) {
1640 // mask future tokens
1641 if (p0 > p1) {
1642 goto skip;
1643 }
1644
1645 // M-RoPE causal mask
1646 if (is_2d) {
1647 if (p0 == p1) {
1648 const auto & p0_ext = cells.ext_get(j);
1649
1650 if (p0_ext.is_2d_gt(p1_x, p1_y)) {
1651 goto skip;
1652 }
1653 }
1654 }
1655 }
1656
1657 // apply SWA if any
1658 if (swa) {
1659 if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
1660 goto skip;
1661 }
1662 }
1663
1664 if (alibi) {
1665 data[idst + j] = llama_cast<T>(static_cast<float>(-std::abs(p0 - p1)));
1666 } else {
1667 data[idst + j] = mask_keep;
1668 }
1669
1670 continue;
1671skip:
1672 data[idst + j] = mask_drop;
1673 }
1674 }
1675 }
1676}
1677
1678template<typename T, bool causal, bool swa, bool is_2d>
1679static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
1680 const bool alibi = args.hparams.use_alibi;
1681 if (alibi) {
1682 set_input_kq_mask_impl<T, causal, swa, is_2d, true> (args, data);
1683 } else {
1684 set_input_kq_mask_impl<T, causal, swa, is_2d, false>(args, data);
1685 }
1686}
1687
1688template<typename T, bool causal, bool swa>
1689static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
1690 const bool is_2d = args.ubatch->is_pos_2d();
1691 if (is_2d) {
1692 set_input_kq_mask_impl<T, causal, swa, true> (args, data);
1693 } else {
1694 set_input_kq_mask_impl<T, causal, swa, false>(args, data);
1695 }
1696}
1697
1698template<typename T, bool causal>
1699static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
1700 const bool swa = args.swa_type != LLAMA_SWA_TYPE_NONE;
1701 if (swa) {
1702 set_input_kq_mask_impl<T, causal, true> (args, data);
1703 } else {
1704 set_input_kq_mask_impl<T, causal, false>(args, data);
1705 }
1706}
1707
1708template<typename T>
1709static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data, bool causal_attn) {
1710 if (causal_attn) {
1711 set_input_kq_mask_impl<T, true> (args, data);
1712 } else {
1713 set_input_kq_mask_impl<T, false>(args, data);
1714 }
1715}
1716
1717void llama_kv_cache::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
1718 const uint32_t n_tokens = ubatch->n_tokens;
1719
1720 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1720, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1721
1722 const int64_t n_kv = dst->ne[0];
1723 const int64_t n_stream = dst->ne[3]; // num streams in the current ubatch
1724
1725 GGML_ASSERT(n_tokens%n_stream == 0)if (!(n_tokens%n_stream == 0)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1725, "GGML_ASSERT(%s) failed", "n_tokens%n_stream == 0");
1726
1727 // n_tps == n_tokens_per_stream
1728 const int64_t n_tps = n_tokens/n_stream;
1729
1730 //const int64_t t_start = ggml_time_us();
1731
1732 const args_set_input_kq_mask args = {
1733 /*.hparams =*/ hparams,
1734 /*.ubatch =*/ ubatch,
1735 /*.v_cells =*/ v_cells,
1736 /*.seq_to_stream =*/ seq_to_stream,
1737 /*.n_swa =*/ n_swa,
1738 /*.swa_type =*/ swa_type,
1739 /*.n_kv =*/ n_kv,
1740 /*.n_stream =*/ n_stream,
1741 /*.n_tps =*/ n_tps,
1742 };
1743
1744 if (dst->type == GGML_TYPE_F16) {
1745 set_input_kq_mask_impl<ggml_fp16_t>(args, (ggml_fp16_t *) dst->data, causal_attn);
1746 } else {
1747 set_input_kq_mask_impl<float>(args, (float *) dst->data, causal_attn);
1748 }
1749
1750 //const int64_t t_end = ggml_time_us();
1751
1752 //LLAMA_LOG_ERROR("%s: kq mask time: %0.3f ms\n", __func__, (t_end - t_start)/1000.0);
1753}
1754
1755void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const {
1756 const int64_t n_tokens = ubatch->n_tokens;
1757
1758 GGML_ASSERT(n_stream == 1 && "TODO: support multiple streams")if (!(n_stream == 1 && "TODO: support multiple streams"
)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1758, "GGML_ASSERT(%s) failed", "n_stream == 1 && \"TODO: support multiple streams\""
);
1759 const auto & cells = v_cells[0];
1760
1761 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1761, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1762 GGML_ASSERT(!ubatch->equal_seqs())if (!(!ubatch->equal_seqs())) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1762, "GGML_ASSERT(%s) failed", "!ubatch->equal_seqs()"); // TODO: use ubatch->n_seqs instead of failing
1763
1764 int32_t * data = (int32_t *) dst->data;
1765
1766 const int32_t n_kv = dst->ne[0];
1767
1768 for (int h = 0; h < 1; ++h) {
1769 for (int i = 0; i < n_tokens; ++i) {
1770 for (int j = 0; j < n_kv; ++j) {
1771 // the position when the cells is empty is irrelevant - it will be masked out later in the attention
1772 const llama_pos p0 = cells.is_empty(j) ? -1 : cells.pos_get(j);
1773
1774 data[h*(n_kv*n_tokens) + i*n_kv + j] = llama_relative_position_bucket(p0, ubatch->pos[i], hparams.n_rel_attn_bkts, false);
1775 }
1776 }
1777 }
1778}
1779
1780void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const {
1781 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1781, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1782
1783 const auto n_rot = dst->ne[0];
1784 GGML_ASSERT(attn_rot_hadamard.count(dst->ne[0]))if (!(attn_rot_hadamard.count(dst->ne[0]))) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1784, "GGML_ASSERT(%s) failed", "attn_rot_hadamard.count(dst->ne[0])"
);
1785
1786 memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
1787}
1788
1789void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
1790 GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer))if (!(ggml_backend_buffer_is_host(dst->buffer))) ggml_abort
("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1790, "GGML_ASSERT(%s) failed", "ggml_backend_buffer_is_host(dst->buffer)"
);
1791
1792 const auto n_rot = dst->ne[0];
1793 GGML_ASSERT(attn_rot_hadamard.count(dst->ne[0]))if (!(attn_rot_hadamard.count(dst->ne[0]))) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1793, "GGML_ASSERT(%s) failed", "attn_rot_hadamard.count(dst->ne[0])"
);
1794
1795 memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
1796}
1797
1798size_t llama_kv_cache::total_size() const {
1799 size_t size = 0;
1800
1801 for (const auto & [_, buf] : ctxs_bufs) {
1802 size += ggml_backend_buffer_get_size(buf.get());
1803 }
1804
1805 return size;
1806}
1807
1808size_t llama_kv_cache::size_k_bytes() const {
1809 size_t size_k_bytes = 0;
1810
1811 for (const auto & layer : layers) {
1812 size_k_bytes += ggml_nbytes(layer.k);
1813 }
1814
1815 return size_k_bytes;
1816}
1817
1818size_t llama_kv_cache::size_v_bytes() const {
1819 size_t size_v_bytes = 0;
1820
1821 for (const auto & layer : layers) {
1822 size_v_bytes += layer.v ? ggml_nbytes(layer.v) : 0;
1823 }
1824
1825 return size_v_bytes;
1826}
1827
1828ggml_tensor * llama_kv_cache::build_rope_shift(
1829 const llama_cparams & cparams,
1830 ggml_context * ctx,
1831 ggml_tensor * cur,
1832 ggml_tensor * shift,
1833 ggml_tensor * rot,
1834 ggml_tensor * factors,
1835 float freq_base,
1836 float freq_scale,
1837 uint32_t il) const {
1838 const auto & n_ctx_orig = cparams.n_ctx_orig_yarn;
1839
1840 const auto & yarn_ext_factor = cparams.yarn_ext_factor;
1841 const auto & yarn_beta_fast = cparams.yarn_beta_fast;
1842 const auto & yarn_beta_slow = cparams.yarn_beta_slow;
1843 const auto & yarn_attn_factor = cparams.yarn_attn_factor;
1844
1845 const auto & n_rot = hparams.n_rot(il);
1846 const auto & rope_type = hparams.rope_type == LLAMA_ROPE_TYPE_MROPE || hparams.rope_type == LLAMA_ROPE_TYPE_IMROPE
1847 // @ngxson : this is a workaround
1848 // for M-RoPE, we want to rotate the whole vector when doing KV shift
1849 // a normal RoPE should work, we just need to use the correct ordering
1850 // ref: https://github.com/ggml-org/llama.cpp/pull/13870
1851 ? LLAMA_ROPE_TYPE_NEOX
1852 : hparams.rope_type;
1853 ggml_tensor * tmp;
1854
1855 if (ggml_is_quantized(cur->type)) {
1856 // dequantize to f32 -> RoPE -> quantize back
1857 tmp = ggml_cast(ctx, cur, GGML_TYPE_F32);
1858
1859 // rotate back
1860 tmp = ggml_mul_mat_aux(ctx, tmp, rot);
1861
1862 tmp = ggml_rope_ext(ctx, tmp,
1863 shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
1864 yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow);
1865
1866 // rotate fwd
1867 tmp = ggml_mul_mat_aux(ctx, tmp, rot);
1868
1869 tmp = ggml_cpy(ctx, tmp, cur);
1870 } else {
1871 // we rotate only the first n_rot dimensions
1872 tmp = ggml_rope_ext_inplace(ctx, cur,
1873 shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
1874 yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow);
1875 }
1876
1877 return tmp;
1878}
1879
1880class llm_graph_input_k_shift : public llm_graph_input_i {
1881public:
1882 llm_graph_input_k_shift(const llama_kv_cache * kv_self) : kv_self(kv_self) {}
1883 virtual ~llm_graph_input_k_shift() = default;
1884
1885 void set_input(const llama_ubatch * ubatch) override;
1886
1887 ggml_tensor * k_shift; // I32 [kv_size*n_stream]
1888
1889 // note: assumes k_rot^2 == I
1890 ggml_tensor * k_rot = nullptr;
1891
1892 const llama_kv_cache * kv_self;
1893};
1894
1895void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) {
1896 GGML_UNUSED(ubatch)(void)(ubatch);
1897
1898 if (k_shift) {
1899 kv_self->set_input_k_shift(k_shift);
1900 }
1901
1902 if (k_rot) {
1903 kv_self->set_input_k_rot(k_rot);
1904 }
1905}
1906
1907ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_context * lctx) const {
1908 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
1909 GGML_ASSERT(!other)if (!(!other)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 1909, "GGML_ASSERT(%s) failed", "!other");
1910
1911 auto * ctx = res->get_ctx();
1912 auto * gf = res->get_gf();
1913
1914 auto inp = std::make_unique<llm_graph_input_k_shift>(this);
1915
1916 inp->k_shift = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, (int64_t) get_size()*n_stream);
1917 ggml_set_input(inp->k_shift);
1918
1919 inp->k_rot = build_input_k_rot(ctx);
1920
1921 const auto & cparams = lctx->get_cparams();
1922
1923 for (const auto & layer : layers) {
1924 const uint32_t il = layer.il;
1925
1926 const int64_t n_head_kv = hparams.n_head_kv(il);
1927 const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
1928
1929 const auto n_rot = hparams.n_rot(il);
1930 const auto n_embd_head_k = hparams.n_embd_head_k(il);
1931 const auto n_embd_nope = hparams.n_lora_kv > 0 ? n_embd_head_k - n_rot : 0;
1932
1933 const float freq_base_l = model.get_rope_freq_base (cparams, il);
1934 const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
1935
1936 ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
1937
1938 ggml_tensor * k =
1939 ggml_view_3d(ctx, layer.k,
1940 n_rot, n_head_kv, get_size()*n_stream,
1941 ggml_row_size(layer.k->type, n_embd_head_k),
1942 ggml_row_size(layer.k->type, n_embd_k_gqa),
1943 ggml_row_size(layer.k->type, n_embd_nope));
1944
1945 ggml_tensor * cur = build_rope_shift(cparams, ctx, k, inp->k_shift, inp->k_rot, rope_factors, freq_base_l, freq_scale_l, il);
1946
1947 ggml_build_forward_expand(gf, cur);
1948 }
1949
1950 res->add_input(std::move(inp));
1951
1952 return gf;
1953}
1954
1955void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
1956 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
1957 if (other) {
1958 return;
1959 }
1960
1961 GGML_UNUSED(flags)(void)(flags);
1962
1963 io.write(&n_stream, sizeof(n_stream));
1964
1965 for (uint32_t s = 0; s < n_stream; ++s) {
1966 cell_ranges_t cr { s, {} };
1967
1968 uint32_t cell_count = 0;
1969
1970 const auto & cells = v_cells[s];
1971
1972 // Count the number of cells with the specified seq_id
1973 // Find all the ranges of cells with this seq id (or all, when -1)
1974 uint32_t cell_range_begin = cells.size();
1975
1976 for (uint32_t i = 0; i < cells.size(); ++i) {
1977 bool add_cell = true;
1978
1979 add_cell = add_cell && !cells.is_empty(i);
1980 add_cell = add_cell && (seq_id == -1 || cells.seq_has(i, seq_id));
1981
1982 // check the cell is not SWA-masked
1983 if (add_cell && seq_id != -1) {
1984 const bool is_masked = llama_hparams::is_masked_swa(n_swa, swa_type, cells.pos_get(i), cells.seq_pos_max(seq_id));
1985
1986 add_cell = !is_masked;
1987 }
1988
1989 if (add_cell) {
1990 ++cell_count;
1991 if (cell_range_begin == cells.size()) {
1992 cell_range_begin = i;
1993 }
1994 } else {
1995 if (cell_range_begin != cells.size()) {
1996 cr.data.emplace_back(cell_range_begin, i);
1997 cell_range_begin = cells.size();
1998 }
1999 }
2000 }
2001
2002 if (cell_range_begin != cells.size()) {
2003 cr.data.emplace_back(cell_range_begin, cells.size());
2004 }
2005
2006 // DEBUG CHECK: Sum of cell counts in ranges should equal the total cell count
2007 uint32_t cell_count_check = 0;
2008 for (const auto & range : cr.data) {
2009 cell_count_check += range.second - range.first;
2010 }
2011 GGML_ASSERT(cell_count == cell_count_check)if (!(cell_count == cell_count_check)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 2011, "GGML_ASSERT(%s) failed", "cell_count == cell_count_check"
);
2012
2013 io.write(&cell_count, sizeof(cell_count));
2014
2015 // skip empty streams
2016 if (cell_count == 0) {
2017 continue;
2018 }
2019
2020 state_write_meta(io, cr, seq_id);
2021 state_write_data(io, cr);
2022 }
2023}
2024
2025void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
2026 // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
2027 if (other) {
2028 return;
2029 }
2030
2031 GGML_UNUSED(flags)(void)(flags);
2032
2033 GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()))if (!(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id
< seq_to_stream.size()))) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 2033, "GGML_ASSERT(%s) failed", "seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())"
);
2034
2035 uint32_t n_stream_cur;
2036 io.read(&n_stream_cur, sizeof(n_stream_cur));
2037 if (n_stream_cur != n_stream) {
2038 throwabort_with_suppression(); if (false) std::runtime_error("n_stream mismatch");
2039 }
2040
2041 for (uint32_t s = 0; s < n_stream; ++s) {
2042 uint32_t cell_count;
2043 io.read(&cell_count, sizeof(cell_count));
2044
2045 if (cell_count == 0) {
2046 continue;
2047 }
2048
2049 const uint32_t strm = seq_id == -1 ? s : seq_to_stream[seq_id];
2050
2051 slot_info sinfo;
2052
2053 bool res = true;
2054 res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id);
2055 res = res && state_read_data(io, strm, cell_count, sinfo);
2056
2057 if (!res) {
2058 if (seq_id == -1) {
2059 clear(true);
2060 } else {
2061 seq_rm(seq_id, -1, -1);
2062 }
2063 throwabort_with_suppression(); if (false) std::runtime_error("failed to restore kv cache");
2064 }
2065 }
2066}
2067
2068void llama_kv_cache::state_write_meta(llama_io_write_i & io, const cell_ranges_t & cr, llama_seq_id seq_id) const {
2069 const auto & cells = v_cells[cr.strm];
2070
2071 for (const auto & range : cr.data) {
2072 for (uint32_t i = range.first; i < range.second; ++i) {
2073 std::vector<llama_seq_id> seq_ids;
2074
2075 for (llama_seq_id cur = 0; cur < (int) n_seq_max; ++cur) {
2076 if (cur == seq_id || seq_id == -1) {
2077 if (cells.seq_has(i, cur)) {
2078 seq_ids.push_back(cur);
2079 }
2080 }
2081 }
2082
2083 const llama_pos pos = cells.pos_get(i);
2084 const uint32_t n_seq_id = seq_ids.size();
2085
2086 io.write(&pos, sizeof(pos));
2087 io.write(&n_seq_id, sizeof(n_seq_id));
2088
2089 if (hparams.n_pos_per_embd() > 1) {
2090 const llama_kv_cell_ext ext = cells.ext_get(i);
2091 io.write(&ext, sizeof(ext));
2092 }
2093
2094 for (const auto & seq_id : seq_ids) {
2095 io.write(&seq_id, sizeof(seq_id));
2096 }
2097 }
2098 }
2099}
2100
2101void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t & cr) const {
2102 const auto & cells = v_cells[cr.strm];
2103
2104 const uint32_t v_trans = this->v_trans ? 1 : 0;
2105 const uint32_t n_layer = layers.size();
2106
2107 io.write(&v_trans, sizeof(v_trans));
2108 io.write(&n_layer, sizeof(n_layer));
2109
2110 // Iterate and write all the keys first, each row is a cell
2111 // Get whole range at a time
2112 for (const auto & layer : layers) {
2113 const uint32_t il = layer.il;
2114
2115 const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
2116
2117 auto * k = layer.k_stream[cr.strm];
2118
2119 // Write key type
2120 const int32_t k_type_i = (int32_t) k->type;
2121 io.write(&k_type_i, sizeof(k_type_i));
2122
2123 // Write row size of key
2124 const uint64_t k_size_row = ggml_row_size(k->type, n_embd_k_gqa);
2125 io.write(&k_size_row, sizeof(k_size_row));
2126
2127 // Read each range of cells of k_size length and write out
2128 for (const auto & range : cr.data) {
2129 const size_t range_size = range.second - range.first;
2130 const size_t buf_size = range_size * k_size_row;
2131 io.write_tensor(k, range.first * k_size_row, buf_size);
2132 }
2133 }
2134
2135 if (!v_trans) {
2136 for (const auto & layer : layers) {
2137 const uint32_t il = layer.il;
2138
2139 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
2140
2141 auto * v = layer.v_stream[cr.strm];
2142 if (!v) {
2143 continue;
2144 }
2145
2146 // Write value type
2147 const int32_t v_type_i = (int32_t) v->type;
2148 io.write(&v_type_i, sizeof(v_type_i));
2149
2150 // Write row size of value
2151 const uint64_t v_size_row = ggml_row_size(v->type, n_embd_v_gqa);
2152 io.write(&v_size_row, sizeof(v_size_row));
2153
2154 // Read each range of cells of v_size length and write out
2155 for (const auto & range : cr.data) {
2156 const size_t range_size = range.second - range.first;
2157 const size_t buf_size = range_size * v_size_row;
2158 io.write_tensor(v, range.first * v_size_row, buf_size);
2159 }
2160 }
2161 } else {
2162 // When v is transposed, we also need the element size and get the element ranges from each row
2163 const uint32_t kv_size = cells.size();
2164
2165 for (const auto & layer : layers) {
2166 const uint32_t il = layer.il;
2167
2168 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
2169
2170 auto * v = layer.v_stream[cr.strm];
2171 if (!v) {
2172 continue;
2173 }
2174
2175 // Write value type
2176 const int32_t v_type_i = (int32_t) v->type;
2177 io.write(&v_type_i, sizeof(v_type_i));
2178
2179 // Write element size
2180 const uint32_t v_size_el = ggml_type_size(v->type);
2181 io.write(&v_size_el, sizeof(v_size_el));
2182
2183 // Write GQA embedding size
2184 io.write(&n_embd_v_gqa, sizeof(n_embd_v_gqa));
2185
2186 // For each row, we get the element values of each cell
2187 for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
2188 // Read each range of cells of v_size_el length and write out
2189 for (const auto & range : cr.data) {
2190 const size_t range_size = range.second - range.first;
2191 const size_t src_offset = (range.first + j * kv_size) * v_size_el;
2192 const size_t buf_size = range_size * v_size_el;
2193 io.write_tensor(v, src_offset, buf_size);
2194 }
2195 }
2196 }
2197 }
2198}
2199
2200bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id) {
2201 auto & cells = v_cells[strm];
2202 auto & head = v_heads[strm];
2203
2204 if (dest_seq_id != -1) {
2205 // single sequence
2206 seq_rm(dest_seq_id, -1, -1);
2207
2208 llama_batch_allocr balloc(hparams.n_pos_per_embd());
2209
2210 llama_ubatch ubatch = balloc.ubatch_reserve(cell_count, 1);
2211
2212 ubatch.seq_id_unq[0] = dest_seq_id;
2213
2214 for (uint32_t i = 0; i < cell_count; ++i) {
2215 llama_pos pos;
2216 uint32_t n_seq_id;
2217
2218 io.read(&pos, sizeof(pos));
2219 io.read(&n_seq_id, sizeof(n_seq_id));
2220
2221 if (n_seq_id != 1) {
2222 LLAMA_LOG_ERROR("%s: invalid seq_id-agnostic kv cell\n", __func__)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: invalid seq_id-agnostic kv cell\n"
, __func__);
2223 return false;
2224 }
2225
2226 if (hparams.n_pos_per_embd() > 1) {
2227 llama_kv_cell_ext ext;
2228 io.read(&ext, sizeof(ext));
2229
2230 ubatch.pos[i + ubatch.n_tokens] = ext.y;
2231 ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
2232 }
2233
2234 // read the sequence id, but directly discard it - we will use dest_seq_id instead
2235 {
2236 llama_seq_id seq_id;
2237 io.read(&seq_id, sizeof(seq_id));
2238 }
2239
2240 ubatch.pos[i] = pos;
2241 ubatch.n_seq_id[i] = n_seq_id;
2242 ubatch.seq_id[i] = &dest_seq_id;
2243 }
2244
2245 sinfo = find_slot(ubatch, false);
2246 if (sinfo.empty()) {
2247 LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: failed to find %d available cells in kv cache\n"
, __func__, cell_count);
2248 return false;
2249 }
2250
2251 // TODO: we cannot yet restore llama_kv_cell_ext as the apply_ubatch() does not support it yet
2252 // see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350
2253 apply_ubatch(sinfo, ubatch);
2254
2255 LLAMA_LOG_DEBUG("%s: cell_count = %d, dest_seq_id = %d\n", __func__, cell_count, dest_seq_id)llama_log_internal(GGML_LOG_LEVEL_DEBUG, "%s: cell_count = %d, dest_seq_id = %d\n"
, __func__, cell_count, dest_seq_id);
2256
2257 // DEBUG CHECK: verify that all cells were allocated and have correct seq_id and pos values
2258 GGML_ASSERT(sinfo.n_stream() == 1)if (!(sinfo.n_stream() == 1)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 2258, "GGML_ASSERT(%s) failed", "sinfo.n_stream() == 1");
2259 GGML_ASSERT(sinfo.idxs[0].size() == cell_count)if (!(sinfo.idxs[0].size() == cell_count)) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 2259, "GGML_ASSERT(%s) failed", "sinfo.idxs[0].size() == cell_count"
);
2260 for (uint32_t i = 0; i < cell_count; ++i) {
2261 const uint32_t idx = sinfo.idxs[0][i];
2262 GGML_ASSERT(cells.pos_get(idx) == ubatch.pos[i])if (!(cells.pos_get(idx) == ubatch.pos[i])) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 2262, "GGML_ASSERT(%s) failed", "cells.pos_get(idx) == ubatch.pos[i]"
);
2263 GGML_ASSERT(cells.seq_has(idx, dest_seq_id))if (!(cells.seq_has(idx, dest_seq_id))) ggml_abort("/root/firefox-clang/third_party/llama.cpp/src/llama-kv-cache.cpp"
, 2263, "GGML_ASSERT(%s) failed", "cells.seq_has(idx, dest_seq_id)"
);
2264 }
2265 } else {
2266 // whole KV cache restore
2267
2268 if (cell_count > cells.size()) {
2269 LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: not enough cells in kv cache\n"
, __func__);
2270 return false;
2271 }
2272
2273 clear(true);
2274
2275 for (uint32_t i = 0; i < cell_count; ++i) {
2276 llama_pos pos;
2277 uint32_t n_seq_id;
2278
2279 io.read(&pos, sizeof(pos));
2280 io.read(&n_seq_id, sizeof(n_seq_id));
2281
2282 cells.pos_set(i, pos);
2283
2284 if (hparams.n_pos_per_embd() > 1) {
2285 llama_kv_cell_ext ext;
2286 io.read(&ext, sizeof(ext));
2287 cells.ext_set(i, ext);
2288 }
2289
2290 for (uint32_t j = 0; j < n_seq_id; ++j) {
2291 llama_seq_id seq_id;
2292 io.read(&seq_id, sizeof(seq_id));
2293
2294 if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
2295 LLAMA_LOG_ERROR("%s: invalid seq_id, %d is out of range [0, %u)\n", __func__, seq_id, n_seq_max)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: invalid seq_id, %d is out of range [0, %u)\n"
, __func__, seq_id, n_seq_max);
2296 return false;
2297 }
2298
2299 cells.seq_add(i, seq_id);
2300 }
2301 }
2302
2303 // Create contiguous slot_info for whole cache restore
2304 sinfo.s0 = strm;
2305 sinfo.s1 = strm;
2306 sinfo.resize(1);
2307 sinfo.strm[0] = strm;
2308 sinfo.idxs[0].resize(cell_count);
2309 for (uint32_t i = 0; i < cell_count; ++i) {
2310 sinfo.idxs[0][i] = i;
2311 }
2312
2313 head = 0;
2314 }
2315
2316 return true;
2317}
2318
2319bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) {
2320 auto & cells = v_cells[strm];
2321
2322 uint32_t v_trans;
2323 uint32_t n_layer;
2324
2325 io.read(&v_trans, sizeof(v_trans));
2326 io.read(&n_layer, sizeof(n_layer));
2327
2328 if (n_layer != layers.size()) {
2329 LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, (uint32_t) layers.size())llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched layer count (%u instead of %u)\n"
, __func__, n_layer, (uint32_t) layers.size());
2330 return false;
2331 }
2332
2333 if (cell_count > cells.size()) {
2334 LLAMA_LOG_ERROR("%s: not enough cells in kv cache to restore state (%u > %u)\n", __func__, cell_count, cells.size())llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: not enough cells in kv cache to restore state (%u > %u)\n"
, __func__, cell_count, cells.size());
2335 return false;
2336 }
2337
2338 if (this->v_trans != (bool) v_trans) {
2339 LLAMA_LOG_ERROR("%s: incompatible V transposition\n", __func__)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: incompatible V transposition\n"
, __func__);
2340 return false;
2341 }
2342
2343 // For each layer, read the keys for each cell, one row is one cell, read as one contiguous block
2344 for (const auto & layer : layers) {
2345 const uint32_t il = layer.il;
2346
2347 const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
2348
2349 auto * k = layer.k_stream[strm];
2350
2351 // Read type of key
2352 int32_t k_type_i_ref;
2353 io.read(&k_type_i_ref, sizeof(k_type_i_ref));
2354 const int32_t k_type_i = (int32_t) k->type;
2355 if (k_type_i != k_type_i_ref) {
2356 LLAMA_LOG_ERROR("%s: mismatched key type (%d != %d, layer %d)\n", __func__, k_type_i, k_type_i_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched key type (%d != %d, layer %d)\n"
, __func__, k_type_i, k_type_i_ref, il);
2357 return false;
2358 }
2359
2360 // Read row size of key
2361 uint64_t k_size_row_ref;
2362 io.read(&k_size_row_ref, sizeof(k_size_row_ref));
2363 const size_t k_size_row = ggml_row_size(k->type, n_embd_k_gqa);
2364 if (k_size_row != k_size_row_ref) {
2365 LLAMA_LOG_ERROR("%s: mismatched key row size (%zu != %zu, layer %d)\n", __func__, k_size_row, (size_t) k_size_row_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched key row size (%zu != %zu, layer %d)\n"
, __func__, k_size_row, (size_t) k_size_row_ref, il);
2366 return false;
2367 }
2368
2369 if (cell_count) {
2370 if (sinfo.is_contiguous()) {
2371 // Fast path: contiguous cells, single memcpy
2372 io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row);
2373 } else {
2374 // Slow path: scatter to non-contiguous positions
2375 for (uint32_t i = 0; i < cell_count; ++i) {
2376 const size_t dst_offset = sinfo.idxs[0][i] * k_size_row;
2377 io.read_tensor(k, dst_offset, k_size_row);
2378 }
2379 }
2380 }
2381 }
2382
2383 if (!this->v_trans) {
2384 for (const auto & layer : layers) {
2385 const uint32_t il = layer.il;
2386
2387 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
2388
2389 auto * v = layer.v_stream[strm];
2390 if (!v) {
2391 continue;
2392 }
2393
2394 // Read type of value
2395 int32_t v_type_i_ref;
2396 io.read(&v_type_i_ref, sizeof(v_type_i_ref));
2397 const int32_t v_type_i = (int32_t) v->type;
2398 if (v_type_i != v_type_i_ref) {
2399 LLAMA_LOG_ERROR("%s: mismatched value type (%d != %d, layer %d)\n", __func__, v_type_i, v_type_i_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched value type (%d != %d, layer %d)\n"
, __func__, v_type_i, v_type_i_ref, il);
2400 return false;
2401 }
2402
2403 // Read row size of value
2404 uint64_t v_size_row_ref;
2405 io.read(&v_size_row_ref, sizeof(v_size_row_ref));
2406 const size_t v_size_row = ggml_row_size(v->type, n_embd_v_gqa);
2407 if (v_size_row != v_size_row_ref) {
2408 LLAMA_LOG_ERROR("%s: mismatched value row size (%zu != %zu, layer %d)\n", __func__, v_size_row, (size_t) v_size_row_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched value row size (%zu != %zu, layer %d)\n"
, __func__, v_size_row, (size_t) v_size_row_ref, il);
2409 return false;
2410 }
2411
2412 if (cell_count) {
2413 if (sinfo.is_contiguous()) {
2414 // Fast path: contiguous cells, single memcpy
2415 io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row);
2416 } else {
2417 // Slow path: scatter to non-contiguous positions
2418 for (uint32_t i = 0; i < cell_count; ++i) {
2419 const size_t dst_offset = sinfo.idxs[0][i] * v_size_row;
2420 io.read_tensor(v, dst_offset, v_size_row);
2421 }
2422 }
2423 }
2424 }
2425 } else {
2426 // For each layer, read the values for each cell (transposed)
2427 for (const auto & layer : layers) {
2428 const uint32_t il = layer.il;
2429
2430 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
2431
2432 auto * v = layer.v_stream[strm];
2433 if (!v) {
2434 continue;
2435 }
2436
2437 // Read type of value
2438 int32_t v_type_i_ref;
2439 io.read(&v_type_i_ref, sizeof(v_type_i_ref));
2440 const int32_t v_type_i = (int32_t) v->type;
2441 if (v_type_i != v_type_i_ref) {
2442 LLAMA_LOG_ERROR("%s: mismatched value type (%d != %d, layer %d)\n", __func__, v_type_i, v_type_i_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched value type (%d != %d, layer %d)\n"
, __func__, v_type_i, v_type_i_ref, il);
2443 return false;
2444 }
2445
2446 // Read element size of value
2447 uint32_t v_size_el_ref;
2448 io.read(&v_size_el_ref, sizeof(v_size_el_ref));
2449 const size_t v_size_el = ggml_type_size(v->type);
2450 if (v_size_el != v_size_el_ref) {
2451 LLAMA_LOG_ERROR("%s: mismatched value element size (%zu != %zu, layer %d)\n", __func__, v_size_el, (size_t) v_size_el_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched value element size (%zu != %zu, layer %d)\n"
, __func__, v_size_el, (size_t) v_size_el_ref, il);
2452 return false;
2453 }
2454
2455 // Read GQA embedding size
2456 uint32_t n_embd_v_gqa_ref;
2457 io.read(&n_embd_v_gqa_ref, sizeof(n_embd_v_gqa_ref));
2458 if (n_embd_v_gqa != n_embd_v_gqa_ref) {
2459 LLAMA_LOG_ERROR("%s: mismatched GQA embedding size (%u != %u, layer %d)\n", __func__, n_embd_v_gqa, n_embd_v_gqa_ref, il)llama_log_internal(GGML_LOG_LEVEL_ERROR, "%s: mismatched GQA embedding size (%u != %u, layer %d)\n"
, __func__, n_embd_v_gqa, n_embd_v_gqa_ref, il);
2460 return false;
2461 }
2462
2463 if (cell_count) {
2464 if (sinfo.is_contiguous()) {
2465 // Fast path: contiguous cells
2466 const uint32_t h = sinfo.head();
2467 for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
2468 const size_t dst_offset = (h + j * cells.size()) * v_size_el;
2469 io.read_tensor(v, dst_offset, cell_count * v_size_el);
2470 }
2471 } else {
2472 // Slow path: scatter to non-contiguous positions
2473 for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
2474 for (uint32_t i = 0; i < cell_count; ++i) {
2475 const size_t dst_offset = (sinfo.idxs[0][i] + j * cells.size()) * v_size_el;
2476 io.read_tensor(v, dst_offset, v_size_el);
2477 }
2478 }
2479 }
2480 }
2481 }
2482 }
2483
2484 return true;
2485}
2486
2487//
2488// llama_kv_cache_context
2489//
2490
2491llama_kv_cache_context::llama_kv_cache_context(llama_memory_status status) : status(status) {}
2492
2493llama_kv_cache_context::llama_kv_cache_context(
2494 llama_kv_cache * kv) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv) {
2495 n_kv = kv->get_size();
2496
2497 const uint32_t n_stream = kv->get_n_stream();
2498
2499 // create a dummy slot info - the actual data is irrelevant. we just need to build the graph
2500 sinfos.resize(1);
2501 sinfos[0].s0 = 0;
2502 sinfos[0].s1 = n_stream - 1;
2503 sinfos[0].idxs.resize(n_stream);
2504 for (uint32_t s = 0; s < n_stream; ++s) {
2505 sinfos[0].strm.push_back(s);
2506 sinfos[0].idxs[s].resize(1, 0);
2507 }
2508}
2509
2510llama_kv_cache_context::llama_kv_cache_context(
2511 llama_kv_cache * kv,
2512 llama_context * lctx,
2513 bool do_shift,
2514 stream_copy_info sc_info) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv), lctx(lctx), do_shift(do_shift), sc_info(std::move(sc_info)) {
2515 if (!do_shift && this->sc_info.empty()) {
2516 status = LLAMA_MEMORY_STATUS_NO_UPDATE;
2517 }
2518}
2519
2520llama_kv_cache_context::llama_kv_cache_context(
2521 llama_kv_cache * kv,
2522 llama_kv_cache::slot_info_vec_t sinfos,
2523 std::vector<llama_ubatch> ubatches) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv), sinfos(std::move(sinfos)), ubatches(std::move(ubatches)) {
2524}
2525
2526llama_kv_cache_context::~llama_kv_cache_context() = default;
2527
2528bool llama_kv_cache_context::next() {
2529 assert(status == LLAMA_MEMORY_STATUS_SUCCESS)(static_cast <bool> (status == LLAMA_MEMORY_STATUS_SUCCESS
) ? void (0) : __assert_fail ("status == LLAMA_MEMORY_STATUS_SUCCESS"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
2530
2531 if (++i_cur >= ubatches.size()) {
2532 return false;
2533 }
2534
2535 return true;
2536}
2537
2538bool llama_kv_cache_context::apply() {
2539 assert(!llama_memory_status_is_fail(status))(static_cast <bool> (!llama_memory_status_is_fail(status
)) ? void (0) : __assert_fail ("!llama_memory_status_is_fail(status)"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
2540
2541 // no ubatches -> this is a KV cache update
2542 if (ubatches.empty()) {
2543 kv->update(lctx, do_shift, sc_info);
2544
2545 return true;
2546 }
2547
2548 kv->apply_ubatch(sinfos[i_cur], ubatches[i_cur]);
2549 n_kv = kv->get_n_kv(sinfos[i_cur]);
2550
2551 return true;
2552}
2553
2554llama_memory_status llama_kv_cache_context::get_status() const {
2555 return status;
2556}
2557
2558const llama_ubatch & llama_kv_cache_context::get_ubatch() const {
2559 assert(status == LLAMA_MEMORY_STATUS_SUCCESS)(static_cast <bool> (status == LLAMA_MEMORY_STATUS_SUCCESS
) ? void (0) : __assert_fail ("status == LLAMA_MEMORY_STATUS_SUCCESS"
, __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__
));
2560
2561 return ubatches[i_cur];
2562}
2563
2564uint32_t llama_kv_cache_context::get_n_kv() const {
2565 return n_kv;
2566}
2567
2568ggml_type llama_kv_cache_context::type_k() const {
2569 return kv->type_k();
2570}
2571
2572ggml_type llama_kv_cache_context::type_v() const {
2573 return kv->type_v();
2574}
2575
2576ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) const {
2577 return kv->get_k(ctx, il, n_kv, sinfos[i_cur]);
2578}
2579
2580ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const {
2581 return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
2582}
2583
2584ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
2585 return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
2586}
2587
2588ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const {
2589 return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]);
2590}
2591
2592ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
2593 return kv->build_input_k_idxs(ctx, ubatch);
2594}
2595
2596ggml_tensor * llama_kv_cache_context::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
2597 return kv->build_input_v_idxs(ctx, ubatch);
2598}
2599
2600ggml_tensor * llama_kv_cache_context::build_input_k_rot(ggml_context * ctx) const {
2601 return kv->build_input_k_rot(ctx);
2602}
2603
2604ggml_tensor * llama_kv_cache_context::build_input_v_rot(ggml_context * ctx) const {
2605 return kv->build_input_v_rot(ctx);
2606}
2607
2608void llama_kv_cache_context::set_input_k_shift(ggml_tensor * dst) const {
2609 kv->set_input_k_shift(dst);
2610}
2611
2612void llama_kv_cache_context::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch) const {
2613 kv->set_input_k_idxs(dst, ubatch, sinfos[i_cur]);
2614}
2615
2616void llama_kv_cache_context::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch) const {
2617 kv->set_input_v_idxs(dst, ubatch, sinfos[i_cur]);
2618}
2619
2620void llama_kv_cache_context::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
2621 kv->set_input_kq_mask(dst, ubatch, causal_attn);
2622}
2623
2624void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const {
2625 kv->set_input_pos_bucket(dst, ubatch);
2626}
2627
2628void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const {
2629 kv->set_input_k_rot(dst);
2630}
2631
2632void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const {
2633 kv->set_input_v_rot(dst);
2634}