kv-cache: per-layer KV types, measured bit allocation

docs/turboquant.md

@@ -212,6 +212,104 @@ per-token magnitude scaling — the mechanism does not exist here. The per-chunk
 normalization that makes TurboQuant expensive at runtime is the same property
 that makes it robust over long contexts.
 
+### Per-layer bit allocation
+
+Layers do not benefit equally from the extra bit. RateQuant (arXiv 2605.06675)
+derives an uneven allocation from a fitted distortion curve; the numbers below
+measure it directly instead, on Qwen3.5-9B (8 attention layers, 24 linear-attention
+layers with no KV cache).
+
+Each row upgrades exactly one layer to turbo4 and leaves the other seven on
+turbo3, against a 1.5337 all-turbo3 baseline:
+
+| layer | PPL | delta |
+|-------|-----|-------|
+| 0 | 1.5341 | +0.0004 |
+| 1 | 1.5338 | +0.0001 |
+| 2 | 1.5324 | -0.0013 |
+| 3 | 1.5325 | -0.0012 |
+| 4 | 1.5321 | -0.0016 |
+| 5 | 1.5330 | -0.0007 |
+| 6 | 1.5329 | -0.0008 |
+| 7 | 1.5326 | -0.0011 |
+
+Layer 0 gets *worse* with more bits, which is impossible and puts the noise floor
+at roughly +/-0.0005. Every individual number is therefore within one to three
+times the noise — far too weak to allocate bits from on its own.
+
+The ranking they form, however, is not. Two runs at an identical 4.0 bpw, differing
+only in which four layers are upgraded:
+
+| allocation | bpw | PPL |
+|------------|-----|-----|
+| all turbo3 | 3.500 | 1.5337 |
+| worst four (0,1,5,6) | 4.000 | 1.5322 |
+| best four (2,3,4,7) | 4.000 | 1.5292 |
+| all turbo4 | 4.500 | 1.5284 |
+
+0.0030 separates two configurations that cost exactly the same — six times the
+noise floor. The ranking carries real information even though the measurements it
+was built from individually do not; the errors partly cancel when sorting.
+
+Walking the ranking (best first: 4, 2, 3, 7, 6, 5, 0, 1) gives the quality/size
+curve:
+
+| upgraded layers | bpw | PPL | share of the turbo4 gain |
+|-----------------|-----|-----|--------------------------|
+| 0 | 3.500 | 1.5337 | 0% |
+| 2 | 3.750 | 1.5311 | 49% |
+| 3 | 3.875 | 1.5303 | 64% |
+| 4 | 4.000 | 1.5292 | 85% |
+| 6 | 4.250 | 1.5279 | 109% |
+| 8 | 4.500 | 1.5284 | 100% |
+
+Four of eight layers capture 85% of what all eight provide, at half the extra
+cost. Six reach full turbo4 quality at 4.25 bpw instead of 4.5.
+
+That six beat eight is 0.0005, exactly the noise floor, so the ordering between
+those two is not itself meaningful — but the two remaining layers are 0 and 1,
+the only two whose single-layer measurement came out positive. Two independent
+measurements agree that they gain nothing, which is worth more than either alone.
+
+Two useful operating points, then: 4.0 bpw for 85% of the gain, or 4.25 bpw for
+all of it — against 4.5 bpw for uniform turbo4.
+
+Set with `LLAMA_KV_TYPE_PER_LAYER`, a comma-separated ggml type name per KV layer
+(K and V get the same type). Fewer entries than layers: the last one applies to
+the remainder, with a warning.
+
+```bash
+# best four layers of Qwen3.5-9B at turbo4, rest at turbo3 (4.0 bpw average)
+LLAMA_KV_TYPE_PER_LAYER=turbo3,turbo3,turbo4,turbo4,turbo4,turbo3,turbo3,turbo4 \
+  llama-cli -m model.gguf -fa on --cache-type-k turbo3 --cache-type-v turbo3
+```
+
+The ranking is model-specific — it was measured on this model and should not be
+assumed to transfer. Reproduce it for another model by upgrading one layer at a
+time, as above.
+
+#### K and V are equally sensitive
+
+Layers are not the only axis a bit budget can be spent on: `--cache-type-k` and
+`--cache-type-v` have always allowed K and V to differ. Work on KV quantization
+usually finds K the more sensitive of the two, because outlier channels
+concentrate there. At the same 4.0 bpw:
+
+| allocation | PPL |
+|------------|-----|
+| K=turbo4, V=turbo3 | 1.5319 |
+| K=turbo3, V=turbo4 | 1.5317 |
+| per-layer, best four | 1.5292 |
+
+0.0002 separates the two directions — noise. There is no K/V asymmetry to exploit
+here, which is what the Hadamard transform is for: every output is a signed sum of
+all 128 inputs, so a single outlier channel is spread across the whole vector
+instead of dominating one dimension. The asymmetry other methods exploit has
+already been flattened.
+
+So the layer axis is the productive one. It beats either K/V split by 0.0026 at
+identical cost, five times the noise floor.
+
 ### KV Cache Memory
 
 | KV Type | Bytes per element | Savings vs f16 |

src/llama-kv-cache.cpp

@@ -11,6 +11,7 @@
 #include <cstring>
 #include <limits>
 #include <map>
+#include <sstream>
 #include <stdexcept>
 
 //
