Felipe97/llama-cpp-compiled
01.1k
1#pragma once2 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_dsv480//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 block84// 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 passed86// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r355817381987 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_i110 //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 API141 //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 kept178// 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 completed275 // 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-state281 // 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 the286 // 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-local291 // [persistent_state | current_ubatch_scratch] tensor into rollback292 // 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 rows297 // index the graph-local [persistent_state | current_ubatch_scratch]298 // tensor. For overlapped compression the first half is previous rows299 // and the second half is current rows; a final synthetic zero/-inf row300 // 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_visible317 // 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_i346 //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 API356 //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};407