Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-20 01:31:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

move to a new llama_kv_cache_dsa_iswa

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-19 20:47:00 +02:00
parent 6bb8fd0f86
commit d9de2b2062
9 changed files with 682 additions and 196 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CMakeLists.txt
+1
View File
@@ -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
src/llama-graph.cpp
+82 -55
View File
@@ -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<const llama_kv_cache_context *>(params.mctx);
mctx = static_cast<const llama_kv_cache_context *>(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<const llama_kv_cache_dsa_context *>(params.mctx);
mctx = static_cast<const llama_kv_cache_dsa_context *>(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<const llama_kv_cache_dsa_iswa_context *>(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<const llama_kv_cache_dsa_context *>(mctx);
static std::unique_ptr<llm_graph_input_attn_k_dsa> 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<llm_graph_input_attn_k_dsa>(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<const llama_kv_cache_dsa_context *>(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<const llama_kv_cache_dsa_iswa_context *>(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<llm_graph_input_attn_k>(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<llm_graph_input_attn_k_dsa_iswa>(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<const llama_kv_cache_msa_context *>(mctx);
src/llama-graph.h
+45 -15
View File
@@ -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<llm_graph_input_attn_k_dsa> inp_dsa,
std::unique_ptr<llm_graph_input_attn_k> 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<llm_graph_input_attn_k_dsa> inp_dsa;
std::unique_ptr<llm_graph_input_attn_k> 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(
src/llama-kv-cache-dsa-iswa.cpp
+341
View File
@@ -0,0 +1,341 @@
#include "llama-kv-cache-dsa-iswa.h"
#include "llama-impl.h"
#include "llama-batch.h"
#include "llama-model.h"
#include <algorithm>
#include <cassert>
//
// 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<llama_kv_cache_dsa>(
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<llama_kv_cache>(
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<ggml_backend_buffer_type_t, size_t> llama_kv_cache_dsa_iswa::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> 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<llama_ubatch> 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<llama_kv_cache_dsa_iswa_context>(
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<llama_ubatch> 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<llama_kv_cache_dsa_iswa_context>(
this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(sinfos_swa), std::move(ubatches));
} while (false);
return std::make_unique<llama_kv_cache_dsa_iswa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}
llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_full() {
return std::make_unique<llama_kv_cache_dsa_iswa_context>(this);
}
llama_memory_context_ptr llama_kv_cache_dsa_iswa::init_update(llama_context * lctx, bool optimize) {
return std::make_unique<llama_kv_cache_dsa_iswa_context>(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<llama_ubatch> 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<const llama_kv_cache_dsa_context *>(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<const llama_kv_cache_context *>(ctx_swa.get());
}
src/llama-kv-cache-dsa-iswa.h
+134
View File
@@ -0,0 +1,134 @@
#pragma once
#include "llama-kv-cache-dsa.h"
#include <vector>
//
// 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<ggml_backend_buffer_type_t, size_t> 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<llama_kv_cache_dsa> kv_dsa;
std::unique_ptr<llama_kv_cache> 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<llama_ubatch> 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<llama_ubatch> ubatches;
const llama_memory_context_ptr ctx_dsa;
const llama_memory_context_ptr ctx_swa;
const llama_memory_status status;
};
src/llama-kv-cache-dsa.cpp
+6 -101
View File
@@ -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<llama_kv_cache>(
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<llama_kv_cache>(
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<llama_kv_cache>(
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<ggml_backend_buffer_type_t, size_t> llama_kv_cache_dsa::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> 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<llama_kv_cache_dsa_context>(
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_kv_cache_dsa_context>(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<llama_ubatch> 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<const llama_kv_cache_context *>(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<const llama_kv_cache_context *>(ctx_mla_swa.get());
}
const llama_kv_cache_context * llama_kv_cache_dsa_context::get_lid() const {
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
src/llama-kv-cache-dsa.h
+4 -11
View File
@@ -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<llama_kv_cache> kv_mla;
std::unique_ptr<llama_kv_cache> kv_mla_swa;
std::unique_ptr<llama_kv_cache> 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<llama_ubatch> 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<llama_ubatch> 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;
src/llama-model.cpp
+53 -4
View File
@@ -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);
src/models/dots3note.cpp
+16 -10
View File
@@ -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);