kv-cache: per-channel K-cache mean-centering for Q4_0 (fork #51)

Re-port of the fork's --kv-mean-center feature onto upstream master. Adds
common/kv-mean-center, the arg + LLAMA_EXAMPLE_KV_MEAN_CENTER, the llama.h API,
and the kv-cache/context/graph hooks. Conflicts resolved against upstream's
rewritten kv-cache/context (kept only the path_kv_mean_center guard).

Author: khosravipasha

common/CMakeLists.txt

@@ -81,6 +81,8 @@ add_library(${TARGET}
     imatrix-loader.cpp
     imatrix-loader.h
     json-schema-to-grammar.cpp
+    kv-mean-center.cpp
+    kv-mean-center.h
     llguidance.cpp
     log.cpp
     log.h

common/arg.cpp

@@ -1682,7 +1682,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, int value) {
             params.n_chunks = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_RETRIEVAL}));
+    ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_KV_MEAN_CENTER}));
     add_opt(common_arg({ "-fa", "--flash-attn" }, "[on|off|auto]",
                        string_format("set Flash Attention use ('on', 'off', or 'auto', default: '%s')",
                                      llama_flash_attn_type_name(params.flash_attn_type)),
@@ -2391,6 +2391,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.cache_type_v = kv_cache_type_from_str(value);
         }
     ).set_env("LLAMA_ARG_CACHE_TYPE_V"));
+    add_opt(common_arg(
+        {"--kv-mean-center"}, "FNAME",
+        "path to a K-cache mean-centering bias file (GGUF), generated with tools/kv-mean-center\n"
+        "subtracts a fixed per-(head,channel) bias from K before it is quantized into the cache;\n"
+        "requires --cache-type-k q4_0 (see docs/kv-mean-center.md)",
+        [](common_params & params, const std::string & value) {
+            params.kv_mean_center_path = value;
+        }
+    ).set_env("LLAMA_ARG_KV_MEAN_CENTER"));
     add_opt(common_arg(
         {"--hellaswag"},
         "compute HellaSwag score over random tasks from datafile supplied with -f",
@@ -3044,7 +3053,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.out_file = value;
         }
     ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_CVECTOR_GENERATOR, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_FINETUNE,
-                    LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, LLAMA_EXAMPLE_CLI}));
+                    LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_KV_MEAN_CENTER}));
     add_opt(common_arg(
         {"-ofreq", "--output-frequency"}, "N",
         string_format("output the imatrix every N iterations (default: %d)", params.n_out_freq),

common/common.cpp

@@ -1667,6 +1667,10 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;
 
+    // note: params (and therefore params.kv_mean_center_path) is kept alive by the caller for
+    // at least as long as it takes to call llama_init_from_model() with the returned cparams
+    cparams.path_kv_mean_center = params.kv_mean_center_path.empty() ? nullptr : params.kv_mean_center_path.c_str();
+
     return cparams;
 }
 

common/common.h

@@ -107,6 +107,7 @@ enum llama_example {
     LLAMA_EXAMPLE_EXPORT_GRAPH_OPS,
     LLAMA_EXAMPLE_DOWNLOAD,
     LLAMA_EXAMPLE_TOKENIZE,
+    LLAMA_EXAMPLE_KV_MEAN_CENTER,
 
     LLAMA_EXAMPLE_COUNT,
 };
@@ -576,6 +577,10 @@ struct common_params {
     ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
     ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
 
+    // path to a K-cache mean-centering bias file (GGUF), or empty to disable.
+    // only takes effect when cache_type_k == GGML_TYPE_Q4_0; see docs/kv-mean-center.md
+    std::string kv_mean_center_path = "";
+
     common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;
 
     // multimodal models (see tools/mtmd)

common/kv-mean-center.cpp

@@ -0,0 +1,74 @@
+#include "kv-mean-center.h"
+
+#include "log.h"
+
+#include "ggml.h"
+#include "gguf.h"
+
+#include <cstring>
+
+bool common_kv_mean_center_write(
+        const std::string & fname,
+        const std::vector<common_kv_mean_center_layer> & layers) {
+    size_t n_with_bias = 0;
+    for (const auto & layer : layers) {
+        if (!layer.bias.empty()) {
+            n_with_bias++;
+        }
+    }
+
+    if (n_with_bias == 0) {
+        LOG_ERR("%s: no layers with bias data to write\n", __func__);
+        return false;
+    }
+
+    size_t data_size = 0;
+    for (const auto & layer : layers) {
+        if (!layer.bias.empty()) {
+            data_size += GGML_PAD(ggml_tensor_overhead() + sizeof(float)*layer.bias.size(), GGML_MEM_ALIGN);
+        }
+    }
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ data_size,
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ false,
+    };
+
+    struct ggml_context  * ctx      = ggml_init(params);
+    struct gguf_context   * ctx_gguf = gguf_init_empty();
+
+    if (!ctx || !ctx_gguf) {
+        LOG_ERR("%s: failed to allocate ggml/gguf context\n", __func__);
+        if (ctx)      ggml_free(ctx);
+        if (ctx_gguf) gguf_free(ctx_gguf);
+        return false;
+    }
+
+    gguf_set_val_str(ctx_gguf, "general.type", "kv-mean-center");
+
+    for (const auto & layer : layers) {
+        if (layer.bias.empty()) {
+            continue;
+        }
+
+        const std::string name = "kv_bar.blk." + std::to_string(layer.il) + ".k";
+
+        ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) layer.bias.size());
+        ggml_set_name(t, name.c_str());
+
+        memcpy(t->data, layer.bias.data(), layer.bias.size()*sizeof(float));
+
+        gguf_add_tensor(ctx_gguf, t);
+    }
+
+    const bool ok = gguf_write_to_file(ctx_gguf, fname.c_str(), false);
+    if (!ok) {
+        LOG_ERR("%s: failed to write %s\n", __func__, fname.c_str());
+    }
+
+    gguf_free(ctx_gguf);
+    ggml_free(ctx);
+
+    return ok;
+}

