1#include "llama-memory-hybrid-iswa.h"
2
3#include "llama-impl.h"
4#include "llama-model.h"
5#include "llama-context.h"
6
7//
8// llama_memory_hybrid_iswa
9//
10
11llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
12 const llama_model & model,
13 /* attn */
14 ggml_type type_k,
15 ggml_type type_v,
16 bool v_trans,
17 bool swa_full,
18 uint32_t kv_size,
19 uint32_t n_ubatch,
20 uint32_t n_pad,
21 /* recurrent */
22 ggml_type type_r,
23 ggml_type type_s,
24 uint32_t rs_size,
25 /* common */
26 uint32_t n_seq_max,
27 bool offload,
28 bool unified,
29 /* layer filters */
30 const layer_filter_cb & filter_attn,
31 const layer_filter_cb & filter_recr) :
32 hparams(model.hparams),
33 mem_attn(new llama_kv_cache_iswa(
34 model,
35 type_k,
36 type_v,
37 v_trans,
38 offload,
39 swa_full,
40 unified,
41 kv_size,
42 n_seq_max,
43 n_ubatch,
44 n_pad,
45 filter_attn == nullptr ?
46 [&](int32_t il) { return !hparams.is_recurrent(il); }
47 : filter_attn,
48 nullptr
49 )),
50 mem_recr(new llama_memory_recurrent(
51 model,
52 type_r,
53 type_s,
54 offload,
55 rs_size,
56 n_seq_max,
57 filter_recr == nullptr ?
58 [&](int32_t il) { return hparams.is_recurrent(il); }
59 : filter_recr
60 )) {}
61
62llama_memory_context_ptr llama_memory_hybrid_iswa::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
63 do {
64 balloc.split_reset();
65
66 // follow the recurrent pattern for creating the ubatch splits
67 std::vector<llama_ubatch> ubatches;
68
69 while (true) {
70 llama_ubatch ubatch;
71
72 if (embd_all) {
73 // if all tokens are output, split by sequence
74 ubatch = balloc.split_seq(n_ubatch);
75 } else {
76 // TODO: non-sequential equal split can be done if using unified KV cache
77 // for simplicity, we always use sequential equal split for now
78 ubatch = balloc.split_equal(n_ubatch, true);
79 }
80
81 if (ubatch.n_tokens == 0) {
82 break;
83 }
84
85 ubatches.push_back(std::move(ubatch)); // NOLINT
86 }
87
88 if (balloc.get_n_used() < balloc.get_n_tokens()) {
89 // failed to find a suitable split
90 break;
91 }
92
93 // prepare the recurrent batches first
94 if (!mem_recr->prepare(ubatches)) {
95 // TODO: will the recurrent cache be in an undefined context at this point?
96 LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__);
97 return std::make_unique<llama_memory_hybrid_iswa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
98 }
99
100 // prepare the attention cache (iswa version returns both base and swa slot infos)
101 auto sinfos_base = mem_attn->get_base()->prepare(ubatches);
102 if (sinfos_base.empty()) {
103 LLAMA_LOG_ERROR("%s: failed to prepare attention base ubatches\n", __func__);
104 return std::make_unique<llama_memory_hybrid_iswa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
105 }
106
107 auto sinfos_swa = mem_attn->get_swa()->prepare(ubatches);
108 if (sinfos_swa.empty()) {
109 LLAMA_LOG_ERROR("%s: failed to prepare attention swa ubatches\n", __func__);
110 return std::make_unique<llama_memory_hybrid_iswa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
111 }
112
113 return std::make_unique<llama_memory_hybrid_iswa_context>(
114 this, std::move(sinfos_base), std::move(sinfos_swa), std::move(ubatches));
115 } while(false);
116
117 return std::make_unique<llama_memory_hybrid_iswa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
118}
119
120llama_memory_context_ptr llama_memory_hybrid_iswa::init_full() {
121 return std::make_unique<llama_memory_hybrid_iswa_context>(this);
122}
123
124llama_memory_context_ptr llama_memory_hybrid_iswa::init_update(llama_context * lctx, bool optimize) {
125 return std::make_unique<llama_memory_hybrid_iswa_context>(this, lctx, optimize);
126}
127
128bool llama_memory_hybrid_iswa::get_can_shift() const {
129 // Shifting is trivially supported for recurrent
130 return mem_attn->get_can_shift();
131}
132
133void llama_memory_hybrid_iswa::clear(bool data) {
134 mem_attn->clear(data);
135 mem_recr->clear(data);
136}
137
138bool llama_memory_hybrid_iswa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
139 // Try removing from the recurrent cache first since it may fail. If it does
140 // fail, the cache will not have been mutated.
141 if (!mem_recr->seq_rm(seq_id, p0, p1)) {
142 return false;
143 }
144 return mem_attn->seq_rm(seq_id, p0, p1);
145}
146
147void llama_memory_hybrid_iswa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
148 mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1);
149 mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1);
150}
151
152void llama_memory_hybrid_iswa::seq_keep(llama_seq_id seq_id) {
153 mem_attn->seq_keep(seq_id);
154 mem_recr->seq_keep(seq_id);
155}
156
157void llama_memory_hybrid_iswa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
158 mem_attn->seq_add(seq_id, p0, p1, shift);
159 mem_recr->seq_add(seq_id, p0, p1, shift);
160}
161
162void llama_memory_hybrid_iswa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
163 mem_attn->seq_div(seq_id, p0, p1, d);
164 mem_recr->seq_div(seq_id, p0, p1, d);
165}
166
167llama_pos llama_memory_hybrid_iswa::seq_pos_min(llama_seq_id seq_id) const {
168 // the min of the total cache is the max of the two caches' min values
169 return std::max(mem_attn->seq_pos_min(seq_id), mem_recr->seq_pos_min(seq_id));
170}
171
172llama_pos llama_memory_hybrid_iswa::seq_pos_max(llama_seq_id seq_id) const {
173 // the max of the total cache is the min of the two caches' max values
174 return std::min(mem_attn->seq_pos_max(seq_id), mem_recr->seq_pos_max(seq_id));
175}
176
177std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid_iswa::memory_breakdown() const {
178 std::map<ggml_backend_buffer_type_t, size_t> mb = mem_attn->memory_breakdown();
179 for (const auto & buft_size : mem_recr->memory_breakdown()) {
180 mb[buft_size.first] += buft_size.second;
181 }
182 return mb;
183}
184
185void llama_memory_hybrid_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
186 mem_attn->state_write(io, seq_id, flags);
187 mem_recr->state_write(io, seq_id, flags);
188}
189
190void llama_memory_hybrid_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
191 mem_attn->state_read(io, seq_id, flags);
192 mem_recr->state_read(io, seq_id, flags);
193}
194
195llama_kv_cache_iswa * llama_memory_hybrid_iswa::get_mem_attn() const {
196 return mem_attn.get();
197}
198
199llama_memory_recurrent * llama_memory_hybrid_iswa::get_mem_recr() const {
200 return mem_recr.get();
201}
202
203//
204// llama_memory_hybrid_iswa_context
205//
206
207llama_memory_hybrid_iswa_context::llama_memory_hybrid_iswa_context(llama_memory_status status) : status(status) {}
208
209llama_memory_hybrid_iswa_context::llama_memory_hybrid_iswa_context(llama_memory_hybrid_iswa * mem) :
210 ctx_attn(mem->get_mem_attn()->init_full()),
211 ctx_recr(mem->get_mem_recr()->init_full()),
212 status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {
213}
214
215llama_memory_hybrid_iswa_context::llama_memory_hybrid_iswa_context(
216 llama_memory_hybrid_iswa * mem,
217 llama_context * lctx,
218 bool optimize) :
219 ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)),
220 ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)),
221 status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {
222}
223
224llama_memory_hybrid_iswa_context::llama_memory_hybrid_iswa_context(
225 llama_memory_hybrid_iswa * mem,
226 slot_info_vec_t sinfos_base,
227 slot_info_vec_t sinfos_swa,
228 std::vector<llama_ubatch> ubatches) :
229 ubatches(std::move(ubatches)),
230 // note: here we copy the ubatches. not sure if this is ideal
231 ctx_attn(new llama_kv_cache_iswa_context(mem->get_mem_attn(), std::move(sinfos_base), std::move(sinfos_swa), this->ubatches)),
232 ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)),
233 status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {
234}
235
236bool llama_memory_hybrid_iswa_context::next() {
237 assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
238
239 ctx_attn->next();
240 ctx_recr->next();
241
242 if (++i_next >= ubatches.size()) {
243 return false;
244 }
245
246 return true;
247}
248
249bool llama_memory_hybrid_iswa_context::apply() {
250 assert(!llama_memory_status_is_fail(status));
251
252 bool res = true;
253
254 res = res & ctx_attn->apply();
255 res = res & ctx_recr->apply();
256
257 return res;
258}
259
260llama_memory_status llama_memory_hybrid_iswa_context::get_status() const {
261 return status;
262}
263
264const llama_ubatch & llama_memory_hybrid_iswa_context::get_ubatch() const {
265 assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
266 return ubatches[i_next];
267}
268
269const llama_kv_cache_iswa_context * llama_memory_hybrid_iswa_context::get_attn() const {
270 return static_cast<const llama_kv_cache_iswa_context *>(ctx_attn.get());
271}
272
273const llama_memory_recurrent_context * llama_memory_hybrid_iswa_context::get_recr() const {
274 return static_cast<const llama_memory_recurrent_context *>(ctx_recr.get());
275}