@@ -99,6 +100,51 @@ llama_kv_cache::llama_kv_cache(
 
     const bool is_mla = hparams.is_mla();
 
+    // Experimental: per-layer KV types via LLAMA_KV_TYPE_PER_LAYER.
+    //
+    // Rate-distortion work on KV caches (RateQuant, arXiv 2605.06675) shows that
+    // spending the same average bit budget unevenly across layers beats spending
+    // it uniformly, because the distortion-vs-bits curve differs per layer. The
+    // cache already supports this structurally — tensors are created per layer
+    // below and every consumer reads the type off the tensor — so this only needs
+    // a way to say which layer gets what.
+    //
+    // Format: comma-separated ggml type names, one per KV layer, in order.
+    // Example: LLAMA_KV_TYPE_PER_LAYER=turbo4,turbo3,turbo3,turbo3,turbo3,turbo3,turbo3,turbo4
+    // Fewer entries than layers: the last one is repeated. Unset: type_k/type_v
+    // apply everywhere, as before.
+    std::vector<ggml_type> types_per_layer;
+    if (const char * spec = getenv("LLAMA_KV_TYPE_PER_LAYER")) {
+        std::istringstream ss(spec);
+        std::string name;
+        while (std::getline(ss, name, ',')) {
+            const size_t first = name.find_first_not_of(" \t");
+            if (first == std::string::npos) {
+                continue;
+            }
+            name = name.substr(first, name.find_last_not_of(" \t") - first + 1);
+
+            ggml_type found = GGML_TYPE_COUNT;
+            for (int t = 0; t < GGML_TYPE_COUNT; ++t) {
+                if (name == ggml_type_name((ggml_type) t)) {
+                    found = (ggml_type) t;
+                    break;
+                }
+            }
+            if (found == GGML_TYPE_COUNT) {
+                LLAMA_LOG_ERROR("%s: LLAMA_KV_TYPE_PER_LAYER: unknown type '%s'\n", __func__, name.c_str());
+                throw std::runtime_error("unknown KV type in LLAMA_KV_TYPE_PER_LAYER");
+            }
+            types_per_layer.push_back(found);
+        }
+        if (!types_per_layer.empty()) {
+            LLAMA_LOG_INFO("%s: per-layer KV types active (%d entries)\n", __func__, (int) types_per_layer.size());
+        }
+    }
+
+    // Index into types_per_layer, counted over layers that actually have a cache.
+    uint32_t kv_layer_idx = 0;
+
     for (uint32_t il = 0; il < hparams.n_layer; il++) {
         if (!hparams.has_kv(il)) {
             LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il);
@@ -132,18 +178,29 @@ llama_kv_cache::llama_kv_cache(
             throw std::runtime_error("failed to create ggml context for kv cache");
         }
 
+        ggml_type type_k_il = type_k;
+        ggml_type type_v_il = type_v;
+        if (!types_per_layer.empty()) {
+            const size_t idx = std::min<size_t>(kv_layer_idx, types_per_layer.size() - 1);
+            type_k_il = types_per_layer[idx];
+            type_v_il = types_per_layer[idx];
+            LLAMA_LOG_DEBUG("%s: layer %3d: KV type %s\n", __func__, il, ggml_type_name(type_k_il));
+        }
+        kv_layer_idx++;
+
         // TurboQuant applies the FWHT in fixed 128-element chunks, so head_dim
         // must be a multiple of 128. A 256-dim head is quantized as two
         // independently normalized halves — each half is a unit vector in R^128,
         // which is the distribution the codebooks were fitted for.
-        if (type_k == GGML_TYPE_TURBO3_0 || type_k == GGML_TYPE_TURBO4_0) {
+        // Checked against the effective per-layer type, not the global one.
+        if (type_k_il == GGML_TYPE_TURBO3_0 || type_k_il == GGML_TYPE_TURBO4_0) {
             const uint32_t n_embd_head_k = hparams.n_embd_head_k(il);
             if (n_embd_head_k % 128 != 0) {
                 LLAMA_LOG_ERROR("%s: TurboQuant requires head_dim to be a multiple of 128, got %d (layer %d)\n", __func__, n_embd_head_k, il);
                 throw std::runtime_error("turbo types require head_dim to be a multiple of 128");
             }
         }
-        if (type_v == GGML_TYPE_TURBO3_0 || type_v == GGML_TYPE_TURBO4_0) {
+        if (type_v_il == GGML_TYPE_TURBO3_0 || type_v_il == GGML_TYPE_TURBO4_0) {
             const uint32_t n_embd_head_v = hparams.n_embd_head_v(il);
             if (n_embd_head_v % 128 != 0) {
                 LLAMA_LOG_ERROR("%s: TurboQuant requires head_dim to be a multiple of 128, got %d (layer %d)\n", __func__, n_embd_head_v, il);
@@ -154,8 +211,8 @@ llama_kv_cache::llama_kv_cache(
         const bool has_k = true;
         const bool has_v = !is_mla;
 
-        ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr;
-        ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr;
+        ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k_il, n_embd_k_gqa, kv_size, n_stream) : nullptr;
+        ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v_il, n_embd_v_gqa, kv_size, n_stream) : nullptr;
 
         has_k && ggml_format_name(k, "cache_k_l%d", il);
         has_v && ggml_format_name(v, "cache_v_l%d", il);
@@ -173,6 +230,15 @@ llama_kv_cache::llama_kv_cache(
         layers.push_back({ il, k, v, k_stream, v_stream, });
     }
 
+    // A count mismatch is almost always a mistake — a list written for a different
+    // model, or one that forgot that only some layers have a KV cache. Repeating
+    // the last entry keeps it working, but silently, so say so.
+    if (!types_per_layer.empty() && types_per_layer.size() != kv_layer_idx) {
+        LLAMA_LOG_WARN("%s: LLAMA_KV_TYPE_PER_LAYER has %d entries but this model has %d KV layers; "
+                       "the last entry applies to the remainder\n",
+                       __func__, (int) types_per_layer.size(), (int) kv_layer_idx);
+    }
+
     if (reuse) {
         LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__);