feat: ship 0.2 packed API

AGENTS.md

@@ -1,3 +1,5 @@
+<!-- Humans only — do not edit this file. -->
+
 # quantize
 
 A simple Rust library for learning and experimenting with quantization techniques.

Cargo.toml

@@ -1,7 +1,8 @@
 [package]
 name = "quantize"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
+rust-version = "1.87"
 description = "A simple quantization library."
 license = "MIT"
 repository = "https://github.com/aksheyd/quantize"
@@ -21,13 +22,20 @@ exclude = [
 all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
+[features]
+default = ["std"]
+std = []
+
 [dependencies]
 half = "2"
 
 [dev-dependencies]
 candle-core = "0.9"
 rand = "0.8"
 
+[lints.clippy]
+all = { level = "deny", priority = 10 }
+
 [[example]]
 name = "ch01_simple"
 path = "chapters/ch01_simple.rs"
@@ -48,6 +56,14 @@ path = "chapters/ch04_block.rs"
 name = "ch05_asymmetric"
 path = "chapters/ch05_asymmetric.rs"
 
+[[example]]
+name = "ch06_adaptive"
+path = "chapters/ch06_adaptive.rs"
+
+[[example]]
+name = "ch07_learned"
+path = "chapters/ch07_learned.rs"
+
 [[example]]
 name = "compare"
 path = "benchmarks/compare.rs"
@@ -56,3 +72,6 @@ path = "benchmarks/compare.rs"
 name = "update_readme"
 path = "benchmarks/update_readme.rs"
 
+[[example]]
+name = "throughput"
+path = "benchmarks/throughput.rs"

Justfile

@@ -2,7 +2,13 @@ c:
     cargo run --release --example compare
 
 l:
-    cargo fmt && cargo clippy -- -D warnings
+    cargo fmt && cargo clippy --all-targets -- -D warnings
+
+t:
+    cargo test
+
+b:
+    cargo run --release --example throughput
 
 ur:
-    cargo run --release --example update_readme
+    cargo run --release --example update_readme

README.md

@@ -5,32 +5,56 @@ A simple quantization library
 ## use as a library
 
 ```rust
-use quantize::{quantize, dequantize};
+use quantize::{adaptive, asymmetric, quantize, Scheme};
 
 let weights = [0.42_f32, -0.10, 0.70, -0.50];
-let (scales, codes) = quantize::<f32, 8, 32>(&weights);
-let back = dequantize::<_, 32>(&scales, &codes);
+
+let q = quantize::<f32, 8, 32>(&weights).unwrap();
+let _ = asymmetric::quantize::<f32, 8, 32>(&weights).unwrap();
+let _ = adaptive::quantize::<f32, 32>(&weights, 0.001).unwrap();
+let _ = Scheme::Q4_32.quantize::<f32>(&weights).unwrap();
+
+let back = q.dequantize();
+let _ = q.dot(&weights);
 ```
 
 ---
 
 ### comparison
 
+1024x1024 matrix, 50 iterations.
+
+#### quality
+
 ```
 cargo run --release --example compare
 ```
 
-<!-- comparison:start -->
+Reconstruct, then matmul.
 
-| method | bits/elt | mse (mean) | cosine (mean) |
-| --- | ---: | ---: | ---: |
-| quantize 4b×32 | 4.5 | 0.008339 | 0.995434 |
-| candle Q4_0 | 4.5 | 0.007543 | 0.995825 |
-| quantize 5b×32 | 5.5 | 0.001808 | 0.998995 |
-| candle Q5_0 | 5.5 | 0.001730 | 0.999037 |
-| quantize 8b×32 | 8.5 | 0.000025 | 0.999986 |
-| candle Q8_0 | 8.5 | 0.000025 | 0.999986 |
+<!-- comparison:start -->
 
-_matrix size: 128x128, runs: 10_
+| bits/value | quantize mse | candle mse |
+| ---: | ---: | ---: |
+| 4.5 | 0.066669 | 0.060276 |
+| 5.5 | 0.014429 | 0.013859 |
+| 8.5 | 0.000201 | 0.000201 |
 
 <!-- comparison:end -->
+
+#### speed
+
+```
+cargo run --release --example throughput
+```
+
+Pack and unpack. Apple M4.
+
+| kernel | quantize ns/value | candle ns/value |
+| --- | ---: | ---: |
+| 4-bit quant | 0.29 | 0.46 |
+| 8-bit quant | 0.24 | 0.30 |
+| 4-bit dequant | 0.09 | 0.29 |
+| 8-bit dequant | 0.07 | 0.26 |
+
+_ns = nanosecond, a billionth of a second. smaller is faster._

benchmarks/compare.rs

@@ -1,12 +1,12 @@
 //! Quantization comparison playground.
 //!
 //! Runs `RUNS` fresh matmuls for every method registered in
-//! `harness/methods.rs` and prints MSE / cosine as mean ± std.
+//! `harness/methods.rs` and prints mean MSE.
 
 mod harness;
 use harness::{Comparison, Harness};
 
-const MATRIX_SIZE: usize = 128;
+const MATRIX_SIZE: usize = 1024;
 const RUNS: usize = 50;
 
 fn main() -> candle_core::Result<()> {
@@ -16,26 +16,11 @@ fn main() -> candle_core::Result<()> {
 }
 
 fn print_report(r: &Comparison) {
-    println!("matrix_size = {}, runs = {}\n", r.matrix_size, r.runs);
-    println!(
-        "{:<16}{:>10}{:>22}{:>22}",
-        "method", "bits/elt", "mse (mean ± std)", "cosine (mean ± std)",
-    );
-    println!("{:-<16}{:->10}{:->22}{:->22}", "", "", "", "");
-
-    let mut prev_bits = 0.0_f32;
+    println!("matrix_size = {MATRIX_SIZE}, runs = {RUNS}\n");
+    println!("{:<12}{:>14}", "bits/value", "mse");
+    println!("{:-<12}{:->14}", "", "");
     for m in &r.methods {
-        // Blank line between precision tiers (>1 bit jump) for easier scanning.
-        if prev_bits > 0.0 && m.bits_per_element - prev_bits > 1.0 {
-            println!();
-        }
-        prev_bits = m.bits_per_element;
-        let mse = format!("{:.6} ± {:.6}", m.stats.mse_mean, m.stats.mse_std);
-        let cos = format!("{:.6} ± {:.6}", m.stats.cosine_mean, m.stats.cosine_std);
-        println!(
-            "{:<16}{:>10.2}{:>22}{:>22}",
-            m.name, m.bits_per_element, mse, cos
-        );
+        println!("{:<12.1}{:>14.6}", m.bits_per_element, m.mse);
     }
-    println!("\nbits/elt = total storage per element (data + amortized scale).");
+    println!("\nbits/value = storage for one number, including its scale.");
 }

benchmarks/harness/methods.rs

@@ -29,7 +29,6 @@ impl Bits {
 }
 
 pub(super) struct Method {
-    pub name: &'static str,
     pub bits_per_element: Bits,
     pub eval: EvalFn,
 }
@@ -39,11 +38,11 @@ pub(super) fn methods() -> Vec<Method> {
     use Bits::*;
     use GgmlDType::{Q4_0, Q5_0, Q8_0};
     vec![
-        Method { name: "quantize 4b×32",  bits_per_element: Quantize(4),      eval: eval_quantize::<4> },
-        Method { name: "candle Q4_0",     bits_per_element: Candle(Q4_0),     eval: eval_q4_0       },
-        Method { name: "quantize 5b×32",  bits_per_element: Quantize(5),      eval: eval_quantize::<5> },
-        Method { name: "candle Q5_0",     bits_per_element: Candle(Q5_0),     eval: eval_q5_0       },
-        Method { name: "quantize 8b×32",  bits_per_element: Quantize(8),      eval: eval_quantize::<8> },
-        Method { name: "candle Q8_0",     bits_per_element: Candle(Q8_0),     eval: eval_q8_0       },
+        Method { bits_per_element: Quantize(4), eval: eval_quantize::<4> },
+        Method { bits_per_element: Candle(Q4_0), eval: eval_q4_0 },
+        Method { bits_per_element: Quantize(5), eval: eval_quantize::<5> },
+        Method { bits_per_element: Candle(Q5_0), eval: eval_q5_0 },
+        Method { bits_per_element: Quantize(8), eval: eval_quantize::<8> },
+        Method { bits_per_element: Candle(Q8_0), eval: eval_q8_0 },
     ]
 }

benchmarks/harness/metrics.rs

@@ -1,8 +1,4 @@
 //! Error metrics comparing a predicted vector against the expected one.
-//! The smaller the error, the better the prediction.
-//!
-//! predicted refers to the output of the model,
-//! while expected is the ground truth.
 
 pub(super) fn mse(predicted: &[f32], expected: &[f32]) -> f32 {
     predicted
@@ -12,14 +8,3 @@ pub(super) fn mse(predicted: &[f32], expected: &[f32]) -> f32 {
         .sum::<f32>()
         / predicted.len() as f32
 }
-
-pub(super) fn cosine(predicted: &[f32], expected: &[f32]) -> f32 {
-    let dot: f32 = predicted.iter().zip(expected).map(|(p, e)| p * e).sum();
-    let norm_predicted: f32 = predicted.iter().map(|p| p * p).sum::<f32>().sqrt();
-    let norm_expected: f32 = expected.iter().map(|e| e * e).sum::<f32>().sqrt();
-    if norm_predicted == 0.0 || norm_expected == 0.0 {
-        0.0
-    } else {
-        dot / (norm_predicted * norm_expected)
-    }
-}

benchmarks/harness/mod.rs

@@ -1,33 +1,21 @@
 //! Test harness. `Harness::new(matrix_size, runs).run()` generates `runs`
-//! fresh random matrices, evaluates every method, returns mean ± std stats.
+//! fresh random matrices, evaluates every method, returns mean MSE.
 
 mod methods;
 mod metrics;
 mod new;
 mod quant;
 mod run;
 mod sample;
-mod stats;
 
 use candle_core::{Device, Tensor};
 
-#[allow(dead_code)]
-pub struct Stats {
-    pub mse_mean: f32,
-    pub mse_std: f32,
-    pub cosine_mean: f32,
-    pub cosine_std: f32,
-}
-
 pub struct MethodReport {
-    pub name: &'static str,
     pub bits_per_element: f32,
-    pub stats: Stats,
+    pub mse: f32,
 }
 
 pub struct Comparison {
-    pub matrix_size: usize,
-    pub runs: usize,
     pub methods: Vec<MethodReport>,
 }
 
@@ -37,8 +25,6 @@ pub struct Harness {
     device: Device,
 }
 
-// One run's worth of fresh data. Submodules can see the private fields
-// because Rust lets descendant modules access the parent's private items.
 struct Sample {
     matrix_size: usize,
     matrix_a: Vec<f32>,

benchmarks/harness/quant.rs

@@ -8,11 +8,12 @@ use candle_core::{
     Device, Result, Tensor,
 };
 use half::f16;
-use quantize::{dequantize, quantize};
+use quantize::quantize;
 
 fn roundtrip_block<const BITS: u32>(values: &[f32]) -> Vec<f32> {
-    let (scales, codes) = quantize::<f16, BITS, 32>(values);
-    dequantize::<_, 32>(&scales, &codes)
+    quantize::<f16, BITS, 32>(values)
+        .expect("valid bits/block")
+        .dequantize()
 }
 
 fn matmul(a: Vec<f32>, b: Vec<f32>, n: usize, d: &Device) -> Result<Vec<f32>> {

benchmarks/harness/run.rs

@@ -1,45 +1,35 @@
 //! Outer loop: for each of `runs` fresh samples, evaluate every method,
-//! accumulate (mse, cosine), then collapse each method's samples to mean ± std.
+//! then collapse each method's MSE samples to a mean.
 
-use super::{
-    methods::methods,
-    metrics::{cosine, mse},
-    stats::mean_std,
-    Comparison, Harness, MethodReport, Stats,
-};
+use super::{methods::methods, metrics::mse, Comparison, Harness, MethodReport};
 use candle_core::Result;
 
 impl Harness {
     pub fn run(&self) -> Result<Comparison> {
         let methods = methods();
         let elements = self.matrix_size * self.matrix_size;
         let mut mses: Vec<Vec<f32>> = vec![Vec::with_capacity(self.runs); methods.len()];
-        let mut coss: Vec<Vec<f32>> = vec![Vec::with_capacity(self.runs); methods.len()];
 
         for _ in 0..self.runs {
             let s = self.sample()?;
             for (i, m) in methods.iter().enumerate() {
                 let predicted = (m.eval)(&s, &self.device)?;
                 mses[i].push(mse(&predicted, &s.ground_truth));
-                coss[i].push(cosine(&predicted, &s.ground_truth));
             }
         }
 
-        #[rustfmt::skip]
-        let methods = methods.iter().enumerate().map(|(i, m)| {
-            let (mse_mean, mse_std) = mean_std(&mses[i]);
-            let (cosine_mean, cosine_std) = mean_std(&coss[i]);
-            MethodReport {
-                name: m.name,
-                bits_per_element: m.bits_per_element.evaluate(elements),
-                stats: Stats { mse_mean, mse_std, cosine_mean, cosine_std },
-            }
-        }).collect();
+        let methods = methods
+            .iter()
+            .enumerate()
+            .map(|(i, m)| {
+                let n = mses[i].len() as f32;
+                MethodReport {
+                    bits_per_element: m.bits_per_element.evaluate(elements),
+                    mse: mses[i].iter().sum::<f32>() / n,
+                }
+            })
+            .collect();
 
-        Ok(Comparison {
-            matrix_size: self.matrix_size,
-            runs: self.runs,
-            methods,
-        })
+        Ok(Comparison { methods })
     }
 }