Commit 0138937e3 for llama.cpp
commit 0138937e3f987f70e90e97d022c73ecb9bec037c
Author: Niklas Wenzel <dev@nikwen.de>
Date: Sat Oct 31 23:30:49 2026 +0100
llama : fix rotation of attention tensors in layers with shared KV (#31406)
* llama : fix rotation of layers that reuse another layer's KV cache
* llama : use KV source rotation only for attn_* tensors
diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp
index 9b4244567..f05b5d6bb 100644
--- a/src/llama-hparams.cpp
+++ b/src/llama-hparams.cpp
@@ -319,20 +319,7 @@ uint32_t llama_hparams::n_embd_head_v_mla() const {
}
bool llama_hparams::has_kv(uint32_t il) const {
- if (kv_source(il) != il) {
- return false;
- }
-
- if (n_layer_kv_from_start >= 0) {
- if (il < (uint32_t) n_layer_kv_from_start) {
- return true;
- }
-
- return false;
- }
-
- // by default, all layers have kv
- return true;
+ return kv_source(il) == il;
}
uint32_t llama_hparams::kv_source(uint32_t il) const {
diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp
index 5c2d7a822..a08e173e9 100644
--- a/src/llama-model-saver.cpp
+++ b/src/llama-model-saver.cpp
@@ -307,10 +307,12 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail);
add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true);
- for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
- if (hparams.kv_source(il) != il) {
- add_kv(LLM_KV_ATTENTION_KV_SOURCE_MAP, hparams.kv_source_impl, true);
- break;
+ if (hparams.n_layer_kv_from_start < 0) {
+ for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+ if (hparams.kv_source(il) != il) {
+ add_kv(LLM_KV_ATTENTION_KV_SOURCE_MAP, hparams.kv_source_impl, true);
+ break;
+ }
}
}
add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true);
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 83aed0144..c4a5b64fd 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -486,7 +486,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
GGML_ASSERT(length_prefix != std::string::npos);
prefix = tensor_name.substr(0, length_prefix + 1);
il = std::stoull(tensor_name.substr(4, length_prefix));
- rotation = get_il_eff(il) % ud->n_devices;
+ const bool is_attn = tensor_name.compare(prefix.size(), 5, "attn_") == 0;
+ // attention heads must be on the same device as the KV of the source layer
+ rotation = get_il_eff(is_attn ? hparams.kv_source(il) : il) % ud->n_devices;
} else if (tensor_name.substr(0, 6) == "cache_") {
const size_t layer_index_start = tensor_name.find("_l", 6);
GGML_ASSERT(layer_index_start != std::string::npos);
@@ -3280,14 +3282,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
llama_kv_cache::layer_share_cb share = nullptr;
if (arch == LLM_ARCH_GEMMA3N || arch == LLM_ARCH_GEMMA4) {
- reuse = [&](uint32_t il) {
- GGML_ASSERT(hparams.n_layer_kv_from_start >= 2);
-
- if (il >= (uint32_t)hparams.n_layer_kv_from_start) {
- return hparams.n_layer_kv_from_start - (hparams.is_swa(il) ? 2 : 1);
- }
-
- return -1;
+ reuse = [&](int32_t il) -> int32_t {
+ return hparams.has_kv(il) ? -1 : (int32_t) hparams.kv_source(il);
};
}
diff --git a/src/models/gemma3n.cpp b/src/models/gemma3n.cpp
index 7848f9f59..7f12b1747 100644
--- a/src/models/gemma3n.cpp
+++ b/src/models/gemma3n.cpp
@@ -7,6 +7,11 @@ void llama_model_gemma3n::load_arch_hparams(llama_model_loader & ml) {
hparams.n_layer_kv_from_start = 20;
hparams.f_attention_scale = 1.0f;
+ // shared layers read the KV of the last layer with own KV and the same SWA type
+ for (uint32_t il = hparams.n_layer_kv_from_start; il < hparams.n_layer_all; il++) {
+ hparams.kv_source_impl[il] = hparams.n_layer_kv_from_start - (hparams.is_swa(il) ? 2 : 1);
+ }
+
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp
index 42d6a7d47..d0b2b3a5b 100644
--- a/src/models/gemma4.cpp
+++ b/src/models/gemma4.cpp
@@ -11,6 +11,12 @@ void llama_model_gemma4::load_arch_hparams(llama_model_loader & ml) {
hparams.n_layer_kv_from_start = hparams.n_layer_all - (int32_t)n_kv_shared_layers;
hparams.f_attention_scale = 1.0f; // Gemma4 uses self.scaling = 1.0 (no pre-attn scaling)
+ // shared layers read the KV of the last layer with own KV and the same SWA type
+ GGML_ASSERT(hparams.n_layer_kv_from_start >= 2);
+ for (uint32_t il = hparams.n_layer_kv_from_start; il < hparams.n_layer_all; il++) {
+ hparams.kv_source_impl[il] = hparams.n_layer_kv_from_start - (hparams.is_swa(il) ? 2 : 1);
+ }
+
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all, false);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);