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    for (&value, slot) in row.iter().zip(output.iter_mut()) {
64        let rounded = (value / scale).clamp(-127.0, 127.0).round_ties_even();
65        // The clamp bounds the conversion inside i8, and the symmetric contract additionally
66        // excludes the otherwise-representable -128.
67        *slot = rounded as i8;
68    }
69    scale
70}
71
72/// A weight matrix quantized with per-output-channel symmetric Q8 scales.
73///
74/// Layout is the `nn.Linear` layout the checkpoint stores: `data` is `[n, k]` row-major with one
75/// f32 scale per output row. Quantized once at hydration; the borrowed f32 tensor is untouched.
76#[derive(Clone, Debug)]
77pub struct QuantizedMatrix {
78    /// Q8 bytes, `[n, k]` row-major, each value in `[-127, 127]`.
79    pub data: Vec<i8>,
80    /// One symmetric scale per output row, `[n]`.
81    pub scales: Vec<f32>,
82    /// Output rows.
83    pub n: usize,
84    /// Reduction length of one output element.
85    pub k: usize,
86}
87
88impl QuantizedMatrix {
89    /// Stacks matrices with a shared reduction length into one taller matrix.
90    ///
91    /// Row bytes and scales are byte-identical to quantizing each part separately — this exists
92    /// so fused projections (QKV, gate‖up) can run as ONE kernel dispatch while every output
93    /// row keeps exactly the per-channel quantization it would have had alone.
94    ///
95    /// # Panics
96    ///
97    /// Panics if the parts disagree on `k` or the list is empty.
98    #[must_use]
99    pub fn concat_rows(parts: &[&Self]) -> Self {
100        let k = parts.first().expect("at least one part").k;
101        assert!(parts.iter().all(|part| part.k == k), "parts must share k");
102        let n = parts.iter().map(|part| part.n).sum();
103        let mut data = Vec::with_capacity(n * k);
104        let mut scales = Vec::with_capacity(n);
105        for part in parts {
106            data.extend_from_slice(&part.data);
107            scales.extend_from_slice(&part.scales);
108        }
109        Self { data, scales, n, k }
110    }
111
112    /// Quantizes an `[n, k]` f32 weight matrix one output channel at a time.
113    ///
114    /// # Panics
115    ///
116    /// Panics if `weight.len() != n * k` or any value is non-finite.
117    #[must_use]
118    pub fn quantize(weight: &[f32], n: usize, k: usize) -> Self {
119        assert_eq!(weight.len(), n * k, "weight must be [n, k]");
120        let mut data = vec![0_i8; n * k];
121        let mut scales = vec![0.0_f32; n];
122        for ((weight_row, data_row), scale) in weight
123            .chunks_exact(k)
124            .zip(data.chunks_exact_mut(k))
125            .zip(scales.iter_mut())
126        {
127            *scale = quantize_row_q8(weight_row, data_row);
128        }
129        Self { data, scales, n, k }
130    }
131}
132
133/// An executable int8 dot-product route.
134///
135/// Every variant is exactly equal in i32 to `Scalar` on every input. `NeonSdot` exists only on
136/// aarch64 builds with the `neon-dotprod` feature and is dispatchable only where the CPU reports
137/// FEAT_DotProd at runtime.
138#[derive(Clone, Copy, Debug, Eq, PartialEq)]
139pub enum Int8Tier {
140    /// Portable left-to-right checked-free scalar loop; the reference every tier must equal.
141    Scalar,
142    /// Portable eight-lane loop, retained ONLY as an A/B datapoint: measured ~15x SLOWER than
143    /// `Scalar` at m=1 on M4 Pro (NE-001) — the manual lane structure defeats LLVM's
144    /// autovectorizer, while the plain `Scalar` shape vectorizes to memory bandwidth. Never the
145    /// dispatch default.
146    Autovec,
147    /// Hand SDOT island (aarch64 + FEAT_DotProd), four 16-byte accumulator streams.
148    NeonSdot,
149}
150
151impl Int8Tier {
152    /// Stable machine-readable route name.
153    #[must_use]
154    pub const fn as_str(self) -> &'static str {
155        match self {
156            Self::Scalar => "scalar",
157            Self::Autovec => "autovec",
158            Self::NeonSdot => "neon-sdot",
159        }
160    }
161
162    /// Every tier this build can execute on the running silicon, scalar first.
163    #[must_use]
164    pub fn available() -> Vec<Self> {
165        let mut tiers = vec![Self::Scalar, Self::Autovec];
166        if neon_sdot_available() {
167            tiers.push(Self::NeonSdot);
168        }
169        tiers
170    }
171
172    /// The route the int8 path dispatches by default, honoring the `FTTS_INT8_TIER` override.
173    ///
174    /// The override exists for interleaved A/B measurement (`scalar` / `autovec` / `neon-sdot`);
175    /// an unavailable or unrecognized override falls back to the measured default rather than
176    /// panicking mid-synthesis. Until a per-shape KernelPlan lands, the default is `NeonSdot`
177    /// where FEAT_DotProd exists, else `Scalar`. Measured on M4 Pro (2026-08-08, shape bench,
178    /// noisy shared host, indicative): plain `Scalar` autovectorizes to ~50 GB/s and ties SDOT
179    /// at m=1 — NE-INH-003 reconfirmed — while the hand-shaped `Autovec` lane loop defeats the
180    /// vectorizer and loses ~15x; it stays only as an A/B datapoint.
181    #[must_use]
182    pub fn dispatch() -> Self {
183        match std::env::var("FTTS_INT8_TIER").as_deref() {
184            Ok("scalar") => Self::Scalar,
185            Ok("autovec") => Self::Autovec,
186            Ok("neon-sdot") if neon_sdot_available() => Self::NeonSdot,
187            _ if neon_sdot_available() => Self::NeonSdot,
188            _ => Self::Scalar,
189        }
190    }
191}
192
193/// Which quantized linear op class the armed route runs.
194///
195/// `W8A8` quantizes activations per row and uses the exact-i32 int8 dot — fastest, but the
196/// activation rounding perturbs logits enough that seeded sampling can draw different tokens
197/// than f32. `W8A16` keeps activations f32 and dequantizes weights in-register — the same
198/// one-byte-per-weight memory traffic, no activation error, so the output tracks the f32
199/// reference much more closely. Its f32 accumulation is lane-ordered (not the reference's
200/// left-to-right order): this is a lossy route already, so reduction-order freedom is part of
201/// the deal, and the fidelity gate is measured downstream, not asserted bitwise.
202#[derive(Clone, Copy, Debug, Eq, PartialEq)]
203pub enum QuantLinearMode {
204    /// Int8 activations times int8 weights, exact i32 accumulation.
205    W8A8(Int8Tier),
206    /// f32 activations times dequantized int8 weights, lane-ordered f32 accumulation.
207    W8A16,
208}
209
210impl QuantLinearMode {
211    /// Stable machine-readable mode name.
212    #[must_use]
213    pub const fn as_str(self) -> &'static str {
214        match self {
215            Self::W8A8(_) => "w8a8",
216            Self::W8A16 => "w8a16",
217        }
218    }
219}
220
221/// W8A16 linear: f32 activations `[m, k]` times a [`QuantizedMatrix`] `[n, k]` producing
222/// f32 `[m, n]`.
223///
224/// Eight independent f32 FMA lanes per dot product, weights widened from i8 in-register; the
225/// per-output-channel scale multiplies once after accumulation, mirroring the W8A8 dequant
226/// order. Weight-stationary loop, like [`linear_q8`].
227///
228/// # Panics
229///
230/// Panics on any shape mismatch.
231pub fn linear_w8a16(
232    x: &[f32],
233    weight: &QuantizedMatrix,
234    bias: Option<&[f32]>,
235    m: usize,
236    out: &mut [f32],
237) {
238    let (n, k) = (weight.n, weight.k);
239    assert_eq!(x.len(), m * k, "x must be [m, k]");
240    assert_eq!(out.len(), m * n, "out must be [m, n]");
241    if let Some(bias) = bias {
242        assert_eq!(bias.len(), n, "bias must be [n]");
243    }
244    for col in 0..n {
245        let w_row = &weight.data[col * k..(col + 1) * k];
246        let w_scale = weight.scales[col];
247        let bias_term = bias.map(|b| b[col]);
248        for row in 0..m {
249            let x_row = &x[row * k..(row + 1) * k];
250            let acc = dot_w8a16(x_row, w_row);
251            let value = acc * w_scale;
252            out[row * n + col] = bias_term.map_or(value, |b| value + b);
253        }
254    }
255}
256
257/// Eight-lane f32 dot of an f32 row against an i8 weight row, widened in-register.
258fn dot_w8a16(x: &[f32], w: &[i8]) -> f32 {
259    const LANES: usize = 8;
260    let mut lanes = [0.0_f32; LANES];
261    let chunks = x.len() / LANES;
262    for chunk in 0..chunks {
263        let base = chunk * LANES;
264        for lane in 0..LANES {
265            lanes[lane] = f32::from(w[base + lane]).mul_add(x[base + lane], lanes[lane]);
266        }
267    }
268    let mut sum: f32 = lanes.iter().sum();
269    for index in chunks * LANES..x.len() {
270        sum = f32::from(w[index]).mul_add(x[index], sum);
271    }
272    sum
273}
274
275/// The armed quantized-linear mode for the talker/microdecoder route.
276///
277/// `FTTS_INT8=1` or `w8a8` selects the int8-dot route; `FTTS_INT8=w8a16` selects the
278/// weight-only route. Anything else means the caller should not be arming quantization at all
279/// (the kill-switch check happens before this is consulted).
280#[must_use]
281pub fn quant_mode_from_environment() -> QuantLinearMode {
282    match std::env::var("FTTS_INT8").as_deref() {
283        Ok("w8a16") => QuantLinearMode::W8A16,
284        _ => QuantLinearMode::W8A8(autotuned_plan().decode_gemv),
285    }
286}
287
288/// Runs one quantized linear in the selected mode; the drop-in used by the armed model paths.
289pub fn quant_linear(
290    mode: QuantLinearMode,
291    x: &[f32],
292    weight: &QuantizedMatrix,
293    bias: Option<&[f32]>,
294    m: usize,
295    out: &mut [f32],
296) {
297    match mode {
298        QuantLinearMode::W8A8(tier) => linear_q8_dynamic(x, weight, bias, m, out, tier),
299        QuantLinearMode::W8A16 => linear_w8a16(x, weight, bias, m, out),
300    }
301}
302
303/// The measured per-regime route assignment, decided once per process.
304///
305/// v0 of the KernelPlan: two regimes, no persistence (`.fttspack` owns that when it lands).
306/// Safe to decide by noisy measurement because every tier produces bit-identical output — a
307/// wrong pick costs microseconds, never correctness.
308#[derive(Clone, Copy, Debug, Eq, PartialEq)]
309pub struct KernelPlanV0 {
310    /// Route for m=1 decode GEMVs (the talker/microdecoder step shape).
311    pub decode_gemv: Int8Tier,
312    /// Route for batched GEMMs (prefill, seq-16 verify, offline codec).
313    pub batch_gemm: Int8Tier,
314}
315
316/// Measures each available tier at the two live regimes and returns the winners.
317///
318/// Decided once per process and cached. `FTTS_INT8_TIER` overrides both regimes — the A/B
319/// override must pin the route it names, not merely suggest it. Cost: a few milliseconds of
320/// synthetic dots at the model's real reduction lengths.
321pub fn autotuned_plan() -> KernelPlanV0 {
322    static PLAN: OnceLock<KernelPlanV0> = OnceLock::new();
323    *PLAN.get_or_init(|| {
324        if std::env::var("FTTS_INT8_TIER").is_ok() {
325            let forced = Int8Tier::dispatch();
326            return KernelPlanV0 {
327                decode_gemv: forced,
328                batch_gemm: forced,
329            };
330        }
331        if let Some(cached) = load_persisted_plan() {
332            return cached;
333        }
334        let plan = KernelPlanV0 {
335            // Talker/microdecoder decode: one activation row against tall matrices; K = 1024
336            // and 3072 are the real reduction lengths, 256 output rows keep the probe cheap
337            // while streaming enough weight bytes to reach the bandwidth regime.
338            decode_gemv: fastest_tier(&[(1, 1024, 256), (1, 3072, 256)]),
339            // Verify/prefill/codec batches: sixteen rows, same reduction lengths.
340            batch_gemm: fastest_tier(&[(16, 1024, 128), (16, 3072, 64)]),
341        };
342        persist_plan(plan);
343        plan
344    })
345}
346
347/// Where the measured plan is cached between runs: the pre-`.fttspack` v0 of the per-machine
348/// execution cache. Losing or corrupting this file only costs a re-measurement.
349fn plan_cache_path() -> Option<std::path::PathBuf> {
350    std::env::var_os("HOME")
351        .map(|home| std::path::PathBuf::from(home).join(".cache/franken_tts/kernel_plan_v0.txt"))
352}
353
354/// The cache key: anything here changing invalidates the measurement.
355fn plan_cache_key() -> String {
356    let tiers: Vec<&str> = Int8Tier::available().iter().map(|t| t.as_str()).collect();
357    format!(
358        "v0|crate={}|tiers={}",
359        env!("CARGO_PKG_VERSION"),
360        tiers.join(",")
361    )
362}
363
364fn load_persisted_plan() -> Option<KernelPlanV0> {
365    let text = std::fs::read_to_string(plan_cache_path()?).ok()?;
366    let mut lines = text.lines();
367    if lines.next()? != plan_cache_key() {
368        return None;
369    }
370    let parse = |line: &str| match line {
371        "scalar" => Some(Int8Tier::Scalar),
372        "autovec" => Some(Int8Tier::Autovec),
373        "neon-sdot" if neon_sdot_available() => Some(Int8Tier::NeonSdot),
374        _ => None,
375    };
376    Some(KernelPlanV0 {
377        decode_gemv: parse(lines.next()?)?,
378        batch_gemm: parse(lines.next()?)?,
379    })
380}
381
382fn persist_plan(plan: KernelPlanV0) {
383    let Some(path) = plan_cache_path() else {
384        return;
385    };
386    if let Some(parent) = path.parent() {
387        let _ = std::fs::create_dir_all(parent);
388    }
389    // Best-effort: an unwritable cache directory must never fail synthesis.
390    let _ = std::fs::write(
391        path,
392        format!(
393            "{}\n{}\n{}\n",
394            plan_cache_key(),
395            plan.decode_gemv.as_str(),
396            plan.batch_gemm.as_str()
397        ),
398    );
399}
400
401/// Times every available tier over the given `(m, k, n)` probes; median of three rounds each,
402/// summed across probes, smallest total wins. Ties break toward the earlier tier in
403/// [`Int8Tier::available`] order (scalar first — the simpler route).
404fn fastest_tier(probes: &[(usize, usize, usize)]) -> Int8Tier {
405    use std::time::Instant;
406    let tiers = Int8Tier::available();
407    let mut best = (tiers[0], f64::MAX);
408    for &tier in &tiers {
409        let mut total = 0.0_f64;
410        for &(m, k, n) in probes {
411            let x_q: Vec<i8> = (0..m * k).map(|i| ((i * 37 + 11) % 255) as i8).collect();
412            let x_scales = vec![1.0_f32; m];
413            let weight = QuantizedMatrix {
414                data: (0..n * k).map(|i| ((i * 29 + 5) % 255) as i8).collect(),
415                scales: vec![1.0_f32; n],
416                n,
417                k,
418            };
419            let mut out = vec![0.0_f32; m * n];
420            let mut rounds: Vec<f64> = (0..3)
421                .map(|_| {
422                    let start = Instant::now();
423                    linear_q8(&x_q, &x_scales, &weight, None, m, &mut out, tier);
424                    start.elapsed().as_secs_f64()
425                })
426                .collect();
427            rounds.sort_by(f64::total_cmp);
428            total += rounds[1];
429        }
430        if total < best.1 {
431            best = (tier, total);
432        }
433    }
434    best.0
435}
436
437/// Whether the SDOT island is compiled in and the CPU reports FEAT_DotProd.
438#[must_use]
439pub fn neon_sdot_available() -> bool {
440    #[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
441    {
442        neon_dotprod::available()
443    }
444    #[cfg(not(all(target_arch = "aarch64", feature = "neon-dotprod")))]
445    {
446        false
447    }
448}
449
450/// Exact i32 dot product of two Q8 rows over the selected route.
451///
452/// # Panics
453///
454/// Panics if the lengths differ, or if `NeonSdot` is requested where it is not executable.
455#[must_use]
456pub fn dot_i32(a: &[i8], b: &[i8], tier: Int8Tier) -> i32 {
457    assert_eq!(a.len(), b.len(), "int8 dot inputs must match");
458    match tier {
459        Int8Tier::Scalar => dot_i32_scalar(a, b),
460        Int8Tier::Autovec => dot_i32_autovec(a, b),
461        Int8Tier::NeonSdot => dot_i32_neon_or_panic(a, b),
462    }
463}
464
465#[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
466fn dot_i32_neon_or_panic(a: &[i8], b: &[i8]) -> i32 {
467    assert!(
468        neon_dotprod::available(),
469        "neon-sdot route selected without FEAT_DotProd"
470    );
471    neon_dotprod::dot_i32(a, b)
472}
473
474#[cfg(not(all(target_arch = "aarch64", feature = "neon-dotprod")))]
475fn dot_i32_neon_or_panic(_a: &[i8], _b: &[i8]) -> i32 {
476    panic!("neon-sdot route selected on a build without the island");
477}
478
479fn dot_i32_scalar(a: &[i8], b: &[i8]) -> i32 {
480    let mut sum = 0_i32;
481    for index in 0..a.len() {
482        sum += i32::from(a[index]) * i32::from(b[index]);
483    }
484    sum
485}
486
487/// Eight independent i32 lanes over fixed-width chunks; LLVM autovectorizes this shape into
488/// widening multiply-accumulate sequences (and SDOT where the target baseline carries it).
489/// Integer addition is associative, so the result is exactly [`dot_i32_scalar`]'s.
490fn dot_i32_autovec(a: &[i8], b: &[i8]) -> i32 {
491    const LANES: usize = 8;
492    let mut lanes = [0_i32; LANES];
493    let chunks = a.len() / LANES;
494    for chunk in 0..chunks {
495        let base = chunk * LANES;
496        for lane in 0..LANES {
497            lanes[lane] += i32::from(a[base + lane]) * i32::from(b[base + lane]);
498        }
499    }
500    let mut sum: i32 = lanes.iter().sum();
501    for index in chunks * LANES..a.len() {
502        sum += i32::from(a[index]) * i32::from(b[index]);
503    }
504    sum
505}
506
507#[cfg(all(target_arch = "aarch64", feature = "neon-dotprod"))]
508mod neon_dotprod {
509    //! The audited SDOT island. Named per the crate law: feature-gated, runtime-detected,
510    //! bit-identical scalar fallback in the parent module, every load bounds-checked by loop
511    //! structure, every unsafe operation carrying a SAFETY note.
512
513    use core::arch::aarch64::{vaddq_s32, vaddvq_s32, vdotq_s32, vdupq_n_s32, vld1q_s8};
514
515    /// Whether the running CPU reports FEAT_DotProd.
516    #[must_use]
517    pub fn available() -> bool {
518        std::arch::is_aarch64_feature_detected!("dotprod")
519    }
520
521    /// Exact i32 dot product via SDOT, four accumulator streams over 64-byte blocks.
522    ///
523    /// # Panics
524    ///
525    /// Panics (in the caller) unless [`available`] returned true; lengths are asserted equal by
526    /// [`super::dot_i32`].
527    #[must_use]
528    pub fn dot_i32(a: &[i8], b: &[i8]) -> i32 {
529        debug_assert!(available(), "SDOT island entered without FEAT_DotProd");
530        // SAFETY: `dot_i32_sdot` requires NEON + FEAT_DotProd, which `available()` has confirmed
531        // on this CPU at every dispatch site (asserted in `super::dot_i32`, debug-asserted here).
532        unsafe { dot_i32_sdot(a, b) }
533    }
534
535    #[target_feature(enable = "neon,dotprod")]
536    unsafe fn dot_i32_sdot(a: &[i8], b: &[i8]) -> i32 {
537        let len = a.len();
538        let a_ptr = a.as_ptr();
539        let b_ptr = b.as_ptr();
540        let mut acc0 = vdupq_n_s32(0);
541        let mut acc1 = vdupq_n_s32(0);
542        let mut acc2 = vdupq_n_s32(0);
543        let mut acc3 = vdupq_n_s32(0);
544        let mut index = 0_usize;
545        while index + 64 <= len {
546            // SAFETY: `index + 64 <= len` bounds all four 16-byte loads inside both slices,
547            // whose lengths are equal by the caller's assertion. `vld1q_s8` has no alignment
548            // requirement beyond byte alignment.
549            unsafe {
550                acc0 = vdotq_s32(acc0, vld1q_s8(a_ptr.add(index)), vld1q_s8(b_ptr.add(index)));
551                acc1 = vdotq_s32(
552                    acc1,
553                    vld1q_s8(a_ptr.add(index + 16)),
554                    vld1q_s8(b_ptr.add(index + 16)),
555                );
556                acc2 = vdotq_s32(
557                    acc2,
558                    vld1q_s8(a_ptr.add(index + 32)),
559                    vld1q_s8(b_ptr.add(index + 32)),
560                );
561                acc3 = vdotq_s32(
562                    acc3,
563                    vld1q_s8(a_ptr.add(index + 48)),
564                    vld1q_s8(b_ptr.add(index + 48)),
565                );
566            }
567            index += 64;
568        }
569        while index + 16 <= len {
570            // SAFETY: `index + 16 <= len` bounds this 16-byte load inside both slices.
571            unsafe {
572                acc0 = vdotq_s32(acc0, vld1q_s8(a_ptr.add(index)), vld1q_s8(b_ptr.add(index)));
573            }
574            index += 16;
575        }
576        let mut sum = vaddvq_s32(vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3)));
577        while index < len {
578            sum += i32::from(a[index]) * i32::from(b[index]);
579            index += 1;
580        }
581        sum
582    }
583}
584
585/// W8A8 linear: quantized activations `[m, k]` times a [`QuantizedMatrix`] `[n, k]`, producing
586/// f32 `[m, n]`.
587///
588/// `x_scales` carries one dynamic activation scale per row of `x_q`. The i32 accumulator is
589/// exact on every tier; dequantization applies `acc as f32 * (x_scale * w_scale)` in exactly
590/// that order on every tier, so the f32 output of any two tiers is bit-identical, not merely
591/// close. Bias (only `text_projection` carries one) is added after dequantization.
592///
593/// # Panics
594///
595/// Panics on any shape mismatch.
596#[allow(clippy::too_many_arguments)]
597pub fn linear_q8(
598    x_q: &[i8],
599    x_scales: &[f32],
600    weight: &QuantizedMatrix,
601    bias: Option<&[f32]>,
602    m: usize,
603    out: &mut [f32],
604    tier: Int8Tier,
605) {
606    let (n, k) = (weight.n, weight.k);
607    assert_eq!(x_q.len(), m * k, "x_q must be [m, k]");
608    assert_eq!(x_scales.len(), m, "x_scales must be [m]");
609    assert_eq!(out.len(), m * n, "out must be [m, n]");
610    if let Some(bias) = bias {
611        assert_eq!(bias.len(), n, "bias must be [n]");
612    }
613    // Large operations fan out across the persistent team when one exists; the partitioned
614    // result is bit-identical per element, so this is purely a speed dispatch. Small matrices
615    // stay serial — the dispatch handshake would cost more than the work.
616    if n * k >= TEAM_WORK_THRESHOLD_BYTES
617        && let Some(team) = crate::team::armed()
618    {
619        team.linear_q8(x_q, x_scales, weight, bias, m, out, tier);
620        return;
621    }
622
623    // Weight-stationary loop order: each Q8 weight row is streamed exactly once and reused
624    // across all m activation rows, so an m>1 call (prefill, the seq-16 verify pass) does not
625    // re-read the whole matrix m times. Each output element's dot product is unchanged, so this
626    // ordering is bit-identical to the m-outer form.
627    for col in 0..n {
628        let w_row = &weight.data[col * k..(col + 1) * k];
629        let w_scale = weight.scales[col];
630        let bias_term = bias.map(|b| b[col]);
631        for row in 0..m {
632            let x_row = &x_q[row * k..(row + 1) * k];
633            let acc = dot_i32(x_row, w_row, tier);
634            let value = acc as f32 * (x_scales[row] * w_scale);
635            out[row * n + col] = bias_term.map_or(value, |b| value + b);
636        }
637    }
638}
639
640/// Quantizes an f32 activation matrix `[m, k]` per row and runs [`linear_q8`].
641///
642/// This is the drop-in W8A8 counterpart of `f32ref::linear`: same `[m, k] × [n, k]ᵀ → [m, n]`
643/// layout, same bias placement. The row quantization is the canonical symmetric recipe.
644///
645/// # Panics
646///
647/// Panics on any shape mismatch or a non-finite activation.
648pub fn linear_q8_dynamic(
649    x: &[f32],
650    weight: &QuantizedMatrix,
651    bias: Option<&[f32]>,
652    m: usize,
653    out: &mut [f32],
654    tier: Int8Tier,
655) {
656    let k = weight.k;
657    assert_eq!(x.len(), m * k, "x must be [m, k]");
658    let mut x_q = vec![0_i8; m * k];
659    let mut x_scales = vec![0.0_f32; m];
660    for ((x_row, q_row), scale) in x
661        .chunks_exact(k)
662        .zip(x_q.chunks_exact_mut(k))
663        .zip(x_scales.iter_mut())
664    {
665        *scale = quantize_row_q8(x_row, q_row);
666    }
667    linear_q8(&x_q, &x_scales, weight, bias, m, out, tier);
668}
669
670#[cfg(test)]
671mod tests {
672    use super::*;
673
674    /// Deterministic pseudo-random Q8 bytes (SplitMix64), full `[-127, 127]` range.
675    fn pseudo_random_q8(len: usize, seed: u64) -> Vec<i8> {
676        let mut state = seed;
677        (0..len)
678            .map(|_| {
679                state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
680                let mut z = state;
681                z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
682                z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
683                z ^= z >> 31;
684                // Map to [-127, 127]; never -128, matching the converter contract.
685                ((z % 255) as i32 - 127) as i8
686            })
687            .collect()
688    }
689
690    /// The model's real decode GEMV shapes: (n, k) per §7 of the plan.
691    const MODEL_SHAPES: &[(usize, usize)] = &[
692        (2048, 1024), // q_proj / per-depth heads
693        (1024, 1024), // k_proj / v_proj
694        (1024, 2048), // o_proj
695        (3072, 1024), // gate/up_proj, primary head
696        (1024, 3072), // down_proj (binding talker K)
697    ];
698
699    #[test]
700    fn every_tier_is_exactly_equal_in_i32_at_every_model_shape() {
701        for &(n, k) in MODEL_SHAPES {
702            let a = pseudo_random_q8(k, 0x5eed_0001 ^ (n as u64) << 20 ^ k as u64);
703            let w = pseudo_random_q8(n * k, 0x5eed_0002 ^ (n as u64) << 20 ^ k as u64);
704            for row in [0, n / 2, n - 1] {
705                let w_row = &w[row * k..(row + 1) * k];
706                let reference = dot_i32(&a, w_row, Int8Tier::Scalar);
707                for tier in Int8Tier::available() {
708                    assert_eq!(
709                        dot_i32(&a, w_row, tier),
710                        reference,
711                        "tier {} diverged at shape {n}x{k} row {row}",
712                        tier.as_str()
713                    );
714                }
715            }
716        }
717    }
718
719    #[test]
720    fn every_tier_survives_the_all_extreme_reduction_at_the_binding_census_k() {
721        // 127 * 127 * 8192 = 132,120,576 — the S8S8 all-extreme envelope at the largest census K.
722        for k in [2048_usize, 3072, 4608, 7168, 8192] {
723            let a = vec![127_i8; k];
724            let b = vec![127_i8; k];
725            let negative = vec![-127_i8; k];
726            let expected = 127_i64 * 127 * k as i64;
727            for tier in Int8Tier::available() {
728                assert_eq!(
729                    i64::from(dot_i32(&a, &b, tier)),
730                    expected,
731                    "positive all-extreme diverged on {} at K={k}",
732                    tier.as_str()
733                );
734                assert_eq!(
735                    i64::from(dot_i32(&a, &negative, tier)),
736                    -expected,
737                    "negative all-extreme diverged on {} at K={k}",
738                    tier.as_str()
739                );
740            }
741        }
742    }
743
744    #[test]
745    fn tail_lengths_that_defeat_block_boundaries_stay_exact() {
746        // Exercise every SDOT path: <16 (pure tail), 16..64 (single-block loop), 64+tail.
747        for len in [1_usize, 7, 15, 16, 17, 63, 64, 65, 100, 129] {
748            let a = pseudo_random_q8(len, tail_seed(len));
749            let b = pseudo_random_q8(len, tail_seed(len) ^ 1);
750            let reference = dot_i32(&a, &b, Int8Tier::Scalar);
751            for tier in Int8Tier::available() {
752                assert_eq!(
753                    dot_i32(&a, &b, tier),
754                    reference,
755                    "len={len} {}",
756                    tier.as_str()
757                );
758            }
759        }
760    }
761
762    #[test]
763    fn quantizer_matches_the_canonical_converter_semantics() {
764        // Ties-to-even, clamp, zero-row scale, and the -128 exclusion. The cross-crate
765        // byte-identity test against `ftts-artifacts` lives in `ftts-model-qwen`.
766        let row = [
767            -127.0_f32, -126.5, -125.5, -1.5, -0.5, 0.5, 1.5, 125.5, 126.5, 127.0,
768        ];
769        let mut q = [0_i8; 10];
770        let scale = quantize_row_q8(&row, &mut q);
771        assert_eq!(scale.to_bits(), 1.0_f32.to_bits());
772        assert_eq!(q, [-127, -126, -126, -2, 0, 0, 2, 126, 126, 127]);
773
774        let zeros = [0.0_f32; 4];
775        let mut qz = [1_i8; 4];
776        assert_eq!(
777            quantize_row_q8(&zeros, &mut qz).to_bits(),
778            1.0_f32.to_bits()
779        );
780        assert_eq!(qz, [0, 0, 0, 0]);
781
782        let matrix = QuantizedMatrix::quantize(&[2.0, -1.0, 0.0, 3.0], 2, 2);
783        assert_eq!(matrix.scales[0].to_bits(), (2.0_f32 / 127.0).to_bits());
784        assert_eq!(matrix.scales[1].to_bits(), (3.0_f32 / 127.0).to_bits());
785        assert!(matrix.data.iter().all(|&b| b != -128));
786    }
787
788    #[test]
789    fn dynamic_w8a8_linear_tracks_the_f32_reference_within_quant_error() {
790        // Not a parity claim — a sanity bound that the dequant plumbing is wired correctly.
791        let (n, k) = (64_usize, 128_usize);
792        let mut weight = vec![0.0_f32; n * k];
793        let mut x = vec![0.0_f32; k];
794        let mut state = 0x1234_5678_u64;
795        let mut next = || {
796            state = state
797                .wrapping_mul(6_364_136_223_846_793_005)
798                .wrapping_add(1);
799            ((state >> 33) as f32 / (1u64 << 31) as f32) - 1.0
800        };
801        for value in weight.iter_mut() {
802            *value = next();
803        }
804        for value in x.iter_mut() {
805            *value = next();
806        }
807        let quantized = QuantizedMatrix::quantize(&weight, n, k);
808        let mut out_q8 = vec![0.0_f32; n];
809        linear_q8_dynamic(&x, &quantized, None, 1, &mut out_q8, Int8Tier::Autovec);
810
811        let mut out_f32 = vec![0.0_f32; n];
812        crate::f32ref::linear(&x, &weight, None, 1, k, n, &mut out_f32);
813
814        let dot = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>();
815        let cosine = dot(&out_q8, &out_f32)
816            / (dot(&out_q8, &out_q8).sqrt() * dot(&out_f32, &out_f32).sqrt());
817        assert!(
818            cosine > 0.999,
819            "W8A8 dequant plumbing is broken: cosine {cosine}"
820        );
821    }
822
823    #[test]
824    fn tiers_produce_bit_identical_f32_output_not_merely_close() {
825        let (n, k) = (256_usize, 1024_usize);
826        let weight: Vec<f32> = pseudo_random_q8(n * k, 77)
827            .iter()
828            .map(|&b| f32::from(b) / 64.0)
829            .collect();
830        let x: Vec<f32> = pseudo_random_q8(k, 78)
831            .iter()
832            .map(|&b| f32::from(b) / 64.0)
833            .collect();
834        let quantized = QuantizedMatrix::quantize(&weight, n, k);
835        let mut reference = vec![0.0_f32; n];
836        linear_q8_dynamic(&x, &quantized, None, 1, &mut reference, Int8Tier::Scalar);
837        for tier in Int8Tier::available() {
838            let mut out = vec![0.0_f32; n];
839            linear_q8_dynamic(&x, &quantized, None, 1, &mut out, tier);
840            for (index, (a, b)) in reference.iter().zip(&out).enumerate() {
841                assert_eq!(
842                    a.to_bits(),
843                    b.to_bits(),
844                    "tier {} f32 output differs at {index}",
845                    tier.as_str()
846                );
847            }
848        }
849    }
850
851    fn tail_seed(len: usize) -> u64 {
852        0x7a11_0000 ^ len as u64
853    }
854}