Struct QuantizedMatrix

Source

pub struct QuantizedMatrix {
    pub data: Vec<i8>,
    pub scales: Vec<f32>,
    pub n: usize,
    pub k: usize,
}

Expand description

A weight matrix quantized with per-output-channel symmetric Q8 scales.

Layout is the nn.Linear layout the checkpoint stores: data is [n, k] row-major with one f32 scale per output row. Quantized once at hydration; the borrowed f32 tensor is untouched.

Fields§

§data: Vec<i8>

Q8 bytes, [n, k] row-major, each value in [-127, 127].

§scales: Vec<f32>

One symmetric scale per output row, [n].

§n: usize

Output rows.

§k: usize

Reduction length of one output element.

Implementations§

Source §

impl QuantizedMatrix

Source

pub fn concat_rows(parts: &[&Self]) -> Self

Stacks matrices with a shared reduction length into one taller matrix.

Row bytes and scales are byte-identical to quantizing each part separately — this exists so fused projections (QKV, gate‖up) can run as ONE kernel dispatch while every output row keeps exactly the per-channel quantization it would have had alone.

§Panics

Panics if the parts disagree on k or the list is empty.

Source

pub fn quantize(weight: &[f32], n: usize, k: usize) -> Self

Quantizes an [n, k] f32 weight matrix one output channel at a time.

§Panics

Panics if weight.len() != n * k or any value is non-finite.

Examples found in repository ?

examples/int8_shape_bench.rs (line 84)

62fn main() {
63    let tiers: Vec<Int8Tier> = Int8Tier::available();
64    println!("int8 shape bench — tiers available: {:?}", {
65        tiers.iter().map(|t| t.as_str()).collect::<Vec<_>>()
66    });
67    println!(
68        "interleaved rounds={ROUNDS} (+{WARMUP_ROUNDS} warmup); per-sample = mean over calls in one round; cv% over rounds"
69    );
70
71    for &m in &[1_usize, 16] {
72        println!(
73            "\n== m = {m} {} ==",
74            if m == 1 {
75                "(decode GEMV)"
76            } else {
77                "(seq-16 verify GEMM)"
78            }
79        );
80        for &(label, n, k) in SHAPES {
81            let calls: usize = (32 / m).max(2);
82            let weight = pseudo_random_f32(n * k, 0xbe0_0001 ^ (n as u64) << 20 ^ k as u64);
83            let x = pseudo_random_f32(m * k, 0xbe0_0002 ^ (m as u64) << 32 ^ k as u64);
84            let quantized = QuantizedMatrix::quantize(&weight, n, k);
85            let mut out = vec![0.0_f32; m * n];
86            let mut x_q = vec![0_i8; m * k];
87            let mut x_scales = vec![0.0_f32; m];
88
89            // One arm per route, all interleaved inside every round.
90            let mut f32_samples = Vec::with_capacity(ROUNDS);
91            let mut tier_samples: Vec<Vec<f64>> =
92                tiers.iter().map(|_| Vec::with_capacity(ROUNDS)).collect();
93
94            for round in 0..ROUNDS + WARMUP_ROUNDS {
95                // f32 arm
96                let start = Instant::now();
97                for _ in 0..calls {
98                    f32ref::linear(
99                        black_box(&x),
100                        black_box(&weight),
101                        None,
102                        m,
103                        k,
104                        n,
105                        black_box(&mut out),
106                    );
107                }
108                let f32_us = start.elapsed().as_secs_f64() * 1e6 / calls as f64;
109
110                // W8A8 arms, including dynamic activation quantization each call.
111                let mut this_round = Vec::with_capacity(tiers.len());
112                for &tier in &tiers {
113                    let start = Instant::now();
114                    for _ in 0..calls {
115                        for ((x_row, q_row), scale) in x
116                            .chunks_exact(k)
117                            .zip(x_q.chunks_exact_mut(k))
118                            .zip(x_scales.iter_mut())
119                        {
120                            *scale = quantize_row_q8(black_box(x_row), q_row);
121                        }
122                        linear_q8(
123                            black_box(&x_q),
124                            black_box(&x_scales),
125                            black_box(&quantized),
126                            None,
127                            m,
128                            black_box(&mut out),
129                            tier,
130                        );
131                    }
132                    this_round.push(start.elapsed().as_secs_f64() * 1e6 / calls as f64);
133                }
134
135                if round >= WARMUP_ROUNDS {
136                    f32_samples.push(f32_us);
137                    for (samples, sample) in tier_samples.iter_mut().zip(&this_round) {
138                        samples.push(*sample);
139                    }
140                }
141            }
142
143            let f32_stats = stats(&f32_samples);
144            // The f32 reference loops rows outermost, so it streams the weight matrix once per
145            // activation row (m times per call); the q8 kernel is weight-stationary and streams
146            // it exactly once per call. The column reports actual weight bytes moved per second.
147            let f32_bytes = (n * k * 4 * m) as f64;
148            println!(
149                "{label}  f32     {:9.1} us  cv {:4.1}%  ({:5.1} GB/s weight-stream)",
150                f32_stats.mean_us,
151                f32_stats.cv_percent,
152                f32_bytes / (f32_stats.mean_us * 1e-6) / 1e9,
153            );
154            for (tier, samples) in tiers.iter().zip(&tier_samples) {
155                let tier_stats = stats(samples);
156                let q8_bytes = (n * k) as f64;
157                let verdict = if tier_stats.cv_percent > 5.0 || f32_stats.cv_percent > 5.0 {
158                    "REFUSED (cv>5%)"
159                } else {
160                    ""
161                };
162                println!(
163                    "{label}  q8 {:9} {:9.1} us  cv {:4.1}%  ({:5.1} GB/s weight-stream)  x{:.2} vs f32 {verdict}",
164                    tier.as_str(),
165                    tier_stats.mean_us,
166                    tier_stats.cv_percent,
167                    q8_bytes / (tier_stats.mean_us * 1e-6) / 1e9,
168                    f32_stats.mean_us / tier_stats.mean_us,
169                );
170            }
171        }
172    }
173    println!(
174        "\nNOTE: ratios above compare routes inside this tree (self-comparison = maintenance),\nnever a pinned incumbent. cv%>5 rows are refused, not averaged."
175    );
176}