From 7644fd9de3b3ab180ccdb899046a796cbea9e43a Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 17 Jul 2026 13:18:22 +0000 Subject: [PATCH 1/4] dflash: rotate injected K/V cache when using K/V quantization --- src/models/dflash.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index a7b4f4435a8..08d75e52fb5 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,17 @@ 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 (quantized KV only) + 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 (quantized KV only) + 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)); } From ef8398b6483eec6a1a282c75a1d8b0549ae0e4c2 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 17 Jul 2026 15:26:26 +0200 Subject: [PATCH 2/4] Update src/models/dflash.cpp Co-authored-by: Georgi Gerganov --- src/models/dflash.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 08d75e52fb5..15dd5e1725f 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -165,11 +165,15 @@ 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 (quantized KV only) + // 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); } + 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 { From 571393fb032b4fcfbace34a2507c089ae5fc7d31 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 17 Jul 2026 13:29:17 +0000 Subject: [PATCH 3/4] clearer format --- src/models/dflash.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 15dd5e1725f..e4b585c548a 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -177,9 +177,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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 (quantized KV only) - 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); } + // 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)); } From c1d126d0535010b0aa2a05938fd51664881ddfdf Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 17 Jul 2026 13:33:40 +0000 Subject: [PATCH 4/4] remove trailing whitespace --- src/models/dflash.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index e4b585c548a..427eed4594e 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -178,11 +178,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) { + 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); + 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));