From 571d0d540df04f25298d0e159e520d9fc62ed121 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Sat, 18 Jul 2026 15:02:18 +0200 Subject: [PATCH] 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 * clearer format * remove trailing whitespace --------- Co-authored-by: Georgi Gerganov --- src/models/dflash.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index a7b4f4435a..427eed4594 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::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)); }