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 a7b4f4435a8892d519c3351f267d91d0f178b279..427eed4594e359b6fd68a7e24160cbd12dbc5b4a 100644 (file)
--- 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));
             }