feat: --save-kv / --load-kv CLI for "Read Once, Query Forever"

Implements KV cache persistence for Document-Level RAG pattern:

  # Process document once (slow prefill)
  ./build/quant model.gguf -p "long document..." --save-kv doc.kv

  # Query instantly, forever (KV restored in <1s)
  ./build/quant model.gguf -p "question?" --load-kv doc.kv

Implementation:
- Per-layer strided save/load (respects max_seq * kv_dim layout)
- Saves FP32 key cache + FP16/FP32 value cache
- Header: position count + kv_dim for validation
- New prompt appended after loaded KV positions

Verified: 3B model recalls "PHOENIX" from saved context.
35/35 tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: unamedkr

include/turboquant/tq_engine.h

@@ -383,6 +383,9 @@ typedef struct {
     int n_threads;
     float rep_penalty;    /* repetition penalty (default: 1.1, 1.0 = disabled) */
     int rep_window;       /* how many recent tokens to penalize (default: 32) */
+    /* KV cache persistence (Document-Level RAG: read once, query forever) */
+    const char* save_kv_path; /* save KV cache after generation (NULL = don't save) */
+    const char* load_kv_path; /* load pre-computed KV cache before generation (NULL = normal) */
     /* Callback for streaming output */
     void (*on_token)(const char* text, void* user_data);
     void* user_data;

src/engine/tq_generate.c

@@ -245,9 +245,79 @@ int tq_generate(tq_model_t* model, tq_tokenizer_t* tokenizer,
         fprintf(stderr, "\n");
     }
 
-    /* Prefill: process all prompt tokens */
+    /* Load pre-computed KV cache if available (skip prefill) */
+    int pos_after_prefill = n_prompt;
+    if (config->load_kv_path) {
+        FILE* kv_fp = fopen(config->load_kv_path, "rb");
+        if (kv_fp) {
+            int32_t saved_pos = 0;
+            size_t kv_dim_save = 0;
+            fread(&saved_pos, sizeof(int32_t), 1, kv_fp);
+            fread(&kv_dim_save, sizeof(size_t), 1, kv_fp);
+            size_t kv_dim = (size_t)model->config.n_kv_heads * model->config.head_dim;
+            int max_seq = model->config.max_seq_len;
+            size_t layer_stride = (size_t)max_seq * kv_dim;
+            /* Read per-layer, respecting stride */
+            for (int l = 0; l < model->config.n_layers; l++) {
+                if (state->key_cache)
+                    fread(state->key_cache + l * layer_stride, sizeof(float), (size_t)saved_pos * kv_dim, kv_fp);
+                if (state->value_cache_fp16)
+                    fread(state->value_cache_fp16 + l * layer_stride, sizeof(uint16_t), (size_t)saved_pos * kv_dim, kv_fp);
+                else if (state->value_cache)
+                    fread(state->value_cache + l * layer_stride, sizeof(float), (size_t)saved_pos * kv_dim, kv_fp);
+            }
+            fclose(kv_fp);
+            pos_after_prefill = saved_pos;
+            size_t total_bytes = (size_t)model->config.n_layers * saved_pos * kv_dim * (sizeof(float) + (state->value_cache_fp16 ? sizeof(uint16_t) : sizeof(float)));
+            fprintf(stderr, "[load-kv] Loaded %d tokens from %s (%.1f MB)\n",
+                    saved_pos, config->load_kv_path,
+                    (double)total_bytes / (1024.0 * 1024.0));
+        } else {
+            fprintf(stderr, "[load-kv] Cannot open %s, running normal prefill\n", config->load_kv_path);
+        }
+    }
+
+    /* Prefill: process prompt tokens.
+     * If KV was loaded, the loaded context occupies positions [0..pos_after_prefill).
+     * The new prompt is appended starting at pos_after_prefill. */
+    int prefill_start = 0;
+    if (config->load_kv_path && pos_after_prefill > 0) {
+        prefill_start = pos_after_prefill;
+    }
     for (int i = 0; i < n_prompt; i++) {
-        tq_forward(model, state, prompt_tokens[i], i);
+        tq_forward(model, state, prompt_tokens[i], prefill_start + i);
+    }
+    pos_after_prefill = prefill_start + n_prompt;
+
+    /* Save KV cache after prefill if requested */
+    if (config->save_kv_path && pos_after_prefill > 0) {
+        FILE* kv_fp = fopen(config->save_kv_path, "wb");
+        if (kv_fp) {
+            int32_t save_pos = (int32_t)pos_after_prefill;
+            size_t kv_dim = (size_t)model->config.n_kv_heads * model->config.head_dim;
+            int max_seq = model->config.max_seq_len;
+            size_t layer_stride = (size_t)max_seq * kv_dim;
+            fwrite(&save_pos, sizeof(int32_t), 1, kv_fp);
+            fwrite(&kv_dim, sizeof(size_t), 1, kv_fp);
+            /* Write per-layer, only saved_pos positions */
+            size_t total = 0;
+            for (int l = 0; l < model->config.n_layers; l++) {
+                if (state->key_cache) {
+                    fwrite(state->key_cache + l * layer_stride, sizeof(float), (size_t)save_pos * kv_dim, kv_fp);
+                    total += (size_t)save_pos * kv_dim * sizeof(float);
+                }
+                if (state->value_cache_fp16) {
+                    fwrite(state->value_cache_fp16 + l * layer_stride, sizeof(uint16_t), (size_t)save_pos * kv_dim, kv_fp);
+                    total += (size_t)save_pos * kv_dim * sizeof(uint16_t);
+                } else if (state->value_cache) {
+                    fwrite(state->value_cache + l * layer_stride, sizeof(float), (size_t)save_pos * kv_dim, kv_fp);
+                    total += (size_t)save_pos * kv_dim * sizeof(float);
+                }
+            }
+            fclose(kv_fp);
+            fprintf(stderr, "[save-kv] Saved %d tokens to %s (%.1f MB)\n",
+                    save_pos, config->save_kv_path, (double)total / (1024.0 * 1024.0));
+        }
     }
 
     /* Repetition penalty setup */
@@ -290,7 +360,7 @@ int tq_generate(tq_model_t* model, tq_tokenizer_t* tokenizer,
     }
 
     /* Sample first generated token */
-    int pos = n_prompt;
+    int pos = pos_after_prefill;
     unsigned long long rng_state = 42;
     int next_token = tq_sample_topp(state->logits, vocab_size,
                                      config->temperature, config->top_p,

tools/quant.c

@@ -132,6 +132,8 @@ static void print_usage(const char* prog) {
     fprintf(stderr, "  --ctx <N>        Override max context length (default: 4096)\n");
     fprintf(stderr, "  --delta, -D      Enable delta KV compression (store key deltas)\n");
     fprintf(stderr, "  --k-window <N>   Age-based K: recent N tokens FP32, rest quantized\n");
+    fprintf(stderr, "  --save-kv <file> Save KV cache after generation (read once, query forever)\n");
+    fprintf(stderr, "  --load-kv <file> Load pre-computed KV cache (skip prefill)\n");
     fprintf(stderr, "  --version        Print version and exit\n");
     fprintf(stderr, "  --json           JSON output for --ppl (machine-parseable)\n");
     fprintf(stderr, "  --save-logits <f> Save per-token softmax (fp16) to file during --ppl\n");
@@ -195,6 +197,8 @@ int main(int argc, char** argv) {
     int chat_mode = 0;       /* 1 = auto-wrap prompt with chat template */
     const char* save_logits_file = NULL;
     const char* kl_baseline_file = NULL;
+    const char* save_kv_file = NULL;   /* --save-kv: save KV cache after generation */
+    const char* load_kv_file = NULL;   /* --load-kv: load pre-computed KV cache */
 
     for (int i = 1; i < argc; i++) {
         if (argv[i][0] != '-') {
@@ -282,6 +286,10 @@ int main(int argc, char** argv) {
         } else if (strcmp(argv[i], "--version") == 0) {
             print_version();
             return 0;
+        } else if (strcmp(argv[i], "--save-kv") == 0 && i + 1 < argc) {
+            save_kv_file = argv[++i];
+        } else if (strcmp(argv[i], "--load-kv") == 0 && i + 1 < argc) {
+            load_kv_file = argv[++i];
         } else if (strcmp(argv[i], "--save-logits") == 0 && i + 1 < argc) {
             save_logits_file = argv[++i];
         } else if (strcmp(argv[i], "--kl-baseline") == 0 && i + 1 < argc) {
@@ -1255,6 +1263,8 @@ int main(int argc, char** argv) {
     config.delta_kv = delta_kv;
     config.delta_iframe_interval = delta_iframe_int;
     config.k_highres_window = k_highres_window;
+    config.save_kv_path = save_kv_file;
+    config.load_kv_path = load_kv_file;
     config.on_token = print_token;
     config.user_data = NULL;