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        /// Vbit only: byte offset of each row's packed data within the
37        /// tensor blob (`[rows + 1]`, computed once at load — the per-
38        /// matvec prefix scan over row bit-widths was O(rows) each call).
39        vbit_offsets: Vec<usize>,
40    },
41}
42
43/// Prefix-sum of vbit row payload offsets (absolute within the tensor
44/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
45fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
46    let ng = cols / GROUP_SIZE;
47    let bits = &bytes[..rows];
48    let mut offsets = Vec::with_capacity(rows + 1);
49    let mut off = rows + rows * ng * 2;
50    for r in 0..rows {
51        offsets.push(off);
52        off += (cols * bits[r] as usize + 7) / 8;
53    }
54    offsets.push(off);
55    offsets
56}
57
58impl QTensor {
59    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
60        debug_assert_eq!(data.len(), rows * cols);
61        Self::F32 { data, rows, cols }
62    }
63
64    /// Wrap a directory tensor without dequantizing the payload.
65    /// Falls back to dequantized f32 for dtypes without a fused kernel.
66    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
67        // Indexed lookup: the linear directory scan made pipeline build
68        // O(N²) on MoE/skills files with thousands of tensors.
69        let idx = model
70            .tensor_index(name)
71            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
72        let entry = &model.tensors[idx];
73        if entry.shape.len() != 2 {
74            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
75        }
76        let (rows, cols) = (entry.shape[0], entry.shape[1]);
77        let bytes = model.entry_bytes(entry);
78
79        match entry.dtype {
80            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
81                let n = rows * cols;
82                let scales_off = n;
83                let row_scale: Vec<f32> = (0..rows)
84                    .map(|o| {
85                        f16_to_f32(u16::from_le_bytes([
86                            bytes[scales_off + o * 2],
87                            bytes[scales_off + o * 2 + 1],
88                        ]))
89                    })
90                    .collect();
91                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
92                    let col_off = n + rows * 2;
93                    (0..cols)
94                        .map(|i| {
95                            f16_to_f32(u16::from_le_bytes([
96                                bytes[col_off + i * 2],
97                                bytes[col_off + i * 2 + 1],
98                            ]))
99                        })
100                        .collect()
101                } else {
102                    Vec::new()
103                };
104                Ok(Self::Mapped {
105                    model: model.clone(),
106                    idx,
107                    dtype: entry.dtype,
108                    rows,
109                    cols,
110                    row_scale,
111                    col_field,
112                    vbit_offsets: Vec::new(),
113                })
114            }
115            // vbit: fused kernel unpacks variable-bit rows from mmap.
116            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
117                model: model.clone(),
118                idx,
119                dtype: entry.dtype,
120                rows,
121                cols,
122                row_scale: Vec::new(),
123                col_field: Vec::new(),
124                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
125            }),
126            // q4_block: fused kernel reads nibbles straight from mmap —
127            // a 14B q4 file no longer explodes into ×8 f32 RAM.
128            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
129                model: model.clone(),
130                idx,
131                dtype: entry.dtype,
132                rows,
133                cols,
134                row_scale: Vec::new(),
135                col_field: Vec::new(),
136                vbit_offsets: Vec::new(),
137            }),
138            // No fused kernel yet → dequantize once (correct, more RAM).
139            _ => {
140                let mut data = vec![0.0f32; rows * cols];
141                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
142                Ok(Self::from_f32(data, rows, cols))
143            }
144        }
145    }
146
147    pub fn rows(&self) -> usize {
148        match self {
149            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
150        }
151    }
152
153    pub fn cols(&self) -> usize {
154        match self {
155            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
156        }
157    }
158
159    /// Dense f32 view — only for owned tensors. Masked/sparse execution
160    /// paths require it; quantized weights don't support masks yet.
161    pub fn as_f32(&self) -> Option<&[f32]> {
162        match self {
163            Self::F32 { data, .. } => Some(data),
164            Self::Mapped { .. } => None,
165        }
166    }
167
168    fn quant_bytes(&self) -> &[u8] {
169        match self {
170            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
171            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
172        }
173    }
174
175    /// Dequantize one row into `dst` (embedding lookup).
176    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
177        let cols = self.cols();
178        debug_assert_eq!(dst.len(), cols);
179        match self {
180            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
181            Self::Mapped {
182                dtype,
183                row_scale,
184                col_field,
185                vbit_offsets,
186                ..
187            } => {
188                if *dtype == TensorDtype::Q4Block {
189                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
190                    let gpr = cols / GROUP_SIZE;
191                    for gi in 0..gpr {
192                        let g = r * gpr + gi;
193                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
194                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
195                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
196                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
197                        }
198                    }
199                    return;
200                }
201                if *dtype == TensorDtype::Vbit {
202                    let bytes = self.quant_bytes();
203                    let rows = self.rows();
204                    let ng = cols / GROUP_SIZE;
205                    let bits = &bytes[..rows];
206                    let sc_off = rows;
207                    // Precomputed at load — embedding lookup used to scan
208                    // the bit-widths of every preceding row (O(token_id)).
209                    let off = vbit_offsets[r];
210                    let b = bits[r] as usize;
211                    let l = ((1usize << (b - 1)) - 1) as f32;
212                    let data = &bytes[off..];
213                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
214                    for (i, d) in dst.iter_mut().enumerate() {
215                        while nbits < b {
216                            acc = (acc << 8) | data[idx] as u64;
217                            idx += 1;
218                            nbits += 8;
219                        }
220                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
221                        nbits -= b;
222                        let so = (r * ng + i / GROUP_SIZE) * 2;
223                        let sv = f16_to_f32(u16::from_le_bytes([
224                            bytes[sc_off + so],
225                            bytes[sc_off + so + 1],
226                        ]));
227                        *d = (u - l) * sv;
228                    }
229                    return;
230                }
231                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
232                let s = row_scale[r];
233                match dtype {
234                    TensorDtype::Q8Row => {
235                        for (d, &b) in dst.iter_mut().zip(q) {
236                            *d = (b as i8) as f32 * s;
237                        }
238                    }
239                    TensorDtype::Q8_2f => {
240                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
241                            *d = (b as i8) as f32 * s * col_field[i];
242                        }
243                    }
244                    _ => unreachable!(),
245                }
246            }
247        }
248    }
249
250    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
251    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
252    /// false for group-packed q4/vbit (column access would unpack whole
253    /// groups — sparse execution falls back to f32 for those).
254    pub fn sparse_col_ok(&self) -> bool {
255        match self {
256            Self::F32 { .. } => true,
257            Self::Mapped { dtype, .. } => {
258                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
259            }
260        }
261    }
262
263    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
264    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
265    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
266    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
267        let inter = self.cols();
268        let hidden = self.rows();
269        debug_assert_eq!(out.len(), hidden);
270        match self {
271            Self::F32 { data, .. } => {
272                for (k, o) in out.iter_mut().enumerate() {
273                    *o += w * data[k * inter + c];
274                }
275            }
276            Self::Mapped {
277                dtype, row_scale, col_field, ..
278            } => {
279                let q = self.quant_bytes();
280                let colf = if *dtype == TensorDtype::Q8_2f {
281                    col_field[c]
282                } else {
283                    1.0
284                };
285                let wc = w * colf;
286                for (k, o) in out.iter_mut().enumerate() {
287                    let b = q[k * inter + c] as i8 as f32;
288                    *o += wc * b * row_scale[k];
289                }
290            }
291        }
292    }
293
294    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
295    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
296    /// into `scratch` first (rare for active-FFN weights).
297    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
298        let cols = self.cols();
299        match self {
300            Self::F32 { data, .. } => {
301                let row = &data[r * cols..(r + 1) * cols];
302                row.iter().zip(x).map(|(w, v)| w * v).sum()
303            }
304            Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
305                TensorDtype::Q8Row => {
306                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
307                    dot_i8_f32(q, x) * row_scale[r]
308                }
309                TensorDtype::Q8_2f => {
310                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
311                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
312                }
313                _ => {
314                    self.row_f32(r, scratch);
315                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
316                }
317            },
318        }
319    }
320
321    /// `out = W · x` (row-major). F32 delegates to the historical
322    /// bit-exact path; Mapped runs the fused int8 kernel.
323    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
324        match self {
325            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
326            Self::Mapped {
327                model,
328                idx,
329                dtype,
330                rows,
331                cols,
332                row_scale,
333                col_field,
334                vbit_offsets,
335            } => {
336                let _ = (model, idx);
337                if *dtype == TensorDtype::Q4Block {
338                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
339                    return;
340                }
341                if *dtype == TensorDtype::Vbit {
342                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
343                    return;
344                }
345                let xs = prescale(x, col_field, *dtype);
346                // D5: large q8 matrices (lm_head-class) — hybrid
347                // CPU∥GPU: split the rows, both sides compute
348                // SIMULTANEOUSLY (same math, shared prescale).
349                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
350                if *rows >= crate::gpu::min_rows()
351                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
352                    && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
353                    && crate::gpu::enabled_here()
354                {
355                    let frac = std::env::var("CMF_GPU_SPLIT")
356                        .ok()
357                        .and_then(|v| v.parse::<f32>().ok())
358                        .unwrap_or(0.5)
359                        .clamp(0.0, 1.0);
360                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
361                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
362                    let bytes = self.quant_bytes();
363                    let ok = std::thread::scope(|sc| {
364                        let g = sc.spawn(|| {
365                            crate::gpu::q8_matvec_range(
366                                model,
367                                *idx,
368                                cpu_rows,
369                                &row_scale[cpu_rows..],
370                                &xs,
371                                *rows - cpu_rows,
372                                *cols,
373                                out_gpu,
374                            )
375                        });
376                        if cpu_rows > 0 {
377                            qmatvec(
378                                &bytes[..cpu_rows * *cols],
379                                &row_scale[..cpu_rows],
380                                &xs,
381                                cpu_rows,
382                                *cols,
383                                out_cpu,
384                                pool,
385                            );
386                        }
387                        g.join().unwrap_or(false)
388                    });
389                    if ok {
390                        return;
391                    }
392                    // GPU failed — CPU finishes its half.
393                    qmatvec(
394                        &bytes[cpu_rows * *cols..(*rows) * *cols],
395                        &row_scale[cpu_rows..],
396                        &xs,
397                        *rows - cpu_rows,
398                        *cols,
399                        out_gpu,
400                        pool,
401                    );
402                    return;
403                }
404                qmatvec(self.quant_bytes(), row_scale, &xs, *rows, *cols, out, pool);
405            }
406        }
407    }
408
409    /// Fused two-input matvec (MTP verify pair): weights streamed once.
410    pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
411        match self {
412            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
413            Self::Mapped {
414                dtype,
415                rows,
416                cols,
417                row_scale,
418                col_field,
419                vbit_offsets,
420                ..
421            } => {
422                if *dtype == TensorDtype::Q4Block {
423                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
424                    return;
425                }
426                if *dtype == TensorDtype::Vbit {
427                    vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
428                    return;
429                }
430                let x1s = prescale(x1, col_field, *dtype);
431                let x2s = prescale(x2, col_field, *dtype);
432                qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
433            }
434        }
435    }
436}
437
438impl QTensor {
439    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
440    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
441    /// to b matvec calls (same dot kernels in the same order); the win —
442    /// the weight row streams from DRAM once per batch, not b times.
443    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
444        let cols = self.cols();
445        let rows = self.rows();
446        debug_assert_eq!(xs_all.len(), b * cols);
447        debug_assert_eq!(out.len(), b * rows);
448        match self {
449            Self::F32 { data, .. } => {
450                let out_addr = SendMut(out.as_mut_ptr());
451                let run = |start: usize, end: usize| {
452                    for o in start..end {
453                        let row = &data[o * cols..(o + 1) * cols];
454                        for bi in 0..b {
455                            let x = &xs_all[bi * cols..(bi + 1) * cols];
456                            let mut acc = 0f32;
457                            for j in 0..cols {
458                                acc += row[j] * x[j];
459                            }
460                            unsafe { *out_addr.at(bi * rows + o) = acc };
461                        }
462                    }
463                };
464                dispatch_rows(pool, rows, &run);
465            }
466            Self::Mapped {
467                dtype,
468                row_scale,
469                col_field,
470                vbit_offsets,
471                ..
472            } => {
473                if *dtype == TensorDtype::Q4Block {
474                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
475                    return;
476                }
477                if *dtype == TensorDtype::Vbit {
478                    vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
479                    return;
480                }
481                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
482                    .map(|bi| {
483                        prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
484                    })
485                    .collect();
486                // D5: large prefill-batch GEMMs — on the GPU (threshold by
487                // work volume: submission carries b×rows×cols MACs).
488                if b >= 8
489                    && b * rows * cols >= 128_000_000
490                    && crate::gpu::enabled_here()
491                {
492                    if let Self::Mapped { model, idx, .. } = self {
493                        let flat: Vec<f32> =
494                            pre.iter().flat_map(|v| v.iter().copied()).collect();
495                        if crate::gpu::q8_matmat(
496                            model, *idx, row_scale, &flat, b, rows, cols, out)
497                        {
498                            return;
499                        }
500                    }
501                }
502                let q = self.quant_bytes();
503                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
504            }
505        }
506    }
507}
508
509impl QTensor {
510    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
511    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
512    /// barrier instead of N. Per-row math is the exact same kernel as
513    /// `matvec` (bit-identical outputs); only the dispatch is fused.
514    /// Falls back to N sequential matvecs when the set is not a uniform
515    /// q8-family/F32 group or there is no pool.
516    pub fn matvec_many<const N: usize>(
517        ts: [&QTensor; N],
518        x: &[f32],
519        mut outs: [&mut [f32]; N],
520        pool: Option<&Pool>,
521    ) {
522        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
523        let uniform_q8 = ts.iter().all(|t| {
524            matches!(
525                t,
526                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
527            )
528        });
529        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
530        let uniform_q4 = ts
531            .iter()
532            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
533        let uniform_vbit = ts
534            .iter()
535            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Vbit, .. }));
536        let Some(pool) = pool else {
537            for (t, o) in ts.iter().zip(outs.iter_mut()) {
538                t.matvec(x, o, None);
539            }
540            return;
541        };
542        if total_rows < 256 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit) {
543            for (t, o) in ts.iter().zip(outs.iter_mut()) {
544                t.matvec(x, o, Some(pool));
545            }
546            return;
547        }
548
549        if uniform_q4 || uniform_vbit {
550            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
551            // q4/vbit share one activation split — no per-tensor col field.
552            if a8w8_enabled() {
553                let act = split_act(x);
554                let act = &act;
555                if uniform_q4 {
556                    let closures: [_; N] = std::array::from_fn(|i| {
557                        let (packed, scales) =
558                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
559                        let (gpr, cols, out) =
560                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
561                        move |s: usize, e: usize| {
562                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
563                        }
564                    });
565                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
566                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
567                    pool.run_many(&parts);
568                } else {
569                    let closures: [_; N] = std::array::from_fn(|i| {
570                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
571                        let (bytes, rows, cols, out) =
572                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
573                        move |s: usize, e: usize| {
574                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
575                        }
576                    });
577                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
578                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
579                    pool.run_many(&parts);
580                }
581                return;
582            }
583            if uniform_q4 {
584                let closures: [_; N] = std::array::from_fn(|i| {
585                    let (packed, scales) =
586                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
587                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
588                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
589                });
590                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
591                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
592                pool.run_many(&parts);
593            } else {
594                let closures: [_; N] = std::array::from_fn(|i| {
595                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
596                    let (bytes, rows, cols, out) =
597                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
598                    move |s: usize, e: usize| {
599                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
600                    }
601                });
602                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
603                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
604                pool.run_many(&parts);
605            }
606            return;
607        }
608
609        if uniform_f32 {
610            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
611            let closures: [_; N] = std::array::from_fn(|i| {
612                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
613                let out = outs_addr[i];
614                move |start: usize, end: usize| {
615                    for o in start..end {
616                        let row = &data[o * cols..(o + 1) * cols];
617                        let mut sum = 0.0f32;
618                        for j in 0..*cols {
619                            sum += row[j] * x[j];
620                        }
621                        // SAFETY: disjoint (tensor, row) cells per worker.
622                        unsafe { *out.at(o) = sum };
623                    }
624                }
625            });
626            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
627                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
628            pool.run_many(&parts);
629            return;
630        }
631
632        // Uniform q8-family: per-tensor prescale (q8_2f col fields
633        // differ per tensor) + the shared range kernels.
634        struct Ctx<'a> {
635            bytes: &'a [u8],
636            row_scale: &'a [f32],
637            cols: usize,
638            xs: std::borrow::Cow<'a, [f32]>,
639        }
640        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
641            let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
642                unreachable!()
643            };
644            Ctx {
645                bytes: ts[i].quant_bytes(),
646                row_scale,
647                cols: *cols,
648                xs: prescale(x, col_field, *dtype),
649            }
650        });
651        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
652        #[cfg(target_arch = "aarch64")]
653        if sdot_enabled() {
654            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
655            let closures: [_; N] = std::array::from_fn(|i| {
656                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
657                move |start: usize, end: usize| {
658                    q8_range_sdot(c.bytes, c.row_scale, act, c.cols, out, start, end)
659                }
660            });
661            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
662                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
663            pool.run_many(&parts);
664            return;
665        }
666        #[cfg(target_arch = "x86_64")]
667        if avx2_a8w8_enabled() {
668            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
669            let closures: [_; N] = std::array::from_fn(|i| {
670                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
671                move |start: usize, end: usize| {
672                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
673                }
674            });
675            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
676                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
677            pool.run_many(&parts);
678            return;
679        }
680        let closures: [_; N] = std::array::from_fn(|i| {
681            let (c, out) = (&ctxs[i], outs_addr[i]);
682            move |start: usize, end: usize| {
683                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
684            }
685        });
686        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
687            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
688        pool.run_many(&parts);
689    }
690}
691
692impl QTensor {
693    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
694    /// single pool dispatch — the MTP/pair decode path publishes one job
695    /// for Q/K/V (and one for gate+up) instead of one per tensor.
696    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
697    #[allow(clippy::needless_range_loop)]
698    pub fn matvec2_many<const N: usize>(
699        ts: [&QTensor; N],
700        x1: &[f32],
701        x2: &[f32],
702        mut o1s: [&mut [f32]; N],
703        mut o2s: [&mut [f32]; N],
704        pool: Option<&Pool>,
705    ) {
706        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
707        let uniform_q8 = ts.iter().all(|t| {
708            matches!(
709                t,
710                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
711            )
712        });
713        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
714        let uniform_q4 = ts
715            .iter()
716            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
717        let uniform_vbit = ts
718            .iter()
719            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Vbit, .. }));
720        let fusable = pool.is_some()
721            && total_rows >= 256
722            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
723        if !fusable {
724            for i in 0..N {
725                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
726            }
727            return;
728        }
729        let pool = pool.unwrap();
730
731        if uniform_q4 || uniform_vbit {
732            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
733            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
734            // q4/vbit share activation splits — no per-tensor col field.
735            if a8w8_enabled() {
736                let a1 = split_act(x1);
737                let a2 = split_act(x2);
738                let (a1, a2) = (&a1, &a2);
739                if uniform_q4 {
740                    let closures: [_; N] = std::array::from_fn(|i| {
741                        let (packed, scales) =
742                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
743                        let (gpr, cols, o1, o2) =
744                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
745                        move |s: usize, e: usize| {
746                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
747                        }
748                    });
749                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
750                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
751                    pool.run_many(&parts);
752                } else {
753                    let closures: [_; N] = std::array::from_fn(|i| {
754                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
755                        let (bytes, rows, cols, o1, o2) =
756                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
757                        move |s: usize, e: usize| {
758                            vbit_range2_a8w8(
759                                bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
760                            )
761                        }
762                    });
763                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
764                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
765                    pool.run_many(&parts);
766                }
767                return;
768            }
769            if uniform_q4 {
770                let closures: [_; N] = std::array::from_fn(|i| {
771                    let (packed, scales) =
772                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
773                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
774                    move |s: usize, e: usize| {
775                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
776                    }
777                });
778                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
779                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
780                pool.run_many(&parts);
781            } else {
782                let closures: [_; N] = std::array::from_fn(|i| {
783                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
784                    let (bytes, rows, cols, o1, o2) =
785                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
786                    move |s: usize, e: usize| {
787                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
788                    }
789                });
790                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
791                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
792                pool.run_many(&parts);
793            }
794            return;
795        }
796
797        if uniform_f32 {
798            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
799            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
800            let closures: [_; N] = std::array::from_fn(|i| {
801                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
802                let (o1, o2) = (p1[i], p2[i]);
803                move |start: usize, end: usize| {
804                    for o in start..end {
805                        let row = &data[o * cols..(o + 1) * cols];
806                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
807                        for j in 0..*cols {
808                            s1 += row[j] * x1[j];
809                            s2 += row[j] * x2[j];
810                        }
811                        // SAFETY: disjoint (tensor, row) cells per worker.
812                        unsafe {
813                            *o1.at(o) = s1;
814                            *o2.at(o) = s2;
815                        }
816                    }
817                }
818            });
819            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
820                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
821            pool.run_many(&parts);
822            return;
823        }
824
825        struct Ctx<'a> {
826            bytes: &'a [u8],
827            row_scale: &'a [f32],
828            cols: usize,
829            xs1: std::borrow::Cow<'a, [f32]>,
830            xs2: std::borrow::Cow<'a, [f32]>,
831        }
832        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
833            let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
834                unreachable!()
835            };
836            Ctx {
837                bytes: ts[i].quant_bytes(),
838                row_scale,
839                cols: *cols,
840                xs1: prescale(x1, col_field, *dtype),
841                xs2: prescale(x2, col_field, *dtype),
842            }
843        });
844        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
845        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
846        #[cfg(target_arch = "aarch64")]
847        if sdot_enabled() {
848            let acts: [(SplitAct, SplitAct); N] =
849                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
850            let closures: [_; N] = std::array::from_fn(|i| {
851                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
852                move |start: usize, end: usize| {
853                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
854                }
855            });
856            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
857                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
858            pool.run_many(&parts);
859            return;
860        }
861        #[cfg(target_arch = "x86_64")]
862        if avx2_a8w8_enabled() {
863            let acts: [(SplitAct, SplitAct); N] =
864                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
865            let closures: [_; N] = std::array::from_fn(|i| {
866                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
867                move |start: usize, end: usize| {
868                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
869                }
870            });
871            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
872                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
873            pool.run_many(&parts);
874            return;
875        }
876        let closures: [_; N] = std::array::from_fn(|i| {
877            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
878            move |start: usize, end: usize| {
879                q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
880            }
881        });
882        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
883            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
884        pool.run_many(&parts);
885    }
886}
887
888/// Batched q8 kernel: same math as qmatvec, the row makes a single
889/// pass from memory for the whole batch.
890fn qmatmat(
891    q: &[u8],
892    row_scale: &[f32],
893    pre: &[std::borrow::Cow<'_, [f32]>],
894    rows: usize,
895    cols: usize,
896    out: &mut [f32],
897    pool: Option<&Pool>,
898) {
899    let b = pre.len();
900    debug_assert_eq!(out.len(), b * rows);
901    #[cfg(target_arch = "aarch64")]
902    if sdot_enabled() {
903        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
904        let out_addr = SendMut(out.as_mut_ptr());
905        let run = |start: usize, end: usize| {
906            for o in start..end {
907                let row = &q[o * cols..(o + 1) * cols];
908                for (bi, act) in acts.iter().enumerate() {
909                    let v = row_dot_sdot(row, act) * row_scale[o];
910                    unsafe { *out_addr.at(bi * rows + o) = v };
911                }
912            }
913        };
914        dispatch_rows(pool, rows, &run);
915        return;
916    }
917    // x86 A8W8 batch: each weight row streams once for the whole chunk
918    // through the AVX2/VNNI row dot (prefill q8 was the last
919    // aarch64-only batched path).
920    #[cfg(target_arch = "x86_64")]
921    if avx2_a8w8_enabled() {
922        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
923        let out_addr = SendMut(out.as_mut_ptr());
924        let run = |start: usize, end: usize| {
925            for o in start..end {
926                let row = &q[o * cols..(o + 1) * cols];
927                for (bi, act) in acts.iter().enumerate() {
928                    let v = row_dot_avx2(row, act) * row_scale[o];
929                    unsafe { *out_addr.at(bi * rows + o) = v };
930                }
931            }
932        };
933        dispatch_rows(pool, rows, &run);
934        return;
935    }
936    let out_addr = SendMut(out.as_mut_ptr());
937    let run = |start: usize, end: usize| {
938        for o in start..end {
939            let row = &q[o * cols..(o + 1) * cols];
940            for (bi, x) in pre.iter().enumerate() {
941                let mut acc = 0f32;
942                for j in 0..cols {
943                    acc += (row[j] as i8) as f32 * x[j];
944                }
945                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
946            }
947        }
948    };
949    dispatch_rows(pool, rows, &run);
950}
951
952/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
953/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
954fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
955    match pool {
956        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
957        _ => run(0, rows),
958    }
959}
960
961/// Split a q4_block blob into (packed nibbles, f16 group scales).
962fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
963    let groups = rows * cols / GROUP_SIZE;
964    bytes.split_at(groups * 16)
965}
966
967/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
968/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
969/// vbit packs MSB-first, so the HIGH nibble is the even element
970/// (opposite of q4_block's lo-first interleave). Centering is u-7.
971#[inline]
972fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
973    #[cfg(target_arch = "aarch64")]
974    unsafe {
975        return vbit_fill4_neon(data, buf);
976    }
977    #[cfg(target_arch = "x86_64")]
978    if avx2_enabled() {
979        return unsafe { vbit_fill4_avx2(data, buf) };
980    }
981    #[allow(unreachable_code)]
982    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
983        let u = unpack8::<4>(&data[blk * 4..]);
984        for k in 0..8 {
985            chunk[k] = (u[k] - 7) as i8 as u8;
986        }
987    }
988}
989
990#[cfg(target_arch = "aarch64")]
991#[target_feature(enable = "neon")]
992unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
993    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
994    // buf.len()/2 packed bytes (validated at load).
995    unsafe {
996        use core::arch::aarch64::*;
997        let n = buf.len();
998        let mask = vdupq_n_u8(0x0F);
999        let seven = vdupq_n_s8(7);
1000        let mut g = 0usize;
1001        while g * 32 + 32 <= n {
1002            let b = vld1q_u8(data.as_ptr().add(g * 16));
1003            let hi = vshrq_n_u8::<4>(b);
1004            let lo = vandq_u8(b, mask);
1005            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1006            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1007            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1008            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1009            g += 1;
1010        }
1011    }
1012}
1013
1014#[cfg(target_arch = "x86_64")]
1015#[target_feature(enable = "avx2")]
1016unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1017    // SAFETY: see vbit_fill4_neon.
1018    unsafe {
1019        use core::arch::x86_64::*;
1020        let n = buf.len();
1021        let mask = _mm_set1_epi8(0x0F);
1022        let seven = _mm256_set1_epi8(7);
1023        let mut g = 0usize;
1024        while g * 32 + 32 <= n {
1025            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1026            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1027            let lo = _mm_and_si128(b, mask);
1028            let z = _mm256_sub_epi8(
1029                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1030                seven,
1031            );
1032            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1033            g += 1;
1034        }
1035    }
1036}
1037
1038/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
1039/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
1040/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
1041/// into 4 such blocks.
1042#[inline(always)]
1043fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1044    let mut acc = 0u64;
1045    for i in 0..B {
1046        acc = (acc << 8) | data[i] as u64;
1047    }
1048    let mask = (1u64 << B) - 1;
1049    let mut out = [0i32; 8];
1050    for (k, o) in out.iter_mut().enumerate() {
1051        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1052    }
1053    out
1054}
1055
1056/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
1057/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
1058/// MSB-first, byte-padded]. Row data offsets are precomputed at load
1059/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
1060/// overhead on every matvec.
1061#[allow(clippy::too_many_arguments)]
1062fn vbitmatvec(
1063    bytes: &[u8],
1064    offsets: &[usize],
1065    x: &[f32],
1066    rows: usize,
1067    cols: usize,
1068    out: &mut [f32],
1069    pool: Option<&Pool>,
1070) {
1071    debug_assert_eq!(out.len(), rows);
1072    debug_assert_eq!(offsets.len(), rows + 1);
1073
1074    // SDOT path: unpack the row to centered i8 once, then per-group
1075    // int8 dot against the quantized activations — same A8W8 contract
1076    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
1077    if a8w8_enabled() {
1078        let act = split_act(x);
1079        let out_addr = SendMut(out.as_mut_ptr());
1080        let run = move |start: usize, end: usize| {
1081            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1082        };
1083        dispatch_rows(pool, rows, &run);
1084        return;
1085    }
1086
1087    let out_addr = SendMut(out.as_mut_ptr());
1088    let run = move |start: usize, end: usize| {
1089        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1090    };
1091    dispatch_rows(pool, rows, &run);
1092}
1093
1094/// One vbit row range via the A8W8 int8 path — kernel body of
1095/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
1096/// several tensors in one dispatch (b=8 rows go exact f32).
1097#[allow(clippy::too_many_arguments)]
1098fn vbit_range_a8w8(
1099    bytes: &[u8],
1100    offsets: &[usize],
1101    x: &[f32],
1102    act: &SplitAct,
1103    rows: usize,
1104    cols: usize,
1105    out: SendMut,
1106    start: usize,
1107    end: usize,
1108) {
1109    let ng = cols / GROUP_SIZE;
1110    let bits = &bytes[..rows];
1111    let sc_off = rows;
1112    let row_dot = |r: usize| -> f32 {
1113            let b = bits[r] as usize;
1114            let l = ((1i32 << (b - 1)) - 1) as i32;
1115            let mask = (1u64 << b) - 1;
1116            let data = &bytes[offsets[r]..offsets[r + 1]];
1117            if b == 8 {
1118                // u−L reaches 128 → does not fit i8; exact f32 path.
1119                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1120                let mut dot = 0f32;
1121                for g in 0..ng {
1122                    let so = (r * ng + g) * 2;
1123                    let sgf = f16_to_f32(u16::from_le_bytes([
1124                        bytes[sc_off + so],
1125                        bytes[sc_off + so + 1],
1126                    ]));
1127                    let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1128                    let mut gd = 0f32;
1129                    for &xv in xg.iter() {
1130                        if nbits < 8 {
1131                            acc = (acc << 8) | data[idx] as u64;
1132                            idx += 1;
1133                            nbits += 8;
1134                        }
1135                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1136                        nbits -= 8;
1137                        gd += (u - l) as f32 * xv;
1138                    }
1139                    dot += gd * sgf;
1140                }
1141                return dot;
1142            }
1143            // Per-worker scratch: this closure runs for every row of the
1144            // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
1145            // row was measurable pure overhead.
1146            thread_local! {
1147                static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1148                    const { std::cell::RefCell::new(Vec::new()) };
1149            }
1150            #[inline(always)]
1151            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1152                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1153                    let u = unpack8::<B>(&data[blk * B..]);
1154                    for k in 0..8 {
1155                        chunk[k] = (u[k] - l) as i8 as u8;
1156                    }
1157                }
1158            }
1159            let _ = mask;
1160            VBIT_SCRATCH.with(|scratch| {
1161                let mut buf = scratch.borrow_mut();
1162                buf.resize(cols, 0);
1163                match b {
1164                    3 => fill::<3>(data, l, &mut buf),
1165                    4 => vbit_fill4(data, &mut buf),
1166                    5 => fill::<5>(data, l, &mut buf),
1167                    6 => fill::<6>(data, l, &mut buf),
1168                    _ => unreachable!(),
1169                }
1170                let mut dot = 0f32;
1171                for g in 0..ng {
1172                    let so = (r * ng + g) * 2;
1173                    let s = f16_to_f32(u16::from_le_bytes([
1174                        bytes[sc_off + so],
1175                        bytes[sc_off + so + 1],
1176                    ]));
1177                    let d = dot_i8_i8(
1178                        &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1179                        &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1180                    ) as f32
1181                        * act.sx;
1182                    dot += d * s;
1183                }
1184                for &(j, xv) in &act.outliers {
1185                    let so = (r * ng + j / GROUP_SIZE) * 2;
1186                    let s = f16_to_f32(u16::from_le_bytes([
1187                        bytes[sc_off + so],
1188                        bytes[sc_off + so + 1],
1189                    ]));
1190                    // xq is zeroed at outlier slots — add the exact term.
1191                    dot += (buf[j] as i8) as f32 * s * xv;
1192                }
1193                dot
1194            })
1195    };
1196    for r in start..end {
1197        // SAFETY: disjoint row ranges per worker.
1198        unsafe { *out.at(r) = row_dot(r) };
1199    }
1200}
1201
1202/// Exact scalar vbit row range (same extraction, non-SDOT path).
1203#[allow(clippy::too_many_arguments)]
1204fn vbit_range_f32(
1205    bytes: &[u8],
1206    offsets: &[usize],
1207    x: &[f32],
1208    rows: usize,
1209    cols: usize,
1210    out: SendMut,
1211    start: usize,
1212    end: usize,
1213) {
1214    let ng = cols / GROUP_SIZE;
1215    let bits = &bytes[..rows];
1216    let sc_off = rows;
1217    // Per-bit-width specialized inner loops: the compiler unrolls the
1218    // constant shifts (the generic bit-buffer loop was branch-bound —
1219    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
1220    #[inline(always)]
1221    fn dot_row<const B: usize>(
1222        data: &[u8],
1223        bytes: &[u8],
1224        sc_off: usize,
1225        r: usize,
1226        ng: usize,
1227        x: &[f32],
1228    ) -> f32 {
1229        let l = ((1i32 << (B - 1)) - 1) as f32;
1230        let gbytes = GROUP_SIZE * B / 8;
1231        let mut dot = 0f32;
1232        for g in 0..ng {
1233            let so = (r * ng + g) * 2;
1234            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1235            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1236            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1237            let mut gd = 0f32;
1238            for blk in 0..GROUP_SIZE / 8 {
1239                let u = unpack8::<B>(&gd0[blk * B..]);
1240                let xb = &xg[blk * 8..blk * 8 + 8];
1241                for k in 0..8 {
1242                    gd += (u[k] as f32 - l) * xb[k];
1243                }
1244            }
1245            dot += gd * s;
1246        }
1247        dot
1248    }
1249    for r in start..end {
1250        let data = &bytes[offsets[r]..offsets[r + 1]];
1251        let v = match bits[r] {
1252            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1253            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1254            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1255            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1256            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1257            b => unreachable!("vbit bit-width {b} (validated at load)"),
1258        };
1259        // SAFETY: disjoint row ranges per worker.
1260        unsafe { *out.at(r) = v };
1261    }
1262}
1263
1264/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
1265/// and dotted against BOTH activations (MTP verify / pair prefill used
1266/// to run two full matvecs — double weight traffic and double unpack).
1267/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
1268#[allow(clippy::too_many_arguments)]
1269fn vbitmatvec2(
1270    bytes: &[u8],
1271    offsets: &[usize],
1272    x1: &[f32],
1273    x2: &[f32],
1274    rows: usize,
1275    cols: usize,
1276    o1: &mut [f32],
1277    o2: &mut [f32],
1278    pool: Option<&Pool>,
1279) {
1280    debug_assert_eq!(o1.len(), rows);
1281    debug_assert_eq!(o2.len(), rows);
1282
1283    if a8w8_enabled() {
1284        let a1 = split_act(x1);
1285        let a2 = split_act(x2);
1286        let p1 = SendMut(o1.as_mut_ptr());
1287        let p2 = SendMut(o2.as_mut_ptr());
1288        let run = move |start: usize, end: usize| {
1289            vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1290        };
1291        dispatch_rows(pool, rows, &run);
1292        return;
1293    }
1294
1295    let p1 = SendMut(o1.as_mut_ptr());
1296    let p2 = SendMut(o2.as_mut_ptr());
1297    let run = move |start: usize, end: usize| {
1298        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1299    };
1300    dispatch_rows(pool, rows, &run);
1301}
1302
1303/// Two-input vbit row range via the A8W8 int8 path — kernel body of
1304/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
1305/// exact f32 for both lanes, bits streamed once).
1306#[allow(clippy::too_many_arguments)]
1307fn vbit_range2_a8w8(
1308    bytes: &[u8],
1309    offsets: &[usize],
1310    x1: &[f32],
1311    x2: &[f32],
1312    a1: &SplitAct,
1313    a2: &SplitAct,
1314    rows: usize,
1315    cols: usize,
1316    p1: SendMut,
1317    p2: SendMut,
1318    start: usize,
1319    end: usize,
1320) {
1321    let ng = cols / GROUP_SIZE;
1322    let bits = &bytes[..rows];
1323    let sc_off = rows;
1324    let row_dots = |r: usize| -> (f32, f32) {
1325            let b = bits[r] as usize;
1326            let l = (1i32 << (b - 1)) - 1;
1327            let data = &bytes[offsets[r]..offsets[r + 1]];
1328            if b == 8 {
1329                // u−L reaches 128 → does not fit i8; exact f32 path,
1330                // bits still streamed once for both lanes.
1331                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1332                let (mut d1, mut d2) = (0f32, 0f32);
1333                for g in 0..ng {
1334                    let so = (r * ng + g) * 2;
1335                    let sgf = f16_to_f32(u16::from_le_bytes([
1336                        bytes[sc_off + so],
1337                        bytes[sc_off + so + 1],
1338                    ]));
1339                    let (mut g1, mut g2) = (0f32, 0f32);
1340                    for k in 0..GROUP_SIZE {
1341                        if nbits < 8 {
1342                            acc = (acc << 8) | data[idx] as u64;
1343                            idx += 1;
1344                            nbits += 8;
1345                        }
1346                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1347                        nbits -= 8;
1348                        let w = (u - l) as f32;
1349                        g1 += w * x1[g * GROUP_SIZE + k];
1350                        g2 += w * x2[g * GROUP_SIZE + k];
1351                    }
1352                    d1 += g1 * sgf;
1353                    d2 += g2 * sgf;
1354                }
1355                return (d1, d2);
1356            }
1357            thread_local! {
1358                static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1359                    const { std::cell::RefCell::new(Vec::new()) };
1360            }
1361            #[inline(always)]
1362            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1363                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1364                    let u = unpack8::<B>(&data[blk * B..]);
1365                    for k in 0..8 {
1366                        chunk[k] = (u[k] - l) as i8 as u8;
1367                    }
1368                }
1369            }
1370            VBIT_SCRATCH2.with(|scratch| {
1371                let mut buf = scratch.borrow_mut();
1372                buf.resize(cols, 0);
1373                match b {
1374                    3 => fill::<3>(data, l, &mut buf),
1375                    4 => vbit_fill4(data, &mut buf),
1376                    5 => fill::<5>(data, l, &mut buf),
1377                    6 => fill::<6>(data, l, &mut buf),
1378                    _ => unreachable!(),
1379                }
1380                let (mut d1, mut d2) = (0f32, 0f32);
1381                for g in 0..ng {
1382                    let so = (r * ng + g) * 2;
1383                    let s = f16_to_f32(u16::from_le_bytes([
1384                        bytes[sc_off + so],
1385                        bytes[sc_off + so + 1],
1386                    ]));
1387                    let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1388                    let v1 =
1389                        dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1390                    let v2 =
1391                        dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1392                    d1 += v1 * s;
1393                    d2 += v2 * s;
1394                }
1395                for &(j, xv) in &a1.outliers {
1396                    let so = (r * ng + j / GROUP_SIZE) * 2;
1397                    let s = f16_to_f32(u16::from_le_bytes([
1398                        bytes[sc_off + so],
1399                        bytes[sc_off + so + 1],
1400                    ]));
1401                    d1 += (buf[j] as i8) as f32 * s * xv;
1402                }
1403                for &(j, xv) in &a2.outliers {
1404                    let so = (r * ng + j / GROUP_SIZE) * 2;
1405                    let s = f16_to_f32(u16::from_le_bytes([
1406                        bytes[sc_off + so],
1407                        bytes[sc_off + so + 1],
1408                    ]));
1409                    d2 += (buf[j] as i8) as f32 * s * xv;
1410                }
1411                (d1, d2)
1412            })
1413    };
1414    for r in start..end {
1415        let (v1, v2) = row_dots(r);
1416        // SAFETY: disjoint row ranges per worker.
1417        unsafe {
1418            *p1.at(r) = v1;
1419            *p2.at(r) = v2;
1420        }
1421    }
1422}
1423
1424/// Two-input exact scalar vbit row range (same extraction) —
1425/// per-bit-width specialized, two accumulators per row; per-lane
1426/// accumulation order matches `vbitmatvec` exactly.
1427#[allow(clippy::too_many_arguments)]
1428fn vbit_range2_f32(
1429    bytes: &[u8],
1430    offsets: &[usize],
1431    x1: &[f32],
1432    x2: &[f32],
1433    rows: usize,
1434    cols: usize,
1435    p1: SendMut,
1436    p2: SendMut,
1437    start: usize,
1438    end: usize,
1439) {
1440    let ng = cols / GROUP_SIZE;
1441    let bits = &bytes[..rows];
1442    let sc_off = rows;
1443    #[inline(always)]
1444    #[allow(clippy::too_many_arguments)]
1445    fn dot_row2<const B: usize>(
1446        data: &[u8],
1447        bytes: &[u8],
1448        sc_off: usize,
1449        r: usize,
1450        ng: usize,
1451        x1: &[f32],
1452        x2: &[f32],
1453    ) -> (f32, f32) {
1454        let l = ((1i32 << (B - 1)) - 1) as f32;
1455        let gbytes = GROUP_SIZE * B / 8;
1456        let (mut d1, mut d2) = (0f32, 0f32);
1457        for g in 0..ng {
1458            let so = (r * ng + g) * 2;
1459            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1460            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1461            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1462            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1463            let (mut g1, mut g2) = (0f32, 0f32);
1464            for blk in 0..GROUP_SIZE / 8 {
1465                let u = unpack8::<B>(&gd0[blk * B..]);
1466                for k in 0..8 {
1467                    let w = u[k] as f32 - l;
1468                    g1 += w * x1g[blk * 8 + k];
1469                    g2 += w * x2g[blk * 8 + k];
1470                }
1471            }
1472            d1 += g1 * s;
1473            d2 += g2 * s;
1474        }
1475        (d1, d2)
1476    }
1477    for r in start..end {
1478        let data = &bytes[offsets[r]..offsets[r + 1]];
1479        let (v1, v2) = match bits[r] {
1480            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
1481            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
1482            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
1483            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
1484            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
1485            b => unreachable!("vbit bit-width {b} (validated at load)"),
1486        };
1487        // SAFETY: disjoint row ranges per worker.
1488        unsafe {
1489            *p1.at(r) = v1;
1490            *p2.at(r) = v2;
1491        }
1492    }
1493}
1494
1495/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
1496/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
1497/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
1498/// 32-group, exact outlier correction — the same A8W8 contract as q8.
1499/// `CMF_SDOT=0` keeps the exact scalar path.
1500fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
1501    debug_assert_eq!(out.len(), rows);
1502    let (packed, scales) = q4_split(bytes, rows, cols);
1503    let gpr = cols / GROUP_SIZE;
1504    let out_addr = SendMut(out.as_mut_ptr());
1505
1506    if a8w8_enabled() {
1507        let act = split_act(x);
1508        let run = move |start: usize, end: usize| {
1509            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
1510        };
1511        dispatch_rows(pool, rows, &run);
1512        return;
1513    }
1514
1515    let run = move |start: usize, end: usize| {
1516        q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
1517    };
1518    dispatch_rows(pool, rows, &run);
1519}
1520
1521/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
1522/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
1523#[inline]
1524#[allow(unreachable_code)]
1525fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
1526    #[cfg(target_arch = "aarch64")]
1527    unsafe {
1528        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
1529    }
1530    #[cfg(target_arch = "x86_64")]
1531    unsafe {
1532        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
1533    }
1534    let mut acc = 0f32;
1535    for gi in 0..gpr {
1536        let g = g0 + gi;
1537        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1538        let mut d = 0i32;
1539        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1540            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
1541                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
1542        }
1543        acc += d as f32 * s;
1544    }
1545    acc
1546}
1547
1548/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
1549#[inline]
1550#[allow(unreachable_code)]
1551fn dot_q4_row_i8_2(
1552    packed: &[u8],
1553    scales: &[u8],
1554    g0: usize,
1555    gpr: usize,
1556    xq1: &[i8],
1557    xq2: &[i8],
1558) -> (f32, f32) {
1559    #[cfg(target_arch = "aarch64")]
1560    unsafe {
1561        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
1562    }
1563    #[cfg(target_arch = "x86_64")]
1564    unsafe {
1565        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
1566    }
1567    (
1568        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
1569        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
1570    )
1571}
1572
1573/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
1574/// multi-matrix jobs can drive it for several tensors in one dispatch).
1575#[allow(clippy::too_many_arguments)]
1576fn q4_range_a8w8(
1577    packed: &[u8],
1578    scales: &[u8],
1579    gpr: usize,
1580    cols: usize,
1581    act: &SplitAct,
1582    out: SendMut,
1583    start: usize,
1584    end: usize,
1585) {
1586    for r in start..end {
1587        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
1588        // xq is zeroed at outlier slots — add the exact terms.
1589        for &(j, xv) in &act.outliers {
1590            let flat = r * cols + j;
1591            let byte = packed[flat / 2];
1592            let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
1593            let s = f16_to_f32(u16::from_le_bytes([
1594                scales[(flat / GROUP_SIZE) * 2],
1595                scales[(flat / GROUP_SIZE) * 2 + 1],
1596            ]));
1597            acc += ((nib as i32 - 8) as f32) * s * xv;
1598        }
1599        // SAFETY: disjoint row ranges per worker.
1600        unsafe { *out.at(r) = acc };
1601    }
1602}
1603
1604/// Two-input q4 row range via the A8W8 int8 path — kernel body of
1605/// `q4matvec2`, extracted for pair multi-matrix jobs.
1606#[allow(clippy::too_many_arguments)]
1607fn q4_range2_a8w8(
1608    packed: &[u8],
1609    scales: &[u8],
1610    gpr: usize,
1611    cols: usize,
1612    a1: &SplitAct,
1613    a2: &SplitAct,
1614    p1: SendMut,
1615    p2: SendMut,
1616    start: usize,
1617    end: usize,
1618) {
1619    for r in start..end {
1620        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
1621        let mut acc1 = s1 * a1.sx;
1622        let mut acc2 = s2 * a2.sx;
1623        // xq is zeroed at outlier slots — add the exact terms.
1624        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
1625            for &(j, xv) in outliers {
1626                let flat = r * cols + j;
1627                let byte = packed[flat / 2];
1628                let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
1629                let s = f16_to_f32(u16::from_le_bytes([
1630                    scales[(flat / GROUP_SIZE) * 2],
1631                    scales[(flat / GROUP_SIZE) * 2 + 1],
1632                ]));
1633                *acc += ((nib as i32 - 8) as f32) * s * xv;
1634            }
1635        };
1636        fix(&a1.outliers, &mut acc1);
1637        fix(&a2.outliers, &mut acc2);
1638        // SAFETY: disjoint row ranges per worker.
1639        unsafe {
1640            *p1.at(r) = acc1;
1641            *p2.at(r) = acc2;
1642        }
1643    }
1644}
1645
1646/// Exact scalar q4 row range (same extraction, non-SDOT path).
1647fn q4_range_f32(
1648    packed: &[u8],
1649    scales: &[u8],
1650    gpr: usize,
1651    x: &[f32],
1652    out: SendMut,
1653    start: usize,
1654    end: usize,
1655) {
1656    for r in start..end {
1657        let mut acc = 0f32;
1658        for gi in 0..gpr {
1659            let g = r * gpr + gi;
1660            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1661            let pk = &packed[g * 16..(g + 1) * 16];
1662            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1663            let mut ga = 0f32;
1664            for (k, &b) in pk.iter().enumerate() {
1665                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
1666                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
1667            }
1668            acc += ga * s;
1669        }
1670        // SAFETY: disjoint row ranges per worker.
1671        unsafe { *out.at(r) = acc };
1672    }
1673}
1674
1675/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
1676/// dotted against both activations (was: two full matvecs — double
1677/// weight traffic). Per-lane math matches `q4matvec` exactly.
1678#[allow(clippy::too_many_arguments)]
1679fn q4matvec2(
1680    bytes: &[u8],
1681    x1: &[f32],
1682    x2: &[f32],
1683    rows: usize,
1684    cols: usize,
1685    o1: &mut [f32],
1686    o2: &mut [f32],
1687    pool: Option<&Pool>,
1688) {
1689    debug_assert_eq!(o1.len(), rows);
1690    debug_assert_eq!(o2.len(), rows);
1691    let (packed, scales) = q4_split(bytes, rows, cols);
1692    let gpr = cols / GROUP_SIZE;
1693
1694    if a8w8_enabled() {
1695        let a1 = split_act(x1);
1696        let a2 = split_act(x2);
1697        let p1 = SendMut(o1.as_mut_ptr());
1698        let p2 = SendMut(o2.as_mut_ptr());
1699        let run = move |start: usize, end: usize| {
1700            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
1701        };
1702        dispatch_rows(pool, rows, &run);
1703        return;
1704    }
1705
1706    let p1 = SendMut(o1.as_mut_ptr());
1707    let p2 = SendMut(o2.as_mut_ptr());
1708    let run = move |start: usize, end: usize| {
1709        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
1710    };
1711    dispatch_rows(pool, rows, &run);
1712}
1713
1714/// Two-input exact scalar q4 row range (same extraction).
1715#[allow(clippy::too_many_arguments)]
1716fn q4_range2_f32(
1717    packed: &[u8],
1718    scales: &[u8],
1719    gpr: usize,
1720    x1: &[f32],
1721    x2: &[f32],
1722    p1: SendMut,
1723    p2: SendMut,
1724    start: usize,
1725    end: usize,
1726) {
1727    for r in start..end {
1728        let (mut acc1, mut acc2) = (0f32, 0f32);
1729        for gi in 0..gpr {
1730            let g = r * gpr + gi;
1731            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1732            let pk = &packed[g * 16..(g + 1) * 16];
1733            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1734            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1735            let (mut g1, mut g2) = (0f32, 0f32);
1736            for (k, &b) in pk.iter().enumerate() {
1737                let wl = (b & 0x0F) as f32 - 8.0;
1738                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
1739                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
1740                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
1741            }
1742            acc1 += g1 * s;
1743            acc2 += g2 * s;
1744        }
1745        // SAFETY: disjoint row ranges per worker.
1746        unsafe {
1747            *p1.at(r) = acc1;
1748            *p2.at(r) = acc2;
1749        }
1750    }
1751}
1752
1753thread_local! {
1754    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
1755    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
1756    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
1757    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1758}
1759
1760/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
1761/// and dotted against ALL b activations (prefill used to fall back to b
1762/// full matvecs — b× weight traffic and b× nibble decode). Per-position
1763/// math matches `q4matvec` exactly: same group order, same accumulation.
1764/// `out` is row-major [b, rows] like `qmatmat`.
1765#[allow(clippy::too_many_arguments)]
1766fn q4matmat(
1767    bytes: &[u8],
1768    xs_all: &[f32],
1769    b: usize,
1770    rows: usize,
1771    cols: usize,
1772    out: &mut [f32],
1773    pool: Option<&Pool>,
1774) {
1775    debug_assert_eq!(xs_all.len(), b * cols);
1776    debug_assert_eq!(out.len(), b * rows);
1777    let (packed, scales) = q4_split(bytes, rows, cols);
1778    let gpr = cols / GROUP_SIZE;
1779    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1780
1781    if a8w8_enabled() {
1782        let acts: Vec<SplitAct> =
1783            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
1784        let acts = &acts;
1785        let out_addr = SendMut(out.as_mut_ptr());
1786        let run = move |start: usize, end: usize| {
1787            ROW_I8.with(|rb| {
1788                let mut buf = rb.borrow_mut();
1789                buf.resize(cols, 0);
1790                for r in start..end {
1791                    // Unpack the row's nibbles to centered i8 once
1792                    // (element 2k = low nibble, 2k+1 = high — flat order,
1793                    // same as dot_q4_row_sdot's zip).
1794                    for gi in 0..gpr {
1795                        let g = r * gpr + gi;
1796                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1797                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
1798                            buf[gi * GROUP_SIZE + k * 2 + 1] =
1799                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
1800                        }
1801                    }
1802                    for (bi, act) in acts.iter().enumerate() {
1803                        let mut acc = 0f32;
1804                        for gi in 0..gpr {
1805                            let d = dot_i8_i8(
1806                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
1807                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
1808                            );
1809                            acc += d as f32 * gscale(r * gpr + gi);
1810                        }
1811                        acc *= act.sx;
1812                        // xq is zeroed at outlier slots — exact terms.
1813                        for &(j, xv) in &act.outliers {
1814                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
1815                        }
1816                        // SAFETY: disjoint (bi, r) cells per worker row range.
1817                        unsafe { *out_addr.at(bi * rows + r) = acc };
1818                    }
1819                }
1820            })
1821        };
1822        dispatch_rows(pool, rows, &run);
1823        return;
1824    }
1825
1826    let out_addr = SendMut(out.as_mut_ptr());
1827    let run = move |start: usize, end: usize| {
1828        ROW_F32.with(|rb| {
1829            let mut buf = rb.borrow_mut();
1830            buf.resize(cols, 0.0);
1831            for r in start..end {
1832                // Decode raw (nib − 8) values once; scales stay per-group
1833                // so the accumulation order matches q4matvec bit-for-bit.
1834                for gi in 0..gpr {
1835                    let g = r * gpr + gi;
1836                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1837                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
1838                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
1839                    }
1840                }
1841                for bi in 0..b {
1842                    let x = &xs_all[bi * cols..(bi + 1) * cols];
1843                    let mut acc = 0f32;
1844                    for gi in 0..gpr {
1845                        let mut ga = 0f32;
1846                        // Pairwise (lo + hi) addition, matching
1847                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
1848                        // a flat one-per-element loop rounds differently
1849                        // and broke bit-parity on the scalar (x86) path.
1850                        for k in 0..GROUP_SIZE / 2 {
1851                            let e = gi * GROUP_SIZE + k * 2;
1852                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
1853                        }
1854                        acc += ga * gscale(r * gpr + gi);
1855                    }
1856                    // SAFETY: disjoint (bi, r) cells per worker row range.
1857                    unsafe { *out_addr.at(bi * rows + r) = acc };
1858                }
1859            }
1860        })
1861    };
1862    dispatch_rows(pool, rows, &run);
1863}
1864
1865/// Batched vbit matmat: each variable-bit row is decoded from the mmap
1866/// ONCE for the whole microbatch. Same per-position math as
1867/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
1868/// and the scalar path).
1869#[allow(clippy::too_many_arguments)]
1870fn vbitmatmat(
1871    bytes: &[u8],
1872    offsets: &[usize],
1873    xs_all: &[f32],
1874    b: usize,
1875    rows: usize,
1876    cols: usize,
1877    out: &mut [f32],
1878    pool: Option<&Pool>,
1879) {
1880    debug_assert_eq!(xs_all.len(), b * cols);
1881    debug_assert_eq!(out.len(), b * rows);
1882    debug_assert_eq!(offsets.len(), rows + 1);
1883    let ng = cols / GROUP_SIZE;
1884    let bits = &bytes[..rows];
1885    let sc_off = rows;
1886    let gscale = |r: usize, g: usize| {
1887        let so = (r * ng + g) * 2;
1888        f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
1889    };
1890
1891    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
1892    let decode_f32 = |r: usize, dst: &mut [f32]| {
1893        let bw = bits[r] as usize;
1894        let l = ((1i32 << (bw - 1)) - 1) as f32;
1895        let data = &bytes[offsets[r]..offsets[r + 1]];
1896        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1897        for d in dst.iter_mut() {
1898            while nbits < bw {
1899                acc = (acc << 8) | data[idx] as u64;
1900                idx += 1;
1901                nbits += 8;
1902            }
1903            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
1904            nbits -= bw;
1905            *d = u - l;
1906        }
1907    };
1908
1909    if a8w8_enabled() {
1910        let acts: Vec<SplitAct> =
1911            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
1912        let acts = &acts;
1913        let out_addr = SendMut(out.as_mut_ptr());
1914        let run = move |start: usize, end: usize| {
1915            for r in start..end {
1916                let bw = bits[r] as usize;
1917                if bw == 8 {
1918                    // u−L reaches 128 → no i8 path; decode once, exact
1919                    // f32 dots for every position (same as vbitmatvec).
1920                    ROW_F32.with(|rb| {
1921                        let mut buf = rb.borrow_mut();
1922                        buf.resize(cols, 0.0);
1923                        decode_f32(r, &mut buf);
1924                        for bi in 0..b {
1925                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1926                            let mut dot = 0f32;
1927                            for g in 0..ng {
1928                                let mut gd = 0f32;
1929                                for k in 0..GROUP_SIZE {
1930                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
1931                                }
1932                                dot += gd * gscale(r, g);
1933                            }
1934                            // SAFETY: disjoint (bi, r) cells per worker range.
1935                            unsafe { *out_addr.at(bi * rows + r) = dot };
1936                        }
1937                    });
1938                    continue;
1939                }
1940                let l = (1i32 << (bw - 1)) - 1;
1941                let data = &bytes[offsets[r]..offsets[r + 1]];
1942                ROW_I8.with(|rb| {
1943                    let mut buf = rb.borrow_mut();
1944                    buf.resize(cols, 0);
1945                    #[inline(always)]
1946                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1947                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1948                            let u = unpack8::<B>(&data[blk * B..]);
1949                            for k in 0..8 {
1950                                chunk[k] = (u[k] - l) as i8 as u8;
1951                            }
1952                        }
1953                    }
1954                    match bw {
1955                        3 => fill::<3>(data, l, &mut buf),
1956                        4 => vbit_fill4(data, &mut buf),
1957                        5 => fill::<5>(data, l, &mut buf),
1958                        6 => fill::<6>(data, l, &mut buf),
1959                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
1960                    }
1961                    for (bi, act) in acts.iter().enumerate() {
1962                        let mut dot = 0f32;
1963                        for g in 0..ng {
1964                            let d = dot_i8_i8(
1965                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1966                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1967                            ) as f32
1968                                * act.sx;
1969                            dot += d * gscale(r, g);
1970                        }
1971                        for &(j, xv) in &act.outliers {
1972                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
1973                        }
1974                        // SAFETY: disjoint (bi, r) cells per worker range.
1975                        unsafe { *out_addr.at(bi * rows + r) = dot };
1976                    }
1977                });
1978            }
1979        };
1980        dispatch_rows(pool, rows, &run);
1981        return;
1982    }
1983
1984    let out_addr = SendMut(out.as_mut_ptr());
1985    let run = move |start: usize, end: usize| {
1986        ROW_F32.with(|rb| {
1987            let mut buf = rb.borrow_mut();
1988            buf.resize(cols, 0.0);
1989            for r in start..end {
1990                decode_f32(r, &mut buf);
1991                for bi in 0..b {
1992                    let x = &xs_all[bi * cols..(bi + 1) * cols];
1993                    let mut dot = 0f32;
1994                    for g in 0..ng {
1995                        let mut gd = 0f32;
1996                        for k in 0..GROUP_SIZE {
1997                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
1998                        }
1999                        dot += gd * gscale(r, g);
2000                    }
2001                    // SAFETY: disjoint (bi, r) cells per worker range.
2002                    unsafe { *out_addr.at(bi * rows + r) = dot };
2003                }
2004            }
2005        })
2006    };
2007    dispatch_rows(pool, rows, &run);
2008}
2009
2010/// θ col-field fold for q8_2f activations. Borrowed pass-through for
2011/// every other dtype — the old unconditional `x.to_vec()` was a pure
2012/// per-matvec allocation on the q8_row hot path.
2013pub(crate) fn prescale<'a>(
2014    x: &'a [f32],
2015    col_field: &[f32],
2016    dtype: TensorDtype,
2017) -> std::borrow::Cow<'a, [f32]> {
2018    if dtype == TensorDtype::Q8_2f {
2019        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
2020    } else {
2021        std::borrow::Cow::Borrowed(x)
2022    }
2023}
2024
2025// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
2026
2027/// AVX2+FMA available? Default ON when the CPU supports both;
2028/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
2029#[cfg(target_arch = "x86_64")]
2030pub(crate) fn avx2_enabled() -> bool {
2031    use std::sync::OnceLock;
2032    static ON: OnceLock<bool> = OnceLock::new();
2033    *ON.get_or_init(|| {
2034        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
2035            && std::arch::is_x86_feature_detected!("avx2")
2036            && std::arch::is_x86_feature_detected!("fma")
2037    })
2038}
2039
2040/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
2041/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
2042/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
2043/// active either way, they are exact (regrouped sums only).
2044#[cfg(target_arch = "x86_64")]
2045fn avx2_a8w8_enabled() -> bool {
2046    use std::sync::OnceLock;
2047    static ON: OnceLock<bool> = OnceLock::new();
2048    *ON.get_or_init(|| {
2049        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
2050    })
2051}
2052
2053/// A8W8 quantized-activation path available on THIS machine? One
2054/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
2055/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
2056#[inline]
2057pub(crate) fn a8w8_enabled() -> bool {
2058    #[cfg(target_arch = "aarch64")]
2059    {
2060        sdot_enabled()
2061    }
2062    #[cfg(target_arch = "x86_64")]
2063    {
2064        avx2_a8w8_enabled()
2065    }
2066    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
2067    {
2068        false
2069    }
2070}
2071
2072/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
2073/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
2074#[inline]
2075#[allow(unreachable_code)]
2076fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
2077    #[cfg(target_arch = "aarch64")]
2078    unsafe {
2079        return dot_i8_sdot(w, xq);
2080    }
2081    #[cfg(target_arch = "x86_64")]
2082    unsafe {
2083        if avx512vnni_enabled() {
2084            return dot_i8_i8_vnni(w, xq);
2085        }
2086        return dot_i8_i8_avx2(w, xq);
2087    }
2088    w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
2089}
2090
2091/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
2092/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
2093/// `vpdpbusd` encoding.
2094#[cfg(target_arch = "x86_64")]
2095fn avx512vnni_enabled() -> bool {
2096    use std::sync::OnceLock;
2097    static ON: OnceLock<bool> = OnceLock::new();
2098    *ON.get_or_init(|| {
2099        std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
2100            && std::arch::is_x86_feature_detected!("avx512f")
2101            && std::arch::is_x86_feature_detected!("avx512bw")
2102            && std::arch::is_x86_feature_detected!("avx512vl")
2103            && std::arch::is_x86_feature_detected!("avx512vnni")
2104    })
2105}
2106
2107/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
2108/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
2109/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
2110/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
2111#[cfg(target_arch = "x86_64")]
2112#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
2113unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
2114    // SAFETY: callers uphold slice-length contracts (see call sites).
2115    unsafe {
2116        use core::arch::x86_64::*;
2117        let n = w.len();
2118        let mut j = 0usize;
2119        let mut total: i32;
2120        // 4 independent accumulators: vpdpbusd is its own loop-carried
2121        // dependency (~5-cycle latency) — a single-acc loop runs
2122        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
2123        // on Granite Rapids.
2124        {
2125            #[inline(always)]
2126            unsafe fn step(
2127                w: *const u8,
2128                x: *const i8,
2129                acc: core::arch::x86_64::__m512i,
2130            ) -> core::arch::x86_64::__m512i {
2131                unsafe {
2132                    use core::arch::x86_64::*;
2133                    let wv = _mm512_loadu_si512(w as *const _);
2134                    let xv = _mm512_loadu_si512(x as *const _);
2135                    let aw = _mm512_abs_epi8(wv);
2136                    let neg = _mm512_movepi8_mask(wv);
2137                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
2138                    _mm512_dpbusd_epi32(acc, aw, sx)
2139                }
2140            }
2141            let (mut a0, mut a1, mut a2, mut a3) = (
2142                _mm512_setzero_si512(),
2143                _mm512_setzero_si512(),
2144                _mm512_setzero_si512(),
2145                _mm512_setzero_si512(),
2146            );
2147            while j + 256 <= n {
2148                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
2149                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
2150                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
2151                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
2152                j += 256;
2153            }
2154            while j + 64 <= n {
2155                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
2156                j += 64;
2157            }
2158            let s01 = _mm512_add_epi32(a0, a1);
2159            let s23 = _mm512_add_epi32(a2, a3);
2160            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
2161        }
2162        // 32-wide (q4/vbit groups are exactly 32 bytes).
2163        if j + 32 <= n {
2164            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
2165            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
2166            let d = _mm256_dpbusd_epi32(
2167                _mm256_setzero_si256(),
2168                _mm256_abs_epi8(wv),
2169                _mm256_sign_epi8(xv, wv),
2170            );
2171            let hi128 = _mm256_extracti128_si256::<1>(d);
2172            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2173            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2174            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2175            total += _mm_cvtsi128_si32(s32);
2176            j += 32;
2177        }
2178        while j < n {
2179            total += (w[j] as i8) as i32 * xq[j] as i32;
2180            j += 1;
2181        }
2182        total
2183    }
2184}
2185
2186/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
2187#[cfg(target_arch = "x86_64")]
2188#[target_feature(enable = "avx2,fma")]
2189unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
2190    // SAFETY: callers uphold slice-length contracts (see call sites).
2191    unsafe {
2192        use core::arch::x86_64::*;
2193        let n = x.len();
2194        let wp = w.as_ptr();
2195        let xp = x.as_ptr();
2196        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
2197        let mut j = 0usize;
2198        while j + 16 <= n {
2199            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
2200            let lo = _mm256_cvtepi8_epi32(wb);
2201            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
2202            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
2203            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
2204            j += 16;
2205        }
2206        let acc = _mm256_add_ps(a0, a1);
2207        let hi128 = _mm256_extractf128_ps::<1>(acc);
2208        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
2209        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
2210        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
2211        let mut sum = _mm_cvtss_f32(s32);
2212        while j < n {
2213            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
2214            j += 1;
2215        }
2216        sum
2217    }
2218}
2219
2220/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
2221/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
2222/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
2223/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
2224#[cfg(target_arch = "x86_64")]
2225#[target_feature(enable = "avx2")]
2226unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
2227    // SAFETY: callers uphold slice-length contracts (see call sites).
2228    unsafe {
2229        use core::arch::x86_64::*;
2230        let n = w.len();
2231        let ones = _mm256_set1_epi16(1);
2232        let mut acc = _mm256_setzero_si256();
2233        let mut j = 0usize;
2234        while j + 32 <= n {
2235            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
2236            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
2237            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
2238            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
2239            j += 32;
2240        }
2241        let hi128 = _mm256_extracti128_si256::<1>(acc);
2242        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
2243        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2244        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2245        let mut s = _mm_cvtsi128_si32(s32);
2246        while j < n {
2247            s += (w[j] as i8) as i32 * xq[j] as i32;
2248            j += 1;
2249        }
2250        s
2251    }
2252}
2253
2254/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
2255/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
2256/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
2257/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
2258#[cfg(target_arch = "x86_64")]
2259#[inline]
2260fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
2261    let dot = if avx512vnni_enabled() && row.len() >= 64 {
2262        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
2263    } else {
2264        unsafe { dot_i8_i8_avx2(row, &act.xq) }
2265    };
2266    let mut acc = dot as f32 * act.sx;
2267    for &(j, xv) in &act.outliers {
2268        acc += (row[j] as i8) as f32 * xv;
2269    }
2270    acc
2271}
2272
2273/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
2274/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
2275/// a single-acc loop runs latency-bound, measured on Granite Rapids).
2276#[cfg(target_arch = "x86_64")]
2277#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
2278unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
2279    // SAFETY: callers uphold slice-length contracts (see call sites).
2280    unsafe {
2281        use core::arch::x86_64::*;
2282        let n = w.len();
2283        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
2284        #[inline(always)]
2285        unsafe fn step(
2286            w: *const u8,
2287            x: *const i8,
2288            flip: core::arch::x86_64::__m512i,
2289            acc: core::arch::x86_64::__m512i,
2290        ) -> core::arch::x86_64::__m512i {
2291            unsafe {
2292                use core::arch::x86_64::*;
2293                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
2294                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
2295            }
2296        }
2297        let (mut a0, mut a1, mut a2, mut a3) = (
2298            _mm512_setzero_si512(),
2299            _mm512_setzero_si512(),
2300            _mm512_setzero_si512(),
2301            _mm512_setzero_si512(),
2302        );
2303        let mut j = 0usize;
2304        while j + 256 <= n {
2305            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
2306            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
2307            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
2308            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
2309            j += 256;
2310        }
2311        while j + 64 <= n {
2312            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
2313            j += 64;
2314        }
2315        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
2316            _mm512_add_epi32(a0, a1),
2317            _mm512_add_epi32(a2, a3),
2318        ));
2319        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
2320        while j < n {
2321            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
2322            j += 1;
2323        }
2324        total
2325    }
2326}
2327
2328/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
2329/// writer's flat order, same as the NEON vzip pair), maddubs against
2330/// the pre-quantized activation group, × the group's f16 scale. Pair
2331/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
2332/// `dot_q4_row_sdot`.
2333#[cfg(target_arch = "x86_64")]
2334#[target_feature(enable = "avx2")]
2335unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2336    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
2337    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
2338    unsafe {
2339        use core::arch::x86_64::*;
2340        let lomask = _mm_set1_epi8(0x0F);
2341        let eight = _mm256_set1_epi8(8);
2342        let ones = _mm256_set1_epi16(1);
2343        let mut acc = 0f32;
2344        for gi in 0..gpr {
2345            let g = g0 + gi;
2346            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2347            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
2348            let lo = _mm_and_si128(b, lomask);
2349            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2350            let w = _mm256_sub_epi8(
2351                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2352                eight,
2353            );
2354            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2355            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2356            let d = _mm256_madd_epi16(p16, ones);
2357            let hi128 = _mm256_extracti128_si256::<1>(d);
2358            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2359            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2360            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2361            acc += _mm_cvtsi128_si32(s32) as f32 * s;
2362        }
2363        acc
2364    }
2365}
2366
2367/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
2368/// both activations dotted against the same centered i8 register.
2369#[cfg(target_arch = "x86_64")]
2370#[target_feature(enable = "avx2")]
2371unsafe fn dot_q4_row_avx2_2(
2372    packed: &[u8],
2373    scales: &[u8],
2374    g0: usize,
2375    gpr: usize,
2376    xq1: &[i8],
2377    xq2: &[i8],
2378) -> (f32, f32) {
2379    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
2380    unsafe {
2381        use core::arch::x86_64::*;
2382        let lomask = _mm_set1_epi8(0x0F);
2383        let eight = _mm256_set1_epi8(8);
2384        let ones = _mm256_set1_epi16(1);
2385        let (mut acc1, mut acc2) = (0f32, 0f32);
2386        #[inline(always)]
2387        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
2388            unsafe {
2389                use core::arch::x86_64::*;
2390                let hi128 = _mm256_extracti128_si256::<1>(d);
2391                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2392                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2393                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2394                _mm_cvtsi128_si32(s32)
2395            }
2396        }
2397        for gi in 0..gpr {
2398            let g = g0 + gi;
2399            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2400            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
2401            let lo = _mm_and_si128(b, lomask);
2402            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2403            let w = _mm256_sub_epi8(
2404                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2405                eight,
2406            );
2407            let aw = _mm256_abs_epi8(w);
2408            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2409            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2410            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
2411            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
2412            acc1 += hsum(d1) as f32 * s;
2413            acc2 += hsum(d2) as f32 * s;
2414        }
2415        (acc1, acc2)
2416    }
2417}
2418
2419/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
2420#[cfg(target_arch = "x86_64")]
2421fn q8_range_avx2(
2422    q: &[u8],
2423    row_scale: &[f32],
2424    act: &SplitAct,
2425    cols: usize,
2426    out_addr: SendMut,
2427    start: usize,
2428    end: usize,
2429) {
2430    for o in start..end {
2431        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
2432        // SAFETY: disjoint row ranges per worker.
2433        unsafe { *out_addr.at(o) = v };
2434    }
2435}
2436
2437/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
2438#[cfg(target_arch = "x86_64")]
2439#[allow(clippy::too_many_arguments)]
2440fn q8_range2_avx2(
2441    q: &[u8],
2442    row_scale: &[f32],
2443    a1: &SplitAct,
2444    a2: &SplitAct,
2445    cols: usize,
2446    p1: SendMut,
2447    p2: SendMut,
2448    start: usize,
2449    end: usize,
2450) {
2451    for o in start..end {
2452        let row = &q[o * cols..(o + 1) * cols];
2453        // SAFETY: disjoint row ranges per worker.
2454        unsafe {
2455            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
2456            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
2457        }
2458    }
2459}
2460
2461// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
2462
2463/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
2464/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
2465/// (On non-ARM release builds only the test tolerance switch calls it.)
2466#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2467fn sdot_enabled() -> bool {
2468    use std::sync::OnceLock;
2469    static ON: OnceLock<bool> = OnceLock::new();
2470    *ON.get_or_init(|| {
2471        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
2472        #[cfg(target_arch = "aarch64")]
2473        {
2474            want && std::arch::is_aarch64_feature_detected!("dotprod")
2475        }
2476        #[cfg(not(target_arch = "aarch64"))]
2477        {
2478            let _ = want;
2479            false
2480        }
2481    })
2482}
2483
2484/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
2485/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
2486/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
2487/// matvec, shared by all rows/workers.
2488struct SplitAct {
2489    xq: Vec<i8>,
2490    sx: f32,
2491    outliers: Vec<(usize, f32)>,
2492    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
2493    /// `−128·Σx`); one i32 per split, computed once per matvec.
2494    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
2495    xsum: i32,
2496}
2497
2498thread_local! {
2499    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
2500    /// and its hidden-size allocation was steady-state heap churn.
2501    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
2502        const { std::cell::RefCell::new(Vec::new()) };
2503}
2504
2505impl Drop for SplitAct {
2506    fn drop(&mut self) {
2507        let buf = std::mem::take(&mut self.xq);
2508        if buf.capacity() > 0 {
2509            XQ_FREE.with(|f| {
2510                let mut f = f.borrow_mut();
2511                if f.len() < 16 {
2512                    f.push(buf);
2513                }
2514            });
2515        }
2516    }
2517}
2518
2519fn split_act(x: &[f32]) -> SplitAct {
2520    let n = x.len();
2521    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
2522    let thr = 8.0 * rms;
2523    // One pass: collect outliers and the bulk absmax (outliers excluded —
2524    // identical to the old zero-then-fold over a copied buffer, minus the
2525    // full-vector copy).
2526    let mut outliers: Vec<(usize, f32)> = Vec::new();
2527    let mut amax = 0f32;
2528    for (j, &v) in x.iter().enumerate() {
2529        let a = v.abs();
2530        if a > thr {
2531            outliers.push((j, v));
2532        } else if a > amax {
2533            amax = a;
2534        }
2535    }
2536    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
2537    let inv = 1.0 / sx;
2538    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
2539    xq.clear();
2540    xq.reserve(n);
2541    if outliers.is_empty() {
2542        xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
2543    } else {
2544        // Outlier slots quantize to 0 (their exact term is added later).
2545        xq.extend(x.iter().map(|&v| {
2546            if v.abs() > thr {
2547                0
2548            } else {
2549                (v * inv).round().clamp(-127.0, 127.0) as i8
2550            }
2551        }));
2552    }
2553    let xsum = xq.iter().map(|&v| v as i32).sum();
2554    SplitAct { xq, sx, outliers, xsum }
2555}
2556
2557/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
2558/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
2559#[cfg(target_arch = "aarch64")]
2560#[target_feature(enable = "neon,dotprod")]
2561unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
2562    // SAFETY: callers uphold slice-length contracts (see call sites).
2563    unsafe {
2564        use core::arch::aarch64::*;
2565        use core::arch::asm;
2566        let wp = w.as_ptr() as *const i8;
2567        let n = w.len();
2568        let (mut a0, mut a1, mut a2, mut a3) =
2569            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
2570        let mut i = 0;
2571        while i + 64 <= n {
2572            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
2573            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
2574            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
2575            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
2576            asm!(
2577                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2578                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2579                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
2580                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
2581                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
2582                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2583                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
2584                options(pure, nomem, nostack),
2585            );
2586            i += 64;
2587        }
2588        while i + 16 <= n {
2589            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
2590            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
2591                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
2592            i += 16;
2593        }
2594        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
2595        while i < n {
2596            s += (*wp.add(i)) as i32 * xq[i] as i32;
2597            i += 1;
2598        }
2599        s
2600}
2601}
2602
2603/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
2604/// loaded once and reused, 4 independent accumulators hide sdot latency
2605/// (port of vmfcore `dot_i8_sdot_4rows`).
2606#[cfg(target_arch = "aarch64")]
2607#[target_feature(enable = "neon,dotprod")]
2608unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
2609    // SAFETY: callers uphold slice-length contracts (see call sites).
2610    unsafe {
2611        use core::arch::aarch64::*;
2612        use core::arch::asm;
2613        let n = xq.len();
2614        let px = xq.as_ptr();
2615        let (p0, p1, p2, p3) = (
2616            w0.as_ptr() as *const i8,
2617            w1.as_ptr() as *const i8,
2618            w2.as_ptr() as *const i8,
2619            w3.as_ptr() as *const i8,
2620        );
2621        let (mut a0, mut a1, mut a2, mut a3) =
2622            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
2623        let mut i = 0;
2624        while i + 16 <= n {
2625            let x = vld1q_s8(px.add(i));
2626            let v0 = vld1q_s8(p0.add(i));
2627            let v1 = vld1q_s8(p1.add(i));
2628            let v2 = vld1q_s8(p2.add(i));
2629            let v3 = vld1q_s8(p3.add(i));
2630            asm!(
2631                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
2632                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
2633                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
2634                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
2635                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
2636                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
2637                options(pure, nomem, nostack),
2638            );
2639            i += 16;
2640        }
2641        let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
2642        while i < n {
2643            let xi = *px.add(i) as i32;
2644            r[0] += (*p0.add(i)) as i32 * xi;
2645            r[1] += (*p1.add(i)) as i32 * xi;
2646            r[2] += (*p2.add(i)) as i32 * xi;
2647            r[3] += (*p3.add(i)) as i32 * xi;
2648            i += 1;
2649        }
2650        r
2651}
2652}
2653
2654/// One q8 row range via SDOT (4-row blocks + tail) — the body of
2655/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
2656/// SAME kernel for several tensors under one pool dispatch.
2657#[cfg(target_arch = "aarch64")]
2658fn q8_range_sdot(
2659    q: &[u8],
2660    row_scale: &[f32],
2661    act: &SplitAct,
2662    cols: usize,
2663    out_addr: SendMut,
2664    start: usize,
2665    end: usize,
2666) {
2667    let mut o = start;
2668    while o + 4 <= end {
2669        let r = unsafe {
2670            dot_i8_sdot_4rows(
2671                &q[o * cols..(o + 1) * cols],
2672                &q[(o + 1) * cols..(o + 2) * cols],
2673                &q[(o + 2) * cols..(o + 3) * cols],
2674                &q[(o + 3) * cols..(o + 4) * cols],
2675                &act.xq,
2676            )
2677        };
2678        for k in 0..4 {
2679            let mut acc = r[k] as f32 * act.sx;
2680            for &(j, xv) in &act.outliers {
2681                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
2682            }
2683            // SAFETY: disjoint row ranges per worker.
2684            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
2685        }
2686        o += 4;
2687    }
2688    while o < end {
2689        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
2690        unsafe { *out_addr.at(o) = v };
2691        o += 1;
2692    }
2693}
2694
2695/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
2696/// for the fused pair multi-matrix job (`matvec2_many`).
2697#[cfg(target_arch = "aarch64")]
2698#[allow(clippy::too_many_arguments)]
2699fn q8_range2_sdot(
2700    q: &[u8],
2701    row_scale: &[f32],
2702    a1: &SplitAct,
2703    a2: &SplitAct,
2704    cols: usize,
2705    p1: SendMut,
2706    p2: SendMut,
2707    start: usize,
2708    end: usize,
2709) {
2710    for o in start..end {
2711        let row = &q[o * cols..(o + 1) * cols];
2712        // SAFETY: disjoint row ranges per worker.
2713        unsafe {
2714            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
2715            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
2716        }
2717    }
2718}
2719
2720/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
2721#[allow(clippy::too_many_arguments)]
2722fn q8_range2_f32(
2723    q: &[u8],
2724    row_scale: &[f32],
2725    x1: &[f32],
2726    x2: &[f32],
2727    cols: usize,
2728    p1: SendMut,
2729    p2: SendMut,
2730    start: usize,
2731    end: usize,
2732) {
2733    for o in start..end {
2734        let row = &q[o * cols..(o + 1) * cols];
2735        // SAFETY: disjoint row ranges per worker.
2736        unsafe {
2737            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
2738            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
2739        }
2740    }
2741}
2742
2743/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
2744fn q8_range_f32(
2745    q: &[u8],
2746    row_scale: &[f32],
2747    xs: &[f32],
2748    cols: usize,
2749    out_addr: SendMut,
2750    start: usize,
2751    end: usize,
2752) {
2753    for o in start..end {
2754        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
2755        // SAFETY: disjoint row ranges per worker.
2756        unsafe { *out_addr.at(o) = v };
2757    }
2758}
2759
2760/// SDOT row dot with exact outlier correction:
2761/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
2762#[cfg(target_arch = "aarch64")]
2763#[inline]
2764fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
2765    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
2766    for &(j, xv) in &act.outliers {
2767        acc += (row[j] as i8) as f32 * xv;
2768    }
2769    acc
2770}
2771
2772/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
2773/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
2774/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
2775/// the caller multiplies by the activation scale and adds the exact
2776/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
2777/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
2778/// → zip(lo,hi) restores flat order.
2779#[cfg(target_arch = "aarch64")]
2780#[target_feature(enable = "neon,dotprod")]
2781unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2782    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
2783    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
2784    unsafe {
2785        use core::arch::aarch64::*;
2786        use core::arch::asm;
2787        let lomask = vdupq_n_u8(0x0F);
2788        let eight = vdupq_n_s8(8);
2789        let mut acc = 0f32;
2790        for gi in 0..gpr {
2791            let g = g0 + gi;
2792            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2793            let b = vld1q_u8(packed.as_ptr().add(g * 16));
2794            let lo = vandq_u8(b, lomask);
2795            let hi = vshrq_n_u8::<4>(b);
2796            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2797            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2798            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2799            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2800            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2801            asm!(
2802                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2803                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2804                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2805                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2806                options(pure, nomem, nostack),
2807            );
2808            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2809        }
2810        acc
2811    }
2812}
2813
2814/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
2815/// part) happens ONCE per group; both pre-quantized activations are
2816/// dotted against the same centered i8 registers. Per-lane math matches
2817/// `dot_q4_row_sdot` exactly.
2818#[cfg(target_arch = "aarch64")]
2819#[target_feature(enable = "neon,dotprod")]
2820unsafe fn dot_q4_row_sdot2(
2821    packed: &[u8],
2822    scales: &[u8],
2823    g0: usize,
2824    gpr: usize,
2825    xq1: &[i8],
2826    xq2: &[i8],
2827) -> (f32, f32) {
2828    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
2829    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
2830    unsafe {
2831        use core::arch::aarch64::*;
2832        use core::arch::asm;
2833        let lomask = vdupq_n_u8(0x0F);
2834        let eight = vdupq_n_s8(8);
2835        let (mut acc1, mut acc2) = (0f32, 0f32);
2836        for gi in 0..gpr {
2837            let g = g0 + gi;
2838            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2839            let b = vld1q_u8(packed.as_ptr().add(g * 16));
2840            let lo = vandq_u8(b, lomask);
2841            let hi = vshrq_n_u8::<4>(b);
2842            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2843            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2844            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
2845            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
2846            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
2847            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
2848            let (mut a0, mut a1, mut b0, mut b1) =
2849                (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
2850            asm!(
2851                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
2852                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
2853                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
2854                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
2855                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2856                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
2857                e0 = in(vreg) e0, e1 = in(vreg) e1,
2858                x10 = in(vreg) x10, x11 = in(vreg) x11,
2859                x20 = in(vreg) x20, x21 = in(vreg) x21,
2860                options(pure, nomem, nostack),
2861            );
2862            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2863            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
2864        }
2865        (acc1, acc2)
2866    }
2867}
2868
2869// ───────────────────── fused int8 kernels ─────────────────────
2870
2871/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
2872/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
2873#[inline]
2874pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
2875    #[cfg(target_arch = "aarch64")]
2876    unsafe {
2877        return axpy_i8_f32_neon(acc, row, w);
2878    }
2879    #[cfg(target_arch = "x86_64")]
2880    if avx2_enabled() {
2881        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
2882    }
2883    #[allow(unreachable_code)]
2884    {
2885        for (a, &b) in acc.iter_mut().zip(row) {
2886            *a += w * b as f32;
2887        }
2888    }
2889}
2890
2891/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
2892#[cfg(target_arch = "x86_64")]
2893#[target_feature(enable = "avx2,fma")]
2894unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
2895    // SAFETY: callers uphold slice-length contracts (see call sites).
2896    unsafe {
2897        use core::arch::x86_64::*;
2898        let n = acc.len().min(row.len());
2899        let ap = acc.as_mut_ptr();
2900        let rp = row.as_ptr();
2901        let wv = _mm256_set1_ps(w);
2902        let mut j = 0usize;
2903        while j + 16 <= n {
2904            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
2905            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
2906            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
2907            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
2908            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
2909            _mm256_storeu_ps(ap.add(j), v0);
2910            _mm256_storeu_ps(ap.add(j + 8), v1);
2911            j += 16;
2912        }
2913        while j < n {
2914            *ap.add(j) += w * (*rp.add(j)) as f32;
2915            j += 1;
2916        }
2917    }
2918}
2919
2920#[cfg(target_arch = "aarch64")]
2921#[target_feature(enable = "neon")]
2922unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
2923    // SAFETY: callers uphold slice-length contracts (see call sites).
2924    unsafe {
2925        use core::arch::aarch64::*;
2926        let n = acc.len().min(row.len());
2927        let ap = acc.as_mut_ptr();
2928        let rp = row.as_ptr();
2929        let wv = vdupq_n_f32(w);
2930        let mut j = 0usize;
2931        while j + 16 <= n {
2932            let rb = vld1q_s8(rp.add(j));
2933            let lo = vmovl_s8(vget_low_s8(rb));
2934            let hi = vmovl_s8(vget_high_s8(rb));
2935            for (off, half) in [(0, lo), (8, hi)] {
2936                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
2937                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
2938                let o = j + off;
2939                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
2940                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
2941            }
2942            j += 16;
2943        }
2944        while j < n {
2945            *ap.add(j) += w * (*rp.add(j)) as f32;
2946            j += 1;
2947        }
2948}
2949}
2950
2951/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
2952/// ≈9× scalar), scalar elsewhere.
2953#[inline]
2954pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
2955    #[cfg(target_arch = "aarch64")]
2956    unsafe {
2957        return dot_i8_f32_neon(w, x);
2958    }
2959    #[cfg(target_arch = "x86_64")]
2960    if avx2_enabled() {
2961        return unsafe { dot_i8_f32_avx2(w, x) };
2962    }
2963    #[allow(unreachable_code)]
2964    {
2965        let mut sum = 0.0f32;
2966        for (j, &b) in w.iter().enumerate() {
2967            sum += (b as i8) as f32 * x[j];
2968        }
2969        sum
2970    }
2971}
2972
2973/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
2974/// folded into the product (no prescaled copy of x). NEON on aarch64,
2975/// scalar elsewhere. Used by the active-neuron path `row_dot`.
2976#[inline]
2977fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
2978    #[cfg(target_arch = "aarch64")]
2979    unsafe {
2980        return dot_i8_col_f32_neon(w, x, col);
2981    }
2982    #[allow(unreachable_code)]
2983    {
2984        let mut sum = 0.0f32;
2985        for (j, &b) in w.iter().enumerate() {
2986            sum += (b as i8) as f32 * x[j] * col[j];
2987        }
2988        sum
2989    }
2990}
2991
2992#[cfg(target_arch = "aarch64")]
2993#[target_feature(enable = "neon")]
2994unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
2995    // SAFETY: callers uphold slice-length contracts (see call sites).
2996    unsafe {
2997        use core::arch::aarch64::*;
2998        let n = x.len();
2999        let wp = w.as_ptr() as *const i8;
3000        let xp = x.as_ptr();
3001        let cp = col.as_ptr();
3002        let (mut a0, mut a1, mut a2, mut a3) =
3003            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3004        let mut j = 0usize;
3005        while j + 16 <= n {
3006            let wb = vld1q_s8(wp.add(j));
3007            let lo = vmovl_s8(vget_low_s8(wb));
3008            let hi = vmovl_s8(vget_high_s8(wb));
3009            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
3010            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
3011            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
3012            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
3013            a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
3014            a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
3015            a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
3016            a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
3017            j += 16;
3018        }
3019        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
3020        while j < n {
3021            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
3022            j += 1;
3023        }
3024        sum
3025}
3026}
3027
3028#[cfg(target_arch = "aarch64")]
3029#[target_feature(enable = "neon")]
3030unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
3031    // SAFETY: callers uphold slice-length contracts (see call sites).
3032    unsafe {
3033        use core::arch::aarch64::*;
3034        let n = x.len();
3035        let wp = w.as_ptr() as *const i8;
3036        let xp = x.as_ptr();
3037        let (mut a0, mut a1, mut a2, mut a3) =
3038            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3039        let mut j = 0usize;
3040        while j + 16 <= n {
3041            let wb = vld1q_s8(wp.add(j));
3042            let lo = vmovl_s8(vget_low_s8(wb));
3043            let hi = vmovl_s8(vget_high_s8(wb));
3044            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
3045            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
3046            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
3047            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
3048            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
3049            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
3050            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
3051            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
3052            j += 16;
3053        }
3054        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
3055        while j < n {
3056            sum += (*wp.add(j)) as f32 * *xp.add(j);
3057            j += 1;
3058        }
3059        sum
3060}
3061}
3062
3063#[allow(clippy::too_many_arguments)]
3064fn qmatvec(
3065    q: &[u8],
3066    row_scale: &[f32],
3067    xs: &[f32],
3068    rows: usize,
3069    cols: usize,
3070    out: &mut [f32],
3071    pool: Option<&Pool>,
3072) {
3073    debug_assert_eq!(out.len(), rows);
3074
3075    #[cfg(target_arch = "aarch64")]
3076    if sdot_enabled() {
3077        let act = split_act(xs);
3078        let out_addr = SendMut(out.as_mut_ptr());
3079        let run_range =
3080            |start: usize, end: usize| q8_range_sdot(q, row_scale, &act, cols, out_addr, start, end);
3081        match pool {
3082            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3083            _ => run_range(0, rows),
3084        }
3085        return;
3086    }
3087    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
3088    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
3089    #[cfg(target_arch = "x86_64")]
3090    if avx2_a8w8_enabled() {
3091        let act = split_act(xs);
3092        let out_addr = SendMut(out.as_mut_ptr());
3093        let run_range =
3094            |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
3095        match pool {
3096            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3097            _ => run_range(0, rows),
3098        }
3099        return;
3100    }
3101
3102    let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
3103    match pool {
3104        Some(pool) if rows >= 256 => {
3105            let out_addr = SendMut(out.as_mut_ptr());
3106            pool.run(&move |widx, n| {
3107                let chunk = rows.div_ceil(n);
3108                let start = widx * chunk;
3109                let end = (start + chunk).min(rows);
3110                for o in start..end {
3111                    // SAFETY: disjoint row ranges per worker.
3112                    unsafe { *out_addr.at(o) = row_dot(o) };
3113                }
3114            });
3115        }
3116        _ => {
3117            for (o, dst) in out.iter_mut().enumerate() {
3118                *dst = row_dot(o);
3119            }
3120        }
3121    }
3122}
3123
3124#[allow(clippy::too_many_arguments)]
3125fn qmatvec2(
3126    q: &[u8],
3127    row_scale: &[f32],
3128    x1: &[f32],
3129    x2: &[f32],
3130    rows: usize,
3131    cols: usize,
3132    o1: &mut [f32],
3133    o2: &mut [f32],
3134    pool: Option<&Pool>,
3135) {
3136    #[cfg(target_arch = "aarch64")]
3137    if sdot_enabled() {
3138        let a1s = split_act(x1);
3139        let a2s = split_act(x2);
3140        let p1 = SendMut(o1.as_mut_ptr());
3141        let p2 = SendMut(o2.as_mut_ptr());
3142        let run_range = |start: usize, end: usize| {
3143            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
3144        };
3145        match pool {
3146            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3147            _ => run_range(0, rows),
3148        }
3149        return;
3150    }
3151    #[cfg(target_arch = "x86_64")]
3152    if avx2_a8w8_enabled() {
3153        let a1s = split_act(x1);
3154        let a2s = split_act(x2);
3155        let p1 = SendMut(o1.as_mut_ptr());
3156        let p2 = SendMut(o2.as_mut_ptr());
3157        let run_range = |start: usize, end: usize| {
3158            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
3159        };
3160        match pool {
3161            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3162            _ => run_range(0, rows),
3163        }
3164        return;
3165    }
3166
3167    let row_dots = |o: usize| -> (f32, f32) {
3168        let row = &q[o * cols..(o + 1) * cols];
3169        (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
3170    };
3171    match pool {
3172        Some(pool) if rows >= 256 => {
3173            let p1 = SendMut(o1.as_mut_ptr());
3174            let p2 = SendMut(o2.as_mut_ptr());
3175            pool.run(&move |widx, n| {
3176                let chunk = rows.div_ceil(n);
3177                let start = widx * chunk;
3178                let end = (start + chunk).min(rows);
3179                for o in start..end {
3180                    let (s1, s2) = row_dots(o);
3181                    // SAFETY: disjoint row ranges per worker.
3182                    unsafe {
3183                        *p1.at(o) = s1;
3184                        *p2.at(o) = s2;
3185                    }
3186                }
3187            });
3188        }
3189        _ => {
3190            for o in 0..rows {
3191                let (s1, s2) = row_dots(o);
3192                o1[o] = s1;
3193                o2[o] = s2;
3194            }
3195        }
3196    }
3197}
3198
3199#[derive(Clone, Copy)]
3200struct SendMut(*mut f32);
3201unsafe impl Send for SendMut {}
3202unsafe impl Sync for SendMut {}
3203
3204impl SendMut {
3205    #[inline]
3206    fn at(self, i: usize) -> *mut f32 {
3207        unsafe { self.0.add(i) }
3208    }
3209}
3210
3211#[cfg(test)]
3212mod tests {
3213    use super::*;
3214
3215    #[test]
3216    fn f32_matvec_matches_matvec_rows_bitexact() {
3217        let (rows, cols) = (300, 40);
3218        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
3219        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
3220        let qt = QTensor::from_f32(w.clone(), rows, cols);
3221
3222        let mut a = vec![0.0f32; rows];
3223        matvec_rows(None, &w, &x, &mut a);
3224        let mut b = vec![0.0f32; rows];
3225        qt.matvec(&x, &mut b, None);
3226        assert_eq!(a, b);
3227    }
3228
3229    #[test]
3230    fn sdot_kernel_exact_on_grid() {
3231        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
3232        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
3233        // exact f32 dot to float rounding. This isolates kernel
3234        // correctness from quantization noise.
3235        eprintln!("sdot_enabled = {}", sdot_enabled());
3236        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
3237        let w: Vec<u8> = (0..rows * cols)
3238            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
3239            .collect();
3240        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
3241        let x: Vec<f32> = (0..cols)
3242            .map(|i| match i % 3 {
3243                0 => 1.0,
3244                1 => -1.0,
3245                _ => 0.0,
3246            })
3247            .collect();
3248        let mut a = vec![0.0f32; rows];
3249        qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
3250        for o in 0..rows {
3251            let mut acc = 0.0f32;
3252            for j in 0..cols {
3253                acc += (w[o * cols + j] as i8) as f32 * x[j];
3254            }
3255            let expect = acc * scales[o];
3256            assert!(
3257                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
3258                "row {o}: {} vs {expect}",
3259                a[o]
3260            );
3261        }
3262    }
3263
3264    #[test]
3265    fn sdot_a8w8_noise_is_bounded() {
3266        // Off-grid activations: A8 quantization noise must stay small in
3267        // relative L2 over the whole output (realistic accuracy contract;
3268        // vmfcore measured argmax-identical decode on real models).
3269        let (rows, cols) = (16, 512);
3270        let w: Vec<u8> = (0..rows * cols)
3271            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
3272            .collect();
3273        let scales = vec![0.01f32; rows];
3274        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
3275        let mut a = vec![0.0f32; rows];
3276        qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
3277        let (mut num, mut den) = (0f64, 0f64);
3278        for o in 0..rows {
3279            let mut acc = 0.0f32;
3280            for j in 0..cols {
3281                acc += (w[o * cols + j] as i8) as f32 * x[j];
3282            }
3283            let expect = acc * scales[o];
3284            num += ((a[o] - expect) as f64).powi(2);
3285            den += (expect as f64).powi(2);
3286        }
3287        let rel = (num / den.max(1e-12)).sqrt();
3288        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
3289    }
3290
3291    #[test]
3292    fn i8_dot_neon_matches_scalar() {
3293        let n = 100;
3294        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
3295        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
3296        let mut scalar = 0.0f32;
3297        for j in 0..n {
3298            scalar += (w[j] as i8) as f32 * x[j];
3299        }
3300        let fast = dot_i8_f32(&w, &x);
3301        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
3302    }
3303
3304    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
3305    #[test]
3306    fn vbitmatvec_matches_full_dequant() {
3307        let (rows, cols) = (6, 64);
3308        let ng = cols / GROUP_SIZE;
3309        // Hand-craft: bits per row, f16 scales, packed rows.
3310        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
3311        let mut bytes = bits.clone();
3312        for g in 0..rows * ng {
3313            let s = 0.02 + 0.001 * g as f32;
3314            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3315        }
3316        for r in 0..rows {
3317            let b = bits[r] as usize;
3318            let (mut acc, mut nb) = (0u64, 0usize);
3319            let mut rowbytes = Vec::new();
3320            for i in 0..cols {
3321                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
3322                acc = (acc << b) | v;
3323                nb += b;
3324                while nb >= 8 {
3325                    nb -= 8;
3326                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
3327                }
3328            }
3329            if nb > 0 {
3330                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
3331            }
3332            bytes.extend_from_slice(&rowbytes);
3333        }
3334        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
3335
3336        let mut reference = vec![0f32; rows * cols];
3337        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
3338        let mut expect = vec![0f32; rows];
3339        for r in 0..rows {
3340            expect[r] = reference[r * cols..(r + 1) * cols]
3341                .iter()
3342                .zip(&x)
3343                .map(|(w, xv)| w * xv)
3344                .sum();
3345        }
3346        let mut got = vec![0f32; rows];
3347        let offsets = vbit_row_offsets(&bytes, rows, cols);
3348        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
3349        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
3350        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
3351        // the golden-parity gate).
3352        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
3353        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
3354        for r in 0..rows {
3355            assert!(
3356                (got[r] - expect[r]).abs() < tol * scale,
3357                "row {r}: {} vs {}",
3358                got[r],
3359                expect[r]
3360            );
3361        }
3362    }
3363
3364    /// Fused q4 matvec must match the reference full-dequant + dense
3365    /// matvec bit-for-bit in structure (same f32 math, group order).
3366    #[test]
3367    fn q4matvec_matches_full_dequant() {
3368        let (rows, cols) = (8, 64);
3369        let groups = rows * cols / GROUP_SIZE;
3370        // Hand-craft a q4_block blob: nibbles then f16 scales.
3371        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
3372        for i in 0..groups * 16 {
3373            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
3374        }
3375        for g in 0..groups {
3376            let s = 0.01 + 0.003 * g as f32;
3377            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3378        }
3379        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
3380
3381        let mut reference = vec![0.0f32; rows * cols];
3382        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
3383        let mut expect = vec![0.0f32; rows];
3384        for r in 0..rows {
3385            expect[r] = reference[r * cols..(r + 1) * cols]
3386                .iter()
3387                .zip(&x)
3388                .map(|(w, xv)| w * xv)
3389                .sum();
3390        }
3391
3392        let mut got = vec![0.0f32; rows];
3393        q4matvec(&bytes, &x, rows, cols, &mut got, None);
3394        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
3395        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
3396        // in the golden-parity gate).
3397        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
3398        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
3399        for r in 0..rows {
3400            assert!(
3401                (got[r] - expect[r]).abs() < tol * scale,
3402                "row {r}: {} vs {}",
3403                got[r],
3404                expect[r]
3405            );
3406        }
3407    }
3408
3409    /// Fused two-input vbit matvec must equal two single matvecs exactly
3410    /// (same per-lane accumulation order on both scalar and SDOT paths).
3411    #[test]
3412    fn vbitmatvec2_equals_two_singles() {
3413        let (rows, cols) = (6, 64);
3414        let ng = cols / GROUP_SIZE;
3415        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
3416        let mut bytes = bits.clone();
3417        for g in 0..rows * ng {
3418            let s = 0.02 + 0.001 * g as f32;
3419            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3420        }
3421        for r in 0..rows {
3422            let b = bits[r] as usize;
3423            let (mut acc, mut nb) = (0u64, 0usize);
3424            let mut rowbytes = Vec::new();
3425            for i in 0..cols {
3426                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
3427                acc = (acc << b) | v;
3428                nb += b;
3429                while nb >= 8 {
3430                    nb -= 8;
3431                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
3432                }
3433            }
3434            if nb > 0 {
3435                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
3436            }
3437            bytes.extend_from_slice(&rowbytes);
3438        }
3439        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
3440        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
3441        let offsets = vbit_row_offsets(&bytes, rows, cols);
3442
3443        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
3444        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
3445        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
3446        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
3447        vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
3448        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
3449        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
3450    }
3451
3452    /// Fused two-input q4 matvec must equal two single matvecs exactly.
3453    #[test]
3454    fn q4matvec2_equals_two_singles() {
3455        let (rows, cols) = (8, 128);
3456        let groups = rows * cols / GROUP_SIZE;
3457        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
3458        for i in 0..groups * 16 {
3459            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
3460        }
3461        for g in 0..groups {
3462            let s = 0.01 + 0.003 * g as f32;
3463            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3464        }
3465        // Include an outlier channel so the SDOT correction path is
3466        // exercised in the pair kernel too.
3467        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
3468        x1[9] = 250.0;
3469        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
3470
3471        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
3472        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
3473        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
3474        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
3475        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
3476        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
3477        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
3478    }
3479
3480    /// Multi-matrix job must equal separate matvecs exactly — same
3481    /// kernels, only the dispatch is fused.
3482    #[test]
3483    fn matvec_many_equals_separate_matvecs() {
3484        use crate::pool::Pool;
3485        let (r1, r2, cols) = (300, 200, 64);
3486        let mk = |salt: usize, rows: usize| {
3487            QTensor::from_f32(
3488                (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
3489                rows,
3490                cols,
3491            )
3492        };
3493        let (a, b) = (mk(1, r1), mk(5, r2));
3494        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
3495        let pool = Pool::new(3);
3496
3497        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
3498        a.matvec(&x, &mut ea, Some(&pool));
3499        b.matvec(&x, &mut eb, Some(&pool));
3500        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
3501        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
3502        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
3503        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
3504    }
3505
3506    /// Batched q4/vbit matmat must equal per-position matvec calls
3507    /// exactly (the fallback it replaced) — same kernels, same order.
3508    #[test]
3509    fn batched_matmat_equals_per_position_matvec() {
3510        let (rows, cols, b) = (8, 64, 5);
3511        // q4 blob.
3512        let groups = rows * cols / GROUP_SIZE;
3513        let mut q4 = Vec::new();
3514        for i in 0..groups * 16 {
3515            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
3516        }
3517        for g in 0..groups {
3518            q4.extend_from_slice(
3519                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
3520            );
3521        }
3522        // vbit blob (mixed widths incl. 8).
3523        let ng = cols / GROUP_SIZE;
3524        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
3525        let mut vb = bits.clone();
3526        for g in 0..rows * ng {
3527            vb.extend_from_slice(
3528                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
3529            );
3530        }
3531        for r in 0..rows {
3532            let bw = bits[r] as usize;
3533            let (mut acc, mut nb) = (0u64, 0usize);
3534            let mut rowbytes = Vec::new();
3535            for i in 0..cols {
3536                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
3537                acc = (acc << bw) | v;
3538                nb += bw;
3539                while nb >= 8 {
3540                    nb -= 8;
3541                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
3542                }
3543            }
3544            if nb > 0 {
3545                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
3546            }
3547            vb.extend_from_slice(&rowbytes);
3548        }
3549        let offsets = vbit_row_offsets(&vb, rows, cols);
3550
3551        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
3552
3553        // q4: batch vs singles.
3554        let mut got = vec![0f32; b * rows];
3555        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
3556        for bi in 0..b {
3557            let mut expect = vec![0f32; rows];
3558            q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
3559            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
3560        }
3561
3562        // vbit: batch vs singles.
3563        let mut got = vec![0f32; b * rows];
3564        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
3565        for bi in 0..b {
3566            let mut expect = vec![0f32; rows];
3567            vbitmatvec(
3568                &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
3569            );
3570            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
3571        }
3572    }
3573
3574    /// q4 SDOT outlier correction: a single huge activation channel
3575    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
3576    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
3577    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
3578    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
3579    /// can never qualify (8² = n).
3580    #[test]
3581    fn q4matvec_sdot_outlier_exact() {
3582        let (rows, cols) = (4, 128);
3583        let groups = rows * cols / GROUP_SIZE;
3584        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
3585        for i in 0..groups * 16 {
3586            bytes.push(((i * 11 + 5) % 256) as u8);
3587        }
3588        for g in 0..groups {
3589            let s = 0.02 + 0.002 * g as f32;
3590            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3591        }
3592        let mut x: Vec<f32> = (0..cols)
3593            .map(|i| match i % 3 {
3594                0 => 1.0,
3595                1 => -1.0,
3596                _ => 0.0,
3597            })
3598            .collect();
3599        x[17] = 300.0; // ≫ 8·rms → outlier channel
3600
3601        let mut reference = vec![0.0f32; rows * cols];
3602        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
3603        let mut expect = vec![0.0f32; rows];
3604        for r in 0..rows {
3605            expect[r] = reference[r * cols..(r + 1) * cols]
3606                .iter()
3607                .zip(&x)
3608                .map(|(w, xv)| w * xv)
3609                .sum();
3610        }
3611        let mut got = vec![0.0f32; rows];
3612        q4matvec(&bytes, &x, rows, cols, &mut got, None);
3613        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
3614        for r in 0..rows {
3615            assert!(
3616                (got[r] - expect[r]).abs() < 2e-3 * scale,
3617                "row {r}: {} vs {} (outlier term must be exact)",
3618                got[r],
3619                expect[r]
3620            );
3621        }
3622    }
3623}