kv-cache: optional per-channel K-cache mean-centering for Q4_0 by bri-prism · Pull Request #51 · PrismML-Eng/llama.cpp

bri-prism · 2026-07-07T19:35:48Z

Summary

Adds an opt-in K-cache mean-centering feature: a fixed, precomputed per-(kv-head, channel) bias is subtracted from the K vector right before it is quantized into the cache. This is currently scoped strictly to GGML_TYPE_Q4_0 K caches.

The idea

GGML_TYPE_Q4_0 is a symmetric (zero-point-free) block quantizer. 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. Centering the residual around zero before quantization reduces that error for channels with a real, consistent activation bias.

Why this is safe

For a fixed query row q attending over cached keys k_0 ... k_n, if every cached key is centered by the same bias k_bar before being quantized and stored:

q . k_i = q . (k_i - k_bar) + q . k_bar = q . k_i_stored + q . k_bar

The q . k_bar term doesn't depend on the key's position, so it's 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, so nothing else in attention needs to change. This also means it's a zero decode-time-cost win: one subtract at the point of cache write.

tests/test-kv-mean-center.cpp checks this directly: against an unquantized F32 K cache (to avoid confounding the comparison with real quantization error), a synthetic nonzero bias applied through the exact same code path leaves output logits matching the uncentered baseline to fp32 rounding (nmse ~5e-10 in practice).

What's included

Core mechanism (src/llama-kv-cache.{h,cpp}, src/llama-context.cpp, include/llama.h): a new path_kv_mean_center field on llama_context_params, applied in llama_kv_cache::cpy_k() right before the K value is quantized and stored. Gated strictly to k->type == GGML_TYPE_Q4_0; if the flag is set but the cache type doesn't match, context creation hard-fails with a clear error, matching this codebase's existing convention for other cache-type-gated mismatches (e.g. "V cache quantization requires flash_attn").
CLI flag (common/): --kv-mean-center <path>, requires -ctk q4_0.
Bias file format: a small GGUF file with one F32 tensor per layer (kv_bar.blk.<il>.k), shaped to broadcast against the K tensor's natural [n_embd_head, n_head] layout. Writer lives in common/kv-mean-center.* so it's shared between the calibration tool and the tests.
Calibration tool (tools/kv-mean-center/): follows the tools/imatrix convention. Runs a text calibration corpus through the model and captures the K tensor right before it's written into the cache, via a one-line addition of a cb() graph-build tag ("k_cache_in", in llm_graph_context::build_attn()) read through the same backend-scheduler eval-callback mechanism llama-imatrix uses.
Tests (tests/test-kv-mean-center.cpp): regression safety (centering disabled is a byte-for-byte no-op, including for a Q4_0 cache with no bias file loaded), the --kv-mean-center/Q4_0 gate itself, and the softmax-invariance math check described above.
Docs (docs/kv-mean-center.md, tools/kv-mean-center/README.md).

Scope / follow-ups

Only GGML_TYPE_Q4_0 is supported; other quant types are rejected. This isn't an arbitrary narrowing: Q4_0 is the only commonly-used int4 K-cache type that's symmetric (zero-point-free), which is the precondition for this technique to help at all. The other common int4 option, Q4_1, is asymmetric (it already stores a per-block min/bias alongside its scale), so it already absorbs a static per-channel mean the same way MLX's affine quantization does — centering would be a measurable no-op there. Generalizing further (e.g. to Q5_0, another symmetric type) is future work, but there's no other int4 type this would meaningfully apply to.
Only the plain (non-recurrent, non-hybrid, non-MLA/DSA) KV cache is supported.
The k_cache_in calibration hook currently only covers the standard dense/GQA attention path, which is the large majority of architectures but not all of them.
If this fork's optional Hadamard K/Q rotation is also active, the bias is calibrated pre-rotation but applied post-rotation; this stays exactly safe (the invariance argument is basis-independent) but is a less accurate calibration in that specific configuration. Noted in the docs as a follow-up.
Quantifying the real quantization-fidelity improvement on a production-scale model (e.g. a logit-KLD comparison against an uncentered Q4_0 baseline) is a natural follow-up; this PR doesn't include a measured number for any specific trained model.

