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    pub fn rows(&self) -> usize {
253        match self {
254            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
255        }
256    }
257
258    pub fn cols(&self) -> usize {
259        match self {
260            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
261        }
262    }
263
264    /// Dense f32 view — only for owned tensors. Masked/sparse execution
265    /// paths require it; quantized weights don't support masks yet.
266    pub fn as_f32(&self) -> Option<&[f32]> {
267        match self {
268            Self::F32 { data, .. } => Some(data),
269            Self::Mapped { .. } => None,
270        }
271    }
272
273    fn quant_bytes(&self) -> &[u8] {
274        match self {
275            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
276            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
277        }
278    }
279
280    /// Dequantize one row into `dst` (embedding lookup).
281    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
282        let cols = self.cols();
283        debug_assert_eq!(dst.len(), cols);
284        match self {
285            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
286            Self::Mapped {
287                dtype,
288                row_scale,
289                col_field,
290                vbit_offsets,
291                ..
292            } => {
293                if *dtype == TensorDtype::Q4Tiled {
294                    let bytes = self.quant_bytes();
295                    let gpr = cols / GROUP_SIZE;
296                    for gi in 0..gpr {
297                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
298                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
299                        for (k, &b) in tile[2..].iter().enumerate() {
300                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
301                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
302                        }
303                    }
304                    return;
305                }
306                if *dtype == TensorDtype::Q4Block {
307                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
308                    let gpr = cols / GROUP_SIZE;
309                    for gi in 0..gpr {
310                        let g = r * gpr + gi;
311                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
312                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
313                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
314                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
315                        }
316                    }
317                    return;
318                }
319                if *dtype == TensorDtype::Q1 {
320                    let bytes = self.quant_bytes();
321                    let gpr = cols / GROUP_SIZE;
322                    for gi in 0..gpr {
323                        let tile =
324                            &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
325                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
326                        for (j, &b) in tile[2..].iter().enumerate() {
327                            for k in 0..8 {
328                                dst[gi * GROUP_SIZE + j * 8 + k] =
329                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
330                            }
331                        }
332                    }
333                    return;
334                }
335                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
336                    let bytes = self.quant_bytes();
337                    let rows = self.rows();
338                    let ng = cols / GROUP_SIZE;
339                    let bits = &bytes[..rows];
340                    let sc_off = rows;
341                    // Precomputed at load — embedding lookup used to scan
342                    // the bit-widths of every preceding row (O(token_id)).
343                    let off = vbit_offsets[r];
344                    let b = bits[r] as usize;
345                    let l = ((1usize << (b - 1)) - 1) as f32;
346                    let data = &bytes[off..];
347                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
348                    for (i, d) in dst.iter_mut().enumerate() {
349                        while nbits < b {
350                            acc = (acc << 8) | data[idx] as u64;
351                            idx += 1;
352                            nbits += 8;
353                        }
354                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
355                        nbits -= b;
356                        let so = (r * ng + i / GROUP_SIZE) * 2;
357                        let sv = f16_to_f32(u16::from_le_bytes([
358                            bytes[sc_off + so],
359                            bytes[sc_off + so + 1],
360                        ]));
361                        *d = (u - l) * sv;
362                    }
363                    return;
364                }
365                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
366                let s = row_scale[r];
367                match dtype {
368                    TensorDtype::Q8Row => {
369                        for (d, &b) in dst.iter_mut().zip(q) {
370                            *d = (b as i8) as f32 * s;
371                        }
372                    }
373                    TensorDtype::Q8_2f => {
374                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
375                            *d = (b as i8) as f32 * s * col_field[i];
376                        }
377                    }
378                    _ => unreachable!(),
379                }
380            }
381        }
382    }
383
384    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
385    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
386    /// false for group-packed q4/vbit (column access would unpack whole
387    /// groups — sparse execution falls back to f32 for those).
388    pub fn sparse_col_ok(&self) -> bool {
389        match self {
390            Self::F32 { .. } => true,
391            Self::Mapped { dtype, .. } => {
392                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
393            }
394        }
395    }
396
397    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
398    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
399    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
400    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
401        let inter = self.cols();
402        let hidden = self.rows();
403        debug_assert_eq!(out.len(), hidden);
404        match self {
405            Self::F32 { data, .. } => {
406                for (k, o) in out.iter_mut().enumerate() {
407                    *o += w * data[k * inter + c];
408                }
409            }
410            Self::Mapped {
411                dtype, row_scale, col_field, ..
412            } => {
413                let q = self.quant_bytes();
414                let colf = if *dtype == TensorDtype::Q8_2f {
415                    col_field[c]
416                } else {
417                    1.0
418                };
419                let wc = w * colf;
420                for (k, o) in out.iter_mut().enumerate() {
421                    let b = q[k * inter + c] as i8 as f32;
422                    *o += wc * b * row_scale[k];
423                }
424            }
425        }
426    }
427
428    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
429    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
430    /// into `scratch` first (rare for active-FFN weights).
431    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
432        let cols = self.cols();
433        match self {
434            Self::F32 { data, .. } => {
435                let row = &data[r * cols..(r + 1) * cols];
436                row.iter().zip(x).map(|(w, v)| w * v).sum()
437            }
438            Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
439                TensorDtype::Q8Row => {
440                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
441                    dot_i8_f32(q, x) * row_scale[r]
442                }
443                TensorDtype::Q8_2f => {
444                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
445                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
446                }
447                _ => {
448                    self.row_f32(r, scratch);
449                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
450                }
451            },
452        }
453    }
454
455    /// `out = W · x` (row-major). F32 delegates to the historical
456    /// bit-exact path; Mapped runs the fused int8 kernel.
457    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
458        match self {
459            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
460            Self::Mapped {
461                model,
462                idx,
463                dtype,
464                rows,
465                cols,
466                row_scale,
467                col_field,
468                vbit_offsets,
469                repack,
470            } => {
471                let _ = (model, idx);
472                if *dtype == TensorDtype::Q4Block {
473                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
474                    return;
475                }
476                if *dtype == TensorDtype::Q4Tiled {
477                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
478                    return;
479                }
480                if *dtype == TensorDtype::Q1 {
481                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
482                    return;
483                }
484                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
485                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
486                    return;
487                }
488                let xs = prescale(x, col_field, *dtype);
489                // D5: large q8 matrices (lm_head-class) — hybrid
490                // CPU∥GPU: split the rows, both sides compute
491                // SIMULTANEOUSLY (same math, shared prescale).
492                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
493                if *rows >= crate::gpu::min_rows()
494                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
495                    && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
496                    && crate::gpu::enabled_here()
497                {
498                    // Runtime probe: alternate the hybrid against the
499                    // pure-CPU matvec, keep whichever is faster HERE.
500                    let t0 = std::time::Instant::now();
501                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
502                        crate::gpu::ProbeArm::Gpu => {}
503                        crate::gpu::ProbeArm::CpuTimed => {
504                            qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
505                            crate::gpu::probe_record(
506                                crate::gpu::OpClass::Matvec,
507                                false,
508                                t0.elapsed(),
509                            );
510                            return;
511                        }
512                        crate::gpu::ProbeArm::Cpu => {
513                            qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
514                            return;
515                        }
516                    }
517                    let frac = std::env::var("CMF_GPU_SPLIT")
518                        .ok()
519                        .and_then(|v| v.parse::<f32>().ok())
520                        .unwrap_or(0.5)
521                        .clamp(0.0, 1.0);
522                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
523                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
524                    let bytes = self.quant_bytes();
525                    let ok = std::thread::scope(|sc| {
526                        let g = sc.spawn(|| {
527                            crate::gpu::q8_matvec_range(
528                                model,
529                                *idx,
530                                cpu_rows,
531                                &row_scale[cpu_rows..],
532                                &xs,
533                                *rows - cpu_rows,
534                                *cols,
535                                out_gpu,
536                            )
537                        });
538                        if cpu_rows > 0 {
539                            // Repack prefix covers the full groups of the
540                            // CPU half (the split starts at row 0).
541                            let rep_cpu = if repack.is_empty() {
542                                &[][..]
543                            } else {
544                                &repack[..(cpu_rows / 4) * 4 * *cols]
545                            };
546                            qmatvec(
547                                &bytes[..cpu_rows * *cols],
548                                rep_cpu,
549                                &row_scale[..cpu_rows],
550                                &xs,
551                                cpu_rows,
552                                *cols,
553                                out_cpu,
554                                pool,
555                            );
556                        }
557                        g.join().unwrap_or(false)
558                    });
559                    if ok {
560                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
561                        return;
562                    }
563                    // GPU failed — CPU finishes its half (rows rebased —
564                    // group offsets don't line up, mmap layout only).
565                    qmatvec(
566                        &bytes[cpu_rows * *cols..(*rows) * *cols],
567                        &[],
568                        &row_scale[cpu_rows..],
569                        &xs,
570                        *rows - cpu_rows,
571                        *cols,
572                        out_gpu,
573                        pool,
574                    );
575                    return;
576                }
577                qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
578            }
579        }
580    }
581
582    /// Fused two-input matvec (MTP verify pair): weights streamed once.
583    pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
584        match self {
585            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
586            Self::Mapped {
587                dtype,
588                rows,
589                cols,
590                row_scale,
591                col_field,
592                vbit_offsets,
593                ..
594            } => {
595                if *dtype == TensorDtype::Q4Block {
596                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
597                    return;
598                }
599                if *dtype == TensorDtype::Q4Tiled {
600                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
601                    return;
602                }
603                if *dtype == TensorDtype::Q1 {
604                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
605                    return;
606                }
607                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
608                    vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
609                    return;
610                }
611                let x1s = prescale(x1, col_field, *dtype);
612                let x2s = prescale(x2, col_field, *dtype);
613                qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
614            }
615        }
616    }
617}
618
619impl QTensor {
620    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
621    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
622    /// to b matvec calls (same dot kernels in the same order); the win —
623    /// the weight row streams from DRAM once per batch, not b times.
624    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
625        let cols = self.cols();
626        let rows = self.rows();
627        debug_assert_eq!(xs_all.len(), b * cols);
628        debug_assert_eq!(out.len(), b * rows);
629        match self {
630            Self::F32 { data, .. } => {
631                let out_addr = SendMut(out.as_mut_ptr());
632                let run = |start: usize, end: usize| {
633                    for o in start..end {
634                        let row = &data[o * cols..(o + 1) * cols];
635                        for bi in 0..b {
636                            let x = &xs_all[bi * cols..(bi + 1) * cols];
637                            let mut acc = 0f32;
638                            for j in 0..cols {
639                                acc += row[j] * x[j];
640                            }
641                            unsafe { *out_addr.at(bi * rows + o) = acc };
642                        }
643                    }
644                };
645                dispatch_rows(pool, rows, &run);
646            }
647            Self::Mapped {
648                dtype,
649                row_scale,
650                col_field,
651                vbit_offsets,
652                ..
653            } => {
654                if *dtype == TensorDtype::Q4Block {
655                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
656                    return;
657                }
658                if *dtype == TensorDtype::Q4Tiled {
659                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
660                    return;
661                }
662                if *dtype == TensorDtype::Q1 {
663                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
664                    return;
665                }
666                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
667                    vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
668                    return;
669                }
670                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
671                    .map(|bi| {
672                        prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
673                    })
674                    .collect();
675                // D5: large prefill-batch GEMMs — on the GPU (threshold by
676                // work volume: submission carries b×rows×cols MACs).
677                // Runtime probe: the naive GEMM shader + sync readback
678                // lose to the CPU GEMM on slow driver stacks — alternate
679                // both arms and keep the winner.
680                if b >= 8
681                    && b * rows * cols >= 128_000_000
682                    && crate::gpu::enabled_here()
683                {
684                    if let Self::Mapped { model, idx, .. } = self {
685                        let t0 = std::time::Instant::now();
686                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
687                            crate::gpu::ProbeArm::Gpu
688                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
689                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
690                            {
691                                // Cold weights during probing: the upload
692                                // has started, the count runs on the CPU —
693                                // the GPU arm samples on the next touch.
694                                let q = self.quant_bytes();
695                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
696                                return;
697                            }
698                            crate::gpu::ProbeArm::Gpu => {
699                                let flat: Vec<f32> =
700                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
701                                if crate::gpu::q8_matmat(
702                                    model, *idx, row_scale, &flat, b, rows, cols, out)
703                                {
704                                    crate::gpu::probe_record(
705                                        crate::gpu::OpClass::Matmat,
706                                        true,
707                                        t0.elapsed(),
708                                    );
709                                    return;
710                                }
711                            }
712                            crate::gpu::ProbeArm::CpuTimed => {
713                                let q = self.quant_bytes();
714                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
715                                crate::gpu::probe_record(
716                                    crate::gpu::OpClass::Matmat,
717                                    false,
718                                    t0.elapsed(),
719                                );
720                                return;
721                            }
722                            crate::gpu::ProbeArm::Cpu => {}
723                        }
724                    }
725                }
726                let q = self.quant_bytes();
727                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
728            }
729        }
730    }
731}
732
733impl QTensor {
734    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
735    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
736    /// barrier instead of N. Per-row math is the exact same kernel as
737    /// `matvec` (bit-identical outputs); only the dispatch is fused.
738    /// Falls back to N sequential matvecs when the set is not a uniform
739    /// q8-family/F32 group or there is no pool.
740    pub fn matvec_many<const N: usize>(
741        ts: [&QTensor; N],
742        x: &[f32],
743        mut outs: [&mut [f32]; N],
744        pool: Option<&Pool>,
745    ) {
746        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
747        let uniform_q8 = ts.iter().all(|t| {
748            matches!(
749                t,
750                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
751            )
752        });
753        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
754        let uniform_q4 = ts
755            .iter()
756            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
757        let uniform_vbit = ts
758            .iter()
759            .all(|t| matches!(
760                t,
761                Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
762            ));
763        let uniform_q1 = ts
764            .iter()
765            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
766        let Some(pool) = pool else {
767            for (t, o) in ts.iter().zip(outs.iter_mut()) {
768                t.matvec(x, o, None);
769            }
770            return;
771        };
772        if total_rows < 256
773            || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
774        {
775            for (t, o) in ts.iter().zip(outs.iter_mut()) {
776                t.matvec(x, o, Some(pool));
777            }
778            return;
779        }
780
781        if uniform_q1 {
782            // One shared activation split + group sums (q1 has no col
783            // field; the same input feeds every tensor).
784            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
785            if a8w8_enabled() {
786                let act = split_act(x);
787                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
788                let (act, gsum) = (&act, &gsum);
789                let closures: [_; N] = std::array::from_fn(|i| {
790                    let (bytes, gpr, out) =
791                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
792                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
793                });
794                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
795                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
796                pool.run_many(&parts);
797            } else {
798                let closures: [_; N] = std::array::from_fn(|i| {
799                    let (bytes, gpr, out) =
800                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
801                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
802                });
803                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
804                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
805                pool.run_many(&parts);
806            }
807            return;
808        }
809
810        if uniform_q4 || uniform_vbit {
811            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
812            // q4/vbit share one activation split — no per-tensor col field.
813            if a8w8_enabled() {
814                let act = split_act(x);
815                let act = &act;
816                if uniform_q4 {
817                    let closures: [_; N] = std::array::from_fn(|i| {
818                        let (packed, scales) =
819                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
820                        let (gpr, cols, out) =
821                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
822                        move |s: usize, e: usize| {
823                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
824                        }
825                    });
826                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
827                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
828                    pool.run_many(&parts);
829                } else {
830                    let closures: [_; N] = std::array::from_fn(|i| {
831                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
832                        let (bytes, rows, cols, out) =
833                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
834                        move |s: usize, e: usize| {
835                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
836                        }
837                    });
838                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
839                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
840                    pool.run_many(&parts);
841                }
842                return;
843            }
844            if uniform_q4 {
845                let closures: [_; N] = std::array::from_fn(|i| {
846                    let (packed, scales) =
847                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
848                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
849                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
850                });
851                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
852                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
853                pool.run_many(&parts);
854            } else {
855                let closures: [_; N] = std::array::from_fn(|i| {
856                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
857                    let (bytes, rows, cols, out) =
858                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
859                    move |s: usize, e: usize| {
860                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
861                    }
862                });
863                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
864                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
865                pool.run_many(&parts);
866            }
867            return;
868        }
869
870        if uniform_f32 {
871            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
872            let closures: [_; N] = std::array::from_fn(|i| {
873                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
874                let out = outs_addr[i];
875                move |start: usize, end: usize| {
876                    for o in start..end {
877                        let row = &data[o * cols..(o + 1) * cols];
878                        let mut sum = 0.0f32;
879                        for j in 0..*cols {
880                            sum += row[j] * x[j];
881                        }
882                        // SAFETY: disjoint (tensor, row) cells per worker.
883                        unsafe { *out.at(o) = sum };
884                    }
885                }
886            });
887            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
888                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
889            pool.run_many(&parts);
890            return;
891        }
892
893        // Uniform q8-family: per-tensor prescale (q8_2f col fields
894        // differ per tensor) + the shared range kernels.
895        struct Ctx<'a> {
896            bytes: &'a [u8],
897            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
898            rep: &'a [u8],
899            row_scale: &'a [f32],
900            cols: usize,
901            xs: std::borrow::Cow<'a, [f32]>,
902        }
903        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
904            let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
905                unreachable!()
906            };
907            Ctx {
908                bytes: ts[i].quant_bytes(),
909                rep: repack,
910                row_scale,
911                cols: *cols,
912                xs: prescale(x, col_field, *dtype),
913            }
914        });
915        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
916        #[cfg(target_arch = "aarch64")]
917        if sdot_enabled() {
918            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
919            let closures: [_; N] = std::array::from_fn(|i| {
920                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
921                move |start: usize, end: usize| {
922                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
923                }
924            });
925            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
926                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
927            pool.run_many(&parts);
928            return;
929        }
930        #[cfg(target_arch = "x86_64")]
931        if avx2_a8w8_enabled() {
932            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
933            let closures: [_; N] = std::array::from_fn(|i| {
934                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
935                move |start: usize, end: usize| {
936                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
937                }
938            });
939            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
940                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
941            pool.run_many(&parts);
942            return;
943        }
944        let closures: [_; N] = std::array::from_fn(|i| {
945            let (c, out) = (&ctxs[i], outs_addr[i]);
946            move |start: usize, end: usize| {
947                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
948            }
949        });
950        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
951            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
952        pool.run_many(&parts);
953    }
954}
955
956impl QTensor {
957    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
958    /// single pool dispatch — the MTP/pair decode path publishes one job
959    /// for Q/K/V (and one for gate+up) instead of one per tensor.
960    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
961    #[allow(clippy::needless_range_loop)]
962    pub fn matvec2_many<const N: usize>(
963        ts: [&QTensor; N],
964        x1: &[f32],
965        x2: &[f32],
966        mut o1s: [&mut [f32]; N],
967        mut o2s: [&mut [f32]; N],
968        pool: Option<&Pool>,
969    ) {
970        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
971        let uniform_q8 = ts.iter().all(|t| {
972            matches!(
973                t,
974                Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
975            )
976        });
977        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
978        let uniform_q4 = ts
979            .iter()
980            .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
981        let uniform_vbit = ts
982            .iter()
983            .all(|t| matches!(
984                t,
985                Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
986            ));
987        let fusable = pool.is_some()
988            && total_rows >= 256
989            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
990        if !fusable {
991            for i in 0..N {
992                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
993            }
994            return;
995        }
996        let pool = pool.unwrap();
997
998        if uniform_q4 || uniform_vbit {
999            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1000            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1001            // q4/vbit share activation splits — no per-tensor col field.
1002            if a8w8_enabled() {
1003                let a1 = split_act(x1);
1004                let a2 = split_act(x2);
1005                let (a1, a2) = (&a1, &a2);
1006                if uniform_q4 {
1007                    let closures: [_; N] = std::array::from_fn(|i| {
1008                        let (packed, scales) =
1009                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1010                        let (gpr, cols, o1, o2) =
1011                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1012                        move |s: usize, e: usize| {
1013                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1014                        }
1015                    });
1016                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1017                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1018                    pool.run_many(&parts);
1019                } else {
1020                    let closures: [_; N] = std::array::from_fn(|i| {
1021                        let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1022                        let (bytes, rows, cols, o1, o2) =
1023                            (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1024                        move |s: usize, e: usize| {
1025                            vbit_range2_a8w8(
1026                                bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1027                            )
1028                        }
1029                    });
1030                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1031                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1032                    pool.run_many(&parts);
1033                }
1034                return;
1035            }
1036            if uniform_q4 {
1037                let closures: [_; N] = std::array::from_fn(|i| {
1038                    let (packed, scales) =
1039                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1040                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1041                    move |s: usize, e: usize| {
1042                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1043                    }
1044                });
1045                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1046                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1047                pool.run_many(&parts);
1048            } else {
1049                let closures: [_; N] = std::array::from_fn(|i| {
1050                    let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1051                    let (bytes, rows, cols, o1, o2) =
1052                        (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1053                    move |s: usize, e: usize| {
1054                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1055                    }
1056                });
1057                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1058                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1059                pool.run_many(&parts);
1060            }
1061            return;
1062        }
1063
1064        if uniform_f32 {
1065            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1066            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1067            let closures: [_; N] = std::array::from_fn(|i| {
1068                let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1069                let (o1, o2) = (p1[i], p2[i]);
1070                move |start: usize, end: usize| {
1071                    for o in start..end {
1072                        let row = &data[o * cols..(o + 1) * cols];
1073                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
1074                        for j in 0..*cols {
1075                            s1 += row[j] * x1[j];
1076                            s2 += row[j] * x2[j];
1077                        }
1078                        // SAFETY: disjoint (tensor, row) cells per worker.
1079                        unsafe {
1080                            *o1.at(o) = s1;
1081                            *o2.at(o) = s2;
1082                        }
1083                    }
1084                }
1085            });
1086            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1087                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1088            pool.run_many(&parts);
1089            return;
1090        }
1091
1092        struct Ctx<'a> {
1093            bytes: &'a [u8],
1094            row_scale: &'a [f32],
1095            cols: usize,
1096            xs1: std::borrow::Cow<'a, [f32]>,
1097            xs2: std::borrow::Cow<'a, [f32]>,
1098        }
1099        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1100            let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1101                unreachable!()
1102            };
1103            Ctx {
1104                bytes: ts[i].quant_bytes(),
1105                row_scale,
1106                cols: *cols,
1107                xs1: prescale(x1, col_field, *dtype),
1108                xs2: prescale(x2, col_field, *dtype),
1109            }
1110        });
1111        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1112        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1113        #[cfg(target_arch = "aarch64")]
1114        if sdot_enabled() {
1115            let acts: [(SplitAct, SplitAct); N] =
1116                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1117            let closures: [_; N] = std::array::from_fn(|i| {
1118                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1119                move |start: usize, end: usize| {
1120                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1121                }
1122            });
1123            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1124                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1125            pool.run_many(&parts);
1126            return;
1127        }
1128        #[cfg(target_arch = "x86_64")]
1129        if avx2_a8w8_enabled() {
1130            let acts: [(SplitAct, SplitAct); N] =
1131                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1132            let closures: [_; N] = std::array::from_fn(|i| {
1133                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1134                move |start: usize, end: usize| {
1135                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1136                }
1137            });
1138            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1139                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1140            pool.run_many(&parts);
1141            return;
1142        }
1143        let closures: [_; N] = std::array::from_fn(|i| {
1144            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1145            move |start: usize, end: usize| {
1146                q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1147            }
1148        });
1149        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1150            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1151        pool.run_many(&parts);
1152    }
1153}
1154
1155/// Batched q8 kernel: same math as qmatvec, the row makes a single
1156/// pass from memory for the whole batch.
1157fn qmatmat(
1158    q: &[u8],
1159    row_scale: &[f32],
1160    pre: &[std::borrow::Cow<'_, [f32]>],
1161    rows: usize,
1162    cols: usize,
1163    out: &mut [f32],
1164    pool: Option<&Pool>,
1165) {
1166    let b = pre.len();
1167    debug_assert_eq!(out.len(), b * rows);
1168    #[cfg(target_arch = "aarch64")]
1169    if sdot_enabled() {
1170        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1171        let out_addr = SendMut(out.as_mut_ptr());
1172        let run = |start: usize, end: usize| {
1173            for o in start..end {
1174                let row = &q[o * cols..(o + 1) * cols];
1175                for (bi, act) in acts.iter().enumerate() {
1176                    let v = row_dot_sdot(row, act) * row_scale[o];
1177                    unsafe { *out_addr.at(bi * rows + o) = v };
1178                }
1179            }
1180        };
1181        dispatch_rows(pool, rows, &run);
1182        return;
1183    }
1184    // x86 A8W8 batch: each weight row streams once for the whole chunk
1185    // through the AVX2/VNNI row dot (prefill q8 was the last
1186    // aarch64-only batched path).
1187    #[cfg(target_arch = "x86_64")]
1188    if avx2_a8w8_enabled() {
1189        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1190        let out_addr = SendMut(out.as_mut_ptr());
1191        let run = |start: usize, end: usize| {
1192            for o in start..end {
1193                let row = &q[o * cols..(o + 1) * cols];
1194                for (bi, act) in acts.iter().enumerate() {
1195                    let v = row_dot_avx2(row, act) * row_scale[o];
1196                    unsafe { *out_addr.at(bi * rows + o) = v };
1197                }
1198            }
1199        };
1200        dispatch_rows(pool, rows, &run);
1201        return;
1202    }
1203    let out_addr = SendMut(out.as_mut_ptr());
1204    let run = |start: usize, end: usize| {
1205        for o in start..end {
1206            let row = &q[o * cols..(o + 1) * cols];
1207            for (bi, x) in pre.iter().enumerate() {
1208                let mut acc = 0f32;
1209                for j in 0..cols {
1210                    acc += (row[j] as i8) as f32 * x[j];
1211                }
1212                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1213            }
1214        }
1215    };
1216    dispatch_rows(pool, rows, &run);
1217}
1218
1219/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
1220/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
1221fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1222    match pool {
1223        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1224        _ => run(0, rows),
1225    }
1226}
1227
1228/// Split a q4_block blob into (packed nibbles, f16 group scales).
1229fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1230    let groups = rows * cols / GROUP_SIZE;
1231    bytes.split_at(groups * 16)
1232}
1233
1234/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
1235/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
1236/// vbit packs MSB-first, so the HIGH nibble is the even element
1237/// (opposite of q4_block's lo-first interleave). Centering is u-7.
1238#[inline]
1239fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1240    #[cfg(target_arch = "aarch64")]
1241    unsafe {
1242        return vbit_fill4_neon(data, buf);
1243    }
1244    #[cfg(target_arch = "x86_64")]
1245    if avx2_enabled() {
1246        return unsafe { vbit_fill4_avx2(data, buf) };
1247    }
1248    #[allow(unreachable_code)]
1249    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1250        let u = unpack8::<4>(&data[blk * 4..]);
1251        for k in 0..8 {
1252            chunk[k] = (u[k] - 7) as i8 as u8;
1253        }
1254    }
1255}
1256
1257#[cfg(target_arch = "aarch64")]
1258#[target_feature(enable = "neon")]
1259unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1260    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
1261    // buf.len()/2 packed bytes (validated at load).
1262    unsafe {
1263        use core::arch::aarch64::*;
1264        let n = buf.len();
1265        let mask = vdupq_n_u8(0x0F);
1266        let seven = vdupq_n_s8(7);
1267        let mut g = 0usize;
1268        while g * 32 + 32 <= n {
1269            let b = vld1q_u8(data.as_ptr().add(g * 16));
1270            let hi = vshrq_n_u8::<4>(b);
1271            let lo = vandq_u8(b, mask);
1272            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1273            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1274            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1275            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1276            g += 1;
1277        }
1278    }
1279}
1280
1281#[cfg(target_arch = "x86_64")]
1282#[target_feature(enable = "avx2")]
1283unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1284    // SAFETY: see vbit_fill4_neon.
1285    unsafe {
1286        use core::arch::x86_64::*;
1287        let n = buf.len();
1288        let mask = _mm_set1_epi8(0x0F);
1289        let seven = _mm256_set1_epi8(7);
1290        let mut g = 0usize;
1291        while g * 32 + 32 <= n {
1292            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1293            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1294            let lo = _mm_and_si128(b, mask);
1295            let z = _mm256_sub_epi8(
1296                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1297                seven,
1298            );
1299            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1300            g += 1;
1301        }
1302    }
1303}
1304
1305/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
1306/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
1307/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
1308/// into 4 such blocks.
1309#[inline(always)]
1310fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1311    let mut acc = 0u64;
1312    for i in 0..B {
1313        acc = (acc << 8) | data[i] as u64;
1314    }
1315    let mask = (1u64 << B) - 1;
1316    let mut out = [0i32; 8];
1317    for (k, o) in out.iter_mut().enumerate() {
1318        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1319    }
1320    out
1321}
1322
1323/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
1324/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
1325/// MSB-first, byte-padded]. Row data offsets are precomputed at load
1326/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
1327/// overhead on every matvec.
1328#[allow(clippy::too_many_arguments)]
1329fn vbitmatvec(
1330    bytes: &[u8],
1331    offsets: &[usize],
1332    x: &[f32],
1333    rows: usize,
1334    cols: usize,
1335    out: &mut [f32],
1336    pool: Option<&Pool>,
1337) {
1338    debug_assert_eq!(out.len(), rows);
1339    debug_assert_eq!(offsets.len(), rows + 1);
1340
1341    // SDOT path: unpack the row to centered i8 once, then per-group
1342    // int8 dot against the quantized activations — same A8W8 contract
1343    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
1344    if a8w8_enabled() {
1345        let act = split_act(x);
1346        let out_addr = SendMut(out.as_mut_ptr());
1347        let run = move |start: usize, end: usize| {
1348            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1349        };
1350        dispatch_rows(pool, rows, &run);
1351        return;
1352    }
1353
1354    let out_addr = SendMut(out.as_mut_ptr());
1355    let run = move |start: usize, end: usize| {
1356        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1357    };
1358    dispatch_rows(pool, rows, &run);
1359}
1360
1361/// One vbit row range via the A8W8 int8 path — kernel body of
1362/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
1363/// several tensors in one dispatch (b=8 rows go exact f32).
1364#[allow(clippy::too_many_arguments)]
1365fn vbit_range_a8w8(
1366    bytes: &[u8],
1367    offsets: &[usize],
1368    x: &[f32],
1369    act: &SplitAct,
1370    rows: usize,
1371    cols: usize,
1372    out: SendMut,
1373    start: usize,
1374    end: usize,
1375) {
1376    let ng = cols / GROUP_SIZE;
1377    let bits = &bytes[..rows];
1378    let sc_off = rows;
1379    let row_dot = |r: usize| -> f32 {
1380            let b = bits[r] as usize;
1381            let l = ((1i32 << (b - 1)) - 1) as i32;
1382            let mask = (1u64 << b) - 1;
1383            let data = &bytes[offsets[r]..offsets[r + 1]];
1384            if b == 8 {
1385                // u−L reaches 128 → does not fit i8; exact f32 path.
1386                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1387                let mut dot = 0f32;
1388                for g in 0..ng {
1389                    let so = (r * ng + g) * 2;
1390                    let sgf = f16_to_f32(u16::from_le_bytes([
1391                        bytes[sc_off + so],
1392                        bytes[sc_off + so + 1],
1393                    ]));
1394                    let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1395                    let mut gd = 0f32;
1396                    for &xv in xg.iter() {
1397                        if nbits < 8 {
1398                            acc = (acc << 8) | data[idx] as u64;
1399                            idx += 1;
1400                            nbits += 8;
1401                        }
1402                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1403                        nbits -= 8;
1404                        gd += (u - l) as f32 * xv;
1405                    }
1406                    dot += gd * sgf;
1407                }
1408                return dot;
1409            }
1410            // Per-worker scratch: this closure runs for every row of the
1411            // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
1412            // row was measurable pure overhead.
1413            thread_local! {
1414                static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1415                    const { std::cell::RefCell::new(Vec::new()) };
1416            }
1417            #[inline(always)]
1418            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1419                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1420                    let u = unpack8::<B>(&data[blk * B..]);
1421                    for k in 0..8 {
1422                        chunk[k] = (u[k] - l) as i8 as u8;
1423                    }
1424                }
1425            }
1426            let _ = mask;
1427            VBIT_SCRATCH.with(|scratch| {
1428                let mut buf = scratch.borrow_mut();
1429                buf.resize(cols, 0);
1430                match b {
1431                    3 => fill::<3>(data, l, &mut buf),
1432                    4 => vbit_fill4(data, &mut buf),
1433                    5 => fill::<5>(data, l, &mut buf),
1434                    6 => fill::<6>(data, l, &mut buf),
1435                    _ => unreachable!(),
1436                }
1437                let mut dot = 0f32;
1438                for g in 0..ng {
1439                    let so = (r * ng + g) * 2;
1440                    let s = f16_to_f32(u16::from_le_bytes([
1441                        bytes[sc_off + so],
1442                        bytes[sc_off + so + 1],
1443                    ]));
1444                    let d = dot_i8_i8(
1445                        &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1446                        &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1447                    ) as f32
1448                        * act.sx;
1449                    dot += d * s;
1450                }
1451                for &(j, xv) in &act.outliers {
1452                    let so = (r * ng + j / GROUP_SIZE) * 2;
1453                    let s = f16_to_f32(u16::from_le_bytes([
1454                        bytes[sc_off + so],
1455                        bytes[sc_off + so + 1],
1456                    ]));
1457                    // xq is zeroed at outlier slots — add the exact term.
1458                    dot += (buf[j] as i8) as f32 * s * xv;
1459                }
1460                dot
1461            })
1462    };
1463    for r in start..end {
1464        // SAFETY: disjoint row ranges per worker.
1465        unsafe { *out.at(r) = row_dot(r) };
1466    }
1467}
1468
1469/// Exact scalar vbit row range (same extraction, non-SDOT path).
1470#[allow(clippy::too_many_arguments)]
1471fn vbit_range_f32(
1472    bytes: &[u8],
1473    offsets: &[usize],
1474    x: &[f32],
1475    rows: usize,
1476    cols: usize,
1477    out: SendMut,
1478    start: usize,
1479    end: usize,
1480) {
1481    let ng = cols / GROUP_SIZE;
1482    let bits = &bytes[..rows];
1483    let sc_off = rows;
1484    // Per-bit-width specialized inner loops: the compiler unrolls the
1485    // constant shifts (the generic bit-buffer loop was branch-bound —
1486    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
1487    #[inline(always)]
1488    fn dot_row<const B: usize>(
1489        data: &[u8],
1490        bytes: &[u8],
1491        sc_off: usize,
1492        r: usize,
1493        ng: usize,
1494        x: &[f32],
1495    ) -> f32 {
1496        let l = ((1i32 << (B - 1)) - 1) as f32;
1497        let gbytes = GROUP_SIZE * B / 8;
1498        let mut dot = 0f32;
1499        for g in 0..ng {
1500            let so = (r * ng + g) * 2;
1501            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1502            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1503            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1504            let mut gd = 0f32;
1505            for blk in 0..GROUP_SIZE / 8 {
1506                let u = unpack8::<B>(&gd0[blk * B..]);
1507                let xb = &xg[blk * 8..blk * 8 + 8];
1508                for k in 0..8 {
1509                    gd += (u[k] as f32 - l) * xb[k];
1510                }
1511            }
1512            dot += gd * s;
1513        }
1514        dot
1515    }
1516    for r in start..end {
1517        let data = &bytes[offsets[r]..offsets[r + 1]];
1518        let v = match bits[r] {
1519            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1520            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1521            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1522            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1523            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1524            b => unreachable!("vbit bit-width {b} (validated at load)"),
1525        };
1526        // SAFETY: disjoint row ranges per worker.
1527        unsafe { *out.at(r) = v };
1528    }
1529}
1530
1531/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
1532/// and dotted against BOTH activations (MTP verify / pair prefill used
1533/// to run two full matvecs — double weight traffic and double unpack).
1534/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
1535#[allow(clippy::too_many_arguments)]
1536fn vbitmatvec2(
1537    bytes: &[u8],
1538    offsets: &[usize],
1539    x1: &[f32],
1540    x2: &[f32],
1541    rows: usize,
1542    cols: usize,
1543    o1: &mut [f32],
1544    o2: &mut [f32],
1545    pool: Option<&Pool>,
1546) {
1547    debug_assert_eq!(o1.len(), rows);
1548    debug_assert_eq!(o2.len(), rows);
1549
1550    if a8w8_enabled() {
1551        let a1 = split_act(x1);
1552        let a2 = split_act(x2);
1553        let p1 = SendMut(o1.as_mut_ptr());
1554        let p2 = SendMut(o2.as_mut_ptr());
1555        let run = move |start: usize, end: usize| {
1556            vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1557        };
1558        dispatch_rows(pool, rows, &run);
1559        return;
1560    }
1561
1562    let p1 = SendMut(o1.as_mut_ptr());
1563    let p2 = SendMut(o2.as_mut_ptr());
1564    let run = move |start: usize, end: usize| {
1565        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1566    };
1567    dispatch_rows(pool, rows, &run);
1568}
1569
1570/// Two-input vbit row range via the A8W8 int8 path — kernel body of
1571/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
1572/// exact f32 for both lanes, bits streamed once).
1573#[allow(clippy::too_many_arguments)]
1574fn vbit_range2_a8w8(
1575    bytes: &[u8],
1576    offsets: &[usize],
1577    x1: &[f32],
1578    x2: &[f32],
1579    a1: &SplitAct,
1580    a2: &SplitAct,
1581    rows: usize,
1582    cols: usize,
1583    p1: SendMut,
1584    p2: SendMut,
1585    start: usize,
1586    end: usize,
1587) {
1588    let ng = cols / GROUP_SIZE;
1589    let bits = &bytes[..rows];
1590    let sc_off = rows;
1591    let row_dots = |r: usize| -> (f32, f32) {
1592            let b = bits[r] as usize;
1593            let l = (1i32 << (b - 1)) - 1;
1594            let data = &bytes[offsets[r]..offsets[r + 1]];
1595            if b == 8 {
1596                // u−L reaches 128 → does not fit i8; exact f32 path,
1597                // bits still streamed once for both lanes.
1598                let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1599                let (mut d1, mut d2) = (0f32, 0f32);
1600                for g in 0..ng {
1601                    let so = (r * ng + g) * 2;
1602                    let sgf = f16_to_f32(u16::from_le_bytes([
1603                        bytes[sc_off + so],
1604                        bytes[sc_off + so + 1],
1605                    ]));
1606                    let (mut g1, mut g2) = (0f32, 0f32);
1607                    for k in 0..GROUP_SIZE {
1608                        if nbits < 8 {
1609                            acc = (acc << 8) | data[idx] as u64;
1610                            idx += 1;
1611                            nbits += 8;
1612                        }
1613                        let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1614                        nbits -= 8;
1615                        let w = (u - l) as f32;
1616                        g1 += w * x1[g * GROUP_SIZE + k];
1617                        g2 += w * x2[g * GROUP_SIZE + k];
1618                    }
1619                    d1 += g1 * sgf;
1620                    d2 += g2 * sgf;
1621                }
1622                return (d1, d2);
1623            }
1624            thread_local! {
1625                static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1626                    const { std::cell::RefCell::new(Vec::new()) };
1627            }
1628            #[inline(always)]
1629            fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1630                for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1631                    let u = unpack8::<B>(&data[blk * B..]);
1632                    for k in 0..8 {
1633                        chunk[k] = (u[k] - l) as i8 as u8;
1634                    }
1635                }
1636            }
1637            VBIT_SCRATCH2.with(|scratch| {
1638                let mut buf = scratch.borrow_mut();
1639                buf.resize(cols, 0);
1640                match b {
1641                    3 => fill::<3>(data, l, &mut buf),
1642                    4 => vbit_fill4(data, &mut buf),
1643                    5 => fill::<5>(data, l, &mut buf),
1644                    6 => fill::<6>(data, l, &mut buf),
1645                    _ => unreachable!(),
1646                }
1647                let (mut d1, mut d2) = (0f32, 0f32);
1648                for g in 0..ng {
1649                    let so = (r * ng + g) * 2;
1650                    let s = f16_to_f32(u16::from_le_bytes([
1651                        bytes[sc_off + so],
1652                        bytes[sc_off + so + 1],
1653                    ]));
1654                    let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1655                    let v1 =
1656                        dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1657                    let v2 =
1658                        dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1659                    d1 += v1 * s;
1660                    d2 += v2 * s;
1661                }
1662                for &(j, xv) in &a1.outliers {
1663                    let so = (r * ng + j / GROUP_SIZE) * 2;
1664                    let s = f16_to_f32(u16::from_le_bytes([
1665                        bytes[sc_off + so],
1666                        bytes[sc_off + so + 1],
1667                    ]));
1668                    d1 += (buf[j] as i8) as f32 * s * xv;
1669                }
1670                for &(j, xv) in &a2.outliers {
1671                    let so = (r * ng + j / GROUP_SIZE) * 2;
1672                    let s = f16_to_f32(u16::from_le_bytes([
1673                        bytes[sc_off + so],
1674                        bytes[sc_off + so + 1],
1675                    ]));
1676                    d2 += (buf[j] as i8) as f32 * s * xv;
1677                }
1678                (d1, d2)
1679            })
1680    };
1681    for r in start..end {
1682        let (v1, v2) = row_dots(r);
1683        // SAFETY: disjoint row ranges per worker.
1684        unsafe {
1685            *p1.at(r) = v1;
1686            *p2.at(r) = v2;
1687        }
1688    }
1689}
1690
1691/// Two-input exact scalar vbit row range (same extraction) —
1692/// per-bit-width specialized, two accumulators per row; per-lane
1693/// accumulation order matches `vbitmatvec` exactly.
1694#[allow(clippy::too_many_arguments)]
1695fn vbit_range2_f32(
1696    bytes: &[u8],
1697    offsets: &[usize],
1698    x1: &[f32],
1699    x2: &[f32],
1700    rows: usize,
1701    cols: usize,
1702    p1: SendMut,
1703    p2: SendMut,
1704    start: usize,
1705    end: usize,
1706) {
1707    let ng = cols / GROUP_SIZE;
1708    let bits = &bytes[..rows];
1709    let sc_off = rows;
1710    #[inline(always)]
1711    #[allow(clippy::too_many_arguments)]
1712    fn dot_row2<const B: usize>(
1713        data: &[u8],
1714        bytes: &[u8],
1715        sc_off: usize,
1716        r: usize,
1717        ng: usize,
1718        x1: &[f32],
1719        x2: &[f32],
1720    ) -> (f32, f32) {
1721        let l = ((1i32 << (B - 1)) - 1) as f32;
1722        let gbytes = GROUP_SIZE * B / 8;
1723        let (mut d1, mut d2) = (0f32, 0f32);
1724        for g in 0..ng {
1725            let so = (r * ng + g) * 2;
1726            let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1727            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1728            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1729            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1730            let (mut g1, mut g2) = (0f32, 0f32);
1731            for blk in 0..GROUP_SIZE / 8 {
1732                let u = unpack8::<B>(&gd0[blk * B..]);
1733                for k in 0..8 {
1734                    let w = u[k] as f32 - l;
1735                    g1 += w * x1g[blk * 8 + k];
1736                    g2 += w * x2g[blk * 8 + k];
1737                }
1738            }
1739            d1 += g1 * s;
1740            d2 += g2 * s;
1741        }
1742        (d1, d2)
1743    }
1744    for r in start..end {
1745        let data = &bytes[offsets[r]..offsets[r + 1]];
1746        let (v1, v2) = match bits[r] {
1747            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
1748            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
1749            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
1750            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
1751            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
1752            b => unreachable!("vbit bit-width {b} (validated at load)"),
1753        };
1754        // SAFETY: disjoint row ranges per worker.
1755        unsafe {
1756            *p1.at(r) = v1;
1757            *p2.at(r) = v2;
1758        }
1759    }
1760}
1761
1762// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
1763
1764/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
1765/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
1766/// distant streams of the split layout. Values/order identical to the
1767/// split kernels.
1768#[inline]
1769#[allow(unreachable_code)]
1770fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
1771    #[cfg(target_arch = "aarch64")]
1772    unsafe {
1773        return dot_q4t_row_sdot(bytes, r, gpr, xq);
1774    }
1775    #[cfg(target_arch = "x86_64")]
1776    unsafe {
1777        return dot_q4t_row_avx2(bytes, r, gpr, xq);
1778    }
1779    let mut acc = 0f32;
1780    for gi in 0..gpr {
1781        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
1782        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1783        let mut d = 0i32;
1784        for (k, &b) in tile[2..].iter().enumerate() {
1785            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
1786                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
1787        }
1788        acc += d as f32 * s;
1789    }
1790    acc
1791}
1792
1793#[cfg(target_arch = "aarch64")]
1794#[target_feature(enable = "neon,dotprod")]
1795unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
1796    // SAFETY: callers uphold slice-length contracts (18B tile per group,
1797    // xq.len() == gpr·GROUP_SIZE).
1798    unsafe {
1799        use core::arch::aarch64::*;
1800        use core::arch::asm;
1801        let lomask = vdupq_n_u8(0x0F);
1802        let eight = vdupq_n_s8(8);
1803        let mut acc = 0f32;
1804        for gi in 0..gpr {
1805            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
1806            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
1807            let b = vld1q_u8(t.add(2));
1808            let lo = vandq_u8(b, lomask);
1809            let hi = vshrq_n_u8::<4>(b);
1810            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
1811            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
1812            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
1813            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
1814            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
1815            asm!(
1816                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
1817                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
1818                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
1819                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
1820                options(pure, nomem, nostack),
1821            );
1822            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
1823        }
1824        acc
1825    }
1826}
1827
1828#[cfg(target_arch = "x86_64")]
1829#[target_feature(enable = "avx2")]
1830unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
1831    // SAFETY: see dot_q4t_row_sdot.
1832    unsafe {
1833        use core::arch::x86_64::*;
1834        let lomask = _mm_set1_epi8(0x0F);
1835        let eight = _mm256_set1_epi8(8);
1836        let ones = _mm256_set1_epi16(1);
1837        let mut acc = 0f32;
1838        for gi in 0..gpr {
1839            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
1840            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
1841            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
1842            let lo = _mm_and_si128(b, lomask);
1843            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
1844            let w = _mm256_sub_epi8(
1845                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
1846                eight,
1847            );
1848            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
1849            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
1850            let d = _mm256_madd_epi16(p16, ones);
1851            let hi128 = _mm256_extracti128_si256::<1>(d);
1852            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
1853            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
1854            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
1855            acc += _mm_cvtsi128_si32(s32) as f32 * s;
1856        }
1857        acc
1858    }
1859}
1860
1861/// Exact-term correction for A8W8 outliers on a tiled row.
1862#[inline]
1863fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
1864    let gi = j / GROUP_SIZE;
1865    let k = j % GROUP_SIZE;
1866    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
1867    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1868    let byte = tile[2 + k / 2];
1869    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
1870    ((nib as i32 - 8) as f32, s)
1871}
1872
1873/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
1874/// accumulation shape as `q4_range_f32`.
1875#[inline]
1876fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
1877    let mut acc = 0f32;
1878    for gi in 0..gpr {
1879        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
1880        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1881        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1882        let mut ga = 0f32;
1883        for (k, &b) in tile[2..].iter().enumerate() {
1884            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
1885                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
1886        }
1887        acc += ga * s;
1888    }
1889    acc
1890}
1891
1892/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
1893fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
1894    debug_assert_eq!(out.len(), rows);
1895    let gpr = cols / GROUP_SIZE;
1896    let out_addr = SendMut(out.as_mut_ptr());
1897    if a8w8_enabled() {
1898        let act = split_act(x);
1899        let run = move |start: usize, end: usize| {
1900            for r in start..end {
1901                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
1902                for &(j, xv) in &act.outliers {
1903                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
1904                    acc += w * s * xv;
1905                }
1906                // SAFETY: disjoint row ranges per worker.
1907                unsafe { *out_addr.at(r) = acc };
1908            }
1909        };
1910        dispatch_rows(pool, rows, &run);
1911        return;
1912    }
1913    let run = move |start: usize, end: usize| {
1914        for r in start..end {
1915            // SAFETY: disjoint row ranges per worker.
1916            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
1917        }
1918    };
1919    dispatch_rows(pool, rows, &run);
1920}
1921
1922/// Fused two-input q4_tiled matvec (weights read once per pair).
1923#[allow(clippy::too_many_arguments)]
1924fn q4t_matvec2(
1925    bytes: &[u8],
1926    x1: &[f32],
1927    x2: &[f32],
1928    rows: usize,
1929    cols: usize,
1930    o1: &mut [f32],
1931    o2: &mut [f32],
1932    pool: Option<&Pool>,
1933) {
1934    let gpr = cols / GROUP_SIZE;
1935    let p1 = SendMut(o1.as_mut_ptr());
1936    let p2 = SendMut(o2.as_mut_ptr());
1937    if a8w8_enabled() {
1938        let a1 = split_act(x1);
1939        let a2 = split_act(x2);
1940        let run = move |start: usize, end: usize| {
1941            for r in start..end {
1942                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
1943                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
1944                for &(j, xv) in &a1.outliers {
1945                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
1946                    v1 += w * s * xv;
1947                }
1948                for &(j, xv) in &a2.outliers {
1949                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
1950                    v2 += w * s * xv;
1951                }
1952                // SAFETY: disjoint row ranges per worker.
1953                unsafe {
1954                    *p1.at(r) = v1;
1955                    *p2.at(r) = v2;
1956                }
1957            }
1958        };
1959        dispatch_rows(pool, rows, &run);
1960        return;
1961    }
1962    let run = move |start: usize, end: usize| {
1963        for r in start..end {
1964            // SAFETY: disjoint row ranges per worker.
1965            unsafe {
1966                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
1967                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
1968            }
1969        }
1970    };
1971    dispatch_rows(pool, rows, &run);
1972}
1973
1974/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
1975#[allow(clippy::too_many_arguments)]
1976fn q4t_matmat(
1977    bytes: &[u8],
1978    xs_all: &[f32],
1979    b: usize,
1980    rows: usize,
1981    cols: usize,
1982    out: &mut [f32],
1983    pool: Option<&Pool>,
1984) {
1985    debug_assert_eq!(out.len(), b * rows);
1986    let gpr = cols / GROUP_SIZE;
1987    let out_addr = SendMut(out.as_mut_ptr());
1988    if a8w8_enabled() {
1989        let acts: Vec<SplitAct> =
1990            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
1991        let acts = &acts;
1992        let run = move |start: usize, end: usize| {
1993            for r in start..end {
1994                for (bi, act) in acts.iter().enumerate() {
1995                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
1996                    for &(j, xv) in &act.outliers {
1997                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
1998                        acc += w * s * xv;
1999                    }
2000                    // SAFETY: disjoint (bi, r) cells per worker range.
2001                    unsafe { *out_addr.at(bi * rows + r) = acc };
2002                }
2003            }
2004        };
2005        dispatch_rows(pool, rows, &run);
2006        return;
2007    }
2008    let run = move |start: usize, end: usize| {
2009        for r in start..end {
2010            for bi in 0..b {
2011                let x = &xs_all[bi * cols..(bi + 1) * cols];
2012                // SAFETY: disjoint (bi, r) cells per worker range.
2013                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2014            }
2015        }
2016    };
2017    dispatch_rows(pool, rows, &run);
2018}
2019
2020// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
2021// 32-group tile. The kernel family mirrors q4_tiled: one sequential
2022// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
2023// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
2024
2025/// Per-32-group sums of the quantized activation — the ±1 identity's
2026/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
2027/// matvec and reused by every row.
2028fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2029    (0..gpr)
2030        .map(|gi| {
2031            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2032                .iter()
2033                .map(|&v| v as i32)
2034                .sum()
2035        })
2036        .collect()
2037}
2038
2039/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
2040/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
2041/// x86 pass).
2042#[inline]
2043#[allow(unreachable_code)]
2044fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2045    #[cfg(target_arch = "aarch64")]
2046    unsafe {
2047        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2048    }
2049    let _ = gsum;
2050    let mut acc = 0f32;
2051    for gi in 0..gpr {
2052        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2053        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2054        let mut d = 0i32;
2055        for (j, &b) in tile[2..].iter().enumerate() {
2056            for k in 0..8 {
2057                let w = ((b >> k) & 1) as i32 * 2 - 1;
2058                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2059            }
2060        }
2061        acc += d as f32 * s;
2062    }
2063    acc
2064}
2065
2066/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
2067/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
2068/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
2069/// per-group activation sums shared across every row of the matvec.
2070/// Four tiles (128 weights) per iteration: integer dots reduce through
2071/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
2072/// fused f32 multiply-add. Integer math throughout — bit-identical to
2073/// the scalar ±1 reference.
2074#[cfg(target_arch = "aarch64")]
2075#[target_feature(enable = "neon,dotprod")]
2076unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2077    // SAFETY: callers uphold slice-length contracts (6B tile per group,
2078    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
2079    unsafe {
2080        use core::arch::aarch64::*;
2081        use core::arch::asm;
2082        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2083        let m = vld1q_u8(MASKS.as_ptr());
2084        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
2085        macro_rules! tile_dot {
2086            ($t:expr, $x:expr) => {{
2087                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2088                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2089                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2090                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2091                let x0 = vld1q_s8($x);
2092                let x1 = vld1q_s8($x.add(16));
2093                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2094                asm!(
2095                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2096                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2097                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2098                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2099                    options(pure, nomem, nostack),
2100                );
2101                vaddq_s32(a0, a1)
2102            }};
2103        }
2104        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2105        let xp = xq.as_ptr();
2106        let gp = gsum.as_ptr();
2107        let mut accv = vdupq_n_f32(0.0);
2108        let mut gi = 0;
2109        while gi + 4 <= gpr {
2110            let t0 = base.add(gi * Q1_TILE);
2111            let t1 = base.add((gi + 1) * Q1_TILE);
2112            let t2 = base.add((gi + 2) * Q1_TILE);
2113            let t3 = base.add((gi + 3) * Q1_TILE);
2114            let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2115            let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2116            let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2117            let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2118            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
2119            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2120            let g = vld1q_s32(gp.add(gi));
2121            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2122            let sc = [
2123                f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2124                f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2125                f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2126                f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2127            ];
2128            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2129            gi += 4;
2130        }
2131        let mut acc = vaddvq_f32(accv);
2132        while gi < gpr {
2133            let t = base.add(gi * Q1_TILE);
2134            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2135            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2136            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2137            gi += 1;
2138        }
2139        acc
2140    }
2141}
2142
2143/// (weight ±1, scale) of one q1 element — the exact outlier term.
2144#[inline]
2145fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2146    let gi = j / GROUP_SIZE;
2147    let k = j % GROUP_SIZE;
2148    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2149    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2150    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2151    ((bit as i32 * 2 - 1) as f32, s)
2152}
2153
2154/// Exact scalar q1 row (CMF_SDOT=0 contract).
2155#[inline]
2156fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2157    let mut acc = 0f32;
2158    for gi in 0..gpr {
2159        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2160        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2161        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2162        let mut ga = 0f32;
2163        for (j, &b) in tile[2..].iter().enumerate() {
2164            for k in 0..8 {
2165                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2166            }
2167        }
2168        acc += ga * s;
2169    }
2170    acc
2171}
2172
2173/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
2174/// extracted so multi-matrix jobs drive the same kernel).
2175#[allow(clippy::too_many_arguments)]
2176fn q1_range_a8w8(
2177    bytes: &[u8],
2178    gpr: usize,
2179    act: &SplitAct,
2180    gsum: &[i32],
2181    out: SendMut,
2182    start: usize,
2183    end: usize,
2184) {
2185    for r in start..end {
2186        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2187        for &(j, xv) in &act.outliers {
2188            let (w, s) = q1_outlier(bytes, r, gpr, j);
2189            acc += w * s * xv;
2190        }
2191        // SAFETY: disjoint row ranges per worker.
2192        unsafe { *out.at(r) = acc };
2193    }
2194}
2195
2196/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
2197fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2198    for r in start..end {
2199        // SAFETY: disjoint row ranges per worker.
2200        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2201    }
2202}
2203
2204/// Fused q1 matvec (dispatch mirrors `q4t_matvec`).
2205fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2206    debug_assert_eq!(out.len(), rows);
2207    let gpr = cols / GROUP_SIZE;
2208    let out_addr = SendMut(out.as_mut_ptr());
2209    if a8w8_enabled() {
2210        let act = split_act(x);
2211        let gsum = q1_group_sums(&act.xq, gpr);
2212        let (act, gsum) = (&act, &gsum);
2213        let run =
2214            move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2215        dispatch_rows(pool, rows, &run);
2216        return;
2217    }
2218    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2219    dispatch_rows(pool, rows, &run);
2220}
2221
2222/// Fused two-input q1 matvec (weights read once per pair).
2223#[allow(clippy::too_many_arguments)]
2224fn q1_matvec2(
2225    bytes: &[u8],
2226    x1: &[f32],
2227    x2: &[f32],
2228    rows: usize,
2229    cols: usize,
2230    o1: &mut [f32],
2231    o2: &mut [f32],
2232    pool: Option<&Pool>,
2233) {
2234    let gpr = cols / GROUP_SIZE;
2235    let p1 = SendMut(o1.as_mut_ptr());
2236    let p2 = SendMut(o2.as_mut_ptr());
2237    if a8w8_enabled() {
2238        let a1 = split_act(x1);
2239        let a2 = split_act(x2);
2240        let g1 = q1_group_sums(&a1.xq, gpr);
2241        let g2 = q1_group_sums(&a2.xq, gpr);
2242        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2243        let run = move |start: usize, end: usize| {
2244            for r in start..end {
2245                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2246                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2247                for &(j, xv) in &a1.outliers {
2248                    let (w, s) = q1_outlier(bytes, r, gpr, j);
2249                    v1 += w * s * xv;
2250                }
2251                for &(j, xv) in &a2.outliers {
2252                    let (w, s) = q1_outlier(bytes, r, gpr, j);
2253                    v2 += w * s * xv;
2254                }
2255                // SAFETY: disjoint row ranges per worker.
2256                unsafe {
2257                    *p1.at(r) = v1;
2258                    *p2.at(r) = v2;
2259                }
2260            }
2261        };
2262        dispatch_rows(pool, rows, &run);
2263        return;
2264    }
2265    let run = move |start: usize, end: usize| {
2266        for r in start..end {
2267            // SAFETY: disjoint row ranges per worker.
2268            unsafe {
2269                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2270                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2271            }
2272        }
2273    };
2274    dispatch_rows(pool, rows, &run);
2275}
2276
2277/// Batched q1 matmat: each row's tiles stream once per microbatch.
2278#[allow(clippy::too_many_arguments)]
2279fn q1_matmat(
2280    bytes: &[u8],
2281    xs_all: &[f32],
2282    b: usize,
2283    rows: usize,
2284    cols: usize,
2285    out: &mut [f32],
2286    pool: Option<&Pool>,
2287) {
2288    debug_assert_eq!(out.len(), b * rows);
2289    let gpr = cols / GROUP_SIZE;
2290    let out_addr = SendMut(out.as_mut_ptr());
2291    if a8w8_enabled() {
2292        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2293            .map(|bi| {
2294                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2295                let gsum = q1_group_sums(&act.xq, gpr);
2296                (act, gsum)
2297            })
2298            .collect();
2299        let acts = &acts;
2300        let run = move |start: usize, end: usize| {
2301            for r in start..end {
2302                for (bi, (act, gsum)) in acts.iter().enumerate() {
2303                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2304                    for &(j, xv) in &act.outliers {
2305                        let (w, s) = q1_outlier(bytes, r, gpr, j);
2306                        acc += w * s * xv;
2307                    }
2308                    // SAFETY: disjoint (bi, r) cells per worker range.
2309                    unsafe { *out_addr.at(bi * rows + r) = acc };
2310                }
2311            }
2312        };
2313        dispatch_rows(pool, rows, &run);
2314        return;
2315    }
2316    let run = move |start: usize, end: usize| {
2317        for r in start..end {
2318            for bi in 0..b {
2319                let x = &xs_all[bi * cols..(bi + 1) * cols];
2320                // SAFETY: disjoint (bi, r) cells per worker range.
2321                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
2322            }
2323        }
2324    };
2325    dispatch_rows(pool, rows, &run);
2326}
2327
2328/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
2329/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
2330/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
2331/// 32-group, exact outlier correction — the same A8W8 contract as q8.
2332/// `CMF_SDOT=0` keeps the exact scalar path.
2333fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2334    debug_assert_eq!(out.len(), rows);
2335    let (packed, scales) = q4_split(bytes, rows, cols);
2336    let gpr = cols / GROUP_SIZE;
2337    let out_addr = SendMut(out.as_mut_ptr());
2338
2339    if a8w8_enabled() {
2340        let act = split_act(x);
2341        let run = move |start: usize, end: usize| {
2342            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
2343        };
2344        dispatch_rows(pool, rows, &run);
2345        return;
2346    }
2347
2348    let run = move |start: usize, end: usize| {
2349        q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
2350    };
2351    dispatch_rows(pool, rows, &run);
2352}
2353
2354/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
2355/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
2356#[inline]
2357#[allow(unreachable_code)]
2358fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2359    #[cfg(target_arch = "aarch64")]
2360    unsafe {
2361        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
2362    }
2363    #[cfg(target_arch = "x86_64")]
2364    unsafe {
2365        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
2366    }
2367    let mut acc = 0f32;
2368    for gi in 0..gpr {
2369        let g = g0 + gi;
2370        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2371        let mut d = 0i32;
2372        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2373            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2374                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2375        }
2376        acc += d as f32 * s;
2377    }
2378    acc
2379}
2380
2381/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
2382#[inline]
2383#[allow(unreachable_code)]
2384fn dot_q4_row_i8_2(
2385    packed: &[u8],
2386    scales: &[u8],
2387    g0: usize,
2388    gpr: usize,
2389    xq1: &[i8],
2390    xq2: &[i8],
2391) -> (f32, f32) {
2392    #[cfg(target_arch = "aarch64")]
2393    unsafe {
2394        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
2395    }
2396    #[cfg(target_arch = "x86_64")]
2397    unsafe {
2398        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
2399    }
2400    (
2401        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
2402        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
2403    )
2404}
2405
2406/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
2407/// multi-matrix jobs can drive it for several tensors in one dispatch).
2408#[allow(clippy::too_many_arguments)]
2409fn q4_range_a8w8(
2410    packed: &[u8],
2411    scales: &[u8],
2412    gpr: usize,
2413    cols: usize,
2414    act: &SplitAct,
2415    out: SendMut,
2416    start: usize,
2417    end: usize,
2418) {
2419    for r in start..end {
2420        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
2421        // xq is zeroed at outlier slots — add the exact terms.
2422        for &(j, xv) in &act.outliers {
2423            let flat = r * cols + j;
2424            let byte = packed[flat / 2];
2425            let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2426            let s = f16_to_f32(u16::from_le_bytes([
2427                scales[(flat / GROUP_SIZE) * 2],
2428                scales[(flat / GROUP_SIZE) * 2 + 1],
2429            ]));
2430            acc += ((nib as i32 - 8) as f32) * s * xv;
2431        }
2432        // SAFETY: disjoint row ranges per worker.
2433        unsafe { *out.at(r) = acc };
2434    }
2435}
2436
2437/// Two-input q4 row range via the A8W8 int8 path — kernel body of
2438/// `q4matvec2`, extracted for pair multi-matrix jobs.
2439#[allow(clippy::too_many_arguments)]
2440fn q4_range2_a8w8(
2441    packed: &[u8],
2442    scales: &[u8],
2443    gpr: usize,
2444    cols: usize,
2445    a1: &SplitAct,
2446    a2: &SplitAct,
2447    p1: SendMut,
2448    p2: SendMut,
2449    start: usize,
2450    end: usize,
2451) {
2452    for r in start..end {
2453        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
2454        let mut acc1 = s1 * a1.sx;
2455        let mut acc2 = s2 * a2.sx;
2456        // xq is zeroed at outlier slots — add the exact terms.
2457        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
2458            for &(j, xv) in outliers {
2459                let flat = r * cols + j;
2460                let byte = packed[flat / 2];
2461                let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2462                let s = f16_to_f32(u16::from_le_bytes([
2463                    scales[(flat / GROUP_SIZE) * 2],
2464                    scales[(flat / GROUP_SIZE) * 2 + 1],
2465                ]));
2466                *acc += ((nib as i32 - 8) as f32) * s * xv;
2467            }
2468        };
2469        fix(&a1.outliers, &mut acc1);
2470        fix(&a2.outliers, &mut acc2);
2471        // SAFETY: disjoint row ranges per worker.
2472        unsafe {
2473            *p1.at(r) = acc1;
2474            *p2.at(r) = acc2;
2475        }
2476    }
2477}
2478
2479/// Exact scalar q4 row range (same extraction, non-SDOT path).
2480fn q4_range_f32(
2481    packed: &[u8],
2482    scales: &[u8],
2483    gpr: usize,
2484    x: &[f32],
2485    out: SendMut,
2486    start: usize,
2487    end: usize,
2488) {
2489    for r in start..end {
2490        let mut acc = 0f32;
2491        for gi in 0..gpr {
2492            let g = r * gpr + gi;
2493            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2494            let pk = &packed[g * 16..(g + 1) * 16];
2495            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2496            let mut ga = 0f32;
2497            for (k, &b) in pk.iter().enumerate() {
2498                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2499                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2500            }
2501            acc += ga * s;
2502        }
2503        // SAFETY: disjoint row ranges per worker.
2504        unsafe { *out.at(r) = acc };
2505    }
2506}
2507
2508/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
2509/// dotted against both activations (was: two full matvecs — double
2510/// weight traffic). Per-lane math matches `q4matvec` exactly.
2511#[allow(clippy::too_many_arguments)]
2512fn q4matvec2(
2513    bytes: &[u8],
2514    x1: &[f32],
2515    x2: &[f32],
2516    rows: usize,
2517    cols: usize,
2518    o1: &mut [f32],
2519    o2: &mut [f32],
2520    pool: Option<&Pool>,
2521) {
2522    debug_assert_eq!(o1.len(), rows);
2523    debug_assert_eq!(o2.len(), rows);
2524    let (packed, scales) = q4_split(bytes, rows, cols);
2525    let gpr = cols / GROUP_SIZE;
2526
2527    if a8w8_enabled() {
2528        let a1 = split_act(x1);
2529        let a2 = split_act(x2);
2530        let p1 = SendMut(o1.as_mut_ptr());
2531        let p2 = SendMut(o2.as_mut_ptr());
2532        let run = move |start: usize, end: usize| {
2533            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
2534        };
2535        dispatch_rows(pool, rows, &run);
2536        return;
2537    }
2538
2539    let p1 = SendMut(o1.as_mut_ptr());
2540    let p2 = SendMut(o2.as_mut_ptr());
2541    let run = move |start: usize, end: usize| {
2542        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
2543    };
2544    dispatch_rows(pool, rows, &run);
2545}
2546
2547/// Two-input exact scalar q4 row range (same extraction).
2548#[allow(clippy::too_many_arguments)]
2549fn q4_range2_f32(
2550    packed: &[u8],
2551    scales: &[u8],
2552    gpr: usize,
2553    x1: &[f32],
2554    x2: &[f32],
2555    p1: SendMut,
2556    p2: SendMut,
2557    start: usize,
2558    end: usize,
2559) {
2560    for r in start..end {
2561        let (mut acc1, mut acc2) = (0f32, 0f32);
2562        for gi in 0..gpr {
2563            let g = r * gpr + gi;
2564            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2565            let pk = &packed[g * 16..(g + 1) * 16];
2566            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2567            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2568            let (mut g1, mut g2) = (0f32, 0f32);
2569            for (k, &b) in pk.iter().enumerate() {
2570                let wl = (b & 0x0F) as f32 - 8.0;
2571                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
2572                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
2573                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
2574            }
2575            acc1 += g1 * s;
2576            acc2 += g2 * s;
2577        }
2578        // SAFETY: disjoint row ranges per worker.
2579        unsafe {
2580            *p1.at(r) = acc1;
2581            *p2.at(r) = acc2;
2582        }
2583    }
2584}
2585
2586thread_local! {
2587    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
2588    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
2589    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
2590    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2591}
2592
2593/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
2594/// and dotted against ALL b activations (prefill used to fall back to b
2595/// full matvecs — b× weight traffic and b× nibble decode). Per-position
2596/// math matches `q4matvec` exactly: same group order, same accumulation.
2597/// `out` is row-major [b, rows] like `qmatmat`.
2598#[allow(clippy::too_many_arguments)]
2599fn q4matmat(
2600    bytes: &[u8],
2601    xs_all: &[f32],
2602    b: usize,
2603    rows: usize,
2604    cols: usize,
2605    out: &mut [f32],
2606    pool: Option<&Pool>,
2607) {
2608    debug_assert_eq!(xs_all.len(), b * cols);
2609    debug_assert_eq!(out.len(), b * rows);
2610    let (packed, scales) = q4_split(bytes, rows, cols);
2611    let gpr = cols / GROUP_SIZE;
2612    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2613
2614    if a8w8_enabled() {
2615        let acts: Vec<SplitAct> =
2616            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2617        let acts = &acts;
2618        let out_addr = SendMut(out.as_mut_ptr());
2619        let run = move |start: usize, end: usize| {
2620            ROW_I8.with(|rb| {
2621                let mut buf = rb.borrow_mut();
2622                buf.resize(cols, 0);
2623                for r in start..end {
2624                    // Unpack the row's nibbles to centered i8 once
2625                    // (element 2k = low nibble, 2k+1 = high — flat order,
2626                    // same as dot_q4_row_sdot's zip).
2627                    for gi in 0..gpr {
2628                        let g = r * gpr + gi;
2629                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2630                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
2631                            buf[gi * GROUP_SIZE + k * 2 + 1] =
2632                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
2633                        }
2634                    }
2635                    for (bi, act) in acts.iter().enumerate() {
2636                        let mut acc = 0f32;
2637                        for gi in 0..gpr {
2638                            let d = dot_i8_i8(
2639                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
2640                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
2641                            );
2642                            acc += d as f32 * gscale(r * gpr + gi);
2643                        }
2644                        acc *= act.sx;
2645                        // xq is zeroed at outlier slots — exact terms.
2646                        for &(j, xv) in &act.outliers {
2647                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
2648                        }
2649                        // SAFETY: disjoint (bi, r) cells per worker row range.
2650                        unsafe { *out_addr.at(bi * rows + r) = acc };
2651                    }
2652                }
2653            })
2654        };
2655        dispatch_rows(pool, rows, &run);
2656        return;
2657    }
2658
2659    let out_addr = SendMut(out.as_mut_ptr());
2660    let run = move |start: usize, end: usize| {
2661        ROW_F32.with(|rb| {
2662            let mut buf = rb.borrow_mut();
2663            buf.resize(cols, 0.0);
2664            for r in start..end {
2665                // Decode raw (nib − 8) values once; scales stay per-group
2666                // so the accumulation order matches q4matvec bit-for-bit.
2667                for gi in 0..gpr {
2668                    let g = r * gpr + gi;
2669                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2670                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
2671                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
2672                    }
2673                }
2674                for bi in 0..b {
2675                    let x = &xs_all[bi * cols..(bi + 1) * cols];
2676                    let mut acc = 0f32;
2677                    for gi in 0..gpr {
2678                        let mut ga = 0f32;
2679                        // Pairwise (lo + hi) addition, matching
2680                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
2681                        // a flat one-per-element loop rounds differently
2682                        // and broke bit-parity on the scalar (x86) path.
2683                        for k in 0..GROUP_SIZE / 2 {
2684                            let e = gi * GROUP_SIZE + k * 2;
2685                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
2686                        }
2687                        acc += ga * gscale(r * gpr + gi);
2688                    }
2689                    // SAFETY: disjoint (bi, r) cells per worker row range.
2690                    unsafe { *out_addr.at(bi * rows + r) = acc };
2691                }
2692            }
2693        })
2694    };
2695    dispatch_rows(pool, rows, &run);
2696}
2697
2698/// Batched vbit matmat: each variable-bit row is decoded from the mmap
2699/// ONCE for the whole microbatch. Same per-position math as
2700/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
2701/// and the scalar path).
2702#[allow(clippy::too_many_arguments)]
2703fn vbitmatmat(
2704    bytes: &[u8],
2705    offsets: &[usize],
2706    xs_all: &[f32],
2707    b: usize,
2708    rows: usize,
2709    cols: usize,
2710    out: &mut [f32],
2711    pool: Option<&Pool>,
2712) {
2713    debug_assert_eq!(xs_all.len(), b * cols);
2714    debug_assert_eq!(out.len(), b * rows);
2715    debug_assert_eq!(offsets.len(), rows + 1);
2716    let ng = cols / GROUP_SIZE;
2717    let bits = &bytes[..rows];
2718    let sc_off = rows;
2719    let gscale = |r: usize, g: usize| {
2720        let so = (r * ng + g) * 2;
2721        f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
2722    };
2723
2724    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
2725    let decode_f32 = |r: usize, dst: &mut [f32]| {
2726        let bw = bits[r] as usize;
2727        let l = ((1i32 << (bw - 1)) - 1) as f32;
2728        let data = &bytes[offsets[r]..offsets[r + 1]];
2729        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2730        for d in dst.iter_mut() {
2731            while nbits < bw {
2732                acc = (acc << 8) | data[idx] as u64;
2733                idx += 1;
2734                nbits += 8;
2735            }
2736            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
2737            nbits -= bw;
2738            *d = u - l;
2739        }
2740    };
2741
2742    if a8w8_enabled() {
2743        let acts: Vec<SplitAct> =
2744            (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2745        let acts = &acts;
2746        let out_addr = SendMut(out.as_mut_ptr());
2747        let run = move |start: usize, end: usize| {
2748            for r in start..end {
2749                let bw = bits[r] as usize;
2750                if bw == 8 {
2751                    // u−L reaches 128 → no i8 path; decode once, exact
2752                    // f32 dots for every position (same as vbitmatvec).
2753                    ROW_F32.with(|rb| {
2754                        let mut buf = rb.borrow_mut();
2755                        buf.resize(cols, 0.0);
2756                        decode_f32(r, &mut buf);
2757                        for bi in 0..b {
2758                            let x = &xs_all[bi * cols..(bi + 1) * cols];
2759                            let mut dot = 0f32;
2760                            for g in 0..ng {
2761                                let mut gd = 0f32;
2762                                for k in 0..GROUP_SIZE {
2763                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
2764                                }
2765                                dot += gd * gscale(r, g);
2766                            }
2767                            // SAFETY: disjoint (bi, r) cells per worker range.
2768                            unsafe { *out_addr.at(bi * rows + r) = dot };
2769                        }
2770                    });
2771                    continue;
2772                }
2773                let l = (1i32 << (bw - 1)) - 1;
2774                let data = &bytes[offsets[r]..offsets[r + 1]];
2775                ROW_I8.with(|rb| {
2776                    let mut buf = rb.borrow_mut();
2777                    buf.resize(cols, 0);
2778                    #[inline(always)]
2779                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2780                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2781                            let u = unpack8::<B>(&data[blk * B..]);
2782                            for k in 0..8 {
2783                                chunk[k] = (u[k] - l) as i8 as u8;
2784                            }
2785                        }
2786                    }
2787                    match bw {
2788                        3 => fill::<3>(data, l, &mut buf),
2789                        4 => vbit_fill4(data, &mut buf),
2790                        5 => fill::<5>(data, l, &mut buf),
2791                        6 => fill::<6>(data, l, &mut buf),
2792                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
2793                    }
2794                    for (bi, act) in acts.iter().enumerate() {
2795                        let mut dot = 0f32;
2796                        for g in 0..ng {
2797                            let d = dot_i8_i8(
2798                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2799                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2800                            ) as f32
2801                                * act.sx;
2802                            dot += d * gscale(r, g);
2803                        }
2804                        for &(j, xv) in &act.outliers {
2805                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
2806                        }
2807                        // SAFETY: disjoint (bi, r) cells per worker range.
2808                        unsafe { *out_addr.at(bi * rows + r) = dot };
2809                    }
2810                });
2811            }
2812        };
2813        dispatch_rows(pool, rows, &run);
2814        return;
2815    }
2816
2817    let out_addr = SendMut(out.as_mut_ptr());
2818    let run = move |start: usize, end: usize| {
2819        ROW_F32.with(|rb| {
2820            let mut buf = rb.borrow_mut();
2821            buf.resize(cols, 0.0);
2822            for r in start..end {
2823                decode_f32(r, &mut buf);
2824                for bi in 0..b {
2825                    let x = &xs_all[bi * cols..(bi + 1) * cols];
2826                    let mut dot = 0f32;
2827                    for g in 0..ng {
2828                        let mut gd = 0f32;
2829                        for k in 0..GROUP_SIZE {
2830                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
2831                        }
2832                        dot += gd * gscale(r, g);
2833                    }
2834                    // SAFETY: disjoint (bi, r) cells per worker range.
2835                    unsafe { *out_addr.at(bi * rows + r) = dot };
2836                }
2837            }
2838        })
2839    };
2840    dispatch_rows(pool, rows, &run);
2841}
2842
2843/// Build a GPU batch job for a q8-family mapped tensor (primary
2844/// shard): prescaled input + directory coordinates. None → not
2845/// GPU-eligible, caller stays on the CPU.
2846pub(crate) fn gpu_batch_job<'a>(
2847    t: &'a QTensor,
2848    x: &[f32],
2849) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
2850    match t {
2851        QTensor::Mapped {
2852            model,
2853            idx,
2854            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
2855            rows,
2856            cols,
2857            row_scale,
2858            col_field,
2859            ..
2860        } => Some((
2861            model.clone(),
2862            crate::gpu::BatchJob {
2863                idx: *idx,
2864                rows: *rows,
2865                cols: *cols,
2866                row_scale,
2867                xs: prescale(x, col_field, *dt).into_owned(),
2868            },
2869        )),
2870        _ => None,
2871    }
2872}
2873
2874/// θ col-field fold for q8_2f activations. Borrowed pass-through for
2875/// every other dtype — the old unconditional `x.to_vec()` was a pure
2876/// per-matvec allocation on the q8_row hot path.
2877pub(crate) fn prescale<'a>(
2878    x: &'a [f32],
2879    col_field: &[f32],
2880    dtype: TensorDtype,
2881) -> std::borrow::Cow<'a, [f32]> {
2882    if dtype == TensorDtype::Q8_2f {
2883        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
2884    } else {
2885        std::borrow::Cow::Borrowed(x)
2886    }
2887}
2888
2889// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
2890
2891/// AVX2+FMA available? Default ON when the CPU supports both;
2892/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
2893#[cfg(target_arch = "x86_64")]
2894pub(crate) fn avx2_enabled() -> bool {
2895    use std::sync::OnceLock;
2896    static ON: OnceLock<bool> = OnceLock::new();
2897    *ON.get_or_init(|| {
2898        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
2899            && std::arch::is_x86_feature_detected!("avx2")
2900            && std::arch::is_x86_feature_detected!("fma")
2901    })
2902}
2903
2904/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
2905/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
2906/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
2907/// active either way, they are exact (regrouped sums only).
2908#[cfg(target_arch = "x86_64")]
2909fn avx2_a8w8_enabled() -> bool {
2910    use std::sync::OnceLock;
2911    static ON: OnceLock<bool> = OnceLock::new();
2912    *ON.get_or_init(|| {
2913        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
2914    })
2915}
2916
2917/// A8W8 quantized-activation path available on THIS machine? One
2918/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
2919/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
2920#[inline]
2921pub(crate) fn a8w8_enabled() -> bool {
2922    #[cfg(target_arch = "aarch64")]
2923    {
2924        sdot_enabled()
2925    }
2926    #[cfg(target_arch = "x86_64")]
2927    {
2928        avx2_a8w8_enabled()
2929    }
2930    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
2931    {
2932        false
2933    }
2934}
2935
2936/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
2937/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
2938#[inline]
2939#[allow(unreachable_code)]
2940fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
2941    #[cfg(target_arch = "aarch64")]
2942    unsafe {
2943        return dot_i8_sdot(w, xq);
2944    }
2945    #[cfg(target_arch = "x86_64")]
2946    unsafe {
2947        if avx512vnni_enabled() {
2948            return dot_i8_i8_vnni(w, xq);
2949        }
2950        return dot_i8_i8_avx2(w, xq);
2951    }
2952    w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
2953}
2954
2955/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
2956/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
2957/// `vpdpbusd` encoding.
2958#[cfg(target_arch = "x86_64")]
2959fn avx512vnni_enabled() -> bool {
2960    use std::sync::OnceLock;
2961    static ON: OnceLock<bool> = OnceLock::new();
2962    *ON.get_or_init(|| {
2963        std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
2964            && std::arch::is_x86_feature_detected!("avx512f")
2965            && std::arch::is_x86_feature_detected!("avx512bw")
2966            && std::arch::is_x86_feature_detected!("avx512vl")
2967            && std::arch::is_x86_feature_detected!("avx512vnni")
2968    })
2969}
2970
2971/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
2972/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
2973/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
2974/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
2975#[cfg(target_arch = "x86_64")]
2976#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
2977unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
2978    // SAFETY: callers uphold slice-length contracts (see call sites).
2979    unsafe {
2980        use core::arch::x86_64::*;
2981        let n = w.len();
2982        let mut j = 0usize;
2983        let mut total: i32;
2984        // 4 independent accumulators: vpdpbusd is its own loop-carried
2985        // dependency (~5-cycle latency) — a single-acc loop runs
2986        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
2987        // on Granite Rapids.
2988        {
2989            #[inline(always)]
2990            unsafe fn step(
2991                w: *const u8,
2992                x: *const i8,
2993                acc: core::arch::x86_64::__m512i,
2994            ) -> core::arch::x86_64::__m512i {
2995                unsafe {
2996                    use core::arch::x86_64::*;
2997                    let wv = _mm512_loadu_si512(w as *const _);
2998                    let xv = _mm512_loadu_si512(x as *const _);
2999                    let aw = _mm512_abs_epi8(wv);
3000                    let neg = _mm512_movepi8_mask(wv);
3001                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3002                    _mm512_dpbusd_epi32(acc, aw, sx)
3003                }
3004            }
3005            let (mut a0, mut a1, mut a2, mut a3) = (
3006                _mm512_setzero_si512(),
3007                _mm512_setzero_si512(),
3008                _mm512_setzero_si512(),
3009                _mm512_setzero_si512(),
3010            );
3011            while j + 256 <= n {
3012                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3013                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3014                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3015                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3016                j += 256;
3017            }
3018            while j + 64 <= n {
3019                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3020                j += 64;
3021            }
3022            let s01 = _mm512_add_epi32(a0, a1);
3023            let s23 = _mm512_add_epi32(a2, a3);
3024            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3025        }
3026        // 32-wide (q4/vbit groups are exactly 32 bytes).
3027        if j + 32 <= n {
3028            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3029            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3030            let d = _mm256_dpbusd_epi32(
3031                _mm256_setzero_si256(),
3032                _mm256_abs_epi8(wv),
3033                _mm256_sign_epi8(xv, wv),
3034            );
3035            let hi128 = _mm256_extracti128_si256::<1>(d);
3036            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3037            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3038            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3039            total += _mm_cvtsi128_si32(s32);
3040            j += 32;
3041        }
3042        while j < n {
3043            total += (w[j] as i8) as i32 * xq[j] as i32;
3044            j += 1;
3045        }
3046        total
3047    }
3048}
3049
3050/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
3051#[cfg(target_arch = "x86_64")]
3052#[target_feature(enable = "avx2,fma")]
3053unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3054    // SAFETY: callers uphold slice-length contracts (see call sites).
3055    unsafe {
3056        use core::arch::x86_64::*;
3057        let n = x.len();
3058        let wp = w.as_ptr();
3059        let xp = x.as_ptr();
3060        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3061        let mut j = 0usize;
3062        while j + 16 <= n {
3063            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3064            let lo = _mm256_cvtepi8_epi32(wb);
3065            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3066            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3067            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3068            j += 16;
3069        }
3070        let acc = _mm256_add_ps(a0, a1);
3071        let hi128 = _mm256_extractf128_ps::<1>(acc);
3072        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3073        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3074        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3075        let mut sum = _mm_cvtss_f32(s32);
3076        while j < n {
3077            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3078            j += 1;
3079        }
3080        sum
3081    }
3082}
3083
3084/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
3085/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
3086/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
3087/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
3088#[cfg(target_arch = "x86_64")]
3089#[target_feature(enable = "avx2")]
3090unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3091    // SAFETY: callers uphold slice-length contracts (see call sites).
3092    unsafe {
3093        use core::arch::x86_64::*;
3094        let n = w.len();
3095        let ones = _mm256_set1_epi16(1);
3096        let mut acc = _mm256_setzero_si256();
3097        let mut j = 0usize;
3098        while j + 32 <= n {
3099            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3100            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3101            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3102            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3103            j += 32;
3104        }
3105        let hi128 = _mm256_extracti128_si256::<1>(acc);
3106        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3107        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3108        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3109        let mut s = _mm_cvtsi128_si32(s32);
3110        while j < n {
3111            s += (w[j] as i8) as i32 * xq[j] as i32;
3112            j += 1;
3113        }
3114        s
3115    }
3116}
3117
3118/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
3119/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
3120/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
3121/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
3122#[cfg(target_arch = "x86_64")]
3123#[inline]
3124fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
3125    let dot = if avx512vnni_enabled() && row.len() >= 64 {
3126        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
3127    } else {
3128        unsafe { dot_i8_i8_avx2(row, &act.xq) }
3129    };
3130    let mut acc = dot as f32 * act.sx;
3131    for &(j, xv) in &act.outliers {
3132        acc += (row[j] as i8) as f32 * xv;
3133    }
3134    acc
3135}
3136
3137/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
3138/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
3139/// a single-acc loop runs latency-bound, measured on Granite Rapids).
3140#[cfg(target_arch = "x86_64")]
3141#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3142unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3143    // SAFETY: callers uphold slice-length contracts (see call sites).
3144    unsafe {
3145        use core::arch::x86_64::*;
3146        let n = w.len();
3147        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
3148        #[inline(always)]
3149        unsafe fn step(
3150            w: *const u8,
3151            x: *const i8,
3152            flip: core::arch::x86_64::__m512i,
3153            acc: core::arch::x86_64::__m512i,
3154        ) -> core::arch::x86_64::__m512i {
3155            unsafe {
3156                use core::arch::x86_64::*;
3157                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
3158                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
3159            }
3160        }
3161        let (mut a0, mut a1, mut a2, mut a3) = (
3162            _mm512_setzero_si512(),
3163            _mm512_setzero_si512(),
3164            _mm512_setzero_si512(),
3165            _mm512_setzero_si512(),
3166        );
3167        let mut j = 0usize;
3168        while j + 256 <= n {
3169            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3170            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
3171            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
3172            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
3173            j += 256;
3174        }
3175        while j + 64 <= n {
3176            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3177            j += 64;
3178        }
3179        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
3180            _mm512_add_epi32(a0, a1),
3181            _mm512_add_epi32(a2, a3),
3182        ));
3183        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
3184        while j < n {
3185            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
3186            j += 1;
3187        }
3188        total
3189    }
3190}
3191
3192/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
3193/// writer's flat order, same as the NEON vzip pair), maddubs against
3194/// the pre-quantized activation group, × the group's f16 scale. Pair
3195/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
3196/// `dot_q4_row_sdot`.
3197#[cfg(target_arch = "x86_64")]
3198#[target_feature(enable = "avx2")]
3199unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3200    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
3201    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
3202    unsafe {
3203        use core::arch::x86_64::*;
3204        let lomask = _mm_set1_epi8(0x0F);
3205        let eight = _mm256_set1_epi8(8);
3206        let ones = _mm256_set1_epi16(1);
3207        let mut acc = 0f32;
3208        for gi in 0..gpr {
3209            let g = g0 + gi;
3210            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3211            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3212            let lo = _mm_and_si128(b, lomask);
3213            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3214            let w = _mm256_sub_epi8(
3215                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3216                eight,
3217            );
3218            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3219            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3220            let d = _mm256_madd_epi16(p16, ones);
3221            let hi128 = _mm256_extracti128_si256::<1>(d);
3222            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3223            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3224            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3225            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3226        }
3227        acc
3228    }
3229}
3230
3231/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
3232/// both activations dotted against the same centered i8 register.
3233#[cfg(target_arch = "x86_64")]
3234#[target_feature(enable = "avx2")]
3235unsafe fn dot_q4_row_avx2_2(
3236    packed: &[u8],
3237    scales: &[u8],
3238    g0: usize,
3239    gpr: usize,
3240    xq1: &[i8],
3241    xq2: &[i8],
3242) -> (f32, f32) {
3243    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
3244    unsafe {
3245        use core::arch::x86_64::*;
3246        let lomask = _mm_set1_epi8(0x0F);
3247        let eight = _mm256_set1_epi8(8);
3248        let ones = _mm256_set1_epi16(1);
3249        let (mut acc1, mut acc2) = (0f32, 0f32);
3250        #[inline(always)]
3251        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
3252            unsafe {
3253                use core::arch::x86_64::*;
3254                let hi128 = _mm256_extracti128_si256::<1>(d);
3255                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3256                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3257                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3258                _mm_cvtsi128_si32(s32)
3259            }
3260        }
3261        for gi in 0..gpr {
3262            let g = g0 + gi;
3263            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3264            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3265            let lo = _mm_and_si128(b, lomask);
3266            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3267            let w = _mm256_sub_epi8(
3268                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3269                eight,
3270            );
3271            let aw = _mm256_abs_epi8(w);
3272            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3273            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3274            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
3275            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
3276            acc1 += hsum(d1) as f32 * s;
3277            acc2 += hsum(d2) as f32 * s;
3278        }
3279        (acc1, acc2)
3280    }
3281}
3282
3283/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
3284#[cfg(target_arch = "x86_64")]
3285fn q8_range_avx2(
3286    q: &[u8],
3287    row_scale: &[f32],
3288    act: &SplitAct,
3289    cols: usize,
3290    out_addr: SendMut,
3291    start: usize,
3292    end: usize,
3293) {
3294    for o in start..end {
3295        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3296        // SAFETY: disjoint row ranges per worker.
3297        unsafe { *out_addr.at(o) = v };
3298    }
3299}
3300
3301/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
3302#[cfg(target_arch = "x86_64")]
3303#[allow(clippy::too_many_arguments)]
3304fn q8_range2_avx2(
3305    q: &[u8],
3306    row_scale: &[f32],
3307    a1: &SplitAct,
3308    a2: &SplitAct,
3309    cols: usize,
3310    p1: SendMut,
3311    p2: SendMut,
3312    start: usize,
3313    end: usize,
3314) {
3315    for o in start..end {
3316        let row = &q[o * cols..(o + 1) * cols];
3317        // SAFETY: disjoint row ranges per worker.
3318        unsafe {
3319            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
3320            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
3321        }
3322    }
3323}
3324
3325// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
3326
3327/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
3328/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
3329/// (On non-ARM release builds only the test tolerance switch calls it.)
3330#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3331fn sdot_enabled() -> bool {
3332    use std::sync::OnceLock;
3333    static ON: OnceLock<bool> = OnceLock::new();
3334    *ON.get_or_init(|| {
3335        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
3336        #[cfg(target_arch = "aarch64")]
3337        {
3338            want && std::arch::is_aarch64_feature_detected!("dotprod")
3339        }
3340        #[cfg(not(target_arch = "aarch64"))]
3341        {
3342            let _ = want;
3343            false
3344        }
3345    })
3346}
3347
3348/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
3349/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
3350/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
3351/// matvec, shared by all rows/workers.
3352struct SplitAct {
3353    xq: Vec<i8>,
3354    sx: f32,
3355    outliers: Vec<(usize, f32)>,
3356    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
3357    /// `−128·Σx`); one i32 per split, computed once per matvec.
3358    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
3359    xsum: i32,
3360}
3361
3362thread_local! {
3363    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
3364    /// and its hidden-size allocation was steady-state heap churn.
3365    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
3366        const { std::cell::RefCell::new(Vec::new()) };
3367}
3368
3369impl Drop for SplitAct {
3370    fn drop(&mut self) {
3371        let buf = std::mem::take(&mut self.xq);
3372        if buf.capacity() > 0 {
3373            XQ_FREE.with(|f| {
3374                let mut f = f.borrow_mut();
3375                if f.len() < 16 {
3376                    f.push(buf);
3377                }
3378            });
3379        }
3380    }
3381}
3382
3383fn split_act(x: &[f32]) -> SplitAct {
3384    let n = x.len();
3385    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
3386    let thr = 8.0 * rms;
3387    // One pass: collect outliers and the bulk absmax (outliers excluded —
3388    // identical to the old zero-then-fold over a copied buffer, minus the
3389    // full-vector copy).
3390    let mut outliers: Vec<(usize, f32)> = Vec::new();
3391    let mut amax = 0f32;
3392    for (j, &v) in x.iter().enumerate() {
3393        let a = v.abs();
3394        if a > thr {
3395            outliers.push((j, v));
3396        } else if a > amax {
3397            amax = a;
3398        }
3399    }
3400    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
3401    let inv = 1.0 / sx;
3402    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
3403    xq.clear();
3404    xq.reserve(n);
3405    if outliers.is_empty() {
3406        xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
3407    } else {
3408        // Outlier slots quantize to 0 (their exact term is added later).
3409        xq.extend(x.iter().map(|&v| {
3410            if v.abs() > thr {
3411                0
3412            } else {
3413                (v * inv).round().clamp(-127.0, 127.0) as i8
3414            }
3415        }));
3416    }
3417    let xsum = xq.iter().map(|&v| v as i32).sum();
3418    SplitAct { xq, sx, outliers, xsum }
3419}
3420
3421/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
3422/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
3423#[cfg(target_arch = "aarch64")]
3424#[target_feature(enable = "neon,dotprod")]
3425unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
3426    // SAFETY: callers uphold slice-length contracts (see call sites).
3427    unsafe {
3428        use core::arch::aarch64::*;
3429        use core::arch::asm;
3430        let wp = w.as_ptr() as *const i8;
3431        let n = w.len();
3432        let (mut a0, mut a1, mut a2, mut a3) =
3433            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3434        let mut i = 0;
3435        while i + 64 <= n {
3436            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
3437            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
3438            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
3439            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
3440            asm!(
3441                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3442                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3443                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
3444                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
3445                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3446                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3447                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
3448                options(pure, nomem, nostack),
3449            );
3450            i += 64;
3451        }
3452        while i + 16 <= n {
3453            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
3454            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
3455                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
3456            i += 16;
3457        }
3458        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
3459        while i < n {
3460            s += (*wp.add(i)) as i32 * xq[i] as i32;
3461            i += 1;
3462        }
3463        s
3464}
3465}
3466
3467/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
3468/// loaded once and reused, 4 independent accumulators hide sdot latency
3469/// (port of vmfcore `dot_i8_sdot_4rows`).
3470#[cfg(target_arch = "aarch64")]
3471#[target_feature(enable = "neon,dotprod")]
3472unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
3473    // SAFETY: callers uphold slice-length contracts (see call sites).
3474    unsafe {
3475        use core::arch::aarch64::*;
3476        use core::arch::asm;
3477        let n = xq.len();
3478        let px = xq.as_ptr();
3479        let (p0, p1, p2, p3) = (
3480            w0.as_ptr() as *const i8,
3481            w1.as_ptr() as *const i8,
3482            w2.as_ptr() as *const i8,
3483            w3.as_ptr() as *const i8,
3484        );
3485        let (mut a0, mut a1, mut a2, mut a3) =
3486            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3487        let mut i = 0;
3488        while i + 16 <= n {
3489            let x = vld1q_s8(px.add(i));
3490            let v0 = vld1q_s8(p0.add(i));
3491            let v1 = vld1q_s8(p1.add(i));
3492            let v2 = vld1q_s8(p2.add(i));
3493            let v3 = vld1q_s8(p3.add(i));
3494            asm!(
3495                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
3496                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
3497                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
3498                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
3499                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3500                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
3501                options(pure, nomem, nostack),
3502            );
3503            i += 16;
3504        }
3505        let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
3506        while i < n {
3507            let xi = *px.add(i) as i32;
3508            r[0] += (*p0.add(i)) as i32 * xi;
3509            r[1] += (*p1.add(i)) as i32 * xi;
3510            r[2] += (*p2.add(i)) as i32 * xi;
3511            r[3] += (*p3.add(i)) as i32 * xi;
3512            i += 1;
3513        }
3514        r
3515}
3516}
3517
3518/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
3519/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
3520/// line plus the shared activation chunk — a single sequential weight
3521/// stream per worker. Per-row accumulation is the same one-accumulator
3522/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
3523/// are bit-identical to the mmap-layout kernel.
3524#[cfg(target_arch = "aarch64")]
3525#[target_feature(enable = "neon,dotprod")]
3526unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
3527    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
3528    // n % 16 == 0 — guaranteed by the repack gate).
3529    unsafe {
3530        use core::arch::aarch64::*;
3531        use core::arch::asm;
3532        let n = xq.len();
3533        let px = xq.as_ptr();
3534        let pg = g.as_ptr() as *const i8;
3535        let (mut a0, mut a1, mut a2, mut a3) =
3536            (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3537        let mut i = 0;
3538        while i + 16 <= n {
3539            let x = vld1q_s8(px.add(i));
3540            let base = pg.add(4 * i);
3541            let v0 = vld1q_s8(base);
3542            let v1 = vld1q_s8(base.add(16));
3543            let v2 = vld1q_s8(base.add(32));
3544            let v3 = vld1q_s8(base.add(48));
3545            asm!(
3546                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
3547                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
3548                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
3549                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
3550                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3551                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
3552                options(pure, nomem, nostack),
3553            );
3554            i += 16;
3555        }
3556        [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
3557    }
3558}
3559
3560/// One q8 row range via SDOT (4-row blocks + tail) — the body of
3561/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
3562/// SAME kernel for several tensors under one pool dispatch. `rep` — the
3563/// load-time interleaved repack (empty = mmap layout only); rows outside
3564/// full 4-row groups always come from the mmap layout.
3565#[cfg(target_arch = "aarch64")]
3566fn q8_range_sdot(
3567    q: &[u8],
3568    rep: &[u8],
3569    row_scale: &[f32],
3570    act: &SplitAct,
3571    cols: usize,
3572    out_addr: SendMut,
3573    start: usize,
3574    end: usize,
3575) {
3576    let mut o = start;
3577    // Leading rows to the group boundary (repack path only): the pool
3578    // splits row ranges arbitrarily, groups are absolute.
3579    if !rep.is_empty() {
3580        while o < end && o % 4 != 0 {
3581            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3582            unsafe { *out_addr.at(o) = v };
3583            o += 1;
3584        }
3585    }
3586    while o + 4 <= end {
3587        let r = if rep.is_empty() {
3588            unsafe {
3589                dot_i8_sdot_4rows(
3590                    &q[o * cols..(o + 1) * cols],
3591                    &q[(o + 1) * cols..(o + 2) * cols],
3592                    &q[(o + 2) * cols..(o + 3) * cols],
3593                    &q[(o + 3) * cols..(o + 4) * cols],
3594                    &act.xq,
3595                )
3596            }
3597        } else {
3598            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
3599        };
3600        for k in 0..4 {
3601            let mut acc = r[k] as f32 * act.sx;
3602            for &(j, xv) in &act.outliers {
3603                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
3604            }
3605            // SAFETY: disjoint row ranges per worker.
3606            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
3607        }
3608        o += 4;
3609    }
3610    while o < end {
3611        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3612        unsafe { *out_addr.at(o) = v };
3613        o += 1;
3614    }
3615}
3616
3617/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
3618/// for the fused pair multi-matrix job (`matvec2_many`).
3619#[cfg(target_arch = "aarch64")]
3620#[allow(clippy::too_many_arguments)]
3621fn q8_range2_sdot(
3622    q: &[u8],
3623    row_scale: &[f32],
3624    a1: &SplitAct,
3625    a2: &SplitAct,
3626    cols: usize,
3627    p1: SendMut,
3628    p2: SendMut,
3629    start: usize,
3630    end: usize,
3631) {
3632    for o in start..end {
3633        let row = &q[o * cols..(o + 1) * cols];
3634        // SAFETY: disjoint row ranges per worker.
3635        unsafe {
3636            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
3637            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
3638        }
3639    }
3640}
3641
3642/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
3643#[allow(clippy::too_many_arguments)]
3644fn q8_range2_f32(
3645    q: &[u8],
3646    row_scale: &[f32],
3647    x1: &[f32],
3648    x2: &[f32],
3649    cols: usize,
3650    p1: SendMut,
3651    p2: SendMut,
3652    start: usize,
3653    end: usize,
3654) {
3655    for o in start..end {
3656        let row = &q[o * cols..(o + 1) * cols];
3657        // SAFETY: disjoint row ranges per worker.
3658        unsafe {
3659            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
3660            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
3661        }
3662    }
3663}
3664
3665/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
3666fn q8_range_f32(
3667    q: &[u8],
3668    row_scale: &[f32],
3669    xs: &[f32],
3670    cols: usize,
3671    out_addr: SendMut,
3672    start: usize,
3673    end: usize,
3674) {
3675    for o in start..end {
3676        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
3677        // SAFETY: disjoint row ranges per worker.
3678        unsafe { *out_addr.at(o) = v };
3679    }
3680}
3681
3682/// SDOT row dot with exact outlier correction:
3683/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
3684#[cfg(target_arch = "aarch64")]
3685#[inline]
3686fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
3687    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
3688    for &(j, xv) in &act.outliers {
3689        acc += (row[j] as i8) as f32 * xv;
3690    }
3691    acc
3692}
3693
3694/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
3695/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
3696/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
3697/// the caller multiplies by the activation scale and adds the exact
3698/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
3699/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
3700/// → zip(lo,hi) restores flat order.
3701#[cfg(target_arch = "aarch64")]
3702#[target_feature(enable = "neon,dotprod")]
3703unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3704    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
3705    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
3706    unsafe {
3707        use core::arch::aarch64::*;
3708        use core::arch::asm;
3709        let lomask = vdupq_n_u8(0x0F);
3710        let eight = vdupq_n_s8(8);
3711        let mut acc = 0f32;
3712        for gi in 0..gpr {
3713            let g = g0 + gi;
3714            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3715            let b = vld1q_u8(packed.as_ptr().add(g * 16));
3716            let lo = vandq_u8(b, lomask);
3717            let hi = vshrq_n_u8::<4>(b);
3718            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3719            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3720            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3721            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3722            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3723            asm!(
3724                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3725                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3726                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3727                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3728                options(pure, nomem, nostack),
3729            );
3730            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3731        }
3732        acc
3733    }
3734}
3735
3736/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
3737/// part) happens ONCE per group; both pre-quantized activations are
3738/// dotted against the same centered i8 registers. Per-lane math matches
3739/// `dot_q4_row_sdot` exactly.
3740#[cfg(target_arch = "aarch64")]
3741#[target_feature(enable = "neon,dotprod")]
3742unsafe fn dot_q4_row_sdot2(
3743    packed: &[u8],
3744    scales: &[u8],
3745    g0: usize,
3746    gpr: usize,
3747    xq1: &[i8],
3748    xq2: &[i8],
3749) -> (f32, f32) {
3750    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
3751    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
3752    unsafe {
3753        use core::arch::aarch64::*;
3754        use core::arch::asm;
3755        let lomask = vdupq_n_u8(0x0F);
3756        let eight = vdupq_n_s8(8);
3757        let (mut acc1, mut acc2) = (0f32, 0f32);
3758        for gi in 0..gpr {
3759            let g = g0 + gi;
3760            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3761            let b = vld1q_u8(packed.as_ptr().add(g * 16));
3762            let lo = vandq_u8(b, lomask);
3763            let hi = vshrq_n_u8::<4>(b);
3764            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3765            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3766            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
3767            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
3768            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
3769            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
3770            let (mut a0, mut a1, mut b0, mut b1) =
3771                (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3772            asm!(
3773                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
3774                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
3775                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
3776                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
3777                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3778                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
3779                e0 = in(vreg) e0, e1 = in(vreg) e1,
3780                x10 = in(vreg) x10, x11 = in(vreg) x11,
3781                x20 = in(vreg) x20, x21 = in(vreg) x21,
3782                options(pure, nomem, nostack),
3783            );
3784            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3785            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
3786        }
3787        (acc1, acc2)
3788    }
3789}
3790
3791// ───────────────────── fused int8 kernels ─────────────────────
3792
3793/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
3794/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
3795#[inline]
3796pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
3797    #[cfg(target_arch = "aarch64")]
3798    unsafe {
3799        return axpy_i8_f32_neon(acc, row, w);
3800    }
3801    #[cfg(target_arch = "x86_64")]
3802    if avx2_enabled() {
3803        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
3804    }
3805    #[allow(unreachable_code)]
3806    {
3807        for (a, &b) in acc.iter_mut().zip(row) {
3808            *a += w * b as f32;
3809        }
3810    }
3811}
3812
3813/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
3814#[cfg(target_arch = "x86_64")]
3815#[target_feature(enable = "avx2,fma")]
3816unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
3817    // SAFETY: callers uphold slice-length contracts (see call sites).
3818    unsafe {
3819        use core::arch::x86_64::*;
3820        let n = acc.len().min(row.len());
3821        let ap = acc.as_mut_ptr();
3822        let rp = row.as_ptr();
3823        let wv = _mm256_set1_ps(w);
3824        let mut j = 0usize;
3825        while j + 16 <= n {
3826            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
3827            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
3828            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
3829            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
3830            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
3831            _mm256_storeu_ps(ap.add(j), v0);
3832            _mm256_storeu_ps(ap.add(j + 8), v1);
3833            j += 16;
3834        }
3835        while j < n {
3836            *ap.add(j) += w * (*rp.add(j)) as f32;
3837            j += 1;
3838        }
3839    }
3840}
3841
3842#[cfg(target_arch = "aarch64")]
3843#[target_feature(enable = "neon")]
3844unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
3845    // SAFETY: callers uphold slice-length contracts (see call sites).
3846    unsafe {
3847        use core::arch::aarch64::*;
3848        let n = acc.len().min(row.len());
3849        let ap = acc.as_mut_ptr();
3850        let rp = row.as_ptr();
3851        let wv = vdupq_n_f32(w);
3852        let mut j = 0usize;
3853        while j + 16 <= n {
3854            let rb = vld1q_s8(rp.add(j));
3855            let lo = vmovl_s8(vget_low_s8(rb));
3856            let hi = vmovl_s8(vget_high_s8(rb));
3857            for (off, half) in [(0, lo), (8, hi)] {
3858                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
3859                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
3860                let o = j + off;
3861                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
3862                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
3863            }
3864            j += 16;
3865        }
3866        while j < n {
3867            *ap.add(j) += w * (*rp.add(j)) as f32;
3868            j += 1;
3869        }
3870}
3871}
3872
3873/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
3874/// ≈9× scalar), scalar elsewhere.
3875#[inline]
3876pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
3877    #[cfg(target_arch = "aarch64")]
3878    unsafe {
3879        return dot_i8_f32_neon(w, x);
3880    }
3881    #[cfg(target_arch = "x86_64")]
3882    if avx2_enabled() {
3883        return unsafe { dot_i8_f32_avx2(w, x) };
3884    }
3885    #[allow(unreachable_code)]
3886    {
3887        let mut sum = 0.0f32;
3888        for (j, &b) in w.iter().enumerate() {
3889            sum += (b as i8) as f32 * x[j];
3890        }
3891        sum
3892    }
3893}
3894
3895/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
3896/// folded into the product (no prescaled copy of x). NEON on aarch64,
3897/// scalar elsewhere. Used by the active-neuron path `row_dot`.
3898#[inline]
3899fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
3900    #[cfg(target_arch = "aarch64")]
3901    unsafe {
3902        return dot_i8_col_f32_neon(w, x, col);
3903    }
3904    #[allow(unreachable_code)]
3905    {
3906        let mut sum = 0.0f32;
3907        for (j, &b) in w.iter().enumerate() {
3908            sum += (b as i8) as f32 * x[j] * col[j];
3909        }
3910        sum
3911    }
3912}
3913
3914#[cfg(target_arch = "aarch64")]
3915#[target_feature(enable = "neon")]
3916unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
3917    // SAFETY: callers uphold slice-length contracts (see call sites).
3918    unsafe {
3919        use core::arch::aarch64::*;
3920        let n = x.len();
3921        let wp = w.as_ptr() as *const i8;
3922        let xp = x.as_ptr();
3923        let cp = col.as_ptr();
3924        let (mut a0, mut a1, mut a2, mut a3) =
3925            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3926        let mut j = 0usize;
3927        while j + 16 <= n {
3928            let wb = vld1q_s8(wp.add(j));
3929            let lo = vmovl_s8(vget_low_s8(wb));
3930            let hi = vmovl_s8(vget_high_s8(wb));
3931            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
3932            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
3933            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
3934            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
3935            a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
3936            a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
3937            a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
3938            a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
3939            j += 16;
3940        }
3941        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
3942        while j < n {
3943            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
3944            j += 1;
3945        }
3946        sum
3947}
3948}
3949
3950#[cfg(target_arch = "aarch64")]
3951#[target_feature(enable = "neon")]
3952unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
3953    // SAFETY: callers uphold slice-length contracts (see call sites).
3954    unsafe {
3955        use core::arch::aarch64::*;
3956        let n = x.len();
3957        let wp = w.as_ptr() as *const i8;
3958        let xp = x.as_ptr();
3959        let (mut a0, mut a1, mut a2, mut a3) =
3960            (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3961        let mut j = 0usize;
3962        while j + 16 <= n {
3963            let wb = vld1q_s8(wp.add(j));
3964            let lo = vmovl_s8(vget_low_s8(wb));
3965            let hi = vmovl_s8(vget_high_s8(wb));
3966            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
3967            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
3968            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
3969            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
3970            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
3971            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
3972            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
3973            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
3974            j += 16;
3975        }
3976        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
3977        while j < n {
3978            sum += (*wp.add(j)) as f32 * *xp.add(j);
3979            j += 1;
3980        }
3981        sum
3982}
3983}
3984
3985#[allow(clippy::too_many_arguments)]
3986fn qmatvec(
3987    q: &[u8],
3988    rep: &[u8],
3989    row_scale: &[f32],
3990    xs: &[f32],
3991    rows: usize,
3992    cols: usize,
3993    out: &mut [f32],
3994    pool: Option<&Pool>,
3995) {
3996    debug_assert_eq!(out.len(), rows);
3997    #[cfg(not(target_arch = "aarch64"))]
3998    let _ = rep;
3999
4000    #[cfg(target_arch = "aarch64")]
4001    if sdot_enabled() {
4002        let act = split_act(xs);
4003        let out_addr = SendMut(out.as_mut_ptr());
4004        let run_range = |start: usize, end: usize| {
4005            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
4006        };
4007        match pool {
4008            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4009            _ => run_range(0, rows),
4010        }
4011        return;
4012    }
4013    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
4014    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
4015    #[cfg(target_arch = "x86_64")]
4016    if avx2_a8w8_enabled() {
4017        let act = split_act(xs);
4018        let out_addr = SendMut(out.as_mut_ptr());
4019        let run_range =
4020            |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
4021        match pool {
4022            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4023            _ => run_range(0, rows),
4024        }
4025        return;
4026    }
4027
4028    let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
4029    match pool {
4030        Some(pool) if rows >= 256 => {
4031            let out_addr = SendMut(out.as_mut_ptr());
4032            pool.run(&move |widx, n| {
4033                let chunk = rows.div_ceil(n);
4034                let start = widx * chunk;
4035                let end = (start + chunk).min(rows);
4036                for o in start..end {
4037                    // SAFETY: disjoint row ranges per worker.
4038                    unsafe { *out_addr.at(o) = row_dot(o) };
4039                }
4040            });
4041        }
4042        _ => {
4043            for (o, dst) in out.iter_mut().enumerate() {
4044                *dst = row_dot(o);
4045            }
4046        }
4047    }
4048}
4049
4050#[allow(clippy::too_many_arguments)]
4051fn qmatvec2(
4052    q: &[u8],
4053    row_scale: &[f32],
4054    x1: &[f32],
4055    x2: &[f32],
4056    rows: usize,
4057    cols: usize,
4058    o1: &mut [f32],
4059    o2: &mut [f32],
4060    pool: Option<&Pool>,
4061) {
4062    #[cfg(target_arch = "aarch64")]
4063    if sdot_enabled() {
4064        let a1s = split_act(x1);
4065        let a2s = split_act(x2);
4066        let p1 = SendMut(o1.as_mut_ptr());
4067        let p2 = SendMut(o2.as_mut_ptr());
4068        let run_range = |start: usize, end: usize| {
4069            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4070        };
4071        match pool {
4072            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4073            _ => run_range(0, rows),
4074        }
4075        return;
4076    }
4077    #[cfg(target_arch = "x86_64")]
4078    if avx2_a8w8_enabled() {
4079        let a1s = split_act(x1);
4080        let a2s = split_act(x2);
4081        let p1 = SendMut(o1.as_mut_ptr());
4082        let p2 = SendMut(o2.as_mut_ptr());
4083        let run_range = |start: usize, end: usize| {
4084            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4085        };
4086        match pool {
4087            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4088            _ => run_range(0, rows),
4089        }
4090        return;
4091    }
4092
4093    let row_dots = |o: usize| -> (f32, f32) {
4094        let row = &q[o * cols..(o + 1) * cols];
4095        (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
4096    };
4097    match pool {
4098        Some(pool) if rows >= 256 => {
4099            let p1 = SendMut(o1.as_mut_ptr());
4100            let p2 = SendMut(o2.as_mut_ptr());
4101            pool.run(&move |widx, n| {
4102                let chunk = rows.div_ceil(n);
4103                let start = widx * chunk;
4104                let end = (start + chunk).min(rows);
4105                for o in start..end {
4106                    let (s1, s2) = row_dots(o);
4107                    // SAFETY: disjoint row ranges per worker.
4108                    unsafe {
4109                        *p1.at(o) = s1;
4110                        *p2.at(o) = s2;
4111                    }
4112                }
4113            });
4114        }
4115        _ => {
4116            for o in 0..rows {
4117                let (s1, s2) = row_dots(o);
4118                o1[o] = s1;
4119                o2[o] = s2;
4120            }
4121        }
4122    }
4123}
4124
4125#[derive(Clone, Copy)]
4126struct SendMut(*mut f32);
4127unsafe impl Send for SendMut {}
4128unsafe impl Sync for SendMut {}
4129
4130impl SendMut {
4131    #[inline]
4132    fn at(self, i: usize) -> *mut f32 {
4133        unsafe { self.0.add(i) }
4134    }
4135}
4136
4137#[cfg(test)]
4138mod tests {
4139    use super::*;
4140
4141    #[test]
4142    fn f32_matvec_matches_matvec_rows_bitexact() {
4143        let (rows, cols) = (300, 40);
4144        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
4145        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
4146        let qt = QTensor::from_f32(w.clone(), rows, cols);
4147
4148        let mut a = vec![0.0f32; rows];
4149        matvec_rows(None, &w, &x, &mut a);
4150        let mut b = vec![0.0f32; rows];
4151        qt.matvec(&x, &mut b, None);
4152        assert_eq!(a, b);
4153    }
4154
4155    #[test]
4156    fn sdot_kernel_exact_on_grid() {
4157        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
4158        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
4159        // exact f32 dot to float rounding. This isolates kernel
4160        // correctness from quantization noise.
4161        eprintln!("sdot_enabled = {}", sdot_enabled());
4162        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
4163        let w: Vec<u8> = (0..rows * cols)
4164            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4165            .collect();
4166        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
4167        let x: Vec<f32> = (0..cols)
4168            .map(|i| match i % 3 {
4169                0 => 1.0,
4170                1 => -1.0,
4171                _ => 0.0,
4172            })
4173            .collect();
4174        let mut a = vec![0.0f32; rows];
4175        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4176        for o in 0..rows {
4177            let mut acc = 0.0f32;
4178            for j in 0..cols {
4179                acc += (w[o * cols + j] as i8) as f32 * x[j];
4180            }
4181            let expect = acc * scales[o];
4182            assert!(
4183                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
4184                "row {o}: {} vs {expect}",
4185                a[o]
4186            );
4187        }
4188    }
4189
4190    #[test]
4191    fn q1_kernels_match_exact_reference() {
4192        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
4193        let (rows, cols) = (7, 96);
4194        let gpr = cols / GROUP_SIZE;
4195        let mut bytes = Vec::new();
4196        for t in 0..rows * gpr {
4197            let s = 0.01 + (t % 13) as f32 * 0.003;
4198            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4199            for j in 0..4 {
4200                bytes.push(((t * 31 + j * 97) % 251) as u8);
4201            }
4202        }
4203        // On-grid activations (±1, amax 1) → the SDOT path is exact.
4204        let x: Vec<f32> = (0..cols)
4205            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
4206            .collect();
4207        // Reference through the core dequant.
4208        let mut w = vec![0.0f32; rows * cols];
4209        cortiq_core::quant::dequant_q1(&bytes, &mut w);
4210        let mut expect = vec![0.0f32; rows];
4211        for o in 0..rows {
4212            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
4213        }
4214        let mut got = vec![0.0f32; rows];
4215        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
4216        for o in 0..rows {
4217            assert!(
4218                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
4219                "row {o}: {} vs {}",
4220                got[o],
4221                expect[o]
4222            );
4223        }
4224        // Pair and batch paths agree with the single path.
4225        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
4226        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
4227        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
4228        assert_eq!(a1, got);
4229        let mut xs = x.clone();
4230        xs.extend_from_slice(&x2);
4231        let mut mm = vec![0.0f32; 2 * rows];
4232        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
4233        assert_eq!(&mm[..rows], got.as_slice());
4234        assert_eq!(&mm[rows..], a2.as_slice());
4235    }
4236
4237    #[test]
4238    fn repack_is_bit_identical() {
4239        // The interleaved-repack kernel must produce EXACTLY the same
4240        // bits as the mmap-layout kernel: integer accumulation is order-
4241        // exact, the f32 epilogue is identical. Odd rows exercise the
4242        // tail; direct range calls exercise unaligned pool splits.
4243        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
4244        let w: Vec<u8> = (0..rows * cols)
4245            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
4246            .collect();
4247        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
4248        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
4249        let rep = q8_repack_layout(&w, rows, cols);
4250        // Group interleave round-trips.
4251        for g in 0..rows / 4 {
4252            for c in 0..cols / 16 {
4253                for lane in 0..4 {
4254                    assert_eq!(
4255                        &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
4256                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
4257                    );
4258                }
4259            }
4260        }
4261        let mut a = vec![0.0f32; rows];
4262        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4263        let mut b = vec![0.0f32; rows];
4264        qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
4265        assert_eq!(a, b, "full-range repack output diverged");
4266
4267        #[cfg(target_arch = "aarch64")]
4268        if sdot_enabled() {
4269            // Unaligned range split (pool workers get arbitrary bounds).
4270            let act = split_act(&x);
4271            let mut c1 = vec![0.0f32; rows];
4272            let mut c2 = vec![0.0f32; rows];
4273            q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
4274            q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
4275            assert_eq!(c1, c2, "unaligned-range repack output diverged");
4276        }
4277    }
4278
4279    #[test]
4280    fn sdot_a8w8_noise_is_bounded() {
4281        // Off-grid activations: A8 quantization noise must stay small in
4282        // relative L2 over the whole output (realistic accuracy contract;
4283        // vmfcore measured argmax-identical decode on real models).
4284        let (rows, cols) = (16, 512);
4285        let w: Vec<u8> = (0..rows * cols)
4286            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4287            .collect();
4288        let scales = vec![0.01f32; rows];
4289        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
4290        let mut a = vec![0.0f32; rows];
4291        qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4292        let (mut num, mut den) = (0f64, 0f64);
4293        for o in 0..rows {
4294            let mut acc = 0.0f32;
4295            for j in 0..cols {
4296                acc += (w[o * cols + j] as i8) as f32 * x[j];
4297            }
4298            let expect = acc * scales[o];
4299            num += ((a[o] - expect) as f64).powi(2);
4300            den += (expect as f64).powi(2);
4301        }
4302        let rel = (num / den.max(1e-12)).sqrt();
4303        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
4304    }
4305
4306    #[test]
4307    fn i8_dot_neon_matches_scalar() {
4308        let n = 100;
4309        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
4310        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
4311        let mut scalar = 0.0f32;
4312        for j in 0..n {
4313            scalar += (w[j] as i8) as f32 * x[j];
4314        }
4315        let fast = dot_i8_f32(&w, &x);
4316        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
4317    }
4318
4319    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
4320    #[test]
4321    fn vbitmatvec_matches_full_dequant() {
4322        let (rows, cols) = (6, 64);
4323        let ng = cols / GROUP_SIZE;
4324        // Hand-craft: bits per row, f16 scales, packed rows.
4325        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4326        let mut bytes = bits.clone();
4327        for g in 0..rows * ng {
4328            let s = 0.02 + 0.001 * g as f32;
4329            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4330        }
4331        for r in 0..rows {
4332            let b = bits[r] as usize;
4333            let (mut acc, mut nb) = (0u64, 0usize);
4334            let mut rowbytes = Vec::new();
4335            for i in 0..cols {
4336                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4337                acc = (acc << b) | v;
4338                nb += b;
4339                while nb >= 8 {
4340                    nb -= 8;
4341                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4342                }
4343            }
4344            if nb > 0 {
4345                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4346            }
4347            bytes.extend_from_slice(&rowbytes);
4348        }
4349        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4350
4351        let mut reference = vec![0f32; rows * cols];
4352        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
4353        let mut expect = vec![0f32; rows];
4354        for r in 0..rows {
4355            expect[r] = reference[r * cols..(r + 1) * cols]
4356                .iter()
4357                .zip(&x)
4358                .map(|(w, xv)| w * xv)
4359                .sum();
4360        }
4361        let mut got = vec![0f32; rows];
4362        let offsets = vbit_row_offsets(&bytes, rows, cols);
4363        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
4364        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
4365        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
4366        // the golden-parity gate).
4367        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4368        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
4369        for r in 0..rows {
4370            assert!(
4371                (got[r] - expect[r]).abs() < tol * scale,
4372                "row {r}: {} vs {}",
4373                got[r],
4374                expect[r]
4375            );
4376        }
4377    }
4378
4379    /// Fused q4 matvec must match the reference full-dequant + dense
4380    /// matvec bit-for-bit in structure (same f32 math, group order).
4381    #[test]
4382    fn q4matvec_matches_full_dequant() {
4383        let (rows, cols) = (8, 64);
4384        let groups = rows * cols / GROUP_SIZE;
4385        // Hand-craft a q4_block blob: nibbles then f16 scales.
4386        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4387        for i in 0..groups * 16 {
4388            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4389        }
4390        for g in 0..groups {
4391            let s = 0.01 + 0.003 * g as f32;
4392            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4393        }
4394        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4395
4396        let mut reference = vec![0.0f32; rows * cols];
4397        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
4398        let mut expect = vec![0.0f32; rows];
4399        for r in 0..rows {
4400            expect[r] = reference[r * cols..(r + 1) * cols]
4401                .iter()
4402                .zip(&x)
4403                .map(|(w, xv)| w * xv)
4404                .sum();
4405        }
4406
4407        let mut got = vec![0.0f32; rows];
4408        q4matvec(&bytes, &x, rows, cols, &mut got, None);
4409        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
4410        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
4411        // in the golden-parity gate).
4412        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4413        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
4414        for r in 0..rows {
4415            assert!(
4416                (got[r] - expect[r]).abs() < tol * scale,
4417                "row {r}: {} vs {}",
4418                got[r],
4419                expect[r]
4420            );
4421        }
4422    }
4423
4424    /// Fused two-input vbit matvec must equal two single matvecs exactly
4425    /// (same per-lane accumulation order on both scalar and SDOT paths).
4426    #[test]
4427    fn vbitmatvec2_equals_two_singles() {
4428        let (rows, cols) = (6, 64);
4429        let ng = cols / GROUP_SIZE;
4430        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4431        let mut bytes = bits.clone();
4432        for g in 0..rows * ng {
4433            let s = 0.02 + 0.001 * g as f32;
4434            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4435        }
4436        for r in 0..rows {
4437            let b = bits[r] as usize;
4438            let (mut acc, mut nb) = (0u64, 0usize);
4439            let mut rowbytes = Vec::new();
4440            for i in 0..cols {
4441                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4442                acc = (acc << b) | v;
4443                nb += b;
4444                while nb >= 8 {
4445                    nb -= 8;
4446                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4447                }
4448            }
4449            if nb > 0 {
4450                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4451            }
4452            bytes.extend_from_slice(&rowbytes);
4453        }
4454        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4455        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
4456        let offsets = vbit_row_offsets(&bytes, rows, cols);
4457
4458        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4459        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
4460        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
4461        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
4462        vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
4463        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
4464        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
4465    }
4466
4467    /// Fused two-input q4 matvec must equal two single matvecs exactly.
4468    #[test]
4469    fn q4matvec2_equals_two_singles() {
4470        let (rows, cols) = (8, 128);
4471        let groups = rows * cols / GROUP_SIZE;
4472        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4473        for i in 0..groups * 16 {
4474            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4475        }
4476        for g in 0..groups {
4477            let s = 0.01 + 0.003 * g as f32;
4478            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4479        }
4480        // Include an outlier channel so the SDOT correction path is
4481        // exercised in the pair kernel too.
4482        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4483        x1[9] = 250.0;
4484        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
4485
4486        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4487        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
4488        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
4489        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
4490        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
4491        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
4492        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
4493    }
4494
4495    /// Multi-matrix job must equal separate matvecs exactly — same
4496    /// kernels, only the dispatch is fused.
4497    #[test]
4498    fn matvec_many_equals_separate_matvecs() {
4499        use crate::pool::Pool;
4500        let (r1, r2, cols) = (300, 200, 64);
4501        let mk = |salt: usize, rows: usize| {
4502            QTensor::from_f32(
4503                (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
4504                rows,
4505                cols,
4506            )
4507        };
4508        let (a, b) = (mk(1, r1), mk(5, r2));
4509        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
4510        let pool = Pool::new(3);
4511
4512        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
4513        a.matvec(&x, &mut ea, Some(&pool));
4514        b.matvec(&x, &mut eb, Some(&pool));
4515        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
4516        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
4517        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
4518        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
4519    }
4520
4521    /// Batched q4/vbit matmat must equal per-position matvec calls
4522    /// exactly (the fallback it replaced) — same kernels, same order.
4523    #[test]
4524    fn batched_matmat_equals_per_position_matvec() {
4525        let (rows, cols, b) = (8, 64, 5);
4526        // q4 blob.
4527        let groups = rows * cols / GROUP_SIZE;
4528        let mut q4 = Vec::new();
4529        for i in 0..groups * 16 {
4530            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4531        }
4532        for g in 0..groups {
4533            q4.extend_from_slice(
4534                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
4535            );
4536        }
4537        // vbit blob (mixed widths incl. 8).
4538        let ng = cols / GROUP_SIZE;
4539        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
4540        let mut vb = bits.clone();
4541        for g in 0..rows * ng {
4542            vb.extend_from_slice(
4543                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
4544            );
4545        }
4546        for r in 0..rows {
4547            let bw = bits[r] as usize;
4548            let (mut acc, mut nb) = (0u64, 0usize);
4549            let mut rowbytes = Vec::new();
4550            for i in 0..cols {
4551                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
4552                acc = (acc << bw) | v;
4553                nb += bw;
4554                while nb >= 8 {
4555                    nb -= 8;
4556                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
4557                }
4558            }
4559            if nb > 0 {
4560                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4561            }
4562            vb.extend_from_slice(&rowbytes);
4563        }
4564        let offsets = vbit_row_offsets(&vb, rows, cols);
4565
4566        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
4567
4568        // q4: batch vs singles.
4569        let mut got = vec![0f32; b * rows];
4570        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
4571        for bi in 0..b {
4572            let mut expect = vec![0f32; rows];
4573            q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
4574            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
4575        }
4576
4577        // vbit: batch vs singles.
4578        let mut got = vec![0f32; b * rows];
4579        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
4580        for bi in 0..b {
4581            let mut expect = vec![0f32; rows];
4582            vbitmatvec(
4583                &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
4584            );
4585            assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
4586        }
4587    }
4588
4589    /// q4_tiled kernels must produce BIT-identical outputs to the q4
4590    /// split kernels on the same values (same ints, same order — only
4591    /// the byte placement differs).
4592    #[test]
4593    fn q4_tiled_matches_q4_block_bitexact() {
4594        let (rows, cols, b) = (8usize, 128usize, 3usize);
4595        let groups = rows * cols / GROUP_SIZE;
4596        let mut split = Vec::with_capacity(groups * 18);
4597        for i in 0..groups * 16 {
4598            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4599        }
4600        for g in 0..groups {
4601            split.extend_from_slice(
4602                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
4603            );
4604        }
4605        // Re-tile: [scale][nibbles] per group.
4606        let (packed, scales) = split.split_at(groups * 16);
4607        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
4608        for g in 0..groups {
4609            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
4610            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
4611        }
4612
4613        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4614        x1[9] = 250.0; // exercise the outlier path
4615        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
4616
4617        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
4618        q4matvec(&split, &x1, rows, cols, &mut a, None);
4619        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
4620        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
4621
4622        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4623        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
4624        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
4625        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
4626        assert_eq!(a1, t1);
4627        assert_eq!(a2, t2);
4628
4629        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
4630        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
4631        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
4632        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
4633        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
4634    }
4635
4636    /// q4 SDOT outlier correction: a single huge activation channel
4637    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
4638    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
4639    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
4640    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
4641    /// can never qualify (8² = n).
4642    #[test]
4643    fn q4matvec_sdot_outlier_exact() {
4644        let (rows, cols) = (4, 128);
4645        let groups = rows * cols / GROUP_SIZE;
4646        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4647        for i in 0..groups * 16 {
4648            bytes.push(((i * 11 + 5) % 256) as u8);
4649        }
4650        for g in 0..groups {
4651            let s = 0.02 + 0.002 * g as f32;
4652            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4653        }
4654        let mut x: Vec<f32> = (0..cols)
4655            .map(|i| match i % 3 {
4656                0 => 1.0,
4657                1 => -1.0,
4658                _ => 0.0,
4659            })
4660            .collect();
4661        x[17] = 300.0; // ≫ 8·rms → outlier channel
4662
4663        let mut reference = vec![0.0f32; rows * cols];
4664        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
4665        let mut expect = vec![0.0f32; rows];
4666        for r in 0..rows {
4667            expect[r] = reference[r * cols..(r + 1) * cols]
4668                .iter()
4669                .zip(&x)
4670                .map(|(w, xv)| w * xv)
4671                .sum();
4672        }
4673        let mut got = vec![0.0f32; rows];
4674        q4matvec(&bytes, &x, rows, cols, &mut got, None);
4675        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
4676        for r in 0..rows {
4677            assert!(
4678                (got[r] - expect[r]).abs() < 2e-3 * scale,
4679                "row {r}: {} vs {} (outlier term must be exact)",
4680                got[r],
4681                expect[r]
4682            );
4683        }
4684    }
4685}