kv-cache: mean-centering on hybrid models + calibration-basis guard (composes with K rotation)#53
Merged
Merged
Conversation
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.
… 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.
github-actions
Bot
added
documentation
…smatch 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.
bri-prism
changed the title
There was a problem hiding this comment.
Pull request overview
Extends KV-cache mean-centering to hybrid attention caches and prevents calibration/inference rotation-basis mismatches.
Changes:
- Routes bias loading to hybrid attention sub-caches.
- Records and validates Hadamard K-rotation state.
- Updates tests and calibration guidance.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
common/kv-mean-center.cpp |
Writes rotation metadata. |
common/kv-mean-center.h |
Extends writer API. |
docs/kv-mean-center.md |
Documents hybrid support and basis matching. |
src/llama-context.cpp |
Routes hybrid bias loading. |
src/llama-kv-cache.cpp |
Enforces rotation-basis compatibility. |
tests/test-kv-mean-center.cpp |
Adds mismatch rejection coverage. |
tools/kv-mean-center/README.md |
Adds calibration instructions and measurements. |
tools/kv-mean-center/kv-mean-center.cpp |
Detects rotation through graph ancestry. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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.
khosravipasha
approved these changes
Jul 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Three changes to the K-cache mean-centering feature (#51):
--kv-mean-centernow works on hybrid (recurrent + attention) models: the bias load is routed to the hybrid memory's attention sub-cache. Previously every hybrid model failed to load a bias file outright.kv_mean_center.k_rotin the bias GGUF, detected from the captured tensor's graph ancestry), and the loader rejects a bias whose basis does not match the inference-time rotation state. Files predating the flag load with a warning.--cache-type-k(andLLAMA_ATTN_ROT_DISABLE, if any) that you serve with.Why
The hybrid rejection was an avoidable hard failure on a whole model family, since the calibration tool already handles those models.
With that fixed, end-to-end measurement (logit KL divergence vs an F16-cache baseline, 12x512-token held-out chunks, Q4_0 K cache, hybrid-attention model) surfaced a basis subtlety. The rotation only activates for quantized K caches, so a calibration run with the default F16 cache measures the unrotated basis; the collector taps the exact tensor
cpy_k()writes, so a run with-ctk q4_0measures the rotated basis. Applied in the matching basis the two features compose; mismatched, centering makes quantization quality worse than doing nothing:-ctk q4_0)A silent quality regression from an easy-to-make calibration mistake is exactly the kind of thing the loader should catch, hence the recorded basis and the hard reject.
How
llama_contextfalls back tollama_memory_hybrid::get_mem_attn()when the memory module is hybrid. Bias tensors are matched by model layer id and absent layers are skipped, which the loader already handles.k_cache_intensor looking for theattn_inp_k_rotinput (substring match, since views and backend-scheduler split copies decorate the name) and passes the result tocommon_kv_mean_center_write.load_kv_mean_center()compares the stored flag against the cache'sattn_rot_kand fails with a clear message on mismatch.test-kv-mean-centerextended: the gate test now also asserts an unrotated-basis bias is rejected by a rotated (Q4_0) context. All tests pass; every row of the table above was validated live, including both mismatch rejection directions and the legacy-file warning path.