Skip to main content

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}