ftts_kernels/int8.rs
1//! Int8 W8A8 kernels: symmetric per-output-channel Q8 weights times per-row Q8 activations.
2//!
3//! This is the Phase-2/3A quantized projection route for the talker and microdecoder GEMMs.
4//! The numeric contract is S8S8: weights quantized by the canonical symmetric recipe
5//! (`scale = max|row| / 127`, ties-to-even, `[-127, 127]`, `-128` never emitted — identical to
6//! `ftts-artifacts::converter::quantize_output_channel_q8`, byte-for-byte, asserted by a
7//! cross-crate test in `ftts-model-qwen`), activations quantized dynamically per row with the
8//! same recipe. Accumulation is exact i32; the two f32 scales are applied once, after
9//! accumulation, in a fixed multiplication order shared by every tier.
10//!
11//! ## Tier law
12//!
13//! Every tier of [`dot_i32`] is *exactly equal in i32* to [`Int8Tier::Scalar`] on every input —
14//! integer addition is associative, and the overflow selftest proves the all-extreme reduction
15//! fits i32 at every census binding K. A tier is only dispatchable after
16//! [`crate::selftest::run_selftest`] has executed its all-extreme proof rows through the real
17//! kernel function on the running silicon. Do not add a tier here without extending the selftest.
18//!
19//! Inherited prior NE-INH-003 (re-verify per toolchain): on Apple M4, LLVM autovectorization of
20//! the scalar shape beat a hand SDOT micro-tile at m=1. Both routes therefore ship; dispatch
21//! preference is decided by measurement (`FTTS_INT8_TIER` forces a route for A/B), never by
22//! assumption.
23
24use std::sync::OnceLock;
25
26/// Largest absolute Q8 byte the canonical symmetric recipe emits.
27pub const Q8_MAX_ABS: i8 = 127;
28
29/// Weight bytes below which a linear stays on the calling thread even when a team exists.
30///
31/// Every talker/microdecoder projection (2-12 MB) and the codec's ConvNeXt projections clear
32/// this; genuinely small ops don't repay the dispatch handshake.
33const TEAM_WORK_THRESHOLD_BYTES: usize = 512 * 1024;
34
35/// Quantizes one row (weight output channel or activation row) with the canonical symmetric
36/// Q8 recipe.
37///
38/// The returned scale is `max(abs(row)) / 127`; all-zero rows use the explicit scale `1.0` and
39/// emit zero bytes. Values are clamped to `[-127, 127]` and rounded ties-to-even; `-128` is never
40/// emitted. This is the same arithmetic as the offline converter's
41/// `quantize_output_channel_q8`, restated here because the artifact crate depends on this one.
42///
43/// # Panics
44///
45/// Panics if `output.len() != row.len()` or a value is non-finite. A NaN/inf activation reaching
46/// the quantizer means the f32 graph upstream is already corrupt; refusing loudly here beats
47/// synthesizing garbage audio quietly.
48pub fn quantize_row_q8(row: &[f32], output: &mut [i8]) -> f32 {
49 assert_eq!(output.len(), row.len(), "quantize output length mismatch");
50 let mut maximum = 0.0_f32;
51 for (index, &value) in row.iter().enumerate() {
52 assert!(
53 value.is_finite(),
54 "non-finite value {value} at index {index} reached the Q8 quantizer"
55 );
56 maximum = maximum.max(value.abs());
57 }
58 if maximum == 0.0 {
59 output.fill(0);
60 return 1.0;
61 }
62 let scale = maximum / 127.0;
63 if scale == 0.0 {
64 // A subnormal maximum can flush the division to zero; value/scale would then be inf or
65 // NaN. A row this close to zero rounds to the zero row it effectively is.
66 output.fill(0);
67 return 1.0;
68 }
69 for (&value, slot) in row.iter().zip(output.iter_mut()) {
70 let rounded = (value / scale).clamp(-127.0, 127.0).round_ties_even();
71 // The clamp bounds the conversion inside i8, and the symmetric contract additionally
72 // excludes the otherwise-representable -128.
73 *slot = rounded as i8;
74 }
75 scale
76}
77
78/// A weight matrix quantized with per-output-channel symmetric Q8 scales.
79///
80/// Layout is the `nn.Linear` layout the checkpoint stores: `data` is `[n, k]` row-major with one
81/// f32 scale per output row. Quantized once at hydration; the borrowed f32 tensor is untouched.
82#[derive(Clone, Debug)]
83pub struct QuantizedMatrix {
84 /// Q8 bytes, `[n, k]` row-major, each value in `[-127, 127]`.
85 pub data: Vec<i8>,
86 /// One symmetric scale per output row, `[n]`.
87 pub scales: Vec<f32>,
88 /// Output rows.
89 pub n: usize,
90 /// Reduction length of one output element.
91 pub k: usize,
92}
93
94impl QuantizedMatrix {
95 /// Stacks matrices with a shared reduction length into one taller matrix.
96 ///
97 /// Row bytes and scales are byte-identical to quantizing each part separately — this exists
98 /// so fused projections (QKV, gate‖up) can run as ONE kernel dispatch while every output
99 /// row keeps exactly the per-channel quantization it would have had alone.
100 ///
101 /// # Panics
102 ///
103 /// Panics if the parts disagree on `k` or the list is empty.
104 #[must_use]
105 pub fn concat_rows(parts: &[&Self]) -> Self {
106 let k = parts.first().expect("at least one part").k;
107 assert!(parts.iter().all(|part| part.k == k), "parts must share k");
108 let n = parts.iter().map(|part| part.n).sum();
109 let mut data = Vec::with_capacity(n * k);
110 let mut scales = Vec::with_capacity(n);
111 for part in parts {
112 data.extend_from_slice(&part.data);
113 scales.extend_from_slice(&part.scales);
114 }
115 Self { data, scales, n, k }
116 }
117
118 /// Quantizes an `[n, k]` f32 weight matrix one output channel at a time.
119 ///
120 /// # Panics
121 ///
122 /// Panics if `weight.len() != n * k` or any value is non-finite.
123 #[must_use]
124 pub fn quantize(weight: &[f32], n: usize, k: usize) -> Self {
125 assert_eq!(weight.len(), n * k, "weight must be [n, k]");
126 let mut data = vec![0_i8; n * k];
127 let mut scales = vec![0.0_f32; n];
128 for ((weight_row, data_row), scale) in weight
129 .chunks_exact(k)
130 .zip(data.chunks_exact_mut(k))
131 .zip(scales.iter_mut())
132 {
133 *scale = quantize_row_q8(weight_row, data_row);
134 }
135 Self { data, scales, n, k }
136 }
137}
138
139/// An executable int8 dot-product route.
140///
141/// Every variant is exactly equal in i32 to `Scalar` on every input. `NeonSdot` exists only on
142/// aarch64 builds with the `neon-dotprod` feature and is dispatchable only where the CPU reports
143/// FEAT_DotProd at runtime.
144#[derive(Clone, Copy, Debug, Eq, PartialEq)]
145pub enum Int8Tier {
146 /// Portable left-to-right checked-free scalar loop; the reference every tier must equal.
147 Scalar,
148 /// Portable eight-lane loop, retained ONLY as an A/B datapoint: measured ~15x SLOWER than
149 /// `Scalar` at m=1 on M4 Pro (NE-001) — the manual lane structure defeats LLVM's
150 /// autovectorizer, while the plain `Scalar` shape vectorizes to memory bandwidth. Never the
151 /// dispatch default.
152 Autovec,
153 /// Hand SDOT island (aarch64 + FEAT_DotProd), four 16-byte accumulator streams.
154 NeonSdot,
155 /// Hand SIMD128 island (wasm32 + `simd128`), four 16-byte accumulator streams.
156 ///
157 /// The browser's equivalent of [`Self::NeonSdot`]. Unlike aarch64, wasm has no int8 dot for
158 /// the autovectorizer to find, so without this tier a browser runs the byte-at-a-time
159 /// `Scalar` loop.
160 WasmSimd128,
161}
162
163impl Int8Tier {
164 /// Stable machine-readable route name.
165 #[must_use]
166 pub const fn as_str(self) -> &'static str {
167 match self {
168 Self::Scalar => "scalar",
169 Self::Autovec => "autovec",
170 Self::NeonSdot => "neon-sdot",
171 Self::WasmSimd128 => "wasm-simd128",
172 }
173 }
174
175 /// Every tier this build can execute on the running silicon, scalar first.
176 #[must_use]
177 pub fn available() -> Vec<Self> {
178 let mut tiers = vec![Self::Scalar, Self::Autovec];
179 if neon_sdot_available() {
180 tiers.push(Self::NeonSdot);
181 }
182 if wasm_simd128_available() {
183 tiers.push(Self::WasmSimd128);
184 }
185 tiers
186 }
187
188 /// The route the int8 path dispatches by default, honoring the `FTTS_INT8_TIER` override.
189 ///
190 /// The override exists for interleaved A/B measurement (`scalar` / `autovec` / `neon-sdot`);
191 /// an unavailable or unrecognized override falls back to the measured default rather than
192 /// panicking mid-synthesis. Until a per-shape KernelPlan lands, the default is `NeonSdot`
193 /// where FEAT_DotProd exists, else `Scalar`. Measured on M4 Pro (2026-08-08, shape bench,
194 /// noisy shared host, indicative): plain `Scalar` autovectorizes to ~50 GB/s and ties SDOT
195 /// at m=1 — NE-INH-003 reconfirmed — while the hand-shaped `Autovec` lane loop defeats the
196 /// vectorizer and loses ~15x; it stays only as an A/B datapoint.
197 #[must_use]
198 pub fn dispatch() -> Self {
199 // wasm32 first and without consulting the environment: there are no environment variables
200 // in a browser, and unlike aarch64 the fallback here is not a vectorized scalar loop but a
201 // byte-at-a-time one, so the island is the only route worth dispatching.
202 if wasm_simd128_available() {
203 return Self::WasmSimd128;
204 }
205 match std::env::var("FTTS_INT8_TIER").as_deref() {
206 Ok("scalar") => Self::Scalar,
207 Ok("autovec") => Self::Autovec,
208 Ok("neon-sdot") if neon_sdot_available() => Self::NeonSdot,
209 _ if neon_sdot_available() => Self::NeonSdot,
210 _ => Self::Scalar,
211 }
212 }
213}
214
215/// Which quantized linear op class the armed route runs.
216///
217/// `W8A8` quantizes activations per row and uses the exact-i32 int8 dot — fastest, but the
218/// activation rounding perturbs logits enough that seeded sampling can draw different tokens
219/// than f32. `W8A16` keeps activations f32 and dequantizes weights in-register — the same
220/// one-byte-per-weight memory traffic, no activation error, so the output tracks the f32
221/// reference much more closely. Its f32 accumulation is lane-ordered (not the reference's
222/// left-to-right order): this is a lossy route already, so reduction-order freedom is part of
223/// the deal, and the fidelity gate is measured downstream, not asserted bitwise.
224#[derive(Clone, Copy, Debug, Eq, PartialEq)]
225pub enum QuantLinearMode {
226 /// Int8 activations times int8 weights, exact i32 accumulation. Carries the full measured
227 /// plan so batched calls (prefill, the seq-16 verify) route to the batch-regime winner
228 /// instead of inheriting the GEMV winner; every tier is bit-identical per element, so the
229 /// split is purely a speed decision.
230 W8A8(KernelPlanV0),
231 /// f32 activations times dequantized int8 weights, lane-ordered f32 accumulation.
232 W8A16,
233}
234
235impl QuantLinearMode {
236 /// Stable machine-readable mode name.
237 #[must_use]
238 pub const fn as_str(self) -> &'static str {
239 match self {
240 Self::W8A8(_) => "w8a8",
241 Self::W8A16 => "w8a16",
242 }
243 }
244}
245
246/// W8A16 linear: f32 activations `[m, k]` times a [`QuantizedMatrix`] `[n, k]` producing
247/// f32 `[m, n]`.
248///
249/// Eight independent f32 FMA lanes per dot product, weights widened from i8 in-register; the
250/// per-output-channel scale multiplies once after accumulation, mirroring the W8A8 dequant
251/// order. Weight-stationary loop, like [`linear_q8`], and like it fanning large calls out
252/// across the persistent team.
253///
254/// # Panics
255///
256/// Panics on any shape mismatch.
257pub fn linear_w8a16(
258 x: &[f32],
259 weight: &QuantizedMatrix,
260 bias: Option<&[f32]>,
261 m: usize,
262 out: &mut [f32],
263) {
264 let (n, k) = (weight.n, weight.k);
265 assert_eq!(x.len(), m * k, "x must be [m, k]");
266 assert_eq!(out.len(), m * n, "out must be [m, n]");
267 if let Some(bias) = bias {
268 assert_eq!(bias.len(), n, "bias must be [n]");
269 }
270 // Same speed-only fan-out gate as `linear_q8`: every output element is still one
271 // `dot_w8a16` over the same span in the same order, so the partitioned result is
272 // bit-identical per element. Before this gate existed, FTTS_INT8=w8a16 single-cored the
273 // whole model while w8a8 ran on the team.
274 if n * k >= TEAM_WORK_THRESHOLD_BYTES
275 && !crate::team::thread_bypassed()
276 && let Some(team) = crate::team::armed()
277 {
278 team.linear_w8a16(x, weight, bias, m, out);
279 return;
280 }
281 for col in 0..n {
282 let w_row = &weight.data[col * k..(col + 1) * k];
283 let w_scale = weight.scales[col];
284 let bias_term = bias.map(|b| b[col]);
285 for row in 0..m {
286 let x_row = &x[row * k..(row + 1) * k];
287 let acc = dot_w8a16(x_row, w_row);
288 let value = acc * w_scale;
289 out[row * n + col] = bias_term.map_or(value, |b| value + b);
290 }
291 }
292}
293
294/// Eight-lane f32 dot of an f32 row against an i8 weight row, widened in-register.
295///
296/// `pub(crate)` so the team's column-partition worker runs the exact same reduction the
297/// serial loop runs — bit-identity between the two paths rests on sharing this function.
298pub(crate) fn dot_w8a16(x: &[f32], w: &[i8]) -> f32 {
299 const LANES: usize = 8;
300 let mut lanes = [0.0_f32; LANES];
301 let chunks = x.len() / LANES;
302 for chunk in 0..chunks {
303 let base = chunk * LANES;
304 for lane in 0..LANES {
305 lanes[lane] = f32::from(w[base + lane]).mul_add(x[base + lane], lanes[lane]);
306 }
307 }
308 let mut sum: f32 = lanes.iter().sum();
309 for index in chunks * LANES..x.len() {
310 sum = f32::from(w[index]).mul_add(x[index], sum);
311 }
312 sum
313}
314
315/// The armed quantized-linear mode for the talker/microdecoder route.
316///
317/// `FTTS_INT8=1` or `w8a8` selects the int8-dot route; `FTTS_INT8=w8a16` selects the
318/// weight-only route. Anything else means the caller should not be arming quantization at all
319/// (the kill-switch check happens before this is consulted).
320#[must_use]
321pub fn quant_mode_from_environment() -> QuantLinearMode {
322 match std::env::var("FTTS_INT8").as_deref() {
323 Ok("w8a16") => QuantLinearMode::W8A16,
324 _ => QuantLinearMode::W8A8(autotuned_plan()),
325 }
326}
327
328/// Runs one quantized linear in the selected mode; the drop-in used by the armed model paths.
329pub fn quant_linear(
330 mode: QuantLinearMode,
331 x: &[f32],
332 weight: &QuantizedMatrix,
333 bias: Option<&[f32]>,
334 m: usize,
335 out: &mut [f32],
336) {
337 match mode {
338 QuantLinearMode::W8A8(plan) => {
339 // The same m threshold the codec route uses: offline/prefill batches take the
340 // measured batch-regime winner, decode GEMVs the GEMV winner.
341 let tier = if m > 4 {
342 plan.batch_gemm
343 } else {
344 plan.decode_gemv
345 };
346 linear_q8_dynamic(x, weight, bias, m, out, tier);
347 }
348 QuantLinearMode::W8A16 => linear_w8a16(x, weight, bias, m, out),
349 }
350}
351
352/// The measured per-regime route assignment, decided once per process.
353///
354/// v0 of the KernelPlan: two regimes, no persistence (`.fttspack` owns that when it lands).
355/// Safe to decide by noisy measurement because every tier produces bit-identical output — a
356/// wrong pick costs microseconds, never correctness.
357#[derive(Clone, Copy, Debug, Eq, PartialEq)]
358pub struct KernelPlanV0 {
359 /// Route for m=1 decode GEMVs (the talker/microdecoder step shape).
360 pub decode_gemv: Int8Tier,
361 /// Route for batched GEMMs (prefill, seq-16 verify, offline codec).
362 pub batch_gemm: Int8Tier,
363}
364
365impl KernelPlanV0 {
366 /// Both regimes pinned to one tier — the `FTTS_INT8_TIER` A/B form and the shape tests use
367 /// to hold the route fixed.
368 #[must_use]
369 pub const fn pinned(tier: Int8Tier) -> Self {
370 Self {
371 decode_gemv: tier,
372 batch_gemm: tier,
373 }
374 }
375}
376
377/// Measures each available tier at the two live regimes and returns the winners.
378///
379/// Decided once per process and cached. `FTTS_INT8_TIER` overrides both regimes — the A/B
380/// override must pin the route it names, not merely suggest it. Cost: a few milliseconds of
381/// synthetic dots at the model's real reduction lengths.
382pub fn autotuned_plan() -> KernelPlanV0 {
383 static PLAN: OnceLock<KernelPlanV0> = OnceLock::new();
384 *PLAN.get_or_init(|| {
385 // wasm32 is pinned, never measured: `Instant::now` panics as `unreachable` there (no
386 // monotonic clock in std — exactly how the browser playground's first synthesize died),
387 // and there is nothing to choose between anyway. `dispatch()` names the SIMD128 island
388 // when it is compiled in, which it is for every browser build; an earlier revision of
389 // this pinned `Scalar` on the grounds that no other tier existed on wasm, and that
390 // sentence stopped being true the moment the island landed — leaving the fast kernel
391 // built, dispatchable, and never dispatched.
392 #[cfg(target_arch = "wasm32")]
393 {
394 KernelPlanV0::pinned(Int8Tier::dispatch())
395 }
396 #[cfg(not(target_arch = "wasm32"))]
397 {
398 if std::env::var("FTTS_INT8_TIER").is_ok() {
399 return KernelPlanV0::pinned(Int8Tier::dispatch());
400 }
401 if let Some(cached) = load_persisted_plan() {
402 return cached;
403 }
404 let plan = KernelPlanV0 {
405 // Talker/microdecoder decode: one activation row against tall matrices; K = 1024
406 // and 3072 are the real reduction lengths, 256 output rows keep the probe cheap
407 // while streaming enough weight bytes to reach the bandwidth regime.
408 decode_gemv: fastest_tier(&[(1, 1024, 256), (1, 3072, 256)]),
409 // Verify/prefill/codec batches: sixteen rows, same reduction lengths.
410 batch_gemm: fastest_tier(&[(16, 1024, 128), (16, 3072, 64)]),
411 };
412 persist_plan(plan);
413 plan
414 }
415 })
416}
417
418/// Where the measured plan is cached between runs: the pre-`.fttspack` v0 of the per-machine
419/// execution cache. Losing or corrupting this file only costs a re-measurement.
420fn plan_cache_path() -> Option<std::path::PathBuf> {
421 std::env::var_os("HOME")
422 .map(|home| std::path::PathBuf::from(home).join(".cache/franken_tts/kernel_plan_v0.txt"))
423}
424
425/// The cache key: anything here changing invalidates the measurement.
426fn plan_cache_key() -> String {
427 let tiers: Vec<&str> = Int8Tier::available().iter().map(|t| t.as_str()).collect();
428 // `tiers` is runtime ISA detection, so a CPU with different features already misses the
429 // cache; arch/OS and core count additionally invalidate a plan replayed on a same-ISA but
430 // different machine (an NFS or migrated `$HOME`), where the measured winner may differ by
431 // microarchitecture. Cost of a false miss is one ~second re-probe.
432 format!(
433 "v1|crate={}|arch={}-{}|cores={}|tiers={}",
434 env!("CARGO_PKG_VERSION"),
435 std::env::consts::ARCH,
436 std::env::consts::OS,
437 std::thread::available_parallelism().map_or(1, usize::from),
438 tiers.join(",")
439 )
440}
441
442fn load_persisted_plan() -> Option<KernelPlanV0> {
443 // A valid plan file is three short lines; reading it bounded keeps a corrupt or hostile
444 // multi-gigabyte file at this user-writable path from ballooning the process.
445 let text = {
446 use std::io::Read as _;
447 let mut text = String::new();
448 let file = std::fs::File::open(plan_cache_path()?).ok()?;
449 file.take(512).read_to_string(&mut text).ok()?;
450 text
451 };
452 let mut lines = text.lines();
453 if lines.next()? != plan_cache_key() {
454 return None;
455 }
456 let parse = |line: &str| match line {
457 "scalar" => Some(Int8Tier::Scalar),
458 "autovec" => Some(Int8Tier::Autovec),
459 "neon-sdot" if neon_sdot_available() => Some(Int8Tier::NeonSdot),
460 _ => None,
461 };
462 Some(KernelPlanV0 {
463 decode_gemv: parse(lines.next()?)?,
464 batch_gemm: parse(lines.next()?)?,
465 })
466}
467
468fn persist_plan(plan: KernelPlanV0) {
469 let Some(path) = plan_cache_path() else {
470 return;
471 };
472 if let Some(parent) = path.parent() {
473 let _ = std::fs::create_dir_all(parent);
474 }
475 // Best-effort: an unwritable cache directory must never fail synthesis.
476 let _ = std::fs::write(
477 path,
478 format!(
479 "{}\n{}\n{}\n",
480 plan_cache_key(),
481 plan.decode_gemv.as_str(),
482 plan.batch_gemm.as_str()
483 ),
484 );
485}
486
487/// Times every available tier over the given `(m, k, n)` probes; median of three rounds each,
488/// summed across probes, smallest total wins. Ties break toward the earlier tier in
489/// [`Int8Tier::available`] order (scalar first — the simpler route).
490fn fastest_tier(probes: &[(usize, usize, usize)]) -> Int8Tier {
491 use std::time::Instant;
492 let tiers = Int8Tier::available();
493 let mut best = (tiers[0], f64::MAX);
494 for &tier in &tiers {
495 let mut total = 0.0_f64;
496 for &(m, k, n) in probes {
497 // Shifted into [-127, 127]: `as i8` alone wraps 128..=254 to -128..=-2, and -128 is
498 // outside the pinned S8S8 contract this same file declares.
499 let x_q: Vec<i8> = (0..m * k)
500 .map(|i| (((i * 37 + 11) % 255) as i32 - 127) as i8)
501 .collect();
502 let x_scales = vec![1.0_f32; m];
503 let weight = QuantizedMatrix {
504 data: (0..n * k)
505 .map(|i| (((i * 29 + 5) % 255) as i32 - 127) as i8)
506 .collect(),
507 scales: vec![1.0_f32; n],
508 n,
509 k,
510 };
511 let mut out = vec![0.0_f32; m * n];
512 // Serial on purpose: the decode-GEMV probe crosses the team work threshold, so an
513 // un-bypassed call would lazily spawn the team mid-measurement and time dispatch
514 // overhead + contention instead of the tier. Tier choice is about the inner loop;
515 // the team partitions whatever tier wins identically.
516 let mut rounds: Vec<f64> = (0..3)
517 .map(|_| {
518 let start = Instant::now();
519 crate::team::with_team_bypassed(|| {
520 linear_q8(&x_q, &x_scales, &weight, None, m, &mut out, tier);
521 });
522 start.elapsed().as_secs_f64()
523 })
524 .collect();
525 rounds.sort_by(f64::total_cmp);
526 total += rounds[1];
527 }
528 if total < best.1 {
529 best = (tier, total);
530 }
531 }
532 best.0
533}
534
535/// Whether the SDOT island is compiled in and the CPU reports FEAT_DotProd.
536#[must_use]
537pub fn neon_sdot_available() -> bool {
538 #[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
539 {
540 neon_dotprod::available()
541 }
542 #[cfg(not(all(target_arch = "aarch64", feature = "neon-dotprod")))]
543 {
544 false
545 }
546}
547
548/// Whether the SIMD128 island is compiled in.
549///
550/// Compile-time only, deliberately: `simd128` is a wasm target feature, so a module built with it
551/// either instantiates on an engine that has it or is refused outright. There is no partial
552/// support to detect at runtime the way FEAT_DotProd must be.
553#[must_use]
554pub fn wasm_simd128_available() -> bool {
555 cfg!(all(target_arch = "wasm32", target_feature = "simd128"))
556}
557
558/// Exact i32 dot product of two Q8 rows over the selected route.
559///
560/// # Panics
561///
562/// Panics if the lengths differ, or if `NeonSdot` is requested where it is not executable.
563#[must_use]
564pub fn dot_i32(a: &[i8], b: &[i8], tier: Int8Tier) -> i32 {
565 assert_eq!(a.len(), b.len(), "int8 dot inputs must match");
566 match tier {
567 Int8Tier::Scalar => dot_i32_scalar(a, b),
568 Int8Tier::Autovec => dot_i32_autovec(a, b),
569 Int8Tier::NeonSdot => dot_i32_neon_or_panic(a, b),
570 Int8Tier::WasmSimd128 => dot_i32_wasm_or_panic(a, b),
571 }
572}
573
574#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
575fn dot_i32_wasm_or_panic(a: &[i8], b: &[i8]) -> i32 {
576 wasm_simd128::dot_i32(a, b)
577}
578
579#[cfg(not(all(target_arch = "wasm32", target_feature = "simd128")))]
580fn dot_i32_wasm_or_panic(_a: &[i8], _b: &[i8]) -> i32 {
581 panic!("wasm-simd128 route selected on a build without the island");
582}
583
584#[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
585fn dot_i32_neon_or_panic(a: &[i8], b: &[i8]) -> i32 {
586 assert!(
587 neon_dotprod::available(),
588 "neon-sdot route selected without FEAT_DotProd"
589 );
590 neon_dotprod::dot_i32(a, b)
591}
592
593#[cfg(not(all(target_arch = "aarch64", feature = "neon-dotprod")))]
594fn dot_i32_neon_or_panic(_a: &[i8], _b: &[i8]) -> i32 {
595 panic!("neon-sdot route selected on a build without the island");
596}
597
598fn dot_i32_scalar(a: &[i8], b: &[i8]) -> i32 {
599 let mut sum = 0_i32;
600 for index in 0..a.len() {
601 sum += i32::from(a[index]) * i32::from(b[index]);
602 }
603 sum
604}
605
606/// Eight independent i32 lanes over fixed-width chunks; LLVM autovectorizes this shape into
607/// widening multiply-accumulate sequences (and SDOT where the target baseline carries it).
608/// Integer addition is associative, so the result is exactly [`dot_i32_scalar`]'s.
609fn dot_i32_autovec(a: &[i8], b: &[i8]) -> i32 {
610 const LANES: usize = 8;
611 let mut lanes = [0_i32; LANES];
612 let chunks = a.len() / LANES;
613 for chunk in 0..chunks {
614 let base = chunk * LANES;
615 for lane in 0..LANES {
616 lanes[lane] += i32::from(a[base + lane]) * i32::from(b[base + lane]);
617 }
618 }
619 let mut sum: i32 = lanes.iter().sum();
620 for index in chunks * LANES..a.len() {
621 sum += i32::from(a[index]) * i32::from(b[index]);
622 }
623 sum
624}
625
626#[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
627mod neon_dotprod {
628 //! The audited SDOT island. Named per the crate law: feature-gated, runtime-detected,
629 //! bit-identical scalar fallback in the parent module, every load bounds-checked by loop
630 //! structure, every unsafe operation carrying a SAFETY note.
631
632 use core::arch::aarch64::{vaddq_s32, vaddvq_s32, vdotq_s32, vdupq_n_s32, vld1q_s8};
633
634 /// Whether the running CPU reports FEAT_DotProd.
635 #[must_use]
636 pub fn available() -> bool {
637 std::arch::is_aarch64_feature_detected!("dotprod")
638 }
639
640 /// Exact i32 dot product via SDOT, four accumulator streams over 64-byte blocks.
641 ///
642 /// # Panics
643 ///
644 /// Panics (in the caller) unless [`available`] returned true; lengths are asserted equal by
645 /// [`super::dot_i32`].
646 #[must_use]
647 pub fn dot_i32(a: &[i8], b: &[i8]) -> i32 {
648 debug_assert!(available(), "SDOT island entered without FEAT_DotProd");
649 // SAFETY: `dot_i32_sdot` requires NEON + FEAT_DotProd, which `available()` has confirmed
650 // on this CPU at every dispatch site (asserted in `super::dot_i32`, debug-asserted here).
651 unsafe { dot_i32_sdot(a, b) }
652 }
653
654 // SAFETY: callers must have confirmed FEAT_DotProd via `available()` — the sole caller
655 // `dot_i32` above does, and `super::dot_i32` asserts it at the dispatch site. All loads are
656 // bounded by `a.len()`, which the caller asserts equals `b.len()`, and the tail is handled
657 // scalar-side, so no read passes either slice's end.
658 #[target_feature(enable = "neon,dotprod")]
659 unsafe fn dot_i32_sdot(a: &[i8], b: &[i8]) -> i32 {
660 let len = a.len();
661 let a_ptr = a.as_ptr();
662 let b_ptr = b.as_ptr();
663 let mut acc0 = vdupq_n_s32(0);
664 let mut acc1 = vdupq_n_s32(0);
665 let mut acc2 = vdupq_n_s32(0);
666 let mut acc3 = vdupq_n_s32(0);
667 let mut index = 0_usize;
668 while index + 64 <= len {
669 // SAFETY: `index + 64 <= len` bounds all four 16-byte loads inside both slices,
670 // whose lengths are equal by the caller's assertion. `vld1q_s8` has no alignment
671 // requirement beyond byte alignment.
672 unsafe {
673 acc0 = vdotq_s32(acc0, vld1q_s8(a_ptr.add(index)), vld1q_s8(b_ptr.add(index)));
674 acc1 = vdotq_s32(
675 acc1,
676 vld1q_s8(a_ptr.add(index + 16)),
677 vld1q_s8(b_ptr.add(index + 16)),
678 );
679 acc2 = vdotq_s32(
680 acc2,
681 vld1q_s8(a_ptr.add(index + 32)),
682 vld1q_s8(b_ptr.add(index + 32)),
683 );
684 acc3 = vdotq_s32(
685 acc3,
686 vld1q_s8(a_ptr.add(index + 48)),
687 vld1q_s8(b_ptr.add(index + 48)),
688 );
689 }
690 index += 64;
691 }
692 while index + 16 <= len {
693 // SAFETY: `index + 16 <= len` bounds this 16-byte load inside both slices.
694 unsafe {
695 acc0 = vdotq_s32(acc0, vld1q_s8(a_ptr.add(index)), vld1q_s8(b_ptr.add(index)));
696 }
697 index += 16;
698 }
699 let mut sum = vaddvq_s32(vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3)));
700 while index < len {
701 sum += i32::from(a[index]) * i32::from(b[index]);
702 index += 1;
703 }
704 sum
705 }
706}
707
708#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
709mod wasm_simd128 {
710 //! The audited SIMD128 island — the browser's counterpart to the SDOT one above.
711 //!
712 //! Why this exists at all: on wasm32 the dispatch fell through to `Scalar`, and NE-001's
713 //! finding that the scalar shape "vectorizes to memory bandwidth" is an *aarch64* result that
714 //! does not transfer. wasm SIMD128 has no int8 dot instruction for LLVM to pattern-match, so
715 //! the autovectorizer has nothing to reach for and the browser ran a byte-at-a-time loop.
716 //!
717 //! The instruction that replaces it is `i32x4.dot_i16x8_s`: eight i16 products summed
718 //! pairwise into four i32 lanes, one op. Feed it from `i16x8.extend_low/high_i8x16_s` and a
719 //! 16-byte block of int8 costs two widenings per operand, two dots and two adds.
720 //!
721 //! Exactness is free here and that is the point: every product of two `i8` fits `i16`, every
722 //! pairwise sum fits `i32`, and integer addition is associative — so lane order, accumulator
723 //! count and reduction order cannot change the result. This tier is *equal* to `Scalar` in
724 //! i32, not merely close, which is what lets it share the parent module's tier-equality test.
725 //!
726 //! Overflow carries no new obligation: the bound is unchanged from the scalar path at the
727 //! model's real worst-case K (3072 → |sum| ≤ 3072 × 127² ≈ 49.5M, ~43× inside i32).
728 //!
729 //! No runtime detection: `simd128` is a compile-time target feature, and a browser without it
730 //! refuses the module outright rather than mis-executing. Every engine this ships to has had
731 //! it for years (Safari 16.4 / iOS 16.4 being the last holdout).
732
733 use core::arch::wasm32::{
734 i16x8_extend_high_i8x16, i16x8_extend_low_i8x16, i32x4_add, i32x4_dot_i16x8,
735 i32x4_extract_lane, i32x4_splat, v128, v128_load,
736 };
737
738 /// Exact i32 dot product via `i32x4.dot_i16x8_s`, four accumulator streams over 64-byte
739 /// blocks.
740 #[must_use]
741 pub fn dot_i32(a: &[i8], b: &[i8]) -> i32 {
742 // SAFETY: every load below is bounded by the loop conditions against `len`, and the two
743 // slices are asserted equal in length by `super::dot_i32`. `v128_load` requires no
744 // alignment beyond the byte alignment an `&[i8]` already guarantees.
745 unsafe { dot_i32_simd128(a, b) }
746 }
747
748 /// Accumulates one 16-byte block of each operand into `acc`.
749 ///
750 /// # Safety
751 ///
752 /// `a` and `b` must each be valid for a 16-byte read.
753 // SAFETY: the contract above is discharged at both call sites, where the block loop runs only
754 // while `offset + 16 <= len` and `len` is the asserted-common length of the two slices.
755 #[inline]
756 unsafe fn accumulate_block(acc: v128, a: *const i8, b: *const i8) -> v128 {
757 // SAFETY: the caller guarantees both pointers address 16 readable bytes.
758 let (left, right) = unsafe { (v128_load(a.cast()), v128_load(b.cast())) };
759 let low = i32x4_dot_i16x8(i16x8_extend_low_i8x16(left), i16x8_extend_low_i8x16(right));
760 let high = i32x4_dot_i16x8(
761 i16x8_extend_high_i8x16(left),
762 i16x8_extend_high_i8x16(right),
763 );
764 i32x4_add(acc, i32x4_add(low, high))
765 }
766
767 /// Four output columns per pass, sharing one widening of the activation.
768 ///
769 /// Loop order stays weight-stationary — four weight rows are streamed once and reused across
770 /// all `m` activation rows — so this keeps the property the serial form was written for while
771 /// removing the redundant activation widening a per-column dot repeats `n` times.
772 pub fn linear_blocked(
773 x_q: &[i8],
774 x_scales: &[f32],
775 weight: &super::QuantizedMatrix,
776 bias: Option<&[f32]>,
777 m: usize,
778 out: &mut [f32],
779 ) {
780 let (n, k) = (weight.n, weight.k);
781 let mut col = 0;
782 while col + 4 <= n {
783 for row in 0..m {
784 let x_row = &x_q[row * k..(row + 1) * k];
785 // SAFETY: `col + 4 <= n` bounds all four weight rows inside `weight.data`, whose
786 // length is `n * k` by the type's invariant.
787 let acc = unsafe { dot4_simd128(x_row, &weight.data[col * k..], k) };
788 for (lane, accumulated) in acc.iter().enumerate() {
789 let column = col + lane;
790 #[allow(clippy::cast_precision_loss)]
791 let value = *accumulated as f32 * (x_scales[row] * weight.scales[column]);
792 out[row * n + column] = bias.map_or(value, |values| value + values[column]);
793 }
794 }
795 col += 4;
796 }
797 // Columns past the last full block of four fall back to the single-column kernel.
798 while col < n {
799 let w_row = &weight.data[col * k..(col + 1) * k];
800 for row in 0..m {
801 let x_row = &x_q[row * k..(row + 1) * k];
802 #[allow(clippy::cast_precision_loss)]
803 let value = dot_i32(x_row, w_row) as f32 * (x_scales[row] * weight.scales[col]);
804 out[row * n + col] = bias.map_or(value, |values| value + values[col]);
805 }
806 col += 1;
807 }
808 }
809
810 /// Dots one activation row against four consecutive weight rows.
811 ///
812 /// # Safety
813 ///
814 /// `weights` must be valid for `4 * k` readable bytes and `x` for `k`.
815 // SAFETY: the caller enters this path only when four whole weight rows remain (`col + 4 <= n`)
816 // and slices `weights` at `col * k` for `4 * k` bytes, with `x` the full k-length activation
817 // row; every load below is bounded by `index < k` against those same lengths.
818 unsafe fn dot4_simd128(x: &[i8], weights: &[i8], k: usize) -> [i32; 4] {
819 let x_ptr = x.as_ptr();
820 let w_ptr = weights.as_ptr();
821 let mut acc = [i32x4_splat(0); 4];
822 let mut index = 0_usize;
823 while index + 16 <= k {
824 // SAFETY: `index + 16 <= k` bounds the activation load, and each weight row starts at
825 // `lane * k` inside a region the caller guarantees is `4 * k` long.
826 let (low, high) = unsafe {
827 let block = v128_load(x_ptr.add(index).cast());
828 (
829 i16x8_extend_low_i8x16(block),
830 i16x8_extend_high_i8x16(block),
831 )
832 };
833 for (lane, accumulator) in acc.iter_mut().enumerate() {
834 // SAFETY: same bound, offset into this lane's weight row.
835 let w = unsafe { v128_load(w_ptr.add(lane * k + index).cast()) };
836 let products = i32x4_add(
837 i32x4_dot_i16x8(low, i16x8_extend_low_i8x16(w)),
838 i32x4_dot_i16x8(high, i16x8_extend_high_i8x16(w)),
839 );
840 *accumulator = i32x4_add(*accumulator, products);
841 }
842 index += 16;
843 }
844 let mut sums = [0_i32; 4];
845 for (lane, sum) in sums.iter_mut().enumerate() {
846 let total = acc[lane];
847 *sum = i32x4_extract_lane::<0>(total)
848 + i32x4_extract_lane::<1>(total)
849 + i32x4_extract_lane::<2>(total)
850 + i32x4_extract_lane::<3>(total);
851 for tail in index..k {
852 // SAFETY: `tail < k` indexes inside this lane's weight row.
853 let w = unsafe { *w_ptr.add(lane * k + tail) };
854 *sum += i32::from(x[tail]) * i32::from(w);
855 }
856 }
857 sums
858 }
859
860 /// # Safety
861 ///
862 /// `a` and `b` must have equal length; the caller asserts this.
863 // SAFETY: `super::dot_i32` asserts the two lengths are equal before dispatching here, and the
864 // only other caller is the public `dot_i32` wrapper directly above, which forwards the same
865 // pair. Every block load is guarded by `offset + 64 <= len` and the remainder runs scalar.
866 unsafe fn dot_i32_simd128(a: &[i8], b: &[i8]) -> i32 {
867 let len = a.len();
868 let a_ptr = a.as_ptr();
869 let b_ptr = b.as_ptr();
870 let mut acc0 = i32x4_splat(0);
871 let mut acc1 = i32x4_splat(0);
872 let mut acc2 = i32x4_splat(0);
873 let mut acc3 = i32x4_splat(0);
874 let mut index = 0_usize;
875 // Four independent streams so the dependent-add latency of one does not stall the next,
876 // mirroring the SDOT island's blocking.
877 while index + 64 <= len {
878 // SAFETY: `index + 64 <= len` bounds all four 16-byte loads inside both slices.
879 unsafe {
880 acc0 = accumulate_block(acc0, a_ptr.add(index), b_ptr.add(index));
881 acc1 = accumulate_block(acc1, a_ptr.add(index + 16), b_ptr.add(index + 16));
882 acc2 = accumulate_block(acc2, a_ptr.add(index + 32), b_ptr.add(index + 32));
883 acc3 = accumulate_block(acc3, a_ptr.add(index + 48), b_ptr.add(index + 48));
884 }
885 index += 64;
886 }
887 while index + 16 <= len {
888 // SAFETY: `index + 16 <= len` bounds this 16-byte load inside both slices.
889 unsafe {
890 acc0 = accumulate_block(acc0, a_ptr.add(index), b_ptr.add(index));
891 }
892 index += 16;
893 }
894 let total = i32x4_add(i32x4_add(acc0, acc1), i32x4_add(acc2, acc3));
895 let mut sum = i32x4_extract_lane::<0>(total)
896 + i32x4_extract_lane::<1>(total)
897 + i32x4_extract_lane::<2>(total)
898 + i32x4_extract_lane::<3>(total);
899 while index < len {
900 sum += i32::from(a[index]) * i32::from(b[index]);
901 index += 1;
902 }
903 sum
904 }
905}
906
907/// W8A8 linear: quantized activations `[m, k]` times a [`QuantizedMatrix`] `[n, k]`, producing
908/// f32 `[m, n]`.
909///
910/// `x_scales` carries one dynamic activation scale per row of `x_q`. The i32 accumulator is
911/// exact on every tier; dequantization applies `acc as f32 * (x_scale * w_scale)` in exactly
912/// that order on every tier, so the f32 output of any two tiers is bit-identical, not merely
913/// close. Bias (only `text_projection` carries one) is added after dequantization.
914///
915/// # Panics
916///
917/// Panics on any shape mismatch.
918#[allow(clippy::too_many_arguments)]
919pub fn linear_q8(
920 x_q: &[i8],
921 x_scales: &[f32],
922 weight: &QuantizedMatrix,
923 bias: Option<&[f32]>,
924 m: usize,
925 out: &mut [f32],
926 tier: Int8Tier,
927) {
928 let (n, k) = (weight.n, weight.k);
929 assert_eq!(x_q.len(), m * k, "x_q must be [m, k]");
930 assert_eq!(x_scales.len(), m, "x_scales must be [m]");
931 assert_eq!(out.len(), m * n, "out must be [m, n]");
932 if let Some(bias) = bias {
933 assert_eq!(bias.len(), n, "bias must be [n]");
934 }
935 // Large operations fan out across the persistent team when one exists; the partitioned
936 // result is bit-identical per element, so this is purely a speed dispatch. Small matrices
937 // stay serial — the dispatch handshake would cost more than the work.
938 if n * k >= TEAM_WORK_THRESHOLD_BYTES
939 && !crate::team::thread_bypassed()
940 && let Some(team) = crate::team::armed()
941 {
942 team.linear_q8(x_q, x_scales, weight, bias, m, out, tier);
943 return;
944 }
945
946 // Weight-stationary loop order: each Q8 weight row is streamed exactly once and reused
947 // across all m activation rows, so an m>1 call (prefill, the seq-16 verify pass) does not
948 // re-read the whole matrix m times. Each output element's dot product is unchanged, so this
949 // ordering is bit-identical to the m-outer form.
950 // wasm has no int8 dot instruction, so a lone dot spends four of its eight ops per 16 bytes
951 // just widening i8 to i16 — and half of that widening is the *activation*, which is identical
952 // for every output column. Computing four columns per pass hoists it: 26 ops for 64 MACs
953 // instead of 32, which is the register-blocking lever the doctrine warns is the real one
954 // ("the instruction is not the lever; the blocking is"). Bit-identical by construction — the
955 // same per-element i32 dot, only the loop nest changes.
956 #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
957 if matches!(tier, Int8Tier::WasmSimd128) {
958 wasm_simd128::linear_blocked(x_q, x_scales, weight, bias, m, out);
959 return;
960 }
961
962 for col in 0..n {
963 let w_row = &weight.data[col * k..(col + 1) * k];
964 let w_scale = weight.scales[col];
965 let bias_term = bias.map(|b| b[col]);
966 for row in 0..m {
967 let x_row = &x_q[row * k..(row + 1) * k];
968 let acc = dot_i32(x_row, w_row, tier);
969 let value = acc as f32 * (x_scales[row] * w_scale);
970 out[row * n + col] = bias_term.map_or(value, |b| value + b);
971 }
972 }
973}
974
975/// Quantizes an f32 activation matrix `[m, k]` per row and runs [`linear_q8`].
976///
977/// This is the drop-in W8A8 counterpart of `f32ref::linear`: same `[m, k] × [n, k]ᵀ → [m, n]`
978/// layout, same bias placement. The row quantization is the canonical symmetric recipe.
979///
980/// # Panics
981///
982/// Panics on any shape mismatch or a non-finite activation.
983pub fn linear_q8_dynamic(
984 x: &[f32],
985 weight: &QuantizedMatrix,
986 bias: Option<&[f32]>,
987 m: usize,
988 out: &mut [f32],
989 tier: Int8Tier,
990) {
991 let k = weight.k;
992 assert_eq!(x.len(), m * k, "x must be [m, k]");
993 // Thread-local scratch instead of two per-call `vec!`s: this is the armed W8A8 entry the
994 // talker and microdecoder hit ~300 times per frame, and the doctrine pins "no allocator
995 // activity in steady-state decode" as load-bearing. Quantization runs on the caller thread
996 // before any team dispatch, and nothing below re-enters this function, so the borrow is
997 // never contended; the buffers only ever grow, to the largest `[m, k]` this thread has seen,
998 // and are sliced to the exact live extent so stale bytes past it are unreachable.
999 thread_local! {
1000 static QUANT_SCRATCH: std::cell::RefCell<(Vec<i8>, Vec<f32>)> =
1001 const { std::cell::RefCell::new((Vec::new(), Vec::new())) };
1002 }
1003 QUANT_SCRATCH.with(|scratch| {
1004 let mut guard = scratch.borrow_mut();
1005 let (q_buffer, scale_buffer) = &mut *guard;
1006 if q_buffer.len() < m * k {
1007 q_buffer.resize(m * k, 0);
1008 }
1009 if scale_buffer.len() < m {
1010 scale_buffer.resize(m, 0.0);
1011 }
1012 let x_q = &mut q_buffer[..m * k];
1013 let x_scales = &mut scale_buffer[..m];
1014 for ((x_row, q_row), scale) in x
1015 .chunks_exact(k)
1016 .zip(x_q.chunks_exact_mut(k))
1017 .zip(x_scales.iter_mut())
1018 {
1019 *scale = quantize_row_q8(x_row, q_row);
1020 }
1021 linear_q8(x_q, x_scales, weight, bias, m, out, tier);
1022 });
1023}
1024
1025#[cfg(test)]
1026mod tests {
1027 use super::*;
1028
1029 /// Deterministic pseudo-random Q8 bytes (SplitMix64), full `[-127, 127]` range.
1030 fn pseudo_random_q8(len: usize, seed: u64) -> Vec<i8> {
1031 let mut state = seed;
1032 (0..len)
1033 .map(|_| {
1034 state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
1035 let mut z = state;
1036 z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
1037 z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
1038 z ^= z >> 31;
1039 // Map to [-127, 127]; never -128, matching the converter contract.
1040 ((z % 255) as i32 - 127) as i8
1041 })
1042 .collect()
1043 }
1044
1045 /// The model's real decode GEMV shapes: (n, k) per §7 of the plan.
1046 const MODEL_SHAPES: &[(usize, usize)] = &[
1047 (2048, 1024), // q_proj / per-depth heads
1048 (1024, 1024), // k_proj / v_proj
1049 (1024, 2048), // o_proj
1050 (3072, 1024), // gate/up_proj, primary head
1051 (1024, 3072), // down_proj (binding talker K)
1052 ];
1053
1054 #[test]
1055 fn every_tier_is_exactly_equal_in_i32_at_every_model_shape() {
1056 for &(n, k) in MODEL_SHAPES {
1057 let a = pseudo_random_q8(k, 0x5eed_0001 ^ (n as u64) << 20 ^ k as u64);
1058 let w = pseudo_random_q8(n * k, 0x5eed_0002 ^ (n as u64) << 20 ^ k as u64);
1059 for row in [0, n / 2, n - 1] {
1060 let w_row = &w[row * k..(row + 1) * k];
1061 let reference = dot_i32(&a, w_row, Int8Tier::Scalar);
1062 for tier in Int8Tier::available() {
1063 assert_eq!(
1064 dot_i32(&a, w_row, tier),
1065 reference,
1066 "tier {} diverged at shape {n}x{k} row {row}",
1067 tier.as_str()
1068 );
1069 }
1070 }
1071 }
1072 }
1073
1074 #[test]
1075 fn every_tier_survives_the_all_extreme_reduction_at_the_binding_census_k() {
1076 // 127 * 127 * 8192 = 132,120,576 — the S8S8 all-extreme envelope at the largest census K.
1077 for k in [2048_usize, 3072, 4608, 7168, 8192] {
1078 let a = vec![127_i8; k];
1079 let b = vec![127_i8; k];
1080 let negative = vec![-127_i8; k];
1081 let expected = 127_i64 * 127 * k as i64;
1082 for tier in Int8Tier::available() {
1083 assert_eq!(
1084 i64::from(dot_i32(&a, &b, tier)),
1085 expected,
1086 "positive all-extreme diverged on {} at K={k}",
1087 tier.as_str()
1088 );
1089 assert_eq!(
1090 i64::from(dot_i32(&a, &negative, tier)),
1091 -expected,
1092 "negative all-extreme diverged on {} at K={k}",
1093 tier.as_str()
1094 );
1095 }
1096 }
1097 }
1098
1099 #[test]
1100 fn tail_lengths_that_defeat_block_boundaries_stay_exact() {
1101 // Exercise every SDOT path: <16 (pure tail), 16..64 (single-block loop), 64+tail.
1102 for len in [1_usize, 7, 15, 16, 17, 63, 64, 65, 100, 129] {
1103 let a = pseudo_random_q8(len, tail_seed(len));
1104 let b = pseudo_random_q8(len, tail_seed(len) ^ 1);
1105 let reference = dot_i32(&a, &b, Int8Tier::Scalar);
1106 for tier in Int8Tier::available() {
1107 assert_eq!(
1108 dot_i32(&a, &b, tier),
1109 reference,
1110 "len={len} {}",
1111 tier.as_str()
1112 );
1113 }
1114 }
1115 }
1116
1117 #[test]
1118 fn quantizer_matches_the_canonical_converter_semantics() {
1119 // Ties-to-even, clamp, zero-row scale, and the -128 exclusion. The cross-crate
1120 // byte-identity test against `ftts-artifacts` lives in `ftts-model-qwen`.
1121 let row = [
1122 -127.0_f32, -126.5, -125.5, -1.5, -0.5, 0.5, 1.5, 125.5, 126.5, 127.0,
1123 ];
1124 let mut q = [0_i8; 10];
1125 let scale = quantize_row_q8(&row, &mut q);
1126 assert_eq!(scale.to_bits(), 1.0_f32.to_bits());
1127 assert_eq!(q, [-127, -126, -126, -2, 0, 0, 2, 126, 126, 127]);
1128
1129 let zeros = [0.0_f32; 4];
1130 let mut qz = [1_i8; 4];
1131 assert_eq!(
1132 quantize_row_q8(&zeros, &mut qz).to_bits(),
1133 1.0_f32.to_bits()
1134 );
1135 assert_eq!(qz, [0, 0, 0, 0]);
1136
1137 let matrix = QuantizedMatrix::quantize(&[2.0, -1.0, 0.0, 3.0], 2, 2);
1138 assert_eq!(matrix.scales[0].to_bits(), (2.0_f32 / 127.0).to_bits());
1139 assert_eq!(matrix.scales[1].to_bits(), (3.0_f32 / 127.0).to_bits());
1140 assert!(matrix.data.iter().all(|&b| b != -128));
1141 }
1142
1143 #[test]
1144 fn dynamic_w8a8_linear_tracks_the_f32_reference_within_quant_error() {
1145 // Not a parity claim — a sanity bound that the dequant plumbing is wired correctly.
1146 let (n, k) = (64_usize, 128_usize);
1147 let mut weight = vec![0.0_f32; n * k];
1148 let mut x = vec![0.0_f32; k];
1149 let mut state = 0x1234_5678_u64;
1150 let mut next = || {
1151 state = state
1152 .wrapping_mul(6_364_136_223_846_793_005)
1153 .wrapping_add(1);
1154 ((state >> 33) as f32 / (1u64 << 31) as f32) - 1.0
1155 };
1156 for value in weight.iter_mut() {
1157 *value = next();
1158 }
1159 for value in x.iter_mut() {
1160 *value = next();
1161 }
1162 let quantized = QuantizedMatrix::quantize(&weight, n, k);
1163 let mut out_q8 = vec![0.0_f32; n];
1164 linear_q8_dynamic(&x, &quantized, None, 1, &mut out_q8, Int8Tier::Autovec);
1165
1166 let mut out_f32 = vec![0.0_f32; n];
1167 crate::f32ref::linear(&x, &weight, None, 1, k, n, &mut out_f32);
1168
1169 let dot = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>();
1170 let cosine = dot(&out_q8, &out_f32)
1171 / (dot(&out_q8, &out_q8).sqrt() * dot(&out_f32, &out_f32).sqrt());
1172 assert!(
1173 cosine > 0.999,
1174 "W8A8 dequant plumbing is broken: cosine {cosine}"
1175 );
1176 }
1177
1178 #[test]
1179 fn tiers_produce_bit_identical_f32_output_not_merely_close() {
1180 let (n, k) = (256_usize, 1024_usize);
1181 let weight: Vec<f32> = pseudo_random_q8(n * k, 77)
1182 .iter()
1183 .map(|&b| f32::from(b) / 64.0)
1184 .collect();
1185 let x: Vec<f32> = pseudo_random_q8(k, 78)
1186 .iter()
1187 .map(|&b| f32::from(b) / 64.0)
1188 .collect();
1189 let quantized = QuantizedMatrix::quantize(&weight, n, k);
1190 let mut reference = vec![0.0_f32; n];
1191 linear_q8_dynamic(&x, &quantized, None, 1, &mut reference, Int8Tier::Scalar);
1192 for tier in Int8Tier::available() {
1193 let mut out = vec![0.0_f32; n];
1194 linear_q8_dynamic(&x, &quantized, None, 1, &mut out, tier);
1195 for (index, (a, b)) in reference.iter().zip(&out).enumerate() {
1196 assert_eq!(
1197 a.to_bits(),
1198 b.to_bits(),
1199 "tier {} f32 output differs at {index}",
1200 tier.as_str()
1201 );
1202 }
1203 }
1204 }
1205
1206 fn tail_seed(len: usize) -> u64 {
1207 0x7a11_0000 ^ len as u64
1208 }
1209
1210 #[test]
1211 fn quant_scratch_left_oversized_by_a_big_call_never_bleeds_into_a_small_one() {
1212 // The dynamic entry reuses thread-local quant buffers that only grow. Run the largest
1213 // batched shape first so the scratch holds 16×3072 stale bytes, then the m=1 decode
1214 // shape, and demand bit-equality with the same computation through freshly allocated
1215 // buffers via `linear_q8` directly.
1216 let (big_m, big_k, big_n) = (16_usize, 3072_usize, 8_usize);
1217 let big_weight: Vec<f32> = pseudo_random_q8(big_n * big_k, 91)
1218 .iter()
1219 .map(|&b| f32::from(b) / 64.0)
1220 .collect();
1221 let big_x: Vec<f32> = pseudo_random_q8(big_m * big_k, 92)
1222 .iter()
1223 .map(|&b| f32::from(b) / 64.0)
1224 .collect();
1225 let big_quantized = QuantizedMatrix::quantize(&big_weight, big_n, big_k);
1226 let mut big_out = vec![0.0_f32; big_m * big_n];
1227 linear_q8_dynamic(
1228 &big_x,
1229 &big_quantized,
1230 None,
1231 big_m,
1232 &mut big_out,
1233 Int8Tier::Scalar,
1234 );
1235
1236 let (m, k, n) = (1_usize, 1024_usize, 32_usize);
1237 let weight: Vec<f32> = pseudo_random_q8(n * k, 93)
1238 .iter()
1239 .map(|&b| f32::from(b) / 64.0)
1240 .collect();
1241 let x: Vec<f32> = pseudo_random_q8(m * k, 94)
1242 .iter()
1243 .map(|&b| f32::from(b) / 64.0)
1244 .collect();
1245 let quantized = QuantizedMatrix::quantize(&weight, n, k);
1246 let mut via_scratch = vec![0.0_f32; m * n];
1247 linear_q8_dynamic(&x, &quantized, None, m, &mut via_scratch, Int8Tier::Scalar);
1248
1249 let mut fresh_q = vec![0_i8; m * k];
1250 let fresh_scale = quantize_row_q8(&x, &mut fresh_q);
1251 let mut via_fresh = vec![0.0_f32; m * n];
1252 linear_q8(
1253 &fresh_q,
1254 &[fresh_scale],
1255 &quantized,
1256 None,
1257 m,
1258 &mut via_fresh,
1259 Int8Tier::Scalar,
1260 );
1261 for (index, (a, b)) in via_scratch.iter().zip(&via_fresh).enumerate() {
1262 assert_eq!(
1263 a.to_bits(),
1264 b.to_bits(),
1265 "scratch-path output differs from fresh-buffer output at {index}"
1266 );
1267 }
1268 }
1269}