char-fork/llama.cpp

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

Aman GuptaDeepseekV4 MTP + DSpark (#25784)596a5795b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
14.2 KiB406 linesraw
1#pragma once
2
3#include "llama-kv-cache.h"
4#include "llama-kv-cache-iswa.h"
5
6#include <map>
7#include <memory>
8#include <unordered_map>
9#include <vector>
10
11class llama_dsv4_comp_state {
12public:
13    using stream_copy_info = llama_kv_cache::stream_copy_info;
14
15    stream_copy_info sc_info;
16
17    llama_dsv4_comp_state(
18            const llama_model & model,
19            bool            offload,
20            bool            unified,
21            uint32_t        n_seq_max,
22            uint32_t        ratio,
23            uint32_t        state_size,
24            uint32_t        n_embd_state,
25            uint32_t        n_rs_seq,
26            const char    * name,
27        const llama_memory_i::layer_filter_cb & filter);
28
29    void clear(llama_seq_id seq_id, bool data);
30    void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
31    void apply_copies(const stream_copy_info & sc_info) const;
32
33    uint32_t get_ratio()      const;
34    uint32_t get_state_size() const;
35    uint32_t get_n_stream()   const;
36    uint32_t get_n_rs_seq()   const;
37    uint32_t get_n_rows()     const;
38
39    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
40
41    void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, const std::vector<uint32_t> & rs_idx) const;
42    void state_read (llama_io_read_i  & io, llama_seq_id seq_id, llama_state_seq_flags flags);
43
44    ggml_tensor * get_kv       (ggml_context * ctx, int32_t il) const;
45    ggml_tensor * get_score    (ggml_context * ctx, int32_t il) const;
46    ggml_tensor * get_kv_all   (ggml_context * ctx, int32_t il) const;
47    ggml_tensor * get_score_all(ggml_context * ctx, int32_t il) const;
48
49    ggml_tensor * cpy_kv   (ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
50    ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
51
52private:
53    struct layer {
54        uint32_t il;
55
56        ggml_tensor * kv;
57        ggml_tensor * score;
58
59        std::vector<ggml_tensor *> kv_stream;
60        std::vector<ggml_tensor *> score_stream;
61    };
62
63    const uint32_t ratio;
64    const uint32_t state_size;
65    const uint32_t n_embd_state;
66    const uint32_t n_stream;
67    const uint32_t n_rs_seq;
68
69    std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
70
71    std::vector<layer> layers;
72
73    std::unordered_map<int32_t, int32_t> map_layer_ids;
74
75    size_t total_size() const;
76};
77
78//
79// llama_kv_cache_dsv4
80//
81
82// DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches.
83// The compressed caches are storage only; DSV4-specific visibility and block
84// planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4.
85// FIXME: currently the cache only supports non-unified mode even if unified flag is passed
86// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819
87
88class llama_kv_cache_dsv4 : public llama_memory_i {
89public:
90    llama_kv_cache_dsv4(
91            const llama_model & model,
92                    ggml_type   type_k,
93                    ggml_type   type_v,
94                         bool   v_trans,
95                         bool   offload,
96                         bool   swa_full,
97                         bool   unified,
98                     uint32_t   kv_size,
99                     uint32_t   n_seq_max,
100                     uint32_t   n_ubatch,
101                     uint32_t   n_pad,
102                     uint32_t   n_rs_seq,
103        const layer_filter_cb & filter,
104        const  layer_reuse_cb & reuse);
105
106    ~llama_kv_cache_dsv4() = default;
107
108    //
109    // llama_memory_i
110    //
111
112    llama_memory_context_ptr init_batch(
113            llama_batch_allocr & balloc,
114            uint32_t n_ubatch,
115            bool embd_all) override;
116
117    llama_memory_context_ptr init_full() override;
118
119    llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;
120
121    bool get_can_shift() const override;
122
123    void clear(bool data) override;
124
125    bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;
126    void seq_cp  (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;
127    void seq_keep(llama_seq_id seq_id)                                                          override;
128    void seq_add (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, llama_pos shift) override;
129    void seq_div (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, int d) override;
130
131    llama_pos seq_pos_min(llama_seq_id seq_id) const override;
132    llama_pos seq_pos_max(llama_seq_id seq_id) const override;
133
134    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;
135
136    void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
137    void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
138
139    //
140    // llama_kv_cache_dsv4 specific API
141    //
142
143    llama_kv_cache_iswa * get_raw() const;
144    llama_kv_cache      * get_csa() const;
145    llama_kv_cache      * get_hca() const;
146    llama_kv_cache      * get_lid() const;
147    llama_dsv4_comp_state * get_csa_state() const;
148    llama_dsv4_comp_state * get_hca_state() const;
149    llama_dsv4_comp_state * get_lid_state() const;
150
151    uint32_t get_n_rs_seq() const;
152    const std::vector<uint32_t> & get_rs_idx() const;
153    void reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches);
154
155private:
156    llama_hparams hparams_raw;
157    llama_hparams hparams_csa;
158    llama_hparams hparams_hca;
159    llama_hparams hparams_lid;
160
161    const uint32_t n_seq_max;
162    const uint32_t n_rs_seq;
163
164    std::vector<uint32_t> rs_idx;
165
166    std::unique_ptr<llama_kv_cache_iswa> kv_raw;
167    std::unique_ptr<llama_kv_cache>      kv_csa;
168    std::unique_ptr<llama_kv_cache>      kv_hca;
169    std::unique_ptr<llama_kv_cache>      kv_lid;
170    std::unique_ptr<llama_dsv4_comp_state> csa_state;
171    std::unique_ptr<llama_dsv4_comp_state> hca_state;
172    std::unique_ptr<llama_dsv4_comp_state> lid_state;
173
174    void clear_compressed(llama_seq_id seq_id, bool data);
175};
176
177// DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept
178// for generic ISWA bookkeeping, but it has no DSV4 layers to expose here.
179class llama_kv_cache_dsv4_raw_context : public llama_memory_context_i {
180public:
181    using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
182
183    llama_kv_cache_dsv4_raw_context(llama_kv_cache_iswa * kv);
184
185    llama_kv_cache_dsv4_raw_context(
186            llama_kv_cache_iswa * kv,
187            llama_context * lctx,
188            bool optimize);
189
190    llama_kv_cache_dsv4_raw_context(
191            llama_kv_cache_iswa * kv,
192            slot_info_vec_t sinfos_base_write,
193            slot_info_vec_t sinfos_swa_write,
194            slot_info_vec_t sinfos_swa_read,
195            std::vector<llama_ubatch> ubatches,
196            std::vector<llama_ubatch> ubatches_write);
197
198    bool next() override;
199    bool apply() override;
200
201    llama_memory_status get_status() const override;
202    const llama_ubatch & get_ubatch() const override;
203
204    uint32_t get_n_kv() const;
205    uint32_t get_n_write() const;
206
207    ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
208    ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
209
210    ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const;
211    ggml_tensor * build_input_k_rot(ggml_context * ctx) const;
212
213    void set_input_k_idxs(ggml_tensor * dst) const;
214    void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
215    void set_input_k_rot(ggml_tensor * dst) const;
216
217private:
218    size_t i_next = 0;
219
220    llama_kv_cache * kv_swa = nullptr;
221
222    slot_info_vec_t sinfos_write;
223    slot_info_vec_t sinfos_read;
224    std::vector<llama_ubatch> ubatches;
225    std::vector<llama_ubatch> ubatches_write;
226
227    const llama_memory_context_ptr ctx_base_mem;
228    const llama_memory_context_ptr ctx_swa_mem;
229
230    uint32_t n_kv = 0;
231
232    const llama_memory_status status;
233};
234
235// DSV4 compressed KV rows are graph outputs, not normal token KV writes.
236// Keep a small context that exposes K tensors without generic apply() semantics.
237class llama_kv_cache_dsv4_comp_context {
238public:
239    using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
240
241    llama_kv_cache_dsv4_comp_context(llama_kv_cache * kv);
242
243    llama_kv_cache_dsv4_comp_context(
244            llama_kv_cache * kv,
245            slot_info_vec_t sinfos,
246            std::vector<llama_ubatch> ubatches);
247
248    bool next();
249
250    uint32_t get_n_kv() const;
251
252    ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
253    ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
254
255    ggml_tensor * build_input_k_rot(ggml_context * ctx) const;
256    void set_input_k_rot(ggml_tensor * dst) const;
257
258private:
259    llama_kv_cache * kv;
260
261    size_t i_cur = 0;
262    slot_info_vec_t sinfos;
263    std::vector<llama_ubatch> ubatches;
264
265    uint32_t n_kv;
266};
267
268class llama_kv_cache_dsv4_context : public llama_memory_context_i {
269public:
270    using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
271    using stream_copy_info = llama_kv_cache::stream_copy_info;
272
273    struct comp_plan {
274        // Per-ubatch recipe for updating compressor state, committing completed
275        // compressed rows, and masking the compressed attention source.
276
277        // APE row ids, i.e. pos % ratio, for the compressor-state updates.
278        std::vector<int32_t> state_pos;
279
280        // Current-ubatch source row ids and unique persistent-state
281        // destination row ids for deterministic ring-state updates.
282        std::vector<int32_t> state_persist_src_idxs;
283        std::vector<int32_t> state_persist_dst_idxs;
284
285        // Device-side rollback restore copies snapshot planes back to the
286        // current compressor-state plane before the graph reads it.
287        std::vector<int32_t> state_restore_src_idxs;
288        std::vector<int32_t> state_restore_dst_idxs;
289
290        // Device-side rollback snapshots copy rows from the graph-local
291        // [persistent_state | current_ubatch_scratch] tensor into rollback
292        // planes after the graph has computed current-token compressor state.
293        std::vector<int32_t> state_snapshot_src_idxs;
294        std::vector<int32_t> state_snapshot_dst_idxs;
295
296        // Flattened source row ids used for state-backed commits. Source rows
297        // index the graph-local [persistent_state | current_ubatch_scratch]
298        // tensor. For overlapped compression the first half is previous rows
299        // and the second half is current rows; a final synthetic zero/-inf row
300        // may be addressed for the first block's previous half.
301        std::vector<int32_t> state_read_idxs;
302
303        // Final compressed-cache row ids written by state-backed commits.
304        // A non-boundary CSA/LID decode step can target a masked scratch row.
305        std::vector<int64_t> state_write_idxs;
306
307        // RoPE positions for state-backed commits.
308        std::vector<int32_t> state_write_pos;
309
310        // Number of completed compressed rows visible for each query token.
311        std::vector<int32_t> n_visible;
312
313        // Number of streams used by the attention graph for this ubatch.
314        int64_t n_stream = 1;
315
316        // Graph-width for compressed rows. This can be larger than n_visible
317        // so masked padding rows do not force a new graph at every CSA block.
318        int64_t n_kv = 0;
319    };
320
321    llama_kv_cache_dsv4_context(llama_memory_status status);
322
323    llama_kv_cache_dsv4_context(
324            llama_kv_cache_dsv4 * kv);
325
326    llama_kv_cache_dsv4_context(
327            llama_kv_cache_dsv4 * kv,
328            llama_context * lctx,
329            bool optimize,
330            stream_copy_info sc_info_csa,
331            stream_copy_info sc_info_hca,
332            stream_copy_info sc_info_lid);
333
334    llama_kv_cache_dsv4_context(
335            llama_kv_cache_dsv4 * kv,
336            slot_info_vec_t sinfos_raw_base_write,
337            slot_info_vec_t sinfos_raw_swa_write,
338            slot_info_vec_t sinfos_raw_swa_read,
339            std::vector<llama_ubatch> ubatches,
340            std::vector<llama_ubatch> ubatches_raw);
341
342    virtual ~llama_kv_cache_dsv4_context();
343
344    //
345    // llama_memory_context_i
346    //
347
348    bool next()  override;
349    bool apply() override;
350
351    llama_memory_status  get_status() const override;
352    const llama_ubatch & get_ubatch() const override;
353
354    //
355    // llama_kv_cache_dsv4_context specific API
356    //
357
358    const llama_kv_cache_dsv4_raw_context * get_raw() const;
359    const llama_kv_cache_dsv4_comp_context * get_csa() const;
360    const llama_kv_cache_dsv4_comp_context * get_hca() const;
361    const llama_kv_cache_dsv4_comp_context * get_lid() const;
362    const llama_dsv4_comp_state       * get_csa_state() const;
363    const llama_dsv4_comp_state       * get_hca_state() const;
364    const llama_dsv4_comp_state       * get_lid_state() const;
365
366    const comp_plan & get_csa_plan() const;
367    const comp_plan & get_hca_plan() const;
368    const comp_plan & get_lid_plan() const;
369
370    const comp_plan & get_csa_plan(const llama_ubatch & ubatch) const;
371    const comp_plan & get_hca_plan(const llama_ubatch & ubatch) const;
372    const comp_plan & get_lid_plan(const llama_ubatch & ubatch) const;
373
374private:
375    size_t i_next = 0;
376
377    std::vector<llama_ubatch> ubatches;
378
379    std::vector<comp_plan> plans_csa;
380    std::vector<comp_plan> plans_hca;
381    std::vector<comp_plan> plans_lid;
382
383    const std::unique_ptr<llama_kv_cache_dsv4_raw_context> ctx_raw;
384    const llama_memory_context_ptr ctx_csa_mem;
385    const llama_memory_context_ptr ctx_hca_mem;
386    const llama_memory_context_ptr ctx_lid_mem;
387
388    const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_csa;
389    const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_hca;
390    const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_lid;
391
392    llama_dsv4_comp_state * csa_state = nullptr;
393    llama_dsv4_comp_state * hca_state = nullptr;
394    llama_dsv4_comp_state * lid_state = nullptr;
395
396    stream_copy_info sc_info_csa;
397    stream_copy_info sc_info_hca;
398    stream_copy_info sc_info_lid;
399
400    bool reserve_plans = false;
401    mutable comp_plan reserve_plan_csa;
402    mutable comp_plan reserve_plan_hca;
403    mutable comp_plan reserve_plan_lid;
404
405    const llama_memory_status status;
406};