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, Q1_TILE, Q4_TILE};
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        /// q8-family decode repack (load-time, optional): rows in groups
41        /// of 4, interleaved in 16-byte units — one 64-byte line per
42        /// iteration feeds all 4 sdot lanes, ONE sequential weight
43        /// stream per worker instead of four (this is where llama.cpp's
44        /// repacked Q8 kernels get their bandwidth). Empty = off
45        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
46        /// an anonymous copy of the quants for mmap pages that go cold.
47        repack: Vec<u8>,
48    },
49}
50
51/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
52/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
53/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
54/// streams per worker feed the prefetcher MORE memory-level parallelism
55/// than one); kept as an experiment flag for x86, where the tradeoff
56/// may land differently.
57fn repack_enabled() -> bool {
58    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
59    *ON.get_or_init(|| {
60        std::env::var("CMF_REPACK").map(|v| v == "1").unwrap_or(false)
61    })
62}
63
64/// Interleave q8 rows for the decode kernel: group g holds rows
65/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
66/// full groups are packed — tail rows keep reading the mmap layout.
67fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
68    #[cfg(target_arch = "aarch64")]
69    let arch_ok = sdot_enabled();
70    #[cfg(not(target_arch = "aarch64"))]
71    let arch_ok = false;
72    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
73        return Vec::new();
74    }
75    q8_repack_layout(bytes, rows, cols)
76}
77
78/// The pure layout transform behind `q8_repack` (tested directly —
79/// the gate depends on arch and env).
80fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
81    let groups = rows / 4;
82    let mut rep = vec![0u8; groups * 4 * cols];
83    for g in 0..groups {
84        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
85        for c in 0..cols / 16 {
86            for lane in 0..4 {
87                let src = (g * 4 + lane) * cols + c * 16;
88                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
89                    .copy_from_slice(&bytes[src..src + 16]);
90            }
91        }
92    }
93    rep
94}
95
96/// Prefix-sum of vbit row payload offsets (absolute within the tensor
97/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
98fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
99    let ng = cols / GROUP_SIZE;
100    let bits = &bytes[..rows];
101    let mut offsets = Vec::with_capacity(rows + 1);
102    let mut off = rows + rows * ng * 2;
103    for r in 0..rows {
104        offsets.push(off);
105        off += (cols * bits[r] as usize + 7) / 8;
106    }
107    offsets.push(off);
108    offsets
109}
110
111impl QTensor {
112    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
113        debug_assert_eq!(data.len(), rows * cols);
114        Self::F32 { data, rows, cols }
115    }
116
117    /// Wrap a directory tensor without dequantizing the payload.
118    /// Falls back to dequantized f32 for dtypes without a fused kernel.
119    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
120        // Indexed lookup: the linear directory scan made pipeline build
121        // O(N²) on MoE/skills files with thousands of tensors.
122        let idx = model
123            .tensor_index(name)
124            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
125        let entry = &model.tensors[idx];
126        if entry.shape.len() != 2 {
127            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
128        }
129        let (rows, cols) = (entry.shape[0], entry.shape[1]);
130        let bytes = model.entry_bytes(entry);
131
132        match entry.dtype {
133            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
134                let n = rows * cols;
135                let scales_off = n;
136                let row_scale: Vec<f32> = (0..rows)
137                    .map(|o| {
138                        f16_to_f32(u16::from_le_bytes([
139                            bytes[scales_off + o * 2],
140                            bytes[scales_off + o * 2 + 1],
141                        ]))
142                    })
143                    .collect();
144                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
145                    let col_off = n + rows * 2;
146                    (0..cols)
147                        .map(|i| {
148                            f16_to_f32(u16::from_le_bytes([
149                                bytes[col_off + i * 2],
150                                bytes[col_off + i * 2 + 1],
151                            ]))
152                        })
153                        .collect()
154                } else {
155                    Vec::new()
156                };
157                Ok(Self::Mapped {
158                    model: model.clone(),
159                    idx,
160                    dtype: entry.dtype,
161                    rows,
162                    cols,
163                    row_scale,
164                    col_field,
165                    vbit_offsets: Vec::new(),
166                    repack: q8_repack(bytes, rows, cols),
167                })
168            }
169            // vbit: fused kernel unpacks variable-bit rows from mmap.
170            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
171                model: model.clone(),
172                idx,
173                dtype: entry.dtype,
174                rows,
175                cols,
176                row_scale: Vec::new(),
177                col_field: Vec::new(),
178                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
179                repack: Vec::new(),
180            }),
181            // vbit_ro (§4.2): the offset table comes straight from the
182            // file — no load-time prefix scan; kernels are shared with
183            // legacy vbit (they consume absolute offsets either way).
184            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
185                let (_, off_off, packed_off) =
186                    cortiq_core::quant::vbit_ro_sections(rows, cols);
187                let offsets: Vec<usize> = (0..=rows)
188                    .map(|r| {
189                        packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r)
190                    })
191                    .collect();
192                Ok(Self::Mapped {
193                    model: model.clone(),
194                    idx,
195                    dtype: entry.dtype,
196                    rows,
197                    cols,
198                    row_scale: Vec::new(),
199                    col_field: Vec::new(),
200                    vbit_offsets: offsets,
201                    repack: Vec::new(),
202                })
203            }
204            // q4_block: fused kernel reads nibbles straight from mmap —
205            // a 14B q4 file no longer explodes into ×8 f32 RAM.
206            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
207            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
208            // at kernel level over the split layout).
209            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
210                model: model.clone(),
211                idx,
212                dtype: entry.dtype,
213                rows,
214                cols,
215                row_scale: Vec::new(),
216                col_field: Vec::new(),
217                vbit_offsets: Vec::new(),
218                repack: Vec::new(),
219            }),
220            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
221                model: model.clone(),
222                idx,
223                dtype: entry.dtype,
224                rows,
225                cols,
226                row_scale: Vec::new(),
227                col_field: Vec::new(),
228                vbit_offsets: Vec::new(),
229                repack: Vec::new(),
230            }),
231            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
232            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
233                model: model.clone(),
234                idx,
235                dtype: entry.dtype,
236                rows,
237                cols,
238                row_scale: Vec::new(),
239                col_field: Vec::new(),
240                vbit_offsets: Vec::new(),
241                repack: Vec::new(),
242            }),
243            // No fused kernel yet → dequantize once (correct, more RAM).
244            _ => {
245                let mut data = vec![0.0f32; rows * cols];
246                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
247                Ok(Self::from_f32(data, rows, cols))
248            }
249        }
250    }
251
252    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
253    /// compute-bound, so offload pays at much smaller shapes than q8.)
254    pub(crate) fn is_q1(&self) -> bool {
255        matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
256    }
257
258    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
259    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
260    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
261        match self {
262            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
263            _ => None,
264        }
265    }
266
267    /// (directory idx, rows, cols) of a q1-mapped tensor — the
268    /// whole-block GPU path resolves offsets itself.
269    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
270        match self {
271            Self::Mapped { idx, dtype: TensorDtype::Q1, rows, cols, .. } => {
272                Some((*idx, *rows, *cols))
273            }
274            _ => None,
275        }
276    }
277
278    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
279    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
280    /// q8_2f is excluded on purpose: its column field would need a
281    /// prescale stage on the device.
282    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
283        match self {
284            Self::Mapped {
285                idx,
286                dtype: TensorDtype::Q8Row,
287                rows,
288                cols,
289                row_scale,
290                col_field,
291                ..
292            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
293            _ => None,
294        }
295    }
296
297    pub fn rows(&self) -> usize {
298        match self {
299            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
300        }
301    }
302
303    pub fn cols(&self) -> usize {
304        match self {
305            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
306        }
307    }
308
309    /// Dense f32 view — only for owned tensors. Masked/sparse execution
310    /// paths require it; quantized weights don't support masks yet.
311    pub fn as_f32(&self) -> Option<&[f32]> {
312        match self {
313            Self::F32 { data, .. } => Some(data),
314            Self::Mapped { .. } => None,
315        }
316    }
317
318    fn quant_bytes(&self) -> &[u8] {
319        match self {
320            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
321            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
322        }
323    }
324
325    /// Dequantize one row into `dst` (embedding lookup).
326    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
327        let cols = self.cols();
328        debug_assert_eq!(dst.len(), cols);
329        match self {
330            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
331            Self::Mapped {
332                dtype,
333                row_scale,
334                col_field,
335                vbit_offsets,
336                ..
337            } => {
338                if *dtype == TensorDtype::Q4Tiled {
339                    let bytes = self.quant_bytes();
340                    let gpr = cols / GROUP_SIZE;
341                    for gi in 0..gpr {
342                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
343                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
344                        for (k, &b) in tile[2..].iter().enumerate() {
345                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
346                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
347                        }
348                    }
349                    return;
350                }
351                if *dtype == TensorDtype::Q4Block {
352                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
353                    let gpr = cols / GROUP_SIZE;
354                    for gi in 0..gpr {
355                        let g = r * gpr + gi;
356                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
357                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
358                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
359                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
360                        }
361                    }
362                    return;
363                }
364                if *dtype == TensorDtype::Q1 {
365                    let bytes = self.quant_bytes();
366                    let gpr = cols / GROUP_SIZE;
367                    for gi in 0..gpr {
368                        let tile =
369                            &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
370                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
371                        for (j, &b) in tile[2..].iter().enumerate() {
372                            for k in 0..8 {
373                                dst[gi * GROUP_SIZE + j * 8 + k] =
374                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
375                            }
376                        }
377                    }
378                    return;
379                }
380                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
381                    let bytes = self.quant_bytes();
382                    let rows = self.rows();
383                    let ng = cols / GROUP_SIZE;
384                    let bits = &bytes[..rows];
385                    let sc_off = rows;
386                    // Precomputed at load — embedding lookup used to scan
387                    // the bit-widths of every preceding row (O(token_id)).
388                    let off = vbit_offsets[r];
389                    let b = bits[r] as usize;
390                    let l = ((1usize << (b - 1)) - 1) as f32;
391                    let data = &bytes[off..];
392                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
393                    for (i, d) in dst.iter_mut().enumerate() {
394                        while nbits < b {
395                            acc = (acc << 8) | data[idx] as u64;
396                            idx += 1;
397                            nbits += 8;
398                        }
399                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
400                        nbits -= b;
401                        let so = (r * ng + i / GROUP_SIZE) * 2;
402                        let sv = f16_to_f32(u16::from_le_bytes([
403                            bytes[sc_off + so],
404                            bytes[sc_off + so + 1],
405                        ]));
406                        *d = (u - l) * sv;
407                    }
408                    return;
409                }
410                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
411                let s = row_scale[r];
412                match dtype {
413                    TensorDtype::Q8Row => {
414                        for (d, &b) in dst.iter_mut().zip(q) {
415                            *d = (b as i8) as f32 * s;
416                        }
417                    }
418                    TensorDtype::Q8_2f => {
419                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
420                            *d = (b as i8) as f32 * s * col_field[i];
421                        }
422                    }
423                    _ => unreachable!(),
424                }
425            }
426        }
427    }
428
429    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
430    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
431    /// false for group-packed q4/vbit (column access would unpack whole
432    /// groups — sparse execution falls back to f32 for those).
433    pub fn sparse_col_ok(&self) -> bool {
434        match self {
435            Self::F32 { .. } => true,
436            Self::Mapped { dtype, .. } => {
437                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
438            }
439        }
440    }
441
442    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
443    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
444    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
445    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
446        let inter = self.cols();
447        let hidden = self.rows();
448        debug_assert_eq!(out.len(), hidden);
449        match self {
450            Self::F32 { data, .. } => {
451                for (k, o) in out.iter_mut().enumerate() {
452                    *o += w * data[k * inter + c];
453                }
454            }
455            Self::Mapped {
456                dtype, row_scale, col_field, ..
457            } => {
458                let q = self.quant_bytes();
459                let colf = if *dtype == TensorDtype::Q8_2f {
460                    col_field[c]
461                } else {
462                    1.0
463                };
464                let wc = w * colf;
465                for (k, o) in out.iter_mut().enumerate() {
466                    let b = q[k * inter + c] as i8 as f32;
467                    *o += wc * b * row_scale[k];
468                }
469            }
470        }
471    }
472
473    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
474    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
475    /// into `scratch` first (rare for active-FFN weights).
476    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
477        let cols = self.cols();
478        match self {
479            Self::F32 { data, .. } => {
480                let row = &data[r * cols..(r + 1) * cols];
481                row.iter().zip(x).map(|(w, v)| w * v).sum()
482            }
483            Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
484                TensorDtype::Q8Row => {
485                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
486                    dot_i8_f32(q, x) * row_scale[r]
487                }
488                TensorDtype::Q8_2f => {
489                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
490                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
491                }
492                _ => {
493                    self.row_f32(r, scratch);
494                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
495                }
496            },
497        }
498    }
499
500    /// `out = W · x` (row-major). F32 delegates to the historical
501    /// bit-exact path; Mapped runs the fused int8 kernel.
502    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
503        match self {
504            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
505            Self::Mapped {
506                model,
507                idx,
508                dtype,
509                rows,
510                cols,
511                row_scale,
512                col_field,
513                vbit_offsets,
514                repack,
515            } => {
516                let _ = (model, idx);
517                if *dtype == TensorDtype::Q4Block {
518                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
519                    return;
520                }
521                if *dtype == TensorDtype::Q4Tiled {
522                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
523                    return;
524                }
525                if *dtype == TensorDtype::Q1 {
526                    // GPU route for large q1 matvecs (out_proj / lm_head
527                    // class): the CPU q1 kernel is load-port-bound at
528                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
529                    // probe measures both arms and keeps the winner.
530                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
531                        let t0 = std::time::Instant::now();
532                        let arm = if crate::gpu::q1_force() {
533                            crate::gpu::ProbeArm::Gpu
534                        } else {
535                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
536                        };
537                        match arm {
538                            crate::gpu::ProbeArm::Gpu => {
539                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
540                                    crate::gpu::probe_record(
541                                        crate::gpu::OpClass::Matvec,
542                                        true,
543                                        t0.elapsed(),
544                                    );
545                                    return;
546                                }
547                            }
548                            crate::gpu::ProbeArm::CpuTimed => {
549                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
550                                crate::gpu::probe_record(
551                                    crate::gpu::OpClass::Matvec,
552                                    false,
553                                    t0.elapsed(),
554                                );
555                                return;
556                            }
557                            crate::gpu::ProbeArm::Cpu => {}
558                        }
559                    }
560                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
561                    return;
562                }
563                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
564                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
565                    return;
566                }
567                let xs = prescale(x, col_field, *dtype);
568                // D5: large q8 matrices (lm_head-class) — hybrid
569                // CPU∥GPU: split the rows, both sides compute
570                // SIMULTANEOUSLY (same math, shared prescale).
571                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
572                if *rows >= crate::gpu::min_rows()
573                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
574                    && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
575                    && crate::gpu::enabled_here()
576                {
577                    // Runtime probe: alternate the hybrid against the
578                    // pure-CPU matvec, keep whichever is faster HERE.
579                    let t0 = std::time::Instant::now();
580                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
581                        crate::gpu::ProbeArm::Gpu => {}
582                        crate::gpu::ProbeArm::CpuTimed => {
583                            qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
584                            crate::gpu::probe_record(
585                                crate::gpu::OpClass::Matvec,
586                                false,
587                                t0.elapsed(),
588                            );
589                            return;
590                        }
591                        crate::gpu::ProbeArm::Cpu => {
592                            qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
593                            return;
594                        }
595                    }
596                    let frac = std::env::var("CMF_GPU_SPLIT")
597                        .ok()
598                        .and_then(|v| v.parse::<f32>().ok())
599                        .unwrap_or(0.5)
600                        .clamp(0.0, 1.0);
601                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
602                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
603                    let bytes = self.quant_bytes();
604                    let ok = std::thread::scope(|sc| {
605                        let g = sc.spawn(|| {
606                            crate::gpu::q8_matvec_range(
607                                model,
608                                *idx,
609                                cpu_rows,
610                                &row_scale[cpu_rows..],
611                                &xs,
612                                *rows - cpu_rows,
613                                *cols,
614                                out_gpu,
615                            )
616                        });
617                        if cpu_rows > 0 {
618                            // Repack prefix covers the full groups of the
619                            // CPU half (the split starts at row 0).
620                            let rep_cpu = if repack.is_empty() {
621                                &[][..]
622                            } else {
623                                &repack[..(cpu_rows / 4) * 4 * *cols]
624                            };
625                            qmatvec(
626                                &bytes[..cpu_rows * *cols],
627                                rep_cpu,
628                                &row_scale[..cpu_rows],
629                                &xs,
630                                cpu_rows,
631                                *cols,
632                                out_cpu,
633                                pool,
634                            );
635                        }
636                        g.join().unwrap_or(false)
637                    });
638                    if ok {
639                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
640                        return;
641                    }
642                    // GPU failed — CPU finishes its half (rows rebased —
643                    // group offsets don't line up, mmap layout only).
644                    qmatvec(
645                        &bytes[cpu_rows * *cols..(*rows) * *cols],
646                        &[],
647                        &row_scale[cpu_rows..],
648                        &xs,
649                        *rows - cpu_rows,
650                        *cols,
651                        out_gpu,
652                        pool,
653                    );
654                    return;
655                }
656                qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
657            }
658        }
659    }
660
661    /// Fused two-input matvec (MTP verify pair): weights streamed once.
662    pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
663        match self {
664            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
665            Self::Mapped {
666                dtype,
667                rows,
668                cols,
669                row_scale,
670                col_field,
671                vbit_offsets,
672                ..
673            } => {
674                if *dtype == TensorDtype::Q4Block {
675                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
676                    return;
677                }
678                if *dtype == TensorDtype::Q4Tiled {
679                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
680                    return;
681                }
682                if *dtype == TensorDtype::Q1 {
683                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
684                    return;
685                }
686                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
687                    vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
688                    return;
689                }
690                let x1s = prescale(x1, col_field, *dtype);
691                let x2s = prescale(x2, col_field, *dtype);
692                qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
693            }
694        }
695    }
696}
697
698impl QTensor {
699    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
700    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
701    /// to b matvec calls (same dot kernels in the same order); the win —
702    /// the weight row streams from DRAM once per batch, not b times.
703    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
704        let cols = self.cols();
705        let rows = self.rows();
706        debug_assert_eq!(xs_all.len(), b * cols);
707        debug_assert_eq!(out.len(), b * rows);
708        match self {
709            Self::F32 { data, .. } => {
710                let out_addr = SendMut(out.as_mut_ptr());
711                let run = |start: usize, end: usize| {
712                    for o in start..end {
713                        let row = &data[o * cols..(o + 1) * cols];
714                        for bi in 0..b {
715                            let x = &xs_all[bi * cols..(bi + 1) * cols];
716                            let mut acc = 0f32;
717                            for j in 0..cols {
718                                acc += row[j] * x[j];
719                            }
720                            unsafe { *out_addr.at(bi * rows + o) = acc };
721                        }
722                    }
723                };
724                dispatch_rows(pool, rows, &run);
725            }
726            Self::Mapped {
727                dtype,
728                row_scale,
729                col_field,
730                vbit_offsets,
731                ..
732            } => {
733                if *dtype == TensorDtype::Q4Block {
734                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
735                    return;
736                }
737                if *dtype == TensorDtype::Q4Tiled {
738                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
739                    return;
740                }
741                if *dtype == TensorDtype::Q1 {
742                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
743                    return;
744                }
745                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
746                    vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
747                    return;
748                }
749                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
750                    .map(|bi| {
751                        prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
752                    })
753                    .collect();
754                // D5: large prefill-batch GEMMs — on the GPU (threshold by
755                // work volume: submission carries b×rows×cols MACs).
756                // Runtime probe: the naive GEMM shader + sync readback
757                // lose to the CPU GEMM on slow driver stacks — alternate
758                // both arms and keep the winner.
759                if b >= 8
760                    && b * rows * cols >= 128_000_000
761                    && crate::gpu::enabled_here()
762                {
763                    if let Self::Mapped { model, idx, .. } = self {
764                        let t0 = std::time::Instant::now();
765                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
766                            crate::gpu::ProbeArm::Gpu
767                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
768                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
769                            {
770                                // Cold weights during probing: the upload
771                                // has started, the count runs on the CPU —
772                                // the GPU arm samples on the next touch.
773                                let q = self.quant_bytes();
774                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
775                                return;
776                            }
777                            crate::gpu::ProbeArm::Gpu => {
778                                let flat: Vec<f32> =
779                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
780                                if crate::gpu::q8_matmat(
781                                    model, *idx, row_scale, &flat, b, rows, cols, out)
782                                {
783                                    crate::gpu::probe_record(
784                                        crate::gpu::OpClass::Matmat,
785                                        true,
786                                        t0.elapsed(),
787                                    );
788                                    return;
789                                }
790                            }
791                            crate::gpu::ProbeArm::CpuTimed => {
792                                let q = self.quant_bytes();
793                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
794                                crate::gpu::probe_record(
795                                    crate::gpu::OpClass::Matmat,
796                                    false,
797                                    t0.elapsed(),
798                                );
799                                return;
800                            }
801                            crate::gpu::ProbeArm::Cpu => {}
802                        }
803                    }
804                }
805                let q = self.quant_bytes();
806                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
807            }
808        }
809    }
810}
811
812impl QTensor {
813    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
814    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
815    /// barrier instead of N. Per-row math is the exact same kernel as
816    /// `matvec` (bit-identical outputs); only the dispatch is fused.
817    /// Falls back to N sequential matvecs when the set is not a uniform
818    /// q8-family/F32 group or there is no pool.
819    pub fn matvec_many<const N: usize>(
820        ts: [&QTensor; N],
821        x: &[f32],
822        mut outs: [&mut [f32]; N],
823        pool: Option<&Pool>,
824    ) {
825        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
826        let uniform_q8 = ts.iter().all(|t| {
827            matches!(
828                t,
829                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
830            )
831        });
832        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
833        let uniform_q4 = ts
834            .iter()
835            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
836        let uniform_vbit = ts
837            .iter()
838            .all(|t| matches!(
839                t,
840                Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
841            ));
842        let uniform_q1 = ts
843            .iter()
844            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
845        let Some(pool) = pool else {
846            for (t, o) in ts.iter().zip(outs.iter_mut()) {
847                t.matvec(x, o, None);
848            }
849            return;
850        };
851        if total_rows < 256
852            || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
853        {
854            for (t, o) in ts.iter().zip(outs.iter_mut()) {
855                t.matvec(x, o, Some(pool));
856            }
857            return;
858        }
859
860        if uniform_q1 {
861            // One shared activation split + group sums (q1 has no col
862            // field; the same input feeds every tensor).
863            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
864            if a8w8_enabled() {
865                let act = split_act(x);
866                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
867                let (act, gsum) = (&act, &gsum);
868                let closures: [_; N] = std::array::from_fn(|i| {
869                    let (bytes, gpr, out) =
870                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
871                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
872                });
873                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
874                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
875                pool.run_many(&parts);
876            } else {
877                let closures: [_; N] = std::array::from_fn(|i| {
878                    let (bytes, gpr, out) =
879                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
880                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
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            return;
887        }
888
889        if uniform_q4 || uniform_vbit {
890            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
891            // q4/vbit share one activation split — no per-tensor col field.
892            if a8w8_enabled() {
893                let act = split_act(x);
894                let act = &act;
895                if uniform_q4 {
896                    let closures: [_; N] = std::array::from_fn(|i| {
897                        let (packed, scales) =
898                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
899                        let (gpr, cols, out) =
900                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
901                        move |s: usize, e: usize| {
902                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
903                        }
904                    });
905                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
906                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
907                    pool.run_many(&parts);
908                } else {
909                    let closures: [_; N] = std::array::from_fn(|i| {
910                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
911                        let (bytes, rows, cols, out) =
912                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
913                        move |s: usize, e: usize| {
914                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
915                        }
916                    });
917                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
918                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
919                    pool.run_many(&parts);
920                }
921                return;
922            }
923            if uniform_q4 {
924                let closures: [_; N] = std::array::from_fn(|i| {
925                    let (packed, scales) =
926                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
927                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
928                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
929                });
930                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
931                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
932                pool.run_many(&parts);
933            } else {
934                let closures: [_; N] = std::array::from_fn(|i| {
935                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
936                    let (bytes, rows, cols, out) =
937                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
938                    move |s: usize, e: usize| {
939                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
940                    }
941                });
942                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
943                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
944                pool.run_many(&parts);
945            }
946            return;
947        }
948
949        if uniform_f32 {
950            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
951            let closures: [_; N] = std::array::from_fn(|i| {
952                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
953                let out = outs_addr[i];
954                move |start: usize, end: usize| {
955                    for o in start..end {
956                        let row = &data[o * cols..(o + 1) * cols];
957                        let mut sum = 0.0f32;
958                        for j in 0..*cols {
959                            sum += row[j] * x[j];
960                        }
961                        // SAFETY: disjoint (tensor, row) cells per worker.
962                        unsafe { *out.at(o) = sum };
963                    }
964                }
965            });
966            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
967                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
968            pool.run_many(&parts);
969            return;
970        }
971
972        // Uniform q8-family: per-tensor prescale (q8_2f col fields
973        // differ per tensor) + the shared range kernels.
974        struct Ctx<'a> {
975            bytes: &'a [u8],
976            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
977            rep: &'a [u8],
978            row_scale: &'a [f32],
979            cols: usize,
980            xs: std::borrow::Cow<'a, [f32]>,
981        }
982        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
983            let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
984                unreachable!()
985            };
986            Ctx {
987                bytes: ts[i].quant_bytes(),
988                rep: repack,
989                row_scale,
990                cols: *cols,
991                xs: prescale(x, col_field, *dtype),
992            }
993        });
994        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
995        #[cfg(target_arch = "aarch64")]
996        if sdot_enabled() {
997            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
998            let closures: [_; N] = std::array::from_fn(|i| {
999                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1000                move |start: usize, end: usize| {
1001                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1002                }
1003            });
1004            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1005                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1006            pool.run_many(&parts);
1007            return;
1008        }
1009        #[cfg(target_arch = "x86_64")]
1010        if avx2_a8w8_enabled() {
1011            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1012            let closures: [_; N] = std::array::from_fn(|i| {
1013                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1014                move |start: usize, end: usize| {
1015                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1016                }
1017            });
1018            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1019                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1020            pool.run_many(&parts);
1021            return;
1022        }
1023        let closures: [_; N] = std::array::from_fn(|i| {
1024            let (c, out) = (&ctxs[i], outs_addr[i]);
1025            move |start: usize, end: usize| {
1026                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1027            }
1028        });
1029        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1030            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1031        pool.run_many(&parts);
1032    }
1033}
1034
1035impl QTensor {
1036    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
1037    /// single pool dispatch — the MTP/pair decode path publishes one job
1038    /// for Q/K/V (and one for gate+up) instead of one per tensor.
1039    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
1040    #[allow(clippy::needless_range_loop)]
1041    pub fn matvec2_many<const N: usize>(
1042        ts: [&QTensor; N],
1043        x1: &[f32],
1044        x2: &[f32],
1045        mut o1s: [&mut [f32]; N],
1046        mut o2s: [&mut [f32]; N],
1047        pool: Option<&Pool>,
1048    ) {
1049        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1050        let uniform_q8 = ts.iter().all(|t| {
1051            matches!(
1052                t,
1053                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1054            )
1055        });
1056        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1057        let uniform_q4 = ts
1058            .iter()
1059            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1060        let uniform_vbit = ts
1061            .iter()
1062            .all(|t| matches!(
1063                t,
1064                Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1065            ));
1066        let fusable = pool.is_some()
1067            && total_rows >= 256
1068            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1069        if !fusable {
1070            for i in 0..N {
1071                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1072            }
1073            return;
1074        }
1075        let pool = pool.unwrap();
1076
1077        if uniform_q4 || uniform_vbit {
1078            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1079            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1080            // q4/vbit share activation splits — no per-tensor col field.
1081            if a8w8_enabled() {
1082                let a1 = split_act(x1);
1083                let a2 = split_act(x2);
1084                let (a1, a2) = (&a1, &a2);
1085                if uniform_q4 {
1086                    let closures: [_; N] = std::array::from_fn(|i| {
1087                        let (packed, scales) =
1088                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1089                        let (gpr, cols, o1, o2) =
1090                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1091                        move |s: usize, e: usize| {
1092                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1093                        }
1094                    });
1095                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1096                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1097                    pool.run_many(&parts);
1098                } else {
1099                    let closures: [_; N] = std::array::from_fn(|i| {
1100                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1101                        let (bytes, rows, cols, o1, o2) =
1102                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1103                        move |s: usize, e: usize| {
1104                            vbit_range2_a8w8(
1105                                bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1106                            )
1107                        }
1108                    });
1109                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1110                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1111                    pool.run_many(&parts);
1112                }
1113                return;
1114            }
1115            if uniform_q4 {
1116                let closures: [_; N] = std::array::from_fn(|i| {
1117                    let (packed, scales) =
1118                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1119                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1120                    move |s: usize, e: usize| {
1121                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1122                    }
1123                });
1124                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1125                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1126                pool.run_many(&parts);
1127            } else {
1128                let closures: [_; N] = std::array::from_fn(|i| {
1129                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1130                    let (bytes, rows, cols, o1, o2) =
1131                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1132                    move |s: usize, e: usize| {
1133                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1134                    }
1135                });
1136                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1137                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1138                pool.run_many(&parts);
1139            }
1140            return;
1141        }
1142
1143        if uniform_f32 {
1144            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1145            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1146            let closures: [_; N] = std::array::from_fn(|i| {
1147                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1148                let (o1, o2) = (p1[i], p2[i]);
1149                move |start: usize, end: usize| {
1150                    for o in start..end {
1151                        let row = &data[o * cols..(o + 1) * cols];
1152                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
1153                        for j in 0..*cols {
1154                            s1 += row[j] * x1[j];
1155                            s2 += row[j] * x2[j];
1156                        }
1157                        // SAFETY: disjoint (tensor, row) cells per worker.
1158                        unsafe {
1159                            *o1.at(o) = s1;
1160                            *o2.at(o) = s2;
1161                        }
1162                    }
1163                }
1164            });
1165            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1166                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1167            pool.run_many(&parts);
1168            return;
1169        }
1170
1171        struct Ctx<'a> {
1172            bytes: &'a [u8],
1173            row_scale: &'a [f32],
1174            cols: usize,
1175            xs1: std::borrow::Cow<'a, [f32]>,
1176            xs2: std::borrow::Cow<'a, [f32]>,
1177        }
1178        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1179            let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1180                unreachable!()
1181            };
1182            Ctx {
1183                bytes: ts[i].quant_bytes(),
1184                row_scale,
1185                cols: *cols,
1186                xs1: prescale(x1, col_field, *dtype),
1187                xs2: prescale(x2, col_field, *dtype),
1188            }
1189        });
1190        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1191        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1192        #[cfg(target_arch = "aarch64")]
1193        if sdot_enabled() {
1194            let acts: [(SplitAct, SplitAct); N] =
1195                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1196            let closures: [_; N] = std::array::from_fn(|i| {
1197                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1198                move |start: usize, end: usize| {
1199                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1200                }
1201            });
1202            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1203                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1204            pool.run_many(&parts);
1205            return;
1206        }
1207        #[cfg(target_arch = "x86_64")]
1208        if avx2_a8w8_enabled() {
1209            let acts: [(SplitAct, SplitAct); N] =
1210                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1211            let closures: [_; N] = std::array::from_fn(|i| {
1212                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1213                move |start: usize, end: usize| {
1214                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1215                }
1216            });
1217            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1218                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1219            pool.run_many(&parts);
1220            return;
1221        }
1222        let closures: [_; N] = std::array::from_fn(|i| {
1223            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1224            move |start: usize, end: usize| {
1225                q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1226            }
1227        });
1228        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1229            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1230        pool.run_many(&parts);
1231    }
1232}
1233
1234/// Batched q8 kernel: same math as qmatvec, the row makes a single
1235/// pass from memory for the whole batch.
1236/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
1237/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
1238#[cfg(target_os = "macos")]
1239mod accel_blas {
1240    #[link(name = "Accelerate", kind = "framework")]
1241    unsafe extern "C" {
1242        pub fn cblas_sgemm(
1243            order: i32,
1244            trans_a: i32,
1245            trans_b: i32,
1246            m: i32,
1247            n: i32,
1248            k: i32,
1249            alpha: f32,
1250            a: *const f32,
1251            lda: i32,
1252            b: *const f32,
1253            ldb: i32,
1254            beta: f32,
1255            c: *mut f32,
1256            ldc: i32,
1257        );
1258    }
1259}
1260
1261#[cfg(target_os = "macos")]
1262pub(crate) fn accel_gemm_enabled() -> bool {
1263    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1264    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1265}
1266
1267/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
1268/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
1269#[cfg(target_os = "macos")]
1270#[allow(clippy::too_many_arguments)]
1271pub(crate) fn sgemm_rm(
1272    m: usize,
1273    n: usize,
1274    k: usize,
1275    alpha: f32,
1276    a: &[f32],
1277    lda: usize,
1278    b_mat: &[f32],
1279    ldb: usize,
1280    b_rows_are_n: bool,
1281    c: &mut [f32],
1282    ldc: usize,
1283) {
1284    debug_assert!(a.len() >= (m - 1) * lda + k);
1285    debug_assert!(c.len() >= (m - 1) * ldc + n);
1286    unsafe {
1287        accel_blas::cblas_sgemm(
1288            101, // RowMajor
1289            111, // NoTrans A
1290            if b_rows_are_n { 112 } else { 111 },
1291            m as i32,
1292            n as i32,
1293            k as i32,
1294            alpha,
1295            a.as_ptr(),
1296            lda as i32,
1297            b_mat.as_ptr(),
1298            ldb as i32,
1299            0.0,
1300            c.as_mut_ptr(),
1301            ldc as i32,
1302        );
1303    }
1304}
1305
1306/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
1307/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
1308/// on the AMX with one row-major sgemm. Tiles live in cache, weights
1309/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
1310/// logits shift within f32 rounding — tolerance-class, like every
1311/// reduction-order change; decode (M=1) never takes this path.
1312#[cfg(target_os = "macos")]
1313fn qmatmat_accel(
1314    q: &[u8],
1315    row_scale: &[f32],
1316    pre: &[std::borrow::Cow<'_, [f32]>],
1317    rows: usize,
1318    cols: usize,
1319    out: &mut [f32],
1320    pool: Option<&Pool>,
1321) {
1322    // NOTE: double-buffering the dequant against the sgemm (a scoped
1323    // thread driving the pool on tile k+1 while the caller multiplies
1324    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
1325    // multithreaded, and the dequant workers just steal its cores.
1326    const TR: usize = 2048;
1327    let b = pre.len();
1328    thread_local! {
1329        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1330        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1331    }
1332    XPANEL.with(|xp| {
1333        WTILE.with(|wt| {
1334            let mut xpanel = xp.borrow_mut();
1335            xpanel.clear();
1336            for x in pre {
1337                xpanel.extend_from_slice(x);
1338            }
1339            let mut wtile = wt.borrow_mut();
1340            wtile.resize(TR * cols, 0.0);
1341            let mut r0 = 0usize;
1342            while r0 < rows {
1343                let tr = TR.min(rows - r0);
1344                // Dequant the tile (scale folded) — pool-parallel.
1345                let wt_addr = SendMut(wtile.as_mut_ptr());
1346                let run = |start: usize, end: usize| {
1347                    for r in start..end {
1348                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1349                        let s = row_scale[r0 + r];
1350                        // SAFETY: workers cover disjoint r ranges.
1351                        let dst = unsafe {
1352                            std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1353                        };
1354                        for (d, &v) in dst.iter_mut().zip(row) {
1355                            *d = (v as i8) as f32 * s;
1356                        }
1357                    }
1358                };
1359                dispatch_rows(pool, tr, &run);
1360                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
1361                unsafe {
1362                    accel_blas::cblas_sgemm(
1363                        101, // RowMajor
1364                        111, // NoTrans A
1365                        112, // Trans B
1366                        b as i32,
1367                        tr as i32,
1368                        cols as i32,
1369                        1.0,
1370                        xpanel.as_ptr(),
1371                        cols as i32,
1372                        wtile.as_ptr(),
1373                        cols as i32,
1374                        0.0,
1375                        out.as_mut_ptr().add(r0),
1376                        rows as i32,
1377                    );
1378                }
1379                r0 += tr;
1380            }
1381        })
1382    });
1383}
1384
1385fn qmatmat(
1386    q: &[u8],
1387    row_scale: &[f32],
1388    pre: &[std::borrow::Cow<'_, [f32]>],
1389    rows: usize,
1390    cols: usize,
1391    out: &mut [f32],
1392    pool: Option<&Pool>,
1393) {
1394    let b = pre.len();
1395    debug_assert_eq!(out.len(), b * rows);
1396    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
1397    // SDOT loop below peaks near the CPU's dot throughput, an order
1398    // below the matrix units. Small tensors and tiny test models stay
1399    // on the exact integer path.
1400    #[cfg(target_os = "macos")]
1401    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1402        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1403        return;
1404    }
1405    #[cfg(target_arch = "aarch64")]
1406    if sdot_enabled() {
1407        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1408        let out_addr = SendMut(out.as_mut_ptr());
1409        let run = |start: usize, end: usize| {
1410            for o in start..end {
1411                let row = &q[o * cols..(o + 1) * cols];
1412                for (bi, act) in acts.iter().enumerate() {
1413                    let v = row_dot_sdot(row, act) * row_scale[o];
1414                    unsafe { *out_addr.at(bi * rows + o) = v };
1415                }
1416            }
1417        };
1418        dispatch_rows(pool, rows, &run);
1419        return;
1420    }
1421    // x86 A8W8 batch: each weight row streams once for the whole chunk
1422    // through the AVX2/VNNI row dot (prefill q8 was the last
1423    // aarch64-only batched path).
1424    #[cfg(target_arch = "x86_64")]
1425    if avx2_a8w8_enabled() {
1426        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1427        let out_addr = SendMut(out.as_mut_ptr());
1428        let run = |start: usize, end: usize| {
1429            for o in start..end {
1430                let row = &q[o * cols..(o + 1) * cols];
1431                for (bi, act) in acts.iter().enumerate() {
1432                    let v = row_dot_avx2(row, act) * row_scale[o];
1433                    unsafe { *out_addr.at(bi * rows + o) = v };
1434                }
1435            }
1436        };
1437        dispatch_rows(pool, rows, &run);
1438        return;
1439    }
1440    let out_addr = SendMut(out.as_mut_ptr());
1441    let run = |start: usize, end: usize| {
1442        for o in start..end {
1443            let row = &q[o * cols..(o + 1) * cols];
1444            for (bi, x) in pre.iter().enumerate() {
1445                let mut acc = 0f32;
1446                for j in 0..cols {
1447                    acc += (row[j] as i8) as f32 * x[j];
1448                }
1449                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1450            }
1451        }
1452    };
1453    dispatch_rows(pool, rows, &run);
1454}
1455
1456/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
1457/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
1458fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1459    match pool {
1460        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1461        _ => run(0, rows),
1462    }
1463}
1464
1465/// Split a q4_block blob into (packed nibbles, f16 group scales).
1466fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1467    let groups = rows * cols / GROUP_SIZE;
1468    bytes.split_at(groups * 16)
1469}
1470
1471/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
1472/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
1473/// vbit packs MSB-first, so the HIGH nibble is the even element
1474/// (opposite of q4_block's lo-first interleave). Centering is u-7.
1475#[inline]
1476fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1477    #[cfg(target_arch = "aarch64")]
1478    unsafe {
1479        return vbit_fill4_neon(data, buf);
1480    }
1481    #[cfg(target_arch = "x86_64")]
1482    if avx2_enabled() {
1483        return unsafe { vbit_fill4_avx2(data, buf) };
1484    }
1485    #[allow(unreachable_code)]
1486    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1487        let u = unpack8::<4>(&data[blk * 4..]);
1488        for k in 0..8 {
1489            chunk[k] = (u[k] - 7) as i8 as u8;
1490        }
1491    }
1492}
1493
1494#[cfg(target_arch = "aarch64")]
1495#[target_feature(enable = "neon")]
1496unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1497    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
1498    // buf.len()/2 packed bytes (validated at load).
1499    unsafe {
1500        use core::arch::aarch64::*;
1501        let n = buf.len();
1502        let mask = vdupq_n_u8(0x0F);
1503        let seven = vdupq_n_s8(7);
1504        let mut g = 0usize;
1505        while g * 32 + 32 <= n {
1506            let b = vld1q_u8(data.as_ptr().add(g * 16));
1507            let hi = vshrq_n_u8::<4>(b);
1508            let lo = vandq_u8(b, mask);
1509            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1510            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1511            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1512            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1513            g += 1;
1514        }
1515    }
1516}
1517
1518#[cfg(target_arch = "x86_64")]
1519#[target_feature(enable = "avx2")]
1520unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1521    // SAFETY: see vbit_fill4_neon.
1522    unsafe {
1523        use core::arch::x86_64::*;
1524        let n = buf.len();
1525        let mask = _mm_set1_epi8(0x0F);
1526        let seven = _mm256_set1_epi8(7);
1527        let mut g = 0usize;
1528        while g * 32 + 32 <= n {
1529            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1530            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1531            let lo = _mm_and_si128(b, mask);
1532            let z = _mm256_sub_epi8(
1533                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1534                seven,
1535            );
1536            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1537            g += 1;
1538        }
1539    }
1540}
1541
1542/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
1543/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
1544/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
1545/// into 4 such blocks.
1546#[inline(always)]
1547fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1548    let mut acc = 0u64;
1549    for i in 0..B {
1550        acc = (acc << 8) | data[i] as u64;
1551    }
1552    let mask = (1u64 << B) - 1;
1553    let mut out = [0i32; 8];
1554    for (k, o) in out.iter_mut().enumerate() {
1555        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1556    }
1557    out
1558}
1559
1560/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
1561/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
1562/// MSB-first, byte-padded]. Row data offsets are precomputed at load
1563/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
1564/// overhead on every matvec.
1565#[allow(clippy::too_many_arguments)]
1566fn vbitmatvec(
1567    bytes: &[u8],
1568    offsets: &[usize],
1569    x: &[f32],
1570    rows: usize,
1571    cols: usize,
1572    out: &mut [f32],
1573    pool: Option<&Pool>,
1574) {
1575    debug_assert_eq!(out.len(), rows);
1576    debug_assert_eq!(offsets.len(), rows + 1);
1577
1578    // SDOT path: unpack the row to centered i8 once, then per-group
1579    // int8 dot against the quantized activations — same A8W8 contract
1580    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
1581    if a8w8_enabled() {
1582        let act = split_act(x);
1583        let out_addr = SendMut(out.as_mut_ptr());
1584        let run = move |start: usize, end: usize| {
1585            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1586        };
1587        dispatch_rows(pool, rows, &run);
1588        return;
1589    }
1590
1591    let out_addr = SendMut(out.as_mut_ptr());
1592    let run = move |start: usize, end: usize| {
1593        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1594    };
1595    dispatch_rows(pool, rows, &run);
1596}
1597
1598/// One vbit row range via the A8W8 int8 path — kernel body of
1599/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
1600/// several tensors in one dispatch (b=8 rows go exact f32).
1601#[allow(clippy::too_many_arguments)]
1602fn vbit_range_a8w8(
1603    bytes: &[u8],
1604    offsets: &[usize],
1605    x: &[f32],
1606    act: &SplitAct,
1607    rows: usize,
1608    cols: usize,
1609    out: SendMut,
1610    start: usize,
1611    end: usize,
1612) {
1613    let ng = cols / GROUP_SIZE;
1614    let bits = &bytes[..rows];
1615    let sc_off = rows;
1616    let row_dot = |r: usize| -> f32 {
1617            let b = bits[r] as usize;
1618            let l = ((1i32 << (b - 1)) - 1) as i32;
1619            let mask = (1u64 << b) - 1;
1620            let data = &bytes[offsets[r]..offsets[r + 1]];
1621            if b == 8 {
1622                // u−L reaches 128 → does not fit i8; exact f32 path.
1623                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1624                let mut dot = 0f32;
1625                for g in 0..ng {
1626                    let so = (r * ng + g) * 2;
1627                    let sgf = f16_to_f32(u16::from_le_bytes([
1628                        bytes[sc_off + so],
1629                        bytes[sc_off + so + 1],
1630                    ]));
1631                    let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1632                    let mut gd = 0f32;
1633                    for &xv in xg.iter() {
1634                        if nbits < 8 {
1635                            acc = (acc << 8) | data[idx] as u64;
1636                            idx += 1;
1637                            nbits += 8;
1638                        }
1639                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1640                        nbits -= 8;
1641                        gd += (u - l) as f32 * xv;
1642                    }
1643                    dot += gd * sgf;
1644                }
1645                return dot;
1646            }
1647            // Per-worker scratch: this closure runs for every row of the
1648            // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
1649            // row was measurable pure overhead.
1650            thread_local! {
1651                static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1652                    const { std::cell::RefCell::new(Vec::new()) };
1653            }
1654            #[inline(always)]
1655            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1656                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1657                    let u = unpack8::<B>(&data[blk * B..]);
1658                    for k in 0..8 {
1659                        chunk[k] = (u[k] - l) as i8 as u8;
1660                    }
1661                }
1662            }
1663            let _ = mask;
1664            VBIT_SCRATCH.with(|scratch| {
1665                let mut buf = scratch.borrow_mut();
1666                buf.resize(cols, 0);
1667                match b {
1668                    3 => fill::<3>(data, l, &mut buf),
1669                    4 => vbit_fill4(data, &mut buf),
1670                    5 => fill::<5>(data, l, &mut buf),
1671                    6 => fill::<6>(data, l, &mut buf),
1672                    _ => unreachable!(),
1673                }
1674                let mut dot = 0f32;
1675                for g in 0..ng {
1676                    let so = (r * ng + g) * 2;
1677                    let s = f16_to_f32(u16::from_le_bytes([
1678                        bytes[sc_off + so],
1679                        bytes[sc_off + so + 1],
1680                    ]));
1681                    let d = dot_i8_i8(
1682                        &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1683                        &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1684                    ) as f32
1685                        * act.sx;
1686                    dot += d * s;
1687                }
1688                for &(j, xv) in &act.outliers {
1689                    let so = (r * ng + j / GROUP_SIZE) * 2;
1690                    let s = f16_to_f32(u16::from_le_bytes([
1691                        bytes[sc_off + so],
1692                        bytes[sc_off + so + 1],
1693                    ]));
1694                    // xq is zeroed at outlier slots — add the exact term.
1695                    dot += (buf[j] as i8) as f32 * s * xv;
1696                }
1697                dot
1698            })
1699    };
1700    for r in start..end {
1701        // SAFETY: disjoint row ranges per worker.
1702        unsafe { *out.at(r) = row_dot(r) };
1703    }
1704}
1705
1706/// Exact scalar vbit row range (same extraction, non-SDOT path).
1707#[allow(clippy::too_many_arguments)]
1708fn vbit_range_f32(
1709    bytes: &[u8],
1710    offsets: &[usize],
1711    x: &[f32],
1712    rows: usize,
1713    cols: usize,
1714    out: SendMut,
1715    start: usize,
1716    end: usize,
1717) {
1718    let ng = cols / GROUP_SIZE;
1719    let bits = &bytes[..rows];
1720    let sc_off = rows;
1721    // Per-bit-width specialized inner loops: the compiler unrolls the
1722    // constant shifts (the generic bit-buffer loop was branch-bound —
1723    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
1724    #[inline(always)]
1725    fn dot_row<const B: usize>(
1726        data: &[u8],
1727        bytes: &[u8],
1728        sc_off: usize,
1729        r: usize,
1730        ng: usize,
1731        x: &[f32],
1732    ) -> f32 {
1733        let l = ((1i32 << (B - 1)) - 1) as f32;
1734        let gbytes = GROUP_SIZE * B / 8;
1735        let mut dot = 0f32;
1736        for g in 0..ng {
1737            let so = (r * ng + g) * 2;
1738            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1739            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1740            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1741            let mut gd = 0f32;
1742            for blk in 0..GROUP_SIZE / 8 {
1743                let u = unpack8::<B>(&gd0[blk * B..]);
1744                let xb = &xg[blk * 8..blk * 8 + 8];
1745                for k in 0..8 {
1746                    gd += (u[k] as f32 - l) * xb[k];
1747                }
1748            }
1749            dot += gd * s;
1750        }
1751        dot
1752    }
1753    for r in start..end {
1754        let data = &bytes[offsets[r]..offsets[r + 1]];
1755        let v = match bits[r] {
1756            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1757            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1758            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1759            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1760            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1761            b => unreachable!("vbit bit-width {b} (validated at load)"),
1762        };
1763        // SAFETY: disjoint row ranges per worker.
1764        unsafe { *out.at(r) = v };
1765    }
1766}
1767
1768/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
1769/// and dotted against BOTH activations (MTP verify / pair prefill used
1770/// to run two full matvecs — double weight traffic and double unpack).
1771/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
1772#[allow(clippy::too_many_arguments)]
1773fn vbitmatvec2(
1774    bytes: &[u8],
1775    offsets: &[usize],
1776    x1: &[f32],
1777    x2: &[f32],
1778    rows: usize,
1779    cols: usize,
1780    o1: &mut [f32],
1781    o2: &mut [f32],
1782    pool: Option<&Pool>,
1783) {
1784    debug_assert_eq!(o1.len(), rows);
1785    debug_assert_eq!(o2.len(), rows);
1786
1787    if a8w8_enabled() {
1788        let a1 = split_act(x1);
1789        let a2 = split_act(x2);
1790        let p1 = SendMut(o1.as_mut_ptr());
1791        let p2 = SendMut(o2.as_mut_ptr());
1792        let run = move |start: usize, end: usize| {
1793            vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1794        };
1795        dispatch_rows(pool, rows, &run);
1796        return;
1797    }
1798
1799    let p1 = SendMut(o1.as_mut_ptr());
1800    let p2 = SendMut(o2.as_mut_ptr());
1801    let run = move |start: usize, end: usize| {
1802        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1803    };
1804    dispatch_rows(pool, rows, &run);
1805}
1806
1807/// Two-input vbit row range via the A8W8 int8 path — kernel body of
1808/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
1809/// exact f32 for both lanes, bits streamed once).
1810#[allow(clippy::too_many_arguments)]
1811fn vbit_range2_a8w8(
1812    bytes: &[u8],
1813    offsets: &[usize],
1814    x1: &[f32],
1815    x2: &[f32],
1816    a1: &SplitAct,
1817    a2: &SplitAct,
1818    rows: usize,
1819    cols: usize,
1820    p1: SendMut,
1821    p2: SendMut,
1822    start: usize,
1823    end: usize,
1824) {
1825    let ng = cols / GROUP_SIZE;
1826    let bits = &bytes[..rows];
1827    let sc_off = rows;
1828    let row_dots = |r: usize| -> (f32, f32) {
1829            let b = bits[r] as usize;
1830            let l = (1i32 << (b - 1)) - 1;
1831            let data = &bytes[offsets[r]..offsets[r + 1]];
1832            if b == 8 {
1833                // u−L reaches 128 → does not fit i8; exact f32 path,
1834                // bits still streamed once for both lanes.
1835                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1836                let (mut d1, mut d2) = (0f32, 0f32);
1837                for g in 0..ng {
1838                    let so = (r * ng + g) * 2;
1839                    let sgf = f16_to_f32(u16::from_le_bytes([
1840                        bytes[sc_off + so],
1841                        bytes[sc_off + so + 1],
1842                    ]));
1843                    let (mut g1, mut g2) = (0f32, 0f32);
1844                    for k in 0..GROUP_SIZE {
1845                        if nbits < 8 {
1846                            acc = (acc << 8) | data[idx] as u64;
1847                            idx += 1;
1848                            nbits += 8;
1849                        }
1850                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1851                        nbits -= 8;
1852                        let w = (u - l) as f32;
1853                        g1 += w * x1[g * GROUP_SIZE + k];
1854                        g2 += w * x2[g * GROUP_SIZE + k];
1855                    }
1856                    d1 += g1 * sgf;
1857                    d2 += g2 * sgf;
1858                }
1859                return (d1, d2);
1860            }
1861            thread_local! {
1862                static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1863                    const { std::cell::RefCell::new(Vec::new()) };
1864            }
1865            #[inline(always)]
1866            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1867                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1868                    let u = unpack8::<B>(&data[blk * B..]);
1869                    for k in 0..8 {
1870                        chunk[k] = (u[k] - l) as i8 as u8;
1871                    }
1872                }
1873            }
1874            VBIT_SCRATCH2.with(|scratch| {
1875                let mut buf = scratch.borrow_mut();
1876                buf.resize(cols, 0);
1877                match b {
1878                    3 => fill::<3>(data, l, &mut buf),
1879                    4 => vbit_fill4(data, &mut buf),
1880                    5 => fill::<5>(data, l, &mut buf),
1881                    6 => fill::<6>(data, l, &mut buf),
1882                    _ => unreachable!(),
1883                }
1884                let (mut d1, mut d2) = (0f32, 0f32);
1885                for g in 0..ng {
1886                    let so = (r * ng + g) * 2;
1887                    let s = f16_to_f32(u16::from_le_bytes([
1888                        bytes[sc_off + so],
1889                        bytes[sc_off + so + 1],
1890                    ]));
1891                    let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1892                    let v1 =
1893                        dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1894                    let v2 =
1895                        dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1896                    d1 += v1 * s;
1897                    d2 += v2 * s;
1898                }
1899                for &(j, xv) in &a1.outliers {
1900                    let so = (r * ng + j / GROUP_SIZE) * 2;
1901                    let s = f16_to_f32(u16::from_le_bytes([
1902                        bytes[sc_off + so],
1903                        bytes[sc_off + so + 1],
1904                    ]));
1905                    d1 += (buf[j] as i8) as f32 * s * xv;
1906                }
1907                for &(j, xv) in &a2.outliers {
1908                    let so = (r * ng + j / GROUP_SIZE) * 2;
1909                    let s = f16_to_f32(u16::from_le_bytes([
1910                        bytes[sc_off + so],
1911                        bytes[sc_off + so + 1],
1912                    ]));
1913                    d2 += (buf[j] as i8) as f32 * s * xv;
1914                }
1915                (d1, d2)
1916            })
1917    };
1918    for r in start..end {
1919        let (v1, v2) = row_dots(r);
1920        // SAFETY: disjoint row ranges per worker.
1921        unsafe {
1922            *p1.at(r) = v1;
1923            *p2.at(r) = v2;
1924        }
1925    }
1926}
1927
1928/// Two-input exact scalar vbit row range (same extraction) —
1929/// per-bit-width specialized, two accumulators per row; per-lane
1930/// accumulation order matches `vbitmatvec` exactly.
1931#[allow(clippy::too_many_arguments)]
1932fn vbit_range2_f32(
1933    bytes: &[u8],
1934    offsets: &[usize],
1935    x1: &[f32],
1936    x2: &[f32],
1937    rows: usize,
1938    cols: usize,
1939    p1: SendMut,
1940    p2: SendMut,
1941    start: usize,
1942    end: usize,
1943) {
1944    let ng = cols / GROUP_SIZE;
1945    let bits = &bytes[..rows];
1946    let sc_off = rows;
1947    #[inline(always)]
1948    #[allow(clippy::too_many_arguments)]
1949    fn dot_row2<const B: usize>(
1950        data: &[u8],
1951        bytes: &[u8],
1952        sc_off: usize,
1953        r: usize,
1954        ng: usize,
1955        x1: &[f32],
1956        x2: &[f32],
1957    ) -> (f32, f32) {
1958        let l = ((1i32 << (B - 1)) - 1) as f32;
1959        let gbytes = GROUP_SIZE * B / 8;
1960        let (mut d1, mut d2) = (0f32, 0f32);
1961        for g in 0..ng {
1962            let so = (r * ng + g) * 2;
1963            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1964            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1965            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1966            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1967            let (mut g1, mut g2) = (0f32, 0f32);
1968            for blk in 0..GROUP_SIZE / 8 {
1969                let u = unpack8::<B>(&gd0[blk * B..]);
1970                for k in 0..8 {
1971                    let w = u[k] as f32 - l;
1972                    g1 += w * x1g[blk * 8 + k];
1973                    g2 += w * x2g[blk * 8 + k];
1974                }
1975            }
1976            d1 += g1 * s;
1977            d2 += g2 * s;
1978        }
1979        (d1, d2)
1980    }
1981    for r in start..end {
1982        let data = &bytes[offsets[r]..offsets[r + 1]];
1983        let (v1, v2) = match bits[r] {
1984            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
1985            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
1986            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
1987            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
1988            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
1989            b => unreachable!("vbit bit-width {b} (validated at load)"),
1990        };
1991        // SAFETY: disjoint row ranges per worker.
1992        unsafe {
1993            *p1.at(r) = v1;
1994            *p2.at(r) = v2;
1995        }
1996    }
1997}
1998
1999// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
2000
2001/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
2002/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
2003/// distant streams of the split layout. Values/order identical to the
2004/// split kernels.
2005#[inline]
2006#[allow(unreachable_code)]
2007fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2008    #[cfg(target_arch = "aarch64")]
2009    unsafe {
2010        return dot_q4t_row_sdot(bytes, r, gpr, xq);
2011    }
2012    #[cfg(target_arch = "x86_64")]
2013    unsafe {
2014        return dot_q4t_row_avx2(bytes, r, gpr, xq);
2015    }
2016    let mut acc = 0f32;
2017    for gi in 0..gpr {
2018        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2019        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2020        let mut d = 0i32;
2021        for (k, &b) in tile[2..].iter().enumerate() {
2022            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2023                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2024        }
2025        acc += d as f32 * s;
2026    }
2027    acc
2028}
2029
2030#[cfg(target_arch = "aarch64")]
2031#[target_feature(enable = "neon,dotprod")]
2032unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2033    // SAFETY: callers uphold slice-length contracts (18B tile per group,
2034    // xq.len() == gpr·GROUP_SIZE).
2035    unsafe {
2036        use core::arch::aarch64::*;
2037        use core::arch::asm;
2038        let lomask = vdupq_n_u8(0x0F);
2039        let eight = vdupq_n_s8(8);
2040        let mut acc = 0f32;
2041        for gi in 0..gpr {
2042            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2043            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2044            let b = vld1q_u8(t.add(2));
2045            let lo = vandq_u8(b, lomask);
2046            let hi = vshrq_n_u8::<4>(b);
2047            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2048            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2049            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2050            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2051            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2052            asm!(
2053                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2054                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2055                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2056                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2057                options(pure, nomem, nostack),
2058            );
2059            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2060        }
2061        acc
2062    }
2063}
2064
2065#[cfg(target_arch = "x86_64")]
2066#[target_feature(enable = "avx2")]
2067unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2068    // SAFETY: see dot_q4t_row_sdot.
2069    unsafe {
2070        use core::arch::x86_64::*;
2071        let lomask = _mm_set1_epi8(0x0F);
2072        let eight = _mm256_set1_epi8(8);
2073        let ones = _mm256_set1_epi16(1);
2074        let mut acc = 0f32;
2075        for gi in 0..gpr {
2076            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2077            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2078            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2079            let lo = _mm_and_si128(b, lomask);
2080            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2081            let w = _mm256_sub_epi8(
2082                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2083                eight,
2084            );
2085            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2086            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2087            let d = _mm256_madd_epi16(p16, ones);
2088            let hi128 = _mm256_extracti128_si256::<1>(d);
2089            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2090            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2091            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2092            acc += _mm_cvtsi128_si32(s32) as f32 * s;
2093        }
2094        acc
2095    }
2096}
2097
2098/// Exact-term correction for A8W8 outliers on a tiled row.
2099#[inline]
2100fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2101    let gi = j / GROUP_SIZE;
2102    let k = j % GROUP_SIZE;
2103    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2104    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2105    let byte = tile[2 + k / 2];
2106    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2107    ((nib as i32 - 8) as f32, s)
2108}
2109
2110/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
2111/// accumulation shape as `q4_range_f32`.
2112#[inline]
2113fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2114    let mut acc = 0f32;
2115    for gi in 0..gpr {
2116        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2117        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2118        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2119        let mut ga = 0f32;
2120        for (k, &b) in tile[2..].iter().enumerate() {
2121            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2122                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2123        }
2124        acc += ga * s;
2125    }
2126    acc
2127}
2128
2129/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
2130fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2131    debug_assert_eq!(out.len(), rows);
2132    let gpr = cols / GROUP_SIZE;
2133    let out_addr = SendMut(out.as_mut_ptr());
2134    if a8w8_enabled() {
2135        let act = split_act(x);
2136        let run = move |start: usize, end: usize| {
2137            for r in start..end {
2138                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2139                for &(j, xv) in &act.outliers {
2140                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2141                    acc += w * s * xv;
2142                }
2143                // SAFETY: disjoint row ranges per worker.
2144                unsafe { *out_addr.at(r) = acc };
2145            }
2146        };
2147        dispatch_rows(pool, rows, &run);
2148        return;
2149    }
2150    let run = move |start: usize, end: usize| {
2151        for r in start..end {
2152            // SAFETY: disjoint row ranges per worker.
2153            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2154        }
2155    };
2156    dispatch_rows(pool, rows, &run);
2157}
2158
2159/// Fused two-input q4_tiled matvec (weights read once per pair).
2160#[allow(clippy::too_many_arguments)]
2161fn q4t_matvec2(
2162    bytes: &[u8],
2163    x1: &[f32],
2164    x2: &[f32],
2165    rows: usize,
2166    cols: usize,
2167    o1: &mut [f32],
2168    o2: &mut [f32],
2169    pool: Option<&Pool>,
2170) {
2171    let gpr = cols / GROUP_SIZE;
2172    let p1 = SendMut(o1.as_mut_ptr());
2173    let p2 = SendMut(o2.as_mut_ptr());
2174    if a8w8_enabled() {
2175        let a1 = split_act(x1);
2176        let a2 = split_act(x2);
2177        let run = move |start: usize, end: usize| {
2178            for r in start..end {
2179                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2180                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2181                for &(j, xv) in &a1.outliers {
2182                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2183                    v1 += w * s * xv;
2184                }
2185                for &(j, xv) in &a2.outliers {
2186                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2187                    v2 += w * s * xv;
2188                }
2189                // SAFETY: disjoint row ranges per worker.
2190                unsafe {
2191                    *p1.at(r) = v1;
2192                    *p2.at(r) = v2;
2193                }
2194            }
2195        };
2196        dispatch_rows(pool, rows, &run);
2197        return;
2198    }
2199    let run = move |start: usize, end: usize| {
2200        for r in start..end {
2201            // SAFETY: disjoint row ranges per worker.
2202            unsafe {
2203                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2204                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2205            }
2206        }
2207    };
2208    dispatch_rows(pool, rows, &run);
2209}
2210
2211/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
2212#[allow(clippy::too_many_arguments)]
2213fn q4t_matmat(
2214    bytes: &[u8],
2215    xs_all: &[f32],
2216    b: usize,
2217    rows: usize,
2218    cols: usize,
2219    out: &mut [f32],
2220    pool: Option<&Pool>,
2221) {
2222    debug_assert_eq!(out.len(), b * rows);
2223    let gpr = cols / GROUP_SIZE;
2224    let out_addr = SendMut(out.as_mut_ptr());
2225    if a8w8_enabled() {
2226        let acts: Vec<SplitAct> =
2227            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2228        let acts = &acts;
2229        let run = move |start: usize, end: usize| {
2230            for r in start..end {
2231                for (bi, act) in acts.iter().enumerate() {
2232                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2233                    for &(j, xv) in &act.outliers {
2234                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
2235                        acc += w * s * xv;
2236                    }
2237                    // SAFETY: disjoint (bi, r) cells per worker range.
2238                    unsafe { *out_addr.at(bi * rows + r) = acc };
2239                }
2240            }
2241        };
2242        dispatch_rows(pool, rows, &run);
2243        return;
2244    }
2245    let run = move |start: usize, end: usize| {
2246        for r in start..end {
2247            for bi in 0..b {
2248                let x = &xs_all[bi * cols..(bi + 1) * cols];
2249                // SAFETY: disjoint (bi, r) cells per worker range.
2250                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2251            }
2252        }
2253    };
2254    dispatch_rows(pool, rows, &run);
2255}
2256
2257// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
2258// 32-group tile. The kernel family mirrors q4_tiled: one sequential
2259// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
2260// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
2261
2262/// Per-32-group sums of the quantized activation — the ±1 identity's
2263/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
2264/// matvec and reused by every row.
2265fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2266    (0..gpr)
2267        .map(|gi| {
2268            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2269                .iter()
2270                .map(|&v| v as i32)
2271                .sum()
2272        })
2273        .collect()
2274}
2275
2276/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
2277/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
2278/// x86 pass).
2279#[inline]
2280#[allow(unreachable_code)]
2281fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2282    #[cfg(target_arch = "aarch64")]
2283    unsafe {
2284        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2285    }
2286    let _ = gsum;
2287    let mut acc = 0f32;
2288    for gi in 0..gpr {
2289        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2290        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2291        let mut d = 0i32;
2292        for (j, &b) in tile[2..].iter().enumerate() {
2293            for k in 0..8 {
2294                let w = ((b >> k) & 1) as i32 * 2 - 1;
2295                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2296            }
2297        }
2298        acc += d as f32 * s;
2299    }
2300    acc
2301}
2302
2303/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
2304/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
2305/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
2306/// per-group activation sums shared across every row of the matvec.
2307/// Four tiles (128 weights) per iteration: integer dots reduce through
2308/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
2309/// fused f32 multiply-add. Integer math throughout — bit-identical to
2310/// the scalar ±1 reference.
2311#[cfg(target_arch = "aarch64")]
2312#[target_feature(enable = "neon,dotprod")]
2313unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2314    // SAFETY: callers uphold slice-length contracts (6B tile per group,
2315    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
2316    unsafe {
2317        use core::arch::aarch64::*;
2318        use core::arch::asm;
2319        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2320        let m = vld1q_u8(MASKS.as_ptr());
2321        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
2322        macro_rules! tile_dot {
2323            ($t:expr, $x:expr) => {{
2324                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2325                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2326                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2327                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2328                let x0 = vld1q_s8($x);
2329                let x1 = vld1q_s8($x.add(16));
2330                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2331                asm!(
2332                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2333                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2334                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2335                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2336                    options(pure, nomem, nostack),
2337                );
2338                vaddq_s32(a0, a1)
2339            }};
2340        }
2341        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2342        let xp = xq.as_ptr();
2343        let gp = gsum.as_ptr();
2344        let mut accv = vdupq_n_f32(0.0);
2345        let mut gi = 0;
2346        while gi + 4 <= gpr {
2347            let t0 = base.add(gi * Q1_TILE);
2348            let t1 = base.add((gi + 1) * Q1_TILE);
2349            let t2 = base.add((gi + 2) * Q1_TILE);
2350            let t3 = base.add((gi + 3) * Q1_TILE);
2351            let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2352            let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2353            let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2354            let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2355            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
2356            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2357            let g = vld1q_s32(gp.add(gi));
2358            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2359            let sc = [
2360                f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2361                f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2362                f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2363                f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2364            ];
2365            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2366            gi += 4;
2367        }
2368        let mut acc = vaddvq_f32(accv);
2369        while gi < gpr {
2370            let t = base.add(gi * Q1_TILE);
2371            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2372            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2373            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2374            gi += 1;
2375        }
2376        acc
2377    }
2378}
2379
2380/// (weight ±1, scale) of one q1 element — the exact outlier term.
2381#[inline]
2382fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2383    let gi = j / GROUP_SIZE;
2384    let k = j % GROUP_SIZE;
2385    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2386    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2387    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2388    ((bit as i32 * 2 - 1) as f32, s)
2389}
2390
2391/// Exact scalar q1 row (CMF_SDOT=0 contract).
2392#[inline]
2393fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2394    let mut acc = 0f32;
2395    for gi in 0..gpr {
2396        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2397        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2398        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2399        let mut ga = 0f32;
2400        for (j, &b) in tile[2..].iter().enumerate() {
2401            for k in 0..8 {
2402                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2403            }
2404        }
2405        acc += ga * s;
2406    }
2407    acc
2408}
2409
2410/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
2411/// extracted so multi-matrix jobs drive the same kernel).
2412#[allow(clippy::too_many_arguments)]
2413fn q1_range_a8w8(
2414    bytes: &[u8],
2415    gpr: usize,
2416    act: &SplitAct,
2417    gsum: &[i32],
2418    out: SendMut,
2419    start: usize,
2420    end: usize,
2421) {
2422    for r in start..end {
2423        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2424        for &(j, xv) in &act.outliers {
2425            let (w, s) = q1_outlier(bytes, r, gpr, j);
2426            acc += w * s * xv;
2427        }
2428        // SAFETY: disjoint row ranges per worker.
2429        unsafe { *out.at(r) = acc };
2430    }
2431}
2432
2433/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
2434fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2435    for r in start..end {
2436        // SAFETY: disjoint row ranges per worker.
2437        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2438    }
2439}
2440
2441/// Fused q1 matvec (dispatch mirrors `q4t_matvec`).
2442fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2443    debug_assert_eq!(out.len(), rows);
2444    let gpr = cols / GROUP_SIZE;
2445    let out_addr = SendMut(out.as_mut_ptr());
2446    if a8w8_enabled() {
2447        let act = split_act(x);
2448        let gsum = q1_group_sums(&act.xq, gpr);
2449        let (act, gsum) = (&act, &gsum);
2450        let run =
2451            move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2452        dispatch_rows(pool, rows, &run);
2453        return;
2454    }
2455    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2456    dispatch_rows(pool, rows, &run);
2457}
2458
2459/// Fused two-input q1 matvec (weights read once per pair).
2460#[allow(clippy::too_many_arguments)]
2461fn q1_matvec2(
2462    bytes: &[u8],
2463    x1: &[f32],
2464    x2: &[f32],
2465    rows: usize,
2466    cols: usize,
2467    o1: &mut [f32],
2468    o2: &mut [f32],
2469    pool: Option<&Pool>,
2470) {
2471    let gpr = cols / GROUP_SIZE;
2472    let p1 = SendMut(o1.as_mut_ptr());
2473    let p2 = SendMut(o2.as_mut_ptr());
2474    if a8w8_enabled() {
2475        let a1 = split_act(x1);
2476        let a2 = split_act(x2);
2477        let g1 = q1_group_sums(&a1.xq, gpr);
2478        let g2 = q1_group_sums(&a2.xq, gpr);
2479        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2480        let run = move |start: usize, end: usize| {
2481            for r in start..end {
2482                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2483                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2484                for &(j, xv) in &a1.outliers {
2485                    let (w, s) = q1_outlier(bytes, r, gpr, j);
2486                    v1 += w * s * xv;
2487                }
2488                for &(j, xv) in &a2.outliers {
2489                    let (w, s) = q1_outlier(bytes, r, gpr, j);
2490                    v2 += w * s * xv;
2491                }
2492                // SAFETY: disjoint row ranges per worker.
2493                unsafe {
2494                    *p1.at(r) = v1;
2495                    *p2.at(r) = v2;
2496                }
2497            }
2498        };
2499        dispatch_rows(pool, rows, &run);
2500        return;
2501    }
2502    let run = move |start: usize, end: usize| {
2503        for r in start..end {
2504            // SAFETY: disjoint row ranges per worker.
2505            unsafe {
2506                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2507                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2508            }
2509        }
2510    };
2511    dispatch_rows(pool, rows, &run);
2512}
2513
2514/// Batched q1 matmat: each row's tiles stream once per microbatch.
2515#[allow(clippy::too_many_arguments)]
2516fn q1_matmat(
2517    bytes: &[u8],
2518    xs_all: &[f32],
2519    b: usize,
2520    rows: usize,
2521    cols: usize,
2522    out: &mut [f32],
2523    pool: Option<&Pool>,
2524) {
2525    debug_assert_eq!(out.len(), b * rows);
2526    let gpr = cols / GROUP_SIZE;
2527    let out_addr = SendMut(out.as_mut_ptr());
2528    if a8w8_enabled() {
2529        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2530            .map(|bi| {
2531                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2532                let gsum = q1_group_sums(&act.xq, gpr);
2533                (act, gsum)
2534            })
2535            .collect();
2536        let acts = &acts;
2537        let run = move |start: usize, end: usize| {
2538            for r in start..end {
2539                for (bi, (act, gsum)) in acts.iter().enumerate() {
2540                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2541                    for &(j, xv) in &act.outliers {
2542                        let (w, s) = q1_outlier(bytes, r, gpr, j);
2543                        acc += w * s * xv;
2544                    }
2545                    // SAFETY: disjoint (bi, r) cells per worker range.
2546                    unsafe { *out_addr.at(bi * rows + r) = acc };
2547                }
2548            }
2549        };
2550        dispatch_rows(pool, rows, &run);
2551        return;
2552    }
2553    let run = move |start: usize, end: usize| {
2554        for r in start..end {
2555            for bi in 0..b {
2556                let x = &xs_all[bi * cols..(bi + 1) * cols];
2557                // SAFETY: disjoint (bi, r) cells per worker range.
2558                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
2559            }
2560        }
2561    };
2562    dispatch_rows(pool, rows, &run);
2563}
2564
2565/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
2566/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
2567/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
2568/// 32-group, exact outlier correction — the same A8W8 contract as q8.
2569/// `CMF_SDOT=0` keeps the exact scalar path.
2570fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2571    debug_assert_eq!(out.len(), rows);
2572    let (packed, scales) = q4_split(bytes, rows, cols);
2573    let gpr = cols / GROUP_SIZE;
2574    let out_addr = SendMut(out.as_mut_ptr());
2575
2576    if a8w8_enabled() {
2577        let act = split_act(x);
2578        let run = move |start: usize, end: usize| {
2579            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
2580        };
2581        dispatch_rows(pool, rows, &run);
2582        return;
2583    }
2584
2585    let run = move |start: usize, end: usize| {
2586        q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
2587    };
2588    dispatch_rows(pool, rows, &run);
2589}
2590
2591/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
2592/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
2593#[inline]
2594#[allow(unreachable_code)]
2595fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2596    #[cfg(target_arch = "aarch64")]
2597    unsafe {
2598        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
2599    }
2600    #[cfg(target_arch = "x86_64")]
2601    unsafe {
2602        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
2603    }
2604    let mut acc = 0f32;
2605    for gi in 0..gpr {
2606        let g = g0 + gi;
2607        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2608        let mut d = 0i32;
2609        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2610            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2611                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2612        }
2613        acc += d as f32 * s;
2614    }
2615    acc
2616}
2617
2618/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
2619#[inline]
2620#[allow(unreachable_code)]
2621fn dot_q4_row_i8_2(
2622    packed: &[u8],
2623    scales: &[u8],
2624    g0: usize,
2625    gpr: usize,
2626    xq1: &[i8],
2627    xq2: &[i8],
2628) -> (f32, f32) {
2629    #[cfg(target_arch = "aarch64")]
2630    unsafe {
2631        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
2632    }
2633    #[cfg(target_arch = "x86_64")]
2634    unsafe {
2635        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
2636    }
2637    (
2638        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
2639        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
2640    )
2641}
2642
2643/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
2644/// multi-matrix jobs can drive it for several tensors in one dispatch).
2645#[allow(clippy::too_many_arguments)]
2646fn q4_range_a8w8(
2647    packed: &[u8],
2648    scales: &[u8],
2649    gpr: usize,
2650    cols: usize,
2651    act: &SplitAct,
2652    out: SendMut,
2653    start: usize,
2654    end: usize,
2655) {
2656    for r in start..end {
2657        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
2658        // xq is zeroed at outlier slots — add the exact terms.
2659        for &(j, xv) in &act.outliers {
2660            let flat = r * cols + j;
2661            let byte = packed[flat / 2];
2662            let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2663            let s = f16_to_f32(u16::from_le_bytes([
2664                scales[(flat / GROUP_SIZE) * 2],
2665                scales[(flat / GROUP_SIZE) * 2 + 1],
2666            ]));
2667            acc += ((nib as i32 - 8) as f32) * s * xv;
2668        }
2669        // SAFETY: disjoint row ranges per worker.
2670        unsafe { *out.at(r) = acc };
2671    }
2672}
2673
2674/// Two-input q4 row range via the A8W8 int8 path — kernel body of
2675/// `q4matvec2`, extracted for pair multi-matrix jobs.
2676#[allow(clippy::too_many_arguments)]
2677fn q4_range2_a8w8(
2678    packed: &[u8],
2679    scales: &[u8],
2680    gpr: usize,
2681    cols: usize,
2682    a1: &SplitAct,
2683    a2: &SplitAct,
2684    p1: SendMut,
2685    p2: SendMut,
2686    start: usize,
2687    end: usize,
2688) {
2689    for r in start..end {
2690        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
2691        let mut acc1 = s1 * a1.sx;
2692        let mut acc2 = s2 * a2.sx;
2693        // xq is zeroed at outlier slots — add the exact terms.
2694        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
2695            for &(j, xv) in outliers {
2696                let flat = r * cols + j;
2697                let byte = packed[flat / 2];
2698                let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2699                let s = f16_to_f32(u16::from_le_bytes([
2700                    scales[(flat / GROUP_SIZE) * 2],
2701                    scales[(flat / GROUP_SIZE) * 2 + 1],
2702                ]));
2703                *acc += ((nib as i32 - 8) as f32) * s * xv;
2704            }
2705        };
2706        fix(&a1.outliers, &mut acc1);
2707        fix(&a2.outliers, &mut acc2);
2708        // SAFETY: disjoint row ranges per worker.
2709        unsafe {
2710            *p1.at(r) = acc1;
2711            *p2.at(r) = acc2;
2712        }
2713    }
2714}
2715
2716/// Exact scalar q4 row range (same extraction, non-SDOT path).
2717fn q4_range_f32(
2718    packed: &[u8],
2719    scales: &[u8],
2720    gpr: usize,
2721    x: &[f32],
2722    out: SendMut,
2723    start: usize,
2724    end: usize,
2725) {
2726    for r in start..end {
2727        let mut acc = 0f32;
2728        for gi in 0..gpr {
2729            let g = r * gpr + gi;
2730            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2731            let pk = &packed[g * 16..(g + 1) * 16];
2732            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2733            let mut ga = 0f32;
2734            for (k, &b) in pk.iter().enumerate() {
2735                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2736                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2737            }
2738            acc += ga * s;
2739        }
2740        // SAFETY: disjoint row ranges per worker.
2741        unsafe { *out.at(r) = acc };
2742    }
2743}
2744
2745/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
2746/// dotted against both activations (was: two full matvecs — double
2747/// weight traffic). Per-lane math matches `q4matvec` exactly.
2748#[allow(clippy::too_many_arguments)]
2749fn q4matvec2(
2750    bytes: &[u8],
2751    x1: &[f32],
2752    x2: &[f32],
2753    rows: usize,
2754    cols: usize,
2755    o1: &mut [f32],
2756    o2: &mut [f32],
2757    pool: Option<&Pool>,
2758) {
2759    debug_assert_eq!(o1.len(), rows);
2760    debug_assert_eq!(o2.len(), rows);
2761    let (packed, scales) = q4_split(bytes, rows, cols);
2762    let gpr = cols / GROUP_SIZE;
2763
2764    if a8w8_enabled() {
2765        let a1 = split_act(x1);
2766        let a2 = split_act(x2);
2767        let p1 = SendMut(o1.as_mut_ptr());
2768        let p2 = SendMut(o2.as_mut_ptr());
2769        let run = move |start: usize, end: usize| {
2770            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
2771        };
2772        dispatch_rows(pool, rows, &run);
2773        return;
2774    }
2775
2776    let p1 = SendMut(o1.as_mut_ptr());
2777    let p2 = SendMut(o2.as_mut_ptr());
2778    let run = move |start: usize, end: usize| {
2779        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
2780    };
2781    dispatch_rows(pool, rows, &run);
2782}
2783
2784/// Two-input exact scalar q4 row range (same extraction).
2785#[allow(clippy::too_many_arguments)]
2786fn q4_range2_f32(
2787    packed: &[u8],
2788    scales: &[u8],
2789    gpr: usize,
2790    x1: &[f32],
2791    x2: &[f32],
2792    p1: SendMut,
2793    p2: SendMut,
2794    start: usize,
2795    end: usize,
2796) {
2797    for r in start..end {
2798        let (mut acc1, mut acc2) = (0f32, 0f32);
2799        for gi in 0..gpr {
2800            let g = r * gpr + gi;
2801            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2802            let pk = &packed[g * 16..(g + 1) * 16];
2803            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2804            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2805            let (mut g1, mut g2) = (0f32, 0f32);
2806            for (k, &b) in pk.iter().enumerate() {
2807                let wl = (b & 0x0F) as f32 - 8.0;
2808                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
2809                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
2810                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
2811            }
2812            acc1 += g1 * s;
2813            acc2 += g2 * s;
2814        }
2815        // SAFETY: disjoint row ranges per worker.
2816        unsafe {
2817            *p1.at(r) = acc1;
2818            *p2.at(r) = acc2;
2819        }
2820    }
2821}
2822
2823thread_local! {
2824    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
2825    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
2826    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
2827    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2828}
2829
2830/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
2831/// and dotted against ALL b activations (prefill used to fall back to b
2832/// full matvecs — b× weight traffic and b× nibble decode). Per-position
2833/// math matches `q4matvec` exactly: same group order, same accumulation.
2834/// `out` is row-major [b, rows] like `qmatmat`.
2835#[allow(clippy::too_many_arguments)]
2836fn q4matmat(
2837    bytes: &[u8],
2838    xs_all: &[f32],
2839    b: usize,
2840    rows: usize,
2841    cols: usize,
2842    out: &mut [f32],
2843    pool: Option<&Pool>,
2844) {
2845    debug_assert_eq!(xs_all.len(), b * cols);
2846    debug_assert_eq!(out.len(), b * rows);
2847    let (packed, scales) = q4_split(bytes, rows, cols);
2848    let gpr = cols / GROUP_SIZE;
2849    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2850
2851    if a8w8_enabled() {
2852        let acts: Vec<SplitAct> =
2853            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2854        let acts = &acts;
2855        let out_addr = SendMut(out.as_mut_ptr());
2856        let run = move |start: usize, end: usize| {
2857            ROW_I8.with(|rb| {
2858                let mut buf = rb.borrow_mut();
2859                buf.resize(cols, 0);
2860                for r in start..end {
2861                    // Unpack the row's nibbles to centered i8 once
2862                    // (element 2k = low nibble, 2k+1 = high — flat order,
2863                    // same as dot_q4_row_sdot's zip).
2864                    for gi in 0..gpr {
2865                        let g = r * gpr + gi;
2866                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2867                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
2868                            buf[gi * GROUP_SIZE + k * 2 + 1] =
2869                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
2870                        }
2871                    }
2872                    for (bi, act) in acts.iter().enumerate() {
2873                        let mut acc = 0f32;
2874                        for gi in 0..gpr {
2875                            let d = dot_i8_i8(
2876                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
2877                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
2878                            );
2879                            acc += d as f32 * gscale(r * gpr + gi);
2880                        }
2881                        acc *= act.sx;
2882                        // xq is zeroed at outlier slots — exact terms.
2883                        for &(j, xv) in &act.outliers {
2884                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
2885                        }
2886                        // SAFETY: disjoint (bi, r) cells per worker row range.
2887                        unsafe { *out_addr.at(bi * rows + r) = acc };
2888                    }
2889                }
2890            })
2891        };
2892        dispatch_rows(pool, rows, &run);
2893        return;
2894    }
2895
2896    let out_addr = SendMut(out.as_mut_ptr());
2897    let run = move |start: usize, end: usize| {
2898        ROW_F32.with(|rb| {
2899            let mut buf = rb.borrow_mut();
2900            buf.resize(cols, 0.0);
2901            for r in start..end {
2902                // Decode raw (nib − 8) values once; scales stay per-group
2903                // so the accumulation order matches q4matvec bit-for-bit.
2904                for gi in 0..gpr {
2905                    let g = r * gpr + gi;
2906                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2907                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
2908                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
2909                    }
2910                }
2911                for bi in 0..b {
2912                    let x = &xs_all[bi * cols..(bi + 1) * cols];
2913                    let mut acc = 0f32;
2914                    for gi in 0..gpr {
2915                        let mut ga = 0f32;
2916                        // Pairwise (lo + hi) addition, matching
2917                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
2918                        // a flat one-per-element loop rounds differently
2919                        // and broke bit-parity on the scalar (x86) path.
2920                        for k in 0..GROUP_SIZE / 2 {
2921                            let e = gi * GROUP_SIZE + k * 2;
2922                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
2923                        }
2924                        acc += ga * gscale(r * gpr + gi);
2925                    }
2926                    // SAFETY: disjoint (bi, r) cells per worker row range.
2927                    unsafe { *out_addr.at(bi * rows + r) = acc };
2928                }
2929            }
2930        })
2931    };
2932    dispatch_rows(pool, rows, &run);
2933}
2934
2935/// Batched vbit matmat: each variable-bit row is decoded from the mmap
2936/// ONCE for the whole microbatch. Same per-position math as
2937/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
2938/// and the scalar path).
2939#[allow(clippy::too_many_arguments)]
2940fn vbitmatmat(
2941    bytes: &[u8],
2942    offsets: &[usize],
2943    xs_all: &[f32],
2944    b: usize,
2945    rows: usize,
2946    cols: usize,
2947    out: &mut [f32],
2948    pool: Option<&Pool>,
2949) {
2950    debug_assert_eq!(xs_all.len(), b * cols);
2951    debug_assert_eq!(out.len(), b * rows);
2952    debug_assert_eq!(offsets.len(), rows + 1);
2953    let ng = cols / GROUP_SIZE;
2954    let bits = &bytes[..rows];
2955    let sc_off = rows;
2956    let gscale = |r: usize, g: usize| {
2957        let so = (r * ng + g) * 2;
2958        f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
2959    };
2960
2961    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
2962    let decode_f32 = |r: usize, dst: &mut [f32]| {
2963        let bw = bits[r] as usize;
2964        let l = ((1i32 << (bw - 1)) - 1) as f32;
2965        let data = &bytes[offsets[r]..offsets[r + 1]];
2966        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2967        for d in dst.iter_mut() {
2968            while nbits < bw {
2969                acc = (acc << 8) | data[idx] as u64;
2970                idx += 1;
2971                nbits += 8;
2972            }
2973            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
2974            nbits -= bw;
2975            *d = u - l;
2976        }
2977    };
2978
2979    if a8w8_enabled() {
2980        let acts: Vec<SplitAct> =
2981            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2982        let acts = &acts;
2983        let out_addr = SendMut(out.as_mut_ptr());
2984        let run = move |start: usize, end: usize| {
2985            for r in start..end {
2986                let bw = bits[r] as usize;
2987                if bw == 8 {
2988                    // u−L reaches 128 → no i8 path; decode once, exact
2989                    // f32 dots for every position (same as vbitmatvec).
2990                    ROW_F32.with(|rb| {
2991                        let mut buf = rb.borrow_mut();
2992                        buf.resize(cols, 0.0);
2993                        decode_f32(r, &mut buf);
2994                        for bi in 0..b {
2995                            let x = &xs_all[bi * cols..(bi + 1) * cols];
2996                            let mut dot = 0f32;
2997                            for g in 0..ng {
2998                                let mut gd = 0f32;
2999                                for k in 0..GROUP_SIZE {
3000                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3001                                }
3002                                dot += gd * gscale(r, g);
3003                            }
3004                            // SAFETY: disjoint (bi, r) cells per worker range.
3005                            unsafe { *out_addr.at(bi * rows + r) = dot };
3006                        }
3007                    });
3008                    continue;
3009                }
3010                let l = (1i32 << (bw - 1)) - 1;
3011                let data = &bytes[offsets[r]..offsets[r + 1]];
3012                ROW_I8.with(|rb| {
3013                    let mut buf = rb.borrow_mut();
3014                    buf.resize(cols, 0);
3015                    #[inline(always)]
3016                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3017                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3018                            let u = unpack8::<B>(&data[blk * B..]);
3019                            for k in 0..8 {
3020                                chunk[k] = (u[k] - l) as i8 as u8;
3021                            }
3022                        }
3023                    }
3024                    match bw {
3025                        3 => fill::<3>(data, l, &mut buf),
3026                        4 => vbit_fill4(data, &mut buf),
3027                        5 => fill::<5>(data, l, &mut buf),
3028                        6 => fill::<6>(data, l, &mut buf),
3029                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3030                    }
3031                    for (bi, act) in acts.iter().enumerate() {
3032                        let mut dot = 0f32;
3033                        for g in 0..ng {
3034                            let d = dot_i8_i8(
3035                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3036                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3037                            ) as f32
3038                                * act.sx;
3039                            dot += d * gscale(r, g);
3040                        }
3041                        for &(j, xv) in &act.outliers {
3042                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3043                        }
3044                        // SAFETY: disjoint (bi, r) cells per worker range.
3045                        unsafe { *out_addr.at(bi * rows + r) = dot };
3046                    }
3047                });
3048            }
3049        };
3050        dispatch_rows(pool, rows, &run);
3051        return;
3052    }
3053
3054    let out_addr = SendMut(out.as_mut_ptr());
3055    let run = move |start: usize, end: usize| {
3056        ROW_F32.with(|rb| {
3057            let mut buf = rb.borrow_mut();
3058            buf.resize(cols, 0.0);
3059            for r in start..end {
3060                decode_f32(r, &mut buf);
3061                for bi in 0..b {
3062                    let x = &xs_all[bi * cols..(bi + 1) * cols];
3063                    let mut dot = 0f32;
3064                    for g in 0..ng {
3065                        let mut gd = 0f32;
3066                        for k in 0..GROUP_SIZE {
3067                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3068                        }
3069                        dot += gd * gscale(r, g);
3070                    }
3071                    // SAFETY: disjoint (bi, r) cells per worker range.
3072                    unsafe { *out_addr.at(bi * rows + r) = dot };
3073                }
3074            }
3075        })
3076    };
3077    dispatch_rows(pool, rows, &run);
3078}
3079
3080/// Build a GPU batch job for a q8-family mapped tensor (primary
3081/// shard): prescaled input + directory coordinates. None → not
3082/// GPU-eligible, caller stays on the CPU.
3083pub(crate) fn gpu_batch_job<'a>(
3084    t: &'a QTensor,
3085    x: &[f32],
3086) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3087    match t {
3088        QTensor::Mapped {
3089            model,
3090            idx,
3091            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3092            rows,
3093            cols,
3094            row_scale,
3095            col_field,
3096            ..
3097        } => Some((
3098            model.clone(),
3099            crate::gpu::BatchJob {
3100                idx: *idx,
3101                rows: *rows,
3102                cols: *cols,
3103                row_scale,
3104                xs: prescale(x, col_field, *dt).into_owned(),
3105                q1: false,
3106            },
3107        )),
3108        // q1: raw f32 activations, tile-embedded scales.
3109        QTensor::Mapped {
3110            model,
3111            idx,
3112            dtype: TensorDtype::Q1,
3113            rows,
3114            cols,
3115            ..
3116        } => Some((
3117            model.clone(),
3118            crate::gpu::BatchJob {
3119                idx: *idx,
3120                rows: *rows,
3121                cols: *cols,
3122                row_scale: &[],
3123                xs: x.to_vec(),
3124                q1: true,
3125            },
3126        )),
3127        _ => None,
3128    }
3129}
3130
3131/// θ col-field fold for q8_2f activations. Borrowed pass-through for
3132/// every other dtype — the old unconditional `x.to_vec()` was a pure
3133/// per-matvec allocation on the q8_row hot path.
3134pub(crate) fn prescale<'a>(
3135    x: &'a [f32],
3136    col_field: &[f32],
3137    dtype: TensorDtype,
3138) -> std::borrow::Cow<'a, [f32]> {
3139    if dtype == TensorDtype::Q8_2f {
3140        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3141    } else {
3142        std::borrow::Cow::Borrowed(x)
3143    }
3144}
3145
3146// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
3147
3148/// AVX2+FMA available? Default ON when the CPU supports both;
3149/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
3150#[cfg(target_arch = "x86_64")]
3151pub(crate) fn avx2_enabled() -> bool {
3152    use std::sync::OnceLock;
3153    static ON: OnceLock<bool> = OnceLock::new();
3154    *ON.get_or_init(|| {
3155        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3156            && std::arch::is_x86_feature_detected!("avx2")
3157            && std::arch::is_x86_feature_detected!("fma")
3158    })
3159}
3160
3161/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
3162/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
3163/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
3164/// active either way, they are exact (regrouped sums only).
3165#[cfg(target_arch = "x86_64")]
3166fn avx2_a8w8_enabled() -> bool {
3167    use std::sync::OnceLock;
3168    static ON: OnceLock<bool> = OnceLock::new();
3169    *ON.get_or_init(|| {
3170        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3171    })
3172}
3173
3174/// A8W8 quantized-activation path available on THIS machine? One
3175/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
3176/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
3177#[inline]
3178pub(crate) fn a8w8_enabled() -> bool {
3179    #[cfg(target_arch = "aarch64")]
3180    {
3181        sdot_enabled()
3182    }
3183    #[cfg(target_arch = "x86_64")]
3184    {
3185        avx2_a8w8_enabled()
3186    }
3187    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3188    {
3189        false
3190    }
3191}
3192
3193/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
3194/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
3195#[inline]
3196#[allow(unreachable_code)]
3197fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
3198    #[cfg(target_arch = "aarch64")]
3199    unsafe {
3200        return dot_i8_sdot(w, xq);
3201    }
3202    #[cfg(target_arch = "x86_64")]
3203    unsafe {
3204        if avx512vnni_enabled() {
3205            return dot_i8_i8_vnni(w, xq);
3206        }
3207        return dot_i8_i8_avx2(w, xq);
3208    }
3209    w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
3210}
3211
3212/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
3213/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
3214/// `vpdpbusd` encoding.
3215#[cfg(target_arch = "x86_64")]
3216fn avx512vnni_enabled() -> bool {
3217    use std::sync::OnceLock;
3218    static ON: OnceLock<bool> = OnceLock::new();
3219    *ON.get_or_init(|| {
3220        std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
3221            && std::arch::is_x86_feature_detected!("avx512f")
3222            && std::arch::is_x86_feature_detected!("avx512bw")
3223            && std::arch::is_x86_feature_detected!("avx512vl")
3224            && std::arch::is_x86_feature_detected!("avx512vnni")
3225    })
3226}
3227
3228/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
3229/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
3230/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
3231/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
3232#[cfg(target_arch = "x86_64")]
3233#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3234unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3235    // SAFETY: callers uphold slice-length contracts (see call sites).
3236    unsafe {
3237        use core::arch::x86_64::*;
3238        let n = w.len();
3239        let mut j = 0usize;
3240        let mut total: i32;
3241        // 4 independent accumulators: vpdpbusd is its own loop-carried
3242        // dependency (~5-cycle latency) — a single-acc loop runs
3243        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
3244        // on Granite Rapids.
3245        {
3246            #[inline(always)]
3247            unsafe fn step(
3248                w: *const u8,
3249                x: *const i8,
3250                acc: core::arch::x86_64::__m512i,
3251            ) -> core::arch::x86_64::__m512i {
3252                unsafe {
3253                    use core::arch::x86_64::*;
3254                    let wv = _mm512_loadu_si512(w as *const _);
3255                    let xv = _mm512_loadu_si512(x as *const _);
3256                    let aw = _mm512_abs_epi8(wv);
3257                    let neg = _mm512_movepi8_mask(wv);
3258                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3259                    _mm512_dpbusd_epi32(acc, aw, sx)
3260                }
3261            }
3262            let (mut a0, mut a1, mut a2, mut a3) = (
3263                _mm512_setzero_si512(),
3264                _mm512_setzero_si512(),
3265                _mm512_setzero_si512(),
3266                _mm512_setzero_si512(),
3267            );
3268            while j + 256 <= n {
3269                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3270                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3271                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3272                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3273                j += 256;
3274            }
3275            while j + 64 <= n {
3276                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3277                j += 64;
3278            }
3279            let s01 = _mm512_add_epi32(a0, a1);
3280            let s23 = _mm512_add_epi32(a2, a3);
3281            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3282        }
3283        // 32-wide (q4/vbit groups are exactly 32 bytes).
3284        if j + 32 <= n {
3285            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3286            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3287            let d = _mm256_dpbusd_epi32(
3288                _mm256_setzero_si256(),
3289                _mm256_abs_epi8(wv),
3290                _mm256_sign_epi8(xv, wv),
3291            );
3292            let hi128 = _mm256_extracti128_si256::<1>(d);
3293            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3294            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3295            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3296            total += _mm_cvtsi128_si32(s32);
3297            j += 32;
3298        }
3299        while j < n {
3300            total += (w[j] as i8) as i32 * xq[j] as i32;
3301            j += 1;
3302        }
3303        total
3304    }
3305}
3306
3307/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
3308#[cfg(target_arch = "x86_64")]
3309#[target_feature(enable = "avx2,fma")]
3310unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3311    // SAFETY: callers uphold slice-length contracts (see call sites).
3312    unsafe {
3313        use core::arch::x86_64::*;
3314        let n = x.len();
3315        let wp = w.as_ptr();
3316        let xp = x.as_ptr();
3317        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3318        let mut j = 0usize;
3319        while j + 16 <= n {
3320            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3321            let lo = _mm256_cvtepi8_epi32(wb);
3322            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3323            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3324            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3325            j += 16;
3326        }
3327        let acc = _mm256_add_ps(a0, a1);
3328        let hi128 = _mm256_extractf128_ps::<1>(acc);
3329        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3330        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3331        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3332        let mut sum = _mm_cvtss_f32(s32);
3333        while j < n {
3334            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3335            j += 1;
3336        }
3337        sum
3338    }
3339}
3340
3341/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
3342/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
3343/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
3344/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
3345#[cfg(target_arch = "x86_64")]
3346#[target_feature(enable = "avx2")]
3347unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3348    // SAFETY: callers uphold slice-length contracts (see call sites).
3349    unsafe {
3350        use core::arch::x86_64::*;
3351        let n = w.len();
3352        let ones = _mm256_set1_epi16(1);
3353        let mut acc = _mm256_setzero_si256();
3354        let mut j = 0usize;
3355        while j + 32 <= n {
3356            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3357            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3358            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3359            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3360            j += 32;
3361        }
3362        let hi128 = _mm256_extracti128_si256::<1>(acc);
3363        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3364        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3365        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3366        let mut s = _mm_cvtsi128_si32(s32);
3367        while j < n {
3368            s += (w[j] as i8) as i32 * xq[j] as i32;
3369            j += 1;
3370        }
3371        s
3372    }
3373}
3374
3375/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
3376/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
3377/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
3378/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
3379#[cfg(target_arch = "x86_64")]
3380#[inline]
3381fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
3382    let dot = if avx512vnni_enabled() && row.len() >= 64 {
3383        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
3384    } else {
3385        unsafe { dot_i8_i8_avx2(row, &act.xq) }
3386    };
3387    let mut acc = dot as f32 * act.sx;
3388    for &(j, xv) in &act.outliers {
3389        acc += (row[j] as i8) as f32 * xv;
3390    }
3391    acc
3392}
3393
3394/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
3395/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
3396/// a single-acc loop runs latency-bound, measured on Granite Rapids).
3397#[cfg(target_arch = "x86_64")]
3398#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3399unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3400    // SAFETY: callers uphold slice-length contracts (see call sites).
3401    unsafe {
3402        use core::arch::x86_64::*;
3403        let n = w.len();
3404        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
3405        #[inline(always)]
3406        unsafe fn step(
3407            w: *const u8,
3408            x: *const i8,
3409            flip: core::arch::x86_64::__m512i,
3410            acc: core::arch::x86_64::__m512i,
3411        ) -> core::arch::x86_64::__m512i {
3412            unsafe {
3413                use core::arch::x86_64::*;
3414                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
3415                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
3416            }
3417        }
3418        let (mut a0, mut a1, mut a2, mut a3) = (
3419            _mm512_setzero_si512(),
3420            _mm512_setzero_si512(),
3421            _mm512_setzero_si512(),
3422            _mm512_setzero_si512(),
3423        );
3424        let mut j = 0usize;
3425        while j + 256 <= n {
3426            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3427            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
3428            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
3429            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
3430            j += 256;
3431        }
3432        while j + 64 <= n {
3433            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3434            j += 64;
3435        }
3436        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
3437            _mm512_add_epi32(a0, a1),
3438            _mm512_add_epi32(a2, a3),
3439        ));
3440        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
3441        while j < n {
3442            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
3443            j += 1;
3444        }
3445        total
3446    }
3447}
3448
3449/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
3450/// writer's flat order, same as the NEON vzip pair), maddubs against
3451/// the pre-quantized activation group, × the group's f16 scale. Pair
3452/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
3453/// `dot_q4_row_sdot`.
3454#[cfg(target_arch = "x86_64")]
3455#[target_feature(enable = "avx2")]
3456unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3457    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
3458    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
3459    unsafe {
3460        use core::arch::x86_64::*;
3461        let lomask = _mm_set1_epi8(0x0F);
3462        let eight = _mm256_set1_epi8(8);
3463        let ones = _mm256_set1_epi16(1);
3464        let mut acc = 0f32;
3465        for gi in 0..gpr {
3466            let g = g0 + gi;
3467            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3468            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3469            let lo = _mm_and_si128(b, lomask);
3470            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3471            let w = _mm256_sub_epi8(
3472                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3473                eight,
3474            );
3475            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3476            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3477            let d = _mm256_madd_epi16(p16, ones);
3478            let hi128 = _mm256_extracti128_si256::<1>(d);
3479            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3480            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3481            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3482            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3483        }
3484        acc
3485    }
3486}
3487
3488/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
3489/// both activations dotted against the same centered i8 register.
3490#[cfg(target_arch = "x86_64")]
3491#[target_feature(enable = "avx2")]
3492unsafe fn dot_q4_row_avx2_2(
3493    packed: &[u8],
3494    scales: &[u8],
3495    g0: usize,
3496    gpr: usize,
3497    xq1: &[i8],
3498    xq2: &[i8],
3499) -> (f32, f32) {
3500    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
3501    unsafe {
3502        use core::arch::x86_64::*;
3503        let lomask = _mm_set1_epi8(0x0F);
3504        let eight = _mm256_set1_epi8(8);
3505        let ones = _mm256_set1_epi16(1);
3506        let (mut acc1, mut acc2) = (0f32, 0f32);
3507        #[inline(always)]
3508        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
3509            unsafe {
3510                use core::arch::x86_64::*;
3511                let hi128 = _mm256_extracti128_si256::<1>(d);
3512                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3513                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3514                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3515                _mm_cvtsi128_si32(s32)
3516            }
3517        }
3518        for gi in 0..gpr {
3519            let g = g0 + gi;
3520            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3521            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3522            let lo = _mm_and_si128(b, lomask);
3523            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3524            let w = _mm256_sub_epi8(
3525                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3526                eight,
3527            );
3528            let aw = _mm256_abs_epi8(w);
3529            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3530            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3531            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
3532            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
3533            acc1 += hsum(d1) as f32 * s;
3534            acc2 += hsum(d2) as f32 * s;
3535        }
3536        (acc1, acc2)
3537    }
3538}
3539
3540/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
3541#[cfg(target_arch = "x86_64")]
3542fn q8_range_avx2(
3543    q: &[u8],
3544    row_scale: &[f32],
3545    act: &SplitAct,
3546    cols: usize,
3547    out_addr: SendMut,
3548    start: usize,
3549    end: usize,
3550) {
3551    for o in start..end {
3552        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3553        // SAFETY: disjoint row ranges per worker.
3554        unsafe { *out_addr.at(o) = v };
3555    }
3556}
3557
3558/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
3559#[cfg(target_arch = "x86_64")]
3560#[allow(clippy::too_many_arguments)]
3561fn q8_range2_avx2(
3562    q: &[u8],
3563    row_scale: &[f32],
3564    a1: &SplitAct,
3565    a2: &SplitAct,
3566    cols: usize,
3567    p1: SendMut,
3568    p2: SendMut,
3569    start: usize,
3570    end: usize,
3571) {
3572    for o in start..end {
3573        let row = &q[o * cols..(o + 1) * cols];
3574        // SAFETY: disjoint row ranges per worker.
3575        unsafe {
3576            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
3577            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
3578        }
3579    }
3580}
3581
3582// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
3583
3584/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
3585/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
3586/// (On non-ARM release builds only the test tolerance switch calls it.)
3587#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3588fn sdot_enabled() -> bool {
3589    use std::sync::OnceLock;
3590    static ON: OnceLock<bool> = OnceLock::new();
3591    *ON.get_or_init(|| {
3592        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
3593        #[cfg(target_arch = "aarch64")]
3594        {
3595            want && std::arch::is_aarch64_feature_detected!("dotprod")
3596        }
3597        #[cfg(not(target_arch = "aarch64"))]
3598        {
3599            let _ = want;
3600            false
3601        }
3602    })
3603}
3604
3605/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
3606/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
3607/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
3608/// matvec, shared by all rows/workers.
3609struct SplitAct {
3610    xq: Vec<i8>,
3611    sx: f32,
3612    outliers: Vec<(usize, f32)>,
3613    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
3614    /// `−128·Σx`); one i32 per split, computed once per matvec.
3615    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
3616    xsum: i32,
3617}
3618
3619thread_local! {
3620    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
3621    /// and its hidden-size allocation was steady-state heap churn.
3622    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
3623        const { std::cell::RefCell::new(Vec::new()) };
3624}
3625
3626impl Drop for SplitAct {
3627    fn drop(&mut self) {
3628        let buf = std::mem::take(&mut self.xq);
3629        if buf.capacity() > 0 {
3630            XQ_FREE.with(|f| {
3631                let mut f = f.borrow_mut();
3632                if f.len() < 16 {
3633                    f.push(buf);
3634                }
3635            });
3636        }
3637    }
3638}
3639
3640fn split_act(x: &[f32]) -> SplitAct {
3641    let n = x.len();
3642    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
3643    let thr = 8.0 * rms;
3644    // One pass: collect outliers and the bulk absmax (outliers excluded —
3645    // identical to the old zero-then-fold over a copied buffer, minus the
3646    // full-vector copy).
3647    let mut outliers: Vec<(usize, f32)> = Vec::new();
3648    let mut amax = 0f32;
3649    for (j, &v) in x.iter().enumerate() {
3650        let a = v.abs();
3651        if a > thr {
3652            outliers.push((j, v));
3653        } else if a > amax {
3654            amax = a;
3655        }
3656    }
3657    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
3658    let inv = 1.0 / sx;
3659    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
3660    xq.clear();
3661    xq.reserve(n);
3662    if outliers.is_empty() {
3663        xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
3664    } else {
3665        // Outlier slots quantize to 0 (their exact term is added later).
3666        xq.extend(x.iter().map(|&v| {
3667            if v.abs() > thr {
3668                0
3669            } else {
3670                (v * inv).round().clamp(-127.0, 127.0) as i8
3671            }
3672        }));
3673    }
3674    let xsum = xq.iter().map(|&v| v as i32).sum();
3675    SplitAct { xq, sx, outliers, xsum }
3676}
3677
3678/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
3679/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
3680#[cfg(target_arch = "aarch64")]
3681#[target_feature(enable = "neon,dotprod")]
3682unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
3683    // SAFETY: callers uphold slice-length contracts (see call sites).
3684    unsafe {
3685        use core::arch::aarch64::*;
3686        use core::arch::asm;
3687        let wp = w.as_ptr() as *const i8;
3688        let n = w.len();
3689        let (mut a0, mut a1, mut a2, mut a3) =
3690            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3691        let mut i = 0;
3692        while i + 64 <= n {
3693            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
3694            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
3695            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
3696            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
3697            asm!(
3698                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3699                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3700                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
3701                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
3702                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3703                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3704                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
3705                options(pure, nomem, nostack),
3706            );
3707            i += 64;
3708        }
3709        while i + 16 <= n {
3710            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
3711            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
3712                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
3713            i += 16;
3714        }
3715        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
3716        while i < n {
3717            s += (*wp.add(i)) as i32 * xq[i] as i32;
3718            i += 1;
3719        }
3720        s
3721}
3722}
3723
3724/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
3725/// loaded once and reused, 4 independent accumulators hide sdot latency
3726/// (port of vmfcore `dot_i8_sdot_4rows`).
3727#[cfg(target_arch = "aarch64")]
3728#[target_feature(enable = "neon,dotprod")]
3729unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
3730    // SAFETY: callers uphold slice-length contracts (see call sites).
3731    unsafe {
3732        use core::arch::aarch64::*;
3733        use core::arch::asm;
3734        let n = xq.len();
3735        let px = xq.as_ptr();
3736        let (p0, p1, p2, p3) = (
3737            w0.as_ptr() as *const i8,
3738            w1.as_ptr() as *const i8,
3739            w2.as_ptr() as *const i8,
3740            w3.as_ptr() as *const i8,
3741        );
3742        let (mut a0, mut a1, mut a2, mut a3) =
3743            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3744        let mut i = 0;
3745        while i + 16 <= n {
3746            let x = vld1q_s8(px.add(i));
3747            let v0 = vld1q_s8(p0.add(i));
3748            let v1 = vld1q_s8(p1.add(i));
3749            let v2 = vld1q_s8(p2.add(i));
3750            let v3 = vld1q_s8(p3.add(i));
3751            asm!(
3752                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
3753                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
3754                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
3755                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
3756                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3757                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
3758                options(pure, nomem, nostack),
3759            );
3760            i += 16;
3761        }
3762        let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
3763        while i < n {
3764            let xi = *px.add(i) as i32;
3765            r[0] += (*p0.add(i)) as i32 * xi;
3766            r[1] += (*p1.add(i)) as i32 * xi;
3767            r[2] += (*p2.add(i)) as i32 * xi;
3768            r[3] += (*p3.add(i)) as i32 * xi;
3769            i += 1;
3770        }
3771        r
3772}
3773}
3774
3775/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
3776/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
3777/// line plus the shared activation chunk — a single sequential weight
3778/// stream per worker. Per-row accumulation is the same one-accumulator
3779/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
3780/// are bit-identical to the mmap-layout kernel.
3781#[cfg(target_arch = "aarch64")]
3782#[target_feature(enable = "neon,dotprod")]
3783unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
3784    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
3785    // n % 16 == 0 — guaranteed by the repack gate).
3786    unsafe {
3787        use core::arch::aarch64::*;
3788        use core::arch::asm;
3789        let n = xq.len();
3790        let px = xq.as_ptr();
3791        let pg = g.as_ptr() as *const i8;
3792        let (mut a0, mut a1, mut a2, mut a3) =
3793            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3794        let mut i = 0;
3795        while i + 16 <= n {
3796            let x = vld1q_s8(px.add(i));
3797            let base = pg.add(4 * i);
3798            let v0 = vld1q_s8(base);
3799            let v1 = vld1q_s8(base.add(16));
3800            let v2 = vld1q_s8(base.add(32));
3801            let v3 = vld1q_s8(base.add(48));
3802            asm!(
3803                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
3804                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
3805                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
3806                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
3807                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3808                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
3809                options(pure, nomem, nostack),
3810            );
3811            i += 16;
3812        }
3813        [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
3814    }
3815}
3816
3817/// One q8 row range via SDOT (4-row blocks + tail) — the body of
3818/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
3819/// SAME kernel for several tensors under one pool dispatch. `rep` — the
3820/// load-time interleaved repack (empty = mmap layout only); rows outside
3821/// full 4-row groups always come from the mmap layout.
3822#[cfg(target_arch = "aarch64")]
3823fn q8_range_sdot(
3824    q: &[u8],
3825    rep: &[u8],
3826    row_scale: &[f32],
3827    act: &SplitAct,
3828    cols: usize,
3829    out_addr: SendMut,
3830    start: usize,
3831    end: usize,
3832) {
3833    let mut o = start;
3834    // Leading rows to the group boundary (repack path only): the pool
3835    // splits row ranges arbitrarily, groups are absolute.
3836    if !rep.is_empty() {
3837        while o < end && o % 4 != 0 {
3838            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3839            unsafe { *out_addr.at(o) = v };
3840            o += 1;
3841        }
3842    }
3843    while o + 4 <= end {
3844        let r = if rep.is_empty() {
3845            unsafe {
3846                dot_i8_sdot_4rows(
3847                    &q[o * cols..(o + 1) * cols],
3848                    &q[(o + 1) * cols..(o + 2) * cols],
3849                    &q[(o + 2) * cols..(o + 3) * cols],
3850                    &q[(o + 3) * cols..(o + 4) * cols],
3851                    &act.xq,
3852                )
3853            }
3854        } else {
3855            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
3856        };
3857        for k in 0..4 {
3858            let mut acc = r[k] as f32 * act.sx;
3859            for &(j, xv) in &act.outliers {
3860                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
3861            }
3862            // SAFETY: disjoint row ranges per worker.
3863            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
3864        }
3865        o += 4;
3866    }
3867    while o < end {
3868        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3869        unsafe { *out_addr.at(o) = v };
3870        o += 1;
3871    }
3872}
3873
3874/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
3875/// for the fused pair multi-matrix job (`matvec2_many`).
3876#[cfg(target_arch = "aarch64")]
3877#[allow(clippy::too_many_arguments)]
3878fn q8_range2_sdot(
3879    q: &[u8],
3880    row_scale: &[f32],
3881    a1: &SplitAct,
3882    a2: &SplitAct,
3883    cols: usize,
3884    p1: SendMut,
3885    p2: SendMut,
3886    start: usize,
3887    end: usize,
3888) {
3889    for o in start..end {
3890        let row = &q[o * cols..(o + 1) * cols];
3891        // SAFETY: disjoint row ranges per worker.
3892        unsafe {
3893            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
3894            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
3895        }
3896    }
3897}
3898
3899/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
3900#[allow(clippy::too_many_arguments)]
3901fn q8_range2_f32(
3902    q: &[u8],
3903    row_scale: &[f32],
3904    x1: &[f32],
3905    x2: &[f32],
3906    cols: usize,
3907    p1: SendMut,
3908    p2: SendMut,
3909    start: usize,
3910    end: usize,
3911) {
3912    for o in start..end {
3913        let row = &q[o * cols..(o + 1) * cols];
3914        // SAFETY: disjoint row ranges per worker.
3915        unsafe {
3916            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
3917            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
3918        }
3919    }
3920}
3921
3922/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
3923fn q8_range_f32(
3924    q: &[u8],
3925    row_scale: &[f32],
3926    xs: &[f32],
3927    cols: usize,
3928    out_addr: SendMut,
3929    start: usize,
3930    end: usize,
3931) {
3932    for o in start..end {
3933        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
3934        // SAFETY: disjoint row ranges per worker.
3935        unsafe { *out_addr.at(o) = v };
3936    }
3937}
3938
3939/// SDOT row dot with exact outlier correction:
3940/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
3941#[cfg(target_arch = "aarch64")]
3942#[inline]
3943fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
3944    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
3945    for &(j, xv) in &act.outliers {
3946        acc += (row[j] as i8) as f32 * xv;
3947    }
3948    acc
3949}
3950
3951/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
3952/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
3953/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
3954/// the caller multiplies by the activation scale and adds the exact
3955/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
3956/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
3957/// → zip(lo,hi) restores flat order.
3958#[cfg(target_arch = "aarch64")]
3959#[target_feature(enable = "neon,dotprod")]
3960unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3961    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
3962    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
3963    unsafe {
3964        use core::arch::aarch64::*;
3965        use core::arch::asm;
3966        let lomask = vdupq_n_u8(0x0F);
3967        let eight = vdupq_n_s8(8);
3968        let mut acc = 0f32;
3969        for gi in 0..gpr {
3970            let g = g0 + gi;
3971            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3972            let b = vld1q_u8(packed.as_ptr().add(g * 16));
3973            let lo = vandq_u8(b, lomask);
3974            let hi = vshrq_n_u8::<4>(b);
3975            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3976            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3977            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3978            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3979            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3980            asm!(
3981                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3982                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3983                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3984                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3985                options(pure, nomem, nostack),
3986            );
3987            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3988        }
3989        acc
3990    }
3991}
3992
3993/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
3994/// part) happens ONCE per group; both pre-quantized activations are
3995/// dotted against the same centered i8 registers. Per-lane math matches
3996/// `dot_q4_row_sdot` exactly.
3997#[cfg(target_arch = "aarch64")]
3998#[target_feature(enable = "neon,dotprod")]
3999unsafe fn dot_q4_row_sdot2(
4000    packed: &[u8],
4001    scales: &[u8],
4002    g0: usize,
4003    gpr: usize,
4004    xq1: &[i8],
4005    xq2: &[i8],
4006) -> (f32, f32) {
4007    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
4008    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
4009    unsafe {
4010        use core::arch::aarch64::*;
4011        use core::arch::asm;
4012        let lomask = vdupq_n_u8(0x0F);
4013        let eight = vdupq_n_s8(8);
4014        let (mut acc1, mut acc2) = (0f32, 0f32);
4015        for gi in 0..gpr {
4016            let g = g0 + gi;
4017            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4018            let b = vld1q_u8(packed.as_ptr().add(g * 16));
4019            let lo = vandq_u8(b, lomask);
4020            let hi = vshrq_n_u8::<4>(b);
4021            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4022            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4023            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
4024            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
4025            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
4026            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
4027            let (mut a0, mut a1, mut b0, mut b1) =
4028                (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4029            asm!(
4030                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
4031                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
4032                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
4033                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
4034                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4035                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
4036                e0 = in(vreg) e0, e1 = in(vreg) e1,
4037                x10 = in(vreg) x10, x11 = in(vreg) x11,
4038                x20 = in(vreg) x20, x21 = in(vreg) x21,
4039                options(pure, nomem, nostack),
4040            );
4041            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4042            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
4043        }
4044        (acc1, acc2)
4045    }
4046}
4047
4048// ───────────────────── fused int8 kernels ─────────────────────
4049
4050/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
4051/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
4052#[inline]
4053pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
4054    #[cfg(target_arch = "aarch64")]
4055    unsafe {
4056        return axpy_i8_f32_neon(acc, row, w);
4057    }
4058    #[cfg(target_arch = "x86_64")]
4059    if avx2_enabled() {
4060        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
4061    }
4062    #[allow(unreachable_code)]
4063    {
4064        for (a, &b) in acc.iter_mut().zip(row) {
4065            *a += w * b as f32;
4066        }
4067    }
4068}
4069
4070/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
4071#[cfg(target_arch = "x86_64")]
4072#[target_feature(enable = "avx2,fma")]
4073unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
4074    // SAFETY: callers uphold slice-length contracts (see call sites).
4075    unsafe {
4076        use core::arch::x86_64::*;
4077        let n = acc.len().min(row.len());
4078        let ap = acc.as_mut_ptr();
4079        let rp = row.as_ptr();
4080        let wv = _mm256_set1_ps(w);
4081        let mut j = 0usize;
4082        while j + 16 <= n {
4083            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
4084            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
4085            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
4086            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
4087            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
4088            _mm256_storeu_ps(ap.add(j), v0);
4089            _mm256_storeu_ps(ap.add(j + 8), v1);
4090            j += 16;
4091        }
4092        while j < n {
4093            *ap.add(j) += w * (*rp.add(j)) as f32;
4094            j += 1;
4095        }
4096    }
4097}
4098
4099#[cfg(target_arch = "aarch64")]
4100#[target_feature(enable = "neon")]
4101unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
4102    // SAFETY: callers uphold slice-length contracts (see call sites).
4103    unsafe {
4104        use core::arch::aarch64::*;
4105        let n = acc.len().min(row.len());
4106        let ap = acc.as_mut_ptr();
4107        let rp = row.as_ptr();
4108        let wv = vdupq_n_f32(w);
4109        let mut j = 0usize;
4110        while j + 16 <= n {
4111            let rb = vld1q_s8(rp.add(j));
4112            let lo = vmovl_s8(vget_low_s8(rb));
4113            let hi = vmovl_s8(vget_high_s8(rb));
4114            for (off, half) in [(0, lo), (8, hi)] {
4115                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
4116                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
4117                let o = j + off;
4118                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
4119                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
4120            }
4121            j += 16;
4122        }
4123        while j < n {
4124            *ap.add(j) += w * (*rp.add(j)) as f32;
4125            j += 1;
4126        }
4127}
4128}
4129
4130/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
4131/// ≈9× scalar), scalar elsewhere.
4132#[inline]
4133pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
4134    #[cfg(target_arch = "aarch64")]
4135    unsafe {
4136        return dot_i8_f32_neon(w, x);
4137    }
4138    #[cfg(target_arch = "x86_64")]
4139    if avx2_enabled() {
4140        return unsafe { dot_i8_f32_avx2(w, x) };
4141    }
4142    #[allow(unreachable_code)]
4143    {
4144        let mut sum = 0.0f32;
4145        for (j, &b) in w.iter().enumerate() {
4146            sum += (b as i8) as f32 * x[j];
4147        }
4148        sum
4149    }
4150}
4151
4152/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
4153/// folded into the product (no prescaled copy of x). NEON on aarch64,
4154/// scalar elsewhere. Used by the active-neuron path `row_dot`.
4155#[inline]
4156fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4157    #[cfg(target_arch = "aarch64")]
4158    unsafe {
4159        return dot_i8_col_f32_neon(w, x, col);
4160    }
4161    #[allow(unreachable_code)]
4162    {
4163        let mut sum = 0.0f32;
4164        for (j, &b) in w.iter().enumerate() {
4165            sum += (b as i8) as f32 * x[j] * col[j];
4166        }
4167        sum
4168    }
4169}
4170
4171#[cfg(target_arch = "aarch64")]
4172#[target_feature(enable = "neon")]
4173unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4174    // SAFETY: callers uphold slice-length contracts (see call sites).
4175    unsafe {
4176        use core::arch::aarch64::*;
4177        let n = x.len();
4178        let wp = w.as_ptr() as *const i8;
4179        let xp = x.as_ptr();
4180        let cp = col.as_ptr();
4181        let (mut a0, mut a1, mut a2, mut a3) =
4182            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4183        let mut j = 0usize;
4184        while j + 16 <= n {
4185            let wb = vld1q_s8(wp.add(j));
4186            let lo = vmovl_s8(vget_low_s8(wb));
4187            let hi = vmovl_s8(vget_high_s8(wb));
4188            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4189            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4190            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4191            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4192            a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
4193            a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
4194            a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
4195            a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
4196            j += 16;
4197        }
4198        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4199        while j < n {
4200            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
4201            j += 1;
4202        }
4203        sum
4204}
4205}
4206
4207#[cfg(target_arch = "aarch64")]
4208#[target_feature(enable = "neon")]
4209unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
4210    // SAFETY: callers uphold slice-length contracts (see call sites).
4211    unsafe {
4212        use core::arch::aarch64::*;
4213        let n = x.len();
4214        let wp = w.as_ptr() as *const i8;
4215        let xp = x.as_ptr();
4216        let (mut a0, mut a1, mut a2, mut a3) =
4217            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4218        let mut j = 0usize;
4219        while j + 16 <= n {
4220            let wb = vld1q_s8(wp.add(j));
4221            let lo = vmovl_s8(vget_low_s8(wb));
4222            let hi = vmovl_s8(vget_high_s8(wb));
4223            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4224            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4225            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4226            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4227            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
4228            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
4229            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
4230            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
4231            j += 16;
4232        }
4233        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4234        while j < n {
4235            sum += (*wp.add(j)) as f32 * *xp.add(j);
4236            j += 1;
4237        }
4238        sum
4239}
4240}
4241
4242#[allow(clippy::too_many_arguments)]
4243fn qmatvec(
4244    q: &[u8],
4245    rep: &[u8],
4246    row_scale: &[f32],
4247    xs: &[f32],
4248    rows: usize,
4249    cols: usize,
4250    out: &mut [f32],
4251    pool: Option<&Pool>,
4252) {
4253    debug_assert_eq!(out.len(), rows);
4254    #[cfg(not(target_arch = "aarch64"))]
4255    let _ = rep;
4256
4257    #[cfg(target_arch = "aarch64")]
4258    if sdot_enabled() {
4259        let act = split_act(xs);
4260        let out_addr = SendMut(out.as_mut_ptr());
4261        let run_range = |start: usize, end: usize| {
4262            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
4263        };
4264        match pool {
4265            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4266            _ => run_range(0, rows),
4267        }
4268        return;
4269    }
4270    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
4271    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
4272    #[cfg(target_arch = "x86_64")]
4273    if avx2_a8w8_enabled() {
4274        let act = split_act(xs);
4275        let out_addr = SendMut(out.as_mut_ptr());
4276        let run_range =
4277            |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
4278        match pool {
4279            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4280            _ => run_range(0, rows),
4281        }
4282        return;
4283    }
4284
4285    let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
4286    match pool {
4287        Some(pool) if rows >= 256 => {
4288            let out_addr = SendMut(out.as_mut_ptr());
4289            pool.run(&move |widx, n| {
4290                let chunk = rows.div_ceil(n);
4291                let start = widx * chunk;
4292                let end = (start + chunk).min(rows);
4293                for o in start..end {
4294                    // SAFETY: disjoint row ranges per worker.
4295                    unsafe { *out_addr.at(o) = row_dot(o) };
4296                }
4297            });
4298        }
4299        _ => {
4300            for (o, dst) in out.iter_mut().enumerate() {
4301                *dst = row_dot(o);
4302            }
4303        }
4304    }
4305}
4306
4307#[allow(clippy::too_many_arguments)]
4308fn qmatvec2(
4309    q: &[u8],
4310    row_scale: &[f32],
4311    x1: &[f32],
4312    x2: &[f32],
4313    rows: usize,
4314    cols: usize,
4315    o1: &mut [f32],
4316    o2: &mut [f32],
4317    pool: Option<&Pool>,
4318) {
4319    #[cfg(target_arch = "aarch64")]
4320    if sdot_enabled() {
4321        let a1s = split_act(x1);
4322        let a2s = split_act(x2);
4323        let p1 = SendMut(o1.as_mut_ptr());
4324        let p2 = SendMut(o2.as_mut_ptr());
4325        let run_range = |start: usize, end: usize| {
4326            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4327        };
4328        match pool {
4329            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4330            _ => run_range(0, rows),
4331        }
4332        return;
4333    }
4334    #[cfg(target_arch = "x86_64")]
4335    if avx2_a8w8_enabled() {
4336        let a1s = split_act(x1);
4337        let a2s = split_act(x2);
4338        let p1 = SendMut(o1.as_mut_ptr());
4339        let p2 = SendMut(o2.as_mut_ptr());
4340        let run_range = |start: usize, end: usize| {
4341            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4342        };
4343        match pool {
4344            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4345            _ => run_range(0, rows),
4346        }
4347        return;
4348    }
4349
4350    let row_dots = |o: usize| -> (f32, f32) {
4351        let row = &q[o * cols..(o + 1) * cols];
4352        (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
4353    };
4354    match pool {
4355        Some(pool) if rows >= 256 => {
4356            let p1 = SendMut(o1.as_mut_ptr());
4357            let p2 = SendMut(o2.as_mut_ptr());
4358            pool.run(&move |widx, n| {
4359                let chunk = rows.div_ceil(n);
4360                let start = widx * chunk;
4361                let end = (start + chunk).min(rows);
4362                for o in start..end {
4363                    let (s1, s2) = row_dots(o);
4364                    // SAFETY: disjoint row ranges per worker.
4365                    unsafe {
4366                        *p1.at(o) = s1;
4367                        *p2.at(o) = s2;
4368                    }
4369                }
4370            });
4371        }
4372        _ => {
4373            for o in 0..rows {
4374                let (s1, s2) = row_dots(o);
4375                o1[o] = s1;
4376                o2[o] = s2;
4377            }
4378        }
4379    }
4380}
4381
4382#[derive(Clone, Copy)]
4383struct SendMut(*mut f32);
4384unsafe impl Send for SendMut {}
4385unsafe impl Sync for SendMut {}
4386
4387impl SendMut {
4388    #[inline]
4389    fn at(self, i: usize) -> *mut f32 {
4390        unsafe { self.0.add(i) }
4391    }
4392}
4393
4394#[cfg(test)]
4395mod tests {
4396    use super::*;
4397
4398    #[test]
4399    fn f32_matvec_matches_matvec_rows_bitexact() {
4400        let (rows, cols) = (300, 40);
4401        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
4402        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
4403        let qt = QTensor::from_f32(w.clone(), rows, cols);
4404
4405        let mut a = vec![0.0f32; rows];
4406        matvec_rows(None, &w, &x, &mut a);
4407        let mut b = vec![0.0f32; rows];
4408        qt.matvec(&x, &mut b, None);
4409        assert_eq!(a, b);
4410    }
4411
4412    #[test]
4413    fn sdot_kernel_exact_on_grid() {
4414        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
4415        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
4416        // exact f32 dot to float rounding. This isolates kernel
4417        // correctness from quantization noise.
4418        eprintln!("sdot_enabled = {}", sdot_enabled());
4419        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
4420        let w: Vec<u8> = (0..rows * cols)
4421            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4422            .collect();
4423        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
4424        let x: Vec<f32> = (0..cols)
4425            .map(|i| match i % 3 {
4426                0 => 1.0,
4427                1 => -1.0,
4428                _ => 0.0,
4429            })
4430            .collect();
4431        let mut a = vec![0.0f32; rows];
4432        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4433        for o in 0..rows {
4434            let mut acc = 0.0f32;
4435            for j in 0..cols {
4436                acc += (w[o * cols + j] as i8) as f32 * x[j];
4437            }
4438            let expect = acc * scales[o];
4439            assert!(
4440                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
4441                "row {o}: {} vs {expect}",
4442                a[o]
4443            );
4444        }
4445    }
4446
4447    #[test]
4448    fn q1_kernels_match_exact_reference() {
4449        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
4450        let (rows, cols) = (7, 96);
4451        let gpr = cols / GROUP_SIZE;
4452        let mut bytes = Vec::new();
4453        for t in 0..rows * gpr {
4454            let s = 0.01 + (t % 13) as f32 * 0.003;
4455            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4456            for j in 0..4 {
4457                bytes.push(((t * 31 + j * 97) % 251) as u8);
4458            }
4459        }
4460        // On-grid activations (±1, amax 1) → the SDOT path is exact.
4461        let x: Vec<f32> = (0..cols)
4462            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
4463            .collect();
4464        // Reference through the core dequant.
4465        let mut w = vec![0.0f32; rows * cols];
4466        cortiq_core::quant::dequant_q1(&bytes, &mut w);
4467        let mut expect = vec![0.0f32; rows];
4468        for o in 0..rows {
4469            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
4470        }
4471        let mut got = vec![0.0f32; rows];
4472        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
4473        for o in 0..rows {
4474            assert!(
4475                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
4476                "row {o}: {} vs {}",
4477                got[o],
4478                expect[o]
4479            );
4480        }
4481        // Pair and batch paths agree with the single path.
4482        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
4483        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
4484        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
4485        assert_eq!(a1, got);
4486        let mut xs = x.clone();
4487        xs.extend_from_slice(&x2);
4488        let mut mm = vec![0.0f32; 2 * rows];
4489        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
4490        assert_eq!(&mm[..rows], got.as_slice());
4491        assert_eq!(&mm[rows..], a2.as_slice());
4492    }
4493
4494    #[test]
4495    fn repack_is_bit_identical() {
4496        // The interleaved-repack kernel must produce EXACTLY the same
4497        // bits as the mmap-layout kernel: integer accumulation is order-
4498        // exact, the f32 epilogue is identical. Odd rows exercise the
4499        // tail; direct range calls exercise unaligned pool splits.
4500        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
4501        let w: Vec<u8> = (0..rows * cols)
4502            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
4503            .collect();
4504        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
4505        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
4506        let rep = q8_repack_layout(&w, rows, cols);
4507        // Group interleave round-trips.
4508        for g in 0..rows / 4 {
4509            for c in 0..cols / 16 {
4510                for lane in 0..4 {
4511                    assert_eq!(
4512                        &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
4513                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
4514                    );
4515                }
4516            }
4517        }
4518        let mut a = vec![0.0f32; rows];
4519        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4520        let mut b = vec![0.0f32; rows];
4521        qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
4522        assert_eq!(a, b, "full-range repack output diverged");
4523
4524        #[cfg(target_arch = "aarch64")]
4525        if sdot_enabled() {
4526            // Unaligned range split (pool workers get arbitrary bounds).
4527            let act = split_act(&x);
4528            let mut c1 = vec![0.0f32; rows];
4529            let mut c2 = vec![0.0f32; rows];
4530            q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
4531            q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
4532            assert_eq!(c1, c2, "unaligned-range repack output diverged");
4533        }
4534    }
4535
4536    #[test]
4537    fn sdot_a8w8_noise_is_bounded() {
4538        // Off-grid activations: A8 quantization noise must stay small in
4539        // relative L2 over the whole output (realistic accuracy contract;
4540        // vmfcore measured argmax-identical decode on real models).
4541        let (rows, cols) = (16, 512);
4542        let w: Vec<u8> = (0..rows * cols)
4543            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4544            .collect();
4545        let scales = vec![0.01f32; rows];
4546        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
4547        let mut a = vec![0.0f32; rows];
4548        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4549        let (mut num, mut den) = (0f64, 0f64);
4550        for o in 0..rows {
4551            let mut acc = 0.0f32;
4552            for j in 0..cols {
4553                acc += (w[o * cols + j] as i8) as f32 * x[j];
4554            }
4555            let expect = acc * scales[o];
4556            num += ((a[o] - expect) as f64).powi(2);
4557            den += (expect as f64).powi(2);
4558        }
4559        let rel = (num / den.max(1e-12)).sqrt();
4560        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
4561    }
4562
4563    #[test]
4564    fn i8_dot_neon_matches_scalar() {
4565        let n = 100;
4566        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
4567        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
4568        let mut scalar = 0.0f32;
4569        for j in 0..n {
4570            scalar += (w[j] as i8) as f32 * x[j];
4571        }
4572        let fast = dot_i8_f32(&w, &x);
4573        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
4574    }
4575
4576    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
4577    #[test]
4578    fn vbitmatvec_matches_full_dequant() {
4579        let (rows, cols) = (6, 64);
4580        let ng = cols / GROUP_SIZE;
4581        // Hand-craft: bits per row, f16 scales, packed rows.
4582        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4583        let mut bytes = bits.clone();
4584        for g in 0..rows * ng {
4585            let s = 0.02 + 0.001 * g as f32;
4586            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4587        }
4588        for r in 0..rows {
4589            let b = bits[r] as usize;
4590            let (mut acc, mut nb) = (0u64, 0usize);
4591            let mut rowbytes = Vec::new();
4592            for i in 0..cols {
4593                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4594                acc = (acc << b) | v;
4595                nb += b;
4596                while nb >= 8 {
4597                    nb -= 8;
4598                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4599                }
4600            }
4601            if nb > 0 {
4602                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4603            }
4604            bytes.extend_from_slice(&rowbytes);
4605        }
4606        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4607
4608        let mut reference = vec![0f32; rows * cols];
4609        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
4610        let mut expect = vec![0f32; rows];
4611        for r in 0..rows {
4612            expect[r] = reference[r * cols..(r + 1) * cols]
4613                .iter()
4614                .zip(&x)
4615                .map(|(w, xv)| w * xv)
4616                .sum();
4617        }
4618        let mut got = vec![0f32; rows];
4619        let offsets = vbit_row_offsets(&bytes, rows, cols);
4620        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
4621        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
4622        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
4623        // the golden-parity gate).
4624        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4625        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
4626        for r in 0..rows {
4627            assert!(
4628                (got[r] - expect[r]).abs() < tol * scale,
4629                "row {r}: {} vs {}",
4630                got[r],
4631                expect[r]
4632            );
4633        }
4634    }
4635
4636    /// Fused q4 matvec must match the reference full-dequant + dense
4637    /// matvec bit-for-bit in structure (same f32 math, group order).
4638    #[test]
4639    fn q4matvec_matches_full_dequant() {
4640        let (rows, cols) = (8, 64);
4641        let groups = rows * cols / GROUP_SIZE;
4642        // Hand-craft a q4_block blob: nibbles then f16 scales.
4643        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4644        for i in 0..groups * 16 {
4645            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4646        }
4647        for g in 0..groups {
4648            let s = 0.01 + 0.003 * g as f32;
4649            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4650        }
4651        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4652
4653        let mut reference = vec![0.0f32; rows * cols];
4654        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
4655        let mut expect = vec![0.0f32; rows];
4656        for r in 0..rows {
4657            expect[r] = reference[r * cols..(r + 1) * cols]
4658                .iter()
4659                .zip(&x)
4660                .map(|(w, xv)| w * xv)
4661                .sum();
4662        }
4663
4664        let mut got = vec![0.0f32; rows];
4665        q4matvec(&bytes, &x, rows, cols, &mut got, None);
4666        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
4667        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
4668        // in the golden-parity gate).
4669        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4670        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
4671        for r in 0..rows {
4672            assert!(
4673                (got[r] - expect[r]).abs() < tol * scale,
4674                "row {r}: {} vs {}",
4675                got[r],
4676                expect[r]
4677            );
4678        }
4679    }
4680
4681    /// Fused two-input vbit matvec must equal two single matvecs exactly
4682    /// (same per-lane accumulation order on both scalar and SDOT paths).
4683    #[test]
4684    fn vbitmatvec2_equals_two_singles() {
4685        let (rows, cols) = (6, 64);
4686        let ng = cols / GROUP_SIZE;
4687        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4688        let mut bytes = bits.clone();
4689        for g in 0..rows * ng {
4690            let s = 0.02 + 0.001 * g as f32;
4691            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4692        }
4693        for r in 0..rows {
4694            let b = bits[r] as usize;
4695            let (mut acc, mut nb) = (0u64, 0usize);
4696            let mut rowbytes = Vec::new();
4697            for i in 0..cols {
4698                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4699                acc = (acc << b) | v;
4700                nb += b;
4701                while nb >= 8 {
4702                    nb -= 8;
4703                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4704                }
4705            }
4706            if nb > 0 {
4707                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4708            }
4709            bytes.extend_from_slice(&rowbytes);
4710        }
4711        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4712        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
4713        let offsets = vbit_row_offsets(&bytes, rows, cols);
4714
4715        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4716        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
4717        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
4718        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
4719        vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
4720        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
4721        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
4722    }
4723
4724    /// Fused two-input q4 matvec must equal two single matvecs exactly.
4725    #[test]
4726    fn q4matvec2_equals_two_singles() {
4727        let (rows, cols) = (8, 128);
4728        let groups = rows * cols / GROUP_SIZE;
4729        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4730        for i in 0..groups * 16 {
4731            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4732        }
4733        for g in 0..groups {
4734            let s = 0.01 + 0.003 * g as f32;
4735            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4736        }
4737        // Include an outlier channel so the SDOT correction path is
4738        // exercised in the pair kernel too.
4739        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4740        x1[9] = 250.0;
4741        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
4742
4743        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4744        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
4745        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
4746        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
4747        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
4748        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
4749        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
4750    }
4751
4752    /// Multi-matrix job must equal separate matvecs exactly — same
4753    /// kernels, only the dispatch is fused.
4754    #[test]
4755    fn matvec_many_equals_separate_matvecs() {
4756        use crate::pool::Pool;
4757        let (r1, r2, cols) = (300, 200, 64);
4758        let mk = |salt: usize, rows: usize| {
4759            QTensor::from_f32(
4760                (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
4761                rows,
4762                cols,
4763            )
4764        };
4765        let (a, b) = (mk(1, r1), mk(5, r2));
4766        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
4767        let pool = Pool::new(3);
4768
4769        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
4770        a.matvec(&x, &mut ea, Some(&pool));
4771        b.matvec(&x, &mut eb, Some(&pool));
4772        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
4773        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
4774        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
4775        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
4776    }
4777
4778    /// Batched q4/vbit matmat must equal per-position matvec calls
4779    /// exactly (the fallback it replaced) — same kernels, same order.
4780    #[test]
4781    fn batched_matmat_equals_per_position_matvec() {
4782        let (rows, cols, b) = (8, 64, 5);
4783        // q4 blob.
4784        let groups = rows * cols / GROUP_SIZE;
4785        let mut q4 = Vec::new();
4786        for i in 0..groups * 16 {
4787            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4788        }
4789        for g in 0..groups {
4790            q4.extend_from_slice(
4791                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
4792            );
4793        }
4794        // vbit blob (mixed widths incl. 8).
4795        let ng = cols / GROUP_SIZE;
4796        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
4797        let mut vb = bits.clone();
4798        for g in 0..rows * ng {
4799            vb.extend_from_slice(
4800                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
4801            );
4802        }
4803        for r in 0..rows {
4804            let bw = bits[r] as usize;
4805            let (mut acc, mut nb) = (0u64, 0usize);
4806            let mut rowbytes = Vec::new();
4807            for i in 0..cols {
4808                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
4809                acc = (acc << bw) | v;
4810                nb += bw;
4811                while nb >= 8 {
4812                    nb -= 8;
4813                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4814                }
4815            }
4816            if nb > 0 {
4817                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4818            }
4819            vb.extend_from_slice(&rowbytes);
4820        }
4821        let offsets = vbit_row_offsets(&vb, rows, cols);
4822
4823        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
4824
4825        // q4: batch vs singles.
4826        let mut got = vec![0f32; b * rows];
4827        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
4828        for bi in 0..b {
4829            let mut expect = vec![0f32; rows];
4830            q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
4831            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
4832        }
4833
4834        // vbit: batch vs singles.
4835        let mut got = vec![0f32; b * rows];
4836        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
4837        for bi in 0..b {
4838            let mut expect = vec![0f32; rows];
4839            vbitmatvec(
4840                &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
4841            );
4842            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
4843        }
4844    }
4845
4846    /// q4_tiled kernels must produce BIT-identical outputs to the q4
4847    /// split kernels on the same values (same ints, same order — only
4848    /// the byte placement differs).
4849    #[test]
4850    fn q4_tiled_matches_q4_block_bitexact() {
4851        let (rows, cols, b) = (8usize, 128usize, 3usize);
4852        let groups = rows * cols / GROUP_SIZE;
4853        let mut split = Vec::with_capacity(groups * 18);
4854        for i in 0..groups * 16 {
4855            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4856        }
4857        for g in 0..groups {
4858            split.extend_from_slice(
4859                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
4860            );
4861        }
4862        // Re-tile: [scale][nibbles] per group.
4863        let (packed, scales) = split.split_at(groups * 16);
4864        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
4865        for g in 0..groups {
4866            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
4867            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
4868        }
4869
4870        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4871        x1[9] = 250.0; // exercise the outlier path
4872        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
4873
4874        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
4875        q4matvec(&split, &x1, rows, cols, &mut a, None);
4876        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
4877        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
4878
4879        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4880        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
4881        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
4882        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
4883        assert_eq!(a1, t1);
4884        assert_eq!(a2, t2);
4885
4886        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
4887        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
4888        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
4889        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
4890        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
4891    }
4892
4893    /// q4 SDOT outlier correction: a single huge activation channel
4894    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
4895    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
4896    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
4897    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
4898    /// can never qualify (8² = n).
4899    #[test]
4900    fn q4matvec_sdot_outlier_exact() {
4901        let (rows, cols) = (4, 128);
4902        let groups = rows * cols / GROUP_SIZE;
4903        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4904        for i in 0..groups * 16 {
4905            bytes.push(((i * 11 + 5) % 256) as u8);
4906        }
4907        for g in 0..groups {
4908            let s = 0.02 + 0.002 * g as f32;
4909            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4910        }
4911        let mut x: Vec<f32> = (0..cols)
4912            .map(|i| match i % 3 {
4913                0 => 1.0,
4914                1 => -1.0,
4915                _ => 0.0,
4916            })
4917            .collect();
4918        x[17] = 300.0; // ≫ 8·rms → outlier channel
4919
4920        let mut reference = vec![0.0f32; rows * cols];
4921        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
4922        let mut expect = vec![0.0f32; rows];
4923        for r in 0..rows {
4924            expect[r] = reference[r * cols..(r + 1) * cols]
4925                .iter()
4926                .zip(&x)
4927                .map(|(w, xv)| w * xv)
4928                .sum();
4929        }
4930        let mut got = vec![0.0f32; rows];
4931        q4matvec(&bytes, &x, rows, cols, &mut got, None);
4932        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
4933        for r in 0..rows {
4934            assert!(
4935                (got[r] - expect[r]).abs() < 2e-3 * scale,
4936                "row {r}: {} vs {} (outlier term must be exact)",
4937                got[r],
4938                expect[r]
4939            );
4940        }
4941    }
4942}