Commit fb5997274 for llama.cpp
commit fb5997274ec0d0b4d8fe067eb55f87c0ad405478
Author: Anton I <inflyer@hotmail.com>
Date: Fri Oct 30 04:15:06 2026 -0400
kv-cache : keep rotation constants on the GPU (#31073)
* kv-cache : keep rotation constants on the GPU
Assisted-by: Codex
* tests : verify resident KV rotation buffers
Assisted-by: Codex
* kv-cache : document remaining multi-device rotation work
Assisted-by: Codex
* kv-cache : support per-layer rotation buffers
Assisted-by: Codex
* kv-cache : fix hybrid SWA rotation inputs and mixed placement check
Assisted-by: Codex
* models : use per-layer rotation in mistral-large4 and dflash
* tests: drop the rotation buffer tests
---------
Co-authored-by: Pascal <admin@serveurperso.com>
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 7a942663d..43eb46efb 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -554,14 +554,6 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
if (self_kq_mask && self_kq_mask->buffer) {
mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
}
-
- if (self_k_rot && self_k_rot->buffer) {
- mctx->set_input_k_rot(self_k_rot);
- }
-
- if (self_v_rot && self_v_rot->buffer) {
- mctx->set_input_v_rot(self_v_rot);
- }
}
bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) {
@@ -646,11 +638,6 @@ void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
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);
-
- // left unallocated when the indexer does not use the rotation
- if (self_k_rot_lid && self_k_rot_lid->buffer) {
- mctx->get_lid()->set_input_k_rot(self_k_rot_lid);
- }
}
bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) {
@@ -715,22 +702,6 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
}
-
- if (self_k_rot && self_k_rot->buffer) {
- mctx->get_base()->set_input_k_rot(self_k_rot);
- }
-
- if (self_v_rot && self_v_rot->buffer) {
- mctx->get_base()->set_input_v_rot(self_v_rot);
- }
-
- if (self_k_rot_swa && self_k_rot_swa->buffer) {
- mctx->get_swa()->set_input_k_rot(self_k_rot_swa);
- }
-
- if (self_v_rot_swa && self_v_rot_swa->buffer) {
- mctx->get_swa()->set_input_v_rot(self_v_rot_swa);
- }
}
bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
@@ -782,14 +753,6 @@ void llm_graph_input_attn_k_iswa::set_input(const llama_ubatch * ubatch) {
if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
}
-
- if (self_k_rot && self_k_rot->buffer) {
- mctx->get_base()->set_input_k_rot(self_k_rot);
- }
-
- if (self_k_rot_swa && self_k_rot_swa->buffer) {
- mctx->get_swa()->set_input_k_rot(self_k_rot_swa);
- }
}
bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
@@ -1069,10 +1032,6 @@ void llm_graph_input_dsv4_raw::set_input(const llama_ubatch * ubatch) {
if (self_kq_mask && self_kq_mask->buffer) {
mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
}
-
- if (self_k_rot) {
- mctx->set_input_k_rot(self_k_rot);
- }
}
void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) {
@@ -1087,18 +1046,6 @@ void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) {
dsv4_set_comp_inputs(inp_csa, plan_csa, "csa", debug > 0, ubatch->n_tokens, n_stream);
dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream);
dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream);
-
- if (inp_csa.k_rot && inp_csa.k_rot->buffer) {
- mctx->get_csa()->set_input_k_rot(inp_csa.k_rot);
- }
-
- if (inp_hca.k_rot && inp_hca.k_rot->buffer) {
- mctx->get_hca()->set_input_k_rot(inp_hca.k_rot);
- }
-
- if (inp_lid.k_rot && inp_lid.k_rot->buffer) {
- mctx->get_lid()->set_input_k_rot(inp_lid.k_rot);
- }
}
bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) {
@@ -1176,14 +1123,6 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
}
- if (inp_attn->self_k_rot) {
- mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot);
- }
-
- if (inp_attn->self_v_rot) {
- mctx->get_attn()->set_input_v_rot(inp_attn->self_v_rot);
- }
-
const int64_t n_rs = mctx->get_recr()->get_n_rs();
if (inp_rs->s_copy) {
@@ -1287,22 +1226,6 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) {
attn_ctx->get_swa()->set_input_kq_mask(inp_attn->self_kq_mask_swa, ubatch, cparams.causal_attn);
}
- if (inp_attn->self_k_rot) {
- attn_ctx->get_base()->set_input_k_rot(inp_attn->self_k_rot);
- }
-
- if (inp_attn->self_v_rot) {
- attn_ctx->get_base()->set_input_v_rot(inp_attn->self_v_rot);
- }
-
- if (inp_attn->self_k_rot_swa) {
- attn_ctx->get_swa()->set_input_k_rot(inp_attn->self_k_rot_swa);
- }
-
- if (inp_attn->self_v_rot_swa) {
- attn_ctx->get_swa()->set_input_v_rot(inp_attn->self_v_rot_swa);
- }
-
const int64_t n_rs = mctx->get_recr()->get_n_rs();
if (inp_rs->s_copy) {
@@ -3218,8 +3141,8 @@ static std::unique_ptr<llm_graph_input_attn_kv> build_attn_inp_kv_impl(
inp->self_kq_mask_cnv = inp->self_kq_mask;
}
- inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0);
- inp->self_v_rot = mctx_cur->build_input_v_rot(ctx0);
+ inp->self_k_rot = mctx_cur->get_k_rot();
+ inp->self_v_rot = mctx_cur->get_v_rot();
return inp;
}
@@ -3247,13 +3170,13 @@ ggml_tensor * llm_graph_context::build_attn(
int il) const {
GGML_ASSERT(v_mla == nullptr);
- if (inp->self_k_rot) {
- q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);
- k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);
+ if (inp->self_k_rot[il]) {
+ q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot[il]);
+ k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot[il]);
}
- if (inp->self_v_rot) {
- v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
+ if (inp->self_v_rot[il]) {
+ v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot[il]);
}
// these nodes are added to the graph together so that they are not reordered
@@ -3292,8 +3215,8 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
- if (inp->self_v_rot) {
- cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
+ if (inp->self_v_rot[il]) {
+ cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot[il]);
}
if (wo) {
@@ -3526,8 +3449,8 @@ ggml_tensor * llm_graph_context::build_attn(
int il) const {
const bool is_swa = hparams.is_swa(il);
- auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot;
- auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot;
+ auto * k_rot = is_swa ? inp->self_k_rot_swa[il] : inp->self_k_rot[il];
+ auto * v_rot = is_swa ? inp->self_v_rot_swa[il] : inp->self_v_rot[il];
if (k_rot) {
q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
@@ -3621,7 +3544,7 @@ ggml_tensor * llm_graph_context::build_attn(
int il) const {
const bool is_swa = hparams.is_swa(il);
- auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot;
+ auto * k_rot = is_swa ? inp->self_k_rot_swa[il] : inp->self_k_rot[il];
if (k_rot) {
q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
@@ -3758,7 +3681,7 @@ static std::unique_ptr<llm_graph_input_attn_k_dsa> build_attn_inp_k_dsa_impl(
inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy);
inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid;
- inp->self_k_rot_lid = mctx_cur->get_lid()->build_input_k_rot(ctx0);
+ inp->self_k_rot_lid = mctx_cur->get_lid()->get_k_rot();
}
return inp;
@@ -3808,8 +3731,8 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_
inp->self_kq_mask_cnv = inp->self_kq_mask;
}
- inp->self_k_rot = mctx_base->build_input_k_rot(ctx0);
- inp->self_v_rot = mctx_base->build_input_v_rot(ctx0);
+ inp->self_k_rot = mctx_base->get_k_rot();
+ inp->self_v_rot = mctx_base->get_v_rot();
if (msa_enabled) {
inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch);
@@ -3844,11 +3767,11 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
}
- inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
- inp->self_v_rot = mctx_cur->get_base()->build_input_v_rot(ctx0);
+ inp->self_k_rot = mctx_cur->get_base()->get_k_rot();
+ inp->self_v_rot = mctx_cur->get_base()->get_v_rot();
- inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0);
- inp->self_v_rot_swa = mctx_cur->get_swa()->build_input_v_rot(ctx0);
+ inp->self_k_rot_swa = mctx_cur->get_swa()->get_k_rot();
+ inp->self_v_rot_swa = mctx_cur->get_swa()->get_v_rot();
return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));
}
@@ -3874,9 +3797,9 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
}
- inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
+ inp->self_k_rot = mctx_cur->get_base()->get_k_rot();
- inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0);
+ inp->self_k_rot_swa = mctx_cur->get_swa()->get_k_rot();
return (llm_graph_input_attn_k_iswa *) res->add_input(std::move(inp));
}
@@ -3895,15 +3818,15 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
inp_raw->self_kq_mask = dsv4_build_raw_kq_mask(ctx0, raw_ctx, ubatch, cparams, n_stream);
inp_raw->self_kq_mask_cnv = inp_raw->self_kq_mask;
- inp_raw->self_k_rot = raw_ctx->build_input_k_rot(ctx0);
+ inp_raw->self_k_rot = raw_ctx->get_k_rot();
auto inp = std::make_unique<llm_graph_input_dsv4>(cparams, std::move(inp_raw), mctx_cur);
dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", cparams, n_stream);
dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", cparams, n_stream);
dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", cparams, n_stream);
- inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0);
- inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0);
- inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0);
+ inp->inp_csa.k_rot = mctx_cur->get_csa()->get_k_rot();
+ inp->inp_hca.k_rot = mctx_cur->get_hca()->get_k_rot();
+ inp->inp_lid.k_rot = mctx_cur->get_lid()->get_k_rot();
return (llm_graph_input_dsv4 *) res->add_input(std::move(inp));
}
@@ -4088,6 +4011,12 @@ llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa()
inp_attn->self_kq_mask_swa_cnv = inp_attn->self_kq_mask_swa;
}
+ inp_attn->self_k_rot = attn_ctx->get_base()->get_k_rot();
+ inp_attn->self_v_rot = attn_ctx->get_base()->get_v_rot();
+
+ inp_attn->self_k_rot_swa = attn_ctx->get_swa()->get_k_rot();
+ inp_attn->self_v_rot_swa = attn_ctx->get_swa()->get_v_rot();
+
auto inp = std::make_unique<llm_graph_input_mem_hybrid_iswa>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);
return (llm_graph_input_mem_hybrid_iswa *) res->add_input(std::move(inp));
diff --git a/src/llama-graph.h b/src/llama-graph.h
index c1a7d36ee..f9497de5a 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -365,8 +365,8 @@ public:
ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
// note: assumes v_rot^2 == I
- ggml_tensor * self_k_rot = nullptr;
- ggml_tensor * self_v_rot = nullptr;
+ std::vector<ggml_tensor *> self_k_rot;
+ std::vector<ggml_tensor *> self_v_rot;
// note: these have to be copies because in order to be able to reuse a graph, its inputs
// need to carry these parameters with them. otherwise, they can point to freed
@@ -446,7 +446,7 @@ public:
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;
+ std::vector<ggml_tensor *> self_k_rot_lid;
const llama_hparams hparams;
const llama_cparams cparams;
@@ -534,11 +534,11 @@ public:
ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_k_rot = nullptr;
- ggml_tensor * self_v_rot = nullptr;
+ std::vector<ggml_tensor *> self_k_rot;
+ std::vector<ggml_tensor *> self_v_rot;
- ggml_tensor * self_k_rot_swa = nullptr;
- ggml_tensor * self_v_rot_swa = nullptr;
+ std::vector<ggml_tensor *> self_k_rot_swa;
+ std::vector<ggml_tensor *> self_v_rot_swa;
const llama_hparams hparams;
const llama_cparams cparams;
@@ -576,8 +576,8 @@ public:
ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_k_rot = nullptr;
- ggml_tensor * self_k_rot_swa = nullptr;
+ std::vector<ggml_tensor *> self_k_rot;
+ std::vector<ggml_tensor *> self_k_rot_swa;
const llama_hparams hparams;
const llama_cparams cparams;
@@ -606,7 +606,7 @@ public:
ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_k_rot = nullptr;
+ std::vector<ggml_tensor *> self_k_rot;
const llama_cparams cparams;
@@ -629,7 +629,7 @@ public:
ggml_tensor * kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * k_rot = nullptr;
+ std::vector<ggml_tensor *> k_rot;
};
llm_graph_input_dsv4(
diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp
index 027f3c61a..c02e15bdd 100644
--- a/src/llama-kv-cache-dsv4.cpp
+++ b/src/llama-kv-cache-dsv4.cpp
@@ -2074,8 +2074,8 @@ ggml_tensor * llama_kv_cache_dsv4_raw_context::build_input_k_idxs(ggml_context *
return k_idxs;
}
-ggml_tensor * llama_kv_cache_dsv4_raw_context::build_input_k_rot(ggml_context * ctx) const {
- return kv_swa->build_input_k_rot(ctx);
+std::vector<ggml_tensor *> llama_kv_cache_dsv4_raw_context::get_k_rot() const {
+ return kv_swa->get_k_rot();
}
void llama_kv_cache_dsv4_raw_context::set_input_k_idxs(ggml_tensor * dst) const {
@@ -2091,10 +2091,6 @@ void llama_kv_cache_dsv4_raw_context::get_prev_tokens(const llama_ubatch & ubatc
kv_swa->get_prev_tokens(ubatch, n, res);
}
-void llama_kv_cache_dsv4_raw_context::set_input_k_rot(ggml_tensor * dst) const {
- kv_swa->set_input_k_rot(dst);
-}
-
//
// llama_kv_cache_dsv4_comp_context
//
@@ -2146,12 +2142,8 @@ ggml_tensor * llama_kv_cache_dsv4_comp_context::cpy_k(ggml_context * ctx, ggml_t
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
}
-ggml_tensor * llama_kv_cache_dsv4_comp_context::build_input_k_rot(ggml_context * ctx) const {
- return kv->build_input_k_rot(ctx);
-}
-
-void llama_kv_cache_dsv4_comp_context::set_input_k_rot(ggml_tensor * dst) const {
- kv->set_input_k_rot(dst);
+std::vector<ggml_tensor *> llama_kv_cache_dsv4_comp_context::get_k_rot() const {
+ return kv->get_k_rot();
}
//
diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h
index 5a9dcfc55..5adc9fb43 100644
--- a/src/llama-kv-cache-dsv4.h
+++ b/src/llama-kv-cache-dsv4.h
@@ -223,11 +223,10 @@ public:
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const;
- ggml_tensor * build_input_k_rot(ggml_context * ctx) const;
+ std::vector<ggml_tensor *> get_k_rot() const;
void set_input_k_idxs(ggml_tensor * dst) const;
void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
- void set_input_k_rot(ggml_tensor * dst) const;
// token history of the raw (window) cells, for host-side n-gram hashing (deepseek41 Engram)
void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
@@ -270,8 +269,7 @@ public:
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
- ggml_tensor * build_input_k_rot(ggml_context * ctx) const;
- void set_input_k_rot(ggml_tensor * dst) const;
+ std::vector<ggml_tensor *> get_k_rot() const;
private:
llama_kv_cache * kv;
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 00dbf7c94..6d5a57633 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -81,8 +81,8 @@ llama_kv_cache::llama_kv_cache(
const layer_share_cb & share,
const char * name_tag) :
model(model), hparams(hparams), v_trans(v_trans),
- n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_swa(n_swa), swa_type(swa_type),
- ctxs_bufs(no_alloc),
+ n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_swa(n_swa),
+ attn_rot_bufs(no_alloc), swa_type(swa_type), ctxs_bufs(no_alloc),
other(static_cast<llama_kv_cache *>(mem_other)),
v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
v_cells(*v_cells_impl) {
@@ -350,7 +350,6 @@ llama_kv_cache::llama_kv_cache(
LLAMA_LOG_INFO("%s: n_rot_v = %u, n_embd_head_k_all = %d\n", __func__, n_rot_v, n_embd_head_v_all);
// pre-compute the haramard matrices and keep them in host memory
- // TODO: in the future, we can make copies in the backend buffers to avoid host -> device transfers
if (n_rot_k || n_rot_v) {
for (int64_t n = 64; n <= std::max(n_embd_head_k_all, n_embd_head_v_all); n *= 2) {
attn_rot_hadamard[n] = std::vector<float>(n*n);
@@ -370,10 +369,60 @@ llama_kv_cache::llama_kv_cache(
}
}
+ init_rotation(name_tag);
+
const char * LLAMA_KV_CACHE_DEBUG = getenv("LLAMA_KV_CACHE_DEBUG");
debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0;
}
+void llama_kv_cache::init_rotation(const char * name_tag) {
+ if (!(n_rot_k || n_rot_v) || layers.empty()) {
+ return;
+ }
+
+ std::map<ggml_backend_buffer_type_t, ggml_context_ptr> ctx_map;
+ for (auto & layer : layers) {
+ for (bool is_k : { true, false }) {
+ const uint32_t n_rot = is_k ? n_rot_k : n_rot_v;
+ auto * cache = is_k ? layer.k : layer.v;
+ if (!n_rot || !cache) {
+ continue;
+ }
+
+ const auto buft = ggml_backend_buffer_get_type(cache->buffer);
+ auto & ctx = ctx_map[buft];
+ if (!ctx) {
+ ggml_init_params params = { 2*layers.size()*ggml_tensor_overhead(), nullptr, true };
+ ctx.reset(ggml_init(params));
+ if (!ctx) {
+ throw std::runtime_error("failed to create ggml context for rotation matrices");
+ }
+ }
+
+ auto * rot = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, n_rot, n_rot);
+ ggml_format_name(rot, "cache_%s%s_rot_l%d", name_tag, is_k ? "k" : "v", layer.il);
+ (is_k ? layer.k_rot : layer.v_rot) = rot;
+ }
+ }
+
+ for (auto & [buft, ctx] : ctx_map) {
+ auto * buf = attn_rot_bufs.emplace(std::move(ctx), buft, "rotation matrices");
+ ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+ LLAMA_LOG_INFO("%s: %10s rotation buffer size = %8.2f MiB\n", __func__,
+ ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);
+ }
+
+ if (!attn_rot_bufs.no_alloc) {
+ for (const auto & layer : layers) {
+ for (auto * t : { layer.k_rot, layer.v_rot }) {
+ if (t) {
+ ggml_backend_tensor_set(t, attn_rot_hadamard.at(t->ne[0]).data(), 0, ggml_nbytes(t));
+ }
+ }
+ }
+ }
+}
+
void llama_kv_cache::clear(bool data) {
for (uint32_t s = 0; s < n_stream; ++s) {
v_cells[s].reset();
@@ -691,7 +740,13 @@ llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
}
std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
- return ctxs_bufs.memory_breakdown();
+ auto ret = ctxs_bufs.memory_breakdown();
+
+ for (const auto & [buft, size] : attn_rot_bufs.memory_breakdown()) {
+ ret[buft] += size;
+ }
+
+ return ret;
}
llama_memory_context_ptr llama_kv_cache::init_batch(
@@ -1440,27 +1495,19 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama
return v_idxs;
}
-ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const {
- ggml_tensor * res = nullptr;
-
- if (n_rot_k) {
- res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_rot_k, n_rot_k);
- ggml_set_input(res);
- ggml_set_name(res, "attn_inp_k_rot");
+std::vector<ggml_tensor *> llama_kv_cache::get_k_rot() const {
+ std::vector<ggml_tensor *> res(hparams.n_layer_all, nullptr);
+ for (const auto & [il, idx] : map_layer_ids) {
+ res[il] = layers[idx].k_rot;
}
-
return res;
}
-ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const {
- ggml_tensor * res = nullptr;
-
- if (n_rot_v) {
- res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_rot_v, n_rot_v);
- ggml_set_input(res);
- ggml_set_name(res, "attn_inp_v_rot");
+std::vector<ggml_tensor *> llama_kv_cache::get_v_rot() const {
+ std::vector<ggml_tensor *> res(hparams.n_layer_all, nullptr);
+ for (const auto & [il, idx] : map_layer_ids) {
+ res[il] = layers[idx].v_rot;
}
-
return res;
}
@@ -1804,24 +1851,6 @@ void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch
}
}
-void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const {
- GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
-
- const auto n_rot = dst->ne[0];
- GGML_ASSERT(attn_rot_hadamard.count(dst->ne[0]));
-
- memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
-}
-
-void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
- GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
-
- const auto n_rot = dst->ne[0];
- GGML_ASSERT(attn_rot_hadamard.count(dst->ne[0]));
-
- memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
-}
-
bool llama_kv_cache::has_cell_ext() const {
// M-RoPE needs the 2D position, the PLE / Engram n-gram hashes need the token id
return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0;
@@ -1977,9 +2006,6 @@ public:
ggml_tensor * k_shift; // I32 [kv_size*n_stream]
- // note: assumes k_rot^2 == I
- ggml_tensor * k_rot = nullptr;
-
const llama_kv_cache * kv_self;
};
@@ -1989,10 +2015,6 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) {
if (k_shift && k_shift->buffer) {
kv_self->set_input_k_shift(k_shift);
}
-
- if (k_rot && k_rot->buffer) {
- kv_self->set_input_k_rot(k_rot);
- }
}
ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_context * lctx) const {
@@ -2007,8 +2029,6 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
inp->k_shift = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, (int64_t) get_size()*n_stream);
ggml_set_input(inp->k_shift);
- inp->k_rot = build_input_k_rot(ctx);
-
const auto & cparams = lctx->get_cparams();
for (const auto & layer : layers) {
@@ -2037,7 +2057,7 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
ggml_row_size(layer.k->type, n_embd_k_gqa),
ggml_row_size(layer.k->type, n_embd_nope));
- ggml_tensor * cur = build_rope_shift(cparams, ctx, k, inp->k_shift, inp->k_rot, rope_factors, freq_base_l, freq_scale_l, il);
+ ggml_tensor * cur = build_rope_shift(cparams, ctx, k, inp->k_shift, layer.k_rot, rope_factors, freq_base_l, freq_scale_l, il);
ggml_build_forward_expand(gf, cur);
}
@@ -2898,12 +2918,12 @@ ggml_tensor * llama_kv_cache_context::build_input_v_idxs(ggml_context * ctx, con
return kv->build_input_v_idxs(ctx, ubatch);
}
-ggml_tensor * llama_kv_cache_context::build_input_k_rot(ggml_context * ctx) const {
- return kv->build_input_k_rot(ctx);
+std::vector<ggml_tensor *> llama_kv_cache_context::get_k_rot() const {
+ return kv->get_k_rot();
}
-ggml_tensor * llama_kv_cache_context::build_input_v_rot(ggml_context * ctx) const {
- return kv->build_input_v_rot(ctx);
+std::vector<ggml_tensor *> llama_kv_cache_context::get_v_rot() const {
+ return kv->get_v_rot();
}
void llama_kv_cache_context::set_input_k_shift(ggml_tensor * dst) const {
@@ -2926,14 +2946,6 @@ void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama
kv->set_input_pos_bucket(dst, ubatch);
}
-void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const {
- kv->set_input_k_rot(dst);
-}
-
-void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const {
- kv->set_input_v_rot(dst);
-}
-
void llama_kv_cache_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
kv->get_prev_tokens(ubatch, n, res);
}
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 60d2e183c..6dfd0c524 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -225,8 +225,8 @@ public:
ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const;
ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const;
- ggml_tensor * build_input_k_rot(ggml_context * ctx) const;
- ggml_tensor * build_input_v_rot(ggml_context * ctx) const;
+ std::vector<ggml_tensor *> get_k_rot() const;
+ std::vector<ggml_tensor *> get_v_rot() const;
void set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const;
void set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch, const slot_info & sinfo) const;
@@ -236,9 +236,6 @@ public:
void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const;
- void set_input_k_rot(ggml_tensor * dst) const;
- void set_input_v_rot(ggml_tensor * dst) const;
-
// true if llama_kv_cell_ext holds information that has to survive a state save/restore
bool has_cell_ext() const;
@@ -264,6 +261,9 @@ private:
std::vector<ggml_tensor *> k_stream;
std::vector<ggml_tensor *> v_stream;
+
+ ggml_tensor * k_rot = nullptr;
+ ggml_tensor * v_rot = nullptr;
};
bool v_trans = true; // the value tensor is transposed
@@ -286,6 +286,11 @@ private:
// pre-computed hadamard martrices
std::unordered_map<int64_t, std::vector<float>> attn_rot_hadamard;
+ // keep constants separate from mutable KV buffers
+ llama_memory_bufs attn_rot_bufs;
+
+ void init_rotation(const char * name_tag);
+
// env: LLAMA_KV_CACHE_DEBUG
int debug = 0;
@@ -419,8 +424,8 @@ public:
ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const;
ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const;
- ggml_tensor * build_input_k_rot(ggml_context * ctx) const;
- ggml_tensor * build_input_v_rot(ggml_context * ctx) const;
+ std::vector<ggml_tensor *> get_k_rot() const;
+ std::vector<ggml_tensor *> get_v_rot() const;
void set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ubatch) const;
void set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch) const;
@@ -429,9 +434,6 @@ public:
void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const;
- void set_input_k_rot(ggml_tensor * dst) const;
- void set_input_v_rot(ggml_tensor * dst) const;
-
// see llama_kv_cache::get_prev_tokens()
void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp
index 083b32ecf..6ee0d564d 100644
--- a/src/models/deepseek32.cpp
+++ b/src/models/deepseek32.cpp
@@ -243,9 +243,9 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
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_dsa->self_k_rot_lid[il], 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_dsa->self_k_rot_lid[il], indexer_k);
cb(indexer_k, "indexer_k", il);
// store indexer keys to KV cache
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index da51549e4..4ffef08c9 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -605,7 +605,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(
const int64_t nt = cur->ne[1];
GGML_ASSERT(inp_lid.kq_mask);
- GGML_ASSERT(inp_lid.k_rot);
+ GGML_ASSERT(inp_lid.k_rot[il]);
GGML_ASSERT(n_embd_indexer_head >= n_embd_indexer_head_rope);
ggml_tensor * indexer_q = build_lora_mm(layer.indexer_attn_q_b, qr);
@@ -618,7 +618,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(
indexer_q = ggml_rope_set_offset(indexer_q, n_embd_indexer_head_nope);
cb(indexer_q, "lid_q_rope", il);
- indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot);
+ indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot[il]);
cb(indexer_q, "lid_q_rot", il);
ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur);
@@ -730,7 +730,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention(
ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il);
- ggml_tensor * k_rot = inp_attn->self_k_rot;
+ ggml_tensor * k_rot = inp_attn->self_k_rot[il];
if (k_rot) {
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
@@ -786,7 +786,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention(
const auto & inp_hca = inp_dsv4->get_hca();
GGML_ASSERT(inp_hca.kq_mask);
- ggml_tensor * k_rot = inp_attn->self_k_rot;
+ ggml_tensor * k_rot = inp_attn->self_k_rot[il];
if (k_rot) {
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
@@ -839,7 +839,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention(
int il) const {
GGML_ASSERT(hparams.is_swa(il));
- ggml_tensor * k_rot = inp_attn->self_k_rot;
+ ggml_tensor * k_rot = inp_attn->self_k_rot[il];
if (k_rot) {
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
@@ -1002,8 +1002,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
"csa_state_compress",
il);
- if (inp_dsv4->get_csa().k_rot) {
- kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot);
+ if (inp_dsv4->get_csa().k_rot[il]) {
+ kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot[il]);
cb(kv_comp_csa_state, "csa_state_compress_rot", il);
}
@@ -1071,8 +1071,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
"lid_state_compress",
il);
- if (inp_dsv4->get_lid().k_rot) {
- kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot);
+ if (inp_dsv4->get_lid().k_rot[il]) {
+ kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot[il]);
cb(kv_comp_lid_state, "lid_state_compress_rot", il);
}
@@ -1131,8 +1131,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
"hca_state_compress",
il);
- if (inp_dsv4->get_hca().k_rot) {
- kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot);
+ if (inp_dsv4->get_hca().k_rot[il]) {
+ kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot[il]);
cb(kv_comp_hca, "hca_state_compress_rot", il);
}
@@ -1197,7 +1197,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
} else if (ratio == DSV4_CSA_RATIO &&
inp_dsv4->get_csa().kq_mask &&
inp_dsv4->get_lid().kq_mask &&
- inp_dsv4->get_lid().k_rot) {
+ inp_dsv4->get_lid().k_rot[il]) {
out = build_csa_lid_attention(model, inp_dsv4, inp_attn, q, kv, qr, cur, inp_pos, layer.attn_sinks,
1.0f/sqrtf(float(n_embd_head)), il);
} else if (ratio == DSV4_HCA_RATIO &&
diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp
index 97b2cbabe..95c8a36d5 100644
--- a/src/models/deepseek41.cpp
+++ b/src/models/deepseek41.cpp
@@ -632,7 +632,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
GGML_ASSERT(top_k && "DeepSeek-V4.1: top-k source layer produced no top-k");
// a quantized K cache stores the window Hadamard-rotated (as build_raw_attention)
- ggml_tensor * k_rot = inp_attn->self_k_rot;
+ ggml_tensor * k_rot = inp_attn->self_k_rot[il];
if (k_rot) {
q = llama_mul_mat_hadamard(ctx0, q, k_rot);
kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index 194f17ec4..fc1ad69b0 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -980,8 +980,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs();
ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs();
// rotate K/V into the cache's rotated space
- ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa : inp_attn_iswa->self_k_rot;
- ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa : inp_attn_iswa->self_v_rot;
+ ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa[il] : inp_attn_iswa->self_k_rot[il];
+ ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa[il] : inp_attn_iswa->self_v_rot[il];
if (k_rot) {
Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot);
}
@@ -992,11 +992,11 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il));
} else {
// rotate K/V into the cache's rotated space
- if (inp_attn->self_k_rot) {
- Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);
+ if (inp_attn->self_k_rot[il]) {
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot[il]);
}
- if (inp_attn->self_v_rot) {
- Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
+ if (inp_attn->self_v_rot[il]) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot[il]);
}
ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));
ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));
@@ -1249,8 +1249,8 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
kv = ggml_rope_set_offset(kv, n_embd_head_nope);
cb(kv, "kv_injected", il);
- if (inp_attn->self_k_rot_swa) {
- kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
+ if (inp_attn->self_k_rot_swa[il]) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
}
ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
}
@@ -1441,8 +1441,8 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
kv = rope(kv);
cb(kv, "kv_injected", il);
- if (inp_attn->self_k_rot_swa) {
- kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
+ if (inp_attn->self_k_rot_swa[il]) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
}
ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
}
@@ -1686,8 +1686,8 @@ llama_model_dflash::graph_mistral_large4::graph_mistral_large4(const llama_model
ggml_tensor * kv = build_kv(inp_g, il).first;
cb(kv, "kv_injected", il);
- if (inp_attn->self_k_rot_swa) {
- kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
+ if (inp_attn->self_k_rot_swa[il]) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
}
ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
}
diff --git a/src/models/dots3note.cpp b/src/models/dots3note.cpp
index 237a92371..5269c8ff8 100644
--- a/src/models/dots3note.cpp
+++ b/src/models/dots3note.cpp
@@ -230,9 +230,9 @@ 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->get_dsa()->self_k_rot_lid, indexer_q);
+ indexer_q = ggml_mul_mat(ctx0, inp_attn->get_dsa()->self_k_rot_lid[il], indexer_q);
cb(indexer_q, "indexer_q", il);
- indexer_k = ggml_mul_mat(ctx0, inp_attn->get_dsa()->self_k_rot_lid, indexer_k);
+ indexer_k = ggml_mul_mat(ctx0, inp_attn->get_dsa()->self_k_rot_lid[il], indexer_k);
cb(indexer_k, "indexer_k", il);
// store indexer keys to KV cache
diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp
index a9d62c076..9822ef5ee 100644
--- a/src/models/glm-dsa.cpp
+++ b/src/models/glm-dsa.cpp
@@ -277,9 +277,9 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
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_dsa->self_k_rot_lid[il], 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_dsa->self_k_rot_lid[il], indexer_k);
cb(indexer_k, "indexer_k", il);
// store indexer keys to KV cache
diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp
index 73d145dae..1b112b743 100644
--- a/src/models/minimax-m3.cpp
+++ b/src/models/minimax-m3.cpp
@@ -358,12 +358,12 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, ik, inp_attn->get_k_idxs_idx(), il));
ggml_tensor * ik_kv = mctx_idx->get_k(ctx0, il);
- if (inp_attn->self_k_rot) {
- Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);
- Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);
+ if (inp_attn->self_k_rot[il]) {
+ Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot[il]);
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot[il]);
}
- if (inp_attn->self_v_rot) {
- Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
+ if (inp_attn->self_v_rot[il]) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot[il]);
}
// Main branch: store K/V, take cache views
@@ -526,8 +526,8 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
cur = ggml_concat(ctx0, cur, outs[st], 1);
}
}
- if (inp_attn->self_v_rot) {
- cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot);
+ if (inp_attn->self_v_rot[il]) {
+ cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot[il]);
}
cb(cur, "kqv_out", il);
if (model.layers[il].wo) {
diff --git a/src/models/mistral-large4.cpp b/src/models/mistral-large4.cpp
index 9cfdf825b..dc7811d4c 100644
--- a/src/models/mistral-large4.cpp
+++ b/src/models/mistral-large4.cpp
@@ -262,9 +262,9 @@ llama_model_mistral_large4::graph::graph(const llama_model & model, const llm_gr
}
// 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_dsa->self_k_rot_lid[il], 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_dsa->self_k_rot_lid[il], indexer_k);
cb(indexer_k, "indexer_k", il);
// store indexer keys to KV cache
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index c3e2b010d..5de67aa50 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -968,13 +968,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
int il) {
// rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer
// has already scored with its own query in build_qsa_sel, so the selection is unaffected.
- if (inp->self_k_rot) {
- q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);
- k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);
+ if (inp->self_k_rot[il]) {
+ q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot[il]);
+ k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot[il]);
}
- if (inp->self_v_rot) {
- v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
+ if (inp->self_v_rot[il]) {
+ v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot[il]);
}
// these nodes are added to the graph together so that they are not reordered
@@ -1009,8 +1009,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
cb(cur, "kqv_out", il);
// the rotation is its own inverse, so undo it on the value side of the output
- if (inp->self_v_rot) {
- cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
+ if (inp->self_v_rot[il]) {
+ cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot[il]);
}
return cur;