common/kv-mean-center.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+// Shared GGUF file format for the K-cache mean-centering bias vectors produced by
+// tools/kv-mean-center and consumed by llama_kv_cache::load_kv_mean_center() (see
+// docs/kv-mean-center.md).
+//
+// The file stores one F32 tensor per model layer that has a bias, named
+// "kv_bar.blk.<il>.k", holding n_embd_head_k(il) * n_head_kv(il) values laid out as
+// [n_embd_head_k, n_head_kv] (channel-fastest), matching the memory layout of the K
+// tensor at the point it is written into the cache.
+
+// per-layer entry to write; `bias` empty means "no bias for this layer" (skipped on write)
+struct common_kv_mean_center_layer {
+    int32_t            il = -1;
+    std::vector<float> bias; // length n_embd_head_k(il) * n_head_kv(il)
+};
+
+// write a K-cache mean-centering bias file in GGUF format.
+// 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);

docs/kv-mean-center.md

@@ -0,0 +1,98 @@
+# K-cache mean-centering
+
+`--kv-mean-center` is an optional, opt-in feature that subtracts a fixed, precomputed
+per-(kv-head, channel) bias from the K vector at the moment it is written into the KV cache,
+in order to improve quantization fidelity for `GGML_TYPE_Q4_0` K caches.
+
+This is currently scoped to `GGML_TYPE_Q4_0` only.
+
+## The idea
+
+`GGML_TYPE_Q4_0` is a symmetric (zero-point-free) block quantizer: each block is represented as
+`value ~= scale * q`, where `q` is a signed low-bit integer and there is no bias/zero-point term.
+If a given (kv-head, channel) position's real K activations have a nonzero mean across tokens,
+symmetric quantization wastes some of its dynamic range encoding that constant bias, which
+increases quantization error for that channel.
+
+The fix: measure a per-(kv-head, channel) bias `k_bar` ahead of time (see
+[Calibration](#calibration) below), and subtract it from the K vector for every token, right
+before `Q4_0` quantization happens as part of writing it into the cache. Nothing else in
+attention needs to change.
+
+### Why this is safe (softmax-invariance)
+
+For a fixed query row `q` attending over cached keys `k_0 ... k_n` (all in one layer/head), if
+every cached key is centered by the same `k_bar` before being quantized and stored, the true dot
+product decomposes as:
+
+```
+q . k_i = q . (k_i - k_bar) + q . k_bar = q . k_i_stored + q . k_bar
+```
+
+The `q . k_bar` term does not depend on `i` (the key's position) -- it is added identically to
+every logit in that query's row. Softmax is invariant to a constant additive shift applied to
+every logit in the same row (`softmax(x + c) == softmax(x)`), so the attention weights, and
+therefore the rest of the model's output, are unaffected. This means the technique is exactly
+correctness-preserving in infinite precision, and in practice the only observable difference is
+ordinary floating point rounding (see `tests/test-kv-mean-center.cpp`, which checks this directly
+against an unquantized F32 K cache). This is also why it is a zero decode-time-cost win: one
+subtract at the point of cache write, nothing else changes.
+
+The actual benefit is purely on quantization fidelity: centering the residual around zero before
+`Q4_0`'s symmetric quantizer reduces per-channel quantization error for channels that have a real,
+consistent activation bias. Quantifying that improvement on a production-scale model (e.g. via a
+logit-KLD comparison against an uncentered `Q4_0` baseline) is a natural follow-up; this repo does
+not ship a measured number for a specific trained model.
+
+## Usage
+
+1. Generate a bias file with `tools/kv-mean-center` (see its
+   [README](../tools/kv-mean-center/README.md) for details):
+
+   ```
+   ./llama-kv-mean-center -m model.gguf -f calibration-data.txt -o kv-mean-center.gguf
+   ```
+
+2. Load it at inference time, together with a `Q4_0` K cache:
+
+   ```
+   ./llama-cli -m model.gguf -ctk q4_0 --kv-mean-center kv-mean-center.gguf -p "..."
+   ```
+
+`--kv-mean-center` requires `--cache-type-k q4_0`. If the K cache type is anything else, context
+creation fails with a clear error rather than silently doing nothing, matching this codebase's
+existing convention for other cache-type-gated options (e.g. quantized V cache requiring flash
+attention).
+
+## Bias file format
+
+The bias file is a small GGUF file with one F32 1-D tensor per layer that has a bias, named
+`kv_bar.blk.<il>.k`, holding `n_embd_head_k(il) * n_head_kv(il)` values laid out as
+`[n_embd_head_k, n_head_kv]` (channel-fastest). This matches the in-memory layout of the K tensor
+at the point it is written into the cache, so the file can be loaded directly as a small
+broadcastable bias tensor per layer.
+
+## Calibration
+
+`tools/kv-mean-center` computes the bias by running a plain text calibration corpus through the
+model and averaging the K tensor right before it would be written into the cache (via the
+`k_cache_in` tag added to `llm_graph_context::build_attn()`, read through the same backend
+scheduler eval-callback mechanism `llama-imatrix` uses to capture activations). See
+[tools/kv-mean-center/README.md](../tools/kv-mean-center/README.md) for usage.
+
+## Scope and limitations
+
+- 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 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.

include/llama.h

@@ -380,6 +380,14 @@ extern "C" {
         enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]
         enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]
 
+        // optional path to a per-layer K-cache mean-centering bias file (GGUF), or NULL to disable.
+        // the bias is subtracted from the K vector for each (kv-head, channel) right before it is
+        // written into the K cache, which improves quantization fidelity for GGML_TYPE_Q4_0 without
+        // changing attention results (the same constant is added to every logit in a query's row,
+        // which softmax is invariant to). currently only supported when type_k == GGML_TYPE_Q4_0.
+        // see tools/kv-mean-center to generate this file and docs/kv-mean-center.md for details.
+        const char * path_kv_mean_center;
+
         // Abort callback
         // if it returns true, execution of llama_decode() will be aborted
         // currently works only with CPU execution

src/llama-context.cpp

@@ -6,6 +6,7 @@
 #include "llama-impl.h"
 #include "llama-batch.h"
 #include "llama-io.h"
+#include "llama-kv-cache.h"
 #include "llama-memory.h"
 #include "llama-mmap.h"
 #include "llama-model.h"
@@ -389,6 +390,17 @@ 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)");
+            }
+            if (!kv->load_kv_mean_center(params.path_kv_mean_center)) {
+                throw std::runtime_error("failed to load K-cache mean-centering bias file");
+            }
+        }
     }
 
     // init backends
