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