From faa6e99068c41ebfe6f99381ef930fb0690c2e25 Mon Sep 17 00:00:00 2001 From: Brian <288398250+bri-prism@users.noreply.github.com> Date: Fri, 10 Jul 2026 07:33:07 -0700 Subject: [PATCH 1/4] kv-cache: support mean-centering on hybrid-memory models Hybrid (recurrent + attention) models keep a standard llama_kv_cache for their attention sublayers, but llama_init_from_model only accepted path_kv_mean_center when the whole memory module was that cache, so any hybrid model failed to load a bias file. Route the load through get_mem_attn() for hybrid memory; bias tensors are matched by model layer id and layers absent from the attention cache are skipped, which the loader already handles. --- src/llama-context.cpp | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a867c24bda67..9f3f797b389e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -8,6 +8,7 @@ #include "llama-io.h" #include "llama-kv-cache.h" #include "llama-memory.h" +#include "llama-memory-hybrid.h" #include "llama-mmap.h" #include "llama-model.h" #include "llama-ext.h" @@ -323,9 +324,17 @@ llama_context::llama_context( if (params.path_kv_mean_center != nullptr) { auto * kv = dynamic_cast(memory.get()); + if (!kv) { + // hybrid (recurrent + attention) models keep a standard KV cache for their + // attention sublayers; center that one. bias tensors are matched by model + // layer id, so layers absent from the attention cache are simply skipped. + if (auto * hyb = dynamic_cast(memory.get())) { + kv = hyb->get_mem_attn(); + } + } if (!kv) { throw std::runtime_error("path_kv_mean_center is only supported for the standard KV cache " - "(not recurrent, hybrid or MLA/DSA memory types)"); + "(not recurrent or MLA/DSA memory types)"); } if (!kv->load_kv_mean_center(params.path_kv_mean_center)) { throw std::runtime_error("failed to load K-cache mean-centering bias file"); From 8039abc39cda43adaecf1ff8e495dc35b534efbd Mon Sep 17 00:00:00 2001 From: Brian <288398250+bri-prism@users.noreply.github.com> Date: Fri, 10 Jul 2026 07:35:24 -0700 Subject: [PATCH 2/4] kv-mean-center: document that centering must not be combined with the K-cache rotation Measured end to end, the pre-rotation bias applied post-rotation is worse than either feature alone; strengthen the README note into a warning with the measured numbers. --- tools/kv-mean-center/README.md | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tools/kv-mean-center/README.md b/tools/kv-mean-center/README.md index aca105b94f73..774ea2e0e90b 100644 --- a/tools/kv-mean-center/README.md +++ b/tools/kv-mean-center/README.md @@ -41,11 +41,16 @@ at inference time with: `--kv-mean-center` requires `--cache-type-k q4_0`; loading fails with a clear error otherwise (centering is currently only implemented/validated for `GGML_TYPE_Q4_0`). -Note: if the K cache also has this fork's optional Hadamard rotation feature active (automatic for -`GGML_TYPE_Q4_0` caches whose head dimension is a multiple of 64, unless `LLAMA_ATTN_ROT_DISABLE=1`), -the bias measured here is taken in the pre-rotation basis, since that is where the model's `Kcur` -tensor naturally sits before the rotation is applied. The centering subtraction still happens in -whatever basis `cpy_k()` sees (post-rotation, if active), which remains exactly safe (see -docs/kv-mean-center.md), but a mismatched basis means the calibrated bias is a less accurate -estimate of that channel's true post-rotation mean. Recalibrating directly against the -post-rotation representation is a natural follow-up. +**Warning: do not combine centering with the Hadamard K-cache rotation.** If the K cache also has +this fork's optional Hadamard rotation feature active (automatic for `GGML_TYPE_Q4_0` caches whose +head dimension is a multiple of 64, unless `LLAMA_ATTN_ROT_DISABLE=1`), the bias measured here is +taken in the pre-rotation basis, since that is where the model's `Kcur` tensor naturally sits +before the rotation is applied, while the centering subtraction happens in whatever basis +`cpy_k()` sees (post-rotation, if active). That stays exactly safe for attention logits (see +docs/kv-mean-center.md), but measured end-to-end it is worse than either feature alone: after +rotation the true per-channel means are near zero, so subtracting the pre-rotation bias injects a +spurious offset instead of removing one. On a hybrid-attention model with a Q4_0 K cache the +logit KL divergence vs an F16 cache was 0.00144 with rotation alone, 0.00149 with centering alone +(rotation disabled), and 0.0020-0.0021 with both. Use centering only where the rotation is not +active; recalibrating directly against the post-rotation representation is the natural follow-up +that would let the two compose. From 8c013af8df6dec8fa3ce9a5c31ea4bdf02fe2c0c Mon Sep 17 00:00:00 2001 From: Brian <288398250+bri-prism@users.noreply.github.com> Date: Fri, 10 Jul 2026 08:02:00 -0700 Subject: [PATCH 3/4] kv-mean-center: record the calibration basis and reject a rotation mismatch at load The bias lives in the basis the calibration run's K cache used: the collector taps the exact tensor cpy_k() writes, after any Hadamard rotation, and the rotation is gated on the K cache being quantized. A calibration run with the default F16 cache therefore measures the unrotated basis, and applying that bias to a rotated cache measurably degrades quality instead of improving it (KLD vs F16 cache 0.00144 rotation alone vs 0.0020 with the mismatched combination), while a bias calibrated with -ctk q4_0 composes (0.00111, the best of all measured configurations). The tool now detects whether the rotation was active from the captured tensor's ancestry, records it in the output file as kv_mean_center.k_rot, and load_kv_mean_center() rejects a bias whose basis does not match the inference-time rotation state (files predating the flag load with a warning). Docs updated with the calibrate-with-matching-settings rule and the measured numbers. --- common/kv-mean-center.cpp | 4 ++- common/kv-mean-center.h | 6 +++- docs/kv-mean-center.md | 19 ++++++----- src/llama-kv-cache.cpp | 26 ++++++++++++++ tests/test-kv-mean-center.cpp | 12 ++++++- tools/kv-mean-center/README.md | 45 ++++++++++++++++++------- tools/kv-mean-center/kv-mean-center.cpp | 36 ++++++++++++++++++-- 7 files changed, 121 insertions(+), 27 deletions(-) diff --git a/common/kv-mean-center.cpp b/common/kv-mean-center.cpp index cb2365264c3e..cc214f688aca 100644 --- a/common/kv-mean-center.cpp +++ b/common/kv-mean-center.cpp @@ -9,7 +9,8 @@ bool common_kv_mean_center_write( const std::string & fname, - const std::vector & layers) { + const std::vector & layers, + bool k_rot) { size_t n_with_bias = 0; for (const auto & layer : layers) { if (!layer.bias.empty()) { @@ -46,6 +47,7 @@ bool common_kv_mean_center_write( } gguf_set_val_str(ctx_gguf, "general.type", "kv-mean-center"); + gguf_set_val_bool(ctx_gguf, "kv_mean_center.k_rot", k_rot); for (const auto & layer : layers) { if (layer.bias.empty()) { diff --git a/common/kv-mean-center.h b/common/kv-mean-center.h index 74fbf2d99839..9da7ffd2e38b 100644 --- a/common/kv-mean-center.h +++ b/common/kv-mean-center.h @@ -20,7 +20,11 @@ struct common_kv_mean_center_layer { }; // write a K-cache mean-centering bias file in GGUF format. +// `k_rot` records whether the bias was measured with the Hadamard K-cache rotation active +// (stored as the "kv_mean_center.k_rot" KV); the loader refuses a bias whose basis does not +// match the inference-time rotation state. // returns false (and logs an error) on failure. bool common_kv_mean_center_write( const std::string & fname, - const std::vector & layers); + const std::vector & layers, + bool k_rot = false); diff --git a/docs/kv-mean-center.md b/docs/kv-mean-center.md index 44ee4ce3bd8a..ee70e7a6ce52 100644 --- a/docs/kv-mean-center.md +++ b/docs/kv-mean-center.md @@ -84,15 +84,18 @@ scheduler eval-callback mechanism `llama-imatrix` uses to capture activations). - Only `GGML_TYPE_Q4_0` is supported; other K cache types are rejected. Generalizing the mechanism to other quantization types is future work. -- Only the plain (non-recurrent, non-hybrid, non-MLA/DSA) KV cache is supported. +- The plain KV cache and the attention sub-cache of hybrid (recurrent + attention) models are + supported. Recurrent-only and MLA/DSA memory types are not. - The calibration hook (`k_cache_in`) is currently only wired into the standard dense/GQA attention path (`llm_graph_context::build_attn(llm_graph_input_attn_kv *, ...)`), which covers the large majority of architectures. MLA and other specialized attention variants are not covered yet. -- If this fork's optional Hadamard K/Q rotation feature is also active (automatic for `Q4_0` - caches whose head dimension is a multiple of 64, unless `LLAMA_ATTN_ROT_DISABLE=1`), the bias is - calibrated in the pre-rotation basis while it is applied in whatever basis `cpy_k()` sees - (post-rotation, if active). This remains exactly safe (the invariance argument above is - basis-independent), but the calibrated bias is a less accurate estimate of that channel's true - post-rotation mean in that configuration. Calibrating directly against the post-rotation - representation is a natural follow-up. +- The bias lives in the basis the calibration run's K cache used. With this fork's optional + Hadamard K rotation (automatic for quantized K caches whose head dimension is a multiple of 64, + unless `LLAMA_ATTN_ROT_DISABLE=1`), that means: calibrate with the same `--cache-type-k` and + rotation settings you serve with, e.g. `-ctk q4_0` for the common case. Applying a bias in the + wrong basis remains exactly safe for attention logits (the invariance argument above is + basis-independent) but measurably degrades quantization quality instead of improving it, so the + calibration tool records its basis in the file (`kv_mean_center.k_rot`) and the loader rejects + a mismatch. In the matching basis the two features compose; see tools/kv-mean-center/README.md + for measured numbers. diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 09b0689e7b6c..3fa5fc2d6638 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1411,6 +1411,32 @@ bool llama_kv_cache::load_kv_mean_center(const char * path, bool require_q4_0) { return false; } + // the bias is only valid in the basis it was measured in: a bias calibrated with the + // Hadamard K-cache rotation active must be applied with the rotation active, and vice + // versa. a mismatched basis measurably degrades quantization quality instead of + // improving it (see tools/kv-mean-center/README.md), so refuse it outright. + { + const int64_t idx_k_rot = gguf_find_key(ctx_gguf, "kv_mean_center.k_rot"); + if (idx_k_rot >= 0) { + const bool bias_k_rot = gguf_get_val_bool(ctx_gguf, idx_k_rot); + if (bias_k_rot != attn_rot_k) { + LLAMA_LOG_ERROR("%s: bias file %s was calibrated with the K-cache rotation %s, but it is %s " + "for this context - recalibrate with matching cache settings " + "(or set LLAMA_ATTN_ROT_DISABLE=1 consistently in both)\n", + __func__, path, + bias_k_rot ? "active" : "inactive", + attn_rot_k ? "active" : "inactive"); + gguf_free(ctx_gguf); + ggml_free(ctx_data); + return false; + } + } else { + LLAMA_LOG_WARN("%s: bias file %s does not record its calibration basis (kv_mean_center.k_rot); " + "K-cache rotation is %s for this context - a basis mismatch degrades quality\n", + __func__, path, attn_rot_k ? "active" : "inactive"); + } + } + k_bar.resize(layers.size(), nullptr); // one ggml context (+ backend buffer) per unique buffer type, so each bias tensor ends up diff --git a/tests/test-kv-mean-center.cpp b/tests/test-kv-mean-center.cpp index f64f3f8f6ae2..5f7a46f7c15d 100644 --- a/tests/test-kv-mean-center.cpp +++ b/tests/test-kv-mean-center.cpp @@ -229,8 +229,10 @@ static void test_q4_0_gate() { layers.push_back(std::move(layer)); } + // this synthetic model has a 64-wide K head, so a Q4_0 K cache activates the Hadamard + // rotation; the bias must be marked as measured in the rotated basis to be accepted there const std::string tmp_path = "test-kv-mean-center-gate.gguf"; - TEST_ASSERT(common_kv_mean_center_write(tmp_path, layers)); + TEST_ASSERT(common_kv_mean_center_write(tmp_path, layers, /*k_rot=*/true)); // F16 K cache + --kv-mean-center must be rejected outright (llama_init_from_model returns // nullptr), matching this codebase's convention for other cache-type-gated mismatches (e.g. @@ -238,6 +240,14 @@ static void test_q4_0_gate() { llama_context_ptr ctx_bad = build_context(model.get(), GGML_TYPE_F16, tmp_path.c_str()); TEST_ASSERT(ctx_bad == nullptr); + // a bias measured in the unrotated basis must be rejected by a rotated (Q4_0) K cache: + // a basis mismatch degrades quantization quality instead of improving it + const std::string tmp_path_unrot = "test-kv-mean-center-gate-unrot.gguf"; + TEST_ASSERT(common_kv_mean_center_write(tmp_path_unrot, layers, /*k_rot=*/false)); + llama_context_ptr ctx_mismatch = build_context(model.get(), GGML_TYPE_Q4_0, tmp_path_unrot.c_str()); + TEST_ASSERT(ctx_mismatch == nullptr); + remove(tmp_path_unrot.c_str()); + // Q4_0 K cache + --kv-mean-center must succeed end-to-end, through the real public entry // point (not the require_q4_0=false test seam used in test_softmax_invariance) llama_context_ptr ctx_good = build_context(model.get(), GGML_TYPE_Q4_0, tmp_path.c_str()); diff --git a/tools/kv-mean-center/README.md b/tools/kv-mean-center/README.md index 774ea2e0e90b..439e77792117 100644 --- a/tools/kv-mean-center/README.md +++ b/tools/kv-mean-center/README.md @@ -41,16 +41,35 @@ at inference time with: `--kv-mean-center` requires `--cache-type-k q4_0`; loading fails with a clear error otherwise (centering is currently only implemented/validated for `GGML_TYPE_Q4_0`). -**Warning: do not combine centering with the Hadamard K-cache rotation.** If the K cache also has -this fork's optional Hadamard rotation feature active (automatic for `GGML_TYPE_Q4_0` caches whose -head dimension is a multiple of 64, unless `LLAMA_ATTN_ROT_DISABLE=1`), the bias measured here is -taken in the pre-rotation basis, since that is where the model's `Kcur` tensor naturally sits -before the rotation is applied, while the centering subtraction happens in whatever basis -`cpy_k()` sees (post-rotation, if active). That stays exactly safe for attention logits (see -docs/kv-mean-center.md), but measured end-to-end it is worse than either feature alone: after -rotation the true per-channel means are near zero, so subtracting the pre-rotation bias injects a -spurious offset instead of removing one. On a hybrid-attention model with a Q4_0 K cache the -logit KL divergence vs an F16 cache was 0.00144 with rotation alone, 0.00149 with centering alone -(rotation disabled), and 0.0020-0.0021 with both. Use centering only where the rotation is not -active; recalibrating directly against the post-rotation representation is the natural follow-up -that would let the two compose. +## Interaction with the Hadamard K-cache rotation: calibrate with matching cache settings + +This fork's optional Hadamard rotation is active for quantized K caches whose head dimension is a +multiple of 64 (unless `LLAMA_ATTN_ROT_DISABLE=1`). The bias measured by this tool lives in +whatever basis the calibration run's K cache used: the collector taps the exact tensor `cpy_k()` +writes, after any rotation. Since the rotation is gated on the K cache being quantized, a +calibration run with the default F16 cache measures the unrotated basis, and that bias must not be +applied to a rotated (quantized) cache: measured end to end the mismatch is worse than either +feature alone, because the subtracted vector no longer matches the channel means of the basis +being quantized. + +**Rule: calibrate with the same `--cache-type-k` (and `LLAMA_ATTN_ROT_DISABLE`, if any) that you +will serve with.** For the common case that is: + +``` +./llama-kv-mean-center -m model.gguf -f calib.txt -o kv-mean-center.gguf -ctk q4_0 [-fa on] +``` + +The tool records the calibration basis in the output file (`kv_mean_center.k_rot`) and the loader +refuses a bias whose basis does not match the inference-time rotation state, so a mismatch fails +fast instead of silently degrading quality. Measured on a hybrid-attention model, Q4_0 K cache, +logit KL divergence vs an F16-cache baseline over 12x512-token held-out chunks: + +| configuration | mean KLD | +| --- | --- | +| rotation alone (uncentered) | 0.00144 | +| centering alone (rotation disabled, matched basis) | 0.00149 | +| rotation + pre-rotation bias (the mismatch, now rejected) | 0.0020-0.0021 | +| rotation + rotated-basis bias (calibrated with `-ctk q4_0`) | **0.00111** | + +The two features compose once the basis matches, and the composed configuration is the best of +the four. diff --git a/tools/kv-mean-center/kv-mean-center.cpp b/tools/kv-mean-center/kv-mean-center.cpp index 7e08f1c409d8..1ec4244230c1 100644 --- a/tools/kv-mean-center/kv-mean-center.cpp +++ b/tools/kv-mean-center/kv-mean-center.cpp @@ -29,6 +29,8 @@ class kv_mean_collector { std::vector finalize() const; + bool saw_k_rot() const { return m_saw_k_rot; } + private: std::mutex m_mutex; std::vector m_host_buf; @@ -36,8 +38,32 @@ class kv_mean_collector { std::unordered_map> m_sum; // il -> [n_embd_head * n_head] std::unordered_map m_count; // il -> total tokens seen + + // whether the captured K tensors were produced with the Hadamard K-cache rotation + // active, detected from the graph itself: the rotation is a mul_mat against the + // "attn_inp_k_rot" input, so it shows up in the ancestry of "k_cache_in-". + // recorded in the output file so the loader can reject a basis mismatch. + bool m_saw_k_rot = false; }; +// look for the "attn_inp_k_rot" input within a few links of the captured tensor. matched as a +// substring: views append a suffix ("attn_inp_k_rot (reshaped)") and the backend scheduler +// decorates split inputs with a backend prefix and split index ("MTL0#attn_inp_k_rot#0") +static bool tensor_has_k_rot_ancestor(const struct ggml_tensor * t, int depth = 8) { + if (t == nullptr || depth < 0) { + return false; + } + if (strstr(t->name, "attn_inp_k_rot") != nullptr) { + return true; + } + for (int i = 0; i < GGML_MAX_SRC; ++i) { + if (t->src[i] && tensor_has_k_rot_ancestor(t->src[i], depth - 1)) { + return true; + } + } + return false; +} + // "k_cache_in-" -> il, as formatted by llm_graph_context::cb() (ggml_format_name("%s-%d", ...)) static bool parse_k_cache_in_layer(const char * name, int32_t & il) { static const char prefix[] = "k_cache_in-"; @@ -70,6 +96,10 @@ bool kv_mean_collector::collect(struct ggml_tensor * t, bool ask) { std::lock_guard lock(m_mutex); + if (!m_saw_k_rot && tensor_has_k_rot_ancestor(t)) { + m_saw_k_rot = true; + } + // cpy_k() can be fed an F32, F16, or BF16 Kcur depending on backend/compute settings. GGML_ASSERT(t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_F16 || t->type == GGML_TYPE_BF16); GGML_ASSERT(ggml_is_contiguous(t)); @@ -262,12 +292,12 @@ int main(int argc, char ** argv) { return 1; } - if (!common_kv_mean_center_write(params.out_file, layers)) { + if (!common_kv_mean_center_write(params.out_file, layers, g_collector.saw_k_rot())) { return 1; } - LOG_INF("%s: wrote K-cache mean-centering bias for %zu layer(s) to %s\n", - __func__, layers.size(), params.out_file.c_str()); + LOG_INF("%s: wrote K-cache mean-centering bias for %zu layer(s) to %s (measured with K rotation %s)\n", + __func__, layers.size(), params.out_file.c_str(), g_collector.saw_k_rot() ? "active" : "inactive"); llama_backend_free(); From dffd0036a33e192b3666e2fc3ef361de25444b74 Mon Sep 17 00:00:00 2001 From: Brian <288398250+bri-prism@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:32:54 -0700 Subject: [PATCH 4/4] kv-mean-center: address review feedback Cover SWA memory layouts: the bias load now also routes through llama_kv_cache_iswa and llama_memory_hybrid_iswa (base and SWA sub-caches), so hybrid models with sliding-window attention are no longer rejected. Validate the kv_mean_center.k_rot metadata type before reading it, so a malformed bias file produces a loader error instead of an assertion abort. Run the Q4_0 cache-type validation before the basis check, so an unsupported cache type keeps its actionable error even when the file's basis would also mismatch, and hoist it out of the tensor loop. The gate test now uses a basis-matching file for the F16-cache case so it exercises the cache-type gate specifically. --- docs/kv-mean-center.md | 5 +++-- src/llama-context.cpp | 37 +++++++++++++++++++++++------------ src/llama-kv-cache.cpp | 28 +++++++++++++++++++------- tests/test-kv-mean-center.cpp | 10 ++++++---- 4 files changed, 54 insertions(+), 26 deletions(-) diff --git a/docs/kv-mean-center.md b/docs/kv-mean-center.md index ee70e7a6ce52..548f5e4a7598 100644 --- a/docs/kv-mean-center.md +++ b/docs/kv-mean-center.md @@ -84,8 +84,9 @@ scheduler eval-callback mechanism `llama-imatrix` uses to capture activations). - Only `GGML_TYPE_Q4_0` is supported; other K cache types are rejected. Generalizing the mechanism to other quantization types is future work. -- The plain KV cache and the attention sub-cache of hybrid (recurrent + attention) models are - supported. Recurrent-only and MLA/DSA memory types are not. +- Every standard-attention KV cache layout is supported: the plain cache, the base/SWA pair of + sliding-window models, and the attention sub-cache of hybrid (recurrent + attention) models, + with or without SWA. Recurrent-only and MLA/DSA memory types are not. - The calibration hook (`k_cache_in`) is currently only wired into the standard dense/GQA attention path (`llm_graph_context::build_attn(llm_graph_input_attn_kv *, ...)`), which covers the large majority of architectures. MLA and other specialized attention variants diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9f3f797b389e..6ff36f288d14 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -8,7 +8,9 @@ #include "llama-io.h" #include "llama-kv-cache.h" #include "llama-memory.h" +#include "llama-kv-cache-iswa.h" #include "llama-memory-hybrid.h" +#include "llama-memory-hybrid-iswa.h" #include "llama-mmap.h" #include "llama-model.h" #include "llama-ext.h" @@ -323,21 +325,30 @@ llama_context::llama_context( memory.reset(model.create_memory(params_mem, cparams)); if (params.path_kv_mean_center != nullptr) { - auto * kv = dynamic_cast(memory.get()); - if (!kv) { - // hybrid (recurrent + attention) models keep a standard KV cache for their - // attention sublayers; center that one. bias tensors are matched by model - // layer id, so layers absent from the attention cache are simply skipped. - if (auto * hyb = dynamic_cast(memory.get())) { - kv = hyb->get_mem_attn(); - } + // collect every standard llama_kv_cache the memory module keeps for attention + // layers. hybrid (recurrent + attention) models keep one for their attention + // sublayers; SWA variants keep a base/SWA pair. bias tensors are matched by + // model layer id, so layers absent from a given cache are simply skipped. + std::vector kvs; + if (auto * kv = dynamic_cast(memory.get())) { + kvs.push_back(kv); + } else if (auto * kv_iswa = dynamic_cast(memory.get())) { + kvs.push_back(kv_iswa->get_base()); + kvs.push_back(kv_iswa->get_swa()); + } else if (auto * hyb = dynamic_cast(memory.get())) { + kvs.push_back(hyb->get_mem_attn()); + } else if (auto * hyb_iswa = dynamic_cast(memory.get())) { + kvs.push_back(hyb_iswa->get_mem_attn()->get_base()); + kvs.push_back(hyb_iswa->get_mem_attn()->get_swa()); } - if (!kv) { - throw std::runtime_error("path_kv_mean_center is only supported for the standard KV cache " - "(not recurrent or MLA/DSA memory types)"); + if (kvs.empty()) { + throw std::runtime_error("path_kv_mean_center is only supported for standard KV caches " + "(not recurrent-only or MLA/DSA memory types)"); } - if (!kv->load_kv_mean_center(params.path_kv_mean_center)) { - throw std::runtime_error("failed to load K-cache mean-centering bias file"); + for (auto * kv : kvs) { + if (!kv->load_kv_mean_center(params.path_kv_mean_center)) { + throw std::runtime_error("failed to load K-cache mean-centering bias file"); + } } } } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3fa5fc2d6638..fcfb44eea08b 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1411,6 +1411,20 @@ bool llama_kv_cache::load_kv_mean_center(const char * path, bool require_q4_0) { return false; } + // validate the cache type first, so an unsupported cache type gets its actionable error + // even when the bias file's basis would also mismatch + if (require_q4_0) { + for (const auto & layer : layers) { + if (layer.k->type != GGML_TYPE_Q4_0) { + LLAMA_LOG_ERROR("%s: K-cache mean-centering requires K cache type %s, but layer %d has type %s\n", + __func__, ggml_type_name(GGML_TYPE_Q4_0), layer.il, ggml_type_name(layer.k->type)); + gguf_free(ctx_gguf); + ggml_free(ctx_data); + return false; + } + } + } + // the bias is only valid in the basis it was measured in: a bias calibrated with the // Hadamard K-cache rotation active must be applied with the rotation active, and vice // versa. a mismatched basis measurably degrades quantization quality instead of @@ -1418,6 +1432,13 @@ bool llama_kv_cache::load_kv_mean_center(const char * path, bool require_q4_0) { { const int64_t idx_k_rot = gguf_find_key(ctx_gguf, "kv_mean_center.k_rot"); if (idx_k_rot >= 0) { + if (gguf_get_kv_type(ctx_gguf, idx_k_rot) != GGUF_TYPE_BOOL) { + LLAMA_LOG_ERROR("%s: bias file %s has a non-boolean kv_mean_center.k_rot key - malformed file\n", + __func__, path); + gguf_free(ctx_gguf); + ggml_free(ctx_data); + return false; + } const bool bias_k_rot = gguf_get_val_bool(ctx_gguf, idx_k_rot); if (bias_k_rot != attn_rot_k) { LLAMA_LOG_ERROR("%s: bias file %s was calibrated with the K-cache rotation %s, but it is %s " @@ -1478,13 +1499,6 @@ bool llama_kv_cache::load_kv_mean_center(const char * path, bool require_q4_0) { continue; } - if (require_q4_0 && layers[ikv].k->type != GGML_TYPE_Q4_0) { - LLAMA_LOG_ERROR("%s: K-cache mean-centering requires K cache type %s, but layer %d has type %s\n", - __func__, ggml_type_name(GGML_TYPE_Q4_0), il, ggml_type_name(layers[ikv].k->type)); - ok = false; - break; - } - if (src->type != GGML_TYPE_F32) { LLAMA_LOG_ERROR("%s: bias tensor %s must be F32 (got %s)\n", __func__, name.c_str(), ggml_type_name(src->type)); diff --git a/tests/test-kv-mean-center.cpp b/tests/test-kv-mean-center.cpp index 5f7a46f7c15d..880415b6a989 100644 --- a/tests/test-kv-mean-center.cpp +++ b/tests/test-kv-mean-center.cpp @@ -236,14 +236,16 @@ static void test_q4_0_gate() { // F16 K cache + --kv-mean-center must be rejected outright (llama_init_from_model returns // nullptr), matching this codebase's convention for other cache-type-gated mismatches (e.g. - // "V cache quantization requires flash_attn") - llama_context_ptr ctx_bad = build_context(model.get(), GGML_TYPE_F16, tmp_path.c_str()); + // "V cache quantization requires flash_attn"). use a basis-matching file (F16 cache means + // the rotation is off, so k_rot=false matches) so this exercises the cache-type gate + // specifically, not the basis check + const std::string tmp_path_unrot = "test-kv-mean-center-gate-unrot.gguf"; + TEST_ASSERT(common_kv_mean_center_write(tmp_path_unrot, layers, /*k_rot=*/false)); + llama_context_ptr ctx_bad = build_context(model.get(), GGML_TYPE_F16, tmp_path_unrot.c_str()); TEST_ASSERT(ctx_bad == nullptr); // a bias measured in the unrotated basis must be rejected by a rotated (Q4_0) K cache: // a basis mismatch degrades quantization quality instead of improving it - const std::string tmp_path_unrot = "test-kv-mean-center-gate-unrot.gguf"; - TEST_ASSERT(common_kv_mean_center_write(tmp_path_unrot, layers, /*k_rot=*/false)); llama_context_ptr ctx_mismatch = build_context(model.get(), GGML_TYPE_Q4_0, tmp_path_unrot.c_str()); TEST_ASSERT(ctx_mismatch == nullptr); remove(tmp_path_unrot.c_str());