Skip to main content

KvEncoder

cram_codec::kvcodec

Struct KvEncoder 

Source 
pub struct KvEncoder {
    pub params: KvParams,
    /* private fields */
}
Expand description

Append-friendly encoder for one layer’s K or V stream.

Fields§

§params: KvParams

Settings (fixed for the stream).

Implementations§

Source§

impl KvEncoder

Source

pub fn new(params: KvParams, prefill: &[f32], n: usize) -> Self

Start a stream from the prefill [heads][n][hd]; its statistics fix the header.

Examples found in repository?
examples/kvc.rs (line 96)
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 append(&mut self, x: &[f32], n: usize)

Append n new tokens [heads][n][hd]. Sealed blocks are never touched.

Examples found in repository?
examples/kvc.rs (line 99)
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 tokens(&self) -> usize

Tokens seen so far.

Source

pub fn sealed_blocks(&self) -> usize

Number of sealed (final) blocks.

Source

pub fn to_bytes(&self) -> Vec<u8> ⓘ

Serialize: header, sealed blocks, and the tail coded as a short block. Only the tail (< block tokens) is coded here.

Examples found in repository?
examples/kvc.rs (line 102)
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}

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> ByRef<T> for T

Source§

fn by_ref(&self) -> &T

Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Imply<T> for U
where T: ?Sized, U: ?Sized,

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.