Test plan

cmake --build build --target llama (core lib) builds clean
Full project build (llama-cli, llama-server, all tools/tests) builds clean
ctest -R test-kv-mean-center passes (regression safety, gate, invariance)
ctest -R test-llama-archs (existing test, unaffected) passes
Confirmed the 3 unrelated pre-existing test failures in this environment (test-tokenizers-ggml-vocabs, test-quantize-fns, test-quant-type-selection) reproduce identically on the unmodified base branch, i.e. are not caused by this change

GGML_TYPE_Q4_0 is a symmetric quantizer, so a K channel with a real, consistent nonzero mean across tokens wastes dynamic range encoding that constant bias. This adds an opt-in mechanism that subtracts a fixed per-(kv-head, channel) bias from K right before it is written into the cache in llama_kv_cache::cpy_k(), gated strictly to k->type == GGML_TYPE_Q4_0. Subtracting the same bias from every cached key is exactly softmax-invariant: it adds the same constant (q . k_bar) to every logit in a query's row, which softmax does not see. Nothing else in attention needs to change, so this is a zero decode-time-cost quantization-fidelity improvement. llama_context_params gets a new path_kv_mean_center field (default NULL); llama_init_from_model() hard-rejects it when the K cache type isn't Q4_0, matching the existing convention for other cache-type-gated mismatches (e.g. "V cache quantization requires flash_attn"). llama_kv_cache::load_kv_mean_center() loads the bias tensors from a GGUF file and applies them; a require_q4_0 escape hatch (used only by tests) exists to validate the underlying math against an unquantized cache without confounding it with real quantization error. Also tags the K tensor right before cpy_k() with the existing cb() graph-build hook ("k_cache_in"), so calibration tooling can capture exactly the tensor that gets written into the cache regardless of what RoPE/rotation preprocessing a given architecture applies upstream.

Adds the CLI-facing side of K-cache mean-centering: common_params gains kv_mean_center_path (plumbed into llama_context_params via common_context_params_to_llama), and --kv-mean-center takes a path to a bias file generated by tools/kv-mean-center. The bias file format is a small GGUF file with one F32 tensor per layer, 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]. The writer lives in common/kv-mean-center.* so it can be shared between the calibration tool and the test suite (which needs to synthesize a bias file to check the underlying math).

New tool, following the tools/imatrix convention: loads a model, runs a plain text calibration corpus through it in chunks, and captures the "k_cache_in" tensor tagged in llm_graph_context::build_attn() via the same backend-scheduler eval-callback mechanism llama-imatrix uses to capture activations (params.cb_eval). The per-(head,channel) mean across all calibration tokens is written out as a bias file consumable by --kv-mean-center.

Uses the same tiny-synthetic-model machinery as test-llama-archs.cpp (llama_model_saver + llama_model_init_from_user with a deterministic random tensor initializer), trimmed to plain LLM_ARCH_LLAMA, to check: - regression safety: two independent contexts with centering disabled produce bit-for-bit identical logits, and a Q4_0 K cache with no bias file loaded still decodes normally (cpy_k()'s new code path is a true no-op when k_bar is empty). - the --kv-mean-center gate: a non-Q4_0 K cache with a bias file is hard-rejected by llama_init_from_model(), while Q4_0 succeeds end-to-end through a real decode. - the softmax-invariance argument itself, against an unquantized F32 K cache with a synthetic nonzero bias applied through the exact same cpy_k() code path (bypassing the Q4_0 gate via load_kv_mean_center()'s require_q4_0=false test seam): output logits match the uncentered baseline to fp32 rounding (nmse ~5e-10 in practice), confirming the math without confounding it with real quantization error.

Explains the technique, the softmax-invariance argument, usage, the bias file format, and the current scope/limitations (Q4_0-only, plain KV cache only, standard dense/GQA attention path only).

Copilot

Pull request overview

