nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

laguna: compacted sliding-window KV cache (--swa-compress) (#2310) 

* laguna: compacted sliding-window KV cache (--swa-compress)

* llama: fold the model-level ctx-shift gate into get_can_shift

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing 2026-08-13 07:54:36 -05:00 committed by GitHub
parent 2cda8d2daf
commit a10ef3eb00
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 160 additions and 77 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
examples/server/server-context.cpp
View File

@ -316,8 +316,15 @@ void server_context::init() {
{"n_ctx_slot", slot.n_ctx}
});
const int ga_n = params_base.grp_attn_n;
const int ga_w = params_base.grp_attn_w;
int ga_n = params_base.grp_attn_n;
int ga_w = params_base.grp_attn_w;
if (ga_n != 1 && !llama_supports_ctx_shift(slot.ctx)) {
// self-extend re-positions cached rows, which a compacted layer cannot represent
LOG_WARNING("%s\n", "self-extend is not supported by this context's KV cache, it will be disabled");
ga_n = 1;
ga_w = 512;
}
if (ga_n != 1) {
GGML_ASSERT(ga_n > 0 && "ga_n must be positive"); // NOLINT
@ -1818,10 +1825,10 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task)
LOG_WARNING("%s\n", "ctx_shift is not implemented for split mode graph, it will be disabled");
}
}
if (!llama_model_supports_ctx_shift(llama_get_model(slot.ctx))) {
if (!llama_supports_ctx_shift(slot.ctx)) {
if (params_base.ctx_shift) {
params_base.ctx_shift = false;
LOG_WARNING("%s\n", "ctx_shift is not supported by this model's KV cache, it will be disabled");
LOG_WARNING("%s\n", "ctx_shift is not supported by this context's KV cache, it will be disabled");
}
}
{
@ -2951,10 +2958,6 @@ void server_context::process_single_task(server_task&& task) {
send_error(task, "slot save is unsupported for openPangu because per-sequence file state is not implemented", ERROR_TYPE_NOT_SUPPORTED);
break;
}
if (!llama_supports_full_state_io(ctx)) {
send_error(task, "slot save is unsupported with --swa-compress because file-session state is not implemented for compacted contexts", ERROR_TYPE_NOT_SUPPORTED);
break;
}
const size_t token_count = slot->cache_tokens.size();
const int64_t t_start = ggml_time_us();
@ -2998,10 +3001,6 @@ void server_context::process_single_task(server_task&& task) {
queue_tasks.defer(std::move(task));
break;
}
if (!llama_supports_full_state_io(ctx)) {
send_error(task, "slot restore is unsupported with --swa-compress because file-session state is not implemented for compacted contexts", ERROR_TYPE_NOT_SUPPORTED);
break;
}
const int64_t t_start = ggml_time_us();
std::string filename = task.data.at("filename");

5
include/llama.h
View File

@ -724,8 +724,9 @@ extern "C" {
// Currently true for every model; no architecture is excluded from partial KV reuse.
LLAMA_API bool llama_model_supports_partial_kv_reuse(const struct llama_model * model);
// false when the context cannot serialize whole-context or file-session state (--swa-compress); per-sequence buffer state is unaffected
LLAMA_API bool llama_supports_full_state_io(const struct llama_context * ctx);
// The complete answer for a context: the model-level query above, plus the context options it
// cannot see (--swa-compress). Matches the gate the engine applies before a K-shift.
LLAMA_API bool llama_supports_ctx_shift(const struct llama_context * ctx);
LLAMA_API const char * llama_model_arch_string(const struct llama_model * model);

6
src/graphs/build_laguna.cpp
View File

@ -12,7 +12,11 @@ ggml_cgraph * llm_build_context::build_laguna() {
ggml_tensor * KQ_mask = build_inp_KQ_mask();
// Laguna M.1 has only global-attention layers and leaves n_swa at zero; building
// the SWA mask in that case trips the generic SWA precondition.
ggml_tensor * KQ_mask_swa = hparams.n_swa > 0 ? build_inp_KQ_mask_swa() : nullptr;
ggml_tensor * KQ_mask_swa = hparams.n_swa > 0
? kv_self.any_compacted()
? build_swa_mask_for_graph(hparams.n_swa, true)
: build_inp_KQ_mask_swa()
: nullptr;
for (int il = 0; il < n_layer; ++il) {
const bool is_swa = hparams.swa_layers[il];

49
src/llama-build-context.cpp
View File

@ -597,7 +597,7 @@ ggml_tensor * llm_build_context::build_inp_KQ_mask_swa_win(int64_t n_kv_win, boo
}
ggml_tensor * llm_build_context::build_swa_mask_for_graph(uint32_t window, bool compacted, bool * windowed) {
*windowed = false;
if (windowed) *windowed = false;
lctx.swa_window_view = {};
if (window == 0) {
@ -625,7 +625,7 @@ ggml_tensor * llm_build_context::build_swa_mask_for_graph(uint32_t window, bool
view.w_view,
view.win_off,
};
*windowed = true;
if (windowed) *windowed = true;
return build_inp_KQ_mask_swa_win(view.w_view);
}
@ -917,6 +917,7 @@ void llm_build_context::llm_build_kv_store(
const llm_build_cb & cb,
int64_t il) {
const int64_t n_ctx = cparams.n_ctx;
const int32_t n_cache_rows = kv.rows(il);
//const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
@ -952,7 +953,7 @@ void llm_build_context::llm_build_kv_store(
} else {
// note: the V cache is transposed for legacy non-FA layouts
v_cache_view = ggml_view_2d(ctx, kv.v_l[il], n_tokens, n_embd_v_gqa,
( n_ctx)*ggml_element_size(kv.v_l[il]),
(n_cache_rows)*ggml_element_size(kv.v_l[il]),
(kv_head)*ggml_element_size(kv.v_l[il]));
lctx.cache_copies[2*il+1].step = ggml_element_size(kv.v_l[il]);
@ -1992,19 +1993,21 @@ static ggml_tensor * llm_build_kqv(
const llm_build_cb & cb,
int il,
ggml_tensor * sinks = nullptr, int n_swa = 0, int kv_il = -1,
ggml_tensor ** k_cache_view = nullptr, ggml_tensor ** v_cache_view = nullptr) {
ggml_tensor ** k_cache_view = nullptr, ggml_tensor ** v_cache_view = nullptr,
int32_t kv_view_offset = 0) {
const llama_model & model = lctx.model;
const llama_hparams & hparams = lctx.model.hparams;
const llama_cparams & cparams = lctx.cparams;
const int kv_layer = kv_il >= 0 ? kv_il : il;
const int64_t n_ctx = kv.size;
const int32_t n_cache_rows = kv.rows(kv_layer);
const int64_t n_head = hparams.n_head(il);
const int64_t n_head_kv = hparams.n_head_kv(il);
const int64_t n_embd_head_k = hparams.n_embd_head_k(il);
//const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
const int64_t n_embd_head_v = hparams.n_embd_head_v(il);
const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);
const int kv_layer = kv_il >= 0 ? kv_il : il;
struct ggml_tensor * q = ggml_permute(ctx, q_cur, 0, 2, 1, 3);
cb(q, "q", il);
@ -2025,7 +2028,7 @@ static ggml_tensor * llm_build_kqv(
n_embd_head_k, n_kv, n_head_kv,
ggml_row_size(k_cache->type, n_embd_head_k)*n_head_kv, //n_embd_k_gqa),
ggml_row_size(k_cache->type, n_embd_head_k),
0);
(size_t) kv_view_offset*ggml_row_size(k_cache->type, n_embd_head_k)*n_head_kv);
if (k_cache_view) {
*k_cache_view = k;
}
@ -2065,7 +2068,7 @@ static ggml_tensor * llm_build_kqv(
n_embd_head_v, n_kv, n_head_kv,
ggml_row_size(v_cache->type, n_embd_v_gqa),
ggml_row_size(v_cache->type, n_embd_head_v),
0);
(size_t) kv_view_offset*ggml_row_size(v_cache->type, n_embd_v_gqa));
if (v_cache_view) {
*v_cache_view = v;
}
@ -2103,15 +2106,15 @@ static ggml_tensor * llm_build_kqv(
if (kv.v_trans) {
v = ggml_view_3d(ctx, v_cache,
n_kv, n_embd_head_v, n_head_kv,
ggml_element_size(v_cache)*n_ctx,
ggml_element_size(v_cache)*n_ctx*n_embd_head_v,
0);
ggml_element_size(v_cache)*n_cache_rows,
ggml_element_size(v_cache)*n_cache_rows*n_embd_head_v,
(size_t) kv_view_offset*ggml_element_size(v_cache));
} else {
v = ggml_view_3d(ctx, v_cache,
n_embd_head_v, n_kv, n_head_kv,
ggml_row_size(v_cache->type, n_embd_v_gqa),
ggml_row_size(v_cache->type, n_embd_head_v),
0);
(size_t) kv_view_offset*ggml_row_size(v_cache->type, n_embd_v_gqa));
v = ggml_cont(ctx, ggml_transpose(ctx, v));
}
if (v_cache_view) {
@ -2253,7 +2256,8 @@ ggml_tensor * llm_build_context::llm_build_kv(
int32_t n_kv,
float kq_scale,
const llm_build_cb & cb, int il, ggml_tensor * sinks, int n_swa, int kv_il,
ggml_tensor ** k_cache_view, ggml_tensor ** v_cache_view) {
ggml_tensor ** k_cache_view, ggml_tensor ** v_cache_view,
int32_t swa_head) {
const llama_hparams & hparams = lctx.model.hparams;
const llama_cparams & cparams = lctx.cparams;
@ -2283,12 +2287,18 @@ ggml_tensor * llm_build_context::llm_build_kv(
ggml_build_forward_expand(graph, v_cur);
}
const bool compacted = kv.is_compacted(il);
const bool use_swa_window = compacted && lctx.swa_window_view.active;
const int32_t store_head = compacted ? swa_head : kv_head;
const int32_t n_kv_view = use_swa_window ? (int32_t) lctx.swa_window_view.w_view : n_kv;
const int32_t kv_view_offset = use_swa_window ? (int32_t) lctx.swa_window_view.win_off : 0;
if (k_cur || v_cur) {
llm_build_kv_store(lctx, ctx, hparams, cparams, kv, graph, k_cur, v_cur, n_tokens, kv_head, cb, il);
llm_build_kv_store(lctx, ctx, hparams, cparams, kv, graph, k_cur, v_cur, n_tokens, store_head, cb, il);
}
auto cur = llm_build_kqv(ctx, lctx, kv, graph, wo, wo_b, q_cur, kq_mask, n_tokens, n_kv, kq_scale, cb, il, sinks, n_swa, kv_il,
k_cache_view, v_cache_view);
auto cur = llm_build_kqv(ctx, lctx, kv, graph, wo, wo_b, q_cur, kq_mask, n_tokens, n_kv_view, kq_scale, cb, il, sinks, n_swa, kv_il,
k_cache_view, v_cache_view, kv_view_offset);
cb(cur, "kqv_out", il);
return cur;
@ -3442,7 +3452,8 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
if (auto wqkv_gate = model.layers[il].wqkv_gate; wqkv_gate != nullptr) {
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
nullptr, nullptr,
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il);
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il,
nullptr, nullptr, swa_head);
cb(cur, "wqkv", il);
auto gate = llm_build_lora_mm(lctx, ctx0, wqkv_gate, input_normed);
if (model.arch == LLM_ARCH_LAGUNA) {
@ -3483,7 +3494,8 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
} else {
if (gate) {
cur = llm_build_kv(ctx0, lctx, kv_self, gf, nullptr, nullptr,
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il);
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il,
nullptr, nullptr, swa_head);
if (false && cur->ne[1] == 1) { // we need to add GGML_UNARY_OP_SIGMOID to the ops supported by ggml_fused_mul_unary
cur = ggml_fused_mul_unary(ctx0, cur, gate, GGML_UNARY_OP_SIGMOID);
} else {
@ -3500,7 +3512,8 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
} else {
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
model.layers[il].wo, model.layers[il].bo,
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il);
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il,
nullptr, nullptr, swa_head);
}
}

5
src/llama-build-context.h
View File

@ -159,7 +159,7 @@ struct llm_build_context {
ggml_tensor * build_inp_KQ_mask_swa_win(int64_t n_kv_win, bool causal = true);
ggml_tensor * build_swa_mask_for_graph(uint32_t window, bool compacted, bool * windowed);
ggml_tensor * build_swa_mask_for_graph(uint32_t window, bool compacted, bool * windowed = nullptr);
ggml_tensor * build_inp_mean();
@ -501,7 +501,8 @@ struct llm_build_context {
int32_t n_kv,
float kq_scale,
const llm_build_cb & cb, int il, ggml_tensor * sinks = nullptr, int n_swa = 0, int kv_il = -1,
ggml_tensor ** k_cache_view = nullptr, ggml_tensor ** v_cache_view = nullptr);
ggml_tensor ** k_cache_view = nullptr, ggml_tensor ** v_cache_view = nullptr,
int32_t swa_head = -1);
static ggml_tensor * llm_build_ffn(ggml_context * ctx, llama_context & lctx, ggml_tensor * ffn_norm,
ggml_tensor * cur,

6
src/llama-model.cpp
View File

@ -2524,7 +2524,9 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_
}
auto n_head_kv = hparams.n_head_kv(il);
auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k(il)) * n_head_kv*kv_size;
auto v_size = ggml_row_size(type_v, hparams.n_embd_v_gqa(il)) * kv_size;
const uint32_t rows = llama_kv_layer_rows(hparams, il, kv_size, swa_compress && supports_swa_compress(), n_ubatch,
llama_kv_cache::get_padding(flash_attn));
auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k(il)) * n_head_kv*rows;
auto v_size = ggml_row_size(type_v, hparams.n_embd_v_gqa(il)) * rows;
return k_size + v_size;
}

5
src/llama-model.h
View File

@ -602,10 +602,9 @@ struct llama_model {
}
// a compacted sliding-window cache needs the graph to build its KQ mask over the compacted
// layout, and the compacted mask keys on position alone, so it also requires K-only cache
// rows and a single sequence
// layout, and the compacted mask keys on position alone, so it requires a single sequence
bool supports_swa_compress() const {
return arch == LLM_ARCH_OPENPANGU || arch == LLM_ARCH_DEEPSEEK4;
return arch == LLM_ARCH_OPENPANGU || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_LAGUNA;
}
static inline int hadamard_size(int head_size) {

138
src/llama.cpp
View File

@ -760,6 +760,7 @@ bool llama_context::update_cache_copies() {
if (!layer_has_attention_kv(il) || kv_self.k_l[il] == nullptr) {
continue;
}
const int32_t cache_head = kv_self.is_compacted(il) ? (int32_t) kv_self.head_swa : (int32_t) kv_self.head;
auto kl = (ggml_split_tensor_t *)kv_self.k_l[il]->extra;
if (kl) {
GGML_ASSERT(model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN);
@ -776,7 +777,7 @@ bool llama_context::update_cache_copies() {
if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kl->splits[id]) {
return false;
}
c.cpy->view_offs = kv_self.head*c.step;
c.cpy->view_offs = cache_head*c.step;
c.cpy->src[1]->data = (char *)kl->splits[id]->data + c.cpy->view_offs;
c.cpy->data = c.cpy->src[1]->data;
}
@ -787,7 +788,7 @@ bool llama_context::update_cache_copies() {
if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != vl->splits[id]) {
return false;
}
c.cpy->view_offs = kv_self.head*c.step;
c.cpy->view_offs = cache_head*c.step;
c.cpy->src[1]->data = (char *)vl->splits[id]->data + c.cpy->view_offs;
c.cpy->data = c.cpy->src[1]->data;
}
@ -797,7 +798,7 @@ bool llama_context::update_cache_copies() {
printf("%s: K has no copy or is not a copy in layer %d\n", __func__, il);
return false;
}
c.cpy->view_offs = kv_self.head*c.step;
c.cpy->view_offs = cache_head*c.step;
c.cpy->src[1]->data = (char *)kv_self.k_l[il]->data + c.cpy->view_offs;
c.cpy->data = c.cpy->src[1]->data;
if (!kv_self.v_l.empty() && kv_self.v_l[il]) {
@ -806,7 +807,7 @@ bool llama_context::update_cache_copies() {
printf("%s: V has no copy or is not a copy in layer %d\n", __func__, il);
return false;
}
c.cpy->view_offs = kv_self.head*c.step;
c.cpy->view_offs = cache_head*c.step;
c.cpy->src[1]->data = (char *)kv_self.v_l[il]->data + c.cpy->view_offs;
c.cpy->data = c.cpy->src[1]->data;
}
@ -1492,7 +1493,7 @@ static bool llama_kv_cache_init(
if (this_type_k != type_k) {
LLAMA_LOG_INFO("================= Setting K-cache type in layer %2d to %s\n", i, ggml_type_name(this_type_k));
}
int64_t v_ne = int64_t(n_embd_v_row)*kv_size;
int64_t v_ne = int64_t(n_embd_v_row)*cache.rows(i);
auto this_type_v = type_v;
if (model.arch != LLM_ARCH_OPENPANGU && type_v_first != type_v && n_v_first > 0 && i < n_v_first) {
this_type_v = type_v_first;
@ -1512,7 +1513,7 @@ static bool llama_kv_cache_init(
} else if (is_dsv4_k_only) {
k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*cache.rows(i));
} else {
k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size);
k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*cache.rows(i));
v = ggml_new_tensor_1d(ctx, this_type_v, v_ne);
}
@ -1774,6 +1775,14 @@ static void llama_kv_cache_compact_swa(struct llama_context & lctx, uint32_t n_t
const uint32_t src_row = cache.sink_rows + live - W;
const uint32_t dst_row = cache.sink_rows;
auto copy_bytes = [&](ggml_tensor * tensor, size_t src_offset, size_t dst_offset, size_t nbytes) {
if (scratch.size() < nbytes) {
scratch.resize(nbytes);
}
ggml_backend_tensor_get(tensor, scratch.data(), src_offset, nbytes);
ggml_backend_tensor_set(tensor, scratch.data(), dst_offset, nbytes);
};
for (size_t il = 0; il < cache.k_l.size(); ++il) {
if (!cache.is_compacted((int) il) || cache.k_l[il] == nullptr) {
continue;
@ -1782,12 +1791,33 @@ static void llama_kv_cache_compact_swa(struct llama_context & lctx, uint32_t n_t
// kl rows are position-major: kl->ne[1]/rows(il) rows per position (n_head_kv)
const size_t rows_per_pos = (size_t) kl->ne[1] / cache.rows((int) il);
const size_t stride = kl->nb[1] * rows_per_pos;
const size_t nbytes = (size_t) W * stride;
if (scratch.size() < nbytes) {
scratch.resize(nbytes);
copy_bytes(kl, (size_t) src_row*stride, (size_t) dst_row*stride, (size_t) W*stride);
// K-only layouts push nothing to v_l, so it is empty rather than null-filled
ggml_tensor * vl = il < cache.v_l.size() ? cache.v_l[il] : nullptr;
if (vl == nullptr) {
continue;
}
const int32_t n_embd_v_row = llama_kv_v_row_embd(lctx.model, lctx.model.hparams, il);
if (!cache.v_trans) {
const size_t v_stride = ggml_row_size(vl->type, n_embd_v_row);
copy_bytes(vl, (size_t) src_row*v_stride, (size_t) dst_row*v_stride, (size_t) W*v_stride);
} else {
// transposed V is position-minor, so one position is a column, not a row
const size_t v_size_el = ggml_type_size(vl->type);
const size_t v_rows = cache.rows((int) il);
const size_t nbytes = ggml_nbytes(vl);
if (scratch.size() < nbytes) {
scratch.resize(nbytes);
}
ggml_backend_tensor_get(vl, scratch.data(), 0, nbytes);
for (int32_t j = 0; j < n_embd_v_row; ++j) {
uint8_t * row = scratch.data() + (size_t) j*v_rows*v_size_el;
memmove(row + (size_t) dst_row*v_size_el, row + (size_t) src_row*v_size_el, (size_t) W*v_size_el);
}
ggml_backend_tensor_set(vl, scratch.data(), 0, nbytes);
}
ggml_backend_tensor_get(kl, scratch.data(), (size_t) src_row*stride, nbytes);
ggml_backend_tensor_set(kl, scratch.data(), (size_t) dst_row*stride, nbytes);
}
cache.pos_base_swa += (llama_pos) (live - W);
@ -2459,6 +2489,9 @@ static void llama_kv_cache_seq_add(
// re-positioned: rope is baked into the cached k_pe rows and the side state is keyed
// by absolute position. Fail loudly instead of corrupting.
GGML_ASSERT(!cache.s_l_position_strict && "K-shift/context shift is not supported for this model's position-indexed KV cache");
// a compacted layer holds window rows rather than one row per context cell, so the cells this
// would re-position do not describe it
GGML_ASSERT(!cache.any_compacted() && "K-shift/context shift is not supported for a compacted KV cache (--swa-compress)");
uint32_t new_head = cache.size;
@ -2510,6 +2543,7 @@ static void llama_kv_cache_seq_div(
int d) {
// see llama_kv_cache_seq_add: position-strict caches cannot be re-positioned
GGML_ASSERT(!cache.s_l_position_strict && "self-extend/position division is not supported for this model's position-indexed KV cache");
GGML_ASSERT(!cache.any_compacted() && "self-extend/position division is not supported for a compacted KV cache (--swa-compress)");
if (p0 < 0) p0 = 0;
if (p1 < 0) p1 = std::numeric_limits<llama_pos>::max();
@ -7150,10 +7184,12 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) {
//LLAMA_LOG_INFO("(tmp log) KV defrag time: %.3f ms\n", (t_end - t_start)/1000.0);
}
static bool get_can_shift(struct llama_context & lctx) {
bool no_shift = lctx.model.is_mla_model();
no_shift = no_shift || lctx.model.arch == LLM_ARCH_DEEPSEEK4;
static bool get_can_shift(const struct llama_context & lctx) {
bool no_shift = !llama_model_supports_ctx_shift(&lctx.model);
no_shift = no_shift || lctx.model.is_mla_model();
no_shift = no_shift || lctx.model.hparams.rope_type == LLAMA_ROPE_TYPE_IMROPE;
// build_k_shift views n_ctx rows per layer, but a compacted layer holds only its window rows
no_shift = no_shift || lctx.kv_self.any_compacted();
return !no_shift;
}
@ -7957,13 +7993,6 @@ struct llama_context * llama_init_from_model(
return nullptr;
}
if (params.n_seq_max > 1 && params.swa_compress && model->arch == LLM_ARCH_DEEPSEEK4) {
// the compacted window is one position stream (head_swa/pos_base_swa are per-cache, not per-sequence)
LLAMA_LOG_ERROR("%s: --swa-compress supports a single sequence only (requested n_seq_max = %u); run with -np 1\n",
__func__, params.n_seq_max);
return nullptr;
}
if (model->arch == LLM_ARCH_OPENPANGU) {
std::string error_msg;
if (!llama_openpangu_validate_latent_cache_types(params.type_k, params.type_v, &error_msg)) {
@ -8455,6 +8484,13 @@ struct llama_context * llama_init_from_model(
return nullptr;
}
if (params.n_seq_max > 1 && ctx->kv_self.any_compacted()) {
LLAMA_LOG_ERROR("%s: --swa-compress supports a single sequence only (requested n_seq_max = %u); run with -np 1\n",
__func__, params.n_seq_max);
llama_free(ctx);
return nullptr;
}
{
size_t memory_size_k = 0;
size_t memory_size_v = 0;
@ -8705,8 +8741,8 @@ void llama_free(struct llama_context * ctx) {
delete ctx;
}
bool llama_supports_full_state_io(const struct llama_context * ctx) {
return ctx != nullptr && !ctx->kv_self.any_compacted();
bool llama_supports_ctx_shift(const struct llama_context * ctx) {
return ctx != nullptr && get_can_shift(*ctx);
}
const struct llama_vocab* llama_model_get_vocab(const struct llama_model* model) {
@ -9835,7 +9871,7 @@ struct llama_data_write {
continue;
}
// only k_l is compacted; it holds the window at [sink_rows, sink_rows + live_swa())
// compacted K and V hold the window at [sink_rows, sink_rows + live_swa())
if (kv_self.is_compacted((int) il)) {
const size_t live = kv_self.live_swa();
if (live) {
@ -9855,7 +9891,7 @@ struct llama_data_write {
if (v_state == 0) {
for (uint32_t il = 0; il < n_layer; ++il) {
const uint32_t n_embd_v_gqa = llama_kv_v_row_embd(ctx->model, hparams, il);
const bool has_v_cache = kv_self.v_l[il] != nullptr && need_kv;
const bool has_v_cache = kv_self.v_l[il] != nullptr && (need_kv || kv_self.is_compacted((int) il));
// Write value type
const int32_t v_type_i = has_v_cache ? (int32_t) kv_self.v_l[il]->type : -1;
@ -9869,6 +9905,14 @@ struct llama_data_write {
continue;
}
if (kv_self.is_compacted((int) il)) {
const size_t live = kv_self.live_swa();
if (live) {
write_tensor_data(kv_self.v_l[il], kv_self.sink_rows * v_size_row, live * v_size_row, il);
}
continue;
}
// Read each range of cells of v_size length each into tmp_buf and write out
for (const auto & range : cell_ranges) {
const size_t range_size = range.second - range.first;
@ -9879,10 +9923,9 @@ struct llama_data_write {
}
else if (v_state == 1) {
// When v is transposed, we also need the element size and get the element ranges from each row
const uint32_t kv_size = kv_self.size;
for (uint32_t il = 0; il < n_layer; ++il) {
const uint32_t n_embd_v_gqa = llama_kv_v_row_embd(ctx->model, hparams, il);
const bool has_v_cache = kv_self.v_l[il] != nullptr && need_kv;
const bool has_v_cache = kv_self.v_l[il] != nullptr && (need_kv || kv_self.is_compacted((int) il));
// Write value type
const int32_t v_type_i = has_v_cache ? (int32_t) kv_self.v_l[il]->type : -1;
@ -9900,6 +9943,18 @@ struct llama_data_write {
continue;
}
const uint32_t kv_size = kv_self.rows((int) il);
if (kv_self.is_compacted((int) il)) {
const size_t live = kv_self.live_swa();
if (live) {
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
const size_t src_offset = (kv_self.sink_rows + (size_t) j * kv_size) * v_size_el;
write_tensor_data(kv_self.v_l[il], src_offset, live * v_size_el, il);
}
}
continue;
}
// For each row, we get the element values of each cell
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
// Read each range of cells of v_size_el length each into tmp_buf and write out
@ -10513,7 +10568,7 @@ struct llama_data_read {
if (v_state == 0) {
for (uint32_t il = 0; il < n_layer; ++il) {
const uint32_t n_embd_v_gqa = llama_kv_v_row_embd(ctx->model, hparams, il);
const bool has_v_cache = kv_self.v_l[il] != nullptr && need_kv;
const bool has_v_cache = kv_self.v_l[il] != nullptr && (need_kv || kv_self.is_compacted((int) il));
// Read type of value
int32_t v_type_i_ref;
@ -10547,12 +10602,16 @@ struct llama_data_read {
return false;
}
if (cell_count) {
const bool v_compact = kv_self.is_compacted((int) il);
const uint32_t v_rows = v_compact ? kv_self.live_swa() : cell_count;
const uint32_t v_dst = v_compact ? kv_self.sink_rows : kv_self.head;
if (v_rows) {
// Read and set the values for the whole cell range
if (kv_self.v_l[il]->extra) {
read_kv_cache_data_split(ctx, kv_self.v_l[il], read(cell_count * v_size_row), kv_self.head, v_size_row, cell_count, il);
read_kv_cache_data_split(ctx, kv_self.v_l[il], read(v_rows * v_size_row), v_dst, v_size_row, v_rows, il);
} else {
ggml_backend_tensor_set(kv_self.v_l[il], read(cell_count * v_size_row), kv_self.head * v_size_row, cell_count * v_size_row);
ggml_backend_tensor_set(kv_self.v_l[il], read(v_rows * v_size_row), v_dst * v_size_row, v_rows * v_size_row);
}
}
}
@ -10561,7 +10620,7 @@ struct llama_data_read {
// For each layer, read the values for each cell (transposed)
for (uint32_t il = 0; il < n_layer; ++il) {
const uint32_t n_embd_v_gqa = llama_kv_v_row_embd(ctx->model, hparams, il);
const bool has_v_cache = kv_self.v_l[il] != nullptr && need_kv;
const bool has_v_cache = kv_self.v_l[il] != nullptr && (need_kv || kv_self.is_compacted((int) il));
// Read type of value
int32_t v_type_i_ref;
@ -10610,15 +10669,20 @@ struct llama_data_read {
return false;
}
if (cell_count) {
const bool v_compact = kv_self.is_compacted((int) il);
const uint32_t v_rows = v_compact ? kv_self.live_swa() : cell_count;
const uint32_t v_dst = v_compact ? kv_self.sink_rows : kv_self.head;
if (v_rows) {
const size_t v_size_el = ggml_type_size(kv_self.v_l[il]->type);
if (kv_self.v_l[il]->extra) {
throw std::runtime_error("Transposed V cache is not sypported with split mode 'graph'");
}
// For each row in the transposed matrix, read the values for the whole cell range
const uint32_t kv_size = kv_self.rows((int) il);
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
const size_t dst_offset = (kv_self.head + j * kv_self.size) * v_size_el;
ggml_backend_tensor_set(kv_self.v_l[il], read(cell_count * v_size_el), dst_offset, cell_count * v_size_el);
const size_t dst_offset = (v_dst + (size_t) j * kv_size) * v_size_el;
ggml_backend_tensor_set(kv_self.v_l[il], read(v_rows * v_size_el), dst_offset, v_rows * v_size_el);
}
}
}
@ -11392,7 +11456,7 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src,
}
static size_t llama_state_seq_save_file_internal(struct llama_context * ctx, const char * filepath, llama_seq_id seq_id, const llama_token * tokens, size_t n_token_count) {
if (!llama_state_io_supported(ctx, __func__)) {
if (!llama_state_io_supported(ctx, __func__, 0, seq_id)) {
return 0;
}
llama_file file(filepath, "wb");
@ -11416,7 +11480,7 @@ static size_t llama_state_seq_save_file_internal(struct llama_context * ctx, con
}
static size_t llama_state_seq_load_file_internal(struct llama_context * ctx, const char * filepath, llama_seq_id dest_seq_id, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) {
if (!llama_state_io_supported(ctx, __func__)) {
if (!llama_state_io_supported(ctx, __func__, 0, dest_seq_id)) {
return 0;
}
llama_file file(filepath, "rb");