kv-cache: mean-centering on hybrid models + calibration-basis guard (composes with K rotation)

common/kv-mean-center.cpp

@@ -9,7 +9,8 @@
 
 bool common_kv_mean_center_write(
         const std::string & fname,
-        const std::vector<common_kv_mean_center_layer> & layers) {
+        const std::vector<common_kv_mean_center_layer> & 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()) {

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<common_kv_mean_center_layer> & layers);
+        const std::vector<common_kv_mean_center_layer> & layers,
+        bool k_rot = false);

docs/kv-mean-center.md

@@ -84,15 +84,19 @@ 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.
+- 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
   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.

src/llama-context.cpp

@@ -8,6 +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"
@@ -322,13 +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<llama_kv_cache *>(memory.get());
-            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)");
+            // 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<llama_kv_cache *> kvs;
+            if (auto * kv = dynamic_cast<llama_kv_cache *>(memory.get())) {
+                kvs.push_back(kv);
+            } else if (auto * kv_iswa = dynamic_cast<llama_kv_cache_iswa *>(memory.get())) {
+                kvs.push_back(kv_iswa->get_base());
+                kvs.push_back(kv_iswa->get_swa());
+            } else if (auto * hyb = dynamic_cast<llama_memory_hybrid *>(memory.get())) {
+                kvs.push_back(hyb->get_mem_attn());
+            } else if (auto * hyb_iswa = dynamic_cast<llama_memory_hybrid_iswa *>(memory.get())) {
+                kvs.push_back(hyb_iswa->get_mem_attn()->get_base());
+                kvs.push_back(hyb_iswa->get_mem_attn()->get_swa());
             }
-            if (!kv->load_kv_mean_center(params.path_kv_mean_center)) {
-                throw std::runtime_error("failed to load K-cache mean-centering bias file");
+            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)");
+            }
+            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");
+                }
             }
         }
     }

src/llama-kv-cache.cpp

@@ -1411,6 +1411,53 @@ 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
+    // 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) {
+            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 "
+                        "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
@@ -1452,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));

tests/test-kv-mean-center.cpp

@@ -229,15 +229,27 @@ 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.
-    // "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
+    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());

tools/kv-mean-center/README.md

@@ -41,11 +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`).
 
-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.
+## 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.

tools/kv-mean-center/kv-mean-center.cpp

@@ -29,15 +29,41 @@ class kv_mean_collector {
 
     std::vector<common_kv_mean_center_layer> finalize() const;
 
+    bool saw_k_rot() const { return m_saw_k_rot; }
+
 private:
     std::mutex m_mutex;
     std::vector<uint8_t> m_host_buf;
     std::vector<float>   m_f32_buf; // scratch space for F16/BF16 -> F32 conversion
 
     std::unordered_map<int32_t, std::vector<double>> m_sum;   // il -> [n_embd_head * n_head]
     std::unordered_map<int32_t, int64_t>              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-<il>".
+    // 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>" -> 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<std::mutex> 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();