From: Ruixiang Wang <redacted>
Date: Sat, 18 Jul 2026 13:02:18 +0000 (+0200)
Subject: model: rotate injected K/V cache for DFlash (#25823)
X-Git-Tag: upstream/0.0.10438~370
X-Git-Url: https://git.djapps.eu/?a=commitdiff_plain;h=571d0d540df04f25298d0e159e520d9fc62ed121;p=pkg%2Fggml%2Fsources%2Fllama.cpp

model: rotate injected K/V cache for DFlash (#25823)

* dflash: rotate injected K/V cache when using K/V quantization

* Update src/models/dflash.cpp

Co-authored-by: Georgi Gerganov <redacted>

* clearer format

* remove trailing whitespace

---------

Co-authored-by: Georgi Gerganov <redacted>
---

diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index a7b4f4435..427eed459 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -1,5 +1,6 @@
 #include "models.h"
 
+#include "llama-impl.h"
 #include "llama-kv-cache.h"
 #include "llama-kv-cache-iswa.h"
 
@@ -164,9 +165,25 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
                 const auto  * kv     = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base();
                 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;
+                if (k_rot) {
+                    Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot);
+                }
+                if (v_rot) {
+                    Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot);
+                }
                 ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il));
                 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_v_rot) {
+                    Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
+                }
                 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));
             }