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..548f5e4a7598 100644 --- a/docs/kv-mean-center.md +++ b/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. diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a867c24bda67..6ff36f288d14 100644 --- a/src/llama-context.cpp +++ b/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(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 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->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"); + } } } } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 09b0689e7b6c..fcfb44eea08b 100644 --- a/src/llama-kv-cache.cpp +++ b/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)); diff --git a/tests/test-kv-mean-center.cpp b/tests/test-kv-mean-center.cpp index f64f3f8f6ae2..880415b6a989 100644 --- a/tests/test-kv-mean-center.cpp +++ b/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()); diff --git a/tools/kv-mean-center/README.md b/tools/kv-mean-center/README.md index aca105b94f73..439e77792117 100644 --- a/tools/kv-mean-center/README.md +++ b/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. 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();