char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenmodel: add dots3-note (#27060)5a32f7b66

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
10.5 KiB341 linesraw
1#include "llama-kv-cache-dsa-iswa.h"
2
3#include "llama-impl.h"
4#include "llama-batch.h"
5#include "llama-model.h"
6
7#include <algorithm>
8#include <cassert>
9
10//
11// llama_kv_cache_dsa_iswa
12//
13
14llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa(
15        const llama_model & model,
16                ggml_type   type_k,
17                ggml_type   type_v,
18                     bool   v_trans,
19                     bool   offload,
20                     bool   swa_full,
21                     bool   unified,
22                 uint32_t   kv_size,
23                 uint32_t   n_seq_max,
24                 uint32_t   n_ubatch,
25                 uint32_t   n_pad,
26    const layer_filter_cb & filter_mla,
27    const layer_filter_cb & filter_lid,
28    const  layer_reuse_cb & reuse) : unified(unified) {
29
30    const auto & hparams = model.hparams;
31
32    // chain filters
33    const layer_filter_cb filter_dsa = [&](int32_t il) {
34        if (filter_mla && !filter_mla(il)) {
35            return false;
36        }
37
38        return !hparams.is_swa(il);
39    };
40
41    const layer_filter_cb filter_swa = [&](int32_t il) {
42        if (filter_mla && !filter_mla(il)) {
43            return false;
44        }
45
46        return hparams.is_swa(il);
47    };
48
49    const uint32_t size_dsa = kv_size;
50
51    // note: the SWA cache is always padded to 256 for performance
52    //       https://github.com/ggml-org/llama.cpp/issues/17037
53    uint32_t size_swa = GGML_PAD(std::min(size_dsa, hparams.n_swa*(unified ? n_seq_max : 1) + n_ubatch), 256);
54
55    // when using full-size SWA cache, we set the SWA cache size to be equal to the base cache size
56    if (swa_full) {
57        LLAMA_LOG_WARN("%s: using full-size SWA cache (ref: %s)\n",
58                __func__, "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055");
59
60        size_swa = size_dsa;
61    }
62
63    LLAMA_LOG_INFO("%s: creating DSA KV cache, size = %u cells\n", __func__, size_dsa);
64
65    kv_dsa = std::make_unique<llama_kv_cache_dsa>(
66            model, type_k, type_v,
67            v_trans, offload, unified, size_dsa, n_seq_max, n_pad,
68            0, LLAMA_SWA_TYPE_NONE, filter_dsa, filter_lid, reuse);
69
70    LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa);
71
72    kv_swa = std::make_unique<llama_kv_cache>(
73            model, hparams, type_k, type_v,
74            v_trans, offload, unified, size_swa, n_seq_max, n_pad,
75            hparams.n_swa, hparams.swa_type, nullptr, filter_swa, reuse, nullptr);
76}
77
78void llama_kv_cache_dsa_iswa::clear(bool data) {
79    kv_dsa->clear(data);
80    kv_swa->clear(data);
81}
82
83bool llama_kv_cache_dsa_iswa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
84    bool res = true;
85
86    res = res & kv_dsa->seq_rm(seq_id, p0, p1);
87    res = res & kv_swa->seq_rm(seq_id, p0, p1);
88
89    return res;
90}
91
92void llama_kv_cache_dsa_iswa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
93    kv_dsa->seq_cp(seq_id_src, seq_id_dst, p0, p1);
94    kv_swa->seq_cp(seq_id_src, seq_id_dst, p0, p1);
95}
96
97void llama_kv_cache_dsa_iswa::seq_keep(llama_seq_id seq_id) {
98    kv_dsa->seq_keep(seq_id);
99    kv_swa->seq_keep(seq_id);
100}
101
102void llama_kv_cache_dsa_iswa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
103    kv_dsa->seq_add(seq_id, p0, p1, shift);
104    kv_swa->seq_add(seq_id, p0, p1, shift);
105}
106
107void llama_kv_cache_dsa_iswa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
108    kv_dsa->seq_div(seq_id, p0, p1, d);
109    kv_swa->seq_div(seq_id, p0, p1, d);
110}
111
112llama_pos llama_kv_cache_dsa_iswa::seq_pos_min(llama_seq_id seq_id) const {
113    // the DSA cache is a superset of the SWA cache, so we can just check the SWA cache
114    return kv_swa->seq_pos_min(seq_id);
115}
116
117llama_pos llama_kv_cache_dsa_iswa::seq_pos_max(llama_seq_id seq_id) const {
118    return kv_swa->seq_pos_max(seq_id);
119}
120
121std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_dsa_iswa::memory_breakdown() const {
122    std::map<ggml_backend_buffer_type_t, size_t> mb = kv_dsa->memory_breakdown();
123    for (const auto & buft_size : kv_swa->memory_breakdown()) {
124        mb[buft_size.first] += buft_size.second;
125    }
126    return mb;
127}
128
129llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
130    GGML_UNUSED(embd_all);
131
132    // first try simple split
133    do {
134        if (!unified) {
135            // requires equal splits, so we skip the simple split
136            break;
137        }
138
139        balloc.split_reset();
140
141        std::vector<llama_ubatch> ubatches;
142        while (true) {
143            auto ubatch = balloc.split_simple(n_ubatch);
144
145            if (ubatch.n_tokens == 0) {
146                break;
147            }
148
149            ubatches.push_back(std::move(ubatch)); // NOLINT
150        }
151
152        if (balloc.get_n_used() < balloc.get_n_tokens()) {
153            // failed to find a suitable split
154            break;
155        }
156
157        auto sinfos_mla = kv_dsa->get_mla()->prepare(ubatches);
158        if (sinfos_mla.empty()) {
159            break;
160        }
161
162        auto sinfos_lid = kv_dsa->get_lid()->prepare(ubatches);
163        if (sinfos_lid.empty()) {
164            break;
165        }
166
167        auto sinfos_swa = kv_swa->prepare(ubatches);
168        if (sinfos_swa.empty()) {
169            break;
170        }
171
172        assert(sinfos_mla.size() == sinfos_swa.size());
173
174        return std::make_unique<llama_kv_cache_dsa_iswa_context>(
175                this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(sinfos_swa), std::move(ubatches));
176    } while (false);
177
178    // if it fails, try equal split
179    do {
180        balloc.split_reset();
181
182        std::vector<llama_ubatch> ubatches;
183        while (true) {
184            auto ubatch = balloc.split_equal(n_ubatch, !unified, 0);
185
186            if (ubatch.n_tokens == 0) {
187                break;
188            }
189
190            ubatches.push_back(std::move(ubatch)); // NOLINT
191        }
192
193        if (balloc.get_n_used() < balloc.get_n_tokens()) {
194            // failed to find a suitable split
195            break;
196        }
197
198        auto sinfos_mla = kv_dsa->get_mla()->prepare(ubatches);
199        if (sinfos_mla.empty()) {
200            break;
201        }
202
203        auto sinfos_lid = kv_dsa->get_lid()->prepare(ubatches);
204        if (sinfos_lid.empty()) {
205            break;
206        }
207
208        auto sinfos_swa = kv_swa->prepare(ubatches);
209        if (sinfos_swa.empty()) {
210            break;
211        }
212
213        assert(sinfos_mla.size() == sinfos_swa.size());
214
215        return std::make_unique<llama_kv_cache_dsa_iswa_context>(
216                this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(sinfos_swa), std::move(ubatches));
217    } while (false);
218
219    return std::make_unique<llama_kv_cache_dsa_iswa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
220}
221
222llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_full() {
223    return std::make_unique<llama_kv_cache_dsa_iswa_context>(this);
224}
225
226llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_update(llama_context * lctx, bool optimize) {
227    return std::make_unique<llama_kv_cache_dsa_iswa_context>(this, lctx, optimize);
228}
229
230bool llama_kv_cache_dsa_iswa::get_can_shift() const {
231    return kv_dsa->get_can_shift() &&
232           kv_swa->get_can_shift() &&
233           kv_dsa->get_mla()->get_size() == kv_swa->get_size();
234}
235
236void llama_kv_cache_dsa_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
237    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
238        kv_dsa->state_write(io, seq_id, flags);
239    }
240
241    kv_swa->state_write(io, seq_id, flags);
242}
243
244void llama_kv_cache_dsa_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
245    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
246        kv_dsa->state_read(io, seq_id, flags);
247    }
248
249    kv_swa->state_read(io, seq_id, flags);
250}
251
252llama_kv_cache_dsa * llama_kv_cache_dsa_iswa::get_dsa() const {
253    return kv_dsa.get();
254}
255
256llama_kv_cache * llama_kv_cache_dsa_iswa::get_swa() const {
257    return kv_swa.get();
258}
259
260//
261// llama_kv_cache_dsa_iswa_context
262//
263
264llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context(llama_memory_status status) : status(status) {}
265
266llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context(
267        llama_kv_cache_dsa_iswa * kv) :
268    ctx_dsa(kv->get_dsa()->init_full()),
269    ctx_swa(kv->get_swa()->init_full()),
270    status(llama_memory_status_combine(ctx_dsa->get_status(), ctx_swa->get_status())) {
271}
272
273llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context(
274        llama_kv_cache_dsa_iswa * kv,
275        llama_context * lctx,
276        bool optimize) :
277    ctx_dsa(kv->get_dsa()->init_update(lctx, optimize)),
278    ctx_swa(kv->get_swa()->init_update(lctx, optimize)),
279    status(llama_memory_status_combine(ctx_dsa->get_status(), ctx_swa->get_status())) {
280}
281
282llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context(
283        llama_kv_cache_dsa_iswa * kv,
284        slot_info_vec_t sinfos_mla,
285        slot_info_vec_t sinfos_lid,
286        slot_info_vec_t sinfos_swa,
287        std::vector<llama_ubatch> ubatches) :
288    ubatches(std::move(ubatches)),
289    // note: here we copy the ubatches. not sure if this is ideal
290    ctx_dsa(new llama_kv_cache_dsa_context(kv->get_dsa(), std::move(sinfos_mla), std::move(sinfos_lid), this->ubatches)),
291    ctx_swa(new llama_kv_cache_context(kv->get_swa(), std::move(sinfos_swa), this->ubatches)),
292    status(llama_memory_status_combine(ctx_dsa->get_status(), ctx_swa->get_status())) {
293}
294
295llama_kv_cache_dsa_iswa_context:: ~llama_kv_cache_dsa_iswa_context() = default;
296
297bool llama_kv_cache_dsa_iswa_context::next() {
298    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
299
300    ctx_dsa->next();
301    ctx_swa->next();
302
303    if (++i_next >= ubatches.size()) {
304        return false;
305    }
306
307    return true;
308}
309
310bool llama_kv_cache_dsa_iswa_context::apply() {
311    assert(!llama_memory_status_is_fail(status));
312
313    bool res = true;
314
315    res = res & ctx_dsa->apply();
316    res = res & ctx_swa->apply();
317
318    return res;
319}
320
321llama_memory_status llama_kv_cache_dsa_iswa_context::get_status() const {
322    return status;
323}
324
325const llama_ubatch & llama_kv_cache_dsa_iswa_context::get_ubatch() const {
326    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
327
328    return ubatches[i_next];
329}
330
331const llama_kv_cache_dsa_context * llama_kv_cache_dsa_iswa_context::get_dsa() const {
332    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
333
334    return static_cast<const llama_kv_cache_dsa_context *>(ctx_dsa.get());
335}
336
337const llama_kv_cache_context * llama_kv_cache_dsa_iswa_context::get_swa() const {
338    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
339
340    return static_cast<const llama_kv_cache_context *>(ctx_swa.get());
341}