phase 9: KV cache + autoregressive generation

- KVCache: per-layer, per-head storage with append + reconstruct
- forward_with_cache: prefill (full prompt) + decode (single token) modes
- Fixed data layout bug: per-head vectors avoid cross-head interleaving
- CLI updated to use KV cache by default
- bench-gpt2 supports --no-cache flag for comparison

Benchmark results (50 prompts × 20 tokens):
- KV cache vs no-cache: 50/50 bit-identical (cache is correct)
- 18x speedup: TTFT 400→24ms, TBT 407→22ms, throughput 2.5→44 tok/s
- vs HF transformers: 40/50 match (10 are FP divergence, avg logit gap 0.20)

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

crates/xserv-model/src/bin/xserv-cli.rs

@@ -1,21 +1,20 @@
 use std::io::{self, Write};
 use std::path::PathBuf;
-use xserv_model::{GPT2, ModelConfig};
-use xserv_model::loader;
-use xserv_model::gpt2::sample_greedy;
-use xserv_tokenizer::Tokenizer;
+use xserv_model::gpt2::{sample_greedy, KVCache};
+use xserv_model::{loader, GPT2, ModelConfig};
 use xserv_tensor::Device;
+use xserv_tokenizer::Tokenizer;
 
 fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
         eprintln!("Usage: xserv-cli <model-dir> [--max-tokens N]");
-        eprintln!("  model-dir: path to HF model directory (containing model.safetensors, config.json, tokenizer.json)");
         std::process::exit(1);
     }
 
     let model_dir = PathBuf::from(&args[1]);
-    let max_tokens: usize = args.iter()
+    let max_tokens: usize = args
+        .iter()
         .position(|a| a == "--max-tokens")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
@@ -25,26 +24,24 @@ fn main() {
     let info = xserv_cuda::device::device_info(0).unwrap();
     eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
 
-    // Load config
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
-    eprintln!("Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
-        config.model_type, config.num_layers(), config.hidden(),
-        config.num_heads(), config.vocab_size);
+    eprintln!(
+        "Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
+        config.model_type,
+        config.num_layers(),
+        config.hidden(),
+        config.num_heads(),
+        config.vocab_size
+    );
 
-    // Load weights
     eprintln!("Loading weights...");
     let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
     eprintln!("Loaded {} tensors", weights.len());
 
-    // GPT-2 uses weight names without "model." prefix
-    let model = GPT2::from_weights(config, weights);
-
-    // Load tokenizer
+    let model = GPT2::from_weights(config.clone(), weights);
     let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-    eprintln!("Tokenizer loaded (vocab_size={})", tokenizer.vocab_size());
-    eprintln!("Ready.\n");
+    eprintln!("Ready (KV cache enabled).\n");
 
-    // Interactive loop
     loop {
         print!("xserv> ");
         io::stdout().flush().unwrap();
@@ -56,22 +53,27 @@ fn main() {
         if input.is_empty() { continue; }
         if input == "quit" || input == "exit" { break; }
 
-        let mut token_ids = tokenizer.encode(input);
+        let token_ids = tokenizer.encode(input);
+        let mut cache = KVCache::new(
+            config.num_layers(), config.num_heads(), config.head_dim(),
+            Device::Cuda(0),
+        );
+
+        // Prefill
+        let logits = model.forward_with_cache(&token_ids, &mut cache);
+        let mut next = sample_greedy(&logits);
         print!("{input}");
         io::stdout().flush().unwrap();
 
         for _ in 0..max_tokens {
-            let logits = model.forward(&token_ids);
-            let next = sample_greedy(&logits);
-            token_ids.push(next);
-
             let text = tokenizer.decode(&[next]);
             print!("{text}");
             io::stdout().flush().unwrap();
 
-            if tokenizer.eos_token_id() == Some(next) {
-                break;
-            }
+            if tokenizer.eos_token_id() == Some(next) { break; }
+
+            let logits = model.forward_with_cache(&[next], &mut cache);
+            next = sample_greedy(&logits);
         }
         println!();
     }