Gahow Wang
eff26a0898
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>
95 lines
3.3 KiB
Rust
95 lines
3.3 KiB
Rust
// M2a KV-cache decode engine — the token-identical correctness gate.
|
|
//
|
|
// The centerpiece M2 invariant: greedy decode through the KV-cache incremental
|
|
// engine (`xtrain_model::generate_greedy_cached`) must be TOKEN-IDENTICAL to the
|
|
// naive full-recompute greedy (`xtrain_train::sample::generate` at temperature 0),
|
|
// which re-runs the whole forward over the growing prefix each step. Same tokens ⇒
|
|
// the cache + decode-time attention + RoPE-at-position reproduce the full forward.
|
|
//
|
|
// Numerics note: a randomly-initialised model has near-uniform logits, so argmax
|
|
// can be fragile to ~1e-6 differences. This unit gate therefore runs in F32 (the
|
|
// tightest path, and the dtype the eval harness actually uses) on a small model.
|
|
// The headline gate on the trained v12 checkpoint (peaked logits → robust argmax)
|
|
// is run on the GPU box and recorded in docs/18.
|
|
#![cfg(not(no_cuda))]
|
|
|
|
use xtrain_cuda::device;
|
|
use xtrain_model::{Config, TinyTransformer, generate_greedy_cached};
|
|
use xtrain_tensor::{DType, Device};
|
|
|
|
fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
|
|
let mut state = seed
|
|
.wrapping_mul(2862933555777941757)
|
|
.wrapping_add(3037000493);
|
|
(0..n)
|
|
.map(|_| {
|
|
state = state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn build(cfg: Config, device: Device, dtype: DType) -> TinyTransformer {
|
|
let mut seed = 1u64;
|
|
let m = TinyTransformer::new(cfg, device, |shape| {
|
|
seed = seed.wrapping_add(1);
|
|
let n: usize = shape.iter().product();
|
|
if shape.len() == 1 {
|
|
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
|
|
} else {
|
|
fill(n, seed, 0.08)
|
|
}
|
|
});
|
|
m.with_compute_dtype(dtype)
|
|
}
|
|
|
|
// A real GQA config (8 query / 2 kv heads → group 4) to exercise repeat_kv in the
|
|
// decode path; head_dim 16, dim 128, 4 layers.
|
|
fn gqa_cfg() -> Config {
|
|
Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2)
|
|
}
|
|
|
|
#[test]
|
|
fn kv_cache_decode_is_token_identical_to_naive_f32() {
|
|
assert!(
|
|
device::device_count().expect("device count") > 0,
|
|
"no CUDA device"
|
|
);
|
|
device::set_device(0).unwrap();
|
|
let device = Device::Cuda(0);
|
|
|
|
let model = build(gqa_cfg(), device, DType::F32);
|
|
let prompt: Vec<i32> = vec![1, 5, 9, 13, 2, 7];
|
|
let max_new = 24usize;
|
|
|
|
let mut rng = 7u64;
|
|
let naive = xtrain_train::sample::generate(&model, device, &prompt, max_new, 0.0, &mut rng);
|
|
let cached = generate_greedy_cached(&model, device, &prompt, max_new);
|
|
|
|
assert_eq!(
|
|
naive.len(),
|
|
cached.len(),
|
|
"length mismatch: naive {} vs cached {}",
|
|
naive.len(),
|
|
cached.len()
|
|
);
|
|
if naive != cached {
|
|
// Report the first divergence for debugging.
|
|
let first = naive
|
|
.iter()
|
|
.zip(&cached)
|
|
.position(|(a, b)| a != b)
|
|
.unwrap();
|
|
panic!(
|
|
"token divergence at index {first}: naive={:?} cached={:?}\nnaive ={naive:?}\ncached ={cached:?}",
|
|
naive[first], cached[first]
|
|
);
|
|
}
|
|
println!(
|
|
"KV-cache decode token-identical to naive over {} generated tokens (F32, GQA 8/2)",
|
|
max_new
|
|
);
|
|
}
|