char-fork/llama.cpp

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

Aman Guptallama-batch: add n_keep_tail in split_equal for recurrent models (#25278)230ea9d21

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