Skip to main content

cortiq_engine/
qtensor.rs

1//! QTensor — weight tensor with pluggable storage.
2//!
3//! Two backings, one interface:
4//! - `F32`   — owned dense floats (small models, tests). Every operation
5//!   is bit-identical to the historical `&[f32]` code paths.
6//! - `Mapped` — quantized bytes zero-copy from the CMF mmap (`q8_row` /
7//!   `q8_2f`). The matvec is fused: int8 rows × f32 activations, the
8//!   q8_2f column field folds into a pre-scale of the input
9//!   (`x'[i] = col[i]·x[i]`), so the inner loop is the same i8 dot as
10//!   q8_row. This is what lets a 15B file run in a few GB of RSS.
11//!
12//! Extension point: new dtypes = new match arm here, nothing else moves.
13
14use crate::pool::{matvec_rows, matvec_rows2, Pool};
15use cortiq_core::quant::{f16_to_f32, GROUP_SIZE};
16use cortiq_core::{CmfModel, TensorDtype};
17use std::sync::Arc;
18
19pub enum QTensor {
20    F32 {
21        data: Vec<f32>,
22        rows: usize,
23        cols: usize,
24    },
25    Mapped {
26        model: Arc<CmfModel>,
27        /// Index into the model's tensor directory.
28        idx: usize,
29        dtype: TensorDtype,
30        rows: usize,
31        cols: usize,
32        /// Per-row scales, dequantized to f32 up front (tiny).
33        row_scale: Vec<f32>,
34        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
35        col_field: Vec<f32>,
36    },
37}
38
39impl QTensor {
40    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
41        debug_assert_eq!(data.len(), rows * cols);
42        Self::F32 { data, rows, cols }
43    }
44
45    /// Wrap a directory tensor without dequantizing the payload.
46    /// Falls back to dequantized f32 for dtypes without a fused kernel.
47    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
48        let idx = model
49            .tensors
50            .iter()
51            .position(|t| t.name == name)
52            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
53        let entry = &model.tensors[idx];
54        if entry.shape.len() != 2 {
55            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
56        }
57        let (rows, cols) = (entry.shape[0], entry.shape[1]);
58        let bytes = model.entry_bytes(entry);
59
60        match entry.dtype {
61            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
62                let n = rows * cols;
63                let scales_off = n;
64                let row_scale: Vec<f32> = (0..rows)
65                    .map(|o| {
66                        f16_to_f32(u16::from_le_bytes([
67                            bytes[scales_off + o * 2],
68                            bytes[scales_off + o * 2 + 1],
69                        ]))
70                    })
71                    .collect();
72                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
73                    let col_off = n + rows * 2;
74                    (0..cols)
75                        .map(|i| {
76                            f16_to_f32(u16::from_le_bytes([
77                                bytes[col_off + i * 2],
78                                bytes[col_off + i * 2 + 1],
79                            ]))
80                        })
81                        .collect()
82                } else {
83                    Vec::new()
84                };
85                Ok(Self::Mapped {
86                    model: model.clone(),
87                    idx,
88                    dtype: entry.dtype,
89                    rows,
90                    cols,
91                    row_scale,
92                    col_field,
93                })
94            }
95            // vbit: fused kernel unpacks variable-bit rows from mmap.
96            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
97                model: model.clone(),
98                idx,
99                dtype: entry.dtype,
100                rows,
101                cols,
102                row_scale: Vec::new(),
103                col_field: Vec::new(),
104            }),
105            // q4_block: fused kernel reads nibbles straight from mmap —
106            // a 14B q4 file no longer explodes into ×8 f32 RAM.
107            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
108                model: model.clone(),
109                idx,
110                dtype: entry.dtype,
111                rows,
112                cols,
113                row_scale: Vec::new(),
114                col_field: Vec::new(),
115            }),
116            // No fused kernel yet → dequantize once (correct, more RAM).
117            _ => {
118                let mut data = vec![0.0f32; rows * cols];
119                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
120                Ok(Self::from_f32(data, rows, cols))
121            }
122        }
123    }
124
125    pub fn rows(&self) -> usize {
126        match self {
127            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
128        }
129    }
130
131    pub fn cols(&self) -> usize {
132        match self {
133            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
134        }
135    }
136
137    /// Dense f32 view — only for owned tensors. Masked/sparse execution
138    /// paths require it; quantized weights don't support masks yet.
139    pub fn as_f32(&self) -> Option<&[f32]> {
140        match self {
141            Self::F32 { data, .. } => Some(data),
142            Self::Mapped { .. } => None,
143        }
144    }
145
146    fn quant_bytes(&self) -> &[u8] {
147        match self {
148            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
149            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
150        }
151    }
152
153    /// Dequantize one row into `dst` (embedding lookup).
154    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
155        let cols = self.cols();
156        debug_assert_eq!(dst.len(), cols);
157        match self {
158            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
159            Self::Mapped {
160                dtype,
161                row_scale,
162                col_field,
163                ..
164            } => {
165                if *dtype == TensorDtype::Q4Block {
166                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
167                    let gpr = cols / GROUP_SIZE;
168                    for gi in 0..gpr {
169                        let g = r * gpr + gi;
170                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
171                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
172                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
173                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
174                        }
175                    }
176                    return;
177                }
178                if *dtype == TensorDtype::Vbit {
179                    let bytes = self.quant_bytes();
180                    let rows = self.rows();
181                    let ng = cols / GROUP_SIZE;
182                    let bits = &bytes[..rows];
183                    let sc_off = rows;
184                    let mut off = sc_off + rows * ng * 2;
185                    for rr in 0..r {
186                        off += (cols * bits[rr] as usize + 7) / 8;
187                    }
188                    let b = bits[r] as usize;
189                    let l = ((1usize << (b - 1)) - 1) as f32;
190                    let data = &bytes[off..];
191                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
192                    for (i, d) in dst.iter_mut().enumerate() {
193                        while nbits < b {
194                            acc = (acc << 8) | data[idx] as u64;
195                            idx += 1;
196                            nbits += 8;
197                        }
198                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
199                        nbits -= b;
200                        let so = (r * ng + i / GROUP_SIZE) * 2;
201                        let sv = f16_to_f32(u16::from_le_bytes([
202                            bytes[sc_off + so],
203                            bytes[sc_off + so + 1],
204                        ]));
205                        *d = (u - l) * sv;
206                    }
207                    return;
208                }
209                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
210                let s = row_scale[r];
211                match dtype {
212                    TensorDtype::Q8Row => {
213                        for (d, &b) in dst.iter_mut().zip(q) {
214                            *d = (b as i8) as f32 * s;
215                        }
216                    }
217                    TensorDtype::Q8_2f => {
218                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
219                            *d = (b as i8) as f32 * s * col_field[i];
220                        }
221                    }
222                    _ => unreachable!(),
223                }
224            }
225        }
226    }
227
228    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
229    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
230    /// false for group-packed q4/vbit (column access would unpack whole
231    /// groups — sparse execution falls back to f32 for those).
232    pub fn sparse_col_ok(&self) -> bool {
233        match self {
234            Self::F32 { .. } => true,
235            Self::Mapped { dtype, .. } => {
236                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
237            }
238        }
239    }
240
241    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
242    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
243    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
244    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
245        let inter = self.cols();
246        let hidden = self.rows();
247        debug_assert_eq!(out.len(), hidden);
248        match self {
249            Self::F32 { data, .. } => {
250                for (k, o) in out.iter_mut().enumerate() {
251                    *o += w * data[k * inter + c];
252                }
253            }
254            Self::Mapped {
255                dtype, row_scale, col_field, ..
256            } => {
257                let q = self.quant_bytes();
258                let colf = if *dtype == TensorDtype::Q8_2f {
259                    col_field[c]
260                } else {
261                    1.0
262                };
263                let wc = w * colf;
264                for (k, o) in out.iter_mut().enumerate() {
265                    let b = q[k * inter + c] as i8 as f32;
266                    *o += wc * b * row_scale[k];
267                }
268            }
269        }
270    }
271
272    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
273    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
274    /// into `scratch` first (rare for active-FFN weights).
275    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
276        let cols = self.cols();
277        match self {
278            Self::F32 { data, .. } => {
279                let row = &data[r * cols..(r + 1) * cols];
280                row.iter().zip(x).map(|(w, v)| w * v).sum()
281            }
282            Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
283                TensorDtype::Q8Row => {
284                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
285                    let s: f32 = q.iter().zip(x).map(|(&b, v)| (b as i8) as f32 * v).sum();
286                    s * row_scale[r]
287                }
288                TensorDtype::Q8_2f => {
289                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
290                    let s: f32 = q
291                        .iter()
292                        .zip(x)
293                        .zip(col_field)
294                        .map(|((&b, v), c)| (b as i8) as f32 * v * c)
295                        .sum();
296                    s * row_scale[r]
297                }
298                _ => {
299                    self.row_f32(r, scratch);
300                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
301                }
302            },
303        }
304    }
305
306    /// `out = W · x` (row-major). F32 delegates to the historical
307    /// bit-exact path; Mapped runs the fused int8 kernel.
308    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
309        match self {
310            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
311            Self::Mapped {
312                model,
313                idx,
314                dtype,
315                rows,
316                cols,
317                row_scale,
318                col_field,
319            } => {
320                let _ = (model, idx);
321                if *dtype == TensorDtype::Q4Block {
322                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
323                    return;
324                }
325                if *dtype == TensorDtype::Vbit {
326                    vbitmatvec(self.quant_bytes(), x, *rows, *cols, out, pool);
327                    return;
328                }
329                let xs = prescale(x, col_field, *dtype);
330                // D5: large q8 matrices (lm_head-class) — hybrid
331                // CPU∥GPU: split the rows, both sides compute
332                // SIMULTANEOUSLY (same math, shared prescale).
333                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
334                if *rows >= crate::gpu::min_rows()
335                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
336                    && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
337                    && crate::gpu::enabled_here()
338                {
339                    let frac = std::env::var("CMF_GPU_SPLIT")
340                        .ok()
341                        .and_then(|v| v.parse::<f32>().ok())
342                        .unwrap_or(0.5)
343                        .clamp(0.0, 1.0);
344                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
345                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
346                    let bytes = self.quant_bytes();
347                    let ok = std::thread::scope(|sc| {
348                        let g = sc.spawn(|| {
349                            crate::gpu::q8_matvec_range(
350                                model,
351                                *idx,
352                                cpu_rows,
353                                &row_scale[cpu_rows..],
354                                &xs,
355                                *rows - cpu_rows,
356                                *cols,
357                                out_gpu,
358                            )
359                        });
360                        if cpu_rows > 0 {
361                            qmatvec(
362                                &bytes[..cpu_rows * *cols],
363                                &row_scale[..cpu_rows],
364                                &xs,
365                                cpu_rows,
366                                *cols,
367                                out_cpu,
368                                pool,
369                            );
370                        }
371                        g.join().unwrap_or(false)
372                    });
373                    if ok {
374                        return;
375                    }
376                    // GPU failed — CPU finishes its half.
377                    qmatvec(
378                        &bytes[cpu_rows * *cols..(*rows) * *cols],
379                        &row_scale[cpu_rows..],
380                        &xs,
381                        *rows - cpu_rows,
382                        *cols,
383                        out_gpu,
384                        pool,
385                    );
386                    return;
387                }
388                qmatvec(self.quant_bytes(), row_scale, &xs, *rows, *cols, out, pool);
389            }
390        }
391    }
392
393    /// Fused two-input matvec (MTP verify pair): weights streamed once.
394    pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
395        match self {
396            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
397            Self::Mapped {
398                dtype,
399                rows,
400                cols,
401                row_scale,
402                col_field,
403                ..
404            } => {
405                if *dtype == TensorDtype::Q4Block {
406                    q4matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
407                    q4matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
408                    return;
409                }
410                if *dtype == TensorDtype::Vbit {
411                    vbitmatvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
412                    vbitmatvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
413                    return;
414                }
415                let x1s = prescale(x1, col_field, *dtype);
416                let x2s = prescale(x2, col_field, *dtype);
417                qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
418            }
419        }
420    }
421}
422
423impl QTensor {
424    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
425    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
426    /// to b matvec calls (same dot kernels in the same order); the win —
427    /// the weight row streams from DRAM once per batch, not b times.
428    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
429        let cols = self.cols();
430        let rows = self.rows();
431        debug_assert_eq!(xs_all.len(), b * cols);
432        debug_assert_eq!(out.len(), b * rows);
433        match self {
434            Self::F32 { data, .. } => {
435                let out_addr = SendMut(out.as_mut_ptr());
436                let run = |start: usize, end: usize| {
437                    for o in start..end {
438                        let row = &data[o * cols..(o + 1) * cols];
439                        for bi in 0..b {
440                            let x = &xs_all[bi * cols..(bi + 1) * cols];
441                            let mut acc = 0f32;
442                            for j in 0..cols {
443                                acc += row[j] * x[j];
444                            }
445                            unsafe { *out_addr.at(bi * rows + o) = acc };
446                        }
447                    }
448                };
449                dispatch_rows(pool, rows, &run);
450            }
451            Self::Mapped {
452                dtype,
453                row_scale,
454                col_field,
455                ..
456            } => {
457                if matches!(dtype, TensorDtype::Q4Block | TensorDtype::Vbit) {
458                    // No batch kernel — honest element-wise fallback.
459                    for bi in 0..b {
460                        let x = &xs_all[bi * cols..(bi + 1) * cols];
461                        let (head, tail) = out[bi * rows..].split_at_mut(rows);
462                        let _ = tail;
463                        self.matvec(x, head, pool);
464                    }
465                    return;
466                }
467                let pre: Vec<Vec<f32>> = (0..b)
468                    .map(|bi| {
469                        prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
470                    })
471                    .collect();
472                // D5: large prefill-batch GEMMs — on the GPU (threshold by
473                // work volume: submission carries b×rows×cols MACs).
474                if b >= 8
475                    && b * rows * cols >= 128_000_000
476                    && crate::gpu::enabled_here()
477                {
478                    if let Self::Mapped { model, idx, .. } = self {
479                        let flat: Vec<f32> =
480                            pre.iter().flat_map(|v| v.iter().copied()).collect();
481                        if crate::gpu::q8_matmat(
482                            model, *idx, row_scale, &flat, b, rows, cols, out)
483                        {
484                            return;
485                        }
486                    }
487                }
488                let q = self.quant_bytes();
489                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
490            }
491        }
492    }
493}
494
495/// Batched q8 kernel: same math as qmatvec, the row makes a single
496/// pass from memory for the whole batch.
497fn qmatmat(
498    q: &[u8],
499    row_scale: &[f32],
500    pre: &[Vec<f32>],
501    rows: usize,
502    cols: usize,
503    out: &mut [f32],
504    pool: Option<&Pool>,
505) {
506    let b = pre.len();
507    debug_assert_eq!(out.len(), b * rows);
508    #[cfg(target_arch = "aarch64")]
509    if sdot_enabled() {
510        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
511        let out_addr = SendMut(out.as_mut_ptr());
512        let run = |start: usize, end: usize| {
513            for o in start..end {
514                let row = &q[o * cols..(o + 1) * cols];
515                for (bi, act) in acts.iter().enumerate() {
516                    let v = row_dot_sdot(row, act) * row_scale[o];
517                    unsafe { *out_addr.at(bi * rows + o) = v };
518                }
519            }
520        };
521        dispatch_rows(pool, rows, &run);
522        return;
523    }
524    let out_addr = SendMut(out.as_mut_ptr());
525    let run = |start: usize, end: usize| {
526        for o in start..end {
527            let row = &q[o * cols..(o + 1) * cols];
528            for (bi, x) in pre.iter().enumerate() {
529                let mut acc = 0f32;
530                for j in 0..cols {
531                    acc += (row[j] as i8) as f32 * x[j];
532                }
533                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
534            }
535        }
536    };
537    dispatch_rows(pool, rows, &run);
538}
539
540/// Split rows across pool workers (shared qmatvec pattern).
541fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
542    match pool {
543        Some(pool) if rows >= 256 => {
544            pool.run(&|widx, n| {
545                let chunk = rows.div_ceil(n);
546                let start = widx * chunk;
547                let end = (start + chunk).min(rows);
548                if start < end {
549                    run(start, end);
550                }
551            });
552        }
553        _ => run(0, rows),
554    }
555}
556
557/// Split a q4_block blob into (packed nibbles, f16 group scales).
558fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
559    let groups = rows * cols / GROUP_SIZE;
560    bytes.split_at(groups * 16)
561}
562
563/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
564/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
565/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
566/// into 4 such blocks.
567#[inline(always)]
568fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
569    let mut acc = 0u64;
570    for i in 0..B {
571        acc = (acc << 8) | data[i] as u64;
572    }
573    let mask = (1u64 << B) - 1;
574    let mut out = [0i32; 8];
575    for (k, o) in out.iter_mut().enumerate() {
576        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
577    }
578    out
579}
580
581/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
582/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
583/// MSB-first, byte-padded]. Row data offsets are a prefix sum over
584/// bits — computed once per call, then rows split across the pool.
585fn vbitmatvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
586    debug_assert_eq!(out.len(), rows);
587    let ng = cols / GROUP_SIZE;
588    let bits = &bytes[..rows];
589    let sc_off = rows;
590    let data_off = sc_off + rows * ng * 2;
591    let mut offsets = Vec::with_capacity(rows + 1);
592    let mut off = data_off;
593    for r in 0..rows {
594        offsets.push(off);
595        off += (cols * bits[r] as usize + 7) / 8;
596    }
597    offsets.push(off);
598
599    let offsets = &offsets;
600
601    // SDOT path: unpack the row to centered i8 once, then per-group
602    // int8 dot against the quantized activations — same A8W8 contract
603    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
604    #[cfg(target_arch = "aarch64")]
605    if sdot_enabled() {
606        let act = split_act(x);
607        let act = &act;
608        let row_dot = move |r: usize| -> f32 {
609            let b = bits[r] as usize;
610            let l = ((1i32 << (b - 1)) - 1) as i32;
611            let mask = (1u64 << b) - 1;
612            let data = &bytes[offsets[r]..offsets[r + 1]];
613            if b == 8 {
614                // u−L reaches 128 → does not fit i8; exact f32 path.
615                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
616                let mut dot = 0f32;
617                for g in 0..ng {
618                    let so = (r * ng + g) * 2;
619                    let sgf = f16_to_f32(u16::from_le_bytes([
620                        bytes[sc_off + so],
621                        bytes[sc_off + so + 1],
622                    ]));
623                    let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
624                    let mut gd = 0f32;
625                    for &xv in xg.iter() {
626                        if nbits < 8 {
627                            acc = (acc << 8) | data[idx] as u64;
628                            idx += 1;
629                            nbits += 8;
630                        }
631                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
632                        nbits -= 8;
633                        gd += (u - l) as f32 * xv;
634                    }
635                    dot += gd * sgf;
636                }
637                return dot;
638            }
639            let mut buf = vec![0u8; cols]; // centered i8 stored as u8 bits
640            #[inline(always)]
641            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
642                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
643                    let u = unpack8::<B>(&data[blk * B..]);
644                    for k in 0..8 {
645                        chunk[k] = (u[k] - l) as i8 as u8;
646                    }
647                }
648            }
649            let _ = mask;
650            match b {
651                3 => fill::<3>(data, l, &mut buf),
652                4 => fill::<4>(data, l, &mut buf),
653                5 => fill::<5>(data, l, &mut buf),
654                6 => fill::<6>(data, l, &mut buf),
655                _ => unreachable!(),
656            }
657            let mut dot = 0f32;
658            for g in 0..ng {
659                let so = (r * ng + g) * 2;
660                let s = f16_to_f32(u16::from_le_bytes([
661                    bytes[sc_off + so],
662                    bytes[sc_off + so + 1],
663                ]));
664                let d = unsafe {
665                    dot_i8_sdot(
666                        &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
667                        &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
668                    )
669                } as f32
670                    * act.sx;
671                dot += d * s;
672            }
673            for &(j, xv) in &act.outliers {
674                let so = (r * ng + j / GROUP_SIZE) * 2;
675                let s = f16_to_f32(u16::from_le_bytes([
676                    bytes[sc_off + so],
677                    bytes[sc_off + so + 1],
678                ]));
679                // xq is zeroed at outlier slots — add the exact term.
680                dot += (buf[j] as i8) as f32 * s * xv;
681            }
682            dot
683        };
684        match pool {
685            Some(pool) if rows >= 256 => {
686                let out_addr = SendMut(out.as_mut_ptr());
687                pool.run(&move |widx, n| {
688                    let chunk = rows.div_ceil(n);
689                    let start = widx * chunk;
690                    let end = (start + chunk).min(rows);
691                    for r in start..end {
692                        // SAFETY: disjoint row ranges per worker.
693                        unsafe { *out_addr.at(r) = row_dot(r) };
694                    }
695                });
696            }
697            _ => {
698                for (r, dst) in out.iter_mut().enumerate() {
699                    *dst = row_dot(r);
700                }
701            }
702        }
703        return;
704    }
705
706    // Per-bit-width specialized inner loops: the compiler unrolls the
707    // constant shifts (the generic bit-buffer loop was branch-bound —
708    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
709    #[inline(always)]
710    fn dot_row<const B: usize>(
711        data: &[u8],
712        bytes: &[u8],
713        sc_off: usize,
714        r: usize,
715        ng: usize,
716        x: &[f32],
717    ) -> f32 {
718        let l = ((1i32 << (B - 1)) - 1) as f32;
719        let gbytes = GROUP_SIZE * B / 8;
720        let mut dot = 0f32;
721        for g in 0..ng {
722            let so = (r * ng + g) * 2;
723            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
724            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
725            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
726            let mut gd = 0f32;
727            for blk in 0..GROUP_SIZE / 8 {
728                let u = unpack8::<B>(&gd0[blk * B..]);
729                let xb = &xg[blk * 8..blk * 8 + 8];
730                for k in 0..8 {
731                    gd += (u[k] as f32 - l) * xb[k];
732                }
733            }
734            dot += gd * s;
735        }
736        dot
737    }
738    let row_dot = move |r: usize| -> f32 {
739        let data = &bytes[offsets[r]..offsets[r + 1]];
740        match bits[r] {
741            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
742            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
743            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
744            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
745            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
746            b => unreachable!("vbit bit-width {b} (validated at load)"),
747        }
748    };
749    match pool {
750        Some(pool) if rows >= 256 => {
751            let out_addr = SendMut(out.as_mut_ptr());
752            pool.run(&move |widx, n| {
753                let chunk = rows.div_ceil(n);
754                let start = widx * chunk;
755                let end = (start + chunk).min(rows);
756                for r in start..end {
757                    // SAFETY: disjoint row ranges per worker.
758                    unsafe { *out_addr.at(r) = row_dot(r) };
759                }
760            });
761        }
762        _ => {
763            for (r, dst) in out.iter_mut().enumerate() {
764                *dst = row_dot(r);
765            }
766        }
767    }
768}
769
770/// Fused q4_block matvec straight from the mapped bytes. Scalar inner
771/// loop for now — correctness (and the ×8 RAM fix) first; a nibble
772/// SDOT port (vmfcore: +23%) is the follow-up optimization.
773fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
774    debug_assert_eq!(out.len(), rows);
775    let (packed, scales) = q4_split(bytes, rows, cols);
776    let gpr = cols / GROUP_SIZE;
777    let row_dot = |r: usize| -> f32 {
778        let mut acc = 0f32;
779        for gi in 0..gpr {
780            let g = r * gpr + gi;
781            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
782            let pk = &packed[g * 16..(g + 1) * 16];
783            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
784            let mut ga = 0f32;
785            for (k, &b) in pk.iter().enumerate() {
786                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
787                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
788            }
789            acc += ga * s;
790        }
791        acc
792    };
793    match pool {
794        Some(pool) if rows >= 256 => {
795            let out_addr = SendMut(out.as_mut_ptr());
796            pool.run(&move |widx, n| {
797                let chunk = rows.div_ceil(n);
798                let start = widx * chunk;
799                let end = (start + chunk).min(rows);
800                for r in start..end {
801                    // SAFETY: disjoint row ranges per worker.
802                    unsafe { *out_addr.at(r) = row_dot(r) };
803                }
804            });
805        }
806        _ => {
807            for (r, dst) in out.iter_mut().enumerate() {
808                *dst = row_dot(r);
809            }
810        }
811    }
812}
813
814pub(crate) fn prescale(x: &[f32], col_field: &[f32], dtype: TensorDtype) -> Vec<f32> {
815    if dtype == TensorDtype::Q8_2f {
816        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
817    } else {
818        x.to_vec()
819    }
820}
821
822// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
823
824/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
825/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
826fn sdot_enabled() -> bool {
827    use std::sync::OnceLock;
828    static ON: OnceLock<bool> = OnceLock::new();
829    *ON.get_or_init(|| {
830        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
831        #[cfg(target_arch = "aarch64")]
832        {
833            want && std::arch::is_aarch64_feature_detected!("dotprod")
834        }
835        #[cfg(not(target_arch = "aarch64"))]
836        {
837            let _ = want;
838            false
839        }
840    })
841}
842
843/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
844/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
845/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
846/// matvec, shared by all rows/workers.
847struct SplitAct {
848    xq: Vec<i8>,
849    sx: f32,
850    outliers: Vec<(usize, f32)>,
851}
852
853fn split_act(x: &[f32]) -> SplitAct {
854    let n = x.len();
855    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
856    let thr = 8.0 * rms;
857    let outliers: Vec<(usize, f32)> = (0..n)
858        .filter(|&j| x[j].abs() > thr)
859        .map(|j| (j, x[j]))
860        .collect();
861    let mut xb = x.to_vec();
862    for &(j, _) in &outliers {
863        xb[j] = 0.0;
864    }
865    let amax = xb.iter().fold(0f32, |m, &v| m.max(v.abs()));
866    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
867    let inv = 1.0 / sx;
868    let xq = xb
869        .iter()
870        .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8)
871        .collect();
872    SplitAct { xq, sx, outliers }
873}
874
875/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
876/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
877#[cfg(target_arch = "aarch64")]
878#[target_feature(enable = "neon,dotprod")]
879unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
880    // SAFETY: callers uphold slice-length contracts (see call sites).
881    unsafe {
882        use core::arch::aarch64::*;
883        use core::arch::asm;
884        let wp = w.as_ptr() as *const i8;
885        let n = w.len();
886        let (mut a0, mut a1, mut a2, mut a3) =
887            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
888        let mut i = 0;
889        while i + 64 <= n {
890            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
891            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
892            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
893            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
894            asm!(
895                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
896                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
897                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
898                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
899                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
900                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
901                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
902                options(pure, nomem, nostack),
903            );
904            i += 64;
905        }
906        while i + 16 <= n {
907            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
908            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
909                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
910            i += 16;
911        }
912        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
913        while i < n {
914            s += (*wp.add(i)) as i32 * xq[i] as i32;
915            i += 1;
916        }
917        s
918}
919}
920
921/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
922/// loaded once and reused, 4 independent accumulators hide sdot latency
923/// (port of vmfcore `dot_i8_sdot_4rows`).
924#[cfg(target_arch = "aarch64")]
925#[target_feature(enable = "neon,dotprod")]
926unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
927    // SAFETY: callers uphold slice-length contracts (see call sites).
928    unsafe {
929        use core::arch::aarch64::*;
930        use core::arch::asm;
931        let n = xq.len();
932        let px = xq.as_ptr();
933        let (p0, p1, p2, p3) = (
934            w0.as_ptr() as *const i8,
935            w1.as_ptr() as *const i8,
936            w2.as_ptr() as *const i8,
937            w3.as_ptr() as *const i8,
938        );
939        let (mut a0, mut a1, mut a2, mut a3) =
940            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
941        let mut i = 0;
942        while i + 16 <= n {
943            let x = vld1q_s8(px.add(i));
944            let v0 = vld1q_s8(p0.add(i));
945            let v1 = vld1q_s8(p1.add(i));
946            let v2 = vld1q_s8(p2.add(i));
947            let v3 = vld1q_s8(p3.add(i));
948            asm!(
949                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
950                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
951                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
952                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
953                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
954                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
955                options(pure, nomem, nostack),
956            );
957            i += 16;
958        }
959        let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
960        while i < n {
961            let xi = *px.add(i) as i32;
962            r[0] += (*p0.add(i)) as i32 * xi;
963            r[1] += (*p1.add(i)) as i32 * xi;
964            r[2] += (*p2.add(i)) as i32 * xi;
965            r[3] += (*p3.add(i)) as i32 * xi;
966            i += 1;
967        }
968        r
969}
970}
971
972/// SDOT row dot with exact outlier correction:
973/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
974#[cfg(target_arch = "aarch64")]
975#[inline]
976fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
977    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
978    for &(j, xv) in &act.outliers {
979        acc += (row[j] as i8) as f32 * xv;
980    }
981    acc
982}
983
984// ───────────────────── fused int8 kernels ─────────────────────
985
986/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
987/// ≈9× scalar), scalar elsewhere.
988#[inline]
989fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
990    #[cfg(target_arch = "aarch64")]
991    unsafe {
992        return dot_i8_f32_neon(w, x);
993    }
994    #[allow(unreachable_code)]
995    {
996        let mut sum = 0.0f32;
997        for (j, &b) in w.iter().enumerate() {
998            sum += (b as i8) as f32 * x[j];
999        }
1000        sum
1001    }
1002}
1003
1004#[cfg(target_arch = "aarch64")]
1005#[target_feature(enable = "neon")]
1006unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
1007    // SAFETY: callers uphold slice-length contracts (see call sites).
1008    unsafe {
1009        use core::arch::aarch64::*;
1010        let n = x.len();
1011        let wp = w.as_ptr() as *const i8;
1012        let xp = x.as_ptr();
1013        let (mut a0, mut a1, mut a2, mut a3) =
1014            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1015        let mut j = 0usize;
1016        while j + 16 <= n {
1017            let wb = vld1q_s8(wp.add(j));
1018            let lo = vmovl_s8(vget_low_s8(wb));
1019            let hi = vmovl_s8(vget_high_s8(wb));
1020            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
1021            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
1022            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
1023            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
1024            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
1025            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
1026            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
1027            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
1028            j += 16;
1029        }
1030        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
1031        while j < n {
1032            sum += (*wp.add(j)) as f32 * *xp.add(j);
1033            j += 1;
1034        }
1035        sum
1036}
1037}
1038
1039#[allow(clippy::too_many_arguments)]
1040fn qmatvec(
1041    q: &[u8],
1042    row_scale: &[f32],
1043    xs: &[f32],
1044    rows: usize,
1045    cols: usize,
1046    out: &mut [f32],
1047    pool: Option<&Pool>,
1048) {
1049    debug_assert_eq!(out.len(), rows);
1050
1051    #[cfg(target_arch = "aarch64")]
1052    if sdot_enabled() {
1053        let act = split_act(xs);
1054        let out_addr = SendMut(out.as_mut_ptr());
1055        let run_range = |start: usize, end: usize| {
1056            let mut o = start;
1057            while o + 4 <= end {
1058                let r = unsafe {
1059                    dot_i8_sdot_4rows(
1060                        &q[o * cols..(o + 1) * cols],
1061                        &q[(o + 1) * cols..(o + 2) * cols],
1062                        &q[(o + 2) * cols..(o + 3) * cols],
1063                        &q[(o + 3) * cols..(o + 4) * cols],
1064                        &act.xq,
1065                    )
1066                };
1067                for k in 0..4 {
1068                    let mut acc = r[k] as f32 * act.sx;
1069                    for &(j, xv) in &act.outliers {
1070                        acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
1071                    }
1072                    // SAFETY: disjoint row ranges per worker.
1073                    unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
1074                }
1075                o += 4;
1076            }
1077            while o < end {
1078                let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], &act) * row_scale[o];
1079                unsafe { *out_addr.at(o) = v };
1080                o += 1;
1081            }
1082        };
1083        match pool {
1084            Some(pool) if rows >= 256 => {
1085                pool.run(&|widx, n| {
1086                    let chunk = rows.div_ceil(n);
1087                    let start = widx * chunk;
1088                    run_range(start, (start + chunk).min(rows));
1089                });
1090            }
1091            _ => run_range(0, rows),
1092        }
1093        return;
1094    }
1095
1096    let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
1097    match pool {
1098        Some(pool) if rows >= 256 => {
1099            let out_addr = SendMut(out.as_mut_ptr());
1100            pool.run(&move |widx, n| {
1101                let chunk = rows.div_ceil(n);
1102                let start = widx * chunk;
1103                let end = (start + chunk).min(rows);
1104                for o in start..end {
1105                    // SAFETY: disjoint row ranges per worker.
1106                    unsafe { *out_addr.at(o) = row_dot(o) };
1107                }
1108            });
1109        }
1110        _ => {
1111            for (o, dst) in out.iter_mut().enumerate() {
1112                *dst = row_dot(o);
1113            }
1114        }
1115    }
1116}
1117
1118#[allow(clippy::too_many_arguments)]
1119fn qmatvec2(
1120    q: &[u8],
1121    row_scale: &[f32],
1122    x1: &[f32],
1123    x2: &[f32],
1124    rows: usize,
1125    cols: usize,
1126    o1: &mut [f32],
1127    o2: &mut [f32],
1128    pool: Option<&Pool>,
1129) {
1130    #[cfg(target_arch = "aarch64")]
1131    if sdot_enabled() {
1132        let a1s = split_act(x1);
1133        let a2s = split_act(x2);
1134        let p1 = SendMut(o1.as_mut_ptr());
1135        let p2 = SendMut(o2.as_mut_ptr());
1136        let run_range = |start: usize, end: usize| {
1137            for o in start..end {
1138                let row = &q[o * cols..(o + 1) * cols];
1139                // SAFETY: disjoint row ranges per worker.
1140                unsafe {
1141                    *p1.at(o) = row_dot_sdot(row, &a1s) * row_scale[o];
1142                    *p2.at(o) = row_dot_sdot(row, &a2s) * row_scale[o];
1143                }
1144            }
1145        };
1146        match pool {
1147            Some(pool) if rows >= 256 => {
1148                pool.run(&|widx, n| {
1149                    let chunk = rows.div_ceil(n);
1150                    let start = widx * chunk;
1151                    run_range(start, (start + chunk).min(rows));
1152                });
1153            }
1154            _ => run_range(0, rows),
1155        }
1156        return;
1157    }
1158
1159    let row_dots = |o: usize| -> (f32, f32) {
1160        let row = &q[o * cols..(o + 1) * cols];
1161        (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
1162    };
1163    match pool {
1164        Some(pool) if rows >= 256 => {
1165            let p1 = SendMut(o1.as_mut_ptr());
1166            let p2 = SendMut(o2.as_mut_ptr());
1167            pool.run(&move |widx, n| {
1168                let chunk = rows.div_ceil(n);
1169                let start = widx * chunk;
1170                let end = (start + chunk).min(rows);
1171                for o in start..end {
1172                    let (s1, s2) = row_dots(o);
1173                    // SAFETY: disjoint row ranges per worker.
1174                    unsafe {
1175                        *p1.at(o) = s1;
1176                        *p2.at(o) = s2;
1177                    }
1178                }
1179            });
1180        }
1181        _ => {
1182            for o in 0..rows {
1183                let (s1, s2) = row_dots(o);
1184                o1[o] = s1;
1185                o2[o] = s2;
1186            }
1187        }
1188    }
1189}
1190
1191#[derive(Clone, Copy)]
1192struct SendMut(*mut f32);
1193unsafe impl Send for SendMut {}
1194unsafe impl Sync for SendMut {}
1195
1196impl SendMut {
1197    #[inline]
1198    fn at(self, i: usize) -> *mut f32 {
1199        unsafe { self.0.add(i) }
1200    }
1201}
1202
1203#[cfg(test)]
1204mod tests {
1205    use super::*;
1206
1207    #[test]
1208    fn f32_matvec_matches_matvec_rows_bitexact() {
1209        let (rows, cols) = (300, 40);
1210        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
1211        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
1212        let qt = QTensor::from_f32(w.clone(), rows, cols);
1213
1214        let mut a = vec![0.0f32; rows];
1215        matvec_rows(None, &w, &x, &mut a);
1216        let mut b = vec![0.0f32; rows];
1217        qt.matvec(&x, &mut b, None);
1218        assert_eq!(a, b);
1219    }
1220
1221    #[test]
1222    fn sdot_kernel_exact_on_grid() {
1223        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
1224        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
1225        // exact f32 dot to float rounding. This isolates kernel
1226        // correctness from quantization noise.
1227        eprintln!("sdot_enabled = {}", sdot_enabled());
1228        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
1229        let w: Vec<u8> = (0..rows * cols)
1230            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
1231            .collect();
1232        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
1233        let x: Vec<f32> = (0..cols)
1234            .map(|i| match i % 3 {
1235                0 => 1.0,
1236                1 => -1.0,
1237                _ => 0.0,
1238            })
1239            .collect();
1240        let mut a = vec![0.0f32; rows];
1241        qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
1242        for o in 0..rows {
1243            let mut acc = 0.0f32;
1244            for j in 0..cols {
1245                acc += (w[o * cols + j] as i8) as f32 * x[j];
1246            }
1247            let expect = acc * scales[o];
1248            assert!(
1249                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
1250                "row {o}: {} vs {expect}",
1251                a[o]
1252            );
1253        }
1254    }
1255
1256    #[test]
1257    fn sdot_a8w8_noise_is_bounded() {
1258        // Off-grid activations: A8 quantization noise must stay small in
1259        // relative L2 over the whole output (realistic accuracy contract;
1260        // vmfcore measured argmax-identical decode on real models).
1261        let (rows, cols) = (16, 512);
1262        let w: Vec<u8> = (0..rows * cols)
1263            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
1264            .collect();
1265        let scales = vec![0.01f32; rows];
1266        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
1267        let mut a = vec![0.0f32; rows];
1268        qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
1269        let (mut num, mut den) = (0f64, 0f64);
1270        for o in 0..rows {
1271            let mut acc = 0.0f32;
1272            for j in 0..cols {
1273                acc += (w[o * cols + j] as i8) as f32 * x[j];
1274            }
1275            let expect = acc * scales[o];
1276            num += ((a[o] - expect) as f64).powi(2);
1277            den += (expect as f64).powi(2);
1278        }
1279        let rel = (num / den.max(1e-12)).sqrt();
1280        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
1281    }
1282
1283    #[test]
1284    fn i8_dot_neon_matches_scalar() {
1285        let n = 100;
1286        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
1287        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
1288        let mut scalar = 0.0f32;
1289        for j in 0..n {
1290            scalar += (w[j] as i8) as f32 * x[j];
1291        }
1292        let fast = dot_i8_f32(&w, &x);
1293        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
1294    }
1295
1296    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
1297    #[test]
1298    fn vbitmatvec_matches_full_dequant() {
1299        let (rows, cols) = (6, 64);
1300        let ng = cols / GROUP_SIZE;
1301        // Hand-craft: bits per row, f16 scales, packed rows.
1302        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
1303        let mut bytes = bits.clone();
1304        for g in 0..rows * ng {
1305            let s = 0.02 + 0.001 * g as f32;
1306            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
1307        }
1308        for r in 0..rows {
1309            let b = bits[r] as usize;
1310            let (mut acc, mut nb) = (0u64, 0usize);
1311            let mut rowbytes = Vec::new();
1312            for i in 0..cols {
1313                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
1314                acc = (acc << b) | v;
1315                nb += b;
1316                while nb >= 8 {
1317                    nb -= 8;
1318                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
1319                }
1320            }
1321            if nb > 0 {
1322                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
1323            }
1324            bytes.extend_from_slice(&rowbytes);
1325        }
1326        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
1327
1328        let mut reference = vec![0f32; rows * cols];
1329        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
1330        let mut expect = vec![0f32; rows];
1331        for r in 0..rows {
1332            expect[r] = reference[r * cols..(r + 1) * cols]
1333                .iter()
1334                .zip(&x)
1335                .map(|(w, xv)| w * xv)
1336                .sum();
1337        }
1338        let mut got = vec![0f32; rows];
1339        vbitmatvec(&bytes, &x, rows, cols, &mut got, None);
1340        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
1341        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
1342        // the golden-parity gate).
1343        let tol = if sdot_enabled() { 6e-2 } else { 1e-4 };
1344        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
1345        for r in 0..rows {
1346            assert!(
1347                (got[r] - expect[r]).abs() < tol * scale,
1348                "row {r}: {} vs {}",
1349                got[r],
1350                expect[r]
1351            );
1352        }
1353    }
1354
1355    /// Fused q4 matvec must match the reference full-dequant + dense
1356    /// matvec bit-for-bit in structure (same f32 math, group order).
1357    #[test]
1358    fn q4matvec_matches_full_dequant() {
1359        let (rows, cols) = (8, 64);
1360        let groups = rows * cols / GROUP_SIZE;
1361        // Hand-craft a q4_block blob: nibbles then f16 scales.
1362        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
1363        for i in 0..groups * 16 {
1364            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
1365        }
1366        for g in 0..groups {
1367            let s = 0.01 + 0.003 * g as f32;
1368            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
1369        }
1370        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
1371
1372        let mut reference = vec![0.0f32; rows * cols];
1373        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
1374        let mut expect = vec![0.0f32; rows];
1375        for r in 0..rows {
1376            expect[r] = reference[r * cols..(r + 1) * cols]
1377                .iter()
1378                .zip(&x)
1379                .map(|(w, xv)| w * xv)
1380                .sum();
1381        }
1382
1383        let mut got = vec![0.0f32; rows];
1384        q4matvec(&bytes, &x, rows, cols, &mut got, None);
1385        for r in 0..rows {
1386            assert!(
1387                (got[r] - expect[r]).abs() < 1e-4 * expect[r].abs().max(1.0),
1388                "row {r}: {} vs {}",
1389                got[r],
1390                expect[r]
1391            );
1392        }
1393    }
1394}