@@ -3507,6 +3519,7 @@ llama_context_params llama_context_default_params() {
         /*.cb_eval_user_data           =*/ nullptr,
         /*.type_k                      =*/ GGML_TYPE_F16,
         /*.type_v                      =*/ GGML_TYPE_F16,
+        /*.path_kv_mean_center         =*/ nullptr,
         /*.abort_callback              =*/ nullptr,
         /*.abort_callback_data         =*/ nullptr,
         /*.embeddings                  =*/ false,
@@ -3595,6 +3608,13 @@ llama_context * llama_init_from_model(
         }
     }
 
+    if (params.path_kv_mean_center != nullptr && params.type_k != GGML_TYPE_Q4_0) {
+        LLAMA_LOG_ERROR("%s: path_kv_mean_center requires the K cache type to be Q4_0 (got %s)\n",
+                __func__, ggml_type_name(params.type_k));
+        return nullptr;
+    }
+
+
     if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED &&
         params.pooling_type != model->hparams.pooling_type) {
         //user-specified pooling-type is different from the model default

src/llama-graph.cpp

@@ -2779,6 +2779,11 @@ ggml_tensor * llm_graph_context::build_attn(
         const auto & k_idxs = inp->get_k_idxs();
         const auto & v_idxs = inp->get_v_idxs();
 
+        // hook point for KV cache calibration tooling (e.g. tools/kv-mean-center): this is
+        // exactly the K tensor that cpy_k() writes into the cache, after any RoPE/rotation
+        // the architecture applies upstream
+        cb(k_cur, "k_cache_in", il);
+
         ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
         ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
     }