Skip to main content

cortiq_engine/
qtensor.rs

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