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