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