pub fn linear(
x: &[f32],
weight: &[f32],
bias: Option<&[f32]>,
m: usize,
k: usize,
n: usize,
out: &mut [f32],
)Expand description
Row-major matrix-vector/matrix-matrix product in the layout PyTorch Linear stores.
x is [m, k], weight is [n, k] (out-features major, as nn.Linear stores it), and the
result is [m, n]. Bias is optional because every attention/MLP projection in this model is
bias-free; only text_projection carries one.
§Panics
Panics if any slice length disagrees with m, k, n.
Examples found in repository?
examples/int8_shape_bench.rs (lines 98-106)
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}