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