post-train: M2a — KV-cache incremental decode engine (token-identical)

Single-sequence KV-cache decode (xtrain-model/src/decode.rs): per-layer K/V
cache + single-token incremental forward (prefill = first prompt.len() decode
steps, one code path). Mirrors model::block_forward at the raw-Tensor level (no
autograd tape — inference needs no grads), using rope_at + decode_attention.
Cache is host-accumulated token-major f32, rebuilt per step (the honest M2a
baseline; M2b moves it device-side + batched ragged).

Gate (the M2 centerpiece): KV-cache greedy decode is TOKEN-IDENTICAL to the
naive full-recompute greedy — tests/decode_kv.rs (small GQA model, F32, 24
tokens) and corroborated on the v12 1.05B SFT checkpoint (cached eval =
naive eval byte-for-byte: format 100/100, correct 8/100).

eval_arith --cached A/Bs the two paths + reports decode tok/s. Measured on v12
(1.05B, batch 1, F32): the cache win is sequence-length-dependent —
  max_new=32   naive 108 vs cached 111 tok/s  (~1.0x; overhead-bound)
  max_new=128  naive  69 vs cached 133 tok/s  (~1.9x)
  max_new=256  naive OOM     vs cached 129 tok/s
Cached throughput stays ~constant (O(1)/token) while naive decays (O(t)/token,
O(seq^2) graph → OOM at length). Short eval prompts are overhead-bound, so the
cache matters for long rollouts (DPO/GRPO), not the arithmetic eval itself.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-train/src/bin/eval_arith.rs

@@ -103,6 +103,10 @@ fn main() {
     let n_show = flag(&args, "--show", 8usize);
     let prompts_file = flag_value(&args, "--prompts-file").expect("--prompts-file is required");
     let gold_file = flag_value(&args, "--gold-file").expect("--gold-file is required");
+    // M2: decode through the KV-cache incremental engine instead of the naive
+    // full-recompute sampler. Token-identical to the naive path (gated by
+    // tests/decode_kv.rs); this flag also lets us A/B the two for the speedup.
+    let use_cached = args.iter().any(|a| a == "--cached");
 
     // Prompts: skip the `#` header / blank lines and decode escaped newlines so the
     // count and order line up with the gold file.
@@ -148,18 +152,26 @@ fn main() {
     xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
 
     println!(
-        "eval_arith: ckpt {} | {} prompts | max_new {}",
+        "eval_arith: ckpt {} | {} prompts | max_new {} | decode={}",
         ckpt.display(),
         prompts.len(),
-        max_new
+        max_new,
+        if use_cached { "kv-cache" } else { "naive" }
     );
 
     let (mut n_boxed, mut n_correct) = (0usize, 0usize);
     let mut shown = 0usize;
+    let mut gen_tokens = 0usize;
+    let t0 = std::time::Instant::now();
     for (prompt, &gold) in prompts.iter().zip(&golds) {
         let ids: Vec<i32> = tok.encode(prompt).into_iter().map(|t| t as i32).collect();
-        let mut rng = 7u64;
-        let out = generate(&model, device, &ids, max_new, 0.0, &mut rng);
+        let out = if use_cached {
+            xtrain_model::generate_greedy_cached(&model, device, &ids, max_new)
+        } else {
+            let mut rng = 7u64;
+            generate(&model, device, &ids, max_new, 0.0, &mut rng)
+        };
+        gen_tokens += out.len() - ids.len();
         let cont = tok.decode(&out[ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
         let seg = first_answer_segment(&cont);
         if parse_boxed_answer(seg).is_some() {
@@ -176,6 +188,7 @@ fn main() {
         }
     }
 
+    let elapsed = t0.elapsed().as_secs_f64();
     let n = prompts.len() as f64;
     println!(
         "RESULT format(boxed)={}/{} ({:.1}%) | correct={}/{} ({:.1}%)",
@@ -186,4 +199,11 @@ fn main() {
         prompts.len(),
         100.0 * n_correct as f64 / n,
     );
+    println!(
+        "TIMING decode={} | {:.2}s | {} gen tokens | {:.1} tok/s",
+        if use_cached { "kv-cache" } else { "naive" },
+        elapsed,
+        gen_tokens,
+        gen_tokens as f64 / elapsed,
+    );
 }