From 0138937e3f987f70e90e97d022c73ecb9bec037c Mon Sep 17 00:00:00 2001 From: Niklas Wenzel Date: Sat, 31 Oct 2026 23:30:49 +0100 Subject: [PATCH] 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 --- src/llama-hparams.cpp | 15 +-------------- src/llama-model-saver.cpp | 10 ++++++---- src/llama-model.cpp | 14 +++++--------- src/models/gemma3n.cpp | 5 +++++ src/models/gemma4.cpp | 6 ++++++ 5 files changed, 23 insertions(+), 27 deletions(-) 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);