Struct KvStream

Source

pub struct KvStream<'a> {
    pub params: KvParams,
    pub tokens: usize,
    /* private fields */
}

Expand description

Parsed stream: header plus block payload slices.

Fields§

§params: KvParams

Settings recovered from the header.

§tokens: usize

Total tokens.

Implementations§

Source §

impl<'a> KvStream<'a>

Source

pub fn parse(bytes: &'a [u8]) -> Result<Self, DecodeError>

Parse a stream produced by KvEncoder::to_bytes.

Examples found in repository ?

examples/kvc.rs (line 115)

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}

Source

pub fn decode_range(&self, t0: usize, t1: usize) -> Vec<f32>

Decode tokens [t0, t1) as [heads][t1-t0][hd], decoding only the blocks that cover the range.

Examples found in repository ?

examples/kvc.rs (line 129)

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}

Source

pub fn gpu_tables(&self) -> Option<GpuTables>

GPU layout tables for crate::kvgpu: per channel mu (i32 bits), rho, class, and per (block, head) lane (byte offset, byte length, first token, tokens).