From d9de2b2062c1737dfa1e10c38f78fbd084fad33d Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 19 Aug 2026 20:47:00 +0200 Subject: [PATCH] move to a new llama_kv_cache_dsa_iswa --- src/CMakeLists.txt | 1 + src/llama-graph.cpp | 137 +++++++------ src/llama-graph.h | 60 ++++-- src/llama-kv-cache-dsa-iswa.cpp | 341 ++++++++++++++++++++++++++++++++ src/llama-kv-cache-dsa-iswa.h | 134 +++++++++++++ src/llama-kv-cache-dsa.cpp | 107 +--------- src/llama-kv-cache-dsa.h | 15 +- src/llama-model.cpp | 57 +++++- src/models/dots3note.cpp | 26 ++- 9 files changed, 682 insertions(+), 196 deletions(-) create mode 100644 src/llama-kv-cache-dsa-iswa.cpp create mode 100644 src/llama-kv-cache-dsa-iswa.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 39ba3061f7..c6df19f2ec 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -25,6 +25,7 @@ add_library(llama llama-kv-cache.cpp llama-kv-cache-iswa.cpp llama-kv-cache-dsa.cpp + llama-kv-cache-dsa-iswa.cpp llama-kv-cache-msa.cpp llama-kv-cache-dsv4.cpp llama-memory.cpp diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1798eb9929..a570a9b587 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -9,6 +9,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" #include "llama-kv-cache-dsa.h" +#include "llama-kv-cache-dsa-iswa.h" #include "llama-kv-cache-msa.h" #include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" @@ -507,10 +508,12 @@ void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) { } bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) { - const auto * mctx = static_cast(params.mctx); + mctx = static_cast(params.mctx); - this->mctx = mctx; + return can_reuse_impl(params); +} +bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) { bool res = true; res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; @@ -559,12 +562,6 @@ void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) { mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn); - if (mctx->get_mla_swa()) { - mctx->get_mla_swa()->set_input_k_idxs(self_k_idxs_mla_swa, ubatch); - - mctx->get_mla_swa()->set_input_kq_mask(self_kq_mask_mla_swa, ubatch, cparams.causal_attn); - } - mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch); mctx->get_lid()->set_input_kq_mask(self_kq_mask_lid, ubatch, cparams.causal_attn); @@ -573,10 +570,12 @@ void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) { } bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) { - const auto * mctx = static_cast(params.mctx); + mctx = static_cast(params.mctx); - this->mctx = mctx; + return can_reuse_impl(params); +} +bool llm_graph_input_attn_k_dsa::can_reuse_impl(const llm_graph_params & params) { bool res = true; res &= self_k_idxs_mla->ne[0] == params.ubatch.n_tokens; @@ -585,10 +584,24 @@ bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) { res &= can_reuse_kq_mask(self_kq_mask_mla, mctx->get_mla(), params.ubatch, params.cparams); res &= can_reuse_kq_mask(self_kq_mask_lid, mctx->get_lid(), params.ubatch, params.cparams); - if (mctx->get_mla_swa()) { - res &= self_k_idxs_mla_swa && self_k_idxs_mla_swa->ne[0] == params.ubatch.n_tokens; - res &= self_kq_mask_mla_swa && can_reuse_kq_mask(self_kq_mask_mla_swa, mctx->get_mla_swa(), params.ubatch, params.cparams); - } + return res; +} + +void llm_graph_input_attn_k_dsa_iswa::set_input(const llama_ubatch * ubatch) { + inp_dsa->set_input(ubatch); + inp_swa->set_input(ubatch); +} + +bool llm_graph_input_attn_k_dsa_iswa::can_reuse(const llm_graph_params & params) { + mctx = static_cast(params.mctx); + + inp_dsa->mctx = mctx->get_dsa(); + inp_swa->mctx = mctx->get_swa(); + + bool res = true; + + res &= inp_dsa->can_reuse_impl(params); + res &= inp_swa->can_reuse_impl(params); return res; } @@ -2953,57 +2966,48 @@ ggml_tensor * llm_graph_context::build_attn( ggml_build_forward_expand(gf, v_cur); ggml_build_forward_expand(gf, k_cur); - const bool is_swa = hparams.is_swa(il); - - // SWA layers use the dedicated window cache and have no indexer, so no top-k selection - GGML_ASSERT(!is_swa || top_k == nullptr); - - const auto * mctx_cur = is_swa ? inp->mctx->get_mla_swa() : inp->mctx->get_mla(); + const auto * mctx_cur = inp->mctx->get_mla(); // store to KV cache { - const auto & k_idxs = is_swa ? inp->get_k_idxs_mla_swa() : inp->get_k_idxs_mla(); + const auto & k_idxs = inp->get_k_idxs_mla(); ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); } - const auto & kq_mask = is_swa ? inp->get_kq_mask_mla_swa() : inp->get_kq_mask_mla(); + const auto & kq_mask = inp->get_kq_mask_mla(); - ggml_tensor * kq_mask_used = kq_mask; + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); - if (top_k) { - // prepare new kq mask - starts filled with -INFINITY - ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); - // reshape KQ mask into tensor with rows of size 1: - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] - kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] - ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); - // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] - // this will be our source of zero values for unmasking top k mask elements - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); - zeros = ggml_fill(ctx0, zeros, 0.0f); + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); - // modify KQ mask by unmasking elements that are in top_k indices - // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) - ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); - // reshape to restore the original shape of KQ mask: - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] - kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); - - // combine with the original kq mask - kq_mask_used = ggml_add(ctx0, kq_mask_top_k, kq_mask); - } + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_used, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -3232,8 +3236,12 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } -llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { - const auto * mctx_cur = static_cast(mctx); +static std::unique_ptr build_attn_inp_k_dsa_impl( + ggml_context * ctx0, + const llama_ubatch & ubatch, + const llama_hparams & hparams, + const llama_cparams & cparams, + const llama_kv_cache_dsa_context * mctx_cur) { auto inp = std::make_unique(hparams, cparams, mctx_cur); @@ -3244,13 +3252,6 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla; } - if (mctx_cur->get_mla_swa()) { - inp->self_k_idxs_mla_swa = mctx_cur->get_mla_swa()->build_input_k_idxs(ctx0, ubatch); - - inp->self_kq_mask_mla_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla_swa(), ubatch, cparams); - inp->self_kq_mask_mla_swa_cnv = inp->self_kq_mask_mla_swa; - } - { inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch); @@ -3264,9 +3265,35 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { inp->self_k_rot_lid = mctx_cur->get_lid()->build_input_k_rot(ctx0); } + return inp; +} + +llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { + const auto * mctx_cur = static_cast(mctx); + + auto inp = build_attn_inp_k_dsa_impl(ctx0, ubatch, hparams, cparams, mctx_cur); + return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp)); } +llm_graph_input_attn_k_dsa_iswa * llm_graph_context::build_attn_inp_k_dsa_iswa() const { + const auto * mctx_cur = static_cast(mctx); + + auto inp_dsa = build_attn_inp_k_dsa_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_dsa()); + + // build_attn_inp_k_impl rejects SWA caches, so construct the input directly + auto inp_swa = std::make_unique(hparams, cparams, mctx_cur->get_swa()); + + inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); + + inp_swa->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); + inp_swa->self_kq_mask_cnv = inp_swa->self_kq_mask; + + auto inp = std::make_unique(std::move(inp_dsa), std::move(inp_swa), mctx_cur); + + return (llm_graph_input_attn_k_dsa_iswa *) res->add_input(std::move(inp)); +} + llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_enabled) const { const auto * mctx_cur = static_cast(mctx); diff --git a/src/llama-graph.h b/src/llama-graph.h index a8993220eb..b388e028cb 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -23,6 +23,7 @@ struct llama_memory_context_i; class llama_kv_cache_context; class llama_kv_cache_dsa_context; +class llama_kv_cache_dsa_iswa_context; class llama_kv_cache_msa_context; class llama_kv_cache_dsv4_raw_context; class llama_kv_cache_dsv4_context; @@ -374,6 +375,9 @@ public: bool can_reuse(const llm_graph_params & params) override; + // like can_reuse, but does not re-bind mctx + bool can_reuse_impl(const llm_graph_params & params); + ggml_tensor * get_k_idxs() const { return self_k_idxs; } ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } @@ -405,24 +409,22 @@ public: bool can_reuse(const llm_graph_params & params) override; - ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; } - ggml_tensor * get_k_idxs_mla_swa() const { return self_k_idxs_mla_swa; } - ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; } + // like can_reuse, but does not re-bind mctx + bool can_reuse_impl(const llm_graph_params & params); - ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; } - ggml_tensor * get_kq_mask_mla_swa() const { return self_kq_mask_mla_swa_cnv; } - ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; } + ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; } + ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; } - ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch] - ggml_tensor * self_k_idxs_mla_swa = nullptr; // I64 [n_batch] - ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch] + ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; } + ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; } - ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_mla_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_mla_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch] + ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch] + + ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] ggml_tensor * self_k_rot_lid = nullptr; @@ -432,6 +434,32 @@ public: const llama_kv_cache_dsa_context * mctx; }; +// DSA input (full-attention layers + indexer) with K-only input for the SWA layers +class llm_graph_input_attn_k_dsa_iswa : public llm_graph_input_i { +public: + llm_graph_input_attn_k_dsa_iswa( + std::unique_ptr inp_dsa, + std::unique_ptr inp_swa, + const llama_kv_cache_dsa_iswa_context * mctx) : + inp_dsa(std::move(inp_dsa)), + inp_swa(std::move(inp_swa)), + mctx(mctx) { + } + ~llm_graph_input_attn_k_dsa_iswa() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override; + + llm_graph_input_attn_k_dsa * get_dsa() const { return inp_dsa.get(); } + llm_graph_input_attn_k * get_swa() const { return inp_swa.get(); } + + std::unique_ptr inp_dsa; + std::unique_ptr inp_swa; + + const llama_kv_cache_dsa_iswa_context * mctx; +}; + // standard K/V attention input against the base cache, plus destination indices for the indexer key cache class llm_graph_input_attn_kv_msa : public llm_graph_input_attn_kv { public: @@ -1196,6 +1224,8 @@ struct llm_graph_context { llm_graph_input_attn_k_dsa * build_attn_inp_k_dsa() const; + llm_graph_input_attn_k_dsa_iswa * build_attn_inp_k_dsa_iswa() const; + llm_graph_input_attn_kv_msa * build_attn_inp_kv_msa(bool msa_enabled) const; ggml_tensor * build_attn( diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp new file mode 100644 index 0000000000..dc10342a19 --- /dev/null +++ b/src/llama-kv-cache-dsa-iswa.cpp @@ -0,0 +1,341 @@ +#include "llama-kv-cache-dsa-iswa.h" + +#include "llama-impl.h" +#include "llama-batch.h" +#include "llama-model.h" + +#include +#include + +// +// llama_kv_cache_dsa_iswa +// + +llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa( + const llama_model & model, + ggml_type type_k, + ggml_type type_v, + bool v_trans, + bool offload, + bool swa_full, + bool unified, + uint32_t kv_size, + uint32_t n_seq_max, + uint32_t n_ubatch, + uint32_t n_pad, + const layer_filter_cb & filter_mla, + const layer_filter_cb & filter_lid, + const layer_reuse_cb & reuse) : unified(unified) { + + const auto & hparams = model.hparams; + + // chain filters + const layer_filter_cb filter_dsa = [&](int32_t il) { + if (filter_mla && !filter_mla(il)) { + return false; + } + + return !hparams.is_swa(il); + }; + + const layer_filter_cb filter_swa = [&](int32_t il) { + if (filter_mla && !filter_mla(il)) { + return false; + } + + return hparams.is_swa(il); + }; + + const uint32_t size_dsa = kv_size; + + // note: the SWA cache is always padded to 256 for performance + // https://github.com/ggml-org/llama.cpp/issues/17037 + uint32_t size_swa = GGML_PAD(std::min(size_dsa, hparams.n_swa*(unified ? n_seq_max : 1) + n_ubatch), 256); + + // when using full-size SWA cache, we set the SWA cache size to be equal to the base cache size + if (swa_full) { + LLAMA_LOG_WARN("%s: using full-size SWA cache (ref: %s)\n", + __func__, "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); + + size_swa = size_dsa; + } + + LLAMA_LOG_INFO("%s: creating DSA KV cache, size = %u cells\n", __func__, size_dsa); + + kv_dsa = std::make_unique( + model, type_k, type_v, + v_trans, offload, unified, size_dsa, n_seq_max, n_pad, + 0, LLAMA_SWA_TYPE_NONE, filter_dsa, filter_lid, reuse); + + LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); + + kv_swa = std::make_unique( + model, hparams, type_k, type_v, + v_trans, offload, unified, size_swa, n_seq_max, n_pad, + hparams.n_swa, hparams.swa_type, nullptr, filter_swa, reuse, nullptr); +} + +void llama_kv_cache_dsa_iswa::clear(bool data) { + kv_dsa->clear(data); + kv_swa->clear(data); +} + +bool llama_kv_cache_dsa_iswa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + bool res = true; + + res = res & kv_dsa->seq_rm(seq_id, p0, p1); + res = res & kv_swa->seq_rm(seq_id, p0, p1); + + return res; +} + +void 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) { + kv_dsa->seq_cp(seq_id_src, seq_id_dst, p0, p1); + kv_swa->seq_cp(seq_id_src, seq_id_dst, p0, p1); +} + +void llama_kv_cache_dsa_iswa::seq_keep(llama_seq_id seq_id) { + kv_dsa->seq_keep(seq_id); + kv_swa->seq_keep(seq_id); +} + +void llama_kv_cache_dsa_iswa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + kv_dsa->seq_add(seq_id, p0, p1, shift); + kv_swa->seq_add(seq_id, p0, p1, shift); +} + +void llama_kv_cache_dsa_iswa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + kv_dsa->seq_div(seq_id, p0, p1, d); + kv_swa->seq_div(seq_id, p0, p1, d); +} + +llama_pos llama_kv_cache_dsa_iswa::seq_pos_min(llama_seq_id seq_id) const { + // the DSA cache is a superset of the SWA cache, so we can just check the SWA cache + return kv_swa->seq_pos_min(seq_id); +} + +llama_pos llama_kv_cache_dsa_iswa::seq_pos_max(llama_seq_id seq_id) const { + return kv_swa->seq_pos_max(seq_id); +} + +std::map llama_kv_cache_dsa_iswa::memory_breakdown() const { + std::map mb = kv_dsa->memory_breakdown(); + for (const auto & buft_size : kv_swa->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + return mb; +} + +llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { + GGML_UNUSED(embd_all); + + // first try simple split + do { + if (!unified) { + // requires equal splits, so we skip the simple split + break; + } + + balloc.split_reset(); + + std::vector ubatches; + while (true) { + auto ubatch = balloc.split_simple(n_ubatch); + + if (ubatch.n_tokens == 0) { + break; + } + + ubatches.push_back(std::move(ubatch)); // NOLINT + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + // failed to find a suitable split + break; + } + + auto sinfos_mla = kv_dsa->get_mla()->prepare(ubatches); + if (sinfos_mla.empty()) { + break; + } + + auto sinfos_lid = kv_dsa->get_lid()->prepare(ubatches); + if (sinfos_lid.empty()) { + break; + } + + auto sinfos_swa = kv_swa->prepare(ubatches); + if (sinfos_swa.empty()) { + break; + } + + assert(sinfos_mla.size() == sinfos_swa.size()); + + return std::make_unique( + this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(sinfos_swa), std::move(ubatches)); + } while (false); + + // if it fails, try equal split + do { + balloc.split_reset(); + + std::vector ubatches; + while (true) { + auto ubatch = balloc.split_equal(n_ubatch, !unified, 0); + + if (ubatch.n_tokens == 0) { + break; + } + + ubatches.push_back(std::move(ubatch)); // NOLINT + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + // failed to find a suitable split + break; + } + + auto sinfos_mla = kv_dsa->get_mla()->prepare(ubatches); + if (sinfos_mla.empty()) { + break; + } + + auto sinfos_lid = kv_dsa->get_lid()->prepare(ubatches); + if (sinfos_lid.empty()) { + break; + } + + auto sinfos_swa = kv_swa->prepare(ubatches); + if (sinfos_swa.empty()) { + break; + } + + assert(sinfos_mla.size() == sinfos_swa.size()); + + return std::make_unique( + this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(sinfos_swa), std::move(ubatches)); + } while (false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_full() { + return std::make_unique(this); +} + +llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_update(llama_context * lctx, bool optimize) { + return std::make_unique(this, lctx, optimize); +} + +bool llama_kv_cache_dsa_iswa::get_can_shift() const { + return kv_dsa->get_can_shift() && + kv_swa->get_can_shift() && + kv_dsa->get_mla()->get_size() == kv_swa->get_size(); +} + +void llama_kv_cache_dsa_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + kv_dsa->state_write(io, seq_id, flags); + } + + kv_swa->state_write(io, seq_id, flags); +} + +void llama_kv_cache_dsa_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + kv_dsa->state_read(io, seq_id, flags); + } + + kv_swa->state_read(io, seq_id, flags); +} + +llama_kv_cache_dsa * llama_kv_cache_dsa_iswa::get_dsa() const { + return kv_dsa.get(); +} + +llama_kv_cache * llama_kv_cache_dsa_iswa::get_swa() const { + return kv_swa.get(); +} + +// +// llama_kv_cache_dsa_iswa_context +// + +llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context(llama_memory_status status) : status(status) {} + +llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context( + llama_kv_cache_dsa_iswa * kv) : + ctx_dsa(kv->get_dsa()->init_full()), + ctx_swa(kv->get_swa()->init_full()), + status(llama_memory_status_combine(ctx_dsa->get_status(), ctx_swa->get_status())) { +} + +llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context( + llama_kv_cache_dsa_iswa * kv, + llama_context * lctx, + bool optimize) : + ctx_dsa(kv->get_dsa()->init_update(lctx, optimize)), + ctx_swa(kv->get_swa()->init_update(lctx, optimize)), + status(llama_memory_status_combine(ctx_dsa->get_status(), ctx_swa->get_status())) { +} + +llama_kv_cache_dsa_iswa_context::llama_kv_cache_dsa_iswa_context( + llama_kv_cache_dsa_iswa * kv, + slot_info_vec_t sinfos_mla, + slot_info_vec_t sinfos_lid, + slot_info_vec_t sinfos_swa, + std::vector ubatches) : + ubatches(std::move(ubatches)), + // note: here we copy the ubatches. not sure if this is ideal + ctx_dsa(new llama_kv_cache_dsa_context(kv->get_dsa(), std::move(sinfos_mla), std::move(sinfos_lid), this->ubatches)), + ctx_swa(new llama_kv_cache_context(kv->get_swa(), std::move(sinfos_swa), this->ubatches)), + status(llama_memory_status_combine(ctx_dsa->get_status(), ctx_swa->get_status())) { +} + +llama_kv_cache_dsa_iswa_context:: ~llama_kv_cache_dsa_iswa_context() = default; + +bool llama_kv_cache_dsa_iswa_context::next() { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + ctx_dsa->next(); + ctx_swa->next(); + + if (++i_next >= ubatches.size()) { + return false; + } + + return true; +} + +bool llama_kv_cache_dsa_iswa_context::apply() { + assert(!llama_memory_status_is_fail(status)); + + bool res = true; + + res = res & ctx_dsa->apply(); + res = res & ctx_swa->apply(); + + return res; +} + +llama_memory_status llama_kv_cache_dsa_iswa_context::get_status() const { + return status; +} + +const llama_ubatch & llama_kv_cache_dsa_iswa_context::get_ubatch() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ubatches[i_next]; +} + +const llama_kv_cache_dsa_context * llama_kv_cache_dsa_iswa_context::get_dsa() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return static_cast(ctx_dsa.get()); +} + +const llama_kv_cache_context * llama_kv_cache_dsa_iswa_context::get_swa() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return static_cast(ctx_swa.get()); +} diff --git a/src/llama-kv-cache-dsa-iswa.h b/src/llama-kv-cache-dsa-iswa.h new file mode 100644 index 0000000000..28cf95bf05 --- /dev/null +++ b/src/llama-kv-cache-dsa-iswa.h @@ -0,0 +1,134 @@ +#pragma once + +#include "llama-kv-cache-dsa.h" + +#include + +// +// llama_kv_cache_dsa_iswa +// + +// utilizes two child memories: llama_kv_cache_dsa for the full-attention (DSA) layers and llama_kv_cache for the SWA layers + +class llama_kv_cache_dsa_iswa : public llama_memory_i { +public: + llama_kv_cache_dsa_iswa( + const llama_model & model, + ggml_type type_k, + ggml_type type_v, + bool v_trans, + bool offload, + bool swa_full, + bool unified, + uint32_t kv_size, + uint32_t n_seq_max, + uint32_t n_ubatch, + uint32_t n_pad, + const layer_filter_cb & filter_mla, + const layer_filter_cb & filter_lid, + const layer_reuse_cb & reuse); + + ~llama_kv_cache_dsa_iswa() = default; + + // + // llama_memory_i + // + + llama_memory_context_ptr init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) override; + + llama_memory_context_ptr init_full() override; + + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + bool get_can_shift() const override; + + void clear(bool data) override; + + bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; + void seq_keep(llama_seq_id seq_id) override; + void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override; + + llama_pos seq_pos_min(llama_seq_id seq_id) const override; + llama_pos seq_pos_max(llama_seq_id seq_id) const override; + + std::map memory_breakdown() const override; + + // state write/load + + void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; + void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + + // + // llama_kv_cache_dsa_iswa specific API + // + + llama_kv_cache_dsa * get_dsa() const; + llama_kv_cache * get_swa() const; + +private: + const bool unified; + + std::unique_ptr kv_dsa; + std::unique_ptr kv_swa; +}; + +class llama_kv_cache_dsa_iswa_context : public llama_memory_context_i { +public: + using slot_info_vec_t = llama_kv_cache::slot_info_vec_t; + + // used for errors + llama_kv_cache_dsa_iswa_context(llama_memory_status status); + + // used to create a full-cache context + llama_kv_cache_dsa_iswa_context( + llama_kv_cache_dsa_iswa * kv); + + // used to create an update context + llama_kv_cache_dsa_iswa_context( + llama_kv_cache_dsa_iswa * kv, + llama_context * lctx, + bool optimize); + + // used to create a batch processing context from a batch + llama_kv_cache_dsa_iswa_context( + llama_kv_cache_dsa_iswa * kv, + slot_info_vec_t sinfos_mla, + slot_info_vec_t sinfos_lid, + slot_info_vec_t sinfos_swa, + std::vector ubatches); + + virtual ~llama_kv_cache_dsa_iswa_context(); + + // + // llama_memory_context_i + // + + bool next() override; + bool apply() override; + + llama_memory_status get_status() const override; + const llama_ubatch & get_ubatch() const override; + + // + // llama_kv_cache_dsa_iswa_context specific API + // + + const llama_kv_cache_dsa_context * get_dsa() const; + const llama_kv_cache_context * get_swa() const; + +private: + // the index of the next ubatch to process + size_t i_next = 0; + + std::vector ubatches; + + const llama_memory_context_ptr ctx_dsa; + const llama_memory_context_ptr ctx_swa; + + const llama_memory_status status; +}; diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 5192d6197d..96cb045d2e 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -20,7 +20,6 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( bool unified, uint32_t kv_size, uint32_t n_seq_max, - uint32_t n_ubatch, uint32_t n_pad, uint32_t n_swa, llama_swa_type swa_type, @@ -29,41 +28,12 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( const layer_reuse_cb & reuse) : hparams_lid(model.hparams), n_stream(unified ? 1 : n_seq_max) { - const bool has_swa = swa_type != LLAMA_SWA_TYPE_NONE; - - const layer_filter_cb filter_mla_full = [&](int32_t il) { - if (filter_mla && !filter_mla(il)) { - return false; - } - - return !(has_swa && model.hparams.is_swa(il)); - }; - - const layer_filter_cb filter_mla_swa = [&](int32_t il) { - if (filter_mla && !filter_mla(il)) { - return false; - } - - return model.hparams.is_swa(il); - }; - LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size); kv_mla = std::make_unique( model, model.hparams, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, - 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_mla_full, reuse, nullptr); - - if (has_swa) { - const uint32_t size_swa = GGML_PAD(std::min(kv_size, n_swa*(unified ? n_seq_max : 1) + n_ubatch), 256); - - LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); - - kv_mla_swa = std::make_unique( - model, model.hparams, type_k, type_v, - v_trans, offload, unified, size_swa, n_seq_max, n_pad, - n_swa, swa_type, nullptr, filter_mla_swa, reuse, nullptr); - } + n_swa, swa_type, nullptr, filter_mla, reuse, nullptr); // we use llama_kv_cache for caching indexer keys // by hand-tweaking some hparams we fool it to create @@ -77,18 +47,14 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); - // the lightning indexer only exists on full-attention layers, so no SWA handling here kv_lid = std::make_unique( model, hparams_lid, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, - 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, reuse, nullptr); + n_swa, swa_type, nullptr, filter_lid, reuse, nullptr); } void llama_kv_cache_dsa::clear(bool data) { kv_mla->clear(data); - if (kv_mla_swa) { - kv_mla_swa->clear(data); - } kv_lid->clear(data); } @@ -96,9 +62,6 @@ bool llama_kv_cache_dsa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) bool res = true; res = res & kv_mla->seq_rm(seq_id, p0, p1); - if (kv_mla_swa) { - res = res & kv_mla_swa->seq_rm(seq_id, p0, p1); - } res = res & kv_lid->seq_rm(seq_id, p0, p1); return res; @@ -106,33 +69,21 @@ bool llama_kv_cache_dsa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) void llama_kv_cache_dsa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { kv_mla->seq_cp(seq_id_src, seq_id_dst, p0, p1); - if (kv_mla_swa) { - kv_mla_swa->seq_cp(seq_id_src, seq_id_dst, p0, p1); - } kv_lid->seq_cp(seq_id_src, seq_id_dst, p0, p1); } void llama_kv_cache_dsa::seq_keep(llama_seq_id seq_id) { kv_mla->seq_keep(seq_id); - if (kv_mla_swa) { - kv_mla_swa->seq_keep(seq_id); - } kv_lid->seq_keep(seq_id); } void llama_kv_cache_dsa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { kv_mla->seq_add(seq_id, p0, p1, shift); - if (kv_mla_swa) { - kv_mla_swa->seq_add(seq_id, p0, p1, shift); - } kv_lid->seq_add(seq_id, p0, p1, shift); } void llama_kv_cache_dsa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { kv_mla->seq_div(seq_id, p0, p1, d); - if (kv_mla_swa) { - kv_mla_swa->seq_div(seq_id, p0, p1, d); - } kv_lid->seq_div(seq_id, p0, p1, d); } @@ -146,11 +97,6 @@ llama_pos llama_kv_cache_dsa::seq_pos_max(llama_seq_id seq_id) const { std::map llama_kv_cache_dsa::memory_breakdown() const { std::map mb = kv_mla->memory_breakdown(); - if (kv_mla_swa) { - for (const auto & buft_size : kv_mla_swa->memory_breakdown()) { - mb[buft_size.first] += buft_size.second; - } - } for (const auto & buft_size : kv_lid->memory_breakdown()) { mb[buft_size.first] += buft_size.second; } @@ -187,14 +133,6 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch( break; } - llama_kv_cache::slot_info_vec_t sinfos_mla_swa; - if (kv_mla_swa) { - sinfos_mla_swa = kv_mla_swa->prepare(ubatches); - if (sinfos_mla_swa.empty()) { - break; - } - } - auto sinfos_lid = kv_lid->prepare(ubatches); if (sinfos_lid.empty()) { break; @@ -203,7 +141,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch( assert(sinfos_mla.size() == sinfos_lid.size()); return std::make_unique( - this, std::move(sinfos_mla), std::move(sinfos_mla_swa), std::move(sinfos_lid), std::move(ubatches)); + this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(ubatches)); } while (false); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); @@ -219,24 +157,17 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_update(llama_context * lctx, b bool llama_kv_cache_dsa::get_can_shift() const { return kv_mla->get_can_shift() && - (!kv_mla_swa || kv_mla_swa->get_can_shift()) && kv_lid->get_can_shift() && kv_mla->get_size() == kv_lid->get_size(); } void llama_kv_cache_dsa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { kv_mla->state_write(io, seq_id, flags); - if (kv_mla_swa) { - kv_mla_swa->state_write(io, seq_id, flags); - } kv_lid->state_write(io, seq_id, flags); } void llama_kv_cache_dsa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { kv_mla->state_read(io, seq_id, flags); - if (kv_mla_swa) { - kv_mla_swa->state_read(io, seq_id, flags); - } kv_lid->state_read(io, seq_id, flags); } @@ -244,10 +175,6 @@ llama_kv_cache * llama_kv_cache_dsa::get_mla() const { return kv_mla.get(); } -llama_kv_cache * llama_kv_cache_dsa::get_mla_swa() const { - return kv_mla_swa.get(); -} - llama_kv_cache * llama_kv_cache_dsa::get_lid() const { return kv_lid.get(); } @@ -261,11 +188,8 @@ llama_kv_cache_dsa_context::llama_kv_cache_dsa_context(llama_memory_status statu llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( llama_kv_cache_dsa * kv) : ctx_mla(kv->get_mla()->init_full()), - ctx_mla_swa(kv->get_mla_swa() ? kv->get_mla_swa()->init_full() : nullptr), ctx_lid(kv->get_lid()->init_full()), - status(llama_memory_status_combine( - llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status()), - ctx_mla_swa ? ctx_mla_swa->get_status() : LLAMA_MEMORY_STATUS_NO_UPDATE)) { + status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) { } llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( @@ -273,27 +197,20 @@ llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( llama_context * lctx, bool optimize) : ctx_mla(kv->get_mla()->init_update(lctx, optimize)), - ctx_mla_swa(kv->get_mla_swa() ? kv->get_mla_swa()->init_update(lctx, optimize) : nullptr), ctx_lid(kv->get_lid()->init_update(lctx, optimize)), - status(llama_memory_status_combine( - llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status()), - ctx_mla_swa ? ctx_mla_swa->get_status() : LLAMA_MEMORY_STATUS_NO_UPDATE)) { + status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) { } llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( llama_kv_cache_dsa * kv, slot_info_vec_t sinfos_mla, - slot_info_vec_t sinfos_mla_swa, slot_info_vec_t sinfos_lid, std::vector ubatches) : ubatches(std::move(ubatches)), // note: here we copy the ubatches. not sure if this is ideal ctx_mla(new llama_kv_cache_context(kv->get_mla(), std::move(sinfos_mla), this->ubatches)), - ctx_mla_swa(kv->get_mla_swa() ? new llama_kv_cache_context(kv->get_mla_swa(), std::move(sinfos_mla_swa), this->ubatches) : nullptr), ctx_lid(new llama_kv_cache_context(kv->get_lid(), std::move(sinfos_lid), this->ubatches)), - status(llama_memory_status_combine( - llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status()), - ctx_mla_swa ? ctx_mla_swa->get_status() : LLAMA_MEMORY_STATUS_NO_UPDATE)) { + status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) { } llama_kv_cache_dsa_context:: ~llama_kv_cache_dsa_context() = default; @@ -302,9 +219,6 @@ bool llama_kv_cache_dsa_context::next() { assert(status == LLAMA_MEMORY_STATUS_SUCCESS); ctx_mla->next(); - if (ctx_mla_swa) { - ctx_mla_swa->next(); - } ctx_lid->next(); if (++i_next >= ubatches.size()) { @@ -320,9 +234,6 @@ bool llama_kv_cache_dsa_context::apply() { bool res = true; res = res & ctx_mla->apply(); - if (ctx_mla_swa) { - res = res & ctx_mla_swa->apply(); - } res = res & ctx_lid->apply(); return res; @@ -344,12 +255,6 @@ const llama_kv_cache_context * llama_kv_cache_dsa_context::get_mla() const { return static_cast(ctx_mla.get()); } -const llama_kv_cache_context * llama_kv_cache_dsa_context::get_mla_swa() const { - assert(status == LLAMA_MEMORY_STATUS_SUCCESS); - - return static_cast(ctx_mla_swa.get()); -} - const llama_kv_cache_context * llama_kv_cache_dsa_context::get_lid() const { assert(status == LLAMA_MEMORY_STATUS_SUCCESS); diff --git a/src/llama-kv-cache-dsa.h b/src/llama-kv-cache-dsa.h index b7b403a763..e74fc4d910 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -11,7 +11,6 @@ // utilizes two instances of llama_kv_cache: // - the first instance is for caching key tensors of the model, // - the second instance is for caching lightning indexer key tensors -// when swa_type != NONE, a third instance holds the SWA layers (iswa-style, window-sized) class llama_kv_cache_dsa : public llama_memory_i { public: @@ -24,7 +23,6 @@ public: bool unified, uint32_t kv_size, uint32_t n_seq_max, - uint32_t n_ubatch, uint32_t n_pad, uint32_t n_swa, llama_swa_type swa_type, @@ -71,9 +69,8 @@ public: // llama_kv_cache_dsa specific API // - llama_kv_cache * get_mla() const; - llama_kv_cache * get_mla_swa() const; // null when the model has no SWA layers - llama_kv_cache * get_lid() const; + llama_kv_cache * get_mla() const; + llama_kv_cache * get_lid() const; private: // we keep indexer KV cache hparams instance here as llama_kv_cache stores only reference to it @@ -81,7 +78,6 @@ private: const uint32_t n_stream = 1; std::unique_ptr kv_mla; - std::unique_ptr kv_mla_swa; std::unique_ptr kv_lid; }; @@ -106,7 +102,6 @@ public: llama_kv_cache_dsa_context( llama_kv_cache_dsa * kv, slot_info_vec_t sinfos_base, - slot_info_vec_t sinfos_swa, slot_info_vec_t sinfos_ik, std::vector ubatches); @@ -126,9 +121,8 @@ public: // llama_kv_cache_dsa_context specific API // - const llama_kv_cache_context * get_mla() const; - const llama_kv_cache_context * get_mla_swa() const; // null when the model has no SWA layers - const llama_kv_cache_context * get_lid() const; + const llama_kv_cache_context * get_mla() const; + const llama_kv_cache_context * get_lid() const; private: //llama_kv_cache_dsa * kv; @@ -139,7 +133,6 @@ private: std::vector ubatches; const llama_memory_context_ptr ctx_mla; - const llama_memory_context_ptr ctx_mla_swa; const llama_memory_context_ptr ctx_lid; const llama_memory_status status; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index faaa69ceed..9a2fb13c16 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -11,6 +11,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" #include "llama-kv-cache-dsa.h" +#include "llama-kv-cache-dsa-iswa.h" #include "llama-kv-cache-msa.h" #include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" @@ -2140,7 +2141,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_GLM_DSA: case LLM_ARCH_DEEPSEEK32: - case LLM_ARCH_DOTS3NOTE: { if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) { // The NextN/MTP draft head runs dense MLA (no DSA indexer), so the @@ -2173,8 +2173,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (hparams.n_layer_nextn > 0) { filter_mla = [&](uint32_t il) { return il < hparams.n_layer(); }; } - const bool lid_by_types = arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE; - llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (!lid_by_types || hparams.is_indexer_full(il)); }; + llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (arch != LLM_ARCH_GLM_DSA || hparams.is_indexer_full(il)); }; res = new llama_kv_cache_dsa( *this, @@ -2185,7 +2184,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.kv_unified, cparams.n_ctx_seq, cparams.n_seq_max, - cparams.n_ubatch, 1, hparams.n_swa, hparams.swa_type, @@ -2194,6 +2192,57 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } break; + case LLM_ARCH_DOTS3NOTE: + { + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); + + if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) { + // MTP draft context: plain attention KV cache holding only the nextn layer + llama_kv_cache::layer_filter_cb filter = + [&](uint32_t il) { return il >= hparams.n_layer(); }; + + res = new llama_kv_cache( + *this, + hparams, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + 1, + hparams.n_swa, + hparams.swa_type, + nullptr, + filter, + nullptr, + nullptr); + } else { + // main context: DSA cache for the trunk full-attention layers plus a window-sized SWA cache + llama_kv_cache::layer_filter_cb filter_mla = nullptr; + if (hparams.n_layer_nextn > 0) { + filter_mla = [&](uint32_t il) { return il < hparams.n_layer(); }; + } + llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_indexer_full(il); }; + + res = new llama_kv_cache_dsa_iswa( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + params.swa_full, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + 1, + filter_mla, + filter_lid, + nullptr); + } + } break; case LLM_ARCH_DEEPSEEK4: { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); diff --git a/src/models/dots3note.cpp b/src/models/dots3note.cpp index e7fc43d1b8..00a008c2c9 100644 --- a/src/models/dots3note.cpp +++ b/src/models/dots3note.cpp @@ -173,7 +173,7 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p ggml_tensor * inp_pos = build_inp_pos(); - llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa(); + llm_graph_input_attn_k_dsa_iswa * inp_attn = build_attn_inp_k_dsa_iswa(); ggml_tensor * inp_out_ids = build_inp_out_ids(); @@ -234,14 +234,14 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p cb(indexer_k, "indexer_k", il); // perform Hadamard transform on indexer q and k - indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q); + indexer_q = ggml_mul_mat(ctx0, inp_attn->get_dsa()->self_k_rot_lid, indexer_q); cb(indexer_q, "indexer_q", il); - indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k); + indexer_k = ggml_mul_mat(ctx0, inp_attn->get_dsa()->self_k_rot_lid, indexer_k); cb(indexer_k, "indexer_k", il); // store indexer keys to KV cache - const auto * mctx_lid = inp_attn_dsa->mctx->get_lid(); - const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid(); + const auto * mctx_lid = inp_attn->get_dsa()->mctx->get_lid(); + const auto & k_idxs_lid = inp_attn->get_dsa()->get_k_idxs_lid(); ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il)); ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur); @@ -260,7 +260,7 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p ggml_tensor * indexer_score = nullptr; if (cparams.fused_lid) { - indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid()); + indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn->get_dsa()->get_kq_mask_lid()); cb(indexer_score, "indexer_score", il); res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il}); } else { @@ -290,7 +290,7 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); cb(indexer_score, "indexer_score", il); - ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid(); + ggml_tensor * indexer_kq_mask = inp_attn->get_dsa()->get_kq_mask_lid(); indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask); cb(indexer_score, "indexer_score", il); } @@ -377,9 +377,15 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p cb(Vcur, "Vcur", il); // apply the head-wise output gate before o_proj, so wo stays out of build_attn - cur = build_attn(inp_attn_dsa, - nullptr, nullptr, nullptr, - Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il); + if (is_swa) { + cur = build_attn(inp_attn->get_swa(), + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, kq_scale, il); + } else { + cur = build_attn(inp_attn->get_dsa(), + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il); + } cb(cur, "attn_out", il); ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);