diff --git a/docs/turboquant.md b/docs/turboquant.md index e03e20b0c..13224fd59 100644 --- a/docs/turboquant.md +++ b/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 | diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index dd061faec..43476298f 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include // @@ -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 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(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__);