Skip to main content

cortiq_engine/
qtensor.rs

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