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(cfg!(target_os = "android"))
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            // q1t (ternary + outlier overlay): fused per-row dequant kernel
244            // reads straight from mmap — a 12B q1t stays ~its file size in
245            // RAM instead of dequantizing to ~48 GB of f32.
246            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
247                model: model.clone(),
248                idx,
249                dtype: entry.dtype,
250                rows,
251                cols,
252                row_scale: Vec::new(),
253                col_field: Vec::new(),
254                vbit_offsets: Vec::new(),
255                repack: Vec::new(),
256            }),
257            // No fused kernel yet → dequantize once (correct, more RAM).
258            _ => {
259                let mut data = vec![0.0f32; rows * cols];
260                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
261                Ok(Self::from_f32(data, rows, cols))
262            }
263        }
264    }
265
266    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
267    /// compute-bound, so offload pays at much smaller shapes than q8.)
268    pub(crate) fn is_q1(&self) -> bool {
269        matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
270    }
271
272    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
273    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
274    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
275        match self {
276            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
277            _ => None,
278        }
279    }
280
281    /// (directory idx, rows, cols) of a q1-mapped tensor — the
282    /// whole-block GPU path resolves offsets itself.
283    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
284    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
285    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
286    /// Named `q1_parts` for historical reasons.
287    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
288        match self {
289            Self::Mapped {
290                idx,
291                dtype: TensorDtype::Q1 | TensorDtype::Q1T | TensorDtype::Q4Block | TensorDtype::Q8Row,
292                rows,
293                cols,
294                ..
295            } => Some((*idx, *rows, *cols)),
296            _ => None,
297        }
298    }
299
300    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
301    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
302    /// q8_2f is excluded on purpose: its column field would need a
303    /// prescale stage on the device.
304    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
305        match self {
306            Self::Mapped {
307                idx,
308                dtype: TensorDtype::Q8Row,
309                rows,
310                cols,
311                row_scale,
312                col_field,
313                ..
314            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
315            _ => None,
316        }
317    }
318
319    pub fn rows(&self) -> usize {
320        match self {
321            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
322        }
323    }
324
325    pub fn cols(&self) -> usize {
326        match self {
327            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
328        }
329    }
330
331    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
332    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
333    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
334        match self {
335            Self::Mapped { model, idx, dtype: TensorDtype::Q1, .. } => Some((model, *idx)),
336            _ => None,
337        }
338    }
339
340    /// (model, idx, kind, row_scale) for a graph-capable mapped weight. kind:
341    /// 0=q8_row (per-row scales), 1=q1, 2=q4_tiled, 3=q1t (tile-embedded, no
342    /// rs). None for dtypes the token graph does not handle (q8_2f/q4_block/vbit).
343    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
344        match self {
345            Self::Mapped { model, idx, dtype: TensorDtype::Q8Row, row_scale, .. } => {
346                Some((model, *idx, 0, row_scale.as_slice()))
347            }
348            Self::Mapped { model, idx, dtype: TensorDtype::Q1, .. } => Some((model, *idx, 1, &[])),
349            Self::Mapped { model, idx, dtype: TensorDtype::Q4Tiled, .. } => Some((model, *idx, 2, &[])),
350            Self::Mapped { model, idx, dtype: TensorDtype::Q1T, .. } => Some((model, *idx, 3, &[])),
351            _ => None,
352        }
353    }
354
355    /// Dense f32 view — only for owned tensors. Masked/sparse execution
356    /// paths require it; quantized weights don't support masks yet.
357    pub fn as_f32(&self) -> Option<&[f32]> {
358        match self {
359            Self::F32 { data, .. } => Some(data),
360            Self::Mapped { .. } => None,
361        }
362    }
363
364    fn quant_bytes(&self) -> &[u8] {
365        match self {
366            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
367            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
368        }
369    }
370
371    /// Dequantize one row into `dst` (embedding lookup).
372    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
373        let cols = self.cols();
374        debug_assert_eq!(dst.len(), cols);
375        match self {
376            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
377            Self::Mapped {
378                dtype,
379                row_scale,
380                col_field,
381                vbit_offsets,
382                ..
383            } => {
384                if *dtype == TensorDtype::Q4Tiled {
385                    let bytes = self.quant_bytes();
386                    let gpr = cols / GROUP_SIZE;
387                    for gi in 0..gpr {
388                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
389                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
390                        for (k, &b) in tile[2..].iter().enumerate() {
391                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
392                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
393                        }
394                    }
395                    return;
396                }
397                if *dtype == TensorDtype::Q4Block {
398                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
399                    let gpr = cols / GROUP_SIZE;
400                    for gi in 0..gpr {
401                        let g = r * gpr + gi;
402                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
403                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
404                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
405                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
406                        }
407                    }
408                    return;
409                }
410                if *dtype == TensorDtype::Q1 {
411                    let bytes = self.quant_bytes();
412                    let gpr = cols / GROUP_SIZE;
413                    for gi in 0..gpr {
414                        let tile =
415                            &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
416                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
417                        for (j, &b) in tile[2..].iter().enumerate() {
418                            for k in 0..8 {
419                                dst[gi * GROUP_SIZE + j * 8 + k] =
420                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
421                            }
422                        }
423                    }
424                    return;
425                }
426                if *dtype == TensorDtype::Q1T {
427                    let bytes = self.quant_bytes();
428                    let gpr = cols / GROUP_SIZE;
429                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
430                    for gi in 0..gpr {
431                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
432                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
433                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
434                        for k in 0..GROUP_SIZE {
435                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k) {
436                                1 => s,
437                                2 => -s,
438                                _ => 0.0,
439                            };
440                        }
441                    }
442                    // Overlay
443                    let rows = self.rows();
444                    let entries = base_len + (rows + 1) * 4;
445                    if entries <= bytes.len() {
446                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
447                        let r0 = u32::from_le_bytes([ptrs[r * 4], ptrs[r * 4 + 1], ptrs[r * 4 + 2], ptrs[r * 4 + 3]]) as usize;
448                        let r1 = u32::from_le_bytes([ptrs[(r + 1) * 4], ptrs[(r + 1) * 4 + 1], ptrs[(r + 1) * 4 + 2], ptrs[(r + 1) * 4 + 3]]) as usize;
449                        let off = entries + r0 * 4;
450                        for i in 0..r1 - r0 {
451                            let item = &bytes[off + i * 4..off + i * 4 + 4];
452                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
453                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([item[2], item[3]]));
454                            if c < cols {
455                                dst[c] = v;
456                            }
457                        }
458                    }
459                    return;
460                }
461                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
462                    let bytes = self.quant_bytes();
463                    let rows = self.rows();
464                    let ng = cols / GROUP_SIZE;
465                    let bits = &bytes[..rows];
466                    let sc_off = rows;
467                    // Precomputed at load — embedding lookup used to scan
468                    // the bit-widths of every preceding row (O(token_id)).
469                    let off = vbit_offsets[r];
470                    let b = bits[r] as usize;
471                    let l = ((1usize << (b - 1)) - 1) as f32;
472                    let data = &bytes[off..];
473                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
474                    for (i, d) in dst.iter_mut().enumerate() {
475                        while nbits < b {
476                            acc = (acc << 8) | data[idx] as u64;
477                            idx += 1;
478                            nbits += 8;
479                        }
480                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
481                        nbits -= b;
482                        let so = (r * ng + i / GROUP_SIZE) * 2;
483                        let sv = f16_to_f32(u16::from_le_bytes([
484                            bytes[sc_off + so],
485                            bytes[sc_off + so + 1],
486                        ]));
487                        *d = (u - l) * sv;
488                    }
489                    return;
490                }
491                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
492                let s = row_scale[r];
493                match dtype {
494                    TensorDtype::Q8Row => {
495                        for (d, &b) in dst.iter_mut().zip(q) {
496                            *d = (b as i8) as f32 * s;
497                        }
498                    }
499                    TensorDtype::Q8_2f => {
500                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
501                            *d = (b as i8) as f32 * s * col_field[i];
502                        }
503                    }
504                    _ => unreachable!(),
505                }
506            }
507        }
508    }
509
510    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
511    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
512    /// false for group-packed q4/vbit (column access would unpack whole
513    /// groups — sparse execution falls back to f32 for those).
514    pub fn sparse_col_ok(&self) -> bool {
515        match self {
516            Self::F32 { .. } => true,
517            Self::Mapped { dtype, .. } => {
518                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
519            }
520        }
521    }
522
523    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
524    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
525    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
526    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
527        let inter = self.cols();
528        let hidden = self.rows();
529        debug_assert_eq!(out.len(), hidden);
530        match self {
531            Self::F32 { data, .. } => {
532                for (k, o) in out.iter_mut().enumerate() {
533                    *o += w * data[k * inter + c];
534                }
535            }
536            Self::Mapped {
537                dtype, row_scale, col_field, ..
538            } => {
539                let q = self.quant_bytes();
540                let colf = if *dtype == TensorDtype::Q8_2f {
541                    col_field[c]
542                } else {
543                    1.0
544                };
545                let wc = w * colf;
546                for (k, o) in out.iter_mut().enumerate() {
547                    let b = q[k * inter + c] as i8 as f32;
548                    *o += wc * b * row_scale[k];
549                }
550            }
551        }
552    }
553
554    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
555    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
556    /// into `scratch` first (rare for active-FFN weights).
557    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
558        let cols = self.cols();
559        match self {
560            Self::F32 { data, .. } => {
561                let row = &data[r * cols..(r + 1) * cols];
562                row.iter().zip(x).map(|(w, v)| w * v).sum()
563            }
564            Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
565                TensorDtype::Q8Row => {
566                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
567                    dot_i8_f32(q, x) * row_scale[r]
568                }
569                TensorDtype::Q8_2f => {
570                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
571                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
572                }
573                _ => {
574                    self.row_f32(r, scratch);
575                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
576                }
577            },
578        }
579    }
580
581    /// `out = W · x` (row-major). F32 delegates to the historical
582    /// bit-exact path; Mapped runs the fused int8 kernel.
583    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
584        match self {
585            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
586            Self::Mapped {
587                model,
588                idx,
589                dtype,
590                rows,
591                cols,
592                row_scale,
593                col_field,
594                vbit_offsets,
595                repack,
596            } => {
597                let _ = (model, idx);
598                if *dtype == TensorDtype::Q4Block {
599                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
600                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
601                    // the winner; Metal returns false → the CPU kernel below.
602                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
603                        let t0 = std::time::Instant::now();
604                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
605                            crate::gpu::ProbeArm::Gpu => {
606                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
607                                    crate::gpu::probe_record(
608                                        crate::gpu::OpClass::Matvec,
609                                        true,
610                                        t0.elapsed(),
611                                    );
612                                    return;
613                                }
614                            }
615                            crate::gpu::ProbeArm::CpuTimed => {
616                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
617                                crate::gpu::probe_record(
618                                    crate::gpu::OpClass::Matvec,
619                                    false,
620                                    t0.elapsed(),
621                                );
622                                return;
623                            }
624                            crate::gpu::ProbeArm::Cpu => {}
625                        }
626                    }
627                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
628                    return;
629                }
630                if *dtype == TensorDtype::Q4Tiled {
631                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
632                    return;
633                }
634                if *dtype == TensorDtype::Q1 {
635                    // GPU route for large q1 matvecs (out_proj / lm_head
636                    // class): the CPU q1 kernel is load-port-bound at
637                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
638                    // probe measures both arms and keeps the winner.
639                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
640                        let t0 = std::time::Instant::now();
641                        let arm = if crate::gpu::q1_force() {
642                            crate::gpu::ProbeArm::Gpu
643                        } else {
644                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
645                        };
646                        match arm {
647                            crate::gpu::ProbeArm::Gpu => {
648                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
649                                    crate::gpu::probe_record(
650                                        crate::gpu::OpClass::Matvec,
651                                        true,
652                                        t0.elapsed(),
653                                    );
654                                    return;
655                                }
656                            }
657                            crate::gpu::ProbeArm::CpuTimed => {
658                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
659                                crate::gpu::probe_record(
660                                    crate::gpu::OpClass::Matvec,
661                                    false,
662                                    t0.elapsed(),
663                                );
664                                return;
665                            }
666                            crate::gpu::ProbeArm::Cpu => {}
667                        }
668                    }
669                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
670                    return;
671                }
672                if *dtype == TensorDtype::Q1T {
673                    // GPU route for large q1t matvecs: the ternary BASE dot runs
674                    // on the GPU (load-port-bound on CPU, like q1), then the
675                    // sparse overlay is added on the CPU. Probe keeps the winner.
676                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
677                        let t0 = std::time::Instant::now();
678                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
679                            crate::gpu::ProbeArm::Gpu => {
680                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
681                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
682                                    crate::gpu::probe_record(
683                                        crate::gpu::OpClass::Matvec,
684                                        true,
685                                        t0.elapsed(),
686                                    );
687                                    return;
688                                }
689                            }
690                            crate::gpu::ProbeArm::CpuTimed => {
691                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
692                                crate::gpu::probe_record(
693                                    crate::gpu::OpClass::Matvec,
694                                    false,
695                                    t0.elapsed(),
696                                );
697                                return;
698                            }
699                            crate::gpu::ProbeArm::Cpu => {}
700                        }
701                    }
702                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
703                    return;
704                }
705                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
706                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
707                    return;
708                }
709                let xs = prescale(x, col_field, *dtype);
710                // D5: large q8 matrices (lm_head-class) — hybrid
711                // CPU∥GPU: split the rows, both sides compute
712                // SIMULTANEOUSLY (same math, shared prescale).
713                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
714                if *rows >= crate::gpu::min_rows()
715                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
716                    && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
717                    && crate::gpu::enabled_here()
718                {
719                    // Runtime probe: alternate the hybrid against the
720                    // pure-CPU matvec, keep whichever is faster HERE.
721                    let t0 = std::time::Instant::now();
722                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
723                        crate::gpu::ProbeArm::Gpu => {}
724                        crate::gpu::ProbeArm::CpuTimed => {
725                            qmatvec(self.quant_bytes(), repack, row_scale, x, col_field, *dtype, *rows, *cols, out, pool);
726                            crate::gpu::probe_record(
727                                crate::gpu::OpClass::Matvec,
728                                false,
729                                t0.elapsed(),
730                            );
731                            return;
732                        }
733                        crate::gpu::ProbeArm::Cpu => {
734                            qmatvec(self.quant_bytes(), repack, row_scale, x, col_field, *dtype, *rows, *cols, out, pool);
735                            return;
736                        }
737                    }
738                    let frac = std::env::var("CMF_GPU_SPLIT")
739                        .ok()
740                        .and_then(|v| v.parse::<f32>().ok())
741                        .unwrap_or(0.5)
742                        .clamp(0.0, 1.0);
743                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
744                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
745                    let bytes = self.quant_bytes();
746                    let ok = std::thread::scope(|sc| {
747                        let g = sc.spawn(|| {
748                            crate::gpu::q8_matvec_range(
749                                model,
750                                *idx,
751                                cpu_rows,
752                                &row_scale[cpu_rows..],
753                                &xs,
754                                *rows - cpu_rows,
755                                *cols,
756                                out_gpu,
757                            )
758                        });
759                        if cpu_rows > 0 {
760                            // Repack prefix covers the full groups of the
761                            // CPU half (the split starts at row 0).
762                            let rep_cpu = if repack.is_empty() {
763                                &[][..]
764                            } else {
765                                &repack[..(cpu_rows / 4) * 4 * *cols]
766                            };
767                            qmatvec(
768                                &bytes[..cpu_rows * *cols],
769                                rep_cpu,
770                                &row_scale[..cpu_rows],
771                                x,
772                                col_field,
773                                *dtype,
774                                cpu_rows,
775                                *cols,
776                                out_cpu,
777                                pool,
778                            );
779                        }
780                        g.join().unwrap_or(false)
781                    });
782                    if ok {
783                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
784                        return;
785                    }
786                    // GPU failed — CPU finishes its half (rows rebased —
787                    // group offsets don't line up, mmap layout only).
788                    qmatvec(
789                        &bytes[cpu_rows * *cols..(*rows) * *cols],
790                        &[],
791                        &row_scale[cpu_rows..],
792                        x,
793                        col_field,
794                        *dtype,
795                        *rows - cpu_rows,
796                        *cols,
797                        out_gpu,
798                        pool,
799                    );
800                    return;
801                }
802                qmatvec(self.quant_bytes(), repack, row_scale, x, col_field, *dtype, *rows, *cols, out, pool);
803            }
804        }
805    }
806
807    /// Fused two-input matvec (MTP verify pair): weights streamed once.
808    pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
809        match self {
810            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
811            Self::Mapped {
812                dtype,
813                rows,
814                cols,
815                row_scale,
816                col_field,
817                vbit_offsets,
818                ..
819            } => {
820                if *dtype == TensorDtype::Q4Block {
821                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
822                    return;
823                }
824                if *dtype == TensorDtype::Q4Tiled {
825                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
826                    return;
827                }
828                if *dtype == TensorDtype::Q1 {
829                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
830                    return;
831                }
832                if *dtype == TensorDtype::Q1T {
833                    // No fused ternary pair kernel; the two passes are still
834                    // the fused decode+dot each. (Q1T lacks a row_scale array —
835                    // scales live inline in the tiles — so it must not fall
836                    // through to the q8 qmatvec2 below.)
837                    q1t_matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
838                    q1t_matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
839                    return;
840                }
841                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
842                    vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
843                    return;
844                }
845                qmatvec2(self.quant_bytes(), row_scale, x1, x2, col_field, *dtype, *rows, *cols, o1, o2, pool);
846            }
847        }
848    }
849}
850
851impl QTensor {
852    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
853    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
854    /// to b matvec calls (same dot kernels in the same order); the win —
855    /// the weight row streams from DRAM once per batch, not b times.
856    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
857        let cols = self.cols();
858        let rows = self.rows();
859        debug_assert_eq!(xs_all.len(), b * cols);
860        debug_assert_eq!(out.len(), b * rows);
861        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
862        // Mapped tensors carry a directory name; the check is a relaxed
863        // atomic load, free when not calibrating.
864        if crate::gptq_capture::capturing() {
865            if let Self::Mapped { model, idx, .. } = self {
866                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
867            }
868        }
869        match self {
870            Self::F32 { data, .. } => {
871                let out_addr = SendMut(out.as_mut_ptr());
872                let run = |start: usize, end: usize| {
873                    for o in start..end {
874                        let row = &data[o * cols..(o + 1) * cols];
875                        for bi in 0..b {
876                            let x = &xs_all[bi * cols..(bi + 1) * cols];
877                            let mut acc = 0f32;
878                            for j in 0..cols {
879                                acc += row[j] * x[j];
880                            }
881                            unsafe { *out_addr.at(bi * rows + o) = acc };
882                        }
883                    }
884                };
885                dispatch_rows(pool, rows, &run);
886            }
887            Self::Mapped {
888                dtype,
889                row_scale,
890                col_field,
891                vbit_offsets,
892                ..
893            } => {
894                if *dtype == TensorDtype::Q4Block {
895                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
896                    return;
897                }
898                if *dtype == TensorDtype::Q4Tiled {
899                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
900                    return;
901                }
902                if *dtype == TensorDtype::Q1 {
903                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
904                    // device); the probe keeps whichever beats the CPU matmat.
905                    if b >= 32
906                        && b * rows * cols >= 128_000_000
907                        && cols % 64 == 0
908                        && crate::gpu::enabled_here()
909                    {
910                        if let Self::Mapped { model, idx, .. } = self {
911                            let t0 = std::time::Instant::now();
912                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
913                                crate::gpu::ProbeArm::Gpu => {
914                                    if crate::gpu::q1_matmat(model, *idx, xs_all, b, rows, cols, out) {
915                                        crate::gpu::probe_record(
916                                            crate::gpu::OpClass::Matmat,
917                                            true,
918                                            t0.elapsed(),
919                                        );
920                                        return;
921                                    }
922                                }
923                                crate::gpu::ProbeArm::CpuTimed => {
924                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
925                                    crate::gpu::probe_record(
926                                        crate::gpu::OpClass::Matmat,
927                                        false,
928                                        t0.elapsed(),
929                                    );
930                                    return;
931                                }
932                                crate::gpu::ProbeArm::Cpu => {}
933                            }
934                        }
935                    }
936                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
937                    return;
938                }
939                if *dtype == TensorDtype::Q1T {
940                    // GPU batched GEMM for wide prefill (base + overlay on the
941                    // device); probe keeps the winner vs the CPU matmat.
942                    if b >= 32
943                        && b * rows * cols >= 128_000_000
944                        && crate::gpu::enabled_here()
945                    {
946                        if let Self::Mapped { model, idx, .. } = self {
947                            let t0 = std::time::Instant::now();
948                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
949                                crate::gpu::ProbeArm::Gpu => {
950                                    if crate::gpu::q1t_matmat(model, *idx, xs_all, b, rows, cols, out) {
951                                        crate::gpu::probe_record(
952                                            crate::gpu::OpClass::Matmat,
953                                            true,
954                                            t0.elapsed(),
955                                        );
956                                        return;
957                                    }
958                                }
959                                crate::gpu::ProbeArm::CpuTimed => {
960                                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
961                                    crate::gpu::probe_record(
962                                        crate::gpu::OpClass::Matmat,
963                                        false,
964                                        t0.elapsed(),
965                                    );
966                                    return;
967                                }
968                                crate::gpu::ProbeArm::Cpu => {}
969                            }
970                        }
971                    }
972                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
973                    return;
974                }
975                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
976                    vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
977                    return;
978                }
979                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
980                    .map(|bi| {
981                        prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
982                    })
983                    .collect();
984                // D5: large prefill-batch GEMMs — on the GPU (threshold by
985                // work volume: submission carries b×rows×cols MACs).
986                // Runtime probe: the naive GEMM shader + sync readback
987                // lose to the CPU GEMM on slow driver stacks — alternate
988                // both arms and keep the winner.
989                if b >= 8
990                    && b * rows * cols >= 128_000_000
991                    && crate::gpu::enabled_here()
992                {
993                    if let Self::Mapped { model, idx, .. } = self {
994                        let t0 = std::time::Instant::now();
995                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
996                            crate::gpu::ProbeArm::Gpu
997                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
998                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
999                            {
1000                                // Cold weights during probing: the upload
1001                                // has started, the count runs on the CPU —
1002                                // the GPU arm samples on the next touch.
1003                                let q = self.quant_bytes();
1004                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1005                                return;
1006                            }
1007                            crate::gpu::ProbeArm::Gpu => {
1008                                let flat: Vec<f32> =
1009                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1010                                if crate::gpu::q8_matmat(
1011                                    model, *idx, row_scale, &flat, b, rows, cols, out)
1012                                {
1013                                    crate::gpu::probe_record(
1014                                        crate::gpu::OpClass::Matmat,
1015                                        true,
1016                                        t0.elapsed(),
1017                                    );
1018                                    return;
1019                                }
1020                            }
1021                            crate::gpu::ProbeArm::CpuTimed => {
1022                                let q = self.quant_bytes();
1023                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1024                                crate::gpu::probe_record(
1025                                    crate::gpu::OpClass::Matmat,
1026                                    false,
1027                                    t0.elapsed(),
1028                                );
1029                                return;
1030                            }
1031                            crate::gpu::ProbeArm::Cpu => {}
1032                        }
1033                    }
1034                }
1035                let q = self.quant_bytes();
1036                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1037            }
1038        }
1039    }
1040}
1041
1042impl QTensor {
1043    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1044    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1045    /// barrier instead of N. Per-row math is the exact same kernel as
1046    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1047    /// Falls back to N sequential matvecs when the set is not a uniform
1048    /// q8-family/F32 group or there is no pool.
1049    pub fn matvec_many<const N: usize>(
1050        ts: [&QTensor; N],
1051        x: &[f32],
1052        mut outs: [&mut [f32]; N],
1053        pool: Option<&Pool>,
1054    ) {
1055        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1056        let uniform_q8 = ts.iter().all(|t| {
1057            matches!(
1058                t,
1059                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1060            )
1061        });
1062        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1063        let uniform_q4 = ts
1064            .iter()
1065            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1066        let uniform_vbit = ts
1067            .iter()
1068            .all(|t| matches!(
1069                t,
1070                Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1071            ));
1072        let uniform_q1 = ts
1073            .iter()
1074            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
1075        let Some(pool) = pool else {
1076            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1077                t.matvec(x, o, None);
1078            }
1079            return;
1080        };
1081        if total_rows < 256
1082            || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
1083        {
1084            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1085                t.matvec(x, o, Some(pool));
1086            }
1087            return;
1088        }
1089
1090        if uniform_q1 {
1091            // One shared activation split + group sums (q1 has no col
1092            // field; the same input feeds every tensor).
1093            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1094            if a8w8_enabled() {
1095                let act = split_act(x);
1096                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1097                let (act, gsum) = (&act, &gsum);
1098                let closures: [_; N] = std::array::from_fn(|i| {
1099                    let (bytes, gpr, out) =
1100                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1101                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1102                });
1103                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1104                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1105                pool.run_many(&parts);
1106            } else {
1107                let closures: [_; N] = std::array::from_fn(|i| {
1108                    let (bytes, gpr, out) =
1109                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1110                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1111                });
1112                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1113                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1114                pool.run_many(&parts);
1115            }
1116            return;
1117        }
1118
1119        if uniform_q4 || uniform_vbit {
1120            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1121            // q4/vbit share one activation split — no per-tensor col field.
1122            if a8w8_enabled() {
1123                let act = split_act(x);
1124                let act = &act;
1125                if uniform_q4 {
1126                    let closures: [_; N] = std::array::from_fn(|i| {
1127                        let (packed, scales) =
1128                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1129                        let (gpr, cols, out) =
1130                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1131                        move |s: usize, e: usize| {
1132                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1133                        }
1134                    });
1135                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1136                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1137                    pool.run_many(&parts);
1138                } else {
1139                    let closures: [_; N] = std::array::from_fn(|i| {
1140                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1141                        let (bytes, rows, cols, out) =
1142                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1143                        move |s: usize, e: usize| {
1144                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1145                        }
1146                    });
1147                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1148                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1149                    pool.run_many(&parts);
1150                }
1151                return;
1152            }
1153            if uniform_q4 {
1154                let closures: [_; N] = std::array::from_fn(|i| {
1155                    let (packed, scales) =
1156                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1157                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1158                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1159                });
1160                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1161                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1162                pool.run_many(&parts);
1163            } else {
1164                let closures: [_; N] = std::array::from_fn(|i| {
1165                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1166                    let (bytes, rows, cols, out) =
1167                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1168                    move |s: usize, e: usize| {
1169                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1170                    }
1171                });
1172                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1173                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1174                pool.run_many(&parts);
1175            }
1176            return;
1177        }
1178
1179        if uniform_f32 {
1180            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1181            let closures: [_; N] = std::array::from_fn(|i| {
1182                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1183                let out = outs_addr[i];
1184                move |start: usize, end: usize| {
1185                    for o in start..end {
1186                        let row = &data[o * cols..(o + 1) * cols];
1187                        let mut sum = 0.0f32;
1188                        for j in 0..*cols {
1189                            sum += row[j] * x[j];
1190                        }
1191                        // SAFETY: disjoint (tensor, row) cells per worker.
1192                        unsafe { *out.at(o) = sum };
1193                    }
1194                }
1195            });
1196            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1197                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1198            pool.run_many(&parts);
1199            return;
1200        }
1201
1202        // Uniform q8-family: per-tensor prescale (q8_2f col fields
1203        // differ per tensor) + the shared range kernels.
1204        struct Ctx<'a> {
1205            bytes: &'a [u8],
1206            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1207            rep: &'a [u8],
1208            row_scale: &'a [f32],
1209            cols: usize,
1210            xs: std::borrow::Cow<'a, [f32]>,
1211        }
1212        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1213            let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
1214                unreachable!()
1215            };
1216            Ctx {
1217                bytes: ts[i].quant_bytes(),
1218                rep: repack,
1219                row_scale,
1220                cols: *cols,
1221                xs: prescale(x, col_field, *dtype),
1222            }
1223        });
1224        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1225        #[cfg(target_arch = "aarch64")]
1226        if sdot_enabled() {
1227            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1228            let closures: [_; N] = std::array::from_fn(|i| {
1229                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1230                move |start: usize, end: usize| {
1231                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1232                }
1233            });
1234            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1235                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1236            pool.run_many(&parts);
1237            return;
1238        }
1239        #[cfg(target_arch = "x86_64")]
1240        if avx2_a8w8_enabled() {
1241            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1242            let closures: [_; N] = std::array::from_fn(|i| {
1243                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1244                move |start: usize, end: usize| {
1245                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1246                }
1247            });
1248            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1249                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1250            pool.run_many(&parts);
1251            return;
1252        }
1253        let closures: [_; N] = std::array::from_fn(|i| {
1254            let (c, out) = (&ctxs[i], outs_addr[i]);
1255            move |start: usize, end: usize| {
1256                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1257            }
1258        });
1259        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1260            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1261        pool.run_many(&parts);
1262    }
1263}
1264
1265impl QTensor {
1266    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
1267    /// single pool dispatch — the MTP/pair decode path publishes one job
1268    /// for Q/K/V (and one for gate+up) instead of one per tensor.
1269    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
1270    #[allow(clippy::needless_range_loop)]
1271    pub fn matvec2_many<const N: usize>(
1272        ts: [&QTensor; N],
1273        x1: &[f32],
1274        x2: &[f32],
1275        mut o1s: [&mut [f32]; N],
1276        mut o2s: [&mut [f32]; N],
1277        pool: Option<&Pool>,
1278    ) {
1279        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1280        let uniform_q8 = ts.iter().all(|t| {
1281            matches!(
1282                t,
1283                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1284            )
1285        });
1286        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1287        let uniform_q4 = ts
1288            .iter()
1289            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1290        let uniform_vbit = ts
1291            .iter()
1292            .all(|t| matches!(
1293                t,
1294                Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1295            ));
1296        let fusable = pool.is_some()
1297            && total_rows >= 256
1298            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1299        if !fusable {
1300            for i in 0..N {
1301                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1302            }
1303            return;
1304        }
1305        let pool = pool.unwrap();
1306
1307        if uniform_q4 || uniform_vbit {
1308            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1309            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1310            // q4/vbit share activation splits — no per-tensor col field.
1311            if a8w8_enabled() {
1312                let a1 = split_act(x1);
1313                let a2 = split_act(x2);
1314                let (a1, a2) = (&a1, &a2);
1315                if uniform_q4 {
1316                    let closures: [_; N] = std::array::from_fn(|i| {
1317                        let (packed, scales) =
1318                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1319                        let (gpr, cols, o1, o2) =
1320                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1321                        move |s: usize, e: usize| {
1322                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1323                        }
1324                    });
1325                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1326                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1327                    pool.run_many(&parts);
1328                } else {
1329                    let closures: [_; N] = std::array::from_fn(|i| {
1330                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1331                        let (bytes, rows, cols, o1, o2) =
1332                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1333                        move |s: usize, e: usize| {
1334                            vbit_range2_a8w8(
1335                                bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1336                            )
1337                        }
1338                    });
1339                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1340                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1341                    pool.run_many(&parts);
1342                }
1343                return;
1344            }
1345            if uniform_q4 {
1346                let closures: [_; N] = std::array::from_fn(|i| {
1347                    let (packed, scales) =
1348                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1349                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1350                    move |s: usize, e: usize| {
1351                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1352                    }
1353                });
1354                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1355                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1356                pool.run_many(&parts);
1357            } else {
1358                let closures: [_; N] = std::array::from_fn(|i| {
1359                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1360                    let (bytes, rows, cols, o1, o2) =
1361                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1362                    move |s: usize, e: usize| {
1363                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1364                    }
1365                });
1366                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1367                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1368                pool.run_many(&parts);
1369            }
1370            return;
1371        }
1372
1373        if uniform_f32 {
1374            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1375            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1376            let closures: [_; N] = std::array::from_fn(|i| {
1377                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1378                let (o1, o2) = (p1[i], p2[i]);
1379                move |start: usize, end: usize| {
1380                    for o in start..end {
1381                        let row = &data[o * cols..(o + 1) * cols];
1382                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
1383                        for j in 0..*cols {
1384                            s1 += row[j] * x1[j];
1385                            s2 += row[j] * x2[j];
1386                        }
1387                        // SAFETY: disjoint (tensor, row) cells per worker.
1388                        unsafe {
1389                            *o1.at(o) = s1;
1390                            *o2.at(o) = s2;
1391                        }
1392                    }
1393                }
1394            });
1395            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1396                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1397            pool.run_many(&parts);
1398            return;
1399        }
1400
1401        struct Ctx<'a> {
1402            bytes: &'a [u8],
1403            row_scale: &'a [f32],
1404            cols: usize,
1405            xs1: std::borrow::Cow<'a, [f32]>,
1406            xs2: std::borrow::Cow<'a, [f32]>,
1407        }
1408        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1409            let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1410                unreachable!()
1411            };
1412            Ctx {
1413                bytes: ts[i].quant_bytes(),
1414                row_scale,
1415                cols: *cols,
1416                xs1: prescale(x1, col_field, *dtype),
1417                xs2: prescale(x2, col_field, *dtype),
1418            }
1419        });
1420        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1421        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1422        #[cfg(target_arch = "aarch64")]
1423        if sdot_enabled() {
1424            let acts: [(SplitAct, SplitAct); N] =
1425                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1426            let closures: [_; N] = std::array::from_fn(|i| {
1427                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1428                move |start: usize, end: usize| {
1429                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1430                }
1431            });
1432            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1433                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1434            pool.run_many(&parts);
1435            return;
1436        }
1437        #[cfg(target_arch = "x86_64")]
1438        if avx2_a8w8_enabled() {
1439            let acts: [(SplitAct, SplitAct); N] =
1440                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1441            let closures: [_; N] = std::array::from_fn(|i| {
1442                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1443                move |start: usize, end: usize| {
1444                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1445                }
1446            });
1447            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1448                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1449            pool.run_many(&parts);
1450            return;
1451        }
1452        let closures: [_; N] = std::array::from_fn(|i| {
1453            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1454            move |start: usize, end: usize| {
1455                q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1456            }
1457        });
1458        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1459            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1460        pool.run_many(&parts);
1461    }
1462}
1463
1464/// Batched q8 kernel: same math as qmatvec, the row makes a single
1465/// pass from memory for the whole batch.
1466/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
1467/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
1468#[cfg(target_os = "macos")]
1469mod accel_blas {
1470    #[link(name = "Accelerate", kind = "framework")]
1471    unsafe extern "C" {
1472        pub fn cblas_sgemm(
1473            order: i32,
1474            trans_a: i32,
1475            trans_b: i32,
1476            m: i32,
1477            n: i32,
1478            k: i32,
1479            alpha: f32,
1480            a: *const f32,
1481            lda: i32,
1482            b: *const f32,
1483            ldb: i32,
1484            beta: f32,
1485            c: *mut f32,
1486            ldc: i32,
1487        );
1488    }
1489}
1490
1491#[cfg(target_os = "macos")]
1492pub(crate) fn accel_gemm_enabled() -> bool {
1493    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1494    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1495}
1496
1497/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
1498/// same entry point, so the batched-attention path opens on mobile.
1499#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1500pub(crate) fn accel_gemm_enabled() -> bool {
1501    true
1502}
1503
1504/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
1505/// micro-kernel with A broadcast against B panels — the mobile stand-in
1506/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
1507/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
1508/// k = head_dim or context), and the goal is removing the per-position
1509/// quadratic wall, not peak GEMM.
1510#[cfg(target_arch = "aarch64")]
1511#[allow(clippy::too_many_arguments)]
1512pub(crate) fn neon_gemm_rm(
1513    m: usize,
1514    n: usize,
1515    k: usize,
1516    alpha: f32,
1517    a: &[f32],
1518    lda: usize,
1519    b_mat: &[f32],
1520    ldb: usize,
1521    b_rows_are_n: bool,
1522    c: &mut [f32],
1523    ldc: usize,
1524) {
1525    debug_assert!(a.len() >= (m - 1) * lda + k);
1526    debug_assert!(c.len() >= (m - 1) * ldc + n);
1527    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
1528    unsafe {
1529        use core::arch::aarch64::*;
1530        let mut i = 0usize;
1531        while i < m {
1532            let mi = (m - i).min(4);
1533            let mut j = 0usize;
1534            while j < n {
1535                let nj = (n - j).min(8);
1536                if mi == 4 && nj == 8 {
1537                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1538                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1539                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1540                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1541                    for p in 0..k {
1542                        let (b0, b1) = if b_rows_are_n {
1543                            // B is [n, k]: column p of Bᵀ = element p of
1544                            // eight consecutive B rows — gathered.
1545                            let base = b_mat.as_ptr().add(j * ldb + p);
1546                            let g = |o: usize| *base.add(o * ldb);
1547                            (
1548                                [g(0), g(1), g(2), g(3)],
1549                                [g(4), g(5), g(6), g(7)],
1550                            )
1551                        } else {
1552                            let base = b_mat.as_ptr().add(p * ldb + j);
1553                            (
1554                                [*base, *base.add(1), *base.add(2), *base.add(3)],
1555                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
1556                            )
1557                        };
1558                        let bv0 = vld1q_f32(b0.as_ptr());
1559                        let bv1 = vld1q_f32(b1.as_ptr());
1560                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
1561                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
1562                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
1563                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
1564                        c0a = vfmaq_f32(c0a, a0, bv0);
1565                        c0b = vfmaq_f32(c0b, a0, bv1);
1566                        c1a = vfmaq_f32(c1a, a1, bv0);
1567                        c1b = vfmaq_f32(c1b, a1, bv1);
1568                        c2a = vfmaq_f32(c2a, a2, bv0);
1569                        c2b = vfmaq_f32(c2b, a2, bv1);
1570                        c3a = vfmaq_f32(c3a, a3, bv0);
1571                        c3b = vfmaq_f32(c3b, a3, bv1);
1572                    }
1573                    let al = vdupq_n_f32(alpha);
1574                    for (r, (ca, cb)) in
1575                        [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)].iter().enumerate()
1576                    {
1577                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
1578                        vst1q_f32(dst, vmulq_f32(*ca, al));
1579                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
1580                    }
1581                } else {
1582                    for r in 0..mi {
1583                        for q in 0..nj {
1584                            let mut acc = 0f32;
1585                            for p in 0..k {
1586                                let bv = if b_rows_are_n {
1587                                    b_mat[(j + q) * ldb + p]
1588                                } else {
1589                                    b_mat[p * ldb + j + q]
1590                                };
1591                                acc += a[(i + r) * lda + p] * bv;
1592                            }
1593                            c[(i + r) * ldc + j + q] = acc * alpha;
1594                        }
1595                    }
1596                }
1597                j += nj;
1598            }
1599            i += mi;
1600        }
1601    }
1602}
1603
1604/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
1605#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1606#[allow(clippy::too_many_arguments)]
1607pub(crate) fn sgemm_rm(
1608    m: usize,
1609    n: usize,
1610    k: usize,
1611    alpha: f32,
1612    a: &[f32],
1613    lda: usize,
1614    b_mat: &[f32],
1615    ldb: usize,
1616    b_rows_are_n: bool,
1617    c: &mut [f32],
1618    ldc: usize,
1619) {
1620    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1621}
1622
1623/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
1624/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
1625#[cfg(target_os = "macos")]
1626#[allow(clippy::too_many_arguments)]
1627pub(crate) fn sgemm_rm(
1628    m: usize,
1629    n: usize,
1630    k: usize,
1631    alpha: f32,
1632    a: &[f32],
1633    lda: usize,
1634    b_mat: &[f32],
1635    ldb: usize,
1636    b_rows_are_n: bool,
1637    c: &mut [f32],
1638    ldc: usize,
1639) {
1640    debug_assert!(a.len() >= (m - 1) * lda + k);
1641    debug_assert!(c.len() >= (m - 1) * ldc + n);
1642    // Test hook: route the attention GEMMs through the portable NEON
1643    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
1644    // measured without a phone in the loop. (Intel macOS has no NEON —
1645    // the hook is a no-op there, Accelerate continues below.)
1646    #[cfg(target_arch = "aarch64")]
1647    if std::env::var("CMF_FORCE_NEON_GEMM").map(|v| v == "1").unwrap_or(false) {
1648        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1649    }
1650    unsafe {
1651        accel_blas::cblas_sgemm(
1652            101, // RowMajor
1653            111, // NoTrans A
1654            if b_rows_are_n { 112 } else { 111 },
1655            m as i32,
1656            n as i32,
1657            k as i32,
1658            alpha,
1659            a.as_ptr(),
1660            lda as i32,
1661            b_mat.as_ptr(),
1662            ldb as i32,
1663            0.0,
1664            c.as_mut_ptr(),
1665            ldc as i32,
1666        );
1667    }
1668}
1669
1670/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
1671/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
1672/// on the AMX with one row-major sgemm. Tiles live in cache, weights
1673/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
1674/// logits shift within f32 rounding — tolerance-class, like every
1675/// reduction-order change; decode (M=1) never takes this path.
1676#[cfg(target_os = "macos")]
1677fn qmatmat_accel(
1678    q: &[u8],
1679    row_scale: &[f32],
1680    pre: &[std::borrow::Cow<'_, [f32]>],
1681    rows: usize,
1682    cols: usize,
1683    out: &mut [f32],
1684    pool: Option<&Pool>,
1685) {
1686    // NOTE: double-buffering the dequant against the sgemm (a scoped
1687    // thread driving the pool on tile k+1 while the caller multiplies
1688    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
1689    // multithreaded, and the dequant workers just steal its cores.
1690    const TR: usize = 2048;
1691    let b = pre.len();
1692    thread_local! {
1693        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1694        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1695    }
1696    XPANEL.with(|xp| {
1697        WTILE.with(|wt| {
1698            let mut xpanel = xp.borrow_mut();
1699            xpanel.clear();
1700            for x in pre {
1701                xpanel.extend_from_slice(x);
1702            }
1703            let mut wtile = wt.borrow_mut();
1704            wtile.resize(TR * cols, 0.0);
1705            let mut r0 = 0usize;
1706            while r0 < rows {
1707                let tr = TR.min(rows - r0);
1708                // Dequant the tile (scale folded) — pool-parallel.
1709                let wt_addr = SendMut(wtile.as_mut_ptr());
1710                let run = |start: usize, end: usize| {
1711                    for r in start..end {
1712                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1713                        let s = row_scale[r0 + r];
1714                        // SAFETY: workers cover disjoint r ranges.
1715                        let dst = unsafe {
1716                            std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1717                        };
1718                        for (d, &v) in dst.iter_mut().zip(row) {
1719                            *d = (v as i8) as f32 * s;
1720                        }
1721                    }
1722                };
1723                dispatch_rows(pool, tr, &run);
1724                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
1725                unsafe {
1726                    accel_blas::cblas_sgemm(
1727                        101, // RowMajor
1728                        111, // NoTrans A
1729                        112, // Trans B
1730                        b as i32,
1731                        tr as i32,
1732                        cols as i32,
1733                        1.0,
1734                        xpanel.as_ptr(),
1735                        cols as i32,
1736                        wtile.as_ptr(),
1737                        cols as i32,
1738                        0.0,
1739                        out.as_mut_ptr().add(r0),
1740                        rows as i32,
1741                    );
1742                }
1743                r0 += tr;
1744            }
1745        })
1746    });
1747}
1748
1749fn qmatmat(
1750    q: &[u8],
1751    row_scale: &[f32],
1752    pre: &[std::borrow::Cow<'_, [f32]>],
1753    rows: usize,
1754    cols: usize,
1755    out: &mut [f32],
1756    pool: Option<&Pool>,
1757) {
1758    let b = pre.len();
1759    debug_assert_eq!(out.len(), b * rows);
1760    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
1761    // SDOT loop below peaks near the CPU's dot throughput, an order
1762    // below the matrix units. Small tensors and tiny test models stay
1763    // on the exact integer path.
1764    #[cfg(target_os = "macos")]
1765    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1766        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1767        return;
1768    }
1769    #[cfg(target_arch = "aarch64")]
1770    if sdot_enabled() {
1771        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1772        let out_addr = SendMut(out.as_mut_ptr());
1773        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
1774        // path IS the ARM prefill GEMM off Apple silicon).
1775        let blocked_ok =
1776            std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1777        let use_i8mm = i8mm_enabled();
1778        if blocked_ok {
1779            let run = |start: usize, end: usize| {
1780                let mut o = start;
1781                while o < end {
1782                    if o + 2 <= end {
1783                        let r0 = &q[o * cols..(o + 1) * cols];
1784                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1785                        let mut bi = 0usize;
1786                        while bi + 4 <= acts.len() {
1787                            let xs = [
1788                                acts[bi].xq.as_slice(),
1789                                acts[bi + 1].xq.as_slice(),
1790                                acts[bi + 2].xq.as_slice(),
1791                                acts[bi + 3].xq.as_slice(),
1792                            ];
1793                            let d = if use_i8mm {
1794                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
1795                            } else {
1796                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
1797                            };
1798                            for (r, row) in [r0, r1].into_iter().enumerate() {
1799                                for k in 0..4 {
1800                                    let act = &acts[bi + k];
1801                                    let mut v = d[r][k] as f32 * act.sx;
1802                                    for &(j, xv) in &act.outliers {
1803                                        v += (row[j] as i8) as f32 * xv;
1804                                    }
1805                                    unsafe {
1806                                        *out_addr.at((bi + k) * rows + o + r) =
1807                                            v * row_scale[o + r]
1808                                    };
1809                                }
1810                            }
1811                            bi += 4;
1812                        }
1813                        while bi < acts.len() {
1814                            for (r, row) in [r0, r1].into_iter().enumerate() {
1815                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
1816                                unsafe { *out_addr.at(bi * rows + o + r) = v };
1817                            }
1818                            bi += 1;
1819                        }
1820                        o += 2;
1821                    } else {
1822                        let row = &q[o * cols..(o + 1) * cols];
1823                        for (bi, act) in acts.iter().enumerate() {
1824                            let v = row_dot_sdot(row, act) * row_scale[o];
1825                            unsafe { *out_addr.at(bi * rows + o) = v };
1826                        }
1827                        o += 1;
1828                    }
1829                }
1830            };
1831            dispatch_rows(pool, rows, &run);
1832            return;
1833        }
1834        let run = |start: usize, end: usize| {
1835            for o in start..end {
1836                let row = &q[o * cols..(o + 1) * cols];
1837                for (bi, act) in acts.iter().enumerate() {
1838                    let v = row_dot_sdot(row, act) * row_scale[o];
1839                    unsafe { *out_addr.at(bi * rows + o) = v };
1840                }
1841            }
1842        };
1843        dispatch_rows(pool, rows, &run);
1844        return;
1845    }
1846    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
1847    // (roadmap P0: two weight rows' abs() stay in registers across four
1848    // activation streams); VNNI machines keep the per-row bias-trick
1849    // dot, which is already throughput-bound there.
1850    #[cfg(target_arch = "x86_64")]
1851    if avx2_a8w8_enabled() {
1852        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1853        let out_addr = SendMut(out.as_mut_ptr());
1854        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
1855        // A/B on noisy shared-vCPU hosts).
1856        let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1857        if !avx512vnni_enabled() && blocked_ok {
1858            let run = |start: usize, end: usize| {
1859                let mut o = start;
1860                while o < end {
1861                    if o + 2 <= end {
1862                        let r0 = &q[o * cols..(o + 1) * cols];
1863                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1864                        let mut bi = 0usize;
1865                        while bi + 4 <= acts.len() {
1866                            let xs = [
1867                                acts[bi].xq.as_slice(),
1868                                acts[bi + 1].xq.as_slice(),
1869                                acts[bi + 2].xq.as_slice(),
1870                                acts[bi + 3].xq.as_slice(),
1871                            ];
1872                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1873                            for (r, row) in [r0, r1].into_iter().enumerate() {
1874                                for k in 0..4 {
1875                                    let act = &acts[bi + k];
1876                                    let mut v = d[r][k] as f32 * act.sx;
1877                                    for &(j, xv) in &act.outliers {
1878                                        v += (row[j] as i8) as f32 * xv;
1879                                    }
1880                                    unsafe {
1881                                        *out_addr.at((bi + k) * rows + o + r) =
1882                                            v * row_scale[o + r]
1883                                    };
1884                                }
1885                            }
1886                            bi += 4;
1887                        }
1888                        while bi < acts.len() {
1889                            for (r, row) in [r0, r1].into_iter().enumerate() {
1890                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1891                                unsafe { *out_addr.at(bi * rows + o + r) = v };
1892                            }
1893                            bi += 1;
1894                        }
1895                        o += 2;
1896                    } else {
1897                        let row = &q[o * cols..(o + 1) * cols];
1898                        for (bi, act) in acts.iter().enumerate() {
1899                            let v = row_dot_avx2(row, act) * row_scale[o];
1900                            unsafe { *out_addr.at(bi * rows + o) = v };
1901                        }
1902                        o += 1;
1903                    }
1904                }
1905            };
1906            dispatch_rows(pool, rows, &run);
1907            return;
1908        }
1909        let run = |start: usize, end: usize| {
1910            for o in start..end {
1911                let row = &q[o * cols..(o + 1) * cols];
1912                for (bi, act) in acts.iter().enumerate() {
1913                    let v = row_dot_avx2(row, act) * row_scale[o];
1914                    unsafe { *out_addr.at(bi * rows + o) = v };
1915                }
1916            }
1917        };
1918        dispatch_rows(pool, rows, &run);
1919        return;
1920    }
1921    let out_addr = SendMut(out.as_mut_ptr());
1922    let run = |start: usize, end: usize| {
1923        for o in start..end {
1924            let row = &q[o * cols..(o + 1) * cols];
1925            for (bi, x) in pre.iter().enumerate() {
1926                let mut acc = 0f32;
1927                for j in 0..cols {
1928                    acc += (row[j] as i8) as f32 * x[j];
1929                }
1930                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1931            }
1932        }
1933    };
1934    dispatch_rows(pool, rows, &run);
1935}
1936
1937/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
1938/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
1939fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1940    match pool {
1941        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1942        _ => run(0, rows),
1943    }
1944}
1945
1946/// Split a q4_block blob into (packed nibbles, f16 group scales).
1947fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1948    let groups = rows * cols / GROUP_SIZE;
1949    bytes.split_at(groups * 16)
1950}
1951
1952/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
1953/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
1954/// vbit packs MSB-first, so the HIGH nibble is the even element
1955/// (opposite of q4_block's lo-first interleave). Centering is u-7.
1956#[inline]
1957fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1958    #[cfg(target_arch = "aarch64")]
1959    unsafe {
1960        return vbit_fill4_neon(data, buf);
1961    }
1962    #[cfg(target_arch = "x86_64")]
1963    if avx2_enabled() {
1964        return unsafe { vbit_fill4_avx2(data, buf) };
1965    }
1966    #[allow(unreachable_code)]
1967    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1968        let u = unpack8::<4>(&data[blk * 4..]);
1969        for k in 0..8 {
1970            chunk[k] = (u[k] - 7) as i8 as u8;
1971        }
1972    }
1973}
1974
1975#[cfg(target_arch = "aarch64")]
1976#[target_feature(enable = "neon")]
1977unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1978    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
1979    // buf.len()/2 packed bytes (validated at load).
1980    unsafe {
1981        use core::arch::aarch64::*;
1982        let n = buf.len();
1983        let mask = vdupq_n_u8(0x0F);
1984        let seven = vdupq_n_s8(7);
1985        let mut g = 0usize;
1986        while g * 32 + 32 <= n {
1987            let b = vld1q_u8(data.as_ptr().add(g * 16));
1988            let hi = vshrq_n_u8::<4>(b);
1989            let lo = vandq_u8(b, mask);
1990            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1991            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1992            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1993            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1994            g += 1;
1995        }
1996    }
1997}
1998
1999#[cfg(target_arch = "x86_64")]
2000#[target_feature(enable = "avx2")]
2001unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2002    // SAFETY: see vbit_fill4_neon.
2003    unsafe {
2004        use core::arch::x86_64::*;
2005        let n = buf.len();
2006        let mask = _mm_set1_epi8(0x0F);
2007        let seven = _mm256_set1_epi8(7);
2008        let mut g = 0usize;
2009        while g * 32 + 32 <= n {
2010            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2011            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2012            let lo = _mm_and_si128(b, mask);
2013            let z = _mm256_sub_epi8(
2014                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2015                seven,
2016            );
2017            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2018            g += 1;
2019        }
2020    }
2021}
2022
2023/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
2024/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
2025/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
2026/// into 4 such blocks.
2027#[inline(always)]
2028fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2029    let mut acc = 0u64;
2030    for i in 0..B {
2031        acc = (acc << 8) | data[i] as u64;
2032    }
2033    let mask = (1u64 << B) - 1;
2034    let mut out = [0i32; 8];
2035    for (k, o) in out.iter_mut().enumerate() {
2036        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2037    }
2038    out
2039}
2040
2041/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
2042/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
2043/// MSB-first, byte-padded]. Row data offsets are precomputed at load
2044/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
2045/// overhead on every matvec.
2046#[allow(clippy::too_many_arguments)]
2047fn vbitmatvec(
2048    bytes: &[u8],
2049    offsets: &[usize],
2050    x: &[f32],
2051    rows: usize,
2052    cols: usize,
2053    out: &mut [f32],
2054    pool: Option<&Pool>,
2055) {
2056    debug_assert_eq!(out.len(), rows);
2057    debug_assert_eq!(offsets.len(), rows + 1);
2058
2059    // SDOT path: unpack the row to centered i8 once, then per-group
2060    // int8 dot against the quantized activations — same A8W8 contract
2061    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
2062    if a8w8_enabled() {
2063        let act = split_act(x);
2064        let out_addr = SendMut(out.as_mut_ptr());
2065        let run = move |start: usize, end: usize| {
2066            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2067        };
2068        dispatch_rows(pool, rows, &run);
2069        return;
2070    }
2071
2072    let out_addr = SendMut(out.as_mut_ptr());
2073    let run = move |start: usize, end: usize| {
2074        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2075    };
2076    dispatch_rows(pool, rows, &run);
2077}
2078
2079/// One vbit row range via the A8W8 int8 path — kernel body of
2080/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
2081/// several tensors in one dispatch (b=8 rows go exact f32).
2082#[allow(clippy::too_many_arguments)]
2083fn vbit_range_a8w8(
2084    bytes: &[u8],
2085    offsets: &[usize],
2086    x: &[f32],
2087    act: &SplitAct,
2088    rows: usize,
2089    cols: usize,
2090    out: SendMut,
2091    start: usize,
2092    end: usize,
2093) {
2094    let ng = cols / GROUP_SIZE;
2095    let bits = &bytes[..rows];
2096    let sc_off = rows;
2097    let row_dot = |r: usize| -> f32 {
2098            let b = bits[r] as usize;
2099            let l = ((1i32 << (b - 1)) - 1) as i32;
2100            let mask = (1u64 << b) - 1;
2101            let data = &bytes[offsets[r]..offsets[r + 1]];
2102            if b == 8 {
2103                // u−L reaches 128 → does not fit i8; exact f32 path.
2104                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2105                let mut dot = 0f32;
2106                for g in 0..ng {
2107                    let so = (r * ng + g) * 2;
2108                    let sgf = f16_to_f32(u16::from_le_bytes([
2109                        bytes[sc_off + so],
2110                        bytes[sc_off + so + 1],
2111                    ]));
2112                    let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2113                    let mut gd = 0f32;
2114                    for &xv in xg.iter() {
2115                        if nbits < 8 {
2116                            acc = (acc << 8) | data[idx] as u64;
2117                            idx += 1;
2118                            nbits += 8;
2119                        }
2120                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2121                        nbits -= 8;
2122                        gd += (u - l) as f32 * xv;
2123                    }
2124                    dot += gd * sgf;
2125                }
2126                return dot;
2127            }
2128            // Per-worker scratch: this closure runs for every row of the
2129            // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
2130            // row was measurable pure overhead.
2131            thread_local! {
2132                static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2133                    const { std::cell::RefCell::new(Vec::new()) };
2134            }
2135            #[inline(always)]
2136            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2137                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2138                    let u = unpack8::<B>(&data[blk * B..]);
2139                    for k in 0..8 {
2140                        chunk[k] = (u[k] - l) as i8 as u8;
2141                    }
2142                }
2143            }
2144            let _ = mask;
2145            VBIT_SCRATCH.with(|scratch| {
2146                let mut buf = scratch.borrow_mut();
2147                buf.resize(cols, 0);
2148                match b {
2149                    3 => fill::<3>(data, l, &mut buf),
2150                    4 => vbit_fill4(data, &mut buf),
2151                    5 => fill::<5>(data, l, &mut buf),
2152                    6 => fill::<6>(data, l, &mut buf),
2153                    _ => unreachable!(),
2154                }
2155                let mut dot = 0f32;
2156                for g in 0..ng {
2157                    let so = (r * ng + g) * 2;
2158                    let s = f16_to_f32(u16::from_le_bytes([
2159                        bytes[sc_off + so],
2160                        bytes[sc_off + so + 1],
2161                    ]));
2162                    let d = dot_i8_i8(
2163                        &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2164                        &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2165                    ) as f32
2166                        * act.sx;
2167                    dot += d * s;
2168                }
2169                for &(j, xv) in &act.outliers {
2170                    let so = (r * ng + j / GROUP_SIZE) * 2;
2171                    let s = f16_to_f32(u16::from_le_bytes([
2172                        bytes[sc_off + so],
2173                        bytes[sc_off + so + 1],
2174                    ]));
2175                    // xq is zeroed at outlier slots — add the exact term.
2176                    dot += (buf[j] as i8) as f32 * s * xv;
2177                }
2178                dot
2179            })
2180    };
2181    for r in start..end {
2182        // SAFETY: disjoint row ranges per worker.
2183        unsafe { *out.at(r) = row_dot(r) };
2184    }
2185}
2186
2187/// Exact scalar vbit row range (same extraction, non-SDOT path).
2188#[allow(clippy::too_many_arguments)]
2189fn vbit_range_f32(
2190    bytes: &[u8],
2191    offsets: &[usize],
2192    x: &[f32],
2193    rows: usize,
2194    cols: usize,
2195    out: SendMut,
2196    start: usize,
2197    end: usize,
2198) {
2199    let ng = cols / GROUP_SIZE;
2200    let bits = &bytes[..rows];
2201    let sc_off = rows;
2202    // Per-bit-width specialized inner loops: the compiler unrolls the
2203    // constant shifts (the generic bit-buffer loop was branch-bound —
2204    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
2205    #[inline(always)]
2206    fn dot_row<const B: usize>(
2207        data: &[u8],
2208        bytes: &[u8],
2209        sc_off: usize,
2210        r: usize,
2211        ng: usize,
2212        x: &[f32],
2213    ) -> f32 {
2214        let l = ((1i32 << (B - 1)) - 1) as f32;
2215        let gbytes = GROUP_SIZE * B / 8;
2216        let mut dot = 0f32;
2217        for g in 0..ng {
2218            let so = (r * ng + g) * 2;
2219            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2220            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2221            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2222            let mut gd = 0f32;
2223            for blk in 0..GROUP_SIZE / 8 {
2224                let u = unpack8::<B>(&gd0[blk * B..]);
2225                let xb = &xg[blk * 8..blk * 8 + 8];
2226                for k in 0..8 {
2227                    gd += (u[k] as f32 - l) * xb[k];
2228                }
2229            }
2230            dot += gd * s;
2231        }
2232        dot
2233    }
2234    for r in start..end {
2235        let data = &bytes[offsets[r]..offsets[r + 1]];
2236        let v = match bits[r] {
2237            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2238            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2239            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2240            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2241            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2242            b => unreachable!("vbit bit-width {b} (validated at load)"),
2243        };
2244        // SAFETY: disjoint row ranges per worker.
2245        unsafe { *out.at(r) = v };
2246    }
2247}
2248
2249/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
2250/// and dotted against BOTH activations (MTP verify / pair prefill used
2251/// to run two full matvecs — double weight traffic and double unpack).
2252/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
2253#[allow(clippy::too_many_arguments)]
2254fn vbitmatvec2(
2255    bytes: &[u8],
2256    offsets: &[usize],
2257    x1: &[f32],
2258    x2: &[f32],
2259    rows: usize,
2260    cols: usize,
2261    o1: &mut [f32],
2262    o2: &mut [f32],
2263    pool: Option<&Pool>,
2264) {
2265    debug_assert_eq!(o1.len(), rows);
2266    debug_assert_eq!(o2.len(), rows);
2267
2268    if a8w8_enabled() {
2269        let a1 = split_act(x1);
2270        let a2 = split_act(x2);
2271        let p1 = SendMut(o1.as_mut_ptr());
2272        let p2 = SendMut(o2.as_mut_ptr());
2273        let run = move |start: usize, end: usize| {
2274            vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
2275        };
2276        dispatch_rows(pool, rows, &run);
2277        return;
2278    }
2279
2280    let p1 = SendMut(o1.as_mut_ptr());
2281    let p2 = SendMut(o2.as_mut_ptr());
2282    let run = move |start: usize, end: usize| {
2283        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2284    };
2285    dispatch_rows(pool, rows, &run);
2286}
2287
2288/// Two-input vbit row range via the A8W8 int8 path — kernel body of
2289/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
2290/// exact f32 for both lanes, bits streamed once).
2291#[allow(clippy::too_many_arguments)]
2292fn vbit_range2_a8w8(
2293    bytes: &[u8],
2294    offsets: &[usize],
2295    x1: &[f32],
2296    x2: &[f32],
2297    a1: &SplitAct,
2298    a2: &SplitAct,
2299    rows: usize,
2300    cols: usize,
2301    p1: SendMut,
2302    p2: SendMut,
2303    start: usize,
2304    end: usize,
2305) {
2306    let ng = cols / GROUP_SIZE;
2307    let bits = &bytes[..rows];
2308    let sc_off = rows;
2309    let row_dots = |r: usize| -> (f32, f32) {
2310            let b = bits[r] as usize;
2311            let l = (1i32 << (b - 1)) - 1;
2312            let data = &bytes[offsets[r]..offsets[r + 1]];
2313            if b == 8 {
2314                // u−L reaches 128 → does not fit i8; exact f32 path,
2315                // bits still streamed once for both lanes.
2316                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2317                let (mut d1, mut d2) = (0f32, 0f32);
2318                for g in 0..ng {
2319                    let so = (r * ng + g) * 2;
2320                    let sgf = f16_to_f32(u16::from_le_bytes([
2321                        bytes[sc_off + so],
2322                        bytes[sc_off + so + 1],
2323                    ]));
2324                    let (mut g1, mut g2) = (0f32, 0f32);
2325                    for k in 0..GROUP_SIZE {
2326                        if nbits < 8 {
2327                            acc = (acc << 8) | data[idx] as u64;
2328                            idx += 1;
2329                            nbits += 8;
2330                        }
2331                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2332                        nbits -= 8;
2333                        let w = (u - l) as f32;
2334                        g1 += w * x1[g * GROUP_SIZE + k];
2335                        g2 += w * x2[g * GROUP_SIZE + k];
2336                    }
2337                    d1 += g1 * sgf;
2338                    d2 += g2 * sgf;
2339                }
2340                return (d1, d2);
2341            }
2342            thread_local! {
2343                static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2344                    const { std::cell::RefCell::new(Vec::new()) };
2345            }
2346            #[inline(always)]
2347            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2348                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2349                    let u = unpack8::<B>(&data[blk * B..]);
2350                    for k in 0..8 {
2351                        chunk[k] = (u[k] - l) as i8 as u8;
2352                    }
2353                }
2354            }
2355            VBIT_SCRATCH2.with(|scratch| {
2356                let mut buf = scratch.borrow_mut();
2357                buf.resize(cols, 0);
2358                match b {
2359                    3 => fill::<3>(data, l, &mut buf),
2360                    4 => vbit_fill4(data, &mut buf),
2361                    5 => fill::<5>(data, l, &mut buf),
2362                    6 => fill::<6>(data, l, &mut buf),
2363                    _ => unreachable!(),
2364                }
2365                let (mut d1, mut d2) = (0f32, 0f32);
2366                for g in 0..ng {
2367                    let so = (r * ng + g) * 2;
2368                    let s = f16_to_f32(u16::from_le_bytes([
2369                        bytes[sc_off + so],
2370                        bytes[sc_off + so + 1],
2371                    ]));
2372                    let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2373                    let v1 =
2374                        dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2375                    let v2 =
2376                        dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2377                    d1 += v1 * s;
2378                    d2 += v2 * s;
2379                }
2380                for &(j, xv) in &a1.outliers {
2381                    let so = (r * ng + j / GROUP_SIZE) * 2;
2382                    let s = f16_to_f32(u16::from_le_bytes([
2383                        bytes[sc_off + so],
2384                        bytes[sc_off + so + 1],
2385                    ]));
2386                    d1 += (buf[j] as i8) as f32 * s * xv;
2387                }
2388                for &(j, xv) in &a2.outliers {
2389                    let so = (r * ng + j / GROUP_SIZE) * 2;
2390                    let s = f16_to_f32(u16::from_le_bytes([
2391                        bytes[sc_off + so],
2392                        bytes[sc_off + so + 1],
2393                    ]));
2394                    d2 += (buf[j] as i8) as f32 * s * xv;
2395                }
2396                (d1, d2)
2397            })
2398    };
2399    for r in start..end {
2400        let (v1, v2) = row_dots(r);
2401        // SAFETY: disjoint row ranges per worker.
2402        unsafe {
2403            *p1.at(r) = v1;
2404            *p2.at(r) = v2;
2405        }
2406    }
2407}
2408
2409/// Two-input exact scalar vbit row range (same extraction) —
2410/// per-bit-width specialized, two accumulators per row; per-lane
2411/// accumulation order matches `vbitmatvec` exactly.
2412#[allow(clippy::too_many_arguments)]
2413fn vbit_range2_f32(
2414    bytes: &[u8],
2415    offsets: &[usize],
2416    x1: &[f32],
2417    x2: &[f32],
2418    rows: usize,
2419    cols: usize,
2420    p1: SendMut,
2421    p2: SendMut,
2422    start: usize,
2423    end: usize,
2424) {
2425    let ng = cols / GROUP_SIZE;
2426    let bits = &bytes[..rows];
2427    let sc_off = rows;
2428    #[inline(always)]
2429    #[allow(clippy::too_many_arguments)]
2430    fn dot_row2<const B: usize>(
2431        data: &[u8],
2432        bytes: &[u8],
2433        sc_off: usize,
2434        r: usize,
2435        ng: usize,
2436        x1: &[f32],
2437        x2: &[f32],
2438    ) -> (f32, f32) {
2439        let l = ((1i32 << (B - 1)) - 1) as f32;
2440        let gbytes = GROUP_SIZE * B / 8;
2441        let (mut d1, mut d2) = (0f32, 0f32);
2442        for g in 0..ng {
2443            let so = (r * ng + g) * 2;
2444            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2445            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2446            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2447            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2448            let (mut g1, mut g2) = (0f32, 0f32);
2449            for blk in 0..GROUP_SIZE / 8 {
2450                let u = unpack8::<B>(&gd0[blk * B..]);
2451                for k in 0..8 {
2452                    let w = u[k] as f32 - l;
2453                    g1 += w * x1g[blk * 8 + k];
2454                    g2 += w * x2g[blk * 8 + k];
2455                }
2456            }
2457            d1 += g1 * s;
2458            d2 += g2 * s;
2459        }
2460        (d1, d2)
2461    }
2462    for r in start..end {
2463        let data = &bytes[offsets[r]..offsets[r + 1]];
2464        let (v1, v2) = match bits[r] {
2465            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2466            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2467            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2468            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2469            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2470            b => unreachable!("vbit bit-width {b} (validated at load)"),
2471        };
2472        // SAFETY: disjoint row ranges per worker.
2473        unsafe {
2474            *p1.at(r) = v1;
2475            *p2.at(r) = v2;
2476        }
2477    }
2478}
2479
2480// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
2481
2482/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
2483/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
2484/// distant streams of the split layout. Values/order identical to the
2485/// split kernels.
2486#[inline]
2487#[allow(unreachable_code)]
2488fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2489    #[cfg(target_arch = "aarch64")]
2490    unsafe {
2491        return dot_q4t_row_sdot(bytes, r, gpr, xq);
2492    }
2493    #[cfg(target_arch = "x86_64")]
2494    unsafe {
2495        return dot_q4t_row_avx2(bytes, r, gpr, xq);
2496    }
2497    let mut acc = 0f32;
2498    for gi in 0..gpr {
2499        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2500        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2501        let mut d = 0i32;
2502        for (k, &b) in tile[2..].iter().enumerate() {
2503            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2504                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2505        }
2506        acc += d as f32 * s;
2507    }
2508    acc
2509}
2510
2511#[cfg(target_arch = "aarch64")]
2512#[target_feature(enable = "neon,dotprod")]
2513unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2514    // SAFETY: callers uphold slice-length contracts (18B tile per group,
2515    // xq.len() == gpr·GROUP_SIZE).
2516    unsafe {
2517        use core::arch::aarch64::*;
2518        use core::arch::asm;
2519        let lomask = vdupq_n_u8(0x0F);
2520        let eight = vdupq_n_s8(8);
2521        let mut acc = 0f32;
2522        for gi in 0..gpr {
2523            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2524            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2525            let b = vld1q_u8(t.add(2));
2526            let lo = vandq_u8(b, lomask);
2527            let hi = vshrq_n_u8::<4>(b);
2528            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2529            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2530            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2531            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2532            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2533            asm!(
2534                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2535                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2536                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2537                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2538                options(pure, nomem, nostack),
2539            );
2540            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2541        }
2542        acc
2543    }
2544}
2545
2546#[cfg(target_arch = "x86_64")]
2547#[target_feature(enable = "avx2")]
2548unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2549    // SAFETY: see dot_q4t_row_sdot.
2550    unsafe {
2551        use core::arch::x86_64::*;
2552        let lomask = _mm_set1_epi8(0x0F);
2553        let eight = _mm256_set1_epi8(8);
2554        let ones = _mm256_set1_epi16(1);
2555        let mut acc = 0f32;
2556        for gi in 0..gpr {
2557            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2558            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2559            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2560            let lo = _mm_and_si128(b, lomask);
2561            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2562            let w = _mm256_sub_epi8(
2563                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2564                eight,
2565            );
2566            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2567            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2568            let d = _mm256_madd_epi16(p16, ones);
2569            let hi128 = _mm256_extracti128_si256::<1>(d);
2570            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2571            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2572            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2573            acc += _mm_cvtsi128_si32(s32) as f32 * s;
2574        }
2575        acc
2576    }
2577}
2578
2579/// One q4_tiled row against FOUR activation streams: the nibble unpack
2580/// and abs() happen once per group instead of once per (group,
2581/// activation) — the unpack is the dominant per-element cost of the
2582/// tiled format (roadmap P0 portable blocking, q4t leg).
2583#[cfg(target_arch = "x86_64")]
2584#[target_feature(enable = "avx2")]
2585unsafe fn dot_q4t_row_1x4_avx2(
2586    bytes: &[u8],
2587    r: usize,
2588    gpr: usize,
2589    xs: [&[i8]; 4],
2590) -> [f32; 4] {
2591    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
2592    unsafe {
2593        use core::arch::x86_64::*;
2594        let lomask = _mm_set1_epi8(0x0F);
2595        let eight = _mm256_set1_epi8(8);
2596        let ones = _mm256_set1_epi16(1);
2597        let mut acc = [0f32; 4];
2598        for gi in 0..gpr {
2599            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2600            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2601            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2602            let lo = _mm_and_si128(bb, lomask);
2603            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2604            let w = _mm256_sub_epi8(
2605                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2606                eight,
2607            );
2608            let aw = _mm256_abs_epi8(w);
2609            for (k, xq) in xs.iter().enumerate() {
2610                let x =
2611                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2612                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2613                let d = _mm256_madd_epi16(p16, ones);
2614                let hi128 = _mm256_extracti128_si256::<1>(d);
2615                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2616                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2617                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2618                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2619            }
2620        }
2621        acc
2622    }
2623}
2624
2625/// Exact-term correction for A8W8 outliers on a tiled row.
2626#[inline]
2627fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2628    let gi = j / GROUP_SIZE;
2629    let k = j % GROUP_SIZE;
2630    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2631    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2632    let byte = tile[2 + k / 2];
2633    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2634    ((nib as i32 - 8) as f32, s)
2635}
2636
2637/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
2638/// accumulation shape as `q4_range_f32`.
2639#[inline]
2640fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2641    let mut acc = 0f32;
2642    for gi in 0..gpr {
2643        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2644        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2645        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2646        let mut ga = 0f32;
2647        for (k, &b) in tile[2..].iter().enumerate() {
2648            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2649                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2650        }
2651        acc += ga * s;
2652    }
2653    acc
2654}
2655
2656/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
2657fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2658    debug_assert_eq!(out.len(), rows);
2659    let gpr = cols / GROUP_SIZE;
2660    let out_addr = SendMut(out.as_mut_ptr());
2661    if a8w8_enabled() {
2662        let act = split_act(x);
2663        let run = move |start: usize, end: usize| {
2664            for r in start..end {
2665                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2666                for &(j, xv) in &act.outliers {
2667                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2668                    acc += w * s * xv;
2669                }
2670                // SAFETY: disjoint row ranges per worker.
2671                unsafe { *out_addr.at(r) = acc };
2672            }
2673        };
2674        dispatch_rows(pool, rows, &run);
2675        return;
2676    }
2677    let run = move |start: usize, end: usize| {
2678        for r in start..end {
2679            // SAFETY: disjoint row ranges per worker.
2680            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2681        }
2682    };
2683    dispatch_rows(pool, rows, &run);
2684}
2685
2686/// Fused two-input q4_tiled matvec (weights read once per pair).
2687#[allow(clippy::too_many_arguments)]
2688fn q4t_matvec2(
2689    bytes: &[u8],
2690    x1: &[f32],
2691    x2: &[f32],
2692    rows: usize,
2693    cols: usize,
2694    o1: &mut [f32],
2695    o2: &mut [f32],
2696    pool: Option<&Pool>,
2697) {
2698    let gpr = cols / GROUP_SIZE;
2699    let p1 = SendMut(o1.as_mut_ptr());
2700    let p2 = SendMut(o2.as_mut_ptr());
2701    if a8w8_enabled() {
2702        let a1 = split_act(x1);
2703        let a2 = split_act(x2);
2704        let run = move |start: usize, end: usize| {
2705            for r in start..end {
2706                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2707                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2708                for &(j, xv) in &a1.outliers {
2709                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2710                    v1 += w * s * xv;
2711                }
2712                for &(j, xv) in &a2.outliers {
2713                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2714                    v2 += w * s * xv;
2715                }
2716                // SAFETY: disjoint row ranges per worker.
2717                unsafe {
2718                    *p1.at(r) = v1;
2719                    *p2.at(r) = v2;
2720                }
2721            }
2722        };
2723        dispatch_rows(pool, rows, &run);
2724        return;
2725    }
2726    let run = move |start: usize, end: usize| {
2727        for r in start..end {
2728            // SAFETY: disjoint row ranges per worker.
2729            unsafe {
2730                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2731                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2732            }
2733        }
2734    };
2735    dispatch_rows(pool, rows, &run);
2736}
2737
2738/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
2739#[allow(clippy::too_many_arguments)]
2740fn q4t_matmat(
2741    bytes: &[u8],
2742    xs_all: &[f32],
2743    b: usize,
2744    rows: usize,
2745    cols: usize,
2746    out: &mut [f32],
2747    pool: Option<&Pool>,
2748) {
2749    debug_assert_eq!(out.len(), b * rows);
2750    let gpr = cols / GROUP_SIZE;
2751    let out_addr = SendMut(out.as_mut_ptr());
2752    if a8w8_enabled() {
2753        let acts: Vec<SplitAct> =
2754            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2755        let acts = &acts;
2756        #[cfg(target_arch = "x86_64")]
2757        let blocked_ok = avx2_enabled()
2758            && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2759        #[cfg(not(target_arch = "x86_64"))]
2760        let blocked_ok = false;
2761        let run = move |start: usize, end: usize| {
2762            for r in start..end {
2763                let mut bi = 0usize;
2764                #[cfg(target_arch = "x86_64")]
2765                if blocked_ok {
2766                    while bi + 4 <= acts.len() {
2767                        let xs = [
2768                            acts[bi].xq.as_slice(),
2769                            acts[bi + 1].xq.as_slice(),
2770                            acts[bi + 2].xq.as_slice(),
2771                            acts[bi + 3].xq.as_slice(),
2772                        ];
2773                        let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2774                        for k in 0..4 {
2775                            let act = &acts[bi + k];
2776                            let mut acc = d[k] * act.sx;
2777                            for &(j, xv) in &act.outliers {
2778                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2779                                acc += w * sc * xv;
2780                            }
2781                            // SAFETY: disjoint (bi, r) cells per worker.
2782                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2783                        }
2784                        bi += 4;
2785                    }
2786                }
2787                let _ = blocked_ok;
2788                while bi < acts.len() {
2789                    let act = &acts[bi];
2790                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2791                    for &(j, xv) in &act.outliers {
2792                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
2793                        acc += w * s * xv;
2794                    }
2795                    // SAFETY: disjoint (bi, r) cells per worker range.
2796                    unsafe { *out_addr.at(bi * rows + r) = acc };
2797                    bi += 1;
2798                }
2799            }
2800        };
2801        dispatch_rows(pool, rows, &run);
2802        return;
2803    }
2804    let run = move |start: usize, end: usize| {
2805        for r in start..end {
2806            for bi in 0..b {
2807                let x = &xs_all[bi * cols..(bi + 1) * cols];
2808                // SAFETY: disjoint (bi, r) cells per worker range.
2809                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2810            }
2811        }
2812    };
2813    dispatch_rows(pool, rows, &run);
2814}
2815
2816// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
2817// 32-group tile. The kernel family mirrors q4_tiled: one sequential
2818// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
2819// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
2820
2821/// Per-32-group sums of the quantized activation — the ±1 identity's
2822/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
2823/// matvec and reused by every row.
2824fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2825    (0..gpr)
2826        .map(|gi| {
2827            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2828                .iter()
2829                .map(|&v| v as i32)
2830                .sum()
2831        })
2832        .collect()
2833}
2834
2835/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
2836/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
2837/// x86 pass).
2838#[inline]
2839#[allow(unreachable_code)]
2840/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
2841/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
2842/// masked activation sums through maddubs(1, x&mask), and
2843/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
2844#[cfg(target_arch = "x86_64")]
2845#[target_feature(enable = "avx2")]
2846unsafe fn dot_q1_row_avx2(
2847    bytes: &[u8],
2848    r: usize,
2849    gpr: usize,
2850    xq: &[i8],
2851    gsum: &[i32],
2852) -> f32 {
2853    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
2854    unsafe {
2855        use core::arch::x86_64::*;
2856        // Byte j of the mask must replicate bits-byte j/8.
2857        let expand = _mm256_setr_epi8(
2858            0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2859            2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2860        );
2861        let bitsel = _mm256_setr_epi8(
2862            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2863            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2864        );
2865        let ones8 = _mm256_set1_epi8(1);
2866        let ones16 = _mm256_set1_epi16(1);
2867        let mut acc = 0f32;
2868        for gi in 0..gpr {
2869            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2870            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2871            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2872            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2873            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2874            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2875            let sel = _mm256_and_si256(x, mask);
2876            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
2877            let p16 = _mm256_maddubs_epi16(ones8, sel);
2878            let d32 = _mm256_madd_epi16(p16, ones16);
2879            let hi128 = _mm256_extracti128_si256::<1>(d32);
2880            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2881            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2882            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2883            let msum = _mm_cvtsi128_si32(s32);
2884            // The and-select keeps x UN-negated (unlike ARM's −1-mask
2885            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
2886            let d = 2 * msum - gsum[gi];
2887            acc += d as f32 * s;
2888        }
2889        acc
2890    }
2891}
2892
2893/// The blocked 1×4 flavor: the expanded bit mask serves four activation
2894/// streams per group (mask build once, four select+reduce chains).
2895#[cfg(target_arch = "x86_64")]
2896#[target_feature(enable = "avx2")]
2897unsafe fn dot_q1_row_1x4_avx2(
2898    bytes: &[u8],
2899    r: usize,
2900    gpr: usize,
2901    xs: [&[i8]; 4],
2902    gsums: [&[i32]; 4],
2903) -> [f32; 4] {
2904    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
2905    unsafe {
2906        use core::arch::x86_64::*;
2907        let expand = _mm256_setr_epi8(
2908            0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2909            2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2910        );
2911        let bitsel = _mm256_setr_epi8(
2912            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2913            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2914        );
2915        let ones8 = _mm256_set1_epi8(1);
2916        let ones16 = _mm256_set1_epi16(1);
2917        let mut acc = [0f32; 4];
2918        for gi in 0..gpr {
2919            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2920            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2921            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2922            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2923            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2924            for (k, xq) in xs.iter().enumerate() {
2925                let x =
2926                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2927                let sel = _mm256_and_si256(x, mask);
2928                let p16 = _mm256_maddubs_epi16(ones8, sel);
2929                let d32 = _mm256_madd_epi16(p16, ones16);
2930                let hi128 = _mm256_extracti128_si256::<1>(d32);
2931                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2932                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2933                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2934                let msum = _mm_cvtsi128_si32(s32);
2935                let d = 2 * msum - gsums[k][gi];
2936                acc[k] += d as f32 * s;
2937            }
2938        }
2939        acc
2940    }
2941}
2942
2943fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2944    #[cfg(target_arch = "aarch64")]
2945    unsafe {
2946        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2947    }
2948    #[cfg(target_arch = "x86_64")]
2949    if avx2_enabled() {
2950        unsafe {
2951            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
2952        }
2953    }
2954    let _ = gsum;
2955    let mut acc = 0f32;
2956    for gi in 0..gpr {
2957        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2958        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2959        let mut d = 0i32;
2960        for (j, &b) in tile[2..].iter().enumerate() {
2961            for k in 0..8 {
2962                let w = ((b >> k) & 1) as i32 * 2 - 1;
2963                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2964            }
2965        }
2966        acc += d as f32 * s;
2967    }
2968    acc
2969}
2970
2971/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
2972/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
2973/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
2974/// per-group activation sums shared across every row of the matvec.
2975/// Four tiles (128 weights) per iteration: integer dots reduce through
2976/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
2977/// fused f32 multiply-add. Integer math throughout — bit-identical to
2978/// the scalar ±1 reference.
2979#[cfg(target_arch = "aarch64")]
2980#[target_feature(enable = "neon,dotprod")]
2981unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2982    // SAFETY: callers uphold slice-length contracts (6B tile per group,
2983    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
2984    unsafe {
2985        use core::arch::aarch64::*;
2986        use core::arch::asm;
2987        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2988        let m = vld1q_u8(MASKS.as_ptr());
2989        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
2990        macro_rules! tile_dot {
2991            ($t:expr, $x:expr) => {{
2992                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2993                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2994                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2995                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2996                let x0 = vld1q_s8($x);
2997                let x1 = vld1q_s8($x.add(16));
2998                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2999                asm!(
3000                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3001                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3002                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3003                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3004                    options(pure, nomem, nostack),
3005                );
3006                vaddq_s32(a0, a1)
3007            }};
3008        }
3009        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
3010        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
3011        // bit-byte across 8 lanes for vtst, and the four scales gather
3012        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
3013        // 4 branchy software f16 conversions per 128 weights (the
3014        // measured load-port wall of this kernel) become 2 vector
3015        // loads + 9 table lookups. Integer math order is unchanged —
3016        // bit-identical results (FCVTL is exact on every f16).
3017        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3018        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3019        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3020        const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
3021        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3022        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3023        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3024        let isc = vld1_u8(ISC.as_ptr());
3025        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
3026        macro_rules! tile_dot_tbl {
3027            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
3028                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
3029                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
3030                let x0 = vld1q_s8($x);
3031                let x1 = vld1q_s8($x.add(16));
3032                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3033                asm!(
3034                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3035                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3036                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3037                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3038                    options(pure, nomem, nostack),
3039                );
3040                vaddq_s32(a0, a1)
3041            }};
3042        }
3043        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3044        let row_base = r * gpr * Q1_TILE;
3045        let abs_end = bytes.len();
3046        let xp = xq.as_ptr();
3047        let gp = gsum.as_ptr();
3048        let mut accv = vdupq_n_f32(0.0);
3049        let mut gi = 0;
3050        // The second pair load reads 4B past tile gi+3 — stay inside
3051        // the payload slice (only the file's final tiles fall back).
3052        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3053            let t0 = base.add(gi * Q1_TILE);
3054            let ld_a = vld1q_u8(t0);
3055            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3056            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
3057            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
3058            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
3059            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
3060            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
3061            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3062            let g = vld1q_s32(gp.add(gi));
3063            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3064            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3065            let scf: float32x4_t;
3066            asm!(
3067                "fcvtl {o:v}.4s, {i:v}.4h",
3068                o = out(vreg) scf, i = in(vreg) sc16,
3069                options(pure, nomem, nostack),
3070            );
3071            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
3072            gi += 4;
3073        }
3074        let mut acc = vaddvq_f32(accv);
3075        while gi < gpr {
3076            let t = base.add(gi * Q1_TILE);
3077            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3078            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
3079            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
3080            gi += 1;
3081        }
3082        acc
3083    }
3084}
3085
3086/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
3087/// activation streams (prefill amortization — the same idea as the
3088/// AVX2 twin; per stream the group order, fma order and tail match the
3089/// single-row kernel exactly, so batch == matvec bit-for-bit).
3090#[cfg(target_arch = "aarch64")]
3091#[target_feature(enable = "neon,dotprod")]
3092unsafe fn dot_q1_row_1x4_sdot(
3093    bytes: &[u8],
3094    r: usize,
3095    gpr: usize,
3096    xs: [&[i8]; 4],
3097    gs: [&[i32]; 4],
3098) -> [f32; 4] {
3099    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
3100    unsafe {
3101        use core::arch::aarch64::*;
3102        use core::arch::asm;
3103        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3104        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3105        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3106        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3107        const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
3108        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3109        let m = vld1q_u8(MASKS.as_ptr());
3110        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3111        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3112        let isc = vld1_u8(ISC.as_ptr());
3113        macro_rules! sdot2 {
3114            ($w0:expr, $w1:expr, $x:expr) => {{
3115                let x0 = vld1q_s8($x);
3116                let x1 = vld1q_s8($x.add(16));
3117                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3118                asm!(
3119                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3120                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3121                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3122                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
3123                    options(pure, nomem, nostack),
3124                );
3125                vaddq_s32(a0, a1)
3126            }};
3127        }
3128        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3129        let row_base = r * gpr * Q1_TILE;
3130        let abs_end = bytes.len();
3131        let mut accv = [vdupq_n_f32(0.0); 4];
3132        let mut gi = 0;
3133        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3134            let t0 = base.add(gi * Q1_TILE);
3135            let ld_a = vld1q_u8(t0);
3136            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3137            // Unpack ONCE — eight ±mask vectors serve all four streams.
3138            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
3139            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
3140            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
3141            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
3142            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
3143            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
3144            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
3145            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
3146            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3147            let scf: float32x4_t;
3148            asm!(
3149                "fcvtl {o:v}.4s, {i:v}.4h",
3150                o = out(vreg) scf, i = in(vreg) sc16,
3151                options(pure, nomem, nostack),
3152            );
3153            for k in 0..4 {
3154                let xp = xs[k].as_ptr();
3155                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
3156                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
3157                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
3158                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
3159                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3160                let g = vld1q_s32(gs[k].as_ptr().add(gi));
3161                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3162                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
3163            }
3164            gi += 4;
3165        }
3166        let mut acc = [
3167            vaddvq_f32(accv[0]),
3168            vaddvq_f32(accv[1]),
3169            vaddvq_f32(accv[2]),
3170            vaddvq_f32(accv[3]),
3171        ];
3172        while gi < gpr {
3173            let t = base.add(gi * Q1_TILE);
3174            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3175            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
3176            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
3177            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3178            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3179            for k in 0..4 {
3180                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
3181                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
3182            }
3183            gi += 1;
3184        }
3185        acc
3186    }
3187}
3188
3189/// (weight ±1, scale) of one q1 element — the exact outlier term.
3190#[inline]
3191fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3192    let gi = j / GROUP_SIZE;
3193    let k = j % GROUP_SIZE;
3194    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3195    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3196    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
3197    ((bit as i32 * 2 - 1) as f32, s)
3198}
3199
3200/// Exact scalar q1 row (CMF_SDOT=0 contract).
3201#[inline]
3202fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3203    let mut acc = 0f32;
3204    for gi in 0..gpr {
3205        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3206        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3207        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3208        let mut ga = 0f32;
3209        for (j, &b) in tile[2..].iter().enumerate() {
3210            for k in 0..8 {
3211                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
3212            }
3213        }
3214        acc += ga * s;
3215    }
3216    acc
3217}
3218
3219/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
3220/// extracted so multi-matrix jobs drive the same kernel).
3221#[allow(clippy::too_many_arguments)]
3222fn q1_range_a8w8(
3223    bytes: &[u8],
3224    gpr: usize,
3225    act: &SplitAct,
3226    gsum: &[i32],
3227    out: SendMut,
3228    start: usize,
3229    end: usize,
3230) {
3231    for r in start..end {
3232        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3233        for &(j, xv) in &act.outliers {
3234            let (w, s) = q1_outlier(bytes, r, gpr, j);
3235            acc += w * s * xv;
3236        }
3237        // SAFETY: disjoint row ranges per worker.
3238        unsafe { *out.at(r) = acc };
3239    }
3240}
3241
3242/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
3243fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
3244    for r in start..end {
3245        // SAFETY: disjoint row ranges per worker.
3246        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
3247    }
3248}
3249
3250/// q1t per-row overlay locator. After the base (`base_len`) come
3251/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
3252/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
3253/// `(row_ptr offset, entries offset, present)`.
3254fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
3255    let entries = base_len + (rows + 1) * 4;
3256    (base_len, entries, entries <= bytes.len())
3257}
3258
3259/// Read `row_ptr[r]` from the overlay's prefix-sum table.
3260#[inline]
3261fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
3262    let o = rp_off + r * 4;
3263    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
3264}
3265
3266/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
3267/// decoding a q1t code is a table load, not the base-3 divide/modulo per
3268/// weight (division is ~20–40× the cost of a load). Built at compile time.
3269const SIGN5: [[f32; 5]; 256] = {
3270    let mut lut = [[0.0f32; 5]; 256];
3271    let pow3 = [1u16, 3, 9, 27, 81];
3272    let mut byte = 0usize;
3273    while byte < 256 {
3274        let mut i = 0usize;
3275        while i < 5 {
3276            let code = (byte as u16 / pow3[i]) % 3;
3277            lut[byte][i] = if code == 1 {
3278                1.0
3279            } else if code == 2 {
3280                -1.0
3281            } else {
3282                0.0
3283            };
3284            i += 1;
3285        }
3286        byte += 1;
3287    }
3288    lut
3289};
3290
3291/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
3292const SIGN5_I8: [[i8; 5]; 256] = {
3293    let mut lut = [[0i8; 5]; 256];
3294    let pow3 = [1u16, 3, 9, 27, 81];
3295    let mut byte = 0usize;
3296    while byte < 256 {
3297        let mut i = 0usize;
3298        while i < 5 {
3299            let code = (byte as u16 / pow3[i]) % 3;
3300            lut[byte][i] = if code == 1 {
3301                1
3302            } else if code == 2 {
3303                -1
3304            } else {
3305                0
3306            };
3307            i += 1;
3308        }
3309        byte += 1;
3310    }
3311    lut
3312};
3313
3314/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
3315/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
3316/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
3317/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
3318/// buffer is padded to 40). This is the decode/prefill hot inner op.
3319const SIGN5_U64: [u64; 256] = {
3320    let mut lut = [0u64; 256];
3321    let pow3 = [1u16, 3, 9, 27, 81];
3322    let mut byte = 0usize;
3323    while byte < 256 {
3324        let mut v = 0u64;
3325        let mut i = 0usize;
3326        while i < 5 {
3327            let code = (byte as u16 / pow3[i]) % 3;
3328            let s: u8 = if code == 1 {
3329                1
3330            } else if code == 2 {
3331                0xFF
3332            } else {
3333                0
3334            };
3335            v |= (s as u64) << (i * 8);
3336            i += 1;
3337        }
3338        lut[byte] = v;
3339        byte += 1;
3340    }
3341    lut
3342};
3343
3344/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
3345/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
3346/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
3347/// the overlay correction owns that column — no double counting.
3348#[inline]
3349fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
3350    const TILE: usize = cortiq_core::quant::Q1T_TILE;
3351    let off = (r * gpr + j / GROUP_SIZE) * TILE;
3352    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3353    let within = j % GROUP_SIZE;
3354    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
3355}
3356
3357/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
3358/// (integer accumulation is order-independent).
3359#[cfg(target_arch = "aarch64")]
3360#[target_feature(enable = "neon,dotprod")]
3361#[inline]
3362unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
3363    // SAFETY: caller guarantees 32 readable i8 at each pointer.
3364    unsafe {
3365        use core::arch::aarch64::*;
3366        use core::arch::asm;
3367        let w0 = vld1q_s8(w);
3368        let w1 = vld1q_s8(w.add(16));
3369        let x0 = vld1q_s8(x);
3370        let x1 = vld1q_s8(x.add(16));
3371        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3372        asm!(
3373            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3374            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3375            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3376            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3377            options(pure, nomem, nostack),
3378        );
3379        vaddvq_s32(vaddq_s32(a0, a1))
3380    }
3381}
3382
3383/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
3384/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
3385#[cfg(target_arch = "x86_64")]
3386#[target_feature(enable = "avx2")]
3387#[inline]
3388unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
3389    // SAFETY: caller guarantees 32 readable i8 at each pointer.
3390    unsafe {
3391        use core::arch::x86_64::*;
3392        let wv = _mm256_loadu_si256(w as *const __m256i);
3393        let xv = _mm256_loadu_si256(x as *const __m256i);
3394        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3395        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
3396        let hi128 = _mm256_extracti128_si256::<1>(d);
3397        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3398        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3399        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3400        _mm_cvtsi128_si32(s32)
3401    }
3402}
3403
3404/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
3405/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
3406/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
3407/// overwritten by the next; the final 6 padding bytes are unused by the dot.
3408#[inline]
3409fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
3410    debug_assert!(dst.len() >= 40);
3411    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
3412    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
3413    unsafe {
3414        let p = dst.as_mut_ptr();
3415        for bi in 0..7 {
3416            core::ptr::write_unaligned(
3417                p.add(bi * 5) as *mut u64,
3418                SIGN5_U64[*codes.add(bi) as usize],
3419            );
3420        }
3421    }
3422}
3423
3424/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
3425/// row's signs are unpacked once and dotted against every batch input).
3426/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
3427/// reachable; the scalar arm is a non-SIMD-arch fallback.
3428#[inline]
3429fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
3430    #[cfg(target_arch = "aarch64")]
3431    unsafe {
3432        return sdot32_i8(w, x);
3433    }
3434    #[cfg(target_arch = "x86_64")]
3435    unsafe {
3436        return i8dot32_avx2(w, x);
3437    }
3438    #[allow(unreachable_code)]
3439    unsafe {
3440        let mut s = 0i32;
3441        for k in 0..GROUP_SIZE {
3442            s += *w.add(k) as i32 * *x.add(k) as i32;
3443        }
3444        s
3445    }
3446}
3447
3448/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
3449/// `sx`). Signs unpack base-3 → a 32-i8 stack buffer, then one int8 dot per
3450/// group. ARM SDOT.
3451#[cfg(target_arch = "aarch64")]
3452#[target_feature(enable = "neon,dotprod")]
3453unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3454    // SAFETY: 9-byte tile per group, xq.len() == gpr·GROUP_SIZE.
3455    unsafe {
3456        const TILE: usize = cortiq_core::quant::Q1T_TILE;
3457        let mut acc = 0f32;
3458        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
3459        for gi in 0..gpr {
3460            let off = (r * gpr + gi) * TILE;
3461            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3462            q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3463            acc += sdot32_i8(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3464        }
3465        acc
3466    }
3467}
3468
3469/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
3470#[cfg(target_arch = "x86_64")]
3471#[target_feature(enable = "avx2")]
3472unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3473    // SAFETY: as q1t_dot_row_sdot.
3474    unsafe {
3475        const TILE: usize = cortiq_core::quant::Q1T_TILE;
3476        let mut acc = 0f32;
3477        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
3478        for gi in 0..gpr {
3479            let off = (r * gpr + gi) * TILE;
3480            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3481            q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3482            acc += i8dot32_avx2(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3483        }
3484        acc
3485    }
3486}
3487
3488/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
3489/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
3490/// reachable.
3491#[inline]
3492fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3493    #[cfg(target_arch = "aarch64")]
3494    unsafe {
3495        return q1t_dot_row_sdot(bytes, r, gpr, xq);
3496    }
3497    #[cfg(target_arch = "x86_64")]
3498    unsafe {
3499        return q1t_dot_row_avx2(bytes, r, gpr, xq);
3500    }
3501    #[allow(unreachable_code)]
3502    {
3503        const TILE: usize = cortiq_core::quant::Q1T_TILE;
3504        let mut acc = 0f32;
3505        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
3506        for gi in 0..gpr {
3507            let off = (r * gpr + gi) * TILE;
3508            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3509            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
3510            let mut d = 0i32;
3511            for k in 0..GROUP_SIZE {
3512                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
3513            }
3514            acc += d as f32 * s;
3515        }
3516        acc
3517    }
3518}
3519
3520/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
3521/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
3522/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
3523/// base contributes nothing there and this is a plain `value·x`, not
3524/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
3525/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
3526fn q1t_row_outlier_correction(
3527    bytes: &[u8],
3528    r: usize,
3529    rp_off: usize,
3530    entries_off: usize,
3531    has_ov: bool,
3532    x: &[f32],
3533) -> f32 {
3534    if !has_ov {
3535        return 0.0;
3536    }
3537    let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3538    let mut corr = 0f32;
3539    for p in c0..c1 {
3540        let e = entries_off + p * 4;
3541        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3542        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3543        corr += val * x[col];
3544    }
3545    corr
3546}
3547
3548/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
3549/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
3550/// Used by the batched (prefill) path where the decode amortizes over the batch.
3551fn q1t_dequant_row(
3552    bytes: &[u8],
3553    r: usize,
3554    gpr: usize,
3555    rp_off: usize,
3556    entries_off: usize,
3557    has_ov: bool,
3558    buf: &mut [f32],
3559) {
3560    const TILE: usize = cortiq_core::quant::Q1T_TILE;
3561    for g in 0..gpr {
3562        let off = (r * gpr + g) * TILE;
3563        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3564        let codes = &bytes[off + 2..off + TILE];
3565        let bc = g * GROUP_SIZE;
3566        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
3567        for bi in 0..6 {
3568            let lut = &SIGN5[codes[bi] as usize];
3569            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
3570            for i in 0..5 {
3571                d[i] = lut[i] * s;
3572            }
3573        }
3574        let lut = &SIGN5[codes[6] as usize];
3575        buf[bc + 30] = lut[0] * s;
3576        buf[bc + 31] = lut[1] * s;
3577    }
3578    if !has_ov {
3579        return;
3580    }
3581    let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3582    for p in c0..c1 {
3583        let e = entries_off + p * 4;
3584        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3585        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3586    }
3587}
3588
3589/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
3590/// computes the ternary base; the overlay stays on the CPU — its entries are
3591/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
3592fn q1t_add_overlay(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3593    const TILE: usize = cortiq_core::quant::Q1T_TILE;
3594    let gpr = cols / GROUP_SIZE;
3595    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3596    if !has_ov {
3597        return;
3598    }
3599    let out_addr = SendMut(out.as_mut_ptr());
3600    let run = move |start: usize, end: usize| {
3601        for r in start..end {
3602            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3603            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
3604            unsafe { *out_addr.at(r) += corr };
3605        }
3606    };
3607    dispatch_rows(pool, rows, &run);
3608}
3609
3610/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
3611/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
3612/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
3613fn q1t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3614    debug_assert_eq!(out.len(), rows);
3615    const TILE: usize = cortiq_core::quant::Q1T_TILE;
3616    let gpr = cols / GROUP_SIZE;
3617    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3618    let out_addr = SendMut(out.as_mut_ptr());
3619    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
3620    // (`split_act`), activation outliers added back exactly in f32, weight
3621    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
3622    if a8w8_enabled() {
3623        let act = split_act(x);
3624        let act = &act;
3625        let run = move |start: usize, end: usize| {
3626            for r in start..end {
3627                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3628                for &(j, xv) in &act.outliers {
3629                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3630                }
3631                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3632                // SAFETY: disjoint row ranges per worker.
3633                unsafe { *out_addr.at(r) = acc };
3634            }
3635        };
3636        dispatch_rows(pool, rows, &run);
3637        return;
3638    }
3639    let run = move |start: usize, end: usize| {
3640        // Per-group signs, unpacked contiguously so the dot below is a clean
3641        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
3642        // 5-values-per-byte base-3 layout won't SIMD in place.
3643        let mut sg = [0f32; GROUP_SIZE];
3644        for r in start..end {
3645            let mut acc = 0f32;
3646            for g in 0..gpr {
3647                let off = (r * gpr + g) * TILE;
3648                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3649                let codes = &bytes[off + 2..off + TILE];
3650                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
3651                for bi in 0..6 {
3652                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
3653                }
3654                let lut = &SIGN5[codes[6] as usize];
3655                sg[30] = lut[0];
3656                sg[31] = lut[1];
3657                let mut gsum = 0f32;
3658                for k in 0..GROUP_SIZE {
3659                    gsum += sg[k] * xg[k];
3660                }
3661                acc += s * gsum;
3662            }
3663            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3664            unsafe { *out_addr.at(r) = acc };
3665        }
3666    };
3667    dispatch_rows(pool, rows, &run);
3668}
3669
3670/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
3671/// batch against it (amortizes the per-row decode).
3672fn q1t_matmat(
3673    bytes: &[u8],
3674    xs: &[f32],
3675    b: usize,
3676    rows: usize,
3677    cols: usize,
3678    out: &mut [f32],
3679    pool: Option<&Pool>,
3680) {
3681    debug_assert_eq!(out.len(), b * rows);
3682    const TILE: usize = cortiq_core::quant::Q1T_TILE;
3683    let gpr = cols / GROUP_SIZE;
3684    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3685    let out_addr = SendMut(out.as_mut_ptr());
3686    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
3687    // each weight row's signs to i8 ONCE, then int8-dot against every input —
3688    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
3689    if a8w8_enabled() {
3690        let acts: Vec<SplitAct> = (0..b)
3691            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
3692            .collect();
3693        let acts = &acts;
3694        let run = move |start: usize, end: usize| {
3695            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
3696            let mut sc = vec![0f32; gpr]; // per-group scales
3697            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
3698            for r in start..end {
3699                for g in 0..gpr {
3700                    let off = (r * gpr + g) * TILE;
3701                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3702                    q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg[g * GROUP_SIZE..]);
3703                }
3704                for bi in 0..b {
3705                    let act = &acts[bi];
3706                    let mut isum = 0f32;
3707                    for g in 0..gpr {
3708                        let d = q1t_i8dot32(
3709                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
3710                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
3711                        );
3712                        isum += d as f32 * sc[g];
3713                    }
3714                    let mut acc = isum * act.sx;
3715                    for &(j, xv) in &act.outliers {
3716                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3717                    }
3718                    accs[bi] = acc;
3719                }
3720                // Overlay ONCE per row for the whole batch: read each (col, val)
3721                // from mmap a single time (was b× — the re-read dominated prefill)
3722                // and fan it out over the batch via the cached inputs.
3723                if has_ov {
3724                    let (c0, c1) =
3725                        (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3726                    for p in c0..c1 {
3727                        let e = ent_off + p * 4;
3728                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3729                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3730                        for bi in 0..b {
3731                            accs[bi] += val * xs[bi * cols + col];
3732                        }
3733                    }
3734                }
3735                for bi in 0..b {
3736                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
3737                }
3738            }
3739        };
3740        dispatch_rows(pool, rows, &run);
3741        return;
3742    }
3743    let run = move |start: usize, end: usize| {
3744        let mut buf = vec![0f32; cols];
3745        for r in start..end {
3746            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
3747            for bi in 0..b {
3748                let xr = &xs[bi * cols..(bi + 1) * cols];
3749                let mut acc = 0f32;
3750                for j in 0..cols {
3751                    acc += buf[j] * xr[j];
3752                }
3753                unsafe { *out_addr.at(bi * rows + r) = acc };
3754            }
3755        }
3756    };
3757    dispatch_rows(pool, rows, &run);
3758}
3759
3760fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3761    debug_assert_eq!(out.len(), rows);
3762    let gpr = cols / GROUP_SIZE;
3763    let out_addr = SendMut(out.as_mut_ptr());
3764    if a8w8_enabled() {
3765        let act = split_act(x);
3766        let gsum = q1_group_sums(&act.xq, gpr);
3767        let (act, gsum) = (&act, &gsum);
3768        let run =
3769            move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
3770        dispatch_rows(pool, rows, &run);
3771        return;
3772    }
3773    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
3774    dispatch_rows(pool, rows, &run);
3775}
3776
3777/// Fused two-input q1 matvec (weights read once per pair).
3778#[allow(clippy::too_many_arguments)]
3779fn q1_matvec2(
3780    bytes: &[u8],
3781    x1: &[f32],
3782    x2: &[f32],
3783    rows: usize,
3784    cols: usize,
3785    o1: &mut [f32],
3786    o2: &mut [f32],
3787    pool: Option<&Pool>,
3788) {
3789    let gpr = cols / GROUP_SIZE;
3790    let p1 = SendMut(o1.as_mut_ptr());
3791    let p2 = SendMut(o2.as_mut_ptr());
3792    if a8w8_enabled() {
3793        let a1 = split_act(x1);
3794        let a2 = split_act(x2);
3795        let g1 = q1_group_sums(&a1.xq, gpr);
3796        let g2 = q1_group_sums(&a2.xq, gpr);
3797        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
3798        let run = move |start: usize, end: usize| {
3799            for r in start..end {
3800                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
3801                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
3802                for &(j, xv) in &a1.outliers {
3803                    let (w, s) = q1_outlier(bytes, r, gpr, j);
3804                    v1 += w * s * xv;
3805                }
3806                for &(j, xv) in &a2.outliers {
3807                    let (w, s) = q1_outlier(bytes, r, gpr, j);
3808                    v2 += w * s * xv;
3809                }
3810                // SAFETY: disjoint row ranges per worker.
3811                unsafe {
3812                    *p1.at(r) = v1;
3813                    *p2.at(r) = v2;
3814                }
3815            }
3816        };
3817        dispatch_rows(pool, rows, &run);
3818        return;
3819    }
3820    let run = move |start: usize, end: usize| {
3821        for r in start..end {
3822            // SAFETY: disjoint row ranges per worker.
3823            unsafe {
3824                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
3825                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
3826            }
3827        }
3828    };
3829    dispatch_rows(pool, rows, &run);
3830}
3831
3832/// Batched q1 matmat: each row's tiles stream once per microbatch.
3833#[allow(clippy::too_many_arguments)]
3834fn q1_matmat(
3835    bytes: &[u8],
3836    xs_all: &[f32],
3837    b: usize,
3838    rows: usize,
3839    cols: usize,
3840    out: &mut [f32],
3841    pool: Option<&Pool>,
3842) {
3843    debug_assert_eq!(out.len(), b * rows);
3844    let gpr = cols / GROUP_SIZE;
3845    let out_addr = SendMut(out.as_mut_ptr());
3846    if a8w8_enabled() {
3847        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
3848            .map(|bi| {
3849                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
3850                let gsum = q1_group_sums(&act.xq, gpr);
3851                (act, gsum)
3852            })
3853            .collect();
3854        let acts = &acts;
3855        #[cfg(target_arch = "x86_64")]
3856        let blocked_ok = avx2_enabled()
3857            && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3858        #[cfg(target_arch = "aarch64")]
3859        let blocked_ok = sdot_enabled()
3860            && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3861        let run = move |start: usize, end: usize| {
3862            for r in start..end {
3863                let mut bi = 0usize;
3864                // Blocked 1×4: the unpacked bit mask serves four
3865                // activation streams per group.
3866                #[cfg(target_arch = "aarch64")]
3867                if blocked_ok {
3868                    while bi + 4 <= acts.len() {
3869                        let xs = [
3870                            acts[bi].0.xq.as_slice(),
3871                            acts[bi + 1].0.xq.as_slice(),
3872                            acts[bi + 2].0.xq.as_slice(),
3873                            acts[bi + 3].0.xq.as_slice(),
3874                        ];
3875                        let gs = [
3876                            acts[bi].1.as_slice(),
3877                            acts[bi + 1].1.as_slice(),
3878                            acts[bi + 2].1.as_slice(),
3879                            acts[bi + 3].1.as_slice(),
3880                        ];
3881                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
3882                        for k in 0..4 {
3883                            let (act, _) = &acts[bi + k];
3884                            let mut acc = d[k] * act.sx;
3885                            for &(j, xv) in &act.outliers {
3886                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
3887                                acc += w * sc * xv;
3888                            }
3889                            // SAFETY: disjoint (bi, r) cells per worker.
3890                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3891                        }
3892                        bi += 4;
3893                    }
3894                }
3895                #[cfg(target_arch = "x86_64")]
3896                if blocked_ok {
3897                    while bi + 4 <= acts.len() {
3898                        let xs = [
3899                            acts[bi].0.xq.as_slice(),
3900                            acts[bi + 1].0.xq.as_slice(),
3901                            acts[bi + 2].0.xq.as_slice(),
3902                            acts[bi + 3].0.xq.as_slice(),
3903                        ];
3904                        let gs = [
3905                            acts[bi].1.as_slice(),
3906                            acts[bi + 1].1.as_slice(),
3907                            acts[bi + 2].1.as_slice(),
3908                            acts[bi + 3].1.as_slice(),
3909                        ];
3910                        let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
3911                        for k in 0..4 {
3912                            let (act, _) = &acts[bi + k];
3913                            let mut acc = d[k] * act.sx;
3914                            for &(j, xv) in &act.outliers {
3915                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
3916                                acc += w * sc * xv;
3917                            }
3918                            // SAFETY: disjoint (bi, r) cells per worker.
3919                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3920                        }
3921                        bi += 4;
3922                    }
3923                }
3924                while bi < acts.len() {
3925                    let (act, gsum) = &acts[bi];
3926                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3927                    for &(j, xv) in &act.outliers {
3928                        let (w, s) = q1_outlier(bytes, r, gpr, j);
3929                        acc += w * s * xv;
3930                    }
3931                    // SAFETY: disjoint (bi, r) cells per worker range.
3932                    unsafe { *out_addr.at(bi * rows + r) = acc };
3933                    bi += 1;
3934                }
3935            }
3936        };
3937        dispatch_rows(pool, rows, &run);
3938        return;
3939    }
3940    let run = move |start: usize, end: usize| {
3941        for r in start..end {
3942            for bi in 0..b {
3943                let x = &xs_all[bi * cols..(bi + 1) * cols];
3944                // SAFETY: disjoint (bi, r) cells per worker range.
3945                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3946            }
3947        }
3948    };
3949    dispatch_rows(pool, rows, &run);
3950}
3951
3952/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
3953/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
3954/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
3955/// 32-group, exact outlier correction — the same A8W8 contract as q8.
3956/// `CMF_SDOT=0` keeps the exact scalar path.
3957fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3958    debug_assert_eq!(out.len(), rows);
3959    let (packed, scales) = q4_split(bytes, rows, cols);
3960    let gpr = cols / GROUP_SIZE;
3961    let out_addr = SendMut(out.as_mut_ptr());
3962
3963    if a8w8_enabled() {
3964        let act = split_act(x);
3965        let run = move |start: usize, end: usize| {
3966            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3967        };
3968        dispatch_rows(pool, rows, &run);
3969        return;
3970    }
3971
3972    let run = move |start: usize, end: usize| {
3973        q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3974    };
3975    dispatch_rows(pool, rows, &run);
3976}
3977
3978/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
3979/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
3980#[inline]
3981#[allow(unreachable_code)]
3982/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
3983/// streams: the 32-byte weight chunk and its abs() load once per group,
3984/// the per-group f16 scale decodes once — four maddubs+reduce chains
3985/// instead of four full (load, abs, dot) rounds.
3986#[cfg(target_arch = "x86_64")]
3987#[target_feature(enable = "avx2")]
3988unsafe fn dot_q4b_row_1x4_avx2(
3989    buf: &[u8],
3990    scales: &[u8],
3991    g0: usize,
3992    gpr: usize,
3993    xs: [&[i8]; 4],
3994) -> [f32; 4] {
3995    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
3996    unsafe {
3997        use core::arch::x86_64::*;
3998        let ones = _mm256_set1_epi16(1);
3999        let mut acc = [0f32; 4];
4000        for gi in 0..gpr {
4001            let s = f16_to_f32(u16::from_le_bytes([
4002                scales[(g0 + gi) * 2],
4003                scales[(g0 + gi) * 2 + 1],
4004            ]));
4005            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4006            let aw = _mm256_abs_epi8(w);
4007            for (k, xq) in xs.iter().enumerate() {
4008                let x =
4009                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4010                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4011                let d = _mm256_madd_epi16(p16, ones);
4012                let hi128 = _mm256_extracti128_si256::<1>(d);
4013                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4014                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4015                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4016                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
4017            }
4018        }
4019        acc
4020    }
4021}
4022
4023/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
4024/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
4025/// accumulation order (the q4_block flavor applies sx once at the end,
4026/// matching ITS single path; the two conventions are historical and
4027/// each blocked leg must mirror its own).
4028#[cfg(target_arch = "x86_64")]
4029#[target_feature(enable = "avx2")]
4030unsafe fn dot_q4b_row_1x4_sx_avx2(
4031    buf: &[u8],
4032    scales: &[u8],
4033    g0: usize,
4034    gpr: usize,
4035    xs: [&[i8]; 4],
4036    sxs: [f32; 4],
4037) -> [f32; 4] {
4038    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
4039    unsafe {
4040        use core::arch::x86_64::*;
4041        let ones = _mm256_set1_epi16(1);
4042        let mut acc = [0f32; 4];
4043        for gi in 0..gpr {
4044            let s = f16_to_f32(u16::from_le_bytes([
4045                scales[(g0 + gi) * 2],
4046                scales[(g0 + gi) * 2 + 1],
4047            ]));
4048            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4049            let aw = _mm256_abs_epi8(w);
4050            for (k, xq) in xs.iter().enumerate() {
4051                let x =
4052                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4053                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4054                let d = _mm256_madd_epi16(p16, ones);
4055                let hi128 = _mm256_extracti128_si256::<1>(d);
4056                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4057                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4058                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4059                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
4060            }
4061        }
4062        acc
4063    }
4064}
4065
4066fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4067    #[cfg(target_arch = "aarch64")]
4068    unsafe {
4069        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
4070    }
4071    #[cfg(target_arch = "x86_64")]
4072    unsafe {
4073        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
4074    }
4075    let mut acc = 0f32;
4076    for gi in 0..gpr {
4077        let g = g0 + gi;
4078        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4079        let mut d = 0i32;
4080        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4081            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4082                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4083        }
4084        acc += d as f32 * s;
4085    }
4086    acc
4087}
4088
4089/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
4090#[inline]
4091#[allow(unreachable_code)]
4092fn dot_q4_row_i8_2(
4093    packed: &[u8],
4094    scales: &[u8],
4095    g0: usize,
4096    gpr: usize,
4097    xq1: &[i8],
4098    xq2: &[i8],
4099) -> (f32, f32) {
4100    #[cfg(target_arch = "aarch64")]
4101    unsafe {
4102        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
4103    }
4104    #[cfg(target_arch = "x86_64")]
4105    unsafe {
4106        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
4107    }
4108    (
4109        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
4110        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
4111    )
4112}
4113
4114/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
4115/// multi-matrix jobs can drive it for several tensors in one dispatch).
4116#[allow(clippy::too_many_arguments)]
4117fn q4_range_a8w8(
4118    packed: &[u8],
4119    scales: &[u8],
4120    gpr: usize,
4121    cols: usize,
4122    act: &SplitAct,
4123    out: SendMut,
4124    start: usize,
4125    end: usize,
4126) {
4127    for r in start..end {
4128        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
4129        // xq is zeroed at outlier slots — add the exact terms.
4130        for &(j, xv) in &act.outliers {
4131            let flat = r * cols + j;
4132            let byte = packed[flat / 2];
4133            let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4134            let s = f16_to_f32(u16::from_le_bytes([
4135                scales[(flat / GROUP_SIZE) * 2],
4136                scales[(flat / GROUP_SIZE) * 2 + 1],
4137            ]));
4138            acc += ((nib as i32 - 8) as f32) * s * xv;
4139        }
4140        // SAFETY: disjoint row ranges per worker.
4141        unsafe { *out.at(r) = acc };
4142    }
4143}
4144
4145/// Two-input q4 row range via the A8W8 int8 path — kernel body of
4146/// `q4matvec2`, extracted for pair multi-matrix jobs.
4147#[allow(clippy::too_many_arguments)]
4148fn q4_range2_a8w8(
4149    packed: &[u8],
4150    scales: &[u8],
4151    gpr: usize,
4152    cols: usize,
4153    a1: &SplitAct,
4154    a2: &SplitAct,
4155    p1: SendMut,
4156    p2: SendMut,
4157    start: usize,
4158    end: usize,
4159) {
4160    for r in start..end {
4161        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
4162        let mut acc1 = s1 * a1.sx;
4163        let mut acc2 = s2 * a2.sx;
4164        // xq is zeroed at outlier slots — add the exact terms.
4165        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
4166            for &(j, xv) in outliers {
4167                let flat = r * cols + j;
4168                let byte = packed[flat / 2];
4169                let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4170                let s = f16_to_f32(u16::from_le_bytes([
4171                    scales[(flat / GROUP_SIZE) * 2],
4172                    scales[(flat / GROUP_SIZE) * 2 + 1],
4173                ]));
4174                *acc += ((nib as i32 - 8) as f32) * s * xv;
4175            }
4176        };
4177        fix(&a1.outliers, &mut acc1);
4178        fix(&a2.outliers, &mut acc2);
4179        // SAFETY: disjoint row ranges per worker.
4180        unsafe {
4181            *p1.at(r) = acc1;
4182            *p2.at(r) = acc2;
4183        }
4184    }
4185}
4186
4187/// Exact scalar q4 row range (same extraction, non-SDOT path).
4188fn q4_range_f32(
4189    packed: &[u8],
4190    scales: &[u8],
4191    gpr: usize,
4192    x: &[f32],
4193    out: SendMut,
4194    start: usize,
4195    end: usize,
4196) {
4197    for r in start..end {
4198        let mut acc = 0f32;
4199        for gi in 0..gpr {
4200            let g = r * gpr + gi;
4201            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4202            let pk = &packed[g * 16..(g + 1) * 16];
4203            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4204            let mut ga = 0f32;
4205            for (k, &b) in pk.iter().enumerate() {
4206                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4207                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4208            }
4209            acc += ga * s;
4210        }
4211        // SAFETY: disjoint row ranges per worker.
4212        unsafe { *out.at(r) = acc };
4213    }
4214}
4215
4216/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
4217/// dotted against both activations (was: two full matvecs — double
4218/// weight traffic). Per-lane math matches `q4matvec` exactly.
4219#[allow(clippy::too_many_arguments)]
4220fn q4matvec2(
4221    bytes: &[u8],
4222    x1: &[f32],
4223    x2: &[f32],
4224    rows: usize,
4225    cols: usize,
4226    o1: &mut [f32],
4227    o2: &mut [f32],
4228    pool: Option<&Pool>,
4229) {
4230    debug_assert_eq!(o1.len(), rows);
4231    debug_assert_eq!(o2.len(), rows);
4232    let (packed, scales) = q4_split(bytes, rows, cols);
4233    let gpr = cols / GROUP_SIZE;
4234
4235    if a8w8_enabled() {
4236        let a1 = split_act(x1);
4237        let a2 = split_act(x2);
4238        let p1 = SendMut(o1.as_mut_ptr());
4239        let p2 = SendMut(o2.as_mut_ptr());
4240        let run = move |start: usize, end: usize| {
4241            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
4242        };
4243        dispatch_rows(pool, rows, &run);
4244        return;
4245    }
4246
4247    let p1 = SendMut(o1.as_mut_ptr());
4248    let p2 = SendMut(o2.as_mut_ptr());
4249    let run = move |start: usize, end: usize| {
4250        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
4251    };
4252    dispatch_rows(pool, rows, &run);
4253}
4254
4255/// Two-input exact scalar q4 row range (same extraction).
4256#[allow(clippy::too_many_arguments)]
4257fn q4_range2_f32(
4258    packed: &[u8],
4259    scales: &[u8],
4260    gpr: usize,
4261    x1: &[f32],
4262    x2: &[f32],
4263    p1: SendMut,
4264    p2: SendMut,
4265    start: usize,
4266    end: usize,
4267) {
4268    for r in start..end {
4269        let (mut acc1, mut acc2) = (0f32, 0f32);
4270        for gi in 0..gpr {
4271            let g = r * gpr + gi;
4272            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4273            let pk = &packed[g * 16..(g + 1) * 16];
4274            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4275            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4276            let (mut g1, mut g2) = (0f32, 0f32);
4277            for (k, &b) in pk.iter().enumerate() {
4278                let wl = (b & 0x0F) as f32 - 8.0;
4279                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
4280                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
4281                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
4282            }
4283            acc1 += g1 * s;
4284            acc2 += g2 * s;
4285        }
4286        // SAFETY: disjoint row ranges per worker.
4287        unsafe {
4288            *p1.at(r) = acc1;
4289            *p2.at(r) = acc2;
4290        }
4291    }
4292}
4293
4294thread_local! {
4295    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
4296    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
4297    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
4298    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4299}
4300
4301/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
4302/// and dotted against ALL b activations (prefill used to fall back to b
4303/// full matvecs — b× weight traffic and b× nibble decode). Per-position
4304/// math matches `q4matvec` exactly: same group order, same accumulation.
4305/// `out` is row-major [b, rows] like `qmatmat`.
4306#[allow(clippy::too_many_arguments)]
4307fn q4matmat(
4308    bytes: &[u8],
4309    xs_all: &[f32],
4310    b: usize,
4311    rows: usize,
4312    cols: usize,
4313    out: &mut [f32],
4314    pool: Option<&Pool>,
4315) {
4316    debug_assert_eq!(xs_all.len(), b * cols);
4317    debug_assert_eq!(out.len(), b * rows);
4318    let (packed, scales) = q4_split(bytes, rows, cols);
4319    let gpr = cols / GROUP_SIZE;
4320    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4321
4322    if a8w8_enabled() {
4323        let acts: Vec<SplitAct> =
4324            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4325        let acts = &acts;
4326        let out_addr = SendMut(out.as_mut_ptr());
4327        let run = move |start: usize, end: usize| {
4328            ROW_I8.with(|rb| {
4329                let mut buf = rb.borrow_mut();
4330                buf.resize(cols, 0);
4331                for r in start..end {
4332                    // Unpack the row's nibbles to centered i8 once
4333                    // (element 2k = low nibble, 2k+1 = high — flat order,
4334                    // same as dot_q4_row_sdot's zip).
4335                    for gi in 0..gpr {
4336                        let g = r * gpr + gi;
4337                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4338                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
4339                            buf[gi * GROUP_SIZE + k * 2 + 1] =
4340                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
4341                        }
4342                    }
4343                    let mut bi = 0usize;
4344                    #[cfg(target_arch = "x86_64")]
4345                    if avx2_enabled()
4346                        && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4347                    {
4348                        while bi + 4 <= acts.len() {
4349                            let xs = [
4350                                acts[bi].xq.as_slice(),
4351                                acts[bi + 1].xq.as_slice(),
4352                                acts[bi + 2].xq.as_slice(),
4353                                acts[bi + 3].xq.as_slice(),
4354                            ];
4355                            let d = unsafe {
4356                                dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
4357                            };
4358                            for k in 0..4 {
4359                                let act = &acts[bi + k];
4360                                let mut acc = d[k] * act.sx;
4361                                for &(j, xv) in &act.outliers {
4362                                    acc += (buf[j] as i8) as f32
4363                                        * gscale((r * cols + j) / GROUP_SIZE)
4364                                        * xv;
4365                                }
4366                                // SAFETY: disjoint (bi, r) cells per worker.
4367                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4368                            }
4369                            bi += 4;
4370                        }
4371                    }
4372                    while bi < acts.len() {
4373                        let act = &acts[bi];
4374                        let mut acc = 0f32;
4375                        for gi in 0..gpr {
4376                            let d = dot_i8_i8(
4377                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4378                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4379                            );
4380                            acc += d as f32 * gscale(r * gpr + gi);
4381                        }
4382                        acc *= act.sx;
4383                        // xq is zeroed at outlier slots — exact terms.
4384                        for &(j, xv) in &act.outliers {
4385                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
4386                        }
4387                        // SAFETY: disjoint (bi, r) cells per worker row range.
4388                        unsafe { *out_addr.at(bi * rows + r) = acc };
4389                        bi += 1;
4390                    }
4391                }
4392            })
4393        };
4394        dispatch_rows(pool, rows, &run);
4395        return;
4396    }
4397
4398    let out_addr = SendMut(out.as_mut_ptr());
4399    let run = move |start: usize, end: usize| {
4400        ROW_F32.with(|rb| {
4401            let mut buf = rb.borrow_mut();
4402            buf.resize(cols, 0.0);
4403            for r in start..end {
4404                // Decode raw (nib − 8) values once; scales stay per-group
4405                // so the accumulation order matches q4matvec bit-for-bit.
4406                for gi in 0..gpr {
4407                    let g = r * gpr + gi;
4408                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4409                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
4410                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
4411                    }
4412                }
4413                for bi in 0..b {
4414                    let x = &xs_all[bi * cols..(bi + 1) * cols];
4415                    let mut acc = 0f32;
4416                    for gi in 0..gpr {
4417                        let mut ga = 0f32;
4418                        // Pairwise (lo + hi) addition, matching
4419                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
4420                        // a flat one-per-element loop rounds differently
4421                        // and broke bit-parity on the scalar (x86) path.
4422                        for k in 0..GROUP_SIZE / 2 {
4423                            let e = gi * GROUP_SIZE + k * 2;
4424                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
4425                        }
4426                        acc += ga * gscale(r * gpr + gi);
4427                    }
4428                    // SAFETY: disjoint (bi, r) cells per worker row range.
4429                    unsafe { *out_addr.at(bi * rows + r) = acc };
4430                }
4431            }
4432        })
4433    };
4434    dispatch_rows(pool, rows, &run);
4435}
4436
4437/// Batched vbit matmat: each variable-bit row is decoded from the mmap
4438/// ONCE for the whole microbatch. Same per-position math as
4439/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
4440/// and the scalar path).
4441#[allow(clippy::too_many_arguments)]
4442fn vbitmatmat(
4443    bytes: &[u8],
4444    offsets: &[usize],
4445    xs_all: &[f32],
4446    b: usize,
4447    rows: usize,
4448    cols: usize,
4449    out: &mut [f32],
4450    pool: Option<&Pool>,
4451) {
4452    debug_assert_eq!(xs_all.len(), b * cols);
4453    debug_assert_eq!(out.len(), b * rows);
4454    debug_assert_eq!(offsets.len(), rows + 1);
4455    let ng = cols / GROUP_SIZE;
4456    let bits = &bytes[..rows];
4457    let sc_off = rows;
4458    let gscale = |r: usize, g: usize| {
4459        let so = (r * ng + g) * 2;
4460        f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
4461    };
4462
4463    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
4464    let decode_f32 = |r: usize, dst: &mut [f32]| {
4465        let bw = bits[r] as usize;
4466        let l = ((1i32 << (bw - 1)) - 1) as f32;
4467        let data = &bytes[offsets[r]..offsets[r + 1]];
4468        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4469        for d in dst.iter_mut() {
4470            while nbits < bw {
4471                acc = (acc << 8) | data[idx] as u64;
4472                idx += 1;
4473                nbits += 8;
4474            }
4475            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
4476            nbits -= bw;
4477            *d = u - l;
4478        }
4479    };
4480
4481    if a8w8_enabled() {
4482        let acts: Vec<SplitAct> =
4483            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4484        let acts = &acts;
4485        let out_addr = SendMut(out.as_mut_ptr());
4486        let run = move |start: usize, end: usize| {
4487            for r in start..end {
4488                let bw = bits[r] as usize;
4489                if bw == 8 {
4490                    // u−L reaches 128 → no i8 path; decode once, exact
4491                    // f32 dots for every position (same as vbitmatvec).
4492                    ROW_F32.with(|rb| {
4493                        let mut buf = rb.borrow_mut();
4494                        buf.resize(cols, 0.0);
4495                        decode_f32(r, &mut buf);
4496                        for bi in 0..b {
4497                            let x = &xs_all[bi * cols..(bi + 1) * cols];
4498                            let mut dot = 0f32;
4499                            for g in 0..ng {
4500                                let mut gd = 0f32;
4501                                for k in 0..GROUP_SIZE {
4502                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4503                                }
4504                                dot += gd * gscale(r, g);
4505                            }
4506                            // SAFETY: disjoint (bi, r) cells per worker range.
4507                            unsafe { *out_addr.at(bi * rows + r) = dot };
4508                        }
4509                    });
4510                    continue;
4511                }
4512                let l = (1i32 << (bw - 1)) - 1;
4513                let data = &bytes[offsets[r]..offsets[r + 1]];
4514                ROW_I8.with(|rb| {
4515                    let mut buf = rb.borrow_mut();
4516                    buf.resize(cols, 0);
4517                    #[inline(always)]
4518                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4519                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4520                            let u = unpack8::<B>(&data[blk * B..]);
4521                            for k in 0..8 {
4522                                chunk[k] = (u[k] - l) as i8 as u8;
4523                            }
4524                        }
4525                    }
4526                    match bw {
4527                        3 => fill::<3>(data, l, &mut buf),
4528                        4 => vbit_fill4(data, &mut buf),
4529                        5 => fill::<5>(data, l, &mut buf),
4530                        6 => fill::<6>(data, l, &mut buf),
4531                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
4532                    }
4533                    let mut bi = 0usize;
4534                    // The vbit scale table shares q4_block's layout
4535                    // (contiguous f16 per (row·ng + g)), so the same
4536                    // blocked 1×4 kernel serves the decoded row.
4537                    #[cfg(target_arch = "x86_64")]
4538                    if avx2_enabled()
4539                        && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4540                    {
4541                        while bi + 4 <= acts.len() {
4542                            let xs = [
4543                                acts[bi].xq.as_slice(),
4544                                acts[bi + 1].xq.as_slice(),
4545                                acts[bi + 2].xq.as_slice(),
4546                                acts[bi + 3].xq.as_slice(),
4547                            ];
4548                            let sxs = [
4549                                acts[bi].sx,
4550                                acts[bi + 1].sx,
4551                                acts[bi + 2].sx,
4552                                acts[bi + 3].sx,
4553                            ];
4554                            let d = unsafe {
4555                                dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
4556                            };
4557                            for k in 0..4 {
4558                                let act = &acts[bi + k];
4559                                let mut dot = d[k];
4560                                for &(j, xv) in &act.outliers {
4561                                    dot +=
4562                                        (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4563                                }
4564                                // SAFETY: disjoint (bi, r) cells per worker.
4565                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
4566                            }
4567                            bi += 4;
4568                        }
4569                    }
4570                    while bi < acts.len() {
4571                        let act = &acts[bi];
4572                        let mut dot = 0f32;
4573                        for g in 0..ng {
4574                            let d = dot_i8_i8(
4575                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4576                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4577                            ) as f32
4578                                * act.sx;
4579                            dot += d * gscale(r, g);
4580                        }
4581                        for &(j, xv) in &act.outliers {
4582                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4583                        }
4584                        // SAFETY: disjoint (bi, r) cells per worker range.
4585                        unsafe { *out_addr.at(bi * rows + r) = dot };
4586                        bi += 1;
4587                    }
4588                });
4589            }
4590        };
4591        dispatch_rows(pool, rows, &run);
4592        return;
4593    }
4594
4595    let out_addr = SendMut(out.as_mut_ptr());
4596    let run = move |start: usize, end: usize| {
4597        ROW_F32.with(|rb| {
4598            let mut buf = rb.borrow_mut();
4599            buf.resize(cols, 0.0);
4600            for r in start..end {
4601                decode_f32(r, &mut buf);
4602                for bi in 0..b {
4603                    let x = &xs_all[bi * cols..(bi + 1) * cols];
4604                    let mut dot = 0f32;
4605                    for g in 0..ng {
4606                        let mut gd = 0f32;
4607                        for k in 0..GROUP_SIZE {
4608                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4609                        }
4610                        dot += gd * gscale(r, g);
4611                    }
4612                    // SAFETY: disjoint (bi, r) cells per worker range.
4613                    unsafe { *out_addr.at(bi * rows + r) = dot };
4614                }
4615            }
4616        })
4617    };
4618    dispatch_rows(pool, rows, &run);
4619}
4620
4621/// Build a GPU batch job for a q8-family mapped tensor (primary
4622/// shard): prescaled input + directory coordinates. None → not
4623/// GPU-eligible, caller stays on the CPU.
4624pub(crate) fn gpu_batch_job<'a>(
4625    t: &'a QTensor,
4626    x: &[f32],
4627) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
4628    match t {
4629        QTensor::Mapped {
4630            model,
4631            idx,
4632            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
4633            rows,
4634            cols,
4635            row_scale,
4636            col_field,
4637            ..
4638        } => Some((
4639            model.clone(),
4640            crate::gpu::BatchJob {
4641                idx: *idx,
4642                rows: *rows,
4643                cols: *cols,
4644                row_scale,
4645                xs: prescale(x, col_field, *dt).into_owned(),
4646                q1: false,
4647            },
4648        )),
4649        // q1: raw f32 activations, tile-embedded scales.
4650        QTensor::Mapped {
4651            model,
4652            idx,
4653            dtype: TensorDtype::Q1,
4654            rows,
4655            cols,
4656            ..
4657        } => Some((
4658            model.clone(),
4659            crate::gpu::BatchJob {
4660                idx: *idx,
4661                rows: *rows,
4662                cols: *cols,
4663                row_scale: &[],
4664                xs: x.to_vec(),
4665                q1: true,
4666            },
4667        )),
4668        _ => None,
4669    }
4670}
4671
4672thread_local! {
4673    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = std::cell::RefCell::new(Vec::new());
4674    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = std::cell::RefCell::new(Vec::new());
4675}
4676
4677pub(crate) fn prescale<'a>(
4678    x: &'a [f32],
4679    col_field: &[f32],
4680    dtype: TensorDtype,
4681) -> std::borrow::Cow<'a, [f32]> {
4682    if dtype == TensorDtype::Q8_2f {
4683        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
4684    } else {
4685        std::borrow::Cow::Borrowed(x)
4686    }
4687}
4688
4689/// θ col-field fold for q8_2f activations. Borrowed pass-through for
4690/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
4691pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
4692    x: &[f32],
4693    col_field: &[f32],
4694    dtype: TensorDtype,
4695    buf_id: u8,
4696    f: F,
4697) -> R {
4698    if dtype == TensorDtype::Q8_2f {
4699        if buf_id == 1 {
4700            PRESCALE_BUF1.with(|b| {
4701                let mut buf = b.borrow_mut();
4702                buf.clear();
4703                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
4704                f(&buf)
4705            })
4706        } else {
4707            PRESCALE_BUF2.with(|b| {
4708                let mut buf = b.borrow_mut();
4709                buf.clear();
4710                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
4711                f(&buf)
4712            })
4713        }
4714    } else {
4715        f(x)
4716    }
4717}
4718
4719// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
4720
4721/// AVX2+FMA available? Default ON when the CPU supports both;
4722/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
4723#[cfg(target_arch = "x86_64")]
4724pub(crate) fn avx2_enabled() -> bool {
4725    use std::sync::OnceLock;
4726    static ON: OnceLock<bool> = OnceLock::new();
4727    *ON.get_or_init(|| {
4728        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
4729            && std::arch::is_x86_feature_detected!("avx2")
4730            && std::arch::is_x86_feature_detected!("fma")
4731    })
4732}
4733
4734/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
4735/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
4736/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
4737/// active either way, they are exact (regrouped sums only).
4738#[cfg(target_arch = "x86_64")]
4739fn avx2_a8w8_enabled() -> bool {
4740    use std::sync::OnceLock;
4741    static ON: OnceLock<bool> = OnceLock::new();
4742    *ON.get_or_init(|| {
4743        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
4744    })
4745}
4746
4747/// A8W8 quantized-activation path available on THIS machine? One
4748/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
4749/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
4750#[inline]
4751pub(crate) fn a8w8_enabled() -> bool {
4752    #[cfg(target_arch = "aarch64")]
4753    {
4754        sdot_enabled()
4755    }
4756    #[cfg(target_arch = "x86_64")]
4757    {
4758        avx2_a8w8_enabled()
4759    }
4760    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4761    {
4762        false
4763    }
4764}
4765
4766/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
4767/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
4768#[inline]
4769#[allow(unreachable_code)]
4770fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
4771    #[cfg(target_arch = "aarch64")]
4772    unsafe {
4773        return dot_i8_sdot(w, xq);
4774    }
4775    #[cfg(target_arch = "x86_64")]
4776    unsafe {
4777        if avx512vnni_enabled() {
4778            return dot_i8_i8_vnni(w, xq);
4779        }
4780        return dot_i8_i8_avx2(w, xq);
4781    }
4782    w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
4783}
4784
4785/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
4786/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
4787/// `vpdpbusd` encoding.
4788#[cfg(target_arch = "x86_64")]
4789fn avx512vnni_enabled() -> bool {
4790    use std::sync::OnceLock;
4791    static ON: OnceLock<bool> = OnceLock::new();
4792    *ON.get_or_init(|| {
4793        std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
4794            && std::arch::is_x86_feature_detected!("avx512f")
4795            && std::arch::is_x86_feature_detected!("avx512bw")
4796            && std::arch::is_x86_feature_detected!("avx512vl")
4797            && std::arch::is_x86_feature_detected!("avx512vnni")
4798    })
4799}
4800
4801/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
4802/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
4803/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
4804/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
4805#[cfg(target_arch = "x86_64")]
4806#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4807unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4808    // SAFETY: callers uphold slice-length contracts (see call sites).
4809    unsafe {
4810        use core::arch::x86_64::*;
4811        let n = w.len();
4812        let mut j = 0usize;
4813        let mut total: i32;
4814        // 4 independent accumulators: vpdpbusd is its own loop-carried
4815        // dependency (~5-cycle latency) — a single-acc loop runs
4816        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
4817        // on Granite Rapids.
4818        {
4819            #[inline(always)]
4820            unsafe fn step(
4821                w: *const u8,
4822                x: *const i8,
4823                acc: core::arch::x86_64::__m512i,
4824            ) -> core::arch::x86_64::__m512i {
4825                unsafe {
4826                    use core::arch::x86_64::*;
4827                    let wv = _mm512_loadu_si512(w as *const _);
4828                    let xv = _mm512_loadu_si512(x as *const _);
4829                    let aw = _mm512_abs_epi8(wv);
4830                    let neg = _mm512_movepi8_mask(wv);
4831                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
4832                    _mm512_dpbusd_epi32(acc, aw, sx)
4833                }
4834            }
4835            let (mut a0, mut a1, mut a2, mut a3) = (
4836                _mm512_setzero_si512(),
4837                _mm512_setzero_si512(),
4838                _mm512_setzero_si512(),
4839                _mm512_setzero_si512(),
4840            );
4841            while j + 256 <= n {
4842                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4843                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
4844                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
4845                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
4846                j += 256;
4847            }
4848            while j + 64 <= n {
4849                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4850                j += 64;
4851            }
4852            let s01 = _mm512_add_epi32(a0, a1);
4853            let s23 = _mm512_add_epi32(a2, a3);
4854            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
4855        }
4856        // 32-wide (q4/vbit groups are exactly 32 bytes).
4857        if j + 32 <= n {
4858            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4859            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4860            let d = _mm256_dpbusd_epi32(
4861                _mm256_setzero_si256(),
4862                _mm256_abs_epi8(wv),
4863                _mm256_sign_epi8(xv, wv),
4864            );
4865            let hi128 = _mm256_extracti128_si256::<1>(d);
4866            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4867            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4868            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4869            total += _mm_cvtsi128_si32(s32);
4870            j += 32;
4871        }
4872        while j < n {
4873            total += (w[j] as i8) as i32 * xq[j] as i32;
4874            j += 1;
4875        }
4876        total
4877    }
4878}
4879
4880/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
4881#[cfg(target_arch = "x86_64")]
4882#[target_feature(enable = "avx2,fma")]
4883unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
4884    // SAFETY: callers uphold slice-length contracts (see call sites).
4885    unsafe {
4886        use core::arch::x86_64::*;
4887        let n = x.len();
4888        let wp = w.as_ptr();
4889        let xp = x.as_ptr();
4890        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
4891        let mut j = 0usize;
4892        while j + 16 <= n {
4893            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
4894            let lo = _mm256_cvtepi8_epi32(wb);
4895            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
4896            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
4897            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
4898            j += 16;
4899        }
4900        let acc = _mm256_add_ps(a0, a1);
4901        let hi128 = _mm256_extractf128_ps::<1>(acc);
4902        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
4903        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
4904        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
4905        let mut sum = _mm_cvtss_f32(s32);
4906        while j < n {
4907            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
4908            j += 1;
4909        }
4910        sum
4911    }
4912}
4913
4914/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
4915/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
4916/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
4917/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
4918#[cfg(target_arch = "x86_64")]
4919#[target_feature(enable = "avx2")]
4920unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
4921    // SAFETY: callers uphold slice-length contracts (see call sites).
4922    unsafe {
4923        use core::arch::x86_64::*;
4924        let n = w.len();
4925        let ones = _mm256_set1_epi16(1);
4926        let mut acc = _mm256_setzero_si256();
4927        let mut j = 0usize;
4928        while j + 32 <= n {
4929            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4930            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4931            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4932            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
4933            j += 32;
4934        }
4935        let hi128 = _mm256_extracti128_si256::<1>(acc);
4936        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
4937        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4938        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4939        let mut s = _mm_cvtsi128_si32(s32);
4940        while j < n {
4941            s += (w[j] as i8) as i32 * xq[j] as i32;
4942            j += 1;
4943        }
4944        s
4945    }
4946}
4947
4948/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
4949/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
4950/// slice as a combined 2×8 register and meets two activation pairs.
4951#[cfg(target_arch = "aarch64")]
4952#[target_feature(enable = "neon,i8mm")]
4953unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4954    // SAFETY: callers uphold slice-length contracts.
4955    unsafe {
4956        use core::arch::aarch64::*;
4957        use core::arch::asm;
4958        let n = w0.len();
4959        let w0p = w0.as_ptr() as *const i8;
4960        let w1p = w1.as_ptr() as *const i8;
4961        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
4962        // same for x2/x3.
4963        let mut acc01 = vdupq_n_s32(0);
4964        let mut acc23 = vdupq_n_s32(0);
4965        let mut i = 0usize;
4966        while i + 8 <= n {
4967            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
4968            let xb01 =
4969                vcombine_s8(vld1_s8(xs[0].as_ptr().add(i)), vld1_s8(xs[1].as_ptr().add(i)));
4970            let xb23 =
4971                vcombine_s8(vld1_s8(xs[2].as_ptr().add(i)), vld1_s8(xs[3].as_ptr().add(i)));
4972            asm!(
4973                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
4974                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
4975                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
4976                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
4977                options(pure, nomem, nostack),
4978            );
4979            i += 8;
4980        }
4981        let mut out = [[0i32; 4]; 2];
4982        let a01: [i32; 4] = core::mem::transmute(acc01);
4983        let a23: [i32; 4] = core::mem::transmute(acc23);
4984        out[0][0] = a01[0];
4985        out[0][1] = a01[1];
4986        out[1][0] = a01[2];
4987        out[1][1] = a01[3];
4988        out[0][2] = a23[0];
4989        out[0][3] = a23[1];
4990        out[1][2] = a23[2];
4991        out[1][3] = a23[3];
4992        if i < n {
4993            for (k, x) in xs.iter().enumerate() {
4994                for j in i..n {
4995                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4996                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4997                }
4998            }
4999        }
5000        out
5001    }
5002}
5003
5004/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
5005/// registers across four activation streams, eight sdot accumulators.
5006/// (The per-row form re-read each W row once per activation.)
5007#[cfg(target_arch = "aarch64")]
5008#[target_feature(enable = "neon,dotprod")]
5009unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5010    // SAFETY: callers uphold slice-length contracts.
5011    unsafe {
5012        use core::arch::aarch64::*;
5013        use core::arch::asm;
5014        let n = w0.len();
5015        let w0p = w0.as_ptr() as *const i8;
5016        let w1p = w1.as_ptr() as *const i8;
5017        let mut acc = [[vdupq_n_s32(0); 4]; 2];
5018        let mut i = 0usize;
5019        while i + 16 <= n {
5020            let wv0 = vld1q_s8(w0p.add(i));
5021            let wv1 = vld1q_s8(w1p.add(i));
5022            for (k, x) in xs.iter().enumerate() {
5023                let xv = vld1q_s8(x.as_ptr().add(i));
5024                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
5025                asm!(
5026                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
5027                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
5028                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5029                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
5030                    options(pure, nomem, nostack),
5031                );
5032                acc[0][k] = a0;
5033                acc[1][k] = a1;
5034            }
5035            i += 16;
5036        }
5037        let mut out = [[0i32; 4]; 2];
5038        for r in 0..2 {
5039            for k in 0..4 {
5040                out[r][k] = vaddvq_s32(acc[r][k]);
5041            }
5042        }
5043        if i < n {
5044            for (k, x) in xs.iter().enumerate() {
5045                for j in i..n {
5046                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
5047                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
5048                }
5049            }
5050        }
5051        out
5052    }
5053}
5054
5055/// Blocked 2 weight rows × 4 activations for the prefill GEMM
5056/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
5057/// abs() live in registers across all four activation streams; the
5058/// sign-fixup is recomputed per pair (the price of the maddubs trick).
5059/// Returns raw i8·i8 dots; the caller applies scales and outliers.
5060#[cfg(target_arch = "x86_64")]
5061#[target_feature(enable = "avx2")]
5062unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5063    // SAFETY: callers uphold slice-length contracts.
5064    unsafe {
5065        use core::arch::x86_64::*;
5066        let n = w0.len();
5067        let ones = _mm256_set1_epi16(1);
5068        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
5069        let mut j = 0usize;
5070        while j + 32 <= n {
5071            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
5072            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
5073            let aw0 = _mm256_abs_epi8(wv0);
5074            let aw1 = _mm256_abs_epi8(wv1);
5075            for (k, x) in xs.iter().enumerate() {
5076                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
5077                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
5078                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
5079                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
5080                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
5081            }
5082            j += 32;
5083        }
5084        let mut out = [[0i32; 4]; 2];
5085        for r in 0..2 {
5086            for k in 0..4 {
5087                let a = acc[r][k];
5088                let hi128 = _mm256_extracti128_si256::<1>(a);
5089                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
5090                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5091                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5092                out[r][k] = _mm_cvtsi128_si32(s32);
5093            }
5094        }
5095        if j < n {
5096            for (k, x) in xs.iter().enumerate() {
5097                for i in j..n {
5098                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
5099                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
5100                }
5101            }
5102        }
5103        out
5104    }
5105}
5106
5107/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
5108/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
5109/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
5110/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
5111#[cfg(target_arch = "x86_64")]
5112#[inline]
5113fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
5114    let dot = if avx512vnni_enabled() && row.len() >= 64 {
5115        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
5116    } else {
5117        unsafe { dot_i8_i8_avx2(row, &act.xq) }
5118    };
5119    let mut acc = dot as f32 * act.sx;
5120    for &(j, xv) in &act.outliers {
5121        acc += (row[j] as i8) as f32 * xv;
5122    }
5123    acc
5124}
5125
5126/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
5127/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
5128/// a single-acc loop runs latency-bound, measured on Granite Rapids).
5129#[cfg(target_arch = "x86_64")]
5130#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5131unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
5132    // SAFETY: callers uphold slice-length contracts (see call sites).
5133    unsafe {
5134        use core::arch::x86_64::*;
5135        let n = w.len();
5136        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
5137        #[inline(always)]
5138        unsafe fn step(
5139            w: *const u8,
5140            x: *const i8,
5141            flip: core::arch::x86_64::__m512i,
5142            acc: core::arch::x86_64::__m512i,
5143        ) -> core::arch::x86_64::__m512i {
5144            unsafe {
5145                use core::arch::x86_64::*;
5146                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
5147                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
5148            }
5149        }
5150        let (mut a0, mut a1, mut a2, mut a3) = (
5151            _mm512_setzero_si512(),
5152            _mm512_setzero_si512(),
5153            _mm512_setzero_si512(),
5154            _mm512_setzero_si512(),
5155        );
5156        let mut j = 0usize;
5157        while j + 256 <= n {
5158            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5159            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
5160            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
5161            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
5162            j += 256;
5163        }
5164        while j + 64 <= n {
5165            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5166            j += 64;
5167        }
5168        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
5169            _mm512_add_epi32(a0, a1),
5170            _mm512_add_epi32(a2, a3),
5171        ));
5172        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
5173        while j < n {
5174            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
5175            j += 1;
5176        }
5177        total
5178    }
5179}
5180
5181/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
5182/// writer's flat order, same as the NEON vzip pair), maddubs against
5183/// the pre-quantized activation group, × the group's f16 scale. Pair
5184/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
5185/// `dot_q4_row_sdot`.
5186#[cfg(target_arch = "x86_64")]
5187#[target_feature(enable = "avx2")]
5188unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5189    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
5190    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
5191    unsafe {
5192        use core::arch::x86_64::*;
5193        let lomask = _mm_set1_epi8(0x0F);
5194        let eight = _mm256_set1_epi8(8);
5195        let ones = _mm256_set1_epi16(1);
5196        let mut acc = 0f32;
5197        for gi in 0..gpr {
5198            let g = g0 + gi;
5199            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5200            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5201            let lo = _mm_and_si128(b, lomask);
5202            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5203            let w = _mm256_sub_epi8(
5204                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5205                eight,
5206            );
5207            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5208            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5209            let d = _mm256_madd_epi16(p16, ones);
5210            let hi128 = _mm256_extracti128_si256::<1>(d);
5211            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5212            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5213            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5214            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5215        }
5216        acc
5217    }
5218}
5219
5220/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
5221/// both activations dotted against the same centered i8 register.
5222#[cfg(target_arch = "x86_64")]
5223#[target_feature(enable = "avx2")]
5224unsafe fn dot_q4_row_avx2_2(
5225    packed: &[u8],
5226    scales: &[u8],
5227    g0: usize,
5228    gpr: usize,
5229    xq1: &[i8],
5230    xq2: &[i8],
5231) -> (f32, f32) {
5232    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
5233    unsafe {
5234        use core::arch::x86_64::*;
5235        let lomask = _mm_set1_epi8(0x0F);
5236        let eight = _mm256_set1_epi8(8);
5237        let ones = _mm256_set1_epi16(1);
5238        let (mut acc1, mut acc2) = (0f32, 0f32);
5239        #[inline(always)]
5240        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
5241            unsafe {
5242                use core::arch::x86_64::*;
5243                let hi128 = _mm256_extracti128_si256::<1>(d);
5244                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5245                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5246                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5247                _mm_cvtsi128_si32(s32)
5248            }
5249        }
5250        for gi in 0..gpr {
5251            let g = g0 + gi;
5252            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5253            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5254            let lo = _mm_and_si128(b, lomask);
5255            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5256            let w = _mm256_sub_epi8(
5257                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5258                eight,
5259            );
5260            let aw = _mm256_abs_epi8(w);
5261            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5262            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5263            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
5264            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
5265            acc1 += hsum(d1) as f32 * s;
5266            acc2 += hsum(d2) as f32 * s;
5267        }
5268        (acc1, acc2)
5269    }
5270}
5271
5272/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
5273#[cfg(target_arch = "x86_64")]
5274fn q8_range_avx2(
5275    q: &[u8],
5276    row_scale: &[f32],
5277    act: &SplitAct,
5278    cols: usize,
5279    out_addr: SendMut,
5280    start: usize,
5281    end: usize,
5282) {
5283    for o in start..end {
5284        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5285        // SAFETY: disjoint row ranges per worker.
5286        unsafe { *out_addr.at(o) = v };
5287    }
5288}
5289
5290/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
5291#[cfg(target_arch = "x86_64")]
5292#[allow(clippy::too_many_arguments)]
5293fn q8_range2_avx2(
5294    q: &[u8],
5295    row_scale: &[f32],
5296    a1: &SplitAct,
5297    a2: &SplitAct,
5298    cols: usize,
5299    p1: SendMut,
5300    p2: SendMut,
5301    start: usize,
5302    end: usize,
5303) {
5304    for o in start..end {
5305        let row = &q[o * cols..(o + 1) * cols];
5306        // SAFETY: disjoint row ranges per worker.
5307        unsafe {
5308            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
5309            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
5310        }
5311    }
5312}
5313
5314// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
5315
5316/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
5317/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
5318/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
5319/// accumulator dependency chain swamp the MAC advantage, and Apple's
5320/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
5321/// field trials on Cortex-A710/X-class parts with two pipes, where the
5322/// balance may differ; a pre-interleaved weight layout (repack infra)
5323/// is the known path if it ever earns its keep.
5324#[cfg(target_arch = "aarch64")]
5325fn i8mm_enabled() -> bool {
5326    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5327    *ON.get_or_init(|| {
5328        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
5329            && std::arch::is_aarch64_feature_detected!("i8mm")
5330    })
5331}
5332
5333/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
5334/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
5335/// (On non-ARM release builds only the test tolerance switch calls it.)
5336#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
5337fn sdot_enabled() -> bool {
5338    use std::sync::OnceLock;
5339    static ON: OnceLock<bool> = OnceLock::new();
5340    *ON.get_or_init(|| {
5341        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
5342        if !want { return false; }
5343        
5344        #[cfg(target_arch = "aarch64")]
5345        {
5346            if std::arch::is_aarch64_feature_detected!("dotprod") {
5347                return true;
5348            }
5349            #[cfg(target_os = "android")]
5350            {
5351                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
5352                    if cpuinfo.lines().any(|l| (l.starts_with("Features") || l.starts_with("features")) && l.contains("asimddp")) {
5353                        return true;
5354                    }
5355                }
5356            }
5357            false
5358        }
5359        #[cfg(not(target_arch = "aarch64"))]
5360        {
5361            false
5362        }
5363    })
5364}
5365
5366/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
5367/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
5368/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
5369/// matvec, shared by all rows/workers.
5370struct SplitAct {
5371    xq: Vec<i8>,
5372    sx: f32,
5373    outliers: Vec<(usize, f32)>,
5374    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
5375    /// `−128·Σx`); one i32 per split, computed once per matvec.
5376    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
5377    xsum: i32,
5378}
5379
5380thread_local! {
5381    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
5382    /// and its hidden-size allocation was steady-state heap churn.
5383    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
5384        const { std::cell::RefCell::new(Vec::new()) };
5385}
5386
5387impl Drop for SplitAct {
5388    fn drop(&mut self) {
5389        let buf = std::mem::take(&mut self.xq);
5390        if buf.capacity() > 0 {
5391            XQ_FREE.with(|f| {
5392                let mut f = f.borrow_mut();
5393                if f.len() < 16 {
5394                    f.push(buf);
5395                }
5396            });
5397        }
5398    }
5399}
5400
5401fn split_act(x: &[f32]) -> SplitAct {
5402    let n = x.len();
5403    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
5404    let thr = 8.0 * rms;
5405    // One pass: collect outliers and the bulk absmax (outliers excluded —
5406    // identical to the old zero-then-fold over a copied buffer, minus the
5407    // full-vector copy).
5408    let mut outliers: Vec<(usize, f32)> = Vec::new();
5409    let mut amax = 0f32;
5410    for (j, &v) in x.iter().enumerate() {
5411        let a = v.abs();
5412        if a > thr {
5413            outliers.push((j, v));
5414        } else if a > amax {
5415            amax = a;
5416        }
5417    }
5418    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
5419    let inv = 1.0 / sx;
5420    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
5421    xq.clear();
5422    xq.reserve(n);
5423    if outliers.is_empty() {
5424        xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
5425    } else {
5426        // Outlier slots quantize to 0 (their exact term is added later).
5427        xq.extend(x.iter().map(|&v| {
5428            if v.abs() > thr {
5429                0
5430            } else {
5431                (v * inv).round().clamp(-127.0, 127.0) as i8
5432            }
5433        }));
5434    }
5435    let xsum = xq.iter().map(|&v| v as i32).sum();
5436    SplitAct { xq, sx, outliers, xsum }
5437}
5438
5439fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
5440    let n = x.len();
5441    let rms = (x.iter().zip(col).map(|(&a, &c)| { let v = a * c; (v * v) as f64 })
5442        .sum::<f64>() / n.max(1) as f64).sqrt() as f32;
5443    let thr = 8.0 * rms;
5444    
5445    let mut outliers = Vec::new();
5446    let mut amax = 0f32;
5447    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
5448        let v = a * c;
5449        let s = v.abs();
5450        if s > thr {
5451            outliers.push((j, v));
5452        } else if s > amax {
5453            amax = s;
5454        }
5455    }
5456    
5457    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
5458    let inv = 1.0 / sx;
5459    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
5460    xq.clear();
5461    xq.reserve(n);
5462    if outliers.is_empty() {
5463        xq.extend(x.iter().zip(col).map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8));
5464    } else {
5465        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
5466            let v = a * c;
5467            if v.abs() > thr { 0 } else { (v * inv).round().clamp(-127.0, 127.0) as i8 }
5468        }));
5469    }
5470    let xsum = xq.iter().map(|&v| v as i32).sum();
5471    SplitAct { xq, sx, outliers, xsum }
5472}
5473
5474/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
5475/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
5476#[cfg(target_arch = "aarch64")]
5477#[target_feature(enable = "neon,dotprod")]
5478unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
5479    // SAFETY: callers uphold slice-length contracts (see call sites).
5480    unsafe {
5481        use core::arch::aarch64::*;
5482        use core::arch::asm;
5483        let wp = w.as_ptr() as *const i8;
5484        let n = w.len();
5485        let (mut a0, mut a1, mut a2, mut a3) =
5486            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5487        let mut i = 0;
5488        while i + 64 <= n {
5489            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5490            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
5491            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
5492            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
5493            asm!(
5494                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5495                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5496                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
5497                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
5498                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5499                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5500                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
5501                options(pure, nomem, nostack),
5502            );
5503            i += 64;
5504        }
5505        while i + 16 <= n {
5506            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5507            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
5508                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
5509            i += 16;
5510        }
5511        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
5512        while i < n {
5513            s += (*wp.add(i)) as i32 * xq[i] as i32;
5514            i += 1;
5515        }
5516        s
5517}
5518}
5519
5520/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
5521/// loaded once and reused, 4 independent accumulators hide sdot latency
5522/// (port of vmfcore `dot_i8_sdot_4rows`).
5523#[cfg(target_arch = "aarch64")]
5524#[target_feature(enable = "neon,dotprod")]
5525unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
5526    // SAFETY: callers uphold slice-length contracts (see call sites).
5527    unsafe {
5528        use core::arch::aarch64::*;
5529        use core::arch::asm;
5530        let n = xq.len();
5531        let px = xq.as_ptr();
5532        let (p0, p1, p2, p3) = (
5533            w0.as_ptr() as *const i8,
5534            w1.as_ptr() as *const i8,
5535            w2.as_ptr() as *const i8,
5536            w3.as_ptr() as *const i8,
5537        );
5538        let (mut a0, mut a1, mut a2, mut a3) =
5539            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5540        let mut i = 0;
5541        while i + 16 <= n {
5542            let x = vld1q_s8(px.add(i));
5543            let v0 = vld1q_s8(p0.add(i));
5544            let v1 = vld1q_s8(p1.add(i));
5545            let v2 = vld1q_s8(p2.add(i));
5546            let v3 = vld1q_s8(p3.add(i));
5547            asm!(
5548                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5549                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5550                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5551                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5552                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5553                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5554                options(pure, nomem, nostack),
5555            );
5556            i += 16;
5557        }
5558        let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
5559        while i < n {
5560            let xi = *px.add(i) as i32;
5561            r[0] += (*p0.add(i)) as i32 * xi;
5562            r[1] += (*p1.add(i)) as i32 * xi;
5563            r[2] += (*p2.add(i)) as i32 * xi;
5564            r[3] += (*p3.add(i)) as i32 * xi;
5565            i += 1;
5566        }
5567        r
5568}
5569}
5570
5571/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
5572/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
5573/// line plus the shared activation chunk — a single sequential weight
5574/// stream per worker. Per-row accumulation is the same one-accumulator
5575/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
5576/// are bit-identical to the mmap-layout kernel.
5577#[cfg(target_arch = "aarch64")]
5578#[target_feature(enable = "neon,dotprod")]
5579unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
5580    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
5581    // n % 16 == 0 — guaranteed by the repack gate).
5582    unsafe {
5583        use core::arch::aarch64::*;
5584        use core::arch::asm;
5585        let n = xq.len();
5586        let px = xq.as_ptr();
5587        let pg = g.as_ptr() as *const i8;
5588        let (mut a0, mut a1, mut a2, mut a3) =
5589            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5590        let mut i = 0;
5591        while i + 16 <= n {
5592            let x = vld1q_s8(px.add(i));
5593            let base = pg.add(4 * i);
5594            let v0 = vld1q_s8(base);
5595            let v1 = vld1q_s8(base.add(16));
5596            let v2 = vld1q_s8(base.add(32));
5597            let v3 = vld1q_s8(base.add(48));
5598            asm!(
5599                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5600                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5601                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5602                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5603                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5604                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5605                options(pure, nomem, nostack),
5606            );
5607            i += 16;
5608        }
5609        [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
5610    }
5611}
5612
5613/// One q8 row range via SDOT (4-row blocks + tail) — the body of
5614/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
5615/// SAME kernel for several tensors under one pool dispatch. `rep` — the
5616/// load-time interleaved repack (empty = mmap layout only); rows outside
5617/// full 4-row groups always come from the mmap layout.
5618#[cfg(target_arch = "aarch64")]
5619fn q8_range_sdot(
5620    q: &[u8],
5621    rep: &[u8],
5622    row_scale: &[f32],
5623    act: &SplitAct,
5624    cols: usize,
5625    out_addr: SendMut,
5626    start: usize,
5627    end: usize,
5628) {
5629    let mut o = start;
5630    // Leading rows to the group boundary (repack path only): the pool
5631    // splits row ranges arbitrarily, groups are absolute.
5632    if !rep.is_empty() {
5633        while o < end && o % 4 != 0 {
5634            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5635            unsafe { *out_addr.at(o) = v };
5636            o += 1;
5637        }
5638    }
5639    while o + 4 <= end {
5640        let r = if rep.is_empty() {
5641            unsafe {
5642                dot_i8_sdot_4rows(
5643                    &q[o * cols..(o + 1) * cols],
5644                    &q[(o + 1) * cols..(o + 2) * cols],
5645                    &q[(o + 2) * cols..(o + 3) * cols],
5646                    &q[(o + 3) * cols..(o + 4) * cols],
5647                    &act.xq,
5648                )
5649            }
5650        } else {
5651            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
5652        };
5653        for k in 0..4 {
5654            let mut acc = r[k] as f32 * act.sx;
5655            for &(j, xv) in &act.outliers {
5656                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
5657            }
5658            // SAFETY: disjoint row ranges per worker.
5659            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
5660        }
5661        o += 4;
5662    }
5663    while o < end {
5664        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5665        unsafe { *out_addr.at(o) = v };
5666        o += 1;
5667    }
5668}
5669
5670/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
5671/// for the fused pair multi-matrix job (`matvec2_many`).
5672#[cfg(target_arch = "aarch64")]
5673#[allow(clippy::too_many_arguments)]
5674fn q8_range2_sdot(
5675    q: &[u8],
5676    row_scale: &[f32],
5677    a1: &SplitAct,
5678    a2: &SplitAct,
5679    cols: usize,
5680    p1: SendMut,
5681    p2: SendMut,
5682    start: usize,
5683    end: usize,
5684) {
5685    for o in start..end {
5686        let row = &q[o * cols..(o + 1) * cols];
5687        // SAFETY: disjoint row ranges per worker.
5688        unsafe {
5689            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
5690            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
5691        }
5692    }
5693}
5694
5695/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
5696#[allow(clippy::too_many_arguments)]
5697fn q8_range2_f32(
5698    q: &[u8],
5699    row_scale: &[f32],
5700    x1: &[f32],
5701    x2: &[f32],
5702    cols: usize,
5703    p1: SendMut,
5704    p2: SendMut,
5705    start: usize,
5706    end: usize,
5707) {
5708    for o in start..end {
5709        let row = &q[o * cols..(o + 1) * cols];
5710        // SAFETY: disjoint row ranges per worker.
5711        unsafe {
5712            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
5713            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
5714        }
5715    }
5716}
5717
5718/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
5719fn q8_range_f32(
5720    q: &[u8],
5721    row_scale: &[f32],
5722    xs: &[f32],
5723    cols: usize,
5724    out_addr: SendMut,
5725    start: usize,
5726    end: usize,
5727) {
5728    for o in start..end {
5729        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
5730        // SAFETY: disjoint row ranges per worker.
5731        unsafe { *out_addr.at(o) = v };
5732    }
5733}
5734
5735/// SDOT row dot with exact outlier correction:
5736/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
5737#[cfg(target_arch = "aarch64")]
5738#[inline]
5739fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
5740    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
5741    for &(j, xv) in &act.outliers {
5742        acc += (row[j] as i8) as f32 * xv;
5743    }
5744    acc
5745}
5746
5747/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
5748/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
5749/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
5750/// the caller multiplies by the activation scale and adds the exact
5751/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
5752/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
5753/// → zip(lo,hi) restores flat order.
5754#[cfg(target_arch = "aarch64")]
5755#[target_feature(enable = "neon,dotprod")]
5756unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5757    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
5758    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
5759    unsafe {
5760        use core::arch::aarch64::*;
5761        use core::arch::asm;
5762        let lomask = vdupq_n_u8(0x0F);
5763        let eight = vdupq_n_s8(8);
5764        let mut acc = 0f32;
5765        for gi in 0..gpr {
5766            let g = g0 + gi;
5767            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5768            let b = vld1q_u8(packed.as_ptr().add(g * 16));
5769            let lo = vandq_u8(b, lomask);
5770            let hi = vshrq_n_u8::<4>(b);
5771            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5772            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5773            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5774            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5775            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5776            asm!(
5777                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5778                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5779                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5780                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5781                options(pure, nomem, nostack),
5782            );
5783            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5784        }
5785        acc
5786    }
5787}
5788
5789/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
5790/// part) happens ONCE per group; both pre-quantized activations are
5791/// dotted against the same centered i8 registers. Per-lane math matches
5792/// `dot_q4_row_sdot` exactly.
5793#[cfg(target_arch = "aarch64")]
5794#[target_feature(enable = "neon,dotprod")]
5795unsafe fn dot_q4_row_sdot2(
5796    packed: &[u8],
5797    scales: &[u8],
5798    g0: usize,
5799    gpr: usize,
5800    xq1: &[i8],
5801    xq2: &[i8],
5802) -> (f32, f32) {
5803    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
5804    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
5805    unsafe {
5806        use core::arch::aarch64::*;
5807        use core::arch::asm;
5808        let lomask = vdupq_n_u8(0x0F);
5809        let eight = vdupq_n_s8(8);
5810        let (mut acc1, mut acc2) = (0f32, 0f32);
5811        for gi in 0..gpr {
5812            let g = g0 + gi;
5813            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5814            let b = vld1q_u8(packed.as_ptr().add(g * 16));
5815            let lo = vandq_u8(b, lomask);
5816            let hi = vshrq_n_u8::<4>(b);
5817            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5818            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5819            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
5820            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
5821            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
5822            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
5823            let (mut a0, mut a1, mut b0, mut b1) =
5824                (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5825            asm!(
5826                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
5827                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
5828                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
5829                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
5830                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5831                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
5832                e0 = in(vreg) e0, e1 = in(vreg) e1,
5833                x10 = in(vreg) x10, x11 = in(vreg) x11,
5834                x20 = in(vreg) x20, x21 = in(vreg) x21,
5835                options(pure, nomem, nostack),
5836            );
5837            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5838            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
5839        }
5840        (acc1, acc2)
5841    }
5842}
5843
5844// ───────────────────── fused int8 kernels ─────────────────────
5845
5846/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
5847/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
5848#[inline]
5849pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
5850    #[cfg(target_arch = "aarch64")]
5851    unsafe {
5852        return axpy_i8_f32_neon(acc, row, w);
5853    }
5854    #[cfg(target_arch = "x86_64")]
5855    if avx2_enabled() {
5856        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
5857    }
5858    #[allow(unreachable_code)]
5859    {
5860        for (a, &b) in acc.iter_mut().zip(row) {
5861            *a += w * b as f32;
5862        }
5863    }
5864}
5865
5866/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
5867#[cfg(target_arch = "x86_64")]
5868#[target_feature(enable = "avx2,fma")]
5869unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
5870    // SAFETY: callers uphold slice-length contracts (see call sites).
5871    unsafe {
5872        use core::arch::x86_64::*;
5873        let n = acc.len().min(row.len());
5874        let ap = acc.as_mut_ptr();
5875        let rp = row.as_ptr();
5876        let wv = _mm256_set1_ps(w);
5877        let mut j = 0usize;
5878        while j + 16 <= n {
5879            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
5880            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
5881            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
5882            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
5883            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
5884            _mm256_storeu_ps(ap.add(j), v0);
5885            _mm256_storeu_ps(ap.add(j + 8), v1);
5886            j += 16;
5887        }
5888        while j < n {
5889            *ap.add(j) += w * (*rp.add(j)) as f32;
5890            j += 1;
5891        }
5892    }
5893}
5894
5895#[cfg(target_arch = "aarch64")]
5896#[target_feature(enable = "neon")]
5897unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
5898    // SAFETY: callers uphold slice-length contracts (see call sites).
5899    unsafe {
5900        use core::arch::aarch64::*;
5901        let n = acc.len().min(row.len());
5902        let ap = acc.as_mut_ptr();
5903        let rp = row.as_ptr();
5904        let wv = vdupq_n_f32(w);
5905        let mut j = 0usize;
5906        while j + 16 <= n {
5907            let rb = vld1q_s8(rp.add(j));
5908            let lo = vmovl_s8(vget_low_s8(rb));
5909            let hi = vmovl_s8(vget_high_s8(rb));
5910            for (off, half) in [(0, lo), (8, hi)] {
5911                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
5912                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
5913                let o = j + off;
5914                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
5915                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
5916            }
5917            j += 16;
5918        }
5919        while j < n {
5920            *ap.add(j) += w * (*rp.add(j)) as f32;
5921            j += 1;
5922        }
5923}
5924}
5925
5926/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
5927/// ≈9× scalar), scalar elsewhere.
5928#[inline]
5929pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
5930    #[cfg(target_arch = "aarch64")]
5931    unsafe {
5932        return dot_i8_f32_neon(w, x);
5933    }
5934    #[cfg(target_arch = "x86_64")]
5935    if avx2_enabled() {
5936        return unsafe { dot_i8_f32_avx2(w, x) };
5937    }
5938    #[allow(unreachable_code)]
5939    {
5940        let mut sum = 0.0f32;
5941        for (j, &b) in w.iter().enumerate() {
5942            sum += (b as i8) as f32 * x[j];
5943        }
5944        sum
5945    }
5946}
5947
5948/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
5949/// folded into the product (no prescaled copy of x). NEON on aarch64,
5950/// scalar elsewhere. Used by the active-neuron path `row_dot`.
5951#[inline]
5952fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5953    #[cfg(target_arch = "aarch64")]
5954    unsafe {
5955        return dot_i8_col_f32_neon(w, x, col);
5956    }
5957    #[allow(unreachable_code)]
5958    {
5959        let mut sum = 0.0f32;
5960        for (j, &b) in w.iter().enumerate() {
5961            sum += (b as i8) as f32 * x[j] * col[j];
5962        }
5963        sum
5964    }
5965}
5966
5967#[cfg(target_arch = "aarch64")]
5968#[target_feature(enable = "neon")]
5969unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5970    // SAFETY: callers uphold slice-length contracts (see call sites).
5971    unsafe {
5972        use core::arch::aarch64::*;
5973        let n = x.len();
5974        let wp = w.as_ptr() as *const i8;
5975        let xp = x.as_ptr();
5976        let cp = col.as_ptr();
5977        let (mut a0, mut a1, mut a2, mut a3) =
5978            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5979        let mut j = 0usize;
5980        while j + 16 <= n {
5981            let wb = vld1q_s8(wp.add(j));
5982            let lo = vmovl_s8(vget_low_s8(wb));
5983            let hi = vmovl_s8(vget_high_s8(wb));
5984            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5985            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5986            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5987            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5988            a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
5989            a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
5990            a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
5991            a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
5992            j += 16;
5993        }
5994        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5995        while j < n {
5996            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
5997            j += 1;
5998        }
5999        sum
6000}
6001}
6002
6003#[cfg(target_arch = "aarch64")]
6004#[target_feature(enable = "neon")]
6005unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
6006    // SAFETY: callers uphold slice-length contracts (see call sites).
6007    unsafe {
6008        use core::arch::aarch64::*;
6009        let n = x.len();
6010        let wp = w.as_ptr() as *const i8;
6011        let xp = x.as_ptr();
6012        let (mut a0, mut a1, mut a2, mut a3) =
6013            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
6014        let mut j = 0usize;
6015        while j + 16 <= n {
6016            let wb = vld1q_s8(wp.add(j));
6017            let lo = vmovl_s8(vget_low_s8(wb));
6018            let hi = vmovl_s8(vget_high_s8(wb));
6019            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
6020            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
6021            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
6022            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
6023            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
6024            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
6025            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
6026            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
6027            j += 16;
6028        }
6029        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
6030        while j < n {
6031            sum += (*wp.add(j)) as f32 * *xp.add(j);
6032            j += 1;
6033        }
6034        sum
6035}
6036}
6037
6038#[allow(clippy::too_many_arguments)]
6039fn qmatvec(
6040    q: &[u8],
6041    rep: &[u8],
6042    row_scale: &[f32],
6043    x: &[f32],
6044    col_field: &[f32],
6045    dtype: TensorDtype,
6046    rows: usize,
6047    cols: usize,
6048    out: &mut [f32],
6049    pool: Option<&Pool>,
6050) {
6051    debug_assert_eq!(out.len(), rows);
6052    #[cfg(not(target_arch = "aarch64"))]
6053    let _ = rep;
6054
6055    #[cfg(target_arch = "aarch64")]
6056    if sdot_enabled() {
6057        let act = if dtype == TensorDtype::Q8_2f {
6058            split_act_q8_2f(x, col_field)
6059        } else {
6060            split_act(x)
6061        };
6062        let out_addr = SendMut(out.as_mut_ptr());
6063        let run_range = |start: usize, end: usize| {
6064            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
6065        };
6066        match pool {
6067            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6068            _ => run_range(0, rows),
6069        }
6070        return;
6071    }
6072    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
6073    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
6074    #[cfg(target_arch = "x86_64")]
6075    if avx2_a8w8_enabled() {
6076        let act = if dtype == TensorDtype::Q8_2f {
6077            split_act_q8_2f(x, col_field)
6078        } else {
6079            split_act(x)
6080        };
6081        let out_addr = SendMut(out.as_mut_ptr());
6082        let run_range =
6083            |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
6084        match pool {
6085            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6086            _ => run_range(0, rows),
6087        }
6088        return;
6089    }
6090
6091    prescale_with(x, col_field, dtype, 1, |xs| {
6092        let out_addr = SendMut(out.as_mut_ptr());
6093        let run_range = move |start: usize, end: usize| {
6094            for o in start..end {
6095                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
6096                // SAFETY: disjoint row ranges per worker.
6097                unsafe { *out_addr.at(o) = v };
6098            }
6099        };
6100        match pool {
6101            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6102            _ => run_range(0, rows),
6103        }
6104    });
6105}
6106
6107#[allow(clippy::too_many_arguments)]
6108fn qmatvec2(
6109    q: &[u8],
6110    row_scale: &[f32],
6111    x1: &[f32],
6112    x2: &[f32],
6113    col_field: &[f32],
6114    dtype: TensorDtype,
6115    rows: usize,
6116    cols: usize,
6117    o1: &mut [f32],
6118    o2: &mut [f32],
6119    pool: Option<&Pool>,
6120) {
6121    #[cfg(target_arch = "aarch64")]
6122    if sdot_enabled() {
6123        let a1s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x1, col_field) } else { split_act(x1) };
6124        let a2s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x2, col_field) } else { split_act(x2) };
6125        let p1 = SendMut(o1.as_mut_ptr());
6126        let p2 = SendMut(o2.as_mut_ptr());
6127        let run_range = |start: usize, end: usize| {
6128            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6129        };
6130        match pool {
6131            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6132            _ => run_range(0, rows),
6133        }
6134        return;
6135    }
6136    #[cfg(target_arch = "x86_64")]
6137    if avx2_a8w8_enabled() {
6138        let a1s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x1, col_field) } else { split_act(x1) };
6139        let a2s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x2, col_field) } else { split_act(x2) };
6140        let p1 = SendMut(o1.as_mut_ptr());
6141        let p2 = SendMut(o2.as_mut_ptr());
6142        let run_range = |start: usize, end: usize| {
6143            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6144        };
6145        match pool {
6146            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6147            _ => run_range(0, rows),
6148        }
6149        return;
6150    }
6151
6152    prescale_with(x1, col_field, dtype, 1, |x1s| {
6153        prescale_with(x2, col_field, dtype, 2, |x2s| {
6154            let p1 = SendMut(o1.as_mut_ptr());
6155            let p2 = SendMut(o2.as_mut_ptr());
6156            let run_range = move |start: usize, end: usize| {
6157                for o in start..end {
6158                    let row = &q[o * cols..(o + 1) * cols];
6159                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
6160                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
6161                    // SAFETY: disjoint row ranges per worker.
6162                    unsafe {
6163                        *p1.at(o) = s1;
6164                        *p2.at(o) = s2;
6165                    }
6166                }
6167            };
6168            match pool {
6169                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6170                _ => run_range(0, rows),
6171            }
6172        });
6173    });
6174}
6175
6176#[derive(Clone, Copy)]
6177struct SendMut(*mut f32);
6178unsafe impl Send for SendMut {}
6179unsafe impl Sync for SendMut {}
6180
6181impl SendMut {
6182    #[inline]
6183    fn at(self, i: usize) -> *mut f32 {
6184        unsafe { self.0.add(i) }
6185    }
6186}
6187
6188#[cfg(test)]
6189mod tests {
6190    use super::*;
6191
6192    #[test]
6193    fn f32_matvec_matches_matvec_rows_bitexact() {
6194        let (rows, cols) = (300, 40);
6195        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
6196        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
6197        let qt = QTensor::from_f32(w.clone(), rows, cols);
6198
6199        let mut a = vec![0.0f32; rows];
6200        matvec_rows(None, &w, &x, &mut a);
6201        let mut b = vec![0.0f32; rows];
6202        qt.matvec(&x, &mut b, None);
6203        assert_eq!(a, b);
6204    }
6205
6206    #[test]
6207    fn sdot_kernel_exact_on_grid() {
6208        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
6209        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
6210        // exact f32 dot to float rounding. This isolates kernel
6211        // correctness from quantization noise.
6212        eprintln!("sdot_enabled = {}", sdot_enabled());
6213        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
6214        let w: Vec<u8> = (0..rows * cols)
6215            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6216            .collect();
6217        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
6218        let x: Vec<f32> = (0..cols)
6219            .map(|i| match i % 3 {
6220                0 => 1.0,
6221                1 => -1.0,
6222                _ => 0.0,
6223            })
6224            .collect();
6225        let mut a = vec![0.0f32; rows];
6226        qmatvec(&w, &[], &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut a, None);
6227        for o in 0..rows {
6228            let mut acc = 0.0f32;
6229            for j in 0..cols {
6230                acc += (w[o * cols + j] as i8) as f32 * x[j];
6231            }
6232            let expect = acc * scales[o];
6233            assert!(
6234                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6235                "row {o}: {} vs {expect}",
6236                a[o]
6237            );
6238        }
6239    }
6240
6241    #[test]
6242    fn q1_tbl_fast_path_matches_reference() {
6243        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
6244        // row's final 4-tile window trips the 4B-overread guard (the
6245        // payload ends exactly at the last tile) — both paths must
6246        // agree with the dequant reference.
6247        let (rows, cols) = (5, 256);
6248        let gpr = cols / GROUP_SIZE;
6249        let mut bytes = Vec::new();
6250        for t in 0..rows * gpr {
6251            let s = 0.007 + (t % 11) as f32 * 0.004;
6252            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6253            for j in 0..4 {
6254                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
6255            }
6256        }
6257        let x: Vec<f32> = (0..cols)
6258            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
6259            .collect();
6260        let mut w = vec![0.0f32; rows * cols];
6261        cortiq_core::quant::dequant_q1(&bytes, &mut w);
6262        let mut got = vec![0.0f32; rows];
6263        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6264        for o in 0..rows {
6265            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6266            assert!(
6267                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6268                "row {o}: {} vs {expect}",
6269                got[o]
6270            );
6271        }
6272        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
6273        // single-matvec path bit-for-bit.
6274        let b = 5usize;
6275        let mut xs_all = Vec::new();
6276        for bi in 0..b {
6277            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
6278        }
6279        let mut mm = vec![0.0f32; b * rows];
6280        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
6281        for bi in 0..b {
6282            let mut single = vec![0.0f32; rows];
6283            q1_matvec(&bytes, &xs_all[bi * cols..(bi + 1) * cols], rows, cols, &mut single, None);
6284            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
6285        }
6286    }
6287
6288    #[test]
6289    fn q1_kernels_match_exact_reference() {
6290        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
6291        let (rows, cols) = (7, 96);
6292        let gpr = cols / GROUP_SIZE;
6293        let mut bytes = Vec::new();
6294        for t in 0..rows * gpr {
6295            let s = 0.01 + (t % 13) as f32 * 0.003;
6296            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6297            for j in 0..4 {
6298                bytes.push(((t * 31 + j * 97) % 251) as u8);
6299            }
6300        }
6301        // On-grid activations (±1, amax 1) → the SDOT path is exact.
6302        let x: Vec<f32> = (0..cols)
6303            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
6304            .collect();
6305        // Reference through the core dequant.
6306        let mut w = vec![0.0f32; rows * cols];
6307        cortiq_core::quant::dequant_q1(&bytes, &mut w);
6308        let mut expect = vec![0.0f32; rows];
6309        for o in 0..rows {
6310            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6311        }
6312        let mut got = vec![0.0f32; rows];
6313        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6314        for o in 0..rows {
6315            assert!(
6316                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
6317                "row {o}: {} vs {}",
6318                got[o],
6319                expect[o]
6320            );
6321        }
6322        // Pair and batch paths agree with the single path.
6323        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
6324        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
6325        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
6326        assert_eq!(a1, got);
6327        let mut xs = x.clone();
6328        xs.extend_from_slice(&x2);
6329        let mut mm = vec![0.0f32; 2 * rows];
6330        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
6331        assert_eq!(&mm[..rows], got.as_slice());
6332        assert_eq!(&mm[rows..], a2.as_slice());
6333    }
6334
6335    #[test]
6336    fn repack_is_bit_identical() {
6337        // The interleaved-repack kernel must produce EXACTLY the same
6338        // bits as the mmap-layout kernel: integer accumulation is order-
6339        // exact, the f32 epilogue is identical. Odd rows exercise the
6340        // tail; direct range calls exercise unaligned pool splits.
6341        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
6342        let w: Vec<u8> = (0..rows * cols)
6343            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
6344            .collect();
6345        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
6346        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
6347        let rep = q8_repack_layout(&w, rows, cols);
6348        // Group interleave round-trips.
6349        for g in 0..rows / 4 {
6350            for c in 0..cols / 16 {
6351                for lane in 0..4 {
6352                    assert_eq!(
6353                        &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
6354                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
6355                    );
6356                }
6357            }
6358        }
6359        let mut a = vec![0.0f32; rows];
6360        qmatvec(&w, &[], &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut a, None);
6361        let mut b = vec![0.0f32; rows];
6362        qmatvec(&w, &rep, &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut b, None);
6363        assert_eq!(a, b, "full-range repack output diverged");
6364
6365        #[cfg(target_arch = "aarch64")]
6366        if sdot_enabled() {
6367            // Unaligned range split (pool workers get arbitrary bounds).
6368            let act = split_act(&x);
6369            let mut c1 = vec![0.0f32; rows];
6370            let mut c2 = vec![0.0f32; rows];
6371            q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
6372            q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
6373            assert_eq!(c1, c2, "unaligned-range repack output diverged");
6374        }
6375    }
6376
6377    #[test]
6378    fn sdot_a8w8_noise_is_bounded() {
6379        // Off-grid activations: A8 quantization noise must stay small in
6380        // relative L2 over the whole output (realistic accuracy contract;
6381        // vmfcore measured argmax-identical decode on real models).
6382        let (rows, cols) = (16, 512);
6383        let w: Vec<u8> = (0..rows * cols)
6384            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6385            .collect();
6386        let scales = vec![0.01f32; rows];
6387        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
6388        let mut a = vec![0.0f32; rows];
6389        qmatvec(&w, &[], &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut a, None);
6390        let (mut num, mut den) = (0f64, 0f64);
6391        for o in 0..rows {
6392            let mut acc = 0.0f32;
6393            for j in 0..cols {
6394                acc += (w[o * cols + j] as i8) as f32 * x[j];
6395            }
6396            let expect = acc * scales[o];
6397            num += ((a[o] - expect) as f64).powi(2);
6398            den += (expect as f64).powi(2);
6399        }
6400        let rel = (num / den.max(1e-12)).sqrt();
6401        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
6402    }
6403
6404    #[test]
6405    fn i8_dot_neon_matches_scalar() {
6406        let n = 100;
6407        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
6408        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
6409        let mut scalar = 0.0f32;
6410        for j in 0..n {
6411            scalar += (w[j] as i8) as f32 * x[j];
6412        }
6413        let fast = dot_i8_f32(&w, &x);
6414        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
6415    }
6416
6417    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
6418    #[test]
6419    fn vbitmatvec_matches_full_dequant() {
6420        let (rows, cols) = (6, 64);
6421        let ng = cols / GROUP_SIZE;
6422        // Hand-craft: bits per row, f16 scales, packed rows.
6423        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6424        let mut bytes = bits.clone();
6425        for g in 0..rows * ng {
6426            let s = 0.02 + 0.001 * g as f32;
6427            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6428        }
6429        for r in 0..rows {
6430            let b = bits[r] as usize;
6431            let (mut acc, mut nb) = (0u64, 0usize);
6432            let mut rowbytes = Vec::new();
6433            for i in 0..cols {
6434                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6435                acc = (acc << b) | v;
6436                nb += b;
6437                while nb >= 8 {
6438                    nb -= 8;
6439                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
6440                }
6441            }
6442            if nb > 0 {
6443                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6444            }
6445            bytes.extend_from_slice(&rowbytes);
6446        }
6447        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6448
6449        let mut reference = vec![0f32; rows * cols];
6450        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
6451        let mut expect = vec![0f32; rows];
6452        for r in 0..rows {
6453            expect[r] = reference[r * cols..(r + 1) * cols]
6454                .iter()
6455                .zip(&x)
6456                .map(|(w, xv)| w * xv)
6457                .sum();
6458        }
6459        let mut got = vec![0f32; rows];
6460        let offsets = vbit_row_offsets(&bytes, rows, cols);
6461        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
6462        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
6463        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
6464        // the golden-parity gate).
6465        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6466        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
6467        for r in 0..rows {
6468            assert!(
6469                (got[r] - expect[r]).abs() < tol * scale,
6470                "row {r}: {} vs {}",
6471                got[r],
6472                expect[r]
6473            );
6474        }
6475    }
6476
6477    /// Fused q4 matvec must match the reference full-dequant + dense
6478    /// matvec bit-for-bit in structure (same f32 math, group order).
6479    /// vbit matmat: the blocked 1×4 leg must match the per-row path
6480    /// (paired env toggle; larger shape so both code paths engage).
6481    #[test]
6482    #[cfg(target_arch = "x86_64")]
6483    fn vbit_matmat_blocked_matches_per_row() {
6484        let (rows, cols, b) = (64usize, 128usize, 9usize);
6485        let ng = cols / GROUP_SIZE;
6486        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
6487        let mut bytes = bits.clone();
6488        for g in 0..rows * ng {
6489            let sc = 0.02 + 0.0005 * g as f32;
6490            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
6491        }
6492        for r in 0..rows {
6493            let bw = bits[r] as usize;
6494            let (mut acc, mut nb) = (0u64, 0usize);
6495            let mut rowbytes = Vec::new();
6496            for i in 0..cols {
6497                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6498                acc = (acc << bw) | v;
6499                nb += bw;
6500                while nb >= 8 {
6501                    nb -= 8;
6502                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
6503                }
6504            }
6505            if nb > 0 {
6506                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6507            }
6508            bytes.extend_from_slice(&rowbytes);
6509        }
6510        let x: Vec<f32> =
6511            (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
6512        let offsets = vbit_row_offsets(&bytes, rows, cols);
6513        let mut y_a = vec![0f32; b * rows];
6514        let mut y_b = vec![0f32; b * rows];
6515        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
6516        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
6517        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
6518        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
6519        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
6520        let max_d =
6521            y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
6522        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
6523    }
6524
6525    #[test]
6526    fn q4matvec_matches_full_dequant() {
6527        let (rows, cols) = (8, 64);
6528        let groups = rows * cols / GROUP_SIZE;
6529        // Hand-craft a q4_block blob: nibbles then f16 scales.
6530        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6531        for i in 0..groups * 16 {
6532            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6533        }
6534        for g in 0..groups {
6535            let s = 0.01 + 0.003 * g as f32;
6536            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6537        }
6538        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6539
6540        let mut reference = vec![0.0f32; rows * cols];
6541        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6542        let mut expect = vec![0.0f32; rows];
6543        for r in 0..rows {
6544            expect[r] = reference[r * cols..(r + 1) * cols]
6545                .iter()
6546                .zip(&x)
6547                .map(|(w, xv)| w * xv)
6548                .sum();
6549        }
6550
6551        let mut got = vec![0.0f32; rows];
6552        q4matvec(&bytes, &x, rows, cols, &mut got, None);
6553        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
6554        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
6555        // in the golden-parity gate).
6556        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6557        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6558        for r in 0..rows {
6559            assert!(
6560                (got[r] - expect[r]).abs() < tol * scale,
6561                "row {r}: {} vs {}",
6562                got[r],
6563                expect[r]
6564            );
6565        }
6566    }
6567
6568    /// Fused two-input vbit matvec must equal two single matvecs exactly
6569    /// (same per-lane accumulation order on both scalar and SDOT paths).
6570    #[test]
6571    fn vbitmatvec2_equals_two_singles() {
6572        let (rows, cols) = (6, 64);
6573        let ng = cols / GROUP_SIZE;
6574        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6575        let mut bytes = bits.clone();
6576        for g in 0..rows * ng {
6577            let s = 0.02 + 0.001 * g as f32;
6578            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6579        }
6580        for r in 0..rows {
6581            let b = bits[r] as usize;
6582            let (mut acc, mut nb) = (0u64, 0usize);
6583            let mut rowbytes = Vec::new();
6584            for i in 0..cols {
6585                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6586                acc = (acc << b) | v;
6587                nb += b;
6588                while nb >= 8 {
6589                    nb -= 8;
6590                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
6591                }
6592            }
6593            if nb > 0 {
6594                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6595            }
6596            bytes.extend_from_slice(&rowbytes);
6597        }
6598        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6599        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
6600        let offsets = vbit_row_offsets(&bytes, rows, cols);
6601
6602        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6603        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
6604        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
6605        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6606        vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6607        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
6608        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
6609    }
6610
6611    /// Fused two-input q4 matvec must equal two single matvecs exactly.
6612    #[test]
6613    fn q4matvec2_equals_two_singles() {
6614        let (rows, cols) = (8, 128);
6615        let groups = rows * cols / GROUP_SIZE;
6616        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6617        for i in 0..groups * 16 {
6618            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6619        }
6620        for g in 0..groups {
6621            let s = 0.01 + 0.003 * g as f32;
6622            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6623        }
6624        // Include an outlier channel so the SDOT correction path is
6625        // exercised in the pair kernel too.
6626        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6627        x1[9] = 250.0;
6628        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6629
6630        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6631        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
6632        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
6633        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6634        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6635        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
6636        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
6637    }
6638
6639    /// Multi-matrix job must equal separate matvecs exactly — same
6640    /// kernels, only the dispatch is fused.
6641    #[test]
6642    fn matvec_many_equals_separate_matvecs() {
6643        use crate::pool::Pool;
6644        let (r1, r2, cols) = (300, 200, 64);
6645        let mk = |salt: usize, rows: usize| {
6646            QTensor::from_f32(
6647                (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
6648                rows,
6649                cols,
6650            )
6651        };
6652        let (a, b) = (mk(1, r1), mk(5, r2));
6653        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
6654        let pool = Pool::new(3);
6655
6656        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
6657        a.matvec(&x, &mut ea, Some(&pool));
6658        b.matvec(&x, &mut eb, Some(&pool));
6659        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
6660        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
6661        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
6662        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
6663    }
6664
6665    /// Batched q4/vbit matmat must equal per-position matvec calls
6666    /// exactly (the fallback it replaced) — same kernels, same order.
6667    #[test]
6668    fn batched_matmat_equals_per_position_matvec() {
6669        let (rows, cols, b) = (8, 64, 5);
6670        // q4 blob.
6671        let groups = rows * cols / GROUP_SIZE;
6672        let mut q4 = Vec::new();
6673        for i in 0..groups * 16 {
6674            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6675        }
6676        for g in 0..groups {
6677            q4.extend_from_slice(
6678                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6679            );
6680        }
6681        // vbit blob (mixed widths incl. 8).
6682        let ng = cols / GROUP_SIZE;
6683        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
6684        let mut vb = bits.clone();
6685        for g in 0..rows * ng {
6686            vb.extend_from_slice(
6687                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
6688            );
6689        }
6690        for r in 0..rows {
6691            let bw = bits[r] as usize;
6692            let (mut acc, mut nb) = (0u64, 0usize);
6693            let mut rowbytes = Vec::new();
6694            for i in 0..cols {
6695                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6696                acc = (acc << bw) | v;
6697                nb += bw;
6698                while nb >= 8 {
6699                    nb -= 8;
6700                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
6701                }
6702            }
6703            if nb > 0 {
6704                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6705            }
6706            vb.extend_from_slice(&rowbytes);
6707        }
6708        let offsets = vbit_row_offsets(&vb, rows, cols);
6709
6710        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6711
6712        // q4: batch vs singles.
6713        let mut got = vec![0f32; b * rows];
6714        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
6715        for bi in 0..b {
6716            let mut expect = vec![0f32; rows];
6717            q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
6718            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
6719        }
6720
6721        // vbit: batch vs singles.
6722        let mut got = vec![0f32; b * rows];
6723        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
6724        for bi in 0..b {
6725            let mut expect = vec![0f32; rows];
6726            vbitmatvec(
6727                &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
6728            );
6729            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
6730        }
6731    }
6732
6733    /// q4_tiled kernels must produce BIT-identical outputs to the q4
6734    /// split kernels on the same values (same ints, same order — only
6735    /// the byte placement differs).
6736    #[test]
6737    fn q4_tiled_matches_q4_block_bitexact() {
6738        let (rows, cols, b) = (8usize, 128usize, 3usize);
6739        let groups = rows * cols / GROUP_SIZE;
6740        let mut split = Vec::with_capacity(groups * 18);
6741        for i in 0..groups * 16 {
6742            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6743        }
6744        for g in 0..groups {
6745            split.extend_from_slice(
6746                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6747            );
6748        }
6749        // Re-tile: [scale][nibbles] per group.
6750        let (packed, scales) = split.split_at(groups * 16);
6751        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
6752        for g in 0..groups {
6753            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
6754            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
6755        }
6756
6757        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6758        x1[9] = 250.0; // exercise the outlier path
6759        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6760
6761        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
6762        q4matvec(&split, &x1, rows, cols, &mut a, None);
6763        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
6764        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
6765
6766        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6767        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
6768        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
6769        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
6770        assert_eq!(a1, t1);
6771        assert_eq!(a2, t2);
6772
6773        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6774        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
6775        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
6776        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
6777        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
6778    }
6779
6780    /// q4 SDOT outlier correction: a single huge activation channel
6781    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
6782    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
6783    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
6784    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
6785    /// can never qualify (8² = n).
6786    #[test]
6787    fn q4matvec_sdot_outlier_exact() {
6788        let (rows, cols) = (4, 128);
6789        let groups = rows * cols / GROUP_SIZE;
6790        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6791        for i in 0..groups * 16 {
6792            bytes.push(((i * 11 + 5) % 256) as u8);
6793        }
6794        for g in 0..groups {
6795            let s = 0.02 + 0.002 * g as f32;
6796            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6797        }
6798        let mut x: Vec<f32> = (0..cols)
6799            .map(|i| match i % 3 {
6800                0 => 1.0,
6801                1 => -1.0,
6802                _ => 0.0,
6803            })
6804            .collect();
6805        x[17] = 300.0; // ≫ 8·rms → outlier channel
6806
6807        let mut reference = vec![0.0f32; rows * cols];
6808        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6809        let mut expect = vec![0.0f32; rows];
6810        for r in 0..rows {
6811            expect[r] = reference[r * cols..(r + 1) * cols]
6812                .iter()
6813                .zip(&x)
6814                .map(|(w, xv)| w * xv)
6815                .sum();
6816        }
6817        let mut got = vec![0.0f32; rows];
6818        q4matvec(&bytes, &x, rows, cols, &mut got, None);
6819        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6820        for r in 0..rows {
6821            assert!(
6822                (got[r] - expect[r]).abs() < 2e-3 * scale,
6823                "row {r}: {} vs {} (outlier term must be exact)",
6824                got[r],
6825                expect[r]
6826            );
6827        }
6828    }
6829
6830    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
6831    /// including the ternary zero level and the binary-searched outlier
6832    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
6833    #[test]
6834    fn q1t_matvec_matches_reference() {
6835        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
6836        let (rows, cols) = (3usize, 64usize); // gpr = 2
6837        let gpr = cols / GROUP_SIZE;
6838        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
6839        // Overlay (must be sorted by flat index): a few spikes across rows.
6840        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
6841        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
6842        let mut bytes = Vec::new();
6843        for r in 0..rows {
6844            for g in 0..gpr {
6845                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
6846                let mut c = [0u8; 7];
6847                for k in 0..GROUP_SIZE {
6848                    // Encoder invariant: code 0 at outlier positions.
6849                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
6850                        0
6851                    } else {
6852                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
6853                    };
6854                    cortiq_core::quant::q1t_pack(&mut c, k, code);
6855                }
6856                bytes.extend_from_slice(&c);
6857            }
6858        }
6859        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
6860        // row (outliers are sorted by flat index → already grouped by row).
6861        let mut row_ptr = vec![0u32; rows + 1];
6862        for &(idx, _) in &outliers {
6863            row_ptr[idx as usize / cols + 1] += 1;
6864        }
6865        for r in 0..rows {
6866            row_ptr[r + 1] += row_ptr[r];
6867        }
6868        for &p in &row_ptr {
6869            bytes.extend_from_slice(&p.to_le_bytes());
6870        }
6871        for &(idx, v) in &outliers {
6872            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
6873            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
6874        }
6875
6876        let mut refw = vec![0f32; rows * cols];
6877        dequant_q1t(&bytes, rows, cols, &mut refw);
6878        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
6879        // x exactly and matches the f32 reference (same trick as the q1 test).
6880        let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6881        let mut expect = vec![0f32; rows];
6882        for r in 0..rows {
6883            let mut a = 0.0f32;
6884            for j in 0..cols {
6885                a += refw[r * cols + j] * x[j];
6886            }
6887            expect[r] = a;
6888        }
6889        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
6890        let mut got = vec![0f32; rows];
6891        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
6892        for r in 0..rows {
6893            assert!((got[r] - expect[r]).abs() < tol(expect[r]), "row {r}: {} vs {}", got[r], expect[r]);
6894        }
6895        // matmat (b=2, f32 decode path) must agree too.
6896        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
6897        let mut gm = vec![0f32; 2 * rows];
6898        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
6899        for r in 0..rows {
6900            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
6901            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
6902        }
6903        // The fused-pair dispatch (matvec2) for Q1T routes to two q1t_matvec
6904        // passes — same kernel, so both outputs equal the single-vec result.
6905        // (Dispatch is covered end-to-end by the mixed-dtype bench; here we
6906        // confirm the two-pass composition is output-identical.)
6907        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
6908        q1t_matvec(&bytes, &x, rows, cols, &mut p1, None);
6909        q1t_matvec(&bytes, &x, rows, cols, &mut p2, None);
6910        for r in 0..rows {
6911            assert!((p1[r] - expect[r]).abs() < tol(expect[r]) && (p2[r] - expect[r]).abs() < tol(expect[r]));
6912        }
6913    }
6914
6915    // Speed A/B: the base-3-division decode (what the packing commit left in
6916    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
6917    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
6918    #[test]
6919    #[ignore]
6920    fn q1t_matvec_speed() {
6921        use cortiq_core::quant::{f32_to_f16, q1t_code, q1t_pack, Q1T_TILE};
6922        use std::time::Instant;
6923        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
6924        let gpr = cols / GROUP_SIZE;
6925        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
6926        for r in 0..rows {
6927            for g in 0..gpr {
6928                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
6929                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
6930                let mut c = [0u8; 7];
6931                for k in 0..GROUP_SIZE {
6932                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
6933                }
6934                bytes.extend_from_slice(&c);
6935            }
6936        }
6937        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
6938        let mut row_ptr = vec![0u32; rows + 1];
6939        let mut idx = 0usize;
6940        while idx < n {
6941            row_ptr[idx / cols + 1] += 1;
6942            idx += stride;
6943        }
6944        for r in 0..rows {
6945            row_ptr[r + 1] += row_ptr[r];
6946        }
6947        for &p in &row_ptr {
6948            bytes.extend_from_slice(&p.to_le_bytes());
6949        }
6950        let mut idx = 0usize;
6951        while idx < n {
6952            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
6953            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
6954            idx += stride;
6955        }
6956        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
6957        // reference (the A/B is a timing check; values must still agree).
6958        let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6959        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
6960
6961        // "before": base-3 division decode into a buffer, then dot.
6962        let slow = |out: &mut [f32]| {
6963            let mut buf = vec![0f32; cols];
6964            for r in 0..rows {
6965                for g in 0..gpr {
6966                    let off = (r * gpr + g) * Q1T_TILE;
6967                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6968                    let codes = &bytes[off + 2..off + Q1T_TILE];
6969                    for k in 0..GROUP_SIZE {
6970                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
6971                            1 => s,
6972                            2 => -s,
6973                            _ => 0.0,
6974                        };
6975                    }
6976                }
6977                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
6978                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
6979            }
6980        };
6981        let iters = 5;
6982        let mut a = vec![0f32; rows];
6983        slow(&mut a); // warm
6984        let t = Instant::now();
6985        for _ in 0..iters {
6986            slow(&mut a);
6987        }
6988        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6989
6990        let mut b = vec![0f32; rows];
6991        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
6992        let t = Instant::now();
6993        for _ in 0..iters {
6994            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
6995        }
6996        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6997
6998        for r in 0..rows {
6999            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
7000        }
7001        println!(
7002            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
7003            slow_ms / fast_ms
7004        );
7005    }
7006}