Skip to main content

cortiq_engine/
qtensor.rs

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