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