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/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
1268/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
1269#[cfg(target_os = "macos")]
1270#[allow(clippy::too_many_arguments)]
1271pub(crate) fn sgemm_rm(
1272    m: usize,
1273    n: usize,
1274    k: usize,
1275    alpha: f32,
1276    a: &[f32],
1277    lda: usize,
1278    b_mat: &[f32],
1279    ldb: usize,
1280    b_rows_are_n: bool,
1281    c: &mut [f32],
1282    ldc: usize,
1283) {
1284    debug_assert!(a.len() >= (m - 1) * lda + k);
1285    debug_assert!(c.len() >= (m - 1) * ldc + n);
1286    unsafe {
1287        accel_blas::cblas_sgemm(
1288            101, // RowMajor
1289            111, // NoTrans A
1290            if b_rows_are_n { 112 } else { 111 },
1291            m as i32,
1292            n as i32,
1293            k as i32,
1294            alpha,
1295            a.as_ptr(),
1296            lda as i32,
1297            b_mat.as_ptr(),
1298            ldb as i32,
1299            0.0,
1300            c.as_mut_ptr(),
1301            ldc as i32,
1302        );
1303    }
1304}
1305
1306/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
1307/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
1308/// on the AMX with one row-major sgemm. Tiles live in cache, weights
1309/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
1310/// logits shift within f32 rounding — tolerance-class, like every
1311/// reduction-order change; decode (M=1) never takes this path.
1312#[cfg(target_os = "macos")]
1313fn qmatmat_accel(
1314    q: &[u8],
1315    row_scale: &[f32],
1316    pre: &[std::borrow::Cow<'_, [f32]>],
1317    rows: usize,
1318    cols: usize,
1319    out: &mut [f32],
1320    pool: Option<&Pool>,
1321) {
1322    // NOTE: double-buffering the dequant against the sgemm (a scoped
1323    // thread driving the pool on tile k+1 while the caller multiplies
1324    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
1325    // multithreaded, and the dequant workers just steal its cores.
1326    const TR: usize = 2048;
1327    let b = pre.len();
1328    thread_local! {
1329        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1330        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1331    }
1332    XPANEL.with(|xp| {
1333        WTILE.with(|wt| {
1334            let mut xpanel = xp.borrow_mut();
1335            xpanel.clear();
1336            for x in pre {
1337                xpanel.extend_from_slice(x);
1338            }
1339            let mut wtile = wt.borrow_mut();
1340            wtile.resize(TR * cols, 0.0);
1341            let mut r0 = 0usize;
1342            while r0 < rows {
1343                let tr = TR.min(rows - r0);
1344                // Dequant the tile (scale folded) — pool-parallel.
1345                let wt_addr = SendMut(wtile.as_mut_ptr());
1346                let run = |start: usize, end: usize| {
1347                    for r in start..end {
1348                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1349                        let s = row_scale[r0 + r];
1350                        // SAFETY: workers cover disjoint r ranges.
1351                        let dst = unsafe {
1352                            std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1353                        };
1354                        for (d, &v) in dst.iter_mut().zip(row) {
1355                            *d = (v as i8) as f32 * s;
1356                        }
1357                    }
1358                };
1359                dispatch_rows(pool, tr, &run);
1360                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
1361                unsafe {
1362                    accel_blas::cblas_sgemm(
1363                        101, // RowMajor
1364                        111, // NoTrans A
1365                        112, // Trans B
1366                        b as i32,
1367                        tr as i32,
1368                        cols as i32,
1369                        1.0,
1370                        xpanel.as_ptr(),
1371                        cols as i32,
1372                        wtile.as_ptr(),
1373                        cols as i32,
1374                        0.0,
1375                        out.as_mut_ptr().add(r0),
1376                        rows as i32,
1377                    );
1378                }
1379                r0 += tr;
1380            }
1381        })
1382    });
1383}
1384
1385fn qmatmat(
1386    q: &[u8],
1387    row_scale: &[f32],
1388    pre: &[std::borrow::Cow<'_, [f32]>],
1389    rows: usize,
1390    cols: usize,
1391    out: &mut [f32],
1392    pool: Option<&Pool>,
1393) {
1394    let b = pre.len();
1395    debug_assert_eq!(out.len(), b * rows);
1396    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
1397    // SDOT loop below peaks near the CPU's dot throughput, an order
1398    // below the matrix units. Small tensors and tiny test models stay
1399    // on the exact integer path.
1400    #[cfg(target_os = "macos")]
1401    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1402        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1403        return;
1404    }
1405    #[cfg(target_arch = "aarch64")]
1406    if sdot_enabled() {
1407        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1408        let out_addr = SendMut(out.as_mut_ptr());
1409        let run = |start: usize, end: usize| {
1410            for o in start..end {
1411                let row = &q[o * cols..(o + 1) * cols];
1412                for (bi, act) in acts.iter().enumerate() {
1413                    let v = row_dot_sdot(row, act) * row_scale[o];
1414                    unsafe { *out_addr.at(bi * rows + o) = v };
1415                }
1416            }
1417        };
1418        dispatch_rows(pool, rows, &run);
1419        return;
1420    }
1421    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
1422    // (roadmap P0: two weight rows' abs() stay in registers across four
1423    // activation streams); VNNI machines keep the per-row bias-trick
1424    // dot, which is already throughput-bound there.
1425    #[cfg(target_arch = "x86_64")]
1426    if avx2_a8w8_enabled() {
1427        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1428        let out_addr = SendMut(out.as_mut_ptr());
1429        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
1430        // A/B on noisy shared-vCPU hosts).
1431        let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1432        if !avx512vnni_enabled() && blocked_ok {
1433            let run = |start: usize, end: usize| {
1434                let mut o = start;
1435                while o < end {
1436                    if o + 2 <= end {
1437                        let r0 = &q[o * cols..(o + 1) * cols];
1438                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1439                        let mut bi = 0usize;
1440                        while bi + 4 <= acts.len() {
1441                            let xs = [
1442                                acts[bi].xq.as_slice(),
1443                                acts[bi + 1].xq.as_slice(),
1444                                acts[bi + 2].xq.as_slice(),
1445                                acts[bi + 3].xq.as_slice(),
1446                            ];
1447                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1448                            for (r, row) in [r0, r1].into_iter().enumerate() {
1449                                for k in 0..4 {
1450                                    let act = &acts[bi + k];
1451                                    let mut v = d[r][k] as f32 * act.sx;
1452                                    for &(j, xv) in &act.outliers {
1453                                        v += (row[j] as i8) as f32 * xv;
1454                                    }
1455                                    unsafe {
1456                                        *out_addr.at((bi + k) * rows + o + r) =
1457                                            v * row_scale[o + r]
1458                                    };
1459                                }
1460                            }
1461                            bi += 4;
1462                        }
1463                        while bi < acts.len() {
1464                            for (r, row) in [r0, r1].into_iter().enumerate() {
1465                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1466                                unsafe { *out_addr.at(bi * rows + o + r) = v };
1467                            }
1468                            bi += 1;
1469                        }
1470                        o += 2;
1471                    } else {
1472                        let row = &q[o * cols..(o + 1) * cols];
1473                        for (bi, act) in acts.iter().enumerate() {
1474                            let v = row_dot_avx2(row, act) * row_scale[o];
1475                            unsafe { *out_addr.at(bi * rows + o) = v };
1476                        }
1477                        o += 1;
1478                    }
1479                }
1480            };
1481            dispatch_rows(pool, rows, &run);
1482            return;
1483        }
1484        let run = |start: usize, end: usize| {
1485            for o in start..end {
1486                let row = &q[o * cols..(o + 1) * cols];
1487                for (bi, act) in acts.iter().enumerate() {
1488                    let v = row_dot_avx2(row, act) * row_scale[o];
1489                    unsafe { *out_addr.at(bi * rows + o) = v };
1490                }
1491            }
1492        };
1493        dispatch_rows(pool, rows, &run);
1494        return;
1495    }
1496    let out_addr = SendMut(out.as_mut_ptr());
1497    let run = |start: usize, end: usize| {
1498        for o in start..end {
1499            let row = &q[o * cols..(o + 1) * cols];
1500            for (bi, x) in pre.iter().enumerate() {
1501                let mut acc = 0f32;
1502                for j in 0..cols {
1503                    acc += (row[j] as i8) as f32 * x[j];
1504                }
1505                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1506            }
1507        }
1508    };
1509    dispatch_rows(pool, rows, &run);
1510}
1511
1512/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
1513/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
1514fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1515    match pool {
1516        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1517        _ => run(0, rows),
1518    }
1519}
1520
1521/// Split a q4_block blob into (packed nibbles, f16 group scales).
1522fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1523    let groups = rows * cols / GROUP_SIZE;
1524    bytes.split_at(groups * 16)
1525}
1526
1527/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
1528/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
1529/// vbit packs MSB-first, so the HIGH nibble is the even element
1530/// (opposite of q4_block's lo-first interleave). Centering is u-7.
1531#[inline]
1532fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1533    #[cfg(target_arch = "aarch64")]
1534    unsafe {
1535        return vbit_fill4_neon(data, buf);
1536    }
1537    #[cfg(target_arch = "x86_64")]
1538    if avx2_enabled() {
1539        return unsafe { vbit_fill4_avx2(data, buf) };
1540    }
1541    #[allow(unreachable_code)]
1542    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1543        let u = unpack8::<4>(&data[blk * 4..]);
1544        for k in 0..8 {
1545            chunk[k] = (u[k] - 7) as i8 as u8;
1546        }
1547    }
1548}
1549
1550#[cfg(target_arch = "aarch64")]
1551#[target_feature(enable = "neon")]
1552unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1553    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
1554    // buf.len()/2 packed bytes (validated at load).
1555    unsafe {
1556        use core::arch::aarch64::*;
1557        let n = buf.len();
1558        let mask = vdupq_n_u8(0x0F);
1559        let seven = vdupq_n_s8(7);
1560        let mut g = 0usize;
1561        while g * 32 + 32 <= n {
1562            let b = vld1q_u8(data.as_ptr().add(g * 16));
1563            let hi = vshrq_n_u8::<4>(b);
1564            let lo = vandq_u8(b, mask);
1565            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1566            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1567            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1568            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1569            g += 1;
1570        }
1571    }
1572}
1573
1574#[cfg(target_arch = "x86_64")]
1575#[target_feature(enable = "avx2")]
1576unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1577    // SAFETY: see vbit_fill4_neon.
1578    unsafe {
1579        use core::arch::x86_64::*;
1580        let n = buf.len();
1581        let mask = _mm_set1_epi8(0x0F);
1582        let seven = _mm256_set1_epi8(7);
1583        let mut g = 0usize;
1584        while g * 32 + 32 <= n {
1585            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1586            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1587            let lo = _mm_and_si128(b, mask);
1588            let z = _mm256_sub_epi8(
1589                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1590                seven,
1591            );
1592            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1593            g += 1;
1594        }
1595    }
1596}
1597
1598/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
1599/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
1600/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
1601/// into 4 such blocks.
1602#[inline(always)]
1603fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1604    let mut acc = 0u64;
1605    for i in 0..B {
1606        acc = (acc << 8) | data[i] as u64;
1607    }
1608    let mask = (1u64 << B) - 1;
1609    let mut out = [0i32; 8];
1610    for (k, o) in out.iter_mut().enumerate() {
1611        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1612    }
1613    out
1614}
1615
1616/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
1617/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
1618/// MSB-first, byte-padded]. Row data offsets are precomputed at load
1619/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
1620/// overhead on every matvec.
1621#[allow(clippy::too_many_arguments)]
1622fn vbitmatvec(
1623    bytes: &[u8],
1624    offsets: &[usize],
1625    x: &[f32],
1626    rows: usize,
1627    cols: usize,
1628    out: &mut [f32],
1629    pool: Option<&Pool>,
1630) {
1631    debug_assert_eq!(out.len(), rows);
1632    debug_assert_eq!(offsets.len(), rows + 1);
1633
1634    // SDOT path: unpack the row to centered i8 once, then per-group
1635    // int8 dot against the quantized activations — same A8W8 contract
1636    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
1637    if a8w8_enabled() {
1638        let act = split_act(x);
1639        let out_addr = SendMut(out.as_mut_ptr());
1640        let run = move |start: usize, end: usize| {
1641            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1642        };
1643        dispatch_rows(pool, rows, &run);
1644        return;
1645    }
1646
1647    let out_addr = SendMut(out.as_mut_ptr());
1648    let run = move |start: usize, end: usize| {
1649        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1650    };
1651    dispatch_rows(pool, rows, &run);
1652}
1653
1654/// One vbit row range via the A8W8 int8 path — kernel body of
1655/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
1656/// several tensors in one dispatch (b=8 rows go exact f32).
1657#[allow(clippy::too_many_arguments)]
1658fn vbit_range_a8w8(
1659    bytes: &[u8],
1660    offsets: &[usize],
1661    x: &[f32],
1662    act: &SplitAct,
1663    rows: usize,
1664    cols: usize,
1665    out: SendMut,
1666    start: usize,
1667    end: usize,
1668) {
1669    let ng = cols / GROUP_SIZE;
1670    let bits = &bytes[..rows];
1671    let sc_off = rows;
1672    let row_dot = |r: usize| -> f32 {
1673            let b = bits[r] as usize;
1674            let l = ((1i32 << (b - 1)) - 1) as i32;
1675            let mask = (1u64 << b) - 1;
1676            let data = &bytes[offsets[r]..offsets[r + 1]];
1677            if b == 8 {
1678                // u−L reaches 128 → does not fit i8; exact f32 path.
1679                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1680                let mut dot = 0f32;
1681                for g in 0..ng {
1682                    let so = (r * ng + g) * 2;
1683                    let sgf = f16_to_f32(u16::from_le_bytes([
1684                        bytes[sc_off + so],
1685                        bytes[sc_off + so + 1],
1686                    ]));
1687                    let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1688                    let mut gd = 0f32;
1689                    for &xv in xg.iter() {
1690                        if nbits < 8 {
1691                            acc = (acc << 8) | data[idx] as u64;
1692                            idx += 1;
1693                            nbits += 8;
1694                        }
1695                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1696                        nbits -= 8;
1697                        gd += (u - l) as f32 * xv;
1698                    }
1699                    dot += gd * sgf;
1700                }
1701                return dot;
1702            }
1703            // Per-worker scratch: this closure runs for every row of the
1704            // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
1705            // row was measurable pure overhead.
1706            thread_local! {
1707                static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1708                    const { std::cell::RefCell::new(Vec::new()) };
1709            }
1710            #[inline(always)]
1711            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1712                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1713                    let u = unpack8::<B>(&data[blk * B..]);
1714                    for k in 0..8 {
1715                        chunk[k] = (u[k] - l) as i8 as u8;
1716                    }
1717                }
1718            }
1719            let _ = mask;
1720            VBIT_SCRATCH.with(|scratch| {
1721                let mut buf = scratch.borrow_mut();
1722                buf.resize(cols, 0);
1723                match b {
1724                    3 => fill::<3>(data, l, &mut buf),
1725                    4 => vbit_fill4(data, &mut buf),
1726                    5 => fill::<5>(data, l, &mut buf),
1727                    6 => fill::<6>(data, l, &mut buf),
1728                    _ => unreachable!(),
1729                }
1730                let mut dot = 0f32;
1731                for g in 0..ng {
1732                    let so = (r * ng + g) * 2;
1733                    let s = f16_to_f32(u16::from_le_bytes([
1734                        bytes[sc_off + so],
1735                        bytes[sc_off + so + 1],
1736                    ]));
1737                    let d = dot_i8_i8(
1738                        &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1739                        &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1740                    ) as f32
1741                        * act.sx;
1742                    dot += d * s;
1743                }
1744                for &(j, xv) in &act.outliers {
1745                    let so = (r * ng + j / GROUP_SIZE) * 2;
1746                    let s = f16_to_f32(u16::from_le_bytes([
1747                        bytes[sc_off + so],
1748                        bytes[sc_off + so + 1],
1749                    ]));
1750                    // xq is zeroed at outlier slots — add the exact term.
1751                    dot += (buf[j] as i8) as f32 * s * xv;
1752                }
1753                dot
1754            })
1755    };
1756    for r in start..end {
1757        // SAFETY: disjoint row ranges per worker.
1758        unsafe { *out.at(r) = row_dot(r) };
1759    }
1760}
1761
1762/// Exact scalar vbit row range (same extraction, non-SDOT path).
1763#[allow(clippy::too_many_arguments)]
1764fn vbit_range_f32(
1765    bytes: &[u8],
1766    offsets: &[usize],
1767    x: &[f32],
1768    rows: usize,
1769    cols: usize,
1770    out: SendMut,
1771    start: usize,
1772    end: usize,
1773) {
1774    let ng = cols / GROUP_SIZE;
1775    let bits = &bytes[..rows];
1776    let sc_off = rows;
1777    // Per-bit-width specialized inner loops: the compiler unrolls the
1778    // constant shifts (the generic bit-buffer loop was branch-bound —
1779    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
1780    #[inline(always)]
1781    fn dot_row<const B: usize>(
1782        data: &[u8],
1783        bytes: &[u8],
1784        sc_off: usize,
1785        r: usize,
1786        ng: usize,
1787        x: &[f32],
1788    ) -> f32 {
1789        let l = ((1i32 << (B - 1)) - 1) as f32;
1790        let gbytes = GROUP_SIZE * B / 8;
1791        let mut dot = 0f32;
1792        for g in 0..ng {
1793            let so = (r * ng + g) * 2;
1794            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1795            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1796            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1797            let mut gd = 0f32;
1798            for blk in 0..GROUP_SIZE / 8 {
1799                let u = unpack8::<B>(&gd0[blk * B..]);
1800                let xb = &xg[blk * 8..blk * 8 + 8];
1801                for k in 0..8 {
1802                    gd += (u[k] as f32 - l) * xb[k];
1803                }
1804            }
1805            dot += gd * s;
1806        }
1807        dot
1808    }
1809    for r in start..end {
1810        let data = &bytes[offsets[r]..offsets[r + 1]];
1811        let v = match bits[r] {
1812            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1813            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1814            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1815            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1816            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1817            b => unreachable!("vbit bit-width {b} (validated at load)"),
1818        };
1819        // SAFETY: disjoint row ranges per worker.
1820        unsafe { *out.at(r) = v };
1821    }
1822}
1823
1824/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
1825/// and dotted against BOTH activations (MTP verify / pair prefill used
1826/// to run two full matvecs — double weight traffic and double unpack).
1827/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
1828#[allow(clippy::too_many_arguments)]
1829fn vbitmatvec2(
1830    bytes: &[u8],
1831    offsets: &[usize],
1832    x1: &[f32],
1833    x2: &[f32],
1834    rows: usize,
1835    cols: usize,
1836    o1: &mut [f32],
1837    o2: &mut [f32],
1838    pool: Option<&Pool>,
1839) {
1840    debug_assert_eq!(o1.len(), rows);
1841    debug_assert_eq!(o2.len(), rows);
1842
1843    if a8w8_enabled() {
1844        let a1 = split_act(x1);
1845        let a2 = split_act(x2);
1846        let p1 = SendMut(o1.as_mut_ptr());
1847        let p2 = SendMut(o2.as_mut_ptr());
1848        let run = move |start: usize, end: usize| {
1849            vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1850        };
1851        dispatch_rows(pool, rows, &run);
1852        return;
1853    }
1854
1855    let p1 = SendMut(o1.as_mut_ptr());
1856    let p2 = SendMut(o2.as_mut_ptr());
1857    let run = move |start: usize, end: usize| {
1858        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1859    };
1860    dispatch_rows(pool, rows, &run);
1861}
1862
1863/// Two-input vbit row range via the A8W8 int8 path — kernel body of
1864/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
1865/// exact f32 for both lanes, bits streamed once).
1866#[allow(clippy::too_many_arguments)]
1867fn vbit_range2_a8w8(
1868    bytes: &[u8],
1869    offsets: &[usize],
1870    x1: &[f32],
1871    x2: &[f32],
1872    a1: &SplitAct,
1873    a2: &SplitAct,
1874    rows: usize,
1875    cols: usize,
1876    p1: SendMut,
1877    p2: SendMut,
1878    start: usize,
1879    end: usize,
1880) {
1881    let ng = cols / GROUP_SIZE;
1882    let bits = &bytes[..rows];
1883    let sc_off = rows;
1884    let row_dots = |r: usize| -> (f32, f32) {
1885            let b = bits[r] as usize;
1886            let l = (1i32 << (b - 1)) - 1;
1887            let data = &bytes[offsets[r]..offsets[r + 1]];
1888            if b == 8 {
1889                // u−L reaches 128 → does not fit i8; exact f32 path,
1890                // bits still streamed once for both lanes.
1891                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1892                let (mut d1, mut d2) = (0f32, 0f32);
1893                for g in 0..ng {
1894                    let so = (r * ng + g) * 2;
1895                    let sgf = f16_to_f32(u16::from_le_bytes([
1896                        bytes[sc_off + so],
1897                        bytes[sc_off + so + 1],
1898                    ]));
1899                    let (mut g1, mut g2) = (0f32, 0f32);
1900                    for k in 0..GROUP_SIZE {
1901                        if nbits < 8 {
1902                            acc = (acc << 8) | data[idx] as u64;
1903                            idx += 1;
1904                            nbits += 8;
1905                        }
1906                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1907                        nbits -= 8;
1908                        let w = (u - l) as f32;
1909                        g1 += w * x1[g * GROUP_SIZE + k];
1910                        g2 += w * x2[g * GROUP_SIZE + k];
1911                    }
1912                    d1 += g1 * sgf;
1913                    d2 += g2 * sgf;
1914                }
1915                return (d1, d2);
1916            }
1917            thread_local! {
1918                static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1919                    const { std::cell::RefCell::new(Vec::new()) };
1920            }
1921            #[inline(always)]
1922            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1923                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1924                    let u = unpack8::<B>(&data[blk * B..]);
1925                    for k in 0..8 {
1926                        chunk[k] = (u[k] - l) as i8 as u8;
1927                    }
1928                }
1929            }
1930            VBIT_SCRATCH2.with(|scratch| {
1931                let mut buf = scratch.borrow_mut();
1932                buf.resize(cols, 0);
1933                match b {
1934                    3 => fill::<3>(data, l, &mut buf),
1935                    4 => vbit_fill4(data, &mut buf),
1936                    5 => fill::<5>(data, l, &mut buf),
1937                    6 => fill::<6>(data, l, &mut buf),
1938                    _ => unreachable!(),
1939                }
1940                let (mut d1, mut d2) = (0f32, 0f32);
1941                for g in 0..ng {
1942                    let so = (r * ng + g) * 2;
1943                    let s = f16_to_f32(u16::from_le_bytes([
1944                        bytes[sc_off + so],
1945                        bytes[sc_off + so + 1],
1946                    ]));
1947                    let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1948                    let v1 =
1949                        dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1950                    let v2 =
1951                        dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1952                    d1 += v1 * s;
1953                    d2 += v2 * s;
1954                }
1955                for &(j, xv) in &a1.outliers {
1956                    let so = (r * ng + j / GROUP_SIZE) * 2;
1957                    let s = f16_to_f32(u16::from_le_bytes([
1958                        bytes[sc_off + so],
1959                        bytes[sc_off + so + 1],
1960                    ]));
1961                    d1 += (buf[j] as i8) as f32 * s * xv;
1962                }
1963                for &(j, xv) in &a2.outliers {
1964                    let so = (r * ng + j / GROUP_SIZE) * 2;
1965                    let s = f16_to_f32(u16::from_le_bytes([
1966                        bytes[sc_off + so],
1967                        bytes[sc_off + so + 1],
1968                    ]));
1969                    d2 += (buf[j] as i8) as f32 * s * xv;
1970                }
1971                (d1, d2)
1972            })
1973    };
1974    for r in start..end {
1975        let (v1, v2) = row_dots(r);
1976        // SAFETY: disjoint row ranges per worker.
1977        unsafe {
1978            *p1.at(r) = v1;
1979            *p2.at(r) = v2;
1980        }
1981    }
1982}
1983
1984/// Two-input exact scalar vbit row range (same extraction) —
1985/// per-bit-width specialized, two accumulators per row; per-lane
1986/// accumulation order matches `vbitmatvec` exactly.
1987#[allow(clippy::too_many_arguments)]
1988fn vbit_range2_f32(
1989    bytes: &[u8],
1990    offsets: &[usize],
1991    x1: &[f32],
1992    x2: &[f32],
1993    rows: usize,
1994    cols: usize,
1995    p1: SendMut,
1996    p2: SendMut,
1997    start: usize,
1998    end: usize,
1999) {
2000    let ng = cols / GROUP_SIZE;
2001    let bits = &bytes[..rows];
2002    let sc_off = rows;
2003    #[inline(always)]
2004    #[allow(clippy::too_many_arguments)]
2005    fn dot_row2<const B: usize>(
2006        data: &[u8],
2007        bytes: &[u8],
2008        sc_off: usize,
2009        r: usize,
2010        ng: usize,
2011        x1: &[f32],
2012        x2: &[f32],
2013    ) -> (f32, f32) {
2014        let l = ((1i32 << (B - 1)) - 1) as f32;
2015        let gbytes = GROUP_SIZE * B / 8;
2016        let (mut d1, mut d2) = (0f32, 0f32);
2017        for g in 0..ng {
2018            let so = (r * ng + g) * 2;
2019            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2020            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2021            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2022            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2023            let (mut g1, mut g2) = (0f32, 0f32);
2024            for blk in 0..GROUP_SIZE / 8 {
2025                let u = unpack8::<B>(&gd0[blk * B..]);
2026                for k in 0..8 {
2027                    let w = u[k] as f32 - l;
2028                    g1 += w * x1g[blk * 8 + k];
2029                    g2 += w * x2g[blk * 8 + k];
2030                }
2031            }
2032            d1 += g1 * s;
2033            d2 += g2 * s;
2034        }
2035        (d1, d2)
2036    }
2037    for r in start..end {
2038        let data = &bytes[offsets[r]..offsets[r + 1]];
2039        let (v1, v2) = match bits[r] {
2040            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2041            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2042            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2043            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2044            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2045            b => unreachable!("vbit bit-width {b} (validated at load)"),
2046        };
2047        // SAFETY: disjoint row ranges per worker.
2048        unsafe {
2049            *p1.at(r) = v1;
2050            *p2.at(r) = v2;
2051        }
2052    }
2053}
2054
2055// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
2056
2057/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
2058/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
2059/// distant streams of the split layout. Values/order identical to the
2060/// split kernels.
2061#[inline]
2062#[allow(unreachable_code)]
2063fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2064    #[cfg(target_arch = "aarch64")]
2065    unsafe {
2066        return dot_q4t_row_sdot(bytes, r, gpr, xq);
2067    }
2068    #[cfg(target_arch = "x86_64")]
2069    unsafe {
2070        return dot_q4t_row_avx2(bytes, r, gpr, xq);
2071    }
2072    let mut acc = 0f32;
2073    for gi in 0..gpr {
2074        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2075        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2076        let mut d = 0i32;
2077        for (k, &b) in tile[2..].iter().enumerate() {
2078            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2079                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2080        }
2081        acc += d as f32 * s;
2082    }
2083    acc
2084}
2085
2086#[cfg(target_arch = "aarch64")]
2087#[target_feature(enable = "neon,dotprod")]
2088unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2089    // SAFETY: callers uphold slice-length contracts (18B tile per group,
2090    // xq.len() == gpr·GROUP_SIZE).
2091    unsafe {
2092        use core::arch::aarch64::*;
2093        use core::arch::asm;
2094        let lomask = vdupq_n_u8(0x0F);
2095        let eight = vdupq_n_s8(8);
2096        let mut acc = 0f32;
2097        for gi in 0..gpr {
2098            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2099            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2100            let b = vld1q_u8(t.add(2));
2101            let lo = vandq_u8(b, lomask);
2102            let hi = vshrq_n_u8::<4>(b);
2103            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2104            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2105            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2106            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2107            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2108            asm!(
2109                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2110                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2111                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2112                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2113                options(pure, nomem, nostack),
2114            );
2115            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2116        }
2117        acc
2118    }
2119}
2120
2121#[cfg(target_arch = "x86_64")]
2122#[target_feature(enable = "avx2")]
2123unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2124    // SAFETY: see dot_q4t_row_sdot.
2125    unsafe {
2126        use core::arch::x86_64::*;
2127        let lomask = _mm_set1_epi8(0x0F);
2128        let eight = _mm256_set1_epi8(8);
2129        let ones = _mm256_set1_epi16(1);
2130        let mut acc = 0f32;
2131        for gi in 0..gpr {
2132            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2133            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2134            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2135            let lo = _mm_and_si128(b, lomask);
2136            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2137            let w = _mm256_sub_epi8(
2138                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2139                eight,
2140            );
2141            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2142            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2143            let d = _mm256_madd_epi16(p16, ones);
2144            let hi128 = _mm256_extracti128_si256::<1>(d);
2145            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2146            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2147            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2148            acc += _mm_cvtsi128_si32(s32) as f32 * s;
2149        }
2150        acc
2151    }
2152}
2153
2154/// One q4_tiled row against FOUR activation streams: the nibble unpack
2155/// and abs() happen once per group instead of once per (group,
2156/// activation) — the unpack is the dominant per-element cost of the
2157/// tiled format (roadmap P0 portable blocking, q4t leg).
2158#[cfg(target_arch = "x86_64")]
2159#[target_feature(enable = "avx2")]
2160unsafe fn dot_q4t_row_1x4_avx2(
2161    bytes: &[u8],
2162    r: usize,
2163    gpr: usize,
2164    xs: [&[i8]; 4],
2165) -> [f32; 4] {
2166    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
2167    unsafe {
2168        use core::arch::x86_64::*;
2169        let lomask = _mm_set1_epi8(0x0F);
2170        let eight = _mm256_set1_epi8(8);
2171        let ones = _mm256_set1_epi16(1);
2172        let mut acc = [0f32; 4];
2173        for gi in 0..gpr {
2174            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2175            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2176            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2177            let lo = _mm_and_si128(bb, lomask);
2178            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2179            let w = _mm256_sub_epi8(
2180                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2181                eight,
2182            );
2183            let aw = _mm256_abs_epi8(w);
2184            for (k, xq) in xs.iter().enumerate() {
2185                let x =
2186                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2187                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2188                let d = _mm256_madd_epi16(p16, ones);
2189                let hi128 = _mm256_extracti128_si256::<1>(d);
2190                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2191                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2192                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2193                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2194            }
2195        }
2196        acc
2197    }
2198}
2199
2200/// Exact-term correction for A8W8 outliers on a tiled row.
2201#[inline]
2202fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2203    let gi = j / GROUP_SIZE;
2204    let k = j % GROUP_SIZE;
2205    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2206    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2207    let byte = tile[2 + k / 2];
2208    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2209    ((nib as i32 - 8) as f32, s)
2210}
2211
2212/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
2213/// accumulation shape as `q4_range_f32`.
2214#[inline]
2215fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2216    let mut acc = 0f32;
2217    for gi in 0..gpr {
2218        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2219        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2220        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2221        let mut ga = 0f32;
2222        for (k, &b) in tile[2..].iter().enumerate() {
2223            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2224                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2225        }
2226        acc += ga * s;
2227    }
2228    acc
2229}
2230
2231/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
2232fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2233    debug_assert_eq!(out.len(), rows);
2234    let gpr = cols / GROUP_SIZE;
2235    let out_addr = SendMut(out.as_mut_ptr());
2236    if a8w8_enabled() {
2237        let act = split_act(x);
2238        let run = move |start: usize, end: usize| {
2239            for r in start..end {
2240                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2241                for &(j, xv) in &act.outliers {
2242                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2243                    acc += w * s * xv;
2244                }
2245                // SAFETY: disjoint row ranges per worker.
2246                unsafe { *out_addr.at(r) = acc };
2247            }
2248        };
2249        dispatch_rows(pool, rows, &run);
2250        return;
2251    }
2252    let run = move |start: usize, end: usize| {
2253        for r in start..end {
2254            // SAFETY: disjoint row ranges per worker.
2255            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2256        }
2257    };
2258    dispatch_rows(pool, rows, &run);
2259}
2260
2261/// Fused two-input q4_tiled matvec (weights read once per pair).
2262#[allow(clippy::too_many_arguments)]
2263fn q4t_matvec2(
2264    bytes: &[u8],
2265    x1: &[f32],
2266    x2: &[f32],
2267    rows: usize,
2268    cols: usize,
2269    o1: &mut [f32],
2270    o2: &mut [f32],
2271    pool: Option<&Pool>,
2272) {
2273    let gpr = cols / GROUP_SIZE;
2274    let p1 = SendMut(o1.as_mut_ptr());
2275    let p2 = SendMut(o2.as_mut_ptr());
2276    if a8w8_enabled() {
2277        let a1 = split_act(x1);
2278        let a2 = split_act(x2);
2279        let run = move |start: usize, end: usize| {
2280            for r in start..end {
2281                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2282                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2283                for &(j, xv) in &a1.outliers {
2284                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2285                    v1 += w * s * xv;
2286                }
2287                for &(j, xv) in &a2.outliers {
2288                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
2289                    v2 += w * s * xv;
2290                }
2291                // SAFETY: disjoint row ranges per worker.
2292                unsafe {
2293                    *p1.at(r) = v1;
2294                    *p2.at(r) = v2;
2295                }
2296            }
2297        };
2298        dispatch_rows(pool, rows, &run);
2299        return;
2300    }
2301    let run = move |start: usize, end: usize| {
2302        for r in start..end {
2303            // SAFETY: disjoint row ranges per worker.
2304            unsafe {
2305                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2306                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2307            }
2308        }
2309    };
2310    dispatch_rows(pool, rows, &run);
2311}
2312
2313/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
2314#[allow(clippy::too_many_arguments)]
2315fn q4t_matmat(
2316    bytes: &[u8],
2317    xs_all: &[f32],
2318    b: usize,
2319    rows: usize,
2320    cols: usize,
2321    out: &mut [f32],
2322    pool: Option<&Pool>,
2323) {
2324    debug_assert_eq!(out.len(), b * rows);
2325    let gpr = cols / GROUP_SIZE;
2326    let out_addr = SendMut(out.as_mut_ptr());
2327    if a8w8_enabled() {
2328        let acts: Vec<SplitAct> =
2329            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2330        let acts = &acts;
2331        #[cfg(target_arch = "x86_64")]
2332        let blocked_ok = avx2_enabled()
2333            && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2334        #[cfg(not(target_arch = "x86_64"))]
2335        let blocked_ok = false;
2336        let run = move |start: usize, end: usize| {
2337            for r in start..end {
2338                let mut bi = 0usize;
2339                #[cfg(target_arch = "x86_64")]
2340                if blocked_ok {
2341                    while bi + 4 <= acts.len() {
2342                        let xs = [
2343                            acts[bi].xq.as_slice(),
2344                            acts[bi + 1].xq.as_slice(),
2345                            acts[bi + 2].xq.as_slice(),
2346                            acts[bi + 3].xq.as_slice(),
2347                        ];
2348                        let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2349                        for k in 0..4 {
2350                            let act = &acts[bi + k];
2351                            let mut acc = d[k] * act.sx;
2352                            for &(j, xv) in &act.outliers {
2353                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2354                                acc += w * sc * xv;
2355                            }
2356                            // SAFETY: disjoint (bi, r) cells per worker.
2357                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2358                        }
2359                        bi += 4;
2360                    }
2361                }
2362                let _ = blocked_ok;
2363                while bi < acts.len() {
2364                    let act = &acts[bi];
2365                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2366                    for &(j, xv) in &act.outliers {
2367                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
2368                        acc += w * s * xv;
2369                    }
2370                    // SAFETY: disjoint (bi, r) cells per worker range.
2371                    unsafe { *out_addr.at(bi * rows + r) = acc };
2372                    bi += 1;
2373                }
2374            }
2375        };
2376        dispatch_rows(pool, rows, &run);
2377        return;
2378    }
2379    let run = move |start: usize, end: usize| {
2380        for r in start..end {
2381            for bi in 0..b {
2382                let x = &xs_all[bi * cols..(bi + 1) * cols];
2383                // SAFETY: disjoint (bi, r) cells per worker range.
2384                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2385            }
2386        }
2387    };
2388    dispatch_rows(pool, rows, &run);
2389}
2390
2391// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
2392// 32-group tile. The kernel family mirrors q4_tiled: one sequential
2393// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
2394// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
2395
2396/// Per-32-group sums of the quantized activation — the ±1 identity's
2397/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
2398/// matvec and reused by every row.
2399fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2400    (0..gpr)
2401        .map(|gi| {
2402            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2403                .iter()
2404                .map(|&v| v as i32)
2405                .sum()
2406        })
2407        .collect()
2408}
2409
2410/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
2411/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
2412/// x86 pass).
2413#[inline]
2414#[allow(unreachable_code)]
2415fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2416    #[cfg(target_arch = "aarch64")]
2417    unsafe {
2418        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2419    }
2420    let _ = gsum;
2421    let mut acc = 0f32;
2422    for gi in 0..gpr {
2423        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2424        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2425        let mut d = 0i32;
2426        for (j, &b) in tile[2..].iter().enumerate() {
2427            for k in 0..8 {
2428                let w = ((b >> k) & 1) as i32 * 2 - 1;
2429                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2430            }
2431        }
2432        acc += d as f32 * s;
2433    }
2434    acc
2435}
2436
2437/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
2438/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
2439/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
2440/// per-group activation sums shared across every row of the matvec.
2441/// Four tiles (128 weights) per iteration: integer dots reduce through
2442/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
2443/// fused f32 multiply-add. Integer math throughout — bit-identical to
2444/// the scalar ±1 reference.
2445#[cfg(target_arch = "aarch64")]
2446#[target_feature(enable = "neon,dotprod")]
2447unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2448    // SAFETY: callers uphold slice-length contracts (6B tile per group,
2449    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
2450    unsafe {
2451        use core::arch::aarch64::*;
2452        use core::arch::asm;
2453        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2454        let m = vld1q_u8(MASKS.as_ptr());
2455        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
2456        macro_rules! tile_dot {
2457            ($t:expr, $x:expr) => {{
2458                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2459                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2460                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2461                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2462                let x0 = vld1q_s8($x);
2463                let x1 = vld1q_s8($x.add(16));
2464                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2465                asm!(
2466                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2467                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2468                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2469                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2470                    options(pure, nomem, nostack),
2471                );
2472                vaddq_s32(a0, a1)
2473            }};
2474        }
2475        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2476        let xp = xq.as_ptr();
2477        let gp = gsum.as_ptr();
2478        let mut accv = vdupq_n_f32(0.0);
2479        let mut gi = 0;
2480        while gi + 4 <= gpr {
2481            let t0 = base.add(gi * Q1_TILE);
2482            let t1 = base.add((gi + 1) * Q1_TILE);
2483            let t2 = base.add((gi + 2) * Q1_TILE);
2484            let t3 = base.add((gi + 3) * Q1_TILE);
2485            let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2486            let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2487            let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2488            let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2489            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
2490            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2491            let g = vld1q_s32(gp.add(gi));
2492            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2493            let sc = [
2494                f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2495                f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2496                f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2497                f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2498            ];
2499            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2500            gi += 4;
2501        }
2502        let mut acc = vaddvq_f32(accv);
2503        while gi < gpr {
2504            let t = base.add(gi * Q1_TILE);
2505            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2506            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2507            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2508            gi += 1;
2509        }
2510        acc
2511    }
2512}
2513
2514/// (weight ±1, scale) of one q1 element — the exact outlier term.
2515#[inline]
2516fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2517    let gi = j / GROUP_SIZE;
2518    let k = j % GROUP_SIZE;
2519    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2520    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2521    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2522    ((bit as i32 * 2 - 1) as f32, s)
2523}
2524
2525/// Exact scalar q1 row (CMF_SDOT=0 contract).
2526#[inline]
2527fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2528    let mut acc = 0f32;
2529    for gi in 0..gpr {
2530        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2531        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2532        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2533        let mut ga = 0f32;
2534        for (j, &b) in tile[2..].iter().enumerate() {
2535            for k in 0..8 {
2536                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2537            }
2538        }
2539        acc += ga * s;
2540    }
2541    acc
2542}
2543
2544/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
2545/// extracted so multi-matrix jobs drive the same kernel).
2546#[allow(clippy::too_many_arguments)]
2547fn q1_range_a8w8(
2548    bytes: &[u8],
2549    gpr: usize,
2550    act: &SplitAct,
2551    gsum: &[i32],
2552    out: SendMut,
2553    start: usize,
2554    end: usize,
2555) {
2556    for r in start..end {
2557        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2558        for &(j, xv) in &act.outliers {
2559            let (w, s) = q1_outlier(bytes, r, gpr, j);
2560            acc += w * s * xv;
2561        }
2562        // SAFETY: disjoint row ranges per worker.
2563        unsafe { *out.at(r) = acc };
2564    }
2565}
2566
2567/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
2568fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2569    for r in start..end {
2570        // SAFETY: disjoint row ranges per worker.
2571        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2572    }
2573}
2574
2575/// Fused q1 matvec (dispatch mirrors `q4t_matvec`).
2576fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2577    debug_assert_eq!(out.len(), rows);
2578    let gpr = cols / GROUP_SIZE;
2579    let out_addr = SendMut(out.as_mut_ptr());
2580    if a8w8_enabled() {
2581        let act = split_act(x);
2582        let gsum = q1_group_sums(&act.xq, gpr);
2583        let (act, gsum) = (&act, &gsum);
2584        let run =
2585            move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2586        dispatch_rows(pool, rows, &run);
2587        return;
2588    }
2589    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2590    dispatch_rows(pool, rows, &run);
2591}
2592
2593/// Fused two-input q1 matvec (weights read once per pair).
2594#[allow(clippy::too_many_arguments)]
2595fn q1_matvec2(
2596    bytes: &[u8],
2597    x1: &[f32],
2598    x2: &[f32],
2599    rows: usize,
2600    cols: usize,
2601    o1: &mut [f32],
2602    o2: &mut [f32],
2603    pool: Option<&Pool>,
2604) {
2605    let gpr = cols / GROUP_SIZE;
2606    let p1 = SendMut(o1.as_mut_ptr());
2607    let p2 = SendMut(o2.as_mut_ptr());
2608    if a8w8_enabled() {
2609        let a1 = split_act(x1);
2610        let a2 = split_act(x2);
2611        let g1 = q1_group_sums(&a1.xq, gpr);
2612        let g2 = q1_group_sums(&a2.xq, gpr);
2613        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2614        let run = move |start: usize, end: usize| {
2615            for r in start..end {
2616                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2617                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2618                for &(j, xv) in &a1.outliers {
2619                    let (w, s) = q1_outlier(bytes, r, gpr, j);
2620                    v1 += w * s * xv;
2621                }
2622                for &(j, xv) in &a2.outliers {
2623                    let (w, s) = q1_outlier(bytes, r, gpr, j);
2624                    v2 += w * s * xv;
2625                }
2626                // SAFETY: disjoint row ranges per worker.
2627                unsafe {
2628                    *p1.at(r) = v1;
2629                    *p2.at(r) = v2;
2630                }
2631            }
2632        };
2633        dispatch_rows(pool, rows, &run);
2634        return;
2635    }
2636    let run = move |start: usize, end: usize| {
2637        for r in start..end {
2638            // SAFETY: disjoint row ranges per worker.
2639            unsafe {
2640                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2641                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2642            }
2643        }
2644    };
2645    dispatch_rows(pool, rows, &run);
2646}
2647
2648/// Batched q1 matmat: each row's tiles stream once per microbatch.
2649#[allow(clippy::too_many_arguments)]
2650fn q1_matmat(
2651    bytes: &[u8],
2652    xs_all: &[f32],
2653    b: usize,
2654    rows: usize,
2655    cols: usize,
2656    out: &mut [f32],
2657    pool: Option<&Pool>,
2658) {
2659    debug_assert_eq!(out.len(), b * rows);
2660    let gpr = cols / GROUP_SIZE;
2661    let out_addr = SendMut(out.as_mut_ptr());
2662    if a8w8_enabled() {
2663        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2664            .map(|bi| {
2665                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2666                let gsum = q1_group_sums(&act.xq, gpr);
2667                (act, gsum)
2668            })
2669            .collect();
2670        let acts = &acts;
2671        let run = move |start: usize, end: usize| {
2672            for r in start..end {
2673                for (bi, (act, gsum)) in acts.iter().enumerate() {
2674                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2675                    for &(j, xv) in &act.outliers {
2676                        let (w, s) = q1_outlier(bytes, r, gpr, j);
2677                        acc += w * s * xv;
2678                    }
2679                    // SAFETY: disjoint (bi, r) cells per worker range.
2680                    unsafe { *out_addr.at(bi * rows + r) = acc };
2681                }
2682            }
2683        };
2684        dispatch_rows(pool, rows, &run);
2685        return;
2686    }
2687    let run = move |start: usize, end: usize| {
2688        for r in start..end {
2689            for bi in 0..b {
2690                let x = &xs_all[bi * cols..(bi + 1) * cols];
2691                // SAFETY: disjoint (bi, r) cells per worker range.
2692                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
2693            }
2694        }
2695    };
2696    dispatch_rows(pool, rows, &run);
2697}
2698
2699/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
2700/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
2701/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
2702/// 32-group, exact outlier correction — the same A8W8 contract as q8.
2703/// `CMF_SDOT=0` keeps the exact scalar path.
2704fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2705    debug_assert_eq!(out.len(), rows);
2706    let (packed, scales) = q4_split(bytes, rows, cols);
2707    let gpr = cols / GROUP_SIZE;
2708    let out_addr = SendMut(out.as_mut_ptr());
2709
2710    if a8w8_enabled() {
2711        let act = split_act(x);
2712        let run = move |start: usize, end: usize| {
2713            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
2714        };
2715        dispatch_rows(pool, rows, &run);
2716        return;
2717    }
2718
2719    let run = move |start: usize, end: usize| {
2720        q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
2721    };
2722    dispatch_rows(pool, rows, &run);
2723}
2724
2725/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
2726/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
2727#[inline]
2728#[allow(unreachable_code)]
2729/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
2730/// streams: the 32-byte weight chunk and its abs() load once per group,
2731/// the per-group f16 scale decodes once — four maddubs+reduce chains
2732/// instead of four full (load, abs, dot) rounds.
2733#[cfg(target_arch = "x86_64")]
2734#[target_feature(enable = "avx2")]
2735unsafe fn dot_q4b_row_1x4_avx2(
2736    buf: &[u8],
2737    scales: &[u8],
2738    g0: usize,
2739    gpr: usize,
2740    xs: [&[i8]; 4],
2741) -> [f32; 4] {
2742    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
2743    unsafe {
2744        use core::arch::x86_64::*;
2745        let ones = _mm256_set1_epi16(1);
2746        let mut acc = [0f32; 4];
2747        for gi in 0..gpr {
2748            let s = f16_to_f32(u16::from_le_bytes([
2749                scales[(g0 + gi) * 2],
2750                scales[(g0 + gi) * 2 + 1],
2751            ]));
2752            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2753            let aw = _mm256_abs_epi8(w);
2754            for (k, xq) in xs.iter().enumerate() {
2755                let x =
2756                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2757                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2758                let d = _mm256_madd_epi16(p16, ones);
2759                let hi128 = _mm256_extracti128_si256::<1>(d);
2760                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2761                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2762                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2763                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2764            }
2765        }
2766        acc
2767    }
2768}
2769
2770/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
2771/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
2772/// accumulation order (the q4_block flavor applies sx once at the end,
2773/// matching ITS single path; the two conventions are historical and
2774/// each blocked leg must mirror its own).
2775#[cfg(target_arch = "x86_64")]
2776#[target_feature(enable = "avx2")]
2777unsafe fn dot_q4b_row_1x4_sx_avx2(
2778    buf: &[u8],
2779    scales: &[u8],
2780    g0: usize,
2781    gpr: usize,
2782    xs: [&[i8]; 4],
2783    sxs: [f32; 4],
2784) -> [f32; 4] {
2785    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
2786    unsafe {
2787        use core::arch::x86_64::*;
2788        let ones = _mm256_set1_epi16(1);
2789        let mut acc = [0f32; 4];
2790        for gi in 0..gpr {
2791            let s = f16_to_f32(u16::from_le_bytes([
2792                scales[(g0 + gi) * 2],
2793                scales[(g0 + gi) * 2 + 1],
2794            ]));
2795            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2796            let aw = _mm256_abs_epi8(w);
2797            for (k, xq) in xs.iter().enumerate() {
2798                let x =
2799                    _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2800                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2801                let d = _mm256_madd_epi16(p16, ones);
2802                let hi128 = _mm256_extracti128_si256::<1>(d);
2803                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2804                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2805                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2806                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
2807            }
2808        }
2809        acc
2810    }
2811}
2812
2813fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2814    #[cfg(target_arch = "aarch64")]
2815    unsafe {
2816        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
2817    }
2818    #[cfg(target_arch = "x86_64")]
2819    unsafe {
2820        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
2821    }
2822    let mut acc = 0f32;
2823    for gi in 0..gpr {
2824        let g = g0 + gi;
2825        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2826        let mut d = 0i32;
2827        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2828            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2829                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2830        }
2831        acc += d as f32 * s;
2832    }
2833    acc
2834}
2835
2836/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
2837#[inline]
2838#[allow(unreachable_code)]
2839fn dot_q4_row_i8_2(
2840    packed: &[u8],
2841    scales: &[u8],
2842    g0: usize,
2843    gpr: usize,
2844    xq1: &[i8],
2845    xq2: &[i8],
2846) -> (f32, f32) {
2847    #[cfg(target_arch = "aarch64")]
2848    unsafe {
2849        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
2850    }
2851    #[cfg(target_arch = "x86_64")]
2852    unsafe {
2853        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
2854    }
2855    (
2856        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
2857        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
2858    )
2859}
2860
2861/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
2862/// multi-matrix jobs can drive it for several tensors in one dispatch).
2863#[allow(clippy::too_many_arguments)]
2864fn q4_range_a8w8(
2865    packed: &[u8],
2866    scales: &[u8],
2867    gpr: usize,
2868    cols: usize,
2869    act: &SplitAct,
2870    out: SendMut,
2871    start: usize,
2872    end: usize,
2873) {
2874    for r in start..end {
2875        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
2876        // xq is zeroed at outlier slots — add the exact terms.
2877        for &(j, xv) in &act.outliers {
2878            let flat = r * cols + j;
2879            let byte = packed[flat / 2];
2880            let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2881            let s = f16_to_f32(u16::from_le_bytes([
2882                scales[(flat / GROUP_SIZE) * 2],
2883                scales[(flat / GROUP_SIZE) * 2 + 1],
2884            ]));
2885            acc += ((nib as i32 - 8) as f32) * s * xv;
2886        }
2887        // SAFETY: disjoint row ranges per worker.
2888        unsafe { *out.at(r) = acc };
2889    }
2890}
2891
2892/// Two-input q4 row range via the A8W8 int8 path — kernel body of
2893/// `q4matvec2`, extracted for pair multi-matrix jobs.
2894#[allow(clippy::too_many_arguments)]
2895fn q4_range2_a8w8(
2896    packed: &[u8],
2897    scales: &[u8],
2898    gpr: usize,
2899    cols: usize,
2900    a1: &SplitAct,
2901    a2: &SplitAct,
2902    p1: SendMut,
2903    p2: SendMut,
2904    start: usize,
2905    end: usize,
2906) {
2907    for r in start..end {
2908        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
2909        let mut acc1 = s1 * a1.sx;
2910        let mut acc2 = s2 * a2.sx;
2911        // xq is zeroed at outlier slots — add the exact terms.
2912        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
2913            for &(j, xv) in outliers {
2914                let flat = r * cols + j;
2915                let byte = packed[flat / 2];
2916                let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2917                let s = f16_to_f32(u16::from_le_bytes([
2918                    scales[(flat / GROUP_SIZE) * 2],
2919                    scales[(flat / GROUP_SIZE) * 2 + 1],
2920                ]));
2921                *acc += ((nib as i32 - 8) as f32) * s * xv;
2922            }
2923        };
2924        fix(&a1.outliers, &mut acc1);
2925        fix(&a2.outliers, &mut acc2);
2926        // SAFETY: disjoint row ranges per worker.
2927        unsafe {
2928            *p1.at(r) = acc1;
2929            *p2.at(r) = acc2;
2930        }
2931    }
2932}
2933
2934/// Exact scalar q4 row range (same extraction, non-SDOT path).
2935fn q4_range_f32(
2936    packed: &[u8],
2937    scales: &[u8],
2938    gpr: usize,
2939    x: &[f32],
2940    out: SendMut,
2941    start: usize,
2942    end: usize,
2943) {
2944    for r in start..end {
2945        let mut acc = 0f32;
2946        for gi in 0..gpr {
2947            let g = r * gpr + gi;
2948            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2949            let pk = &packed[g * 16..(g + 1) * 16];
2950            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2951            let mut ga = 0f32;
2952            for (k, &b) in pk.iter().enumerate() {
2953                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2954                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2955            }
2956            acc += ga * s;
2957        }
2958        // SAFETY: disjoint row ranges per worker.
2959        unsafe { *out.at(r) = acc };
2960    }
2961}
2962
2963/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
2964/// dotted against both activations (was: two full matvecs — double
2965/// weight traffic). Per-lane math matches `q4matvec` exactly.
2966#[allow(clippy::too_many_arguments)]
2967fn q4matvec2(
2968    bytes: &[u8],
2969    x1: &[f32],
2970    x2: &[f32],
2971    rows: usize,
2972    cols: usize,
2973    o1: &mut [f32],
2974    o2: &mut [f32],
2975    pool: Option<&Pool>,
2976) {
2977    debug_assert_eq!(o1.len(), rows);
2978    debug_assert_eq!(o2.len(), rows);
2979    let (packed, scales) = q4_split(bytes, rows, cols);
2980    let gpr = cols / GROUP_SIZE;
2981
2982    if a8w8_enabled() {
2983        let a1 = split_act(x1);
2984        let a2 = split_act(x2);
2985        let p1 = SendMut(o1.as_mut_ptr());
2986        let p2 = SendMut(o2.as_mut_ptr());
2987        let run = move |start: usize, end: usize| {
2988            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
2989        };
2990        dispatch_rows(pool, rows, &run);
2991        return;
2992    }
2993
2994    let p1 = SendMut(o1.as_mut_ptr());
2995    let p2 = SendMut(o2.as_mut_ptr());
2996    let run = move |start: usize, end: usize| {
2997        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
2998    };
2999    dispatch_rows(pool, rows, &run);
3000}
3001
3002/// Two-input exact scalar q4 row range (same extraction).
3003#[allow(clippy::too_many_arguments)]
3004fn q4_range2_f32(
3005    packed: &[u8],
3006    scales: &[u8],
3007    gpr: usize,
3008    x1: &[f32],
3009    x2: &[f32],
3010    p1: SendMut,
3011    p2: SendMut,
3012    start: usize,
3013    end: usize,
3014) {
3015    for r in start..end {
3016        let (mut acc1, mut acc2) = (0f32, 0f32);
3017        for gi in 0..gpr {
3018            let g = r * gpr + gi;
3019            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3020            let pk = &packed[g * 16..(g + 1) * 16];
3021            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3022            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3023            let (mut g1, mut g2) = (0f32, 0f32);
3024            for (k, &b) in pk.iter().enumerate() {
3025                let wl = (b & 0x0F) as f32 - 8.0;
3026                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
3027                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
3028                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
3029            }
3030            acc1 += g1 * s;
3031            acc2 += g2 * s;
3032        }
3033        // SAFETY: disjoint row ranges per worker.
3034        unsafe {
3035            *p1.at(r) = acc1;
3036            *p2.at(r) = acc2;
3037        }
3038    }
3039}
3040
3041thread_local! {
3042    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
3043    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
3044    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
3045    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3046}
3047
3048/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
3049/// and dotted against ALL b activations (prefill used to fall back to b
3050/// full matvecs — b× weight traffic and b× nibble decode). Per-position
3051/// math matches `q4matvec` exactly: same group order, same accumulation.
3052/// `out` is row-major [b, rows] like `qmatmat`.
3053#[allow(clippy::too_many_arguments)]
3054fn q4matmat(
3055    bytes: &[u8],
3056    xs_all: &[f32],
3057    b: usize,
3058    rows: usize,
3059    cols: usize,
3060    out: &mut [f32],
3061    pool: Option<&Pool>,
3062) {
3063    debug_assert_eq!(xs_all.len(), b * cols);
3064    debug_assert_eq!(out.len(), b * rows);
3065    let (packed, scales) = q4_split(bytes, rows, cols);
3066    let gpr = cols / GROUP_SIZE;
3067    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3068
3069    if a8w8_enabled() {
3070        let acts: Vec<SplitAct> =
3071            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3072        let acts = &acts;
3073        let out_addr = SendMut(out.as_mut_ptr());
3074        let run = move |start: usize, end: usize| {
3075            ROW_I8.with(|rb| {
3076                let mut buf = rb.borrow_mut();
3077                buf.resize(cols, 0);
3078                for r in start..end {
3079                    // Unpack the row's nibbles to centered i8 once
3080                    // (element 2k = low nibble, 2k+1 = high — flat order,
3081                    // same as dot_q4_row_sdot's zip).
3082                    for gi in 0..gpr {
3083                        let g = r * gpr + gi;
3084                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3085                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
3086                            buf[gi * GROUP_SIZE + k * 2 + 1] =
3087                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
3088                        }
3089                    }
3090                    let mut bi = 0usize;
3091                    #[cfg(target_arch = "x86_64")]
3092                    if avx2_enabled()
3093                        && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3094                    {
3095                        while bi + 4 <= acts.len() {
3096                            let xs = [
3097                                acts[bi].xq.as_slice(),
3098                                acts[bi + 1].xq.as_slice(),
3099                                acts[bi + 2].xq.as_slice(),
3100                                acts[bi + 3].xq.as_slice(),
3101                            ];
3102                            let d = unsafe {
3103                                dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
3104                            };
3105                            for k in 0..4 {
3106                                let act = &acts[bi + k];
3107                                let mut acc = d[k] * act.sx;
3108                                for &(j, xv) in &act.outliers {
3109                                    acc += (buf[j] as i8) as f32
3110                                        * gscale((r * cols + j) / GROUP_SIZE)
3111                                        * xv;
3112                                }
3113                                // SAFETY: disjoint (bi, r) cells per worker.
3114                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3115                            }
3116                            bi += 4;
3117                        }
3118                    }
3119                    while bi < acts.len() {
3120                        let act = &acts[bi];
3121                        let mut acc = 0f32;
3122                        for gi in 0..gpr {
3123                            let d = dot_i8_i8(
3124                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3125                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3126                            );
3127                            acc += d as f32 * gscale(r * gpr + gi);
3128                        }
3129                        acc *= act.sx;
3130                        // xq is zeroed at outlier slots — exact terms.
3131                        for &(j, xv) in &act.outliers {
3132                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
3133                        }
3134                        // SAFETY: disjoint (bi, r) cells per worker row range.
3135                        unsafe { *out_addr.at(bi * rows + r) = acc };
3136                        bi += 1;
3137                    }
3138                }
3139            })
3140        };
3141        dispatch_rows(pool, rows, &run);
3142        return;
3143    }
3144
3145    let out_addr = SendMut(out.as_mut_ptr());
3146    let run = move |start: usize, end: usize| {
3147        ROW_F32.with(|rb| {
3148            let mut buf = rb.borrow_mut();
3149            buf.resize(cols, 0.0);
3150            for r in start..end {
3151                // Decode raw (nib − 8) values once; scales stay per-group
3152                // so the accumulation order matches q4matvec bit-for-bit.
3153                for gi in 0..gpr {
3154                    let g = r * gpr + gi;
3155                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3156                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
3157                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
3158                    }
3159                }
3160                for bi in 0..b {
3161                    let x = &xs_all[bi * cols..(bi + 1) * cols];
3162                    let mut acc = 0f32;
3163                    for gi in 0..gpr {
3164                        let mut ga = 0f32;
3165                        // Pairwise (lo + hi) addition, matching
3166                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
3167                        // a flat one-per-element loop rounds differently
3168                        // and broke bit-parity on the scalar (x86) path.
3169                        for k in 0..GROUP_SIZE / 2 {
3170                            let e = gi * GROUP_SIZE + k * 2;
3171                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
3172                        }
3173                        acc += ga * gscale(r * gpr + gi);
3174                    }
3175                    // SAFETY: disjoint (bi, r) cells per worker row range.
3176                    unsafe { *out_addr.at(bi * rows + r) = acc };
3177                }
3178            }
3179        })
3180    };
3181    dispatch_rows(pool, rows, &run);
3182}
3183
3184/// Batched vbit matmat: each variable-bit row is decoded from the mmap
3185/// ONCE for the whole microbatch. Same per-position math as
3186/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
3187/// and the scalar path).
3188#[allow(clippy::too_many_arguments)]
3189fn vbitmatmat(
3190    bytes: &[u8],
3191    offsets: &[usize],
3192    xs_all: &[f32],
3193    b: usize,
3194    rows: usize,
3195    cols: usize,
3196    out: &mut [f32],
3197    pool: Option<&Pool>,
3198) {
3199    debug_assert_eq!(xs_all.len(), b * cols);
3200    debug_assert_eq!(out.len(), b * rows);
3201    debug_assert_eq!(offsets.len(), rows + 1);
3202    let ng = cols / GROUP_SIZE;
3203    let bits = &bytes[..rows];
3204    let sc_off = rows;
3205    let gscale = |r: usize, g: usize| {
3206        let so = (r * ng + g) * 2;
3207        f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
3208    };
3209
3210    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
3211    let decode_f32 = |r: usize, dst: &mut [f32]| {
3212        let bw = bits[r] as usize;
3213        let l = ((1i32 << (bw - 1)) - 1) as f32;
3214        let data = &bytes[offsets[r]..offsets[r + 1]];
3215        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3216        for d in dst.iter_mut() {
3217            while nbits < bw {
3218                acc = (acc << 8) | data[idx] as u64;
3219                idx += 1;
3220                nbits += 8;
3221            }
3222            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
3223            nbits -= bw;
3224            *d = u - l;
3225        }
3226    };
3227
3228    if a8w8_enabled() {
3229        let acts: Vec<SplitAct> =
3230            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3231        let acts = &acts;
3232        let out_addr = SendMut(out.as_mut_ptr());
3233        let run = move |start: usize, end: usize| {
3234            for r in start..end {
3235                let bw = bits[r] as usize;
3236                if bw == 8 {
3237                    // u−L reaches 128 → no i8 path; decode once, exact
3238                    // f32 dots for every position (same as vbitmatvec).
3239                    ROW_F32.with(|rb| {
3240                        let mut buf = rb.borrow_mut();
3241                        buf.resize(cols, 0.0);
3242                        decode_f32(r, &mut buf);
3243                        for bi in 0..b {
3244                            let x = &xs_all[bi * cols..(bi + 1) * cols];
3245                            let mut dot = 0f32;
3246                            for g in 0..ng {
3247                                let mut gd = 0f32;
3248                                for k in 0..GROUP_SIZE {
3249                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3250                                }
3251                                dot += gd * gscale(r, g);
3252                            }
3253                            // SAFETY: disjoint (bi, r) cells per worker range.
3254                            unsafe { *out_addr.at(bi * rows + r) = dot };
3255                        }
3256                    });
3257                    continue;
3258                }
3259                let l = (1i32 << (bw - 1)) - 1;
3260                let data = &bytes[offsets[r]..offsets[r + 1]];
3261                ROW_I8.with(|rb| {
3262                    let mut buf = rb.borrow_mut();
3263                    buf.resize(cols, 0);
3264                    #[inline(always)]
3265                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3266                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3267                            let u = unpack8::<B>(&data[blk * B..]);
3268                            for k in 0..8 {
3269                                chunk[k] = (u[k] - l) as i8 as u8;
3270                            }
3271                        }
3272                    }
3273                    match bw {
3274                        3 => fill::<3>(data, l, &mut buf),
3275                        4 => vbit_fill4(data, &mut buf),
3276                        5 => fill::<5>(data, l, &mut buf),
3277                        6 => fill::<6>(data, l, &mut buf),
3278                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3279                    }
3280                    let mut bi = 0usize;
3281                    // The vbit scale table shares q4_block's layout
3282                    // (contiguous f16 per (row·ng + g)), so the same
3283                    // blocked 1×4 kernel serves the decoded row.
3284                    #[cfg(target_arch = "x86_64")]
3285                    if avx2_enabled()
3286                        && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3287                    {
3288                        while bi + 4 <= acts.len() {
3289                            let xs = [
3290                                acts[bi].xq.as_slice(),
3291                                acts[bi + 1].xq.as_slice(),
3292                                acts[bi + 2].xq.as_slice(),
3293                                acts[bi + 3].xq.as_slice(),
3294                            ];
3295                            let sxs = [
3296                                acts[bi].sx,
3297                                acts[bi + 1].sx,
3298                                acts[bi + 2].sx,
3299                                acts[bi + 3].sx,
3300                            ];
3301                            let d = unsafe {
3302                                dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
3303                            };
3304                            for k in 0..4 {
3305                                let act = &acts[bi + k];
3306                                let mut dot = d[k];
3307                                for &(j, xv) in &act.outliers {
3308                                    dot +=
3309                                        (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3310                                }
3311                                // SAFETY: disjoint (bi, r) cells per worker.
3312                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
3313                            }
3314                            bi += 4;
3315                        }
3316                    }
3317                    while bi < acts.len() {
3318                        let act = &acts[bi];
3319                        let mut dot = 0f32;
3320                        for g in 0..ng {
3321                            let d = dot_i8_i8(
3322                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3323                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3324                            ) as f32
3325                                * act.sx;
3326                            dot += d * gscale(r, g);
3327                        }
3328                        for &(j, xv) in &act.outliers {
3329                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3330                        }
3331                        // SAFETY: disjoint (bi, r) cells per worker range.
3332                        unsafe { *out_addr.at(bi * rows + r) = dot };
3333                        bi += 1;
3334                    }
3335                });
3336            }
3337        };
3338        dispatch_rows(pool, rows, &run);
3339        return;
3340    }
3341
3342    let out_addr = SendMut(out.as_mut_ptr());
3343    let run = move |start: usize, end: usize| {
3344        ROW_F32.with(|rb| {
3345            let mut buf = rb.borrow_mut();
3346            buf.resize(cols, 0.0);
3347            for r in start..end {
3348                decode_f32(r, &mut buf);
3349                for bi in 0..b {
3350                    let x = &xs_all[bi * cols..(bi + 1) * cols];
3351                    let mut dot = 0f32;
3352                    for g in 0..ng {
3353                        let mut gd = 0f32;
3354                        for k in 0..GROUP_SIZE {
3355                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3356                        }
3357                        dot += gd * gscale(r, g);
3358                    }
3359                    // SAFETY: disjoint (bi, r) cells per worker range.
3360                    unsafe { *out_addr.at(bi * rows + r) = dot };
3361                }
3362            }
3363        })
3364    };
3365    dispatch_rows(pool, rows, &run);
3366}
3367
3368/// Build a GPU batch job for a q8-family mapped tensor (primary
3369/// shard): prescaled input + directory coordinates. None → not
3370/// GPU-eligible, caller stays on the CPU.
3371pub(crate) fn gpu_batch_job<'a>(
3372    t: &'a QTensor,
3373    x: &[f32],
3374) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3375    match t {
3376        QTensor::Mapped {
3377            model,
3378            idx,
3379            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3380            rows,
3381            cols,
3382            row_scale,
3383            col_field,
3384            ..
3385        } => Some((
3386            model.clone(),
3387            crate::gpu::BatchJob {
3388                idx: *idx,
3389                rows: *rows,
3390                cols: *cols,
3391                row_scale,
3392                xs: prescale(x, col_field, *dt).into_owned(),
3393                q1: false,
3394            },
3395        )),
3396        // q1: raw f32 activations, tile-embedded scales.
3397        QTensor::Mapped {
3398            model,
3399            idx,
3400            dtype: TensorDtype::Q1,
3401            rows,
3402            cols,
3403            ..
3404        } => Some((
3405            model.clone(),
3406            crate::gpu::BatchJob {
3407                idx: *idx,
3408                rows: *rows,
3409                cols: *cols,
3410                row_scale: &[],
3411                xs: x.to_vec(),
3412                q1: true,
3413            },
3414        )),
3415        _ => None,
3416    }
3417}
3418
3419/// θ col-field fold for q8_2f activations. Borrowed pass-through for
3420/// every other dtype — the old unconditional `x.to_vec()` was a pure
3421/// per-matvec allocation on the q8_row hot path.
3422pub(crate) fn prescale<'a>(
3423    x: &'a [f32],
3424    col_field: &[f32],
3425    dtype: TensorDtype,
3426) -> std::borrow::Cow<'a, [f32]> {
3427    if dtype == TensorDtype::Q8_2f {
3428        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3429    } else {
3430        std::borrow::Cow::Borrowed(x)
3431    }
3432}
3433
3434// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
3435
3436/// AVX2+FMA available? Default ON when the CPU supports both;
3437/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
3438#[cfg(target_arch = "x86_64")]
3439pub(crate) fn avx2_enabled() -> bool {
3440    use std::sync::OnceLock;
3441    static ON: OnceLock<bool> = OnceLock::new();
3442    *ON.get_or_init(|| {
3443        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3444            && std::arch::is_x86_feature_detected!("avx2")
3445            && std::arch::is_x86_feature_detected!("fma")
3446    })
3447}
3448
3449/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
3450/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
3451/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
3452/// active either way, they are exact (regrouped sums only).
3453#[cfg(target_arch = "x86_64")]
3454fn avx2_a8w8_enabled() -> bool {
3455    use std::sync::OnceLock;
3456    static ON: OnceLock<bool> = OnceLock::new();
3457    *ON.get_or_init(|| {
3458        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3459    })
3460}
3461
3462/// A8W8 quantized-activation path available on THIS machine? One
3463/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
3464/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
3465#[inline]
3466pub(crate) fn a8w8_enabled() -> bool {
3467    #[cfg(target_arch = "aarch64")]
3468    {
3469        sdot_enabled()
3470    }
3471    #[cfg(target_arch = "x86_64")]
3472    {
3473        avx2_a8w8_enabled()
3474    }
3475    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3476    {
3477        false
3478    }
3479}
3480
3481/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
3482/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
3483#[inline]
3484#[allow(unreachable_code)]
3485fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
3486    #[cfg(target_arch = "aarch64")]
3487    unsafe {
3488        return dot_i8_sdot(w, xq);
3489    }
3490    #[cfg(target_arch = "x86_64")]
3491    unsafe {
3492        if avx512vnni_enabled() {
3493            return dot_i8_i8_vnni(w, xq);
3494        }
3495        return dot_i8_i8_avx2(w, xq);
3496    }
3497    w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
3498}
3499
3500/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
3501/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
3502/// `vpdpbusd` encoding.
3503#[cfg(target_arch = "x86_64")]
3504fn avx512vnni_enabled() -> bool {
3505    use std::sync::OnceLock;
3506    static ON: OnceLock<bool> = OnceLock::new();
3507    *ON.get_or_init(|| {
3508        std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
3509            && std::arch::is_x86_feature_detected!("avx512f")
3510            && std::arch::is_x86_feature_detected!("avx512bw")
3511            && std::arch::is_x86_feature_detected!("avx512vl")
3512            && std::arch::is_x86_feature_detected!("avx512vnni")
3513    })
3514}
3515
3516/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
3517/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
3518/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
3519/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
3520#[cfg(target_arch = "x86_64")]
3521#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3522unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3523    // SAFETY: callers uphold slice-length contracts (see call sites).
3524    unsafe {
3525        use core::arch::x86_64::*;
3526        let n = w.len();
3527        let mut j = 0usize;
3528        let mut total: i32;
3529        // 4 independent accumulators: vpdpbusd is its own loop-carried
3530        // dependency (~5-cycle latency) — a single-acc loop runs
3531        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
3532        // on Granite Rapids.
3533        {
3534            #[inline(always)]
3535            unsafe fn step(
3536                w: *const u8,
3537                x: *const i8,
3538                acc: core::arch::x86_64::__m512i,
3539            ) -> core::arch::x86_64::__m512i {
3540                unsafe {
3541                    use core::arch::x86_64::*;
3542                    let wv = _mm512_loadu_si512(w as *const _);
3543                    let xv = _mm512_loadu_si512(x as *const _);
3544                    let aw = _mm512_abs_epi8(wv);
3545                    let neg = _mm512_movepi8_mask(wv);
3546                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3547                    _mm512_dpbusd_epi32(acc, aw, sx)
3548                }
3549            }
3550            let (mut a0, mut a1, mut a2, mut a3) = (
3551                _mm512_setzero_si512(),
3552                _mm512_setzero_si512(),
3553                _mm512_setzero_si512(),
3554                _mm512_setzero_si512(),
3555            );
3556            while j + 256 <= n {
3557                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3558                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3559                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3560                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3561                j += 256;
3562            }
3563            while j + 64 <= n {
3564                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3565                j += 64;
3566            }
3567            let s01 = _mm512_add_epi32(a0, a1);
3568            let s23 = _mm512_add_epi32(a2, a3);
3569            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3570        }
3571        // 32-wide (q4/vbit groups are exactly 32 bytes).
3572        if j + 32 <= n {
3573            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3574            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3575            let d = _mm256_dpbusd_epi32(
3576                _mm256_setzero_si256(),
3577                _mm256_abs_epi8(wv),
3578                _mm256_sign_epi8(xv, wv),
3579            );
3580            let hi128 = _mm256_extracti128_si256::<1>(d);
3581            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3582            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3583            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3584            total += _mm_cvtsi128_si32(s32);
3585            j += 32;
3586        }
3587        while j < n {
3588            total += (w[j] as i8) as i32 * xq[j] as i32;
3589            j += 1;
3590        }
3591        total
3592    }
3593}
3594
3595/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
3596#[cfg(target_arch = "x86_64")]
3597#[target_feature(enable = "avx2,fma")]
3598unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3599    // SAFETY: callers uphold slice-length contracts (see call sites).
3600    unsafe {
3601        use core::arch::x86_64::*;
3602        let n = x.len();
3603        let wp = w.as_ptr();
3604        let xp = x.as_ptr();
3605        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3606        let mut j = 0usize;
3607        while j + 16 <= n {
3608            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3609            let lo = _mm256_cvtepi8_epi32(wb);
3610            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3611            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3612            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3613            j += 16;
3614        }
3615        let acc = _mm256_add_ps(a0, a1);
3616        let hi128 = _mm256_extractf128_ps::<1>(acc);
3617        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3618        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3619        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3620        let mut sum = _mm_cvtss_f32(s32);
3621        while j < n {
3622            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3623            j += 1;
3624        }
3625        sum
3626    }
3627}
3628
3629/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
3630/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
3631/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
3632/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
3633#[cfg(target_arch = "x86_64")]
3634#[target_feature(enable = "avx2")]
3635unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3636    // SAFETY: callers uphold slice-length contracts (see call sites).
3637    unsafe {
3638        use core::arch::x86_64::*;
3639        let n = w.len();
3640        let ones = _mm256_set1_epi16(1);
3641        let mut acc = _mm256_setzero_si256();
3642        let mut j = 0usize;
3643        while j + 32 <= n {
3644            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3645            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3646            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3647            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3648            j += 32;
3649        }
3650        let hi128 = _mm256_extracti128_si256::<1>(acc);
3651        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3652        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3653        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3654        let mut s = _mm_cvtsi128_si32(s32);
3655        while j < n {
3656            s += (w[j] as i8) as i32 * xq[j] as i32;
3657            j += 1;
3658        }
3659        s
3660    }
3661}
3662
3663/// Blocked 2 weight rows × 4 activations for the prefill GEMM
3664/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
3665/// abs() live in registers across all four activation streams; the
3666/// sign-fixup is recomputed per pair (the price of the maddubs trick).
3667/// Returns raw i8·i8 dots; the caller applies scales and outliers.
3668#[cfg(target_arch = "x86_64")]
3669#[target_feature(enable = "avx2")]
3670unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
3671    // SAFETY: callers uphold slice-length contracts.
3672    unsafe {
3673        use core::arch::x86_64::*;
3674        let n = w0.len();
3675        let ones = _mm256_set1_epi16(1);
3676        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
3677        let mut j = 0usize;
3678        while j + 32 <= n {
3679            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
3680            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
3681            let aw0 = _mm256_abs_epi8(wv0);
3682            let aw1 = _mm256_abs_epi8(wv1);
3683            for (k, x) in xs.iter().enumerate() {
3684                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
3685                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
3686                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
3687                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
3688                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
3689            }
3690            j += 32;
3691        }
3692        let mut out = [[0i32; 4]; 2];
3693        for r in 0..2 {
3694            for k in 0..4 {
3695                let a = acc[r][k];
3696                let hi128 = _mm256_extracti128_si256::<1>(a);
3697                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
3698                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3699                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3700                out[r][k] = _mm_cvtsi128_si32(s32);
3701            }
3702        }
3703        if j < n {
3704            for (k, x) in xs.iter().enumerate() {
3705                for i in j..n {
3706                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
3707                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
3708                }
3709            }
3710        }
3711        out
3712    }
3713}
3714
3715/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
3716/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
3717/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
3718/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
3719#[cfg(target_arch = "x86_64")]
3720#[inline]
3721fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
3722    let dot = if avx512vnni_enabled() && row.len() >= 64 {
3723        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
3724    } else {
3725        unsafe { dot_i8_i8_avx2(row, &act.xq) }
3726    };
3727    let mut acc = dot as f32 * act.sx;
3728    for &(j, xv) in &act.outliers {
3729        acc += (row[j] as i8) as f32 * xv;
3730    }
3731    acc
3732}
3733
3734/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
3735/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
3736/// a single-acc loop runs latency-bound, measured on Granite Rapids).
3737#[cfg(target_arch = "x86_64")]
3738#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3739unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3740    // SAFETY: callers uphold slice-length contracts (see call sites).
3741    unsafe {
3742        use core::arch::x86_64::*;
3743        let n = w.len();
3744        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
3745        #[inline(always)]
3746        unsafe fn step(
3747            w: *const u8,
3748            x: *const i8,
3749            flip: core::arch::x86_64::__m512i,
3750            acc: core::arch::x86_64::__m512i,
3751        ) -> core::arch::x86_64::__m512i {
3752            unsafe {
3753                use core::arch::x86_64::*;
3754                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
3755                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
3756            }
3757        }
3758        let (mut a0, mut a1, mut a2, mut a3) = (
3759            _mm512_setzero_si512(),
3760            _mm512_setzero_si512(),
3761            _mm512_setzero_si512(),
3762            _mm512_setzero_si512(),
3763        );
3764        let mut j = 0usize;
3765        while j + 256 <= n {
3766            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3767            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
3768            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
3769            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
3770            j += 256;
3771        }
3772        while j + 64 <= n {
3773            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3774            j += 64;
3775        }
3776        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
3777            _mm512_add_epi32(a0, a1),
3778            _mm512_add_epi32(a2, a3),
3779        ));
3780        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
3781        while j < n {
3782            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
3783            j += 1;
3784        }
3785        total
3786    }
3787}
3788
3789/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
3790/// writer's flat order, same as the NEON vzip pair), maddubs against
3791/// the pre-quantized activation group, × the group's f16 scale. Pair
3792/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
3793/// `dot_q4_row_sdot`.
3794#[cfg(target_arch = "x86_64")]
3795#[target_feature(enable = "avx2")]
3796unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3797    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
3798    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
3799    unsafe {
3800        use core::arch::x86_64::*;
3801        let lomask = _mm_set1_epi8(0x0F);
3802        let eight = _mm256_set1_epi8(8);
3803        let ones = _mm256_set1_epi16(1);
3804        let mut acc = 0f32;
3805        for gi in 0..gpr {
3806            let g = g0 + gi;
3807            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3808            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3809            let lo = _mm_and_si128(b, lomask);
3810            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3811            let w = _mm256_sub_epi8(
3812                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3813                eight,
3814            );
3815            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3816            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3817            let d = _mm256_madd_epi16(p16, ones);
3818            let hi128 = _mm256_extracti128_si256::<1>(d);
3819            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3820            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3821            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3822            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3823        }
3824        acc
3825    }
3826}
3827
3828/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
3829/// both activations dotted against the same centered i8 register.
3830#[cfg(target_arch = "x86_64")]
3831#[target_feature(enable = "avx2")]
3832unsafe fn dot_q4_row_avx2_2(
3833    packed: &[u8],
3834    scales: &[u8],
3835    g0: usize,
3836    gpr: usize,
3837    xq1: &[i8],
3838    xq2: &[i8],
3839) -> (f32, f32) {
3840    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
3841    unsafe {
3842        use core::arch::x86_64::*;
3843        let lomask = _mm_set1_epi8(0x0F);
3844        let eight = _mm256_set1_epi8(8);
3845        let ones = _mm256_set1_epi16(1);
3846        let (mut acc1, mut acc2) = (0f32, 0f32);
3847        #[inline(always)]
3848        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
3849            unsafe {
3850                use core::arch::x86_64::*;
3851                let hi128 = _mm256_extracti128_si256::<1>(d);
3852                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3853                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3854                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3855                _mm_cvtsi128_si32(s32)
3856            }
3857        }
3858        for gi in 0..gpr {
3859            let g = g0 + gi;
3860            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3861            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3862            let lo = _mm_and_si128(b, lomask);
3863            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3864            let w = _mm256_sub_epi8(
3865                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3866                eight,
3867            );
3868            let aw = _mm256_abs_epi8(w);
3869            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3870            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3871            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
3872            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
3873            acc1 += hsum(d1) as f32 * s;
3874            acc2 += hsum(d2) as f32 * s;
3875        }
3876        (acc1, acc2)
3877    }
3878}
3879
3880/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
3881#[cfg(target_arch = "x86_64")]
3882fn q8_range_avx2(
3883    q: &[u8],
3884    row_scale: &[f32],
3885    act: &SplitAct,
3886    cols: usize,
3887    out_addr: SendMut,
3888    start: usize,
3889    end: usize,
3890) {
3891    for o in start..end {
3892        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3893        // SAFETY: disjoint row ranges per worker.
3894        unsafe { *out_addr.at(o) = v };
3895    }
3896}
3897
3898/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
3899#[cfg(target_arch = "x86_64")]
3900#[allow(clippy::too_many_arguments)]
3901fn q8_range2_avx2(
3902    q: &[u8],
3903    row_scale: &[f32],
3904    a1: &SplitAct,
3905    a2: &SplitAct,
3906    cols: usize,
3907    p1: SendMut,
3908    p2: SendMut,
3909    start: usize,
3910    end: usize,
3911) {
3912    for o in start..end {
3913        let row = &q[o * cols..(o + 1) * cols];
3914        // SAFETY: disjoint row ranges per worker.
3915        unsafe {
3916            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
3917            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
3918        }
3919    }
3920}
3921
3922// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
3923
3924/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
3925/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
3926/// (On non-ARM release builds only the test tolerance switch calls it.)
3927#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3928fn sdot_enabled() -> bool {
3929    use std::sync::OnceLock;
3930    static ON: OnceLock<bool> = OnceLock::new();
3931    *ON.get_or_init(|| {
3932        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
3933        #[cfg(target_arch = "aarch64")]
3934        {
3935            want && std::arch::is_aarch64_feature_detected!("dotprod")
3936        }
3937        #[cfg(not(target_arch = "aarch64"))]
3938        {
3939            let _ = want;
3940            false
3941        }
3942    })
3943}
3944
3945/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
3946/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
3947/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
3948/// matvec, shared by all rows/workers.
3949struct SplitAct {
3950    xq: Vec<i8>,
3951    sx: f32,
3952    outliers: Vec<(usize, f32)>,
3953    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
3954    /// `−128·Σx`); one i32 per split, computed once per matvec.
3955    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
3956    xsum: i32,
3957}
3958
3959thread_local! {
3960    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
3961    /// and its hidden-size allocation was steady-state heap churn.
3962    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
3963        const { std::cell::RefCell::new(Vec::new()) };
3964}
3965
3966impl Drop for SplitAct {
3967    fn drop(&mut self) {
3968        let buf = std::mem::take(&mut self.xq);
3969        if buf.capacity() > 0 {
3970            XQ_FREE.with(|f| {
3971                let mut f = f.borrow_mut();
3972                if f.len() < 16 {
3973                    f.push(buf);
3974                }
3975            });
3976        }
3977    }
3978}
3979
3980fn split_act(x: &[f32]) -> SplitAct {
3981    let n = x.len();
3982    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
3983    let thr = 8.0 * rms;
3984    // One pass: collect outliers and the bulk absmax (outliers excluded —
3985    // identical to the old zero-then-fold over a copied buffer, minus the
3986    // full-vector copy).
3987    let mut outliers: Vec<(usize, f32)> = Vec::new();
3988    let mut amax = 0f32;
3989    for (j, &v) in x.iter().enumerate() {
3990        let a = v.abs();
3991        if a > thr {
3992            outliers.push((j, v));
3993        } else if a > amax {
3994            amax = a;
3995        }
3996    }
3997    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
3998    let inv = 1.0 / sx;
3999    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
4000    xq.clear();
4001    xq.reserve(n);
4002    if outliers.is_empty() {
4003        xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
4004    } else {
4005        // Outlier slots quantize to 0 (their exact term is added later).
4006        xq.extend(x.iter().map(|&v| {
4007            if v.abs() > thr {
4008                0
4009            } else {
4010                (v * inv).round().clamp(-127.0, 127.0) as i8
4011            }
4012        }));
4013    }
4014    let xsum = xq.iter().map(|&v| v as i32).sum();
4015    SplitAct { xq, sx, outliers, xsum }
4016}
4017
4018/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
4019/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
4020#[cfg(target_arch = "aarch64")]
4021#[target_feature(enable = "neon,dotprod")]
4022unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
4023    // SAFETY: callers uphold slice-length contracts (see call sites).
4024    unsafe {
4025        use core::arch::aarch64::*;
4026        use core::arch::asm;
4027        let wp = w.as_ptr() as *const i8;
4028        let n = w.len();
4029        let (mut a0, mut a1, mut a2, mut a3) =
4030            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4031        let mut i = 0;
4032        while i + 64 <= n {
4033            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4034            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
4035            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
4036            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
4037            asm!(
4038                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4039                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4040                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
4041                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
4042                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4043                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4044                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
4045                options(pure, nomem, nostack),
4046            );
4047            i += 64;
4048        }
4049        while i + 16 <= n {
4050            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4051            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
4052                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
4053            i += 16;
4054        }
4055        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
4056        while i < n {
4057            s += (*wp.add(i)) as i32 * xq[i] as i32;
4058            i += 1;
4059        }
4060        s
4061}
4062}
4063
4064/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
4065/// loaded once and reused, 4 independent accumulators hide sdot latency
4066/// (port of vmfcore `dot_i8_sdot_4rows`).
4067#[cfg(target_arch = "aarch64")]
4068#[target_feature(enable = "neon,dotprod")]
4069unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
4070    // SAFETY: callers uphold slice-length contracts (see call sites).
4071    unsafe {
4072        use core::arch::aarch64::*;
4073        use core::arch::asm;
4074        let n = xq.len();
4075        let px = xq.as_ptr();
4076        let (p0, p1, p2, p3) = (
4077            w0.as_ptr() as *const i8,
4078            w1.as_ptr() as *const i8,
4079            w2.as_ptr() as *const i8,
4080            w3.as_ptr() as *const i8,
4081        );
4082        let (mut a0, mut a1, mut a2, mut a3) =
4083            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4084        let mut i = 0;
4085        while i + 16 <= n {
4086            let x = vld1q_s8(px.add(i));
4087            let v0 = vld1q_s8(p0.add(i));
4088            let v1 = vld1q_s8(p1.add(i));
4089            let v2 = vld1q_s8(p2.add(i));
4090            let v3 = vld1q_s8(p3.add(i));
4091            asm!(
4092                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4093                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4094                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4095                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4096                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4097                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4098                options(pure, nomem, nostack),
4099            );
4100            i += 16;
4101        }
4102        let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
4103        while i < n {
4104            let xi = *px.add(i) as i32;
4105            r[0] += (*p0.add(i)) as i32 * xi;
4106            r[1] += (*p1.add(i)) as i32 * xi;
4107            r[2] += (*p2.add(i)) as i32 * xi;
4108            r[3] += (*p3.add(i)) as i32 * xi;
4109            i += 1;
4110        }
4111        r
4112}
4113}
4114
4115/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
4116/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
4117/// line plus the shared activation chunk — a single sequential weight
4118/// stream per worker. Per-row accumulation is the same one-accumulator
4119/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
4120/// are bit-identical to the mmap-layout kernel.
4121#[cfg(target_arch = "aarch64")]
4122#[target_feature(enable = "neon,dotprod")]
4123unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
4124    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
4125    // n % 16 == 0 — guaranteed by the repack gate).
4126    unsafe {
4127        use core::arch::aarch64::*;
4128        use core::arch::asm;
4129        let n = xq.len();
4130        let px = xq.as_ptr();
4131        let pg = g.as_ptr() as *const i8;
4132        let (mut a0, mut a1, mut a2, mut a3) =
4133            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4134        let mut i = 0;
4135        while i + 16 <= n {
4136            let x = vld1q_s8(px.add(i));
4137            let base = pg.add(4 * i);
4138            let v0 = vld1q_s8(base);
4139            let v1 = vld1q_s8(base.add(16));
4140            let v2 = vld1q_s8(base.add(32));
4141            let v3 = vld1q_s8(base.add(48));
4142            asm!(
4143                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4144                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4145                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4146                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4147                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4148                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4149                options(pure, nomem, nostack),
4150            );
4151            i += 16;
4152        }
4153        [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
4154    }
4155}
4156
4157/// One q8 row range via SDOT (4-row blocks + tail) — the body of
4158/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
4159/// SAME kernel for several tensors under one pool dispatch. `rep` — the
4160/// load-time interleaved repack (empty = mmap layout only); rows outside
4161/// full 4-row groups always come from the mmap layout.
4162#[cfg(target_arch = "aarch64")]
4163fn q8_range_sdot(
4164    q: &[u8],
4165    rep: &[u8],
4166    row_scale: &[f32],
4167    act: &SplitAct,
4168    cols: usize,
4169    out_addr: SendMut,
4170    start: usize,
4171    end: usize,
4172) {
4173    let mut o = start;
4174    // Leading rows to the group boundary (repack path only): the pool
4175    // splits row ranges arbitrarily, groups are absolute.
4176    if !rep.is_empty() {
4177        while o < end && o % 4 != 0 {
4178            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4179            unsafe { *out_addr.at(o) = v };
4180            o += 1;
4181        }
4182    }
4183    while o + 4 <= end {
4184        let r = if rep.is_empty() {
4185            unsafe {
4186                dot_i8_sdot_4rows(
4187                    &q[o * cols..(o + 1) * cols],
4188                    &q[(o + 1) * cols..(o + 2) * cols],
4189                    &q[(o + 2) * cols..(o + 3) * cols],
4190                    &q[(o + 3) * cols..(o + 4) * cols],
4191                    &act.xq,
4192                )
4193            }
4194        } else {
4195            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
4196        };
4197        for k in 0..4 {
4198            let mut acc = r[k] as f32 * act.sx;
4199            for &(j, xv) in &act.outliers {
4200                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
4201            }
4202            // SAFETY: disjoint row ranges per worker.
4203            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
4204        }
4205        o += 4;
4206    }
4207    while o < end {
4208        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4209        unsafe { *out_addr.at(o) = v };
4210        o += 1;
4211    }
4212}
4213
4214/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
4215/// for the fused pair multi-matrix job (`matvec2_many`).
4216#[cfg(target_arch = "aarch64")]
4217#[allow(clippy::too_many_arguments)]
4218fn q8_range2_sdot(
4219    q: &[u8],
4220    row_scale: &[f32],
4221    a1: &SplitAct,
4222    a2: &SplitAct,
4223    cols: usize,
4224    p1: SendMut,
4225    p2: SendMut,
4226    start: usize,
4227    end: usize,
4228) {
4229    for o in start..end {
4230        let row = &q[o * cols..(o + 1) * cols];
4231        // SAFETY: disjoint row ranges per worker.
4232        unsafe {
4233            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
4234            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
4235        }
4236    }
4237}
4238
4239/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
4240#[allow(clippy::too_many_arguments)]
4241fn q8_range2_f32(
4242    q: &[u8],
4243    row_scale: &[f32],
4244    x1: &[f32],
4245    x2: &[f32],
4246    cols: usize,
4247    p1: SendMut,
4248    p2: SendMut,
4249    start: usize,
4250    end: usize,
4251) {
4252    for o in start..end {
4253        let row = &q[o * cols..(o + 1) * cols];
4254        // SAFETY: disjoint row ranges per worker.
4255        unsafe {
4256            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
4257            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
4258        }
4259    }
4260}
4261
4262/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
4263fn q8_range_f32(
4264    q: &[u8],
4265    row_scale: &[f32],
4266    xs: &[f32],
4267    cols: usize,
4268    out_addr: SendMut,
4269    start: usize,
4270    end: usize,
4271) {
4272    for o in start..end {
4273        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
4274        // SAFETY: disjoint row ranges per worker.
4275        unsafe { *out_addr.at(o) = v };
4276    }
4277}
4278
4279/// SDOT row dot with exact outlier correction:
4280/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
4281#[cfg(target_arch = "aarch64")]
4282#[inline]
4283fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
4284    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
4285    for &(j, xv) in &act.outliers {
4286        acc += (row[j] as i8) as f32 * xv;
4287    }
4288    acc
4289}
4290
4291/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
4292/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
4293/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
4294/// the caller multiplies by the activation scale and adds the exact
4295/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
4296/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
4297/// → zip(lo,hi) restores flat order.
4298#[cfg(target_arch = "aarch64")]
4299#[target_feature(enable = "neon,dotprod")]
4300unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4301    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
4302    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
4303    unsafe {
4304        use core::arch::aarch64::*;
4305        use core::arch::asm;
4306        let lomask = vdupq_n_u8(0x0F);
4307        let eight = vdupq_n_s8(8);
4308        let mut acc = 0f32;
4309        for gi in 0..gpr {
4310            let g = g0 + gi;
4311            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4312            let b = vld1q_u8(packed.as_ptr().add(g * 16));
4313            let lo = vandq_u8(b, lomask);
4314            let hi = vshrq_n_u8::<4>(b);
4315            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4316            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4317            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4318            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4319            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4320            asm!(
4321                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4322                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4323                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4324                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4325                options(pure, nomem, nostack),
4326            );
4327            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4328        }
4329        acc
4330    }
4331}
4332
4333/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
4334/// part) happens ONCE per group; both pre-quantized activations are
4335/// dotted against the same centered i8 registers. Per-lane math matches
4336/// `dot_q4_row_sdot` exactly.
4337#[cfg(target_arch = "aarch64")]
4338#[target_feature(enable = "neon,dotprod")]
4339unsafe fn dot_q4_row_sdot2(
4340    packed: &[u8],
4341    scales: &[u8],
4342    g0: usize,
4343    gpr: usize,
4344    xq1: &[i8],
4345    xq2: &[i8],
4346) -> (f32, f32) {
4347    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
4348    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
4349    unsafe {
4350        use core::arch::aarch64::*;
4351        use core::arch::asm;
4352        let lomask = vdupq_n_u8(0x0F);
4353        let eight = vdupq_n_s8(8);
4354        let (mut acc1, mut acc2) = (0f32, 0f32);
4355        for gi in 0..gpr {
4356            let g = g0 + gi;
4357            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4358            let b = vld1q_u8(packed.as_ptr().add(g * 16));
4359            let lo = vandq_u8(b, lomask);
4360            let hi = vshrq_n_u8::<4>(b);
4361            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4362            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4363            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
4364            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
4365            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
4366            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
4367            let (mut a0, mut a1, mut b0, mut b1) =
4368                (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4369            asm!(
4370                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
4371                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
4372                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
4373                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
4374                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4375                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
4376                e0 = in(vreg) e0, e1 = in(vreg) e1,
4377                x10 = in(vreg) x10, x11 = in(vreg) x11,
4378                x20 = in(vreg) x20, x21 = in(vreg) x21,
4379                options(pure, nomem, nostack),
4380            );
4381            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4382            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
4383        }
4384        (acc1, acc2)
4385    }
4386}
4387
4388// ───────────────────── fused int8 kernels ─────────────────────
4389
4390/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
4391/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
4392#[inline]
4393pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
4394    #[cfg(target_arch = "aarch64")]
4395    unsafe {
4396        return axpy_i8_f32_neon(acc, row, w);
4397    }
4398    #[cfg(target_arch = "x86_64")]
4399    if avx2_enabled() {
4400        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
4401    }
4402    #[allow(unreachable_code)]
4403    {
4404        for (a, &b) in acc.iter_mut().zip(row) {
4405            *a += w * b as f32;
4406        }
4407    }
4408}
4409
4410/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
4411#[cfg(target_arch = "x86_64")]
4412#[target_feature(enable = "avx2,fma")]
4413unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
4414    // SAFETY: callers uphold slice-length contracts (see call sites).
4415    unsafe {
4416        use core::arch::x86_64::*;
4417        let n = acc.len().min(row.len());
4418        let ap = acc.as_mut_ptr();
4419        let rp = row.as_ptr();
4420        let wv = _mm256_set1_ps(w);
4421        let mut j = 0usize;
4422        while j + 16 <= n {
4423            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
4424            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
4425            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
4426            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
4427            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
4428            _mm256_storeu_ps(ap.add(j), v0);
4429            _mm256_storeu_ps(ap.add(j + 8), v1);
4430            j += 16;
4431        }
4432        while j < n {
4433            *ap.add(j) += w * (*rp.add(j)) as f32;
4434            j += 1;
4435        }
4436    }
4437}
4438
4439#[cfg(target_arch = "aarch64")]
4440#[target_feature(enable = "neon")]
4441unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
4442    // SAFETY: callers uphold slice-length contracts (see call sites).
4443    unsafe {
4444        use core::arch::aarch64::*;
4445        let n = acc.len().min(row.len());
4446        let ap = acc.as_mut_ptr();
4447        let rp = row.as_ptr();
4448        let wv = vdupq_n_f32(w);
4449        let mut j = 0usize;
4450        while j + 16 <= n {
4451            let rb = vld1q_s8(rp.add(j));
4452            let lo = vmovl_s8(vget_low_s8(rb));
4453            let hi = vmovl_s8(vget_high_s8(rb));
4454            for (off, half) in [(0, lo), (8, hi)] {
4455                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
4456                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
4457                let o = j + off;
4458                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
4459                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
4460            }
4461            j += 16;
4462        }
4463        while j < n {
4464            *ap.add(j) += w * (*rp.add(j)) as f32;
4465            j += 1;
4466        }
4467}
4468}
4469
4470/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
4471/// ≈9× scalar), scalar elsewhere.
4472#[inline]
4473pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
4474    #[cfg(target_arch = "aarch64")]
4475    unsafe {
4476        return dot_i8_f32_neon(w, x);
4477    }
4478    #[cfg(target_arch = "x86_64")]
4479    if avx2_enabled() {
4480        return unsafe { dot_i8_f32_avx2(w, x) };
4481    }
4482    #[allow(unreachable_code)]
4483    {
4484        let mut sum = 0.0f32;
4485        for (j, &b) in w.iter().enumerate() {
4486            sum += (b as i8) as f32 * x[j];
4487        }
4488        sum
4489    }
4490}
4491
4492/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
4493/// folded into the product (no prescaled copy of x). NEON on aarch64,
4494/// scalar elsewhere. Used by the active-neuron path `row_dot`.
4495#[inline]
4496fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4497    #[cfg(target_arch = "aarch64")]
4498    unsafe {
4499        return dot_i8_col_f32_neon(w, x, col);
4500    }
4501    #[allow(unreachable_code)]
4502    {
4503        let mut sum = 0.0f32;
4504        for (j, &b) in w.iter().enumerate() {
4505            sum += (b as i8) as f32 * x[j] * col[j];
4506        }
4507        sum
4508    }
4509}
4510
4511#[cfg(target_arch = "aarch64")]
4512#[target_feature(enable = "neon")]
4513unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4514    // SAFETY: callers uphold slice-length contracts (see call sites).
4515    unsafe {
4516        use core::arch::aarch64::*;
4517        let n = x.len();
4518        let wp = w.as_ptr() as *const i8;
4519        let xp = x.as_ptr();
4520        let cp = col.as_ptr();
4521        let (mut a0, mut a1, mut a2, mut a3) =
4522            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4523        let mut j = 0usize;
4524        while j + 16 <= n {
4525            let wb = vld1q_s8(wp.add(j));
4526            let lo = vmovl_s8(vget_low_s8(wb));
4527            let hi = vmovl_s8(vget_high_s8(wb));
4528            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4529            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4530            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4531            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4532            a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
4533            a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
4534            a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
4535            a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
4536            j += 16;
4537        }
4538        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4539        while j < n {
4540            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
4541            j += 1;
4542        }
4543        sum
4544}
4545}
4546
4547#[cfg(target_arch = "aarch64")]
4548#[target_feature(enable = "neon")]
4549unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
4550    // SAFETY: callers uphold slice-length contracts (see call sites).
4551    unsafe {
4552        use core::arch::aarch64::*;
4553        let n = x.len();
4554        let wp = w.as_ptr() as *const i8;
4555        let xp = x.as_ptr();
4556        let (mut a0, mut a1, mut a2, mut a3) =
4557            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4558        let mut j = 0usize;
4559        while j + 16 <= n {
4560            let wb = vld1q_s8(wp.add(j));
4561            let lo = vmovl_s8(vget_low_s8(wb));
4562            let hi = vmovl_s8(vget_high_s8(wb));
4563            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4564            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4565            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4566            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4567            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
4568            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
4569            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
4570            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
4571            j += 16;
4572        }
4573        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4574        while j < n {
4575            sum += (*wp.add(j)) as f32 * *xp.add(j);
4576            j += 1;
4577        }
4578        sum
4579}
4580}
4581
4582#[allow(clippy::too_many_arguments)]
4583fn qmatvec(
4584    q: &[u8],
4585    rep: &[u8],
4586    row_scale: &[f32],
4587    xs: &[f32],
4588    rows: usize,
4589    cols: usize,
4590    out: &mut [f32],
4591    pool: Option<&Pool>,
4592) {
4593    debug_assert_eq!(out.len(), rows);
4594    #[cfg(not(target_arch = "aarch64"))]
4595    let _ = rep;
4596
4597    #[cfg(target_arch = "aarch64")]
4598    if sdot_enabled() {
4599        let act = split_act(xs);
4600        let out_addr = SendMut(out.as_mut_ptr());
4601        let run_range = |start: usize, end: usize| {
4602            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
4603        };
4604        match pool {
4605            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4606            _ => run_range(0, rows),
4607        }
4608        return;
4609    }
4610    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
4611    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
4612    #[cfg(target_arch = "x86_64")]
4613    if avx2_a8w8_enabled() {
4614        let act = split_act(xs);
4615        let out_addr = SendMut(out.as_mut_ptr());
4616        let run_range =
4617            |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
4618        match pool {
4619            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4620            _ => run_range(0, rows),
4621        }
4622        return;
4623    }
4624
4625    let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
4626    match pool {
4627        Some(pool) if rows >= 256 => {
4628            let out_addr = SendMut(out.as_mut_ptr());
4629            pool.run(&move |widx, n| {
4630                let chunk = rows.div_ceil(n);
4631                let start = widx * chunk;
4632                let end = (start + chunk).min(rows);
4633                for o in start..end {
4634                    // SAFETY: disjoint row ranges per worker.
4635                    unsafe { *out_addr.at(o) = row_dot(o) };
4636                }
4637            });
4638        }
4639        _ => {
4640            for (o, dst) in out.iter_mut().enumerate() {
4641                *dst = row_dot(o);
4642            }
4643        }
4644    }
4645}
4646
4647#[allow(clippy::too_many_arguments)]
4648fn qmatvec2(
4649    q: &[u8],
4650    row_scale: &[f32],
4651    x1: &[f32],
4652    x2: &[f32],
4653    rows: usize,
4654    cols: usize,
4655    o1: &mut [f32],
4656    o2: &mut [f32],
4657    pool: Option<&Pool>,
4658) {
4659    #[cfg(target_arch = "aarch64")]
4660    if sdot_enabled() {
4661        let a1s = split_act(x1);
4662        let a2s = split_act(x2);
4663        let p1 = SendMut(o1.as_mut_ptr());
4664        let p2 = SendMut(o2.as_mut_ptr());
4665        let run_range = |start: usize, end: usize| {
4666            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4667        };
4668        match pool {
4669            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4670            _ => run_range(0, rows),
4671        }
4672        return;
4673    }
4674    #[cfg(target_arch = "x86_64")]
4675    if avx2_a8w8_enabled() {
4676        let a1s = split_act(x1);
4677        let a2s = split_act(x2);
4678        let p1 = SendMut(o1.as_mut_ptr());
4679        let p2 = SendMut(o2.as_mut_ptr());
4680        let run_range = |start: usize, end: usize| {
4681            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4682        };
4683        match pool {
4684            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4685            _ => run_range(0, rows),
4686        }
4687        return;
4688    }
4689
4690    let row_dots = |o: usize| -> (f32, f32) {
4691        let row = &q[o * cols..(o + 1) * cols];
4692        (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
4693    };
4694    match pool {
4695        Some(pool) if rows >= 256 => {
4696            let p1 = SendMut(o1.as_mut_ptr());
4697            let p2 = SendMut(o2.as_mut_ptr());
4698            pool.run(&move |widx, n| {
4699                let chunk = rows.div_ceil(n);
4700                let start = widx * chunk;
4701                let end = (start + chunk).min(rows);
4702                for o in start..end {
4703                    let (s1, s2) = row_dots(o);
4704                    // SAFETY: disjoint row ranges per worker.
4705                    unsafe {
4706                        *p1.at(o) = s1;
4707                        *p2.at(o) = s2;
4708                    }
4709                }
4710            });
4711        }
4712        _ => {
4713            for o in 0..rows {
4714                let (s1, s2) = row_dots(o);
4715                o1[o] = s1;
4716                o2[o] = s2;
4717            }
4718        }
4719    }
4720}
4721
4722#[derive(Clone, Copy)]
4723struct SendMut(*mut f32);
4724unsafe impl Send for SendMut {}
4725unsafe impl Sync for SendMut {}
4726
4727impl SendMut {
4728    #[inline]
4729    fn at(self, i: usize) -> *mut f32 {
4730        unsafe { self.0.add(i) }
4731    }
4732}
4733
4734#[cfg(test)]
4735mod tests {
4736    use super::*;
4737
4738    #[test]
4739    fn f32_matvec_matches_matvec_rows_bitexact() {
4740        let (rows, cols) = (300, 40);
4741        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
4742        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
4743        let qt = QTensor::from_f32(w.clone(), rows, cols);
4744
4745        let mut a = vec![0.0f32; rows];
4746        matvec_rows(None, &w, &x, &mut a);
4747        let mut b = vec![0.0f32; rows];
4748        qt.matvec(&x, &mut b, None);
4749        assert_eq!(a, b);
4750    }
4751
4752    #[test]
4753    fn sdot_kernel_exact_on_grid() {
4754        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
4755        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
4756        // exact f32 dot to float rounding. This isolates kernel
4757        // correctness from quantization noise.
4758        eprintln!("sdot_enabled = {}", sdot_enabled());
4759        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
4760        let w: Vec<u8> = (0..rows * cols)
4761            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4762            .collect();
4763        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
4764        let x: Vec<f32> = (0..cols)
4765            .map(|i| match i % 3 {
4766                0 => 1.0,
4767                1 => -1.0,
4768                _ => 0.0,
4769            })
4770            .collect();
4771        let mut a = vec![0.0f32; rows];
4772        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4773        for o in 0..rows {
4774            let mut acc = 0.0f32;
4775            for j in 0..cols {
4776                acc += (w[o * cols + j] as i8) as f32 * x[j];
4777            }
4778            let expect = acc * scales[o];
4779            assert!(
4780                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
4781                "row {o}: {} vs {expect}",
4782                a[o]
4783            );
4784        }
4785    }
4786
4787    #[test]
4788    fn q1_kernels_match_exact_reference() {
4789        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
4790        let (rows, cols) = (7, 96);
4791        let gpr = cols / GROUP_SIZE;
4792        let mut bytes = Vec::new();
4793        for t in 0..rows * gpr {
4794            let s = 0.01 + (t % 13) as f32 * 0.003;
4795            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4796            for j in 0..4 {
4797                bytes.push(((t * 31 + j * 97) % 251) as u8);
4798            }
4799        }
4800        // On-grid activations (±1, amax 1) → the SDOT path is exact.
4801        let x: Vec<f32> = (0..cols)
4802            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
4803            .collect();
4804        // Reference through the core dequant.
4805        let mut w = vec![0.0f32; rows * cols];
4806        cortiq_core::quant::dequant_q1(&bytes, &mut w);
4807        let mut expect = vec![0.0f32; rows];
4808        for o in 0..rows {
4809            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
4810        }
4811        let mut got = vec![0.0f32; rows];
4812        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
4813        for o in 0..rows {
4814            assert!(
4815                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
4816                "row {o}: {} vs {}",
4817                got[o],
4818                expect[o]
4819            );
4820        }
4821        // Pair and batch paths agree with the single path.
4822        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
4823        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
4824        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
4825        assert_eq!(a1, got);
4826        let mut xs = x.clone();
4827        xs.extend_from_slice(&x2);
4828        let mut mm = vec![0.0f32; 2 * rows];
4829        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
4830        assert_eq!(&mm[..rows], got.as_slice());
4831        assert_eq!(&mm[rows..], a2.as_slice());
4832    }
4833
4834    #[test]
4835    fn repack_is_bit_identical() {
4836        // The interleaved-repack kernel must produce EXACTLY the same
4837        // bits as the mmap-layout kernel: integer accumulation is order-
4838        // exact, the f32 epilogue is identical. Odd rows exercise the
4839        // tail; direct range calls exercise unaligned pool splits.
4840        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
4841        let w: Vec<u8> = (0..rows * cols)
4842            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
4843            .collect();
4844        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
4845        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
4846        let rep = q8_repack_layout(&w, rows, cols);
4847        // Group interleave round-trips.
4848        for g in 0..rows / 4 {
4849            for c in 0..cols / 16 {
4850                for lane in 0..4 {
4851                    assert_eq!(
4852                        &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
4853                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
4854                    );
4855                }
4856            }
4857        }
4858        let mut a = vec![0.0f32; rows];
4859        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4860        let mut b = vec![0.0f32; rows];
4861        qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
4862        assert_eq!(a, b, "full-range repack output diverged");
4863
4864        #[cfg(target_arch = "aarch64")]
4865        if sdot_enabled() {
4866            // Unaligned range split (pool workers get arbitrary bounds).
4867            let act = split_act(&x);
4868            let mut c1 = vec![0.0f32; rows];
4869            let mut c2 = vec![0.0f32; rows];
4870            q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
4871            q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
4872            assert_eq!(c1, c2, "unaligned-range repack output diverged");
4873        }
4874    }
4875
4876    #[test]
4877    fn sdot_a8w8_noise_is_bounded() {
4878        // Off-grid activations: A8 quantization noise must stay small in
4879        // relative L2 over the whole output (realistic accuracy contract;
4880        // vmfcore measured argmax-identical decode on real models).
4881        let (rows, cols) = (16, 512);
4882        let w: Vec<u8> = (0..rows * cols)
4883            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4884            .collect();
4885        let scales = vec![0.01f32; rows];
4886        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
4887        let mut a = vec![0.0f32; rows];
4888        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4889        let (mut num, mut den) = (0f64, 0f64);
4890        for o in 0..rows {
4891            let mut acc = 0.0f32;
4892            for j in 0..cols {
4893                acc += (w[o * cols + j] as i8) as f32 * x[j];
4894            }
4895            let expect = acc * scales[o];
4896            num += ((a[o] - expect) as f64).powi(2);
4897            den += (expect as f64).powi(2);
4898        }
4899        let rel = (num / den.max(1e-12)).sqrt();
4900        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
4901    }
4902
4903    #[test]
4904    fn i8_dot_neon_matches_scalar() {
4905        let n = 100;
4906        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
4907        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
4908        let mut scalar = 0.0f32;
4909        for j in 0..n {
4910            scalar += (w[j] as i8) as f32 * x[j];
4911        }
4912        let fast = dot_i8_f32(&w, &x);
4913        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
4914    }
4915
4916    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
4917    #[test]
4918    fn vbitmatvec_matches_full_dequant() {
4919        let (rows, cols) = (6, 64);
4920        let ng = cols / GROUP_SIZE;
4921        // Hand-craft: bits per row, f16 scales, packed rows.
4922        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4923        let mut bytes = bits.clone();
4924        for g in 0..rows * ng {
4925            let s = 0.02 + 0.001 * g as f32;
4926            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4927        }
4928        for r in 0..rows {
4929            let b = bits[r] as usize;
4930            let (mut acc, mut nb) = (0u64, 0usize);
4931            let mut rowbytes = Vec::new();
4932            for i in 0..cols {
4933                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4934                acc = (acc << b) | v;
4935                nb += b;
4936                while nb >= 8 {
4937                    nb -= 8;
4938                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4939                }
4940            }
4941            if nb > 0 {
4942                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4943            }
4944            bytes.extend_from_slice(&rowbytes);
4945        }
4946        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4947
4948        let mut reference = vec![0f32; rows * cols];
4949        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
4950        let mut expect = vec![0f32; rows];
4951        for r in 0..rows {
4952            expect[r] = reference[r * cols..(r + 1) * cols]
4953                .iter()
4954                .zip(&x)
4955                .map(|(w, xv)| w * xv)
4956                .sum();
4957        }
4958        let mut got = vec![0f32; rows];
4959        let offsets = vbit_row_offsets(&bytes, rows, cols);
4960        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
4961        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
4962        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
4963        // the golden-parity gate).
4964        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4965        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
4966        for r in 0..rows {
4967            assert!(
4968                (got[r] - expect[r]).abs() < tol * scale,
4969                "row {r}: {} vs {}",
4970                got[r],
4971                expect[r]
4972            );
4973        }
4974    }
4975
4976    /// Fused q4 matvec must match the reference full-dequant + dense
4977    /// matvec bit-for-bit in structure (same f32 math, group order).
4978    /// vbit matmat: the blocked 1×4 leg must match the per-row path
4979    /// (paired env toggle; larger shape so both code paths engage).
4980    #[test]
4981    #[cfg(target_arch = "x86_64")]
4982    fn vbit_matmat_blocked_matches_per_row() {
4983        let (rows, cols, b) = (64usize, 128usize, 9usize);
4984        let ng = cols / GROUP_SIZE;
4985        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
4986        let mut bytes = bits.clone();
4987        for g in 0..rows * ng {
4988            let sc = 0.02 + 0.0005 * g as f32;
4989            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
4990        }
4991        for r in 0..rows {
4992            let bw = bits[r] as usize;
4993            let (mut acc, mut nb) = (0u64, 0usize);
4994            let mut rowbytes = Vec::new();
4995            for i in 0..cols {
4996                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
4997                acc = (acc << bw) | v;
4998                nb += bw;
4999                while nb >= 8 {
5000                    nb -= 8;
5001                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
5002                }
5003            }
5004            if nb > 0 {
5005                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5006            }
5007            bytes.extend_from_slice(&rowbytes);
5008        }
5009        let x: Vec<f32> =
5010            (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
5011        let offsets = vbit_row_offsets(&bytes, rows, cols);
5012        let mut y_a = vec![0f32; b * rows];
5013        let mut y_b = vec![0f32; b * rows];
5014        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
5015        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
5016        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
5017        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
5018        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
5019        let max_d =
5020            y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
5021        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
5022    }
5023
5024    #[test]
5025    fn q4matvec_matches_full_dequant() {
5026        let (rows, cols) = (8, 64);
5027        let groups = rows * cols / GROUP_SIZE;
5028        // Hand-craft a q4_block blob: nibbles then f16 scales.
5029        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5030        for i in 0..groups * 16 {
5031            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5032        }
5033        for g in 0..groups {
5034            let s = 0.01 + 0.003 * g as f32;
5035            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5036        }
5037        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5038
5039        let mut reference = vec![0.0f32; rows * cols];
5040        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5041        let mut expect = vec![0.0f32; rows];
5042        for r in 0..rows {
5043            expect[r] = reference[r * cols..(r + 1) * cols]
5044                .iter()
5045                .zip(&x)
5046                .map(|(w, xv)| w * xv)
5047                .sum();
5048        }
5049
5050        let mut got = vec![0.0f32; rows];
5051        q4matvec(&bytes, &x, rows, cols, &mut got, None);
5052        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
5053        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
5054        // in the golden-parity gate).
5055        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5056        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5057        for r in 0..rows {
5058            assert!(
5059                (got[r] - expect[r]).abs() < tol * scale,
5060                "row {r}: {} vs {}",
5061                got[r],
5062                expect[r]
5063            );
5064        }
5065    }
5066
5067    /// Fused two-input vbit matvec must equal two single matvecs exactly
5068    /// (same per-lane accumulation order on both scalar and SDOT paths).
5069    #[test]
5070    fn vbitmatvec2_equals_two_singles() {
5071        let (rows, cols) = (6, 64);
5072        let ng = cols / GROUP_SIZE;
5073        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5074        let mut bytes = bits.clone();
5075        for g in 0..rows * ng {
5076            let s = 0.02 + 0.001 * g as f32;
5077            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5078        }
5079        for r in 0..rows {
5080            let b = bits[r] as usize;
5081            let (mut acc, mut nb) = (0u64, 0usize);
5082            let mut rowbytes = Vec::new();
5083            for i in 0..cols {
5084                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5085                acc = (acc << b) | v;
5086                nb += b;
5087                while nb >= 8 {
5088                    nb -= 8;
5089                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
5090                }
5091            }
5092            if nb > 0 {
5093                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5094            }
5095            bytes.extend_from_slice(&rowbytes);
5096        }
5097        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5098        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
5099        let offsets = vbit_row_offsets(&bytes, rows, cols);
5100
5101        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5102        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
5103        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
5104        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5105        vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5106        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
5107        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
5108    }
5109
5110    /// Fused two-input q4 matvec must equal two single matvecs exactly.
5111    #[test]
5112    fn q4matvec2_equals_two_singles() {
5113        let (rows, cols) = (8, 128);
5114        let groups = rows * cols / GROUP_SIZE;
5115        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5116        for i in 0..groups * 16 {
5117            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5118        }
5119        for g in 0..groups {
5120            let s = 0.01 + 0.003 * g as f32;
5121            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5122        }
5123        // Include an outlier channel so the SDOT correction path is
5124        // exercised in the pair kernel too.
5125        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5126        x1[9] = 250.0;
5127        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5128
5129        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5130        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
5131        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
5132        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5133        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5134        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
5135        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
5136    }
5137
5138    /// Multi-matrix job must equal separate matvecs exactly — same
5139    /// kernels, only the dispatch is fused.
5140    #[test]
5141    fn matvec_many_equals_separate_matvecs() {
5142        use crate::pool::Pool;
5143        let (r1, r2, cols) = (300, 200, 64);
5144        let mk = |salt: usize, rows: usize| {
5145            QTensor::from_f32(
5146                (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
5147                rows,
5148                cols,
5149            )
5150        };
5151        let (a, b) = (mk(1, r1), mk(5, r2));
5152        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
5153        let pool = Pool::new(3);
5154
5155        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
5156        a.matvec(&x, &mut ea, Some(&pool));
5157        b.matvec(&x, &mut eb, Some(&pool));
5158        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
5159        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
5160        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
5161        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
5162    }
5163
5164    /// Batched q4/vbit matmat must equal per-position matvec calls
5165    /// exactly (the fallback it replaced) — same kernels, same order.
5166    #[test]
5167    fn batched_matmat_equals_per_position_matvec() {
5168        let (rows, cols, b) = (8, 64, 5);
5169        // q4 blob.
5170        let groups = rows * cols / GROUP_SIZE;
5171        let mut q4 = Vec::new();
5172        for i in 0..groups * 16 {
5173            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5174        }
5175        for g in 0..groups {
5176            q4.extend_from_slice(
5177                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5178            );
5179        }
5180        // vbit blob (mixed widths incl. 8).
5181        let ng = cols / GROUP_SIZE;
5182        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
5183        let mut vb = bits.clone();
5184        for g in 0..rows * ng {
5185            vb.extend_from_slice(
5186                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
5187            );
5188        }
5189        for r in 0..rows {
5190            let bw = bits[r] as usize;
5191            let (mut acc, mut nb) = (0u64, 0usize);
5192            let mut rowbytes = Vec::new();
5193            for i in 0..cols {
5194                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5195                acc = (acc << bw) | v;
5196                nb += bw;
5197                while nb >= 8 {
5198                    nb -= 8;
5199                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
5200                }
5201            }
5202            if nb > 0 {
5203                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5204            }
5205            vb.extend_from_slice(&rowbytes);
5206        }
5207        let offsets = vbit_row_offsets(&vb, rows, cols);
5208
5209        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5210
5211        // q4: batch vs singles.
5212        let mut got = vec![0f32; b * rows];
5213        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
5214        for bi in 0..b {
5215            let mut expect = vec![0f32; rows];
5216            q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
5217            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
5218        }
5219
5220        // vbit: batch vs singles.
5221        let mut got = vec![0f32; b * rows];
5222        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
5223        for bi in 0..b {
5224            let mut expect = vec![0f32; rows];
5225            vbitmatvec(
5226                &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
5227            );
5228            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
5229        }
5230    }
5231
5232    /// q4_tiled kernels must produce BIT-identical outputs to the q4
5233    /// split kernels on the same values (same ints, same order — only
5234    /// the byte placement differs).
5235    #[test]
5236    fn q4_tiled_matches_q4_block_bitexact() {
5237        let (rows, cols, b) = (8usize, 128usize, 3usize);
5238        let groups = rows * cols / GROUP_SIZE;
5239        let mut split = Vec::with_capacity(groups * 18);
5240        for i in 0..groups * 16 {
5241            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5242        }
5243        for g in 0..groups {
5244            split.extend_from_slice(
5245                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5246            );
5247        }
5248        // Re-tile: [scale][nibbles] per group.
5249        let (packed, scales) = split.split_at(groups * 16);
5250        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
5251        for g in 0..groups {
5252            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
5253            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
5254        }
5255
5256        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5257        x1[9] = 250.0; // exercise the outlier path
5258        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5259
5260        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
5261        q4matvec(&split, &x1, rows, cols, &mut a, None);
5262        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
5263        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
5264
5265        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5266        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
5267        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
5268        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
5269        assert_eq!(a1, t1);
5270        assert_eq!(a2, t2);
5271
5272        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5273        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
5274        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
5275        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
5276        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
5277    }
5278
5279    /// q4 SDOT outlier correction: a single huge activation channel
5280    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
5281    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
5282    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
5283    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
5284    /// can never qualify (8² = n).
5285    #[test]
5286    fn q4matvec_sdot_outlier_exact() {
5287        let (rows, cols) = (4, 128);
5288        let groups = rows * cols / GROUP_SIZE;
5289        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5290        for i in 0..groups * 16 {
5291            bytes.push(((i * 11 + 5) % 256) as u8);
5292        }
5293        for g in 0..groups {
5294            let s = 0.02 + 0.002 * g as f32;
5295            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5296        }
5297        let mut x: Vec<f32> = (0..cols)
5298            .map(|i| match i % 3 {
5299                0 => 1.0,
5300                1 => -1.0,
5301                _ => 0.0,
5302            })
5303            .collect();
5304        x[17] = 300.0; // ≫ 8·rms → outlier channel
5305
5306        let mut reference = vec![0.0f32; rows * cols];
5307        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5308        let mut expect = vec![0.0f32; rows];
5309        for r in 0..rows {
5310            expect[r] = reference[r * cols..(r + 1) * cols]
5311                .iter()
5312                .zip(&x)
5313                .map(|(w, xv)| w * xv)
5314                .sum();
5315        }
5316        let mut got = vec![0.0f32; rows];
5317        q4matvec(&bytes, &x, rows, cols, &mut got, None);
5318        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5319        for r in 0..rows {
5320            assert!(
5321                (got[r] - expect[r]).abs() < 2e-3 * scale,
5322                "row {r}: {} vs {} (outlier term must be exact)",
5323                got[r],
5324                expect[r]
5325            );
5326        }
5327    }
5328}