kvc/
kvc.rs

1//! KV-cache codec CLI (see `cram_codec::kvcodec`).
2//!
3//!     kvc encode IN.f32 OUT --shape HxTxD --kind k|v (--rmse r | --maxabs e | --bpv b)
4//!         [--post] [--sink n] [--sinkdiv n] [--block n] [--noar] [--theta t] [--prefill n]
5//!     kvc decode IN OUT.f32 [--range t0:t1] [--reps n]
6//!
7//! `--prefill n` fixes the header from the first n tokens and appends the rest
8//! one token at a time (the streaming path); default is one shot.
9
10use std::time::Instant;
11
12use cram_codec::kvcodec::{step_for_bpv, Kind, KvEncoder, KvParams, KvStream};
13
14fn usage() -> ! {
15    eprintln!("usage: kvc encode IN OUT --shape HxTxD --kind k|v (--rmse r|--maxabs e|--bpv b) [--post] [--sink n] [--sinkdiv n] [--block n] [--noar] [--theta t] [--prefill n]\n       kvc decode IN OUT [--range t0:t1] [--reps n]");
16    std::process::exit(2)
17}
18
19fn flag(args: &[String], name: &str) -> Option<String> {
20    args.iter()
21        .position(|a| a == name)
22        .map(|i| args.get(i + 1).cloned().unwrap_or_else(|| usage()))
23}
24
25fn read_f32(path: &str) -> Vec<f32> {
26    std::fs::read(path)
27        .unwrap()
28        .as_chunks::<4>()
29        .0
30        .iter()
31        .map(|&c| f32::from_le_bytes(c))
32        .collect()
33}
34
35fn main() {
36    let args: Vec<String> = std::env::args().collect();
37    if args.len() < 4 {
38        usage()
39    }
40    match args[1].as_str() {
41        "encode" => {
42            let x = read_f32(&args[2]);
43            let shape: Vec<usize> = flag(&args, "--shape")
44                .unwrap_or_else(|| usage())
45                .split('x')
46                .map(|s| s.parse().unwrap())
47                .collect();
48            let (h, n, d) = (shape[0], shape[1], shape[2]);
49            let kind = match flag(&args, "--kind").as_deref() {
50                Some("k") => Kind::K,
51                Some("v") => Kind::V,
52                _ => usage(),
53            };
54            let mut p = KvParams::new(kind, h, d, 1.0);
55            p.pre_rope = !args.iter().any(|a| a == "--post");
56            p.ar = !args.iter().any(|a| a == "--noar");
57            p.gpu = args.iter().any(|a| a == "--gpu");
58            if let Some(v) = flag(&args, "--sink") {
59                p.sink = v.parse().unwrap();
60            }
61            if let Some(v) = flag(&args, "--sinkdiv") {
62                p.sink_div = v.parse().unwrap();
63            }
64            if let Some(v) = flag(&args, "--block") {
65                p.block = v.parse().unwrap();
66            }
67            if let Some(v) = flag(&args, "--theta") {
68                p.theta = v.parse().unwrap();
69            }
70            let t = Instant::now();
71            p.step = if let Some(r) = flag(&args, "--rmse") {
72                r.parse::<f64>().unwrap() * 12f64.sqrt()
73            } else if let Some(e) = flag(&args, "--maxabs") {
74                // pre-RoPE pair rotation can move a per-element error by sqrt(2)
75                let e: f64 = e.parse().unwrap();
76                if kind == Kind::K && p.pre_rope {
77                    e * std::f64::consts::SQRT_2
78                } else {
79                    2.0 * e
80                }
81            } else if let Some(b) = flag(&args, "--bpv") {
82                step_for_bpv(p.clone(), &x, n, b.parse().unwrap())
83            } else {
84                usage()
85            };
86            let pre: usize = flag(&args, "--prefill")
87                .map(|v| v.parse().unwrap())
88                .unwrap_or(n);
89            let slice = |t0: usize, t1: usize| -> Vec<f32> {
90                let mut v = Vec::with_capacity(h * (t1 - t0) * d);
91                for hh in 0..h {
92                    v.extend_from_slice(&x[(hh * n + t0) * d..(hh * n + t1) * d]);
93                }
94                v
95            };
96            let mut e = KvEncoder::new(p.clone(), &slice(0, pre), pre);
97            let ta = Instant::now();
98            for t in pre..n {
99                e.append(&slice(t, t + 1), 1);
100            }
101            let append_s = ta.elapsed().as_secs_f64();
102            let bytes = e.to_bytes();
103            let enc = t.elapsed().as_secs_f64();
104            std::fs::write(&args[3], &bytes).unwrap();
105            println!(
106                "{{\"bytes\":{},\"bpv\":{},\"step\":{},\"encode_s\":{enc},\"append_s\":{append_s},\"appended\":{}}}",
107                bytes.len(),
108                8.0 * bytes.len() as f64 / x.len() as f64,
109                p.step,
110                n - pre
111            );
112        }
113        "decode" => {
114            let bytes = std::fs::read(&args[2]).unwrap();
115            let s = KvStream::parse(&bytes).unwrap();
116            let (t0, t1) = match flag(&args, "--range") {
117                Some(r) => {
118                    let (a, b) = r.split_once(':').unwrap();
119                    (a.parse().unwrap(), b.parse().unwrap())
120                }
121                None => (0, s.tokens),
122            };
123            let reps: usize = flag(&args, "--reps")
124                .map(|v| v.parse().unwrap())
125                .unwrap_or(1);
126            let t = Instant::now();
127            let mut y = Vec::new();
128            for _ in 0..reps {
129                y = s.decode_range(t0, t1);
130            }
131            let dec = t.elapsed().as_secs_f64() / reps as f64;
132            let out: Vec<u8> = y.iter().flat_map(|v| v.to_le_bytes()).collect();
133            std::fs::write(&args[3], out).unwrap();
134            println!(
135                "{{\"decode_s\":{dec},\"ns_per_value\":{}}}",
136                1e9 * dec / y.len().max(1) as f64
137            );
138        }
139        _ => usage(),
140    }
141}
kvc/kvc.rs

kvc/
kvc.rs