This PR adds an opt-in “K-cache mean-centering” path that subtracts a precomputed per-(kv-head, channel) bias from K immediately before writing into the KV cache, currently gated to GGML_TYPE_Q4_0, plus tooling to calibrate/write the bias GGUF and tests/docs to validate softmax-invariance and gating.

Changes:

Add llama_context_params::path_kv_mean_center and load/apply a per-layer bias in llama_kv_cache::cpy_k() for Q4_0 K caches.
Introduce a new calibration tool (llama-kv-mean-center) and shared GGUF writer under common/.
Add regression + invariance tests and documentation for the feature and tool.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tools/kv-mean-center/README.md	Documents calibration tool usage and output format.
tools/kv-mean-center/kv-mean-center.cpp	Implements calibration by collecting `k_cache_in-*` tensors via eval callback.
tools/kv-mean-center/CMakeLists.txt	Adds build/install rules for the new calibration tool.
tools/CMakeLists.txt	Registers the new `kv-mean-center` tool directory.
tests/test-kv-mean-center.cpp	Adds regression, gate, and softmax-invariance tests for mean-centering.
tests/CMakeLists.txt	Adds the new test target to the test suite.
src/llama-kv-cache.h	Declares `load_kv_mean_center()` and bias storage in KV cache.
src/llama-kv-cache.cpp	Applies bias in `cpy_k()` and implements GGUF bias loading.
src/llama-graph.cpp	Tags `k_cur` with `k_cache_in` for calibration tooling.
src/llama-context.cpp	Loads bias file during context creation; adds default param initialization and type gate.
include/llama.h	Exposes `path_kv_mean_center` in the public C API context params.
docs/kv-mean-center.md	Adds feature documentation, rationale, file format, and scope notes.
common/kv-mean-center.h	Defines shared bias-layer struct and writer API.
common/kv-mean-center.cpp	Implements GGUF bias writer used by tool/tests.
common/common.h	Adds `LLAMA_EXAMPLE_KV_MEAN_CENTER` and `kv_mean_center_path` param.
common/common.cpp	Plumbs `kv_mean_center_path` into `llama_context_params`.
common/CMakeLists.txt	Builds the new common kv-mean-center writer into `llama-common`.
common/arg.cpp	Adds `--kv-mean-center` CLI flag and includes the new example for output-file routing.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+
+    std::lock_guard<std::mutex> lock(m_mutex);
+
+    GGML_ASSERT(t->type == GGML_TYPE_F32);


+    const float * f = (const float *) data;
+
+    for (int64_t i2 = 0; i2 < n_tokens; ++i2) {
+        for (int64_t i1 = 0; i1 < n_head; ++i1) {
+            for (int64_t i0 = 0; i0 < n_embd_head; ++i0) {
+                sum[i1*n_embd_head + i0] += f[(i2*n_head + i1)*n_embd_head + i0];
+            }


- kv-mean-center.cpp: the K tensor captured at "k_cache_in" can be F16 or BF16 depending on backend/compute settings, not just F32. The collector was asserting F32-only and reinterpreting raw bytes as float*, which either aborts or silently computes a garbage mean on non-F32 backends. Now accepts F32/F16/BF16 and converts to F32 via ggml_fp16_to_fp32_row/ggml_bf16_to_fp32_row before accumulating. - common/arg.cpp: wire --chunks into the LLAMA_EXAMPLE_KV_MEAN_CENTER example set so the flag the tool's own README documents is actually available, instead of being silently filtered out by the shared arg parser. - docs/kv-mean-center.md: replace non-ASCII "~=" and "." characters (was U+2248 and U+00B7) with ASCII equivalents, matching the project's ASCII-only docs convention.

bri-prism · 2026-07-07T21:04:09Z

Addressed the 6 Copilot review comments in 2551475:

tools/kv-mean-center/kv-mean-center.cpp (comments on the F32 assert and the raw-byte reinterpret): the K tensor captured at the "k_cache_in" tag can be F16 or BF16 depending on backend/compute settings, not only F32. The collector now accepts F32/F16/BF16 and converts to F32 with ggml_fp16_to_fp32_row/ggml_bf16_to_fp32_row before accumulating the mean, instead of asserting F32-only and reinterpreting raw bytes as float*.
tools/kv-mean-center/README.md (the --chunks gap): wired --chunks into the LLAMA_EXAMPLE_KV_MEAN_CENTER example set in common/arg.cpp, so the flag the README documents is actually available to this tool rather than being filtered out by the shared arg parser.
docs/kv-mean-center.md (the two non-ASCII characters): replaced the U+2248 ("approx equal") and U+00B7 (middle dot) characters with ASCII equivalents (~= and .), matching this project's ASCII-only docs convention.

On the CI failures: the three failing jobs (ubuntu x64, ubuntu arm64, windows) are pre-existing and unrelated to this PR.

ubuntu x64 and arm64 fail in test-quantize-fns on Q2_0 ("q2_0 absolute quantization error: FAILED", "q2_0 dot product error: FAILED"). This PR only touches GGML_TYPE_Q4_0 code paths and never references Q2_0 anywhere in its diff. I re-verified this directly: checking out the actual pre-PR base commit (0ad1dab, the parent of this PR's first commit) and building test-quantize-fns there reproduces the identical failure with the identical error values (0.008678 / 0.141111), with none of this PR's changes present. Rebuilding at the current PR tip reproduces the exact same numbers again, confirming this PR has zero effect on that test's outcome either way.
windows fails in unit/test_kv_keep_only_active.py::test_clear_and_restore, a server-side idle-slot-cache-clearing integration test with no relationship to K-cache mean-centering.

Not attempting to fix either of these here since they're outside this PR's scope.

khosravipasha

Took a pass, LGTM

some questions

for our fork this makes sense, do we want to eventually upstream? in that case might need a bunch of changes
for the mean centering needs calibration data? tools/kv-mean-center, doest it matter much? Mostly asking for next qestion;
what's plan to integrate into our bonsai-demo repo? Need clear examples and how to use there eventuallly.

also after merngin need to cut a build-release from actions eventually, so the prebuild version of demos also work out of the box.

bri-prism added 5 commits July 7, 2026 12:32

docs: document K-cache mean-centering

4025670

Explains the technique, the softmax-invariance argument, usage, the bias file format, and the current scope/limitations (Q4_0-only, plain KV cache only, standard dense/GQA attention path only).

github-actions Bot added documentation Improvements or additions to documentation examples testing labels Jul 7, 2026

bri-prism requested a review from khosravipasha July 7, 2026 19:38

khosravipasha requested a review from Copilot July 7, 2026 19:44

Copilot started reviewing on behalf of khosravipasha July 7, 2026 19:45 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

khosravipasha approved these changes Jul 9, 2026

View reviewed changes

bri-prism merged commit afc74b7 into prism Jul 9, 2026
10 of 35 checks passed

bri-prism mentioned this pull request Jul 10, 2026

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

Merged

khosravipasha deleted the feat/kv-mean-centering-public branch July 10, 2026 18:11

khosravipasha mentioned this pull request Jul 10, 2026

Guide: Q4_0 4-bit KV cache with per-channel mean-centering (setup + calibration) #54

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

kv-cache: optional per-channel K-cache mean-centering for Q4_0#51

kv-cache: optional per-channel K-cache mean-centering for Q4_0#51
bri-prism merged 6 commits into
prismfrom
feat/kv-mean-centering-public

bri-prism commented Jul 7, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bri-prism commented Jul 7, 2026

Uh oh!

khosravipasha left a comment •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants


		std::lock_guard<std::mutex> lock(m_mutex);

		GGML_ASSERT(t->type == GGML_TYPE_F32);

Uh oh!

Conversation

bri-prism commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

The idea

Why this is safe

What's included

Scope / follow-ups

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bri-prism commented Jul 7, 2026

Uh oh!

khosravipasha left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bri-prism commented Jul 7, 2026 •

edited

Loading

khosravipasha left a comment •

edited

Loading