From fb5997274ec0d0b4d8fe067eb55f87c0ad405478 Mon Sep 17 00:00:00 2001 From: Anton I Date: Fri, 30 Oct 2026 04:15:06 -0400 Subject: [PATCH] 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 --- src/llama-graph.cpp | 133 ++++++++-------------------------- src/llama-graph.h | 22 +++--- src/llama-kv-cache-dsv4.cpp | 16 +--- src/llama-kv-cache-dsv4.h | 6 +- src/llama-kv-cache.cpp | 132 ++++++++++++++++++--------------- src/llama-kv-cache.h | 22 +++--- src/models/deepseek32.cpp | 4 +- src/models/deepseek4.cpp | 24 +++--- src/models/deepseek41.cpp | 2 +- src/models/dflash.cpp | 24 +++--- src/models/dots3note.cpp | 4 +- src/models/glm-dsa.cpp | 4 +- src/models/minimax-m3.cpp | 14 ++-- src/models/mistral-large4.cpp | 4 +- src/models/qwen4exp.cpp | 14 ++-- 15 files changed, 179 insertions(+), 246 deletions(-) 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 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 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(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(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 self_k_rot; + std::vector 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 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 self_k_rot; + std::vector self_v_rot; - ggml_tensor * self_k_rot_swa = nullptr; - ggml_tensor * self_v_rot_swa = nullptr; + std::vector self_k_rot_swa; + std::vector 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 self_k_rot; + std::vector 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 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 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 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 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 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 & 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 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(mem_other)), v_cells_impl(other ? other->v_cells_impl : std::make_shared()), 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(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 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 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 llama_kv_cache::get_k_rot() const { + std::vector 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 llama_kv_cache::get_v_rot() const { + std::vector 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 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 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 & 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 get_k_rot() const; + std::vector 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 k_stream; std::vector 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> 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 get_k_rot() const; + std::vector 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 & 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::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::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;