Skip to main content

cortiq_engine/
qtensor.rs

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