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::{
16    GROUP_SIZE, Q1_TILE, Q4_TILE, Q4TP_NIB, f16_to_f32, q4tp_code, q4tp_ladder, q4tp_sections,
17};
18use cortiq_core::{CmfModel, TensorDtype};
19use std::sync::Arc;
20
21pub enum QTensor {
22    F32 {
23        data: Vec<f32>,
24        rows: usize,
25        cols: usize,
26    },
27    Mapped {
28        model: Arc<CmfModel>,
29        /// Index into the model's tensor directory.
30        idx: usize,
31        dtype: TensorDtype,
32        rows: usize,
33        cols: usize,
34        /// Per-row scales, dequantized to f32 up front (tiny).
35        row_scale: Vec<f32>,
36        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
37        col_field: Vec<f32>,
38        /// Vbit only: byte offset of each row's packed data within the
39        /// tensor blob (`[rows + 1]`, computed once at load — the per-
40        /// matvec prefix scan over row bit-widths was O(rows) each call).
41        vbit_offsets: Vec<usize>,
42        /// q8-family decode repack (load-time, optional): rows in groups
43        /// of 4, interleaved in 16-byte units — one 64-byte line per
44        /// iteration feeds all 4 sdot lanes, ONE sequential weight
45        /// stream per worker instead of four (this is where llama.cpp's
46        /// repacked Q8 kernels get their bandwidth). Empty = off
47        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
48        /// an anonymous copy of the quants for mmap pages that go cold.
49        repack: Vec<u8>,
50    },
51}
52
53/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
54/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
55/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
56/// streams per worker feed the prefetcher MORE memory-level parallelism
57/// than one); kept as an experiment flag for x86, where the tradeoff
58/// may land differently.
59fn repack_enabled() -> bool {
60    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
61    *ON.get_or_init(|| {
62        std::env::var("CMF_REPACK")
63            .map(|v| v == "1")
64            .unwrap_or(cfg!(target_os = "android"))
65    })
66}
67
68/// Interleave q8 rows for the decode kernel: group g holds rows
69/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
70/// full groups are packed — tail rows keep reading the mmap layout.
71fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
72    #[cfg(target_arch = "aarch64")]
73    let arch_ok = sdot_enabled();
74    #[cfg(not(target_arch = "aarch64"))]
75    let arch_ok = false;
76    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
77        return Vec::new();
78    }
79    q8_repack_layout(bytes, rows, cols)
80}
81
82/// The pure layout transform behind `q8_repack` (tested directly —
83/// the gate depends on arch and env).
84fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
85    let groups = rows / 4;
86    let mut rep = vec![0u8; groups * 4 * cols];
87    for g in 0..groups {
88        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
89        for c in 0..cols / 16 {
90            for lane in 0..4 {
91                let src = (g * 4 + lane) * cols + c * 16;
92                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
93                    .copy_from_slice(&bytes[src..src + 16]);
94            }
95        }
96    }
97    rep
98}
99
100/// Prefix-sum of vbit row payload offsets (absolute within the tensor
101/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
102fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
103    let ng = cols / GROUP_SIZE;
104    let bits = &bytes[..rows];
105    let mut offsets = Vec::with_capacity(rows + 1);
106    let mut off = rows + rows * ng * 2;
107    for r in 0..rows {
108        offsets.push(off);
109        off += (cols * bits[r] as usize).div_ceil(8);
110    }
111    offsets.push(off);
112    offsets
113}
114
115impl QTensor {
116    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
117        debug_assert_eq!(data.len(), rows * cols);
118        Self::F32 { data, rows, cols }
119    }
120
121    /// Wrap a directory tensor without dequantizing the payload.
122    /// Falls back to dequantized f32 for dtypes without a fused kernel.
123    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
124        // Indexed lookup: the linear directory scan made pipeline build
125        // O(N²) on MoE/skills files with thousands of tensors.
126        let idx = model
127            .tensor_index(name)
128            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
129        let entry = &model.tensors[idx];
130        if entry.shape.len() != 2 {
131            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
132        }
133        let (rows, cols) = (entry.shape[0], entry.shape[1]);
134        let bytes = model.entry_bytes(entry);
135
136        match entry.dtype {
137            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
138                let n = rows * cols;
139                let scales_off = n;
140                let row_scale: Vec<f32> = (0..rows)
141                    .map(|o| {
142                        f16_to_f32(u16::from_le_bytes([
143                            bytes[scales_off + o * 2],
144                            bytes[scales_off + o * 2 + 1],
145                        ]))
146                    })
147                    .collect();
148                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
149                    let col_off = n + rows * 2;
150                    (0..cols)
151                        .map(|i| {
152                            f16_to_f32(u16::from_le_bytes([
153                                bytes[col_off + i * 2],
154                                bytes[col_off + i * 2 + 1],
155                            ]))
156                        })
157                        .collect()
158                } else {
159                    Vec::new()
160                };
161                Ok(Self::Mapped {
162                    model: model.clone(),
163                    idx,
164                    dtype: entry.dtype,
165                    rows,
166                    cols,
167                    row_scale,
168                    col_field,
169                    vbit_offsets: Vec::new(),
170                    repack: q8_repack(bytes, rows, cols),
171                })
172            }
173            // vbit: fused kernel unpacks variable-bit rows from mmap.
174            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
175                model: model.clone(),
176                idx,
177                dtype: entry.dtype,
178                rows,
179                cols,
180                row_scale: Vec::new(),
181                col_field: Vec::new(),
182                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
183                repack: Vec::new(),
184            }),
185            // vbit_ro (§4.2): the offset table comes straight from the
186            // file — no load-time prefix scan; kernels are shared with
187            // legacy vbit (they consume absolute offsets either way).
188            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
189                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
190                let offsets: Vec<usize> = (0..=rows)
191                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
192                    .collect();
193                Ok(Self::Mapped {
194                    model: model.clone(),
195                    idx,
196                    dtype: entry.dtype,
197                    rows,
198                    cols,
199                    row_scale: Vec::new(),
200                    col_field: Vec::new(),
201                    vbit_offsets: offsets,
202                    repack: Vec::new(),
203                })
204            }
205            // q4_block: fused kernel reads nibbles straight from mmap —
206            // a 14B q4 file no longer explodes into ×8 f32 RAM.
207            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
208            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
209            // at kernel level over the split layout).
210            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
211                model: model.clone(),
212                idx,
213                dtype: entry.dtype,
214                rows,
215                cols,
216                row_scale: Vec::new(),
217                col_field: Vec::new(),
218                vbit_offsets: Vec::new(),
219                repack: Vec::new(),
220            }),
221            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
222            // 7.3% less file than q4t at the same 4-bit grid.
223            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
224                model: model.clone(),
225                idx,
226                dtype: entry.dtype,
227                rows,
228                cols,
229                row_scale: Vec::new(),
230                col_field: Vec::new(),
231                vbit_offsets: Vec::new(),
232                repack: Vec::new(),
233            }),
234            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
235                model: model.clone(),
236                idx,
237                dtype: entry.dtype,
238                rows,
239                cols,
240                row_scale: Vec::new(),
241                col_field: Vec::new(),
242                vbit_offsets: Vec::new(),
243                repack: Vec::new(),
244            }),
245            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
246            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
247                model: model.clone(),
248                idx,
249                dtype: entry.dtype,
250                rows,
251                cols,
252                row_scale: Vec::new(),
253                col_field: Vec::new(),
254                vbit_offsets: Vec::new(),
255                repack: Vec::new(),
256            }),
257            // q1t (ternary + outlier overlay): fused per-row dequant kernel
258            // reads straight from mmap — a 12B q1t stays ~its file size in
259            // RAM instead of dequantizing to ~48 GB of f32.
260            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
261                model: model.clone(),
262                idx,
263                dtype: entry.dtype,
264                rows,
265                cols,
266                row_scale: Vec::new(),
267                col_field: Vec::new(),
268                vbit_offsets: Vec::new(),
269                repack: Vec::new(),
270            }),
271            // No fused kernel yet → dequantize once (correct, more RAM).
272            _ => {
273                let mut data = vec![0.0f32; rows * cols];
274                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
275                Ok(Self::from_f32(data, rows, cols))
276            }
277        }
278    }
279
280    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
281    /// compute-bound, so offload pays at much smaller shapes than q8.)
282    pub(crate) fn is_q1(&self) -> bool {
283        matches!(
284            self,
285            Self::Mapped {
286                dtype: TensorDtype::Q1,
287                ..
288            }
289        )
290    }
291
292    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
293    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
294    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
295        match self {
296            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
297            _ => None,
298        }
299    }
300
301    /// (directory idx, rows, cols) of a q1-mapped tensor — the
302    /// whole-block GPU path resolves offsets itself.
303    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
304    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
305    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
306    /// Named `q1_parts` for historical reasons.
307    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
308        match self {
309            #[cfg(target_os = "macos")]
310            Self::Mapped {
311                dtype: TensorDtype::Q1T,
312                ..
313            } if !crate::gpu::metal_q1t_enabled() => None,
314            Self::Mapped {
315                idx,
316                dtype:
317                    TensorDtype::Q1
318                    | TensorDtype::Q1T
319                    | TensorDtype::Q4Block
320                    | TensorDtype::Q4Tiled
321                    | TensorDtype::Q4TiledP
322                    | TensorDtype::Q8Row
323                    | TensorDtype::Q8_2f,
324                rows,
325                cols,
326                ..
327            } => Some((*idx, *rows, *cols)),
328            _ => None,
329        }
330    }
331
332    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
333    /// chunk-prefill graph takes it in the same 4-tuple slot as
334    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
335    /// inside the 18-byte tiles, and the empty slice is what tells the
336    /// encoder to reach for the q4t kernels.
337    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
338        match self {
339            Self::Mapped {
340                idx,
341                dtype: TensorDtype::Q4Tiled,
342                rows,
343                cols,
344                ..
345            } => Some((*idx, *rows, *cols)),
346            _ => None,
347        }
348    }
349
350    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
351    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
352    /// apart by the tensor's dtype, not by the slot.
353    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
354        match self {
355            Self::Mapped {
356                idx,
357                dtype: TensorDtype::Q4TiledP,
358                rows,
359                cols,
360                ..
361            } => Some((*idx, *rows, *cols)),
362            _ => None,
363        }
364    }
365
366    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
367    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
368    /// q8_2f is excluded on purpose: its column field would need a
369    /// prescale stage on the device.
370    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
371        match self {
372            Self::Mapped {
373                idx,
374                dtype: TensorDtype::Q8Row,
375                rows,
376                cols,
377                row_scale,
378                col_field,
379                ..
380            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
381            _ => None,
382        }
383    }
384
385    pub fn rows(&self) -> usize {
386        match self {
387            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
388        }
389    }
390
391    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
392    /// needs the raw file coordinates of its three projections.
393    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
394        match self {
395            Self::Mapped {
396                model,
397                idx,
398                dtype: TensorDtype::Q4Tiled,
399                ..
400            } => Some((model, *idx)),
401            _ => None,
402        }
403    }
404
405    pub fn cols(&self) -> usize {
406        match self {
407            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
408        }
409    }
410
411    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
412    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
413    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
414        match self {
415            Self::Mapped {
416                model,
417                idx,
418                dtype: TensorDtype::Q1,
419                ..
420            } => Some((model, *idx)),
421            _ => None,
422        }
423    }
424
425    /// (model, idx, kind, row_scale) for a graph-capable mapped weight. kind:
426    /// 0=q8_row (per-row scales), 1=q1, 2=q4_tiled, 3=q1t (tile-embedded, no
427    /// rs). None for dtypes the token graph does not handle (q8_2f/q4_block/vbit).
428    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
429        match self {
430            Self::Mapped {
431                model,
432                idx,
433                dtype: TensorDtype::Q8Row,
434                row_scale,
435                ..
436            } => Some((model, *idx, 0, row_scale.as_slice())),
437            Self::Mapped {
438                model,
439                idx,
440                dtype: TensorDtype::Q1,
441                ..
442            } => Some((model, *idx, 1, &[])),
443            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
444            // the wgpu token graph fed 18B interleaved tiles to the
445            // split-layout q4b kernel — garbage output on q4t models
446            // (caught by an end-to-end answer check on real Vulkan).
447            Self::Mapped {
448                model,
449                idx,
450                dtype: TensorDtype::Q4Tiled,
451                ..
452            } => Some((model, *idx, 5, &[])),
453            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
454            // and feeding them to the q4t kernel is exactly the mistake that
455            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
456            Self::Mapped {
457                model,
458                idx,
459                dtype: TensorDtype::Q4TiledP,
460                ..
461            } => Some((model, *idx, 6, &[])),
462            Self::Mapped {
463                model,
464                idx,
465                dtype: TensorDtype::Q4Block,
466                ..
467            } => Some((model, *idx, 2, &[])),
468            Self::Mapped {
469                model,
470                idx,
471                dtype: TensorDtype::Q1T,
472                ..
473            } => Some((model, *idx, 3, &[])),
474            _ => None,
475        }
476    }
477
478    /// Dense f32 view — only for owned tensors. Masked/sparse execution
479    /// paths require it; quantized weights don't support masks yet.
480    pub fn as_f32(&self) -> Option<&[f32]> {
481        match self {
482            Self::F32 { data, .. } => Some(data),
483            Self::Mapped { .. } => None,
484        }
485    }
486
487    fn quant_bytes(&self) -> &[u8] {
488        match self {
489            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
490            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
491        }
492    }
493
494    /// Dequantize one row into `dst` (embedding lookup).
495    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
496        let cols = self.cols();
497        debug_assert_eq!(dst.len(), cols);
498        match self {
499            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
500            Self::Mapped {
501                dtype,
502                row_scale,
503                col_field,
504                vbit_offsets,
505                ..
506            } => {
507                if *dtype == TensorDtype::Q4Tiled {
508                    let bytes = self.quant_bytes();
509                    let gpr = cols / GROUP_SIZE;
510                    for gi in 0..gpr {
511                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
512                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
513                        for (k, &b) in tile[2..].iter().enumerate() {
514                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
515                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
516                        }
517                    }
518                    return;
519                }
520                if *dtype == TensorDtype::Q4TiledP {
521                    let bytes = self.quant_bytes();
522                    let gpr = cols / GROUP_SIZE;
523                    let v = Q4tpView::new(bytes, self.rows(), cols);
524                    let mut sc = vec![0f32; gpr];
525                    v.scales_into(r, gpr, &mut sc);
526                    for gi in 0..gpr {
527                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
528                        let s = sc[gi];
529                        for (k, &b) in tile.iter().enumerate() {
530                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
531                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
532                        }
533                    }
534                    return;
535                }
536                if *dtype == TensorDtype::Q4Block {
537                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
538                    let gpr = cols / GROUP_SIZE;
539                    for gi in 0..gpr {
540                        let g = r * gpr + gi;
541                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
542                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
543                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
544                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
545                        }
546                    }
547                    return;
548                }
549                if *dtype == TensorDtype::Q1 {
550                    let bytes = self.quant_bytes();
551                    let gpr = cols / GROUP_SIZE;
552                    for gi in 0..gpr {
553                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
554                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
555                        for (j, &b) in tile[2..].iter().enumerate() {
556                            for k in 0..8 {
557                                dst[gi * GROUP_SIZE + j * 8 + k] =
558                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
559                            }
560                        }
561                    }
562                    return;
563                }
564                if *dtype == TensorDtype::Q1T {
565                    let bytes = self.quant_bytes();
566                    let gpr = cols / GROUP_SIZE;
567                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
568                    for gi in 0..gpr {
569                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
570                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
571                            bytes[off],
572                            bytes[off + 1],
573                        ]));
574                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
575                        for k in 0..GROUP_SIZE {
576                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
577                            {
578                                1 => s,
579                                2 => -s,
580                                _ => 0.0,
581                            };
582                        }
583                    }
584                    // Overlay
585                    let rows = self.rows();
586                    let entries = base_len + (rows + 1) * 4;
587                    if entries <= bytes.len() {
588                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
589                        let r0 = u32::from_le_bytes([
590                            ptrs[r * 4],
591                            ptrs[r * 4 + 1],
592                            ptrs[r * 4 + 2],
593                            ptrs[r * 4 + 3],
594                        ]) as usize;
595                        let r1 = u32::from_le_bytes([
596                            ptrs[(r + 1) * 4],
597                            ptrs[(r + 1) * 4 + 1],
598                            ptrs[(r + 1) * 4 + 2],
599                            ptrs[(r + 1) * 4 + 3],
600                        ]) as usize;
601                        let off = entries + r0 * 4;
602                        for i in 0..r1 - r0 {
603                            let item = &bytes[off + i * 4..off + i * 4 + 4];
604                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
605                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
606                                item[2], item[3],
607                            ]));
608                            if c < cols {
609                                dst[c] = v;
610                            }
611                        }
612                    }
613                    return;
614                }
615                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
616                    let bytes = self.quant_bytes();
617                    let rows = self.rows();
618                    let ng = cols / GROUP_SIZE;
619                    let bits = &bytes[..rows];
620                    let sc_off = rows;
621                    // Precomputed at load — embedding lookup used to scan
622                    // the bit-widths of every preceding row (O(token_id)).
623                    let off = vbit_offsets[r];
624                    let b = bits[r] as usize;
625                    let l = ((1usize << (b - 1)) - 1) as f32;
626                    let data = &bytes[off..];
627                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
628                    for (i, d) in dst.iter_mut().enumerate() {
629                        while nbits < b {
630                            acc = (acc << 8) | data[idx] as u64;
631                            idx += 1;
632                            nbits += 8;
633                        }
634                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
635                        nbits -= b;
636                        let so = (r * ng + i / GROUP_SIZE) * 2;
637                        let sv = f16_to_f32(u16::from_le_bytes([
638                            bytes[sc_off + so],
639                            bytes[sc_off + so + 1],
640                        ]));
641                        *d = (u - l) * sv;
642                    }
643                    return;
644                }
645                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
646                let s = row_scale[r];
647                match dtype {
648                    TensorDtype::Q8Row => {
649                        for (d, &b) in dst.iter_mut().zip(q) {
650                            *d = (b as i8) as f32 * s;
651                        }
652                    }
653                    TensorDtype::Q8_2f => {
654                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
655                            *d = (b as i8) as f32 * s * col_field[i];
656                        }
657                    }
658                    _ => unreachable!(),
659                }
660            }
661        }
662    }
663
664    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
665    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
666    /// false for group-packed q4/vbit (column access would unpack whole
667    /// groups — sparse execution falls back to f32 for those).
668    pub fn sparse_col_ok(&self) -> bool {
669        match self {
670            Self::F32 { .. } => true,
671            Self::Mapped { dtype, .. } => {
672                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
673            }
674        }
675    }
676
677    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
678    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
679    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
680    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
681        let inter = self.cols();
682        let hidden = self.rows();
683        debug_assert_eq!(out.len(), hidden);
684        match self {
685            Self::F32 { data, .. } => {
686                for (k, o) in out.iter_mut().enumerate() {
687                    *o += w * data[k * inter + c];
688                }
689            }
690            Self::Mapped {
691                dtype,
692                row_scale,
693                col_field,
694                ..
695            } => {
696                let q = self.quant_bytes();
697                let colf = if *dtype == TensorDtype::Q8_2f {
698                    col_field[c]
699                } else {
700                    1.0
701                };
702                let wc = w * colf;
703                for (k, o) in out.iter_mut().enumerate() {
704                    let b = q[k * inter + c] as i8 as f32;
705                    *o += wc * b * row_scale[k];
706                }
707            }
708        }
709    }
710
711    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
712    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
713    /// into `scratch` first (rare for active-FFN weights).
714    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
715        let cols = self.cols();
716        match self {
717            Self::F32 { data, .. } => {
718                let row = &data[r * cols..(r + 1) * cols];
719                row.iter().zip(x).map(|(w, v)| w * v).sum()
720            }
721            Self::Mapped {
722                dtype,
723                row_scale,
724                col_field,
725                ..
726            } => match dtype {
727                TensorDtype::Q8Row => {
728                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
729                    dot_i8_f32(q, x) * row_scale[r]
730                }
731                TensorDtype::Q8_2f => {
732                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
733                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
734                }
735                _ => {
736                    self.row_f32(r, scratch);
737                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
738                }
739            },
740        }
741    }
742
743    /// `out = W · x` (row-major). F32 delegates to the historical
744    /// bit-exact path; Mapped runs the fused int8 kernel.
745    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
746        match self {
747            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
748            Self::Mapped {
749                model,
750                idx,
751                dtype,
752                rows,
753                cols,
754                row_scale,
755                col_field,
756                vbit_offsets,
757                repack,
758            } => {
759                let _ = (model, idx);
760                if *dtype == TensorDtype::Q4Block {
761                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
762                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
763                    // the winner; Metal returns false → the CPU kernel below.
764                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
765                        let t0 = std::time::Instant::now();
766                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
767                            crate::gpu::ProbeArm::Gpu => {
768                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
769                                    crate::gpu::probe_record(
770                                        crate::gpu::OpClass::Matvec,
771                                        true,
772                                        t0.elapsed(),
773                                    );
774                                    return;
775                                }
776                            }
777                            crate::gpu::ProbeArm::CpuTimed => {
778                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
779                                crate::gpu::probe_record(
780                                    crate::gpu::OpClass::Matvec,
781                                    false,
782                                    t0.elapsed(),
783                                );
784                                return;
785                            }
786                            crate::gpu::ProbeArm::Cpu => {}
787                        }
788                    }
789                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
790                    return;
791                }
792                if *dtype == TensorDtype::Q4Tiled {
793                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
794                    return;
795                }
796                if *dtype == TensorDtype::Q4TiledP {
797                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
798                    return;
799                }
800                if *dtype == TensorDtype::Q1 {
801                    // GPU route for large q1 matvecs (out_proj / lm_head
802                    // class): the CPU q1 kernel is load-port-bound at
803                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
804                    // probe measures both arms and keeps the winner.
805                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
806                        let t0 = std::time::Instant::now();
807                        let arm = if crate::gpu::q1_force() {
808                            crate::gpu::ProbeArm::Gpu
809                        } else {
810                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
811                        };
812                        match arm {
813                            crate::gpu::ProbeArm::Gpu => {
814                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
815                                    crate::gpu::probe_record(
816                                        crate::gpu::OpClass::Matvec,
817                                        true,
818                                        t0.elapsed(),
819                                    );
820                                    return;
821                                }
822                            }
823                            crate::gpu::ProbeArm::CpuTimed => {
824                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
825                                crate::gpu::probe_record(
826                                    crate::gpu::OpClass::Matvec,
827                                    false,
828                                    t0.elapsed(),
829                                );
830                                return;
831                            }
832                            crate::gpu::ProbeArm::Cpu => {}
833                        }
834                    }
835                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
836                    return;
837                }
838                if *dtype == TensorDtype::Q1T {
839                    // GPU route for large q1t matvecs: the ternary BASE dot runs
840                    // on the GPU (load-port-bound on CPU, like q1), then the
841                    // sparse overlay is added on the CPU. Probe keeps the winner.
842                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
843                        let t0 = std::time::Instant::now();
844                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
845                            crate::gpu::ProbeArm::Gpu => {
846                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
847                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
848                                    crate::gpu::probe_record(
849                                        crate::gpu::OpClass::Matvec,
850                                        true,
851                                        t0.elapsed(),
852                                    );
853                                    return;
854                                }
855                            }
856                            crate::gpu::ProbeArm::CpuTimed => {
857                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
858                                crate::gpu::probe_record(
859                                    crate::gpu::OpClass::Matvec,
860                                    false,
861                                    t0.elapsed(),
862                                );
863                                return;
864                            }
865                            crate::gpu::ProbeArm::Cpu => {}
866                        }
867                    }
868                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
869                    return;
870                }
871                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
872                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
873                    return;
874                }
875                let xs = prescale(x, col_field, *dtype);
876                // D5: large q8 matrices (lm_head-class) — hybrid
877                // CPU∥GPU: split the rows, both sides compute
878                // SIMULTANEOUSLY (same math, shared prescale).
879                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
880                if *rows >= crate::gpu::min_rows()
881                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
882                    && std::env::var("CMF_GPU_LMHEAD")
883                        .map(|v| v != "0")
884                        .unwrap_or(true)
885                    && crate::gpu::enabled_here()
886                {
887                    // Runtime probe: alternate the hybrid against the
888                    // pure-CPU matvec, keep whichever is faster HERE.
889                    let t0 = std::time::Instant::now();
890                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
891                        crate::gpu::ProbeArm::Gpu => {}
892                        crate::gpu::ProbeArm::CpuTimed => {
893                            qmatvec(
894                                self.quant_bytes(),
895                                repack,
896                                row_scale,
897                                x,
898                                col_field,
899                                *dtype,
900                                *rows,
901                                *cols,
902                                out,
903                                pool,
904                            );
905                            crate::gpu::probe_record(
906                                crate::gpu::OpClass::Matvec,
907                                false,
908                                t0.elapsed(),
909                            );
910                            return;
911                        }
912                        crate::gpu::ProbeArm::Cpu => {
913                            qmatvec(
914                                self.quant_bytes(),
915                                repack,
916                                row_scale,
917                                x,
918                                col_field,
919                                *dtype,
920                                *rows,
921                                *cols,
922                                out,
923                                pool,
924                            );
925                            return;
926                        }
927                    }
928                    let frac = std::env::var("CMF_GPU_SPLIT")
929                        .ok()
930                        .and_then(|v| v.parse::<f32>().ok())
931                        .unwrap_or(0.5)
932                        .clamp(0.0, 1.0);
933                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
934                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
935                    let bytes = self.quant_bytes();
936                    let ok = std::thread::scope(|sc| {
937                        let g = sc.spawn(|| {
938                            crate::gpu::q8_matvec_range(
939                                model,
940                                *idx,
941                                cpu_rows,
942                                &row_scale[cpu_rows..],
943                                &xs,
944                                *rows - cpu_rows,
945                                *cols,
946                                out_gpu,
947                            )
948                        });
949                        if cpu_rows > 0 {
950                            // Repack prefix covers the full groups of the
951                            // CPU half (the split starts at row 0).
952                            let rep_cpu = if repack.is_empty() {
953                                &[][..]
954                            } else {
955                                &repack[..(cpu_rows / 4) * 4 * *cols]
956                            };
957                            qmatvec(
958                                &bytes[..cpu_rows * *cols],
959                                rep_cpu,
960                                &row_scale[..cpu_rows],
961                                x,
962                                col_field,
963                                *dtype,
964                                cpu_rows,
965                                *cols,
966                                out_cpu,
967                                pool,
968                            );
969                        }
970                        g.join().unwrap_or(false)
971                    });
972                    if ok {
973                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
974                        return;
975                    }
976                    // GPU failed — CPU finishes its half (rows rebased —
977                    // group offsets don't line up, mmap layout only).
978                    qmatvec(
979                        &bytes[cpu_rows * *cols..(*rows) * *cols],
980                        &[],
981                        &row_scale[cpu_rows..],
982                        x,
983                        col_field,
984                        *dtype,
985                        *rows - cpu_rows,
986                        *cols,
987                        out_gpu,
988                        pool,
989                    );
990                    return;
991                }
992                qmatvec(
993                    self.quant_bytes(),
994                    repack,
995                    row_scale,
996                    x,
997                    col_field,
998                    *dtype,
999                    *rows,
1000                    *cols,
1001                    out,
1002                    pool,
1003                );
1004            }
1005        }
1006    }
1007
1008    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1009    pub fn matvec2(
1010        &self,
1011        x1: &[f32],
1012        x2: &[f32],
1013        o1: &mut [f32],
1014        o2: &mut [f32],
1015        pool: Option<&Pool>,
1016    ) {
1017        match self {
1018            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1019            Self::Mapped {
1020                dtype,
1021                rows,
1022                cols,
1023                row_scale,
1024                col_field,
1025                vbit_offsets,
1026                ..
1027            } => {
1028                if *dtype == TensorDtype::Q4Block {
1029                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1030                    return;
1031                }
1032                if *dtype == TensorDtype::Q4Tiled {
1033                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1034                    return;
1035                }
1036                if *dtype == TensorDtype::Q4TiledP {
1037                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1038                    return;
1039                }
1040                if *dtype == TensorDtype::Q1 {
1041                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1042                    return;
1043                }
1044                if *dtype == TensorDtype::Q1T {
1045                    // Fused ternary pair: one row pass, the register
1046                    // unpack shared across both streams on ARM. (Q1T
1047                    // lacks a row_scale array — scales live inline in
1048                    // the tiles — so it must not fall through to the
1049                    // q8 qmatvec2 below.)
1050                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1051                    return;
1052                }
1053                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1054                    vbitmatvec2(
1055                        self.quant_bytes(),
1056                        vbit_offsets,
1057                        x1,
1058                        x2,
1059                        *rows,
1060                        *cols,
1061                        o1,
1062                        o2,
1063                        pool,
1064                    );
1065                    return;
1066                }
1067                qmatvec2(
1068                    self.quant_bytes(),
1069                    row_scale,
1070                    x1,
1071                    x2,
1072                    col_field,
1073                    *dtype,
1074                    *rows,
1075                    *cols,
1076                    o1,
1077                    o2,
1078                    pool,
1079                );
1080            }
1081        }
1082    }
1083}
1084
1085impl QTensor {
1086    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1087    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1088    /// to b matvec calls (same dot kernels in the same order); the win —
1089    /// the weight row streams from DRAM once per batch, not b times.
1090    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1091        let cols = self.cols();
1092        let rows = self.rows();
1093        debug_assert_eq!(xs_all.len(), b * cols);
1094        debug_assert_eq!(out.len(), b * rows);
1095        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1096        // Mapped tensors carry a directory name; the check is a relaxed
1097        // atomic load, free when not calibrating.
1098        if crate::gptq_capture::capturing() {
1099            if let Self::Mapped { model, idx, .. } = self {
1100                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1101            }
1102        }
1103        match self {
1104            Self::F32 { data, .. } => {
1105                let out_addr = SendMut(out.as_mut_ptr());
1106                let run = |start: usize, end: usize| {
1107                    for o in start..end {
1108                        let row = &data[o * cols..(o + 1) * cols];
1109                        for bi in 0..b {
1110                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1111                            let mut acc = 0f32;
1112                            for j in 0..cols {
1113                                acc += row[j] * x[j];
1114                            }
1115                            unsafe { *out_addr.at(bi * rows + o) = acc };
1116                        }
1117                    }
1118                };
1119                dispatch_rows(pool, rows, &run);
1120            }
1121            Self::Mapped {
1122                dtype,
1123                row_scale,
1124                col_field,
1125                vbit_offsets,
1126                ..
1127            } => {
1128                if *dtype == TensorDtype::Q4Block {
1129                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1130                    return;
1131                }
1132                if *dtype == TensorDtype::Q4TiledP {
1133                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1134                    // device); the probe keeps whichever beats the CPU arm.
1135                    // Narrow (prompt-encode) and wide (DiT) batches probe
1136                    // as separate classes — the regimes have opposite
1137                    // winners and one shared verdict locked the wrong arm.
1138                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1139                    // (a fair-condition op is ≤~100 ms even at 1024px)
1140                    // means the device is contended by another process
1141                    // (e.g. a simulator) — verdicts are per-process, so
1142                    // without the bail the whole render crawls behind
1143                    // someone else's queue.
1144                    if b >= 32
1145                        && b * rows * cols >= 128_000_000
1146                        && cols % 32 == 0
1147                        && !crate::gpu::mm_killed()
1148                        && crate::gpu::enabled_here()
1149                    {
1150                        let class = if b >= 128 {
1151                            crate::gpu::OpClass::MatmatWide
1152                        } else {
1153                            crate::gpu::OpClass::Matmat
1154                        };
1155                        if let Self::Mapped { model, idx, .. } = self {
1156                            let t0 = std::time::Instant::now();
1157                            match crate::gpu::probe_arm(class) {
1158                                crate::gpu::ProbeArm::Gpu => {
1159                                    if crate::gpu::q4tp_matmat(
1160                                        model, *idx, xs_all, b, rows, cols, out,
1161                                    ) {
1162                                        let el = t0.elapsed();
1163                                        // Work-proportional budget: ~8× the
1164                                        // fair-device estimate (+20 ms slack).
1165                                        // An absolute cap missed the worst
1166                                        // case — contended ops sit at
1167                                        // 100–240 ms each and still bury a
1168                                        // render whose fair op is 3–9 ms.
1169                                        // Cold ops (first PSO build, buffer
1170                                        // alloc) are exempt: a one-off
1171                                        // ~50 ms compile is not contention.
1172                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1173                                        let budget = std::time::Duration::from_secs_f64(
1174                                            flops / 1.5e12 * 8.0 + 0.020,
1175                                        );
1176                                        if el > budget && !crate::gpu::probe_was_cold() {
1177                                            tracing::warn!(
1178                                                "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1179                                                 device contended, CPU for the rest of the process"
1180                                            );
1181                                            crate::gpu::mm_kill();
1182                                        }
1183                                        crate::gpu::probe_record(class, true, el);
1184                                        return;
1185                                    }
1186                                }
1187                                crate::gpu::ProbeArm::CpuTimed => {
1188                                    q4tp_matmat(
1189                                        self.quant_bytes(),
1190                                        xs_all,
1191                                        b,
1192                                        rows,
1193                                        cols,
1194                                        out,
1195                                        pool,
1196                                    );
1197                                    crate::gpu::probe_record(class, false, t0.elapsed());
1198                                    return;
1199                                }
1200                                crate::gpu::ProbeArm::Cpu => {}
1201                            }
1202                        }
1203                    }
1204                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1205                    return;
1206                }
1207                if *dtype == TensorDtype::Q4Tiled {
1208                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1209                    // device); the probe keeps whichever beats the CPU arm.
1210                    // Narrow (prompt-encode) and wide (DiT) batches probe
1211                    // as separate classes — the regimes have opposite
1212                    // winners and one shared verdict locked the wrong arm.
1213                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1214                    // (a fair-condition op is ≤~100 ms even at 1024px)
1215                    // means the device is contended by another process
1216                    // (e.g. a simulator) — verdicts are per-process, so
1217                    // without the bail the whole render crawls behind
1218                    // someone else's queue.
1219                    if b >= 32
1220                        && b * rows * cols >= 128_000_000
1221                        && cols % 32 == 0
1222                        && !crate::gpu::mm_killed()
1223                        && crate::gpu::enabled_here()
1224                    {
1225                        let class = if b >= 128 {
1226                            crate::gpu::OpClass::MatmatWide
1227                        } else {
1228                            crate::gpu::OpClass::Matmat
1229                        };
1230                        if let Self::Mapped { model, idx, .. } = self {
1231                            let t0 = std::time::Instant::now();
1232                            match crate::gpu::probe_arm(class) {
1233                                crate::gpu::ProbeArm::Gpu => {
1234                                    if crate::gpu::q4t_matmat(
1235                                        model, *idx, xs_all, b, rows, cols, out,
1236                                    ) {
1237                                        let el = t0.elapsed();
1238                                        // Work-proportional budget: ~8× the
1239                                        // fair-device estimate (+20 ms slack).
1240                                        // An absolute cap missed the worst
1241                                        // case — contended ops sit at
1242                                        // 100–240 ms each and still bury a
1243                                        // render whose fair op is 3–9 ms.
1244                                        // Cold ops (first PSO build, buffer
1245                                        // alloc) are exempt: a one-off
1246                                        // ~50 ms compile is not contention.
1247                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1248                                        let budget = std::time::Duration::from_secs_f64(
1249                                            flops / 1.5e12 * 8.0 + 0.020,
1250                                        );
1251                                        if el > budget && !crate::gpu::probe_was_cold() {
1252                                            tracing::warn!(
1253                                                "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1254                                                 device contended, CPU for the rest of the process"
1255                                            );
1256                                            crate::gpu::mm_kill();
1257                                        }
1258                                        crate::gpu::probe_record(class, true, el);
1259                                        return;
1260                                    }
1261                                }
1262                                crate::gpu::ProbeArm::CpuTimed => {
1263                                    q4t_matmat(
1264                                        self.quant_bytes(),
1265                                        xs_all,
1266                                        b,
1267                                        rows,
1268                                        cols,
1269                                        out,
1270                                        pool,
1271                                    );
1272                                    crate::gpu::probe_record(class, false, t0.elapsed());
1273                                    return;
1274                                }
1275                                crate::gpu::ProbeArm::Cpu => {}
1276                            }
1277                        }
1278                    }
1279                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1280                    return;
1281                }
1282                if *dtype == TensorDtype::Q1 {
1283                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1284                    // device); the probe keeps whichever beats the CPU matmat.
1285                    if b >= 32
1286                        && b * rows * cols >= 128_000_000
1287                        && cols % 64 == 0
1288                        && crate::gpu::enabled_here()
1289                    {
1290                        if let Self::Mapped { model, idx, .. } = self {
1291                            let t0 = std::time::Instant::now();
1292                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1293                                crate::gpu::ProbeArm::Gpu => {
1294                                    if crate::gpu::q1_matmat(
1295                                        model, *idx, xs_all, b, rows, cols, out,
1296                                    ) {
1297                                        crate::gpu::probe_record(
1298                                            crate::gpu::OpClass::Matmat,
1299                                            true,
1300                                            t0.elapsed(),
1301                                        );
1302                                        return;
1303                                    }
1304                                }
1305                                crate::gpu::ProbeArm::CpuTimed => {
1306                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1307                                    crate::gpu::probe_record(
1308                                        crate::gpu::OpClass::Matmat,
1309                                        false,
1310                                        t0.elapsed(),
1311                                    );
1312                                    return;
1313                                }
1314                                crate::gpu::ProbeArm::Cpu => {}
1315                            }
1316                        }
1317                    }
1318                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1319                    return;
1320                }
1321                if *dtype == TensorDtype::Q1T {
1322                    // GPU batched GEMM for wide prefill (base + overlay on the
1323                    // device); probe keeps the winner vs the CPU matmat.
1324                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1325                        if let Self::Mapped { model, idx, .. } = self {
1326                            let t0 = std::time::Instant::now();
1327                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1328                                crate::gpu::ProbeArm::Gpu => {
1329                                    if crate::gpu::q1t_matmat(
1330                                        model, *idx, xs_all, b, rows, cols, out,
1331                                    ) {
1332                                        crate::gpu::probe_record(
1333                                            crate::gpu::OpClass::Matmat,
1334                                            true,
1335                                            t0.elapsed(),
1336                                        );
1337                                        return;
1338                                    }
1339                                }
1340                                crate::gpu::ProbeArm::CpuTimed => {
1341                                    q1t_matmat(
1342                                        self.quant_bytes(),
1343                                        xs_all,
1344                                        b,
1345                                        rows,
1346                                        cols,
1347                                        out,
1348                                        pool,
1349                                    );
1350                                    crate::gpu::probe_record(
1351                                        crate::gpu::OpClass::Matmat,
1352                                        false,
1353                                        t0.elapsed(),
1354                                    );
1355                                    return;
1356                                }
1357                                crate::gpu::ProbeArm::Cpu => {}
1358                            }
1359                        }
1360                    }
1361                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1362                    return;
1363                }
1364                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1365                    vbitmatmat(
1366                        self.quant_bytes(),
1367                        vbit_offsets,
1368                        xs_all,
1369                        b,
1370                        rows,
1371                        cols,
1372                        out,
1373                        pool,
1374                    );
1375                    return;
1376                }
1377                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1378                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1379                    .collect();
1380                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1381                // work volume: submission carries b×rows×cols MACs).
1382                // Runtime probe: the naive GEMM shader + sync readback
1383                // lose to the CPU GEMM on slow driver stacks — alternate
1384                // both arms and keep the winner.
1385                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1386                    if let Self::Mapped { model, idx, .. } = self {
1387                        let t0 = std::time::Instant::now();
1388                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1389                            crate::gpu::ProbeArm::Gpu
1390                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1391                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1392                            {
1393                                // Cold weights during probing: the upload
1394                                // has started, the count runs on the CPU —
1395                                // the GPU arm samples on the next touch.
1396                                let q = self.quant_bytes();
1397                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1398                                return;
1399                            }
1400                            crate::gpu::ProbeArm::Gpu => {
1401                                let flat: Vec<f32> =
1402                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1403                                if crate::gpu::q8_matmat(
1404                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1405                                ) {
1406                                    crate::gpu::probe_record(
1407                                        crate::gpu::OpClass::Matmat,
1408                                        true,
1409                                        t0.elapsed(),
1410                                    );
1411                                    return;
1412                                }
1413                            }
1414                            crate::gpu::ProbeArm::CpuTimed => {
1415                                let q = self.quant_bytes();
1416                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1417                                crate::gpu::probe_record(
1418                                    crate::gpu::OpClass::Matmat,
1419                                    false,
1420                                    t0.elapsed(),
1421                                );
1422                                return;
1423                            }
1424                            crate::gpu::ProbeArm::Cpu => {}
1425                        }
1426                    }
1427                }
1428                let q = self.quant_bytes();
1429                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1430            }
1431        }
1432    }
1433}
1434
1435impl QTensor {
1436    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1437    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1438    /// barrier instead of N. Per-row math is the exact same kernel as
1439    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1440    /// Falls back to N sequential matvecs when the set is not a uniform
1441    /// q8-family/F32 group or there is no pool.
1442    pub fn matvec_many<const N: usize>(
1443        ts: [&QTensor; N],
1444        x: &[f32],
1445        mut outs: [&mut [f32]; N],
1446        pool: Option<&Pool>,
1447    ) {
1448        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1449        let uniform_q8 = ts.iter().all(|t| {
1450            matches!(
1451                t,
1452                Self::Mapped {
1453                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1454                    ..
1455                }
1456            )
1457        });
1458        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1459        let uniform_q4 = ts.iter().all(|t| {
1460            matches!(
1461                t,
1462                Self::Mapped {
1463                    dtype: TensorDtype::Q4Block,
1464                    ..
1465                }
1466            )
1467        });
1468        let uniform_vbit = ts.iter().all(|t| {
1469            matches!(
1470                t,
1471                Self::Mapped {
1472                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1473                    ..
1474                }
1475            )
1476        });
1477        let uniform_q1 = ts.iter().all(|t| {
1478            matches!(
1479                t,
1480                Self::Mapped {
1481                    dtype: TensorDtype::Q1,
1482                    ..
1483                }
1484            )
1485        });
1486        let uniform_q1t = ts.iter().all(|t| {
1487            matches!(
1488                t,
1489                Self::Mapped {
1490                    dtype: TensorDtype::Q1T,
1491                    ..
1492                }
1493            )
1494        });
1495        let Some(pool) = pool else {
1496            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1497                t.matvec(x, o, None);
1498            }
1499            return;
1500        };
1501        if total_rows < 256
1502            || !(uniform_q8
1503                || uniform_f32
1504                || uniform_q4
1505                || uniform_vbit
1506                || uniform_q1
1507                || uniform_q1t)
1508        {
1509            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1510                t.matvec(x, o, Some(pool));
1511            }
1512            return;
1513        }
1514
1515        if uniform_q1 {
1516            // One shared activation split + group sums (q1 has no col
1517            // field; the same input feeds every tensor).
1518            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1519            if a8w8_enabled() {
1520                let act = split_act(x);
1521                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1522                let (act, gsum) = (&act, &gsum);
1523                let closures: [_; N] = std::array::from_fn(|i| {
1524                    let (bytes, gpr, out) =
1525                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1526                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1527                });
1528                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1529                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1530                pool.run_many(&parts);
1531            } else {
1532                let closures: [_; N] = std::array::from_fn(|i| {
1533                    let (bytes, gpr, out) =
1534                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1535                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1536                });
1537                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1538                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1539                pool.run_many(&parts);
1540            }
1541            return;
1542        }
1543
1544        if uniform_q1t {
1545            // Q1T batched: one shared activation split + overlay decode,
1546            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1547            // and N−1 redundant split_act calls per layer).
1548            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1549            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1550            if a8w8_enabled() {
1551                let act = split_act(x);
1552                let act = &act;
1553                let x_ref = x;
1554                let closures: [_; N] = std::array::from_fn(|i| {
1555                    let bytes = ts[i].quant_bytes();
1556                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1557                    let gpr = cols / GROUP_SIZE;
1558                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1559                    let out = outs_addr[i];
1560                    move |s: usize, e: usize| {
1561                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1562                    }
1563                });
1564                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1565                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1566                pool.run_many(&parts);
1567            } else {
1568                let x_ref = x;
1569                let closures: [_; N] = std::array::from_fn(|i| {
1570                    let bytes = ts[i].quant_bytes();
1571                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1572                    let gpr = cols / GROUP_SIZE;
1573                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1574                    let out = outs_addr[i];
1575                    move |s: usize, e: usize| {
1576                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1577                    }
1578                });
1579                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1580                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1581                pool.run_many(&parts);
1582            }
1583            return;
1584        }
1585
1586        if uniform_q4 || uniform_vbit {
1587            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1588            // q4/vbit share one activation split — no per-tensor col field.
1589            if a8w8_enabled() {
1590                let act = split_act(x);
1591                let act = &act;
1592                if uniform_q4 {
1593                    let closures: [_; N] = std::array::from_fn(|i| {
1594                        let (packed, scales) =
1595                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1596                        let (gpr, cols, out) =
1597                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1598                        move |s: usize, e: usize| {
1599                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1600                        }
1601                    });
1602                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1603                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1604                    pool.run_many(&parts);
1605                } else {
1606                    let closures: [_; N] = std::array::from_fn(|i| {
1607                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1608                            unreachable!()
1609                        };
1610                        let (bytes, rows, cols, out) = (
1611                            ts[i].quant_bytes(),
1612                            ts[i].rows(),
1613                            ts[i].cols(),
1614                            outs_addr[i],
1615                        );
1616                        move |s: usize, e: usize| {
1617                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1618                        }
1619                    });
1620                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1621                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1622                    pool.run_many(&parts);
1623                }
1624                return;
1625            }
1626            if uniform_q4 {
1627                let closures: [_; N] = std::array::from_fn(|i| {
1628                    let (packed, scales) =
1629                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1630                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1631                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1632                });
1633                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1634                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1635                pool.run_many(&parts);
1636            } else {
1637                let closures: [_; N] = std::array::from_fn(|i| {
1638                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1639                        unreachable!()
1640                    };
1641                    let (bytes, rows, cols, out) = (
1642                        ts[i].quant_bytes(),
1643                        ts[i].rows(),
1644                        ts[i].cols(),
1645                        outs_addr[i],
1646                    );
1647                    move |s: usize, e: usize| {
1648                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1649                    }
1650                });
1651                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1652                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1653                pool.run_many(&parts);
1654            }
1655            return;
1656        }
1657
1658        if uniform_f32 {
1659            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1660            let closures: [_; N] = std::array::from_fn(|i| {
1661                let Self::F32 { data, cols, .. } = ts[i] else {
1662                    unreachable!()
1663                };
1664                let out = outs_addr[i];
1665                move |start: usize, end: usize| {
1666                    for o in start..end {
1667                        let row = &data[o * cols..(o + 1) * cols];
1668                        let mut sum = 0.0f32;
1669                        for j in 0..*cols {
1670                            sum += row[j] * x[j];
1671                        }
1672                        // SAFETY: disjoint (tensor, row) cells per worker.
1673                        unsafe { *out.at(o) = sum };
1674                    }
1675                }
1676            });
1677            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1678                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1679            pool.run_many(&parts);
1680            return;
1681        }
1682
1683        // Uniform q8-family: per-tensor prescale (q8_2f col fields
1684        // differ per tensor) + the shared range kernels.
1685        struct Ctx<'a> {
1686            bytes: &'a [u8],
1687            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1688            rep: &'a [u8],
1689            row_scale: &'a [f32],
1690            cols: usize,
1691            xs: std::borrow::Cow<'a, [f32]>,
1692        }
1693        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1694            let Self::Mapped {
1695                dtype,
1696                cols,
1697                row_scale,
1698                col_field,
1699                repack,
1700                ..
1701            } = ts[i]
1702            else {
1703                unreachable!()
1704            };
1705            Ctx {
1706                bytes: ts[i].quant_bytes(),
1707                rep: repack,
1708                row_scale,
1709                cols: *cols,
1710                xs: prescale(x, col_field, *dtype),
1711            }
1712        });
1713        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1714        #[cfg(target_arch = "aarch64")]
1715        if sdot_enabled() {
1716            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1717            let closures: [_; N] = std::array::from_fn(|i| {
1718                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1719                move |start: usize, end: usize| {
1720                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1721                }
1722            });
1723            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1724                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1725            pool.run_many(&parts);
1726            return;
1727        }
1728        #[cfg(target_arch = "x86_64")]
1729        if avx2_a8w8_enabled() {
1730            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1731            let closures: [_; N] = std::array::from_fn(|i| {
1732                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1733                move |start: usize, end: usize| {
1734                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1735                }
1736            });
1737            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1738                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1739            pool.run_many(&parts);
1740            return;
1741        }
1742        let closures: [_; N] = std::array::from_fn(|i| {
1743            let (c, out) = (&ctxs[i], outs_addr[i]);
1744            move |start: usize, end: usize| {
1745                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1746            }
1747        });
1748        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1749            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1750        pool.run_many(&parts);
1751    }
1752}
1753
1754impl QTensor {
1755    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
1756    /// single pool dispatch — the MTP/pair decode path publishes one job
1757    /// for Q/K/V (and one for gate+up) instead of one per tensor.
1758    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
1759    #[allow(clippy::needless_range_loop)]
1760    pub fn matvec2_many<const N: usize>(
1761        ts: [&QTensor; N],
1762        x1: &[f32],
1763        x2: &[f32],
1764        mut o1s: [&mut [f32]; N],
1765        mut o2s: [&mut [f32]; N],
1766        pool: Option<&Pool>,
1767    ) {
1768        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1769        let uniform_q8 = ts.iter().all(|t| {
1770            matches!(
1771                t,
1772                Self::Mapped {
1773                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1774                    ..
1775                }
1776            )
1777        });
1778        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1779        let uniform_q4 = ts.iter().all(|t| {
1780            matches!(
1781                t,
1782                Self::Mapped {
1783                    dtype: TensorDtype::Q4Block,
1784                    ..
1785                }
1786            )
1787        });
1788        let uniform_vbit = ts.iter().all(|t| {
1789            matches!(
1790                t,
1791                Self::Mapped {
1792                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1793                    ..
1794                }
1795            )
1796        });
1797        let fusable = pool.is_some()
1798            && total_rows >= 256
1799            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1800        if !fusable {
1801            for i in 0..N {
1802                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1803            }
1804            return;
1805        }
1806        let pool = pool.unwrap();
1807
1808        if uniform_q4 || uniform_vbit {
1809            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1810            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1811            // q4/vbit share activation splits — no per-tensor col field.
1812            if a8w8_enabled() {
1813                let a1 = split_act(x1);
1814                let a2 = split_act(x2);
1815                let (a1, a2) = (&a1, &a2);
1816                if uniform_q4 {
1817                    let closures: [_; N] = std::array::from_fn(|i| {
1818                        let (packed, scales) =
1819                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1820                        let (gpr, cols, o1, o2) =
1821                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1822                        move |s: usize, e: usize| {
1823                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1824                        }
1825                    });
1826                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1827                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1828                    pool.run_many(&parts);
1829                } else {
1830                    let closures: [_; N] = std::array::from_fn(|i| {
1831                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1832                            unreachable!()
1833                        };
1834                        let (bytes, rows, cols, o1, o2) = (
1835                            ts[i].quant_bytes(),
1836                            ts[i].rows(),
1837                            ts[i].cols(),
1838                            p1[i],
1839                            p2[i],
1840                        );
1841                        move |s: usize, e: usize| {
1842                            vbit_range2_a8w8(
1843                                bytes,
1844                                vbit_offsets,
1845                                x1,
1846                                x2,
1847                                a1,
1848                                a2,
1849                                rows,
1850                                cols,
1851                                o1,
1852                                o2,
1853                                s,
1854                                e,
1855                            )
1856                        }
1857                    });
1858                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1859                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1860                    pool.run_many(&parts);
1861                }
1862                return;
1863            }
1864            if uniform_q4 {
1865                let closures: [_; N] = std::array::from_fn(|i| {
1866                    let (packed, scales) =
1867                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1868                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1869                    move |s: usize, e: usize| {
1870                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1871                    }
1872                });
1873                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1874                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1875                pool.run_many(&parts);
1876            } else {
1877                let closures: [_; N] = std::array::from_fn(|i| {
1878                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1879                        unreachable!()
1880                    };
1881                    let (bytes, rows, cols, o1, o2) = (
1882                        ts[i].quant_bytes(),
1883                        ts[i].rows(),
1884                        ts[i].cols(),
1885                        p1[i],
1886                        p2[i],
1887                    );
1888                    move |s: usize, e: usize| {
1889                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1890                    }
1891                });
1892                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1893                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1894                pool.run_many(&parts);
1895            }
1896            return;
1897        }
1898
1899        if uniform_f32 {
1900            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1901            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1902            let closures: [_; N] = std::array::from_fn(|i| {
1903                let Self::F32 { data, cols, .. } = ts[i] else {
1904                    unreachable!()
1905                };
1906                let (o1, o2) = (p1[i], p2[i]);
1907                move |start: usize, end: usize| {
1908                    for o in start..end {
1909                        let row = &data[o * cols..(o + 1) * cols];
1910                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
1911                        for j in 0..*cols {
1912                            s1 += row[j] * x1[j];
1913                            s2 += row[j] * x2[j];
1914                        }
1915                        // SAFETY: disjoint (tensor, row) cells per worker.
1916                        unsafe {
1917                            *o1.at(o) = s1;
1918                            *o2.at(o) = s2;
1919                        }
1920                    }
1921                }
1922            });
1923            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1924                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1925            pool.run_many(&parts);
1926            return;
1927        }
1928
1929        struct Ctx<'a> {
1930            bytes: &'a [u8],
1931            row_scale: &'a [f32],
1932            cols: usize,
1933            xs1: std::borrow::Cow<'a, [f32]>,
1934            xs2: std::borrow::Cow<'a, [f32]>,
1935        }
1936        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1937            let Self::Mapped {
1938                dtype,
1939                cols,
1940                row_scale,
1941                col_field,
1942                ..
1943            } = ts[i]
1944            else {
1945                unreachable!()
1946            };
1947            Ctx {
1948                bytes: ts[i].quant_bytes(),
1949                row_scale,
1950                cols: *cols,
1951                xs1: prescale(x1, col_field, *dtype),
1952                xs2: prescale(x2, col_field, *dtype),
1953            }
1954        });
1955        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1956        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1957        #[cfg(target_arch = "aarch64")]
1958        if sdot_enabled() {
1959            let acts: [(SplitAct, SplitAct); N] =
1960                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1961            let closures: [_; N] = std::array::from_fn(|i| {
1962                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1963                move |start: usize, end: usize| {
1964                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1965                }
1966            });
1967            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1968                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1969            pool.run_many(&parts);
1970            return;
1971        }
1972        #[cfg(target_arch = "x86_64")]
1973        if avx2_a8w8_enabled() {
1974            let acts: [(SplitAct, SplitAct); N] =
1975                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1976            let closures: [_; N] = std::array::from_fn(|i| {
1977                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1978                move |start: usize, end: usize| {
1979                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1980                }
1981            });
1982            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1983                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1984            pool.run_many(&parts);
1985            return;
1986        }
1987        let closures: [_; N] = std::array::from_fn(|i| {
1988            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1989            move |start: usize, end: usize| {
1990                q8_range2_f32(
1991                    c.bytes,
1992                    c.row_scale,
1993                    &c.xs1,
1994                    &c.xs2,
1995                    c.cols,
1996                    o1,
1997                    o2,
1998                    start,
1999                    end,
2000                )
2001            }
2002        });
2003        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2004            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2005        pool.run_many(&parts);
2006    }
2007
2008    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2009    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2010    /// no intermediate g/u buffers, no separate silu pass. Falls back
2011    /// (returns false) for unsupported dtype combos.
2012    pub fn matvec_silu_mul(
2013        gate: &QTensor,
2014        up: &QTensor,
2015        x: &[f32],
2016        out: &mut [f32],
2017        pool: Option<&Pool>,
2018    ) -> bool {
2019        let inter = gate.rows();
2020        debug_assert_eq!(up.rows(), inter);
2021        debug_assert_eq!(out.len(), inter);
2022        debug_assert_eq!(gate.cols(), up.cols());
2023        if !a8w8_enabled() {
2024            return false;
2025        }
2026        let act = split_act(x);
2027        let act = &act;
2028        let x_ref = x;
2029        let out_addr = SendMut(out.as_mut_ptr());
2030
2031        match (gate, up) {
2032            // Q4Block gate + Q4Block up (most common mobile q4 models)
2033            (
2034                Self::Mapped {
2035                    dtype: TensorDtype::Q4Block,
2036                    ..
2037                },
2038                Self::Mapped {
2039                    dtype: TensorDtype::Q4Block,
2040                    ..
2041                },
2042            ) => {
2043                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2044                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2045                let gpr = gate.cols() / GROUP_SIZE;
2046                let cols = gate.cols();
2047                let run = move |start: usize, end: usize| {
2048                    for r in start..end {
2049                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2050                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2051                        for &(j, xv) in &act.outliers {
2052                            let flat = r * cols + j;
2053                            let gb = gp[flat / 2];
2054                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2055                            let gsc = f16_to_f32(u16::from_le_bytes([
2056                                gs[(flat / GROUP_SIZE) * 2],
2057                                gs[(flat / GROUP_SIZE) * 2 + 1],
2058                            ]));
2059                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2060                            let ub = up_p[flat / 2];
2061                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2062                            let usc = f16_to_f32(u16::from_le_bytes([
2063                                up_s[(flat / GROUP_SIZE) * 2],
2064                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2065                            ]));
2066                            uv += ((un as i32 - 8) as f32) * usc * xv;
2067                        }
2068                        let silu_g = gv / (1.0 + (-gv).exp());
2069                        // SAFETY: disjoint row ranges per worker.
2070                        unsafe { *out_addr.at(r) = silu_g * uv };
2071                    }
2072                };
2073                dispatch_rows(pool, inter, &run);
2074                true
2075            }
2076            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2077            // streams sequential, silu·mul fused (same per-row math as
2078            // `q4t_matvec`).
2079            (
2080                Self::Mapped {
2081                    dtype: TensorDtype::Q4Tiled,
2082                    ..
2083                },
2084                Self::Mapped {
2085                    dtype: TensorDtype::Q4Tiled,
2086                    ..
2087                },
2088            ) => {
2089                let g_bytes = gate.quant_bytes();
2090                let u_bytes = up.quant_bytes();
2091                let gpr = gate.cols() / GROUP_SIZE;
2092                let run = move |start: usize, end: usize| {
2093                    for r in start..end {
2094                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2095                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2096                        for &(j, xv) in &act.outliers {
2097                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2098                            gv += w * s * xv;
2099                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2100                            uv += w * s * xv;
2101                        }
2102                        let silu_g = gv / (1.0 + (-gv).exp());
2103                        // SAFETY: disjoint row ranges per worker.
2104                        unsafe { *out_addr.at(r) = silu_g * uv };
2105                    }
2106                };
2107                dispatch_rows(pool, inter, &run);
2108                true
2109            }
2110            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2111            // each row's two ladders built once and spent on both streams.
2112            (
2113                Self::Mapped {
2114                    dtype: TensorDtype::Q4TiledP,
2115                    ..
2116                },
2117                Self::Mapped {
2118                    dtype: TensorDtype::Q4TiledP,
2119                    ..
2120                },
2121            ) => {
2122                let cols = gate.cols();
2123                let gpr = cols / GROUP_SIZE;
2124                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2125                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2126                let run = |start: usize, end: usize| {
2127                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2128                    for r in start..end {
2129                        gv_view.scales_into(r, gpr, &mut gsc);
2130                        uv_view.scales_into(r, gpr, &mut usc);
2131                        let mut gv =
2132                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2133                        let mut uv =
2134                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2135                        for &(j, xv) in &act.outliers {
2136                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2137                            gv += w * s * xv;
2138                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2139                            uv += w * s * xv;
2140                        }
2141                        let silu_g = gv / (1.0 + (-gv).exp());
2142                        // SAFETY: disjoint row ranges per worker.
2143                        unsafe { *out_addr.at(r) = silu_g * uv };
2144                    }
2145                };
2146                dispatch_rows(pool, inter, &run);
2147                true
2148            }
2149            // Q1T gate + Q1T up
2150            (
2151                Self::Mapped {
2152                    dtype: TensorDtype::Q1T,
2153                    ..
2154                },
2155                Self::Mapped {
2156                    dtype: TensorDtype::Q1T,
2157                    ..
2158                },
2159            ) => {
2160                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2161                let g_bytes = gate.quant_bytes();
2162                let u_bytes = up.quant_bytes();
2163                let gpr = gate.cols() / GROUP_SIZE;
2164                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2165                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2166                let run = move |start: usize, end: usize| {
2167                    for r in start..end {
2168                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2169                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2170                        for &(j, xv) in &act.outliers {
2171                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2172                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2173                        }
2174                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2175                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2176                        let silu_g = gv / (1.0 + (-gv).exp());
2177                        // SAFETY: disjoint row ranges per worker.
2178                        unsafe { *out_addr.at(r) = silu_g * uv };
2179                    }
2180                };
2181                dispatch_rows(pool, inter, &run);
2182                true
2183            }
2184            _ => false,
2185        }
2186    }
2187}
2188
2189/// Batched q8 kernel: same math as qmatvec, the row makes a single
2190/// pass from memory for the whole batch.
2191/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2192/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2193#[cfg(target_os = "macos")]
2194mod accel_blas {
2195    #[link(name = "Accelerate", kind = "framework")]
2196    unsafe extern "C" {
2197        pub fn cblas_sgemm(
2198            order: i32,
2199            trans_a: i32,
2200            trans_b: i32,
2201            m: i32,
2202            n: i32,
2203            k: i32,
2204            alpha: f32,
2205            a: *const f32,
2206            lda: i32,
2207            b: *const f32,
2208            ldb: i32,
2209            beta: f32,
2210            c: *mut f32,
2211            ldc: i32,
2212        );
2213    }
2214}
2215
2216#[cfg(target_os = "macos")]
2217pub(crate) fn accel_gemm_enabled() -> bool {
2218    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2219    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2220}
2221
2222/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2223/// same entry point, so the batched-attention path opens on mobile.
2224#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2225pub(crate) fn accel_gemm_enabled() -> bool {
2226    true
2227}
2228
2229/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2230/// micro-kernel with A broadcast against B panels — the mobile stand-in
2231/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2232/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2233/// k = head_dim or context), and the goal is removing the per-position
2234/// quadratic wall, not peak GEMM.
2235#[cfg(target_arch = "aarch64")]
2236#[allow(clippy::too_many_arguments)]
2237pub(crate) fn neon_gemm_rm(
2238    m: usize,
2239    n: usize,
2240    k: usize,
2241    alpha: f32,
2242    a: &[f32],
2243    lda: usize,
2244    b_mat: &[f32],
2245    ldb: usize,
2246    b_rows_are_n: bool,
2247    c: &mut [f32],
2248    ldc: usize,
2249) {
2250    debug_assert!(a.len() >= (m - 1) * lda + k);
2251    debug_assert!(c.len() >= (m - 1) * ldc + n);
2252    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2253    unsafe {
2254        use core::arch::aarch64::*;
2255        let mut i = 0usize;
2256        while i < m {
2257            let mi = (m - i).min(4);
2258            let mut j = 0usize;
2259            while j < n {
2260                let nj = (n - j).min(8);
2261                if mi == 4 && nj == 8 {
2262                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2263                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2264                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2265                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2266                    for p in 0..k {
2267                        let (b0, b1) = if b_rows_are_n {
2268                            // B is [n, k]: column p of Bᵀ = element p of
2269                            // eight consecutive B rows — gathered.
2270                            let base = b_mat.as_ptr().add(j * ldb + p);
2271                            let g = |o: usize| *base.add(o * ldb);
2272                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2273                        } else {
2274                            let base = b_mat.as_ptr().add(p * ldb + j);
2275                            (
2276                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2277                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2278                            )
2279                        };
2280                        let bv0 = vld1q_f32(b0.as_ptr());
2281                        let bv1 = vld1q_f32(b1.as_ptr());
2282                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2283                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2284                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2285                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2286                        c0a = vfmaq_f32(c0a, a0, bv0);
2287                        c0b = vfmaq_f32(c0b, a0, bv1);
2288                        c1a = vfmaq_f32(c1a, a1, bv0);
2289                        c1b = vfmaq_f32(c1b, a1, bv1);
2290                        c2a = vfmaq_f32(c2a, a2, bv0);
2291                        c2b = vfmaq_f32(c2b, a2, bv1);
2292                        c3a = vfmaq_f32(c3a, a3, bv0);
2293                        c3b = vfmaq_f32(c3b, a3, bv1);
2294                    }
2295                    let al = vdupq_n_f32(alpha);
2296                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2297                        .iter()
2298                        .enumerate()
2299                    {
2300                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2301                        vst1q_f32(dst, vmulq_f32(*ca, al));
2302                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2303                    }
2304                } else {
2305                    for r in 0..mi {
2306                        for q in 0..nj {
2307                            let mut acc = 0f32;
2308                            for p in 0..k {
2309                                let bv = if b_rows_are_n {
2310                                    b_mat[(j + q) * ldb + p]
2311                                } else {
2312                                    b_mat[p * ldb + j + q]
2313                                };
2314                                acc += a[(i + r) * lda + p] * bv;
2315                            }
2316                            c[(i + r) * ldc + j + q] = acc * alpha;
2317                        }
2318                    }
2319                }
2320                j += nj;
2321            }
2322            i += mi;
2323        }
2324    }
2325}
2326
2327/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2328#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2329#[allow(clippy::too_many_arguments)]
2330pub(crate) fn sgemm_rm(
2331    m: usize,
2332    n: usize,
2333    k: usize,
2334    alpha: f32,
2335    a: &[f32],
2336    lda: usize,
2337    b_mat: &[f32],
2338    ldb: usize,
2339    b_rows_are_n: bool,
2340    c: &mut [f32],
2341    ldc: usize,
2342) {
2343    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2344}
2345
2346/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
2347/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
2348#[cfg(target_os = "macos")]
2349#[allow(clippy::too_many_arguments)]
2350pub(crate) fn sgemm_rm(
2351    m: usize,
2352    n: usize,
2353    k: usize,
2354    alpha: f32,
2355    a: &[f32],
2356    lda: usize,
2357    b_mat: &[f32],
2358    ldb: usize,
2359    b_rows_are_n: bool,
2360    c: &mut [f32],
2361    ldc: usize,
2362) {
2363    debug_assert!(a.len() >= (m - 1) * lda + k);
2364    debug_assert!(c.len() >= (m - 1) * ldc + n);
2365    // Test hook: route the attention GEMMs through the portable NEON
2366    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
2367    // measured without a phone in the loop. (Intel macOS has no NEON —
2368    // the hook is a no-op there, Accelerate continues below.)
2369    #[cfg(target_arch = "aarch64")]
2370    if std::env::var("CMF_FORCE_NEON_GEMM")
2371        .map(|v| v == "1")
2372        .unwrap_or(false)
2373    {
2374        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2375    }
2376    unsafe {
2377        accel_blas::cblas_sgemm(
2378            101, // RowMajor
2379            111, // NoTrans A
2380            if b_rows_are_n { 112 } else { 111 },
2381            m as i32,
2382            n as i32,
2383            k as i32,
2384            alpha,
2385            a.as_ptr(),
2386            lda as i32,
2387            b_mat.as_ptr(),
2388            ldb as i32,
2389            0.0,
2390            c.as_mut_ptr(),
2391            ldc as i32,
2392        );
2393    }
2394}
2395
2396/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
2397/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
2398/// on the AMX with one row-major sgemm. Tiles live in cache, weights
2399/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
2400/// logits shift within f32 rounding — tolerance-class, like every
2401/// reduction-order change; decode (M=1) never takes this path.
2402#[cfg(target_os = "macos")]
2403fn qmatmat_accel(
2404    q: &[u8],
2405    row_scale: &[f32],
2406    pre: &[std::borrow::Cow<'_, [f32]>],
2407    rows: usize,
2408    cols: usize,
2409    out: &mut [f32],
2410    pool: Option<&Pool>,
2411) {
2412    // NOTE: double-buffering the dequant against the sgemm (a scoped
2413    // thread driving the pool on tile k+1 while the caller multiplies
2414    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
2415    // multithreaded, and the dequant workers just steal its cores.
2416    const TR: usize = 2048;
2417    let b = pre.len();
2418    thread_local! {
2419        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2420        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2421    }
2422    XPANEL.with(|xp| {
2423        WTILE.with(|wt| {
2424            let mut xpanel = xp.borrow_mut();
2425            xpanel.clear();
2426            for x in pre {
2427                xpanel.extend_from_slice(x);
2428            }
2429            let mut wtile = wt.borrow_mut();
2430            wtile.resize(TR * cols, 0.0);
2431            let mut r0 = 0usize;
2432            while r0 < rows {
2433                let tr = TR.min(rows - r0);
2434                // Dequant the tile (scale folded) — pool-parallel.
2435                let wt_addr = SendMut(wtile.as_mut_ptr());
2436                let run = |start: usize, end: usize| {
2437                    for r in start..end {
2438                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2439                        let s = row_scale[r0 + r];
2440                        // SAFETY: workers cover disjoint r ranges.
2441                        let dst =
2442                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2443                        for (d, &v) in dst.iter_mut().zip(row) {
2444                            *d = (v as i8) as f32 * s;
2445                        }
2446                    }
2447                };
2448                dispatch_rows(pool, tr, &run);
2449                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
2450                unsafe {
2451                    accel_blas::cblas_sgemm(
2452                        101, // RowMajor
2453                        111, // NoTrans A
2454                        112, // Trans B
2455                        b as i32,
2456                        tr as i32,
2457                        cols as i32,
2458                        1.0,
2459                        xpanel.as_ptr(),
2460                        cols as i32,
2461                        wtile.as_ptr(),
2462                        cols as i32,
2463                        0.0,
2464                        out.as_mut_ptr().add(r0),
2465                        rows as i32,
2466                    );
2467                }
2468                r0 += tr;
2469            }
2470        })
2471    });
2472}
2473
2474fn qmatmat(
2475    q: &[u8],
2476    row_scale: &[f32],
2477    pre: &[std::borrow::Cow<'_, [f32]>],
2478    rows: usize,
2479    cols: usize,
2480    out: &mut [f32],
2481    pool: Option<&Pool>,
2482) {
2483    let b = pre.len();
2484    debug_assert_eq!(out.len(), b * rows);
2485    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
2486    // SDOT loop below peaks near the CPU's dot throughput, an order
2487    // below the matrix units. Small tensors and tiny test models stay
2488    // on the exact integer path.
2489    #[cfg(target_os = "macos")]
2490    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2491        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2492        return;
2493    }
2494    #[cfg(target_arch = "aarch64")]
2495    if sdot_enabled() {
2496        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2497        let out_addr = SendMut(out.as_mut_ptr());
2498        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
2499        // path IS the ARM prefill GEMM off Apple silicon).
2500        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2501            .map(|v| v != "0")
2502            .unwrap_or(true);
2503        let use_i8mm = i8mm_enabled();
2504        if blocked_ok {
2505            let run = |start: usize, end: usize| {
2506                let mut o = start;
2507                while o < end {
2508                    if o + 2 <= end {
2509                        let r0 = &q[o * cols..(o + 1) * cols];
2510                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2511                        let mut bi = 0usize;
2512                        while bi + 4 <= acts.len() {
2513                            let xs = [
2514                                acts[bi].xq.as_slice(),
2515                                acts[bi + 1].xq.as_slice(),
2516                                acts[bi + 2].xq.as_slice(),
2517                                acts[bi + 3].xq.as_slice(),
2518                            ];
2519                            let d = if use_i8mm {
2520                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2521                            } else {
2522                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2523                            };
2524                            for (r, row) in [r0, r1].into_iter().enumerate() {
2525                                for k in 0..4 {
2526                                    let act = &acts[bi + k];
2527                                    let mut v = d[r][k] as f32 * act.sx;
2528                                    for &(j, xv) in &act.outliers {
2529                                        v += (row[j] as i8) as f32 * xv;
2530                                    }
2531                                    unsafe {
2532                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2533                                    };
2534                                }
2535                            }
2536                            bi += 4;
2537                        }
2538                        while bi < acts.len() {
2539                            for (r, row) in [r0, r1].into_iter().enumerate() {
2540                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2541                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2542                            }
2543                            bi += 1;
2544                        }
2545                        o += 2;
2546                    } else {
2547                        let row = &q[o * cols..(o + 1) * cols];
2548                        for (bi, act) in acts.iter().enumerate() {
2549                            let v = row_dot_sdot(row, act) * row_scale[o];
2550                            unsafe { *out_addr.at(bi * rows + o) = v };
2551                        }
2552                        o += 1;
2553                    }
2554                }
2555            };
2556            dispatch_rows(pool, rows, &run);
2557            return;
2558        }
2559        let run = |start: usize, end: usize| {
2560            for o in start..end {
2561                let row = &q[o * cols..(o + 1) * cols];
2562                for (bi, act) in acts.iter().enumerate() {
2563                    let v = row_dot_sdot(row, act) * row_scale[o];
2564                    unsafe { *out_addr.at(bi * rows + o) = v };
2565                }
2566            }
2567        };
2568        dispatch_rows(pool, rows, &run);
2569        return;
2570    }
2571    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
2572    // (roadmap P0: two weight rows' abs() stay in registers across four
2573    // activation streams); VNNI machines keep the per-row bias-trick
2574    // dot, which is already throughput-bound there.
2575    #[cfg(target_arch = "x86_64")]
2576    if avx2_a8w8_enabled() {
2577        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2578        let out_addr = SendMut(out.as_mut_ptr());
2579        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
2580        // A/B on noisy shared-vCPU hosts).
2581        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2582            .map(|v| v != "0")
2583            .unwrap_or(true);
2584        if !avx512vnni_enabled() && blocked_ok {
2585            let run = |start: usize, end: usize| {
2586                let mut o = start;
2587                while o < end {
2588                    if o + 2 <= end {
2589                        let r0 = &q[o * cols..(o + 1) * cols];
2590                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2591                        let mut bi = 0usize;
2592                        while bi + 4 <= acts.len() {
2593                            let xs = [
2594                                acts[bi].xq.as_slice(),
2595                                acts[bi + 1].xq.as_slice(),
2596                                acts[bi + 2].xq.as_slice(),
2597                                acts[bi + 3].xq.as_slice(),
2598                            ];
2599                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
2600                            for (r, row) in [r0, r1].into_iter().enumerate() {
2601                                for k in 0..4 {
2602                                    let act = &acts[bi + k];
2603                                    let mut v = d[r][k] as f32 * act.sx;
2604                                    for &(j, xv) in &act.outliers {
2605                                        v += (row[j] as i8) as f32 * xv;
2606                                    }
2607                                    unsafe {
2608                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2609                                    };
2610                                }
2611                            }
2612                            bi += 4;
2613                        }
2614                        while bi < acts.len() {
2615                            for (r, row) in [r0, r1].into_iter().enumerate() {
2616                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
2617                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2618                            }
2619                            bi += 1;
2620                        }
2621                        o += 2;
2622                    } else {
2623                        let row = &q[o * cols..(o + 1) * cols];
2624                        for (bi, act) in acts.iter().enumerate() {
2625                            let v = row_dot_avx2(row, act) * row_scale[o];
2626                            unsafe { *out_addr.at(bi * rows + o) = v };
2627                        }
2628                        o += 1;
2629                    }
2630                }
2631            };
2632            dispatch_rows(pool, rows, &run);
2633            return;
2634        }
2635        let run = |start: usize, end: usize| {
2636            for o in start..end {
2637                let row = &q[o * cols..(o + 1) * cols];
2638                for (bi, act) in acts.iter().enumerate() {
2639                    let v = row_dot_avx2(row, act) * row_scale[o];
2640                    unsafe { *out_addr.at(bi * rows + o) = v };
2641                }
2642            }
2643        };
2644        dispatch_rows(pool, rows, &run);
2645        return;
2646    }
2647    let out_addr = SendMut(out.as_mut_ptr());
2648    let run = |start: usize, end: usize| {
2649        for o in start..end {
2650            let row = &q[o * cols..(o + 1) * cols];
2651            for (bi, x) in pre.iter().enumerate() {
2652                let mut acc = 0f32;
2653                for j in 0..cols {
2654                    acc += (row[j] as i8) as f32 * x[j];
2655                }
2656                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
2657            }
2658        }
2659    };
2660    dispatch_rows(pool, rows, &run);
2661}
2662
2663/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
2664/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
2665fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
2666    match pool {
2667        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
2668        _ => run(0, rows),
2669    }
2670}
2671
2672/// Split a q4_block blob into (packed nibbles, f16 group scales).
2673fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
2674    let groups = rows * cols / GROUP_SIZE;
2675    bytes.split_at(groups * 16)
2676}
2677
2678/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
2679/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
2680/// vbit packs MSB-first, so the HIGH nibble is the even element
2681/// (opposite of q4_block's lo-first interleave). Centering is u-7.
2682#[inline]
2683fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
2684    #[cfg(target_arch = "aarch64")]
2685    unsafe {
2686        return vbit_fill4_neon(data, buf);
2687    }
2688    #[cfg(target_arch = "x86_64")]
2689    if avx2_enabled() {
2690        return unsafe { vbit_fill4_avx2(data, buf) };
2691    }
2692    #[allow(unreachable_code)]
2693    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2694        let u = unpack8::<4>(&data[blk * 4..]);
2695        for k in 0..8 {
2696            chunk[k] = (u[k] - 7) as i8 as u8;
2697        }
2698    }
2699}
2700
2701#[cfg(target_arch = "aarch64")]
2702#[target_feature(enable = "neon")]
2703unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
2704    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
2705    // buf.len()/2 packed bytes (validated at load).
2706    unsafe {
2707        use core::arch::aarch64::*;
2708        let n = buf.len();
2709        let mask = vdupq_n_u8(0x0F);
2710        let seven = vdupq_n_s8(7);
2711        let mut g = 0usize;
2712        while g * 32 + 32 <= n {
2713            let b = vld1q_u8(data.as_ptr().add(g * 16));
2714            let hi = vshrq_n_u8::<4>(b);
2715            let lo = vandq_u8(b, mask);
2716            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
2717            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
2718            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
2719            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
2720            g += 1;
2721        }
2722    }
2723}
2724
2725#[cfg(target_arch = "x86_64")]
2726#[target_feature(enable = "avx2")]
2727unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2728    // SAFETY: see vbit_fill4_neon.
2729    unsafe {
2730        use core::arch::x86_64::*;
2731        let n = buf.len();
2732        let mask = _mm_set1_epi8(0x0F);
2733        let seven = _mm256_set1_epi8(7);
2734        let mut g = 0usize;
2735        while g * 32 + 32 <= n {
2736            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2737            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2738            let lo = _mm_and_si128(b, mask);
2739            let z = _mm256_sub_epi8(
2740                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2741                seven,
2742            );
2743            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2744            g += 1;
2745        }
2746    }
2747}
2748
2749/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
2750/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
2751/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
2752/// into 4 such blocks.
2753#[inline(always)]
2754fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2755    let mut acc = 0u64;
2756    for i in 0..B {
2757        acc = (acc << 8) | data[i] as u64;
2758    }
2759    let mask = (1u64 << B) - 1;
2760    let mut out = [0i32; 8];
2761    for (k, o) in out.iter_mut().enumerate() {
2762        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2763    }
2764    out
2765}
2766
2767/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
2768/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
2769/// MSB-first, byte-padded]. Row data offsets are precomputed at load
2770/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
2771/// overhead on every matvec.
2772#[allow(clippy::too_many_arguments)]
2773fn vbitmatvec(
2774    bytes: &[u8],
2775    offsets: &[usize],
2776    x: &[f32],
2777    rows: usize,
2778    cols: usize,
2779    out: &mut [f32],
2780    pool: Option<&Pool>,
2781) {
2782    debug_assert_eq!(out.len(), rows);
2783    debug_assert_eq!(offsets.len(), rows + 1);
2784
2785    // SDOT path: unpack the row to centered i8 once, then per-group
2786    // int8 dot against the quantized activations — same A8W8 contract
2787    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
2788    if a8w8_enabled() {
2789        let act = split_act(x);
2790        let out_addr = SendMut(out.as_mut_ptr());
2791        let run = move |start: usize, end: usize| {
2792            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2793        };
2794        dispatch_rows(pool, rows, &run);
2795        return;
2796    }
2797
2798    let out_addr = SendMut(out.as_mut_ptr());
2799    let run = move |start: usize, end: usize| {
2800        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2801    };
2802    dispatch_rows(pool, rows, &run);
2803}
2804
2805/// One vbit row range via the A8W8 int8 path — kernel body of
2806/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
2807/// several tensors in one dispatch (b=8 rows go exact f32).
2808#[allow(clippy::too_many_arguments)]
2809fn vbit_range_a8w8(
2810    bytes: &[u8],
2811    offsets: &[usize],
2812    x: &[f32],
2813    act: &SplitAct,
2814    rows: usize,
2815    cols: usize,
2816    out: SendMut,
2817    start: usize,
2818    end: usize,
2819) {
2820    let ng = cols / GROUP_SIZE;
2821    let bits = &bytes[..rows];
2822    let sc_off = rows;
2823    let row_dot = |r: usize| -> f32 {
2824        let b = bits[r] as usize;
2825        let l = (1i32 << (b - 1)) - 1;
2826        let mask = (1u64 << b) - 1;
2827        let data = &bytes[offsets[r]..offsets[r + 1]];
2828        if b == 8 {
2829            // u−L reaches 128 → does not fit i8; exact f32 path.
2830            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2831            let mut dot = 0f32;
2832            for g in 0..ng {
2833                let so = (r * ng + g) * 2;
2834                let sgf = f16_to_f32(u16::from_le_bytes([
2835                    bytes[sc_off + so],
2836                    bytes[sc_off + so + 1],
2837                ]));
2838                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2839                let mut gd = 0f32;
2840                for &xv in xg.iter() {
2841                    if nbits < 8 {
2842                        acc = (acc << 8) | data[idx] as u64;
2843                        idx += 1;
2844                        nbits += 8;
2845                    }
2846                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2847                    nbits -= 8;
2848                    gd += (u - l) as f32 * xv;
2849                }
2850                dot += gd * sgf;
2851            }
2852            return dot;
2853        }
2854        // Per-worker scratch: this closure runs for every row of the
2855        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
2856        // row was measurable pure overhead.
2857        thread_local! {
2858            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2859                const { std::cell::RefCell::new(Vec::new()) };
2860        }
2861        #[inline(always)]
2862        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2863            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2864                let u = unpack8::<B>(&data[blk * B..]);
2865                for k in 0..8 {
2866                    chunk[k] = (u[k] - l) as i8 as u8;
2867                }
2868            }
2869        }
2870        let _ = mask;
2871        VBIT_SCRATCH.with(|scratch| {
2872            let mut buf = scratch.borrow_mut();
2873            buf.resize(cols, 0);
2874            match b {
2875                3 => fill::<3>(data, l, &mut buf),
2876                4 => vbit_fill4(data, &mut buf),
2877                5 => fill::<5>(data, l, &mut buf),
2878                6 => fill::<6>(data, l, &mut buf),
2879                _ => unreachable!(),
2880            }
2881            let mut dot = 0f32;
2882            for g in 0..ng {
2883                let so = (r * ng + g) * 2;
2884                let s = f16_to_f32(u16::from_le_bytes([
2885                    bytes[sc_off + so],
2886                    bytes[sc_off + so + 1],
2887                ]));
2888                let d = dot_i8_i8(
2889                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2890                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2891                ) as f32
2892                    * act.sx;
2893                dot += d * s;
2894            }
2895            for &(j, xv) in &act.outliers {
2896                let so = (r * ng + j / GROUP_SIZE) * 2;
2897                let s = f16_to_f32(u16::from_le_bytes([
2898                    bytes[sc_off + so],
2899                    bytes[sc_off + so + 1],
2900                ]));
2901                // xq is zeroed at outlier slots — add the exact term.
2902                dot += (buf[j] as i8) as f32 * s * xv;
2903            }
2904            dot
2905        })
2906    };
2907    for r in start..end {
2908        // SAFETY: disjoint row ranges per worker.
2909        unsafe { *out.at(r) = row_dot(r) };
2910    }
2911}
2912
2913/// Exact scalar vbit row range (same extraction, non-SDOT path).
2914#[allow(clippy::too_many_arguments)]
2915fn vbit_range_f32(
2916    bytes: &[u8],
2917    offsets: &[usize],
2918    x: &[f32],
2919    rows: usize,
2920    cols: usize,
2921    out: SendMut,
2922    start: usize,
2923    end: usize,
2924) {
2925    let ng = cols / GROUP_SIZE;
2926    let bits = &bytes[..rows];
2927    let sc_off = rows;
2928    // Per-bit-width specialized inner loops: the compiler unrolls the
2929    // constant shifts (the generic bit-buffer loop was branch-bound —
2930    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
2931    #[inline(always)]
2932    fn dot_row<const B: usize>(
2933        data: &[u8],
2934        bytes: &[u8],
2935        sc_off: usize,
2936        r: usize,
2937        ng: usize,
2938        x: &[f32],
2939    ) -> f32 {
2940        let l = ((1i32 << (B - 1)) - 1) as f32;
2941        let gbytes = GROUP_SIZE * B / 8;
2942        let mut dot = 0f32;
2943        for g in 0..ng {
2944            let so = (r * ng + g) * 2;
2945            let s = f16_to_f32(u16::from_le_bytes([
2946                bytes[sc_off + so],
2947                bytes[sc_off + so + 1],
2948            ]));
2949            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2950            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2951            let mut gd = 0f32;
2952            for blk in 0..GROUP_SIZE / 8 {
2953                let u = unpack8::<B>(&gd0[blk * B..]);
2954                let xb = &xg[blk * 8..blk * 8 + 8];
2955                for k in 0..8 {
2956                    gd += (u[k] as f32 - l) * xb[k];
2957                }
2958            }
2959            dot += gd * s;
2960        }
2961        dot
2962    }
2963    for r in start..end {
2964        let data = &bytes[offsets[r]..offsets[r + 1]];
2965        let v = match bits[r] {
2966            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2967            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2968            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2969            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2970            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2971            b => unreachable!("vbit bit-width {b} (validated at load)"),
2972        };
2973        // SAFETY: disjoint row ranges per worker.
2974        unsafe { *out.at(r) = v };
2975    }
2976}
2977
2978/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
2979/// and dotted against BOTH activations (MTP verify / pair prefill used
2980/// to run two full matvecs — double weight traffic and double unpack).
2981/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
2982#[allow(clippy::too_many_arguments)]
2983fn vbitmatvec2(
2984    bytes: &[u8],
2985    offsets: &[usize],
2986    x1: &[f32],
2987    x2: &[f32],
2988    rows: usize,
2989    cols: usize,
2990    o1: &mut [f32],
2991    o2: &mut [f32],
2992    pool: Option<&Pool>,
2993) {
2994    debug_assert_eq!(o1.len(), rows);
2995    debug_assert_eq!(o2.len(), rows);
2996
2997    if a8w8_enabled() {
2998        let a1 = split_act(x1);
2999        let a2 = split_act(x2);
3000        let p1 = SendMut(o1.as_mut_ptr());
3001        let p2 = SendMut(o2.as_mut_ptr());
3002        let run = move |start: usize, end: usize| {
3003            vbit_range2_a8w8(
3004                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3005            )
3006        };
3007        dispatch_rows(pool, rows, &run);
3008        return;
3009    }
3010
3011    let p1 = SendMut(o1.as_mut_ptr());
3012    let p2 = SendMut(o2.as_mut_ptr());
3013    let run = move |start: usize, end: usize| {
3014        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3015    };
3016    dispatch_rows(pool, rows, &run);
3017}
3018
3019/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3020/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3021/// exact f32 for both lanes, bits streamed once).
3022#[allow(clippy::too_many_arguments)]
3023fn vbit_range2_a8w8(
3024    bytes: &[u8],
3025    offsets: &[usize],
3026    x1: &[f32],
3027    x2: &[f32],
3028    a1: &SplitAct,
3029    a2: &SplitAct,
3030    rows: usize,
3031    cols: usize,
3032    p1: SendMut,
3033    p2: SendMut,
3034    start: usize,
3035    end: usize,
3036) {
3037    let ng = cols / GROUP_SIZE;
3038    let bits = &bytes[..rows];
3039    let sc_off = rows;
3040    let row_dots = |r: usize| -> (f32, f32) {
3041        let b = bits[r] as usize;
3042        let l = (1i32 << (b - 1)) - 1;
3043        let data = &bytes[offsets[r]..offsets[r + 1]];
3044        if b == 8 {
3045            // u−L reaches 128 → does not fit i8; exact f32 path,
3046            // bits still streamed once for both lanes.
3047            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3048            let (mut d1, mut d2) = (0f32, 0f32);
3049            for g in 0..ng {
3050                let so = (r * ng + g) * 2;
3051                let sgf = f16_to_f32(u16::from_le_bytes([
3052                    bytes[sc_off + so],
3053                    bytes[sc_off + so + 1],
3054                ]));
3055                let (mut g1, mut g2) = (0f32, 0f32);
3056                for k in 0..GROUP_SIZE {
3057                    if nbits < 8 {
3058                        acc = (acc << 8) | data[idx] as u64;
3059                        idx += 1;
3060                        nbits += 8;
3061                    }
3062                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3063                    nbits -= 8;
3064                    let w = (u - l) as f32;
3065                    g1 += w * x1[g * GROUP_SIZE + k];
3066                    g2 += w * x2[g * GROUP_SIZE + k];
3067                }
3068                d1 += g1 * sgf;
3069                d2 += g2 * sgf;
3070            }
3071            return (d1, d2);
3072        }
3073        thread_local! {
3074            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3075                const { std::cell::RefCell::new(Vec::new()) };
3076        }
3077        #[inline(always)]
3078        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3079            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3080                let u = unpack8::<B>(&data[blk * B..]);
3081                for k in 0..8 {
3082                    chunk[k] = (u[k] - l) as i8 as u8;
3083                }
3084            }
3085        }
3086        VBIT_SCRATCH2.with(|scratch| {
3087            let mut buf = scratch.borrow_mut();
3088            buf.resize(cols, 0);
3089            match b {
3090                3 => fill::<3>(data, l, &mut buf),
3091                4 => vbit_fill4(data, &mut buf),
3092                5 => fill::<5>(data, l, &mut buf),
3093                6 => fill::<6>(data, l, &mut buf),
3094                _ => unreachable!(),
3095            }
3096            let (mut d1, mut d2) = (0f32, 0f32);
3097            for g in 0..ng {
3098                let so = (r * ng + g) * 2;
3099                let s = f16_to_f32(u16::from_le_bytes([
3100                    bytes[sc_off + so],
3101                    bytes[sc_off + so + 1],
3102                ]));
3103                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3104                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3105                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3106                d1 += v1 * s;
3107                d2 += v2 * s;
3108            }
3109            for &(j, xv) in &a1.outliers {
3110                let so = (r * ng + j / GROUP_SIZE) * 2;
3111                let s = f16_to_f32(u16::from_le_bytes([
3112                    bytes[sc_off + so],
3113                    bytes[sc_off + so + 1],
3114                ]));
3115                d1 += (buf[j] as i8) as f32 * s * xv;
3116            }
3117            for &(j, xv) in &a2.outliers {
3118                let so = (r * ng + j / GROUP_SIZE) * 2;
3119                let s = f16_to_f32(u16::from_le_bytes([
3120                    bytes[sc_off + so],
3121                    bytes[sc_off + so + 1],
3122                ]));
3123                d2 += (buf[j] as i8) as f32 * s * xv;
3124            }
3125            (d1, d2)
3126        })
3127    };
3128    for r in start..end {
3129        let (v1, v2) = row_dots(r);
3130        // SAFETY: disjoint row ranges per worker.
3131        unsafe {
3132            *p1.at(r) = v1;
3133            *p2.at(r) = v2;
3134        }
3135    }
3136}
3137
3138/// Two-input exact scalar vbit row range (same extraction) —
3139/// per-bit-width specialized, two accumulators per row; per-lane
3140/// accumulation order matches `vbitmatvec` exactly.
3141#[allow(clippy::too_many_arguments)]
3142fn vbit_range2_f32(
3143    bytes: &[u8],
3144    offsets: &[usize],
3145    x1: &[f32],
3146    x2: &[f32],
3147    rows: usize,
3148    cols: usize,
3149    p1: SendMut,
3150    p2: SendMut,
3151    start: usize,
3152    end: usize,
3153) {
3154    let ng = cols / GROUP_SIZE;
3155    let bits = &bytes[..rows];
3156    let sc_off = rows;
3157    #[inline(always)]
3158    #[allow(clippy::too_many_arguments)]
3159    fn dot_row2<const B: usize>(
3160        data: &[u8],
3161        bytes: &[u8],
3162        sc_off: usize,
3163        r: usize,
3164        ng: usize,
3165        x1: &[f32],
3166        x2: &[f32],
3167    ) -> (f32, f32) {
3168        let l = ((1i32 << (B - 1)) - 1) as f32;
3169        let gbytes = GROUP_SIZE * B / 8;
3170        let (mut d1, mut d2) = (0f32, 0f32);
3171        for g in 0..ng {
3172            let so = (r * ng + g) * 2;
3173            let s = f16_to_f32(u16::from_le_bytes([
3174                bytes[sc_off + so],
3175                bytes[sc_off + so + 1],
3176            ]));
3177            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3178            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3179            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3180            let (mut g1, mut g2) = (0f32, 0f32);
3181            for blk in 0..GROUP_SIZE / 8 {
3182                let u = unpack8::<B>(&gd0[blk * B..]);
3183                for k in 0..8 {
3184                    let w = u[k] as f32 - l;
3185                    g1 += w * x1g[blk * 8 + k];
3186                    g2 += w * x2g[blk * 8 + k];
3187                }
3188            }
3189            d1 += g1 * s;
3190            d2 += g2 * s;
3191        }
3192        (d1, d2)
3193    }
3194    for r in start..end {
3195        let data = &bytes[offsets[r]..offsets[r + 1]];
3196        let (v1, v2) = match bits[r] {
3197            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3198            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3199            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3200            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3201            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3202            b => unreachable!("vbit bit-width {b} (validated at load)"),
3203        };
3204        // SAFETY: disjoint row ranges per worker.
3205        unsafe {
3206            *p1.at(r) = v1;
3207            *p2.at(r) = v2;
3208        }
3209    }
3210}
3211
3212// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3213
3214/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3215/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3216/// distant streams of the split layout. Values/order identical to the
3217/// split kernels.
3218#[inline]
3219#[allow(unreachable_code)]
3220fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3221    #[cfg(target_arch = "aarch64")]
3222    unsafe {
3223        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3224    }
3225    #[cfg(target_arch = "x86_64")]
3226    unsafe {
3227        if vnni_tiles_enabled() {
3228            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3229        }
3230        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3231    }
3232    let mut acc = 0f32;
3233    for gi in 0..gpr {
3234        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3235        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3236        let mut d = 0i32;
3237        for (k, &b) in tile[2..].iter().enumerate() {
3238            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3239                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3240        }
3241        acc += d as f32 * s;
3242    }
3243    acc
3244}
3245
3246#[cfg(target_arch = "aarch64")]
3247#[target_feature(enable = "neon,dotprod")]
3248unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3249    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3250    // xq.len() == gpr·GROUP_SIZE).
3251    unsafe {
3252        use core::arch::aarch64::*;
3253        use core::arch::asm;
3254        let lomask = vdupq_n_u8(0x0F);
3255        let eight = vdupq_n_s8(8);
3256        let mut acc = 0f32;
3257        for gi in 0..gpr {
3258            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3259            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3260            let b = vld1q_u8(t.add(2));
3261            let lo = vandq_u8(b, lomask);
3262            let hi = vshrq_n_u8::<4>(b);
3263            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3264            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3265            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3266            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3267            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3268            asm!(
3269                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3270                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3271                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3272                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3273                options(pure, nomem, nostack),
3274            );
3275            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3276        }
3277        acc
3278    }
3279}
3280
3281#[cfg(target_arch = "x86_64")]
3282#[target_feature(enable = "avx2")]
3283unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3284    // SAFETY: see dot_q4t_row_sdot.
3285    unsafe {
3286        use core::arch::x86_64::*;
3287        let lomask = _mm_set1_epi8(0x0F);
3288        let eight = _mm256_set1_epi8(8);
3289        let ones = _mm256_set1_epi16(1);
3290        let mut acc = 0f32;
3291        for gi in 0..gpr {
3292            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3293            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3294            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3295            let lo = _mm_and_si128(b, lomask);
3296            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3297            let w = _mm256_sub_epi8(
3298                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3299                eight,
3300            );
3301            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3302            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3303            let d = _mm256_madd_epi16(p16, ones);
3304            let hi128 = _mm256_extracti128_si256::<1>(d);
3305            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3306            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3307            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3308            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3309        }
3310        acc
3311    }
3312}
3313
3314/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
3315/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
3316/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
3317#[cfg(target_arch = "x86_64")]
3318#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3319unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3320    // SAFETY: see dot_q4t_row_sdot.
3321    unsafe {
3322        use core::arch::x86_64::*;
3323        let lomask = _mm_set1_epi8(0x0F);
3324        let eight = _mm256_set1_epi8(8);
3325        let mut acc = 0f32;
3326        for gi in 0..gpr {
3327            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3328            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3329            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3330            let lo = _mm_and_si128(b, lomask);
3331            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3332            let w = _mm256_sub_epi8(
3333                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3334                eight,
3335            );
3336            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3337            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3338            acc += d as f32 * s;
3339        }
3340        acc
3341    }
3342}
3343
3344/// One q4_tiled row against FOUR activation streams: the nibble unpack
3345/// and abs() happen once per group instead of once per (group,
3346/// activation) — the unpack is the dominant per-element cost of the
3347/// tiled format (roadmap P0 portable blocking, q4t leg).
3348#[cfg(target_arch = "x86_64")]
3349// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
3350// to a libm call per lane — measured 2x slower than the reduction this
3351// kernel replaces. The runtime gate (`avx2_enabled`) already requires
3352// both features, so declaring it here is safe.
3353#[target_feature(enable = "avx2,fma")]
3354unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3355    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3356    unsafe {
3357        use core::arch::x86_64::*;
3358        let lomask = _mm_set1_epi8(0x0F);
3359        let eight = _mm256_set1_epi8(8);
3360        let ones = _mm256_set1_epi16(1);
3361        // One f32 accumulator VECTOR per activation, reduced once at the
3362        // end. Folding each group's i32 lanes to a scalar inside the loop
3363        // costs an extracti128 + three shift/add + a movd — a cross-lane
3364        // dependency chain per (group, activation), 288 of them per row at
3365        // cols=2304. The per-group scale is what forces a float
3366        // accumulator; it does not force a horizontal sum.
3367        //
3368        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
3369        // indexed by a loop variable LLVM keeps them in memory and every
3370        // group pays four 32-byte loads and stores. That alone made this
3371        // kernel 2x SLOWER than the per-group reduction it replaces
3372        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
3373        let mut f0 = _mm256_setzero_ps();
3374        let mut f1 = _mm256_setzero_ps();
3375        let mut f2 = _mm256_setzero_ps();
3376        let mut f3 = _mm256_setzero_ps();
3377        for gi in 0..gpr {
3378            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3379            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3380            let sv = _mm256_set1_ps(s);
3381            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3382            let lo = _mm_and_si128(bb, lomask);
3383            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3384            let w = _mm256_sub_epi8(
3385                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3386                eight,
3387            );
3388            let aw = _mm256_abs_epi8(w);
3389            let off = gi * GROUP_SIZE;
3390            let dot = |xq: &[i8]| {
3391                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3392                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3393                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3394            };
3395            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3396            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3397            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3398            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3399        }
3400        [
3401            hsum256_ps(f0),
3402            hsum256_ps(f1),
3403            hsum256_ps(f2),
3404            hsum256_ps(f3),
3405        ]
3406    }
3407}
3408
3409/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
3410/// blocked kernels pay, once per row instead of once per group.
3411#[cfg(target_arch = "x86_64")]
3412#[target_feature(enable = "avx2")]
3413#[inline]
3414unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3415    // SAFETY: pure register arithmetic on the caller's vector.
3416    unsafe {
3417        use core::arch::x86_64::*;
3418        let hi = _mm256_extractf128_ps::<1>(v);
3419        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3420        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3421        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3422        _mm_cvtss_f32(s)
3423    }
3424}
3425
3426/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
3427#[cfg(target_arch = "x86_64")]
3428#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3429unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3430    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3431    unsafe {
3432        use core::arch::x86_64::*;
3433        let lomask = _mm_set1_epi8(0x0F);
3434        let eight = _mm256_set1_epi8(8);
3435        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
3436        // one cross-lane reduction per row, not per (group, activation).
3437        let mut f0 = _mm256_setzero_ps();
3438        let mut f1 = _mm256_setzero_ps();
3439        let mut f2 = _mm256_setzero_ps();
3440        let mut f3 = _mm256_setzero_ps();
3441        for gi in 0..gpr {
3442            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3443            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3444            let sv = _mm256_set1_ps(s);
3445            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3446            let lo = _mm_and_si128(bb, lomask);
3447            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3448            let w = _mm256_sub_epi8(
3449                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3450                eight,
3451            );
3452            let aw = _mm256_abs_epi8(w);
3453            let off = gi * GROUP_SIZE;
3454            let dot = |xq: &[i8]| {
3455                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3456                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3457                    _mm256_setzero_si256(),
3458                    aw,
3459                    _mm256_sign_epi8(x, w),
3460                ))
3461            };
3462            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3463            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3464            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3465            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3466        }
3467        let acc = [
3468            hsum256_ps(f0),
3469            hsum256_ps(f1),
3470            hsum256_ps(f2),
3471            hsum256_ps(f3),
3472        ];
3473        acc
3474    }
3475}
3476
3477/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
3478/// serves FOUR activation streams. Per stream the group order and f32
3479/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
3480/// bit-for-bit.
3481#[cfg(target_arch = "aarch64")]
3482#[target_feature(enable = "neon,dotprod")]
3483unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3484    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3485    unsafe {
3486        use core::arch::aarch64::*;
3487        use core::arch::asm;
3488        let lomask = vdupq_n_u8(0x0F);
3489        let eight = vdupq_n_s8(8);
3490        let mut acc = [0f32; 4];
3491        for gi in 0..gpr {
3492            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3493            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3494            let b = vld1q_u8(t.add(2));
3495            let lo = vandq_u8(b, lomask);
3496            let hi = vshrq_n_u8::<4>(b);
3497            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3498            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3499            for (k, xq) in xs.iter().enumerate() {
3500                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3501                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3502                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3503                asm!(
3504                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3505                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3506                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3507                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3508                    options(pure, nomem, nostack),
3509                );
3510                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3511            }
3512        }
3513        acc
3514    }
3515}
3516
3517/// Exact-term correction for A8W8 outliers on a tiled row.
3518#[inline]
3519fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3520    let gi = j / GROUP_SIZE;
3521    let k = j % GROUP_SIZE;
3522    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3523    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3524    let byte = tile[2 + k / 2];
3525    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3526    ((nib as i32 - 8) as f32, s)
3527}
3528
3529/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
3530/// accumulation shape as `q4_range_f32`.
3531#[inline]
3532fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3533    let mut acc = 0f32;
3534    for gi in 0..gpr {
3535        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3536        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3537        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3538        let mut ga = 0f32;
3539        for (k, &b) in tile[2..].iter().enumerate() {
3540            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3541                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3542        }
3543        acc += ga * s;
3544    }
3545    acc
3546}
3547
3548/// Split view of a `q4tp` payload. The three planes are resolved once per
3549/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
3550/// the row loop would put a division on the hot path for nothing.
3551struct Q4tpView<'a> {
3552    nib: &'a [u8],
3553    params: &'a [u8],
3554    codes: &'a [u8],
3555    stride: usize,
3556}
3557
3558impl<'a> Q4tpView<'a> {
3559    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
3560        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
3561        Self {
3562            nib: &bytes[..params_off],
3563            params: &bytes[params_off..codes_off],
3564            codes: &bytes[codes_off..],
3565            stride,
3566        }
3567    }
3568
3569    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
3570    ///
3571    /// Doing this once per row — rather than decoding a 5-bit code inside the
3572    /// tile loop — is what makes the format free at runtime. Random access to
3573    /// a packed 5-bit field costs a division, two bounds checks and a branch;
3574    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
3575    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
3576    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
3577    #[inline]
3578    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
3579        let tab = q4tp_ladder(self.params, r);
3580        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
3581        let (mut acc, mut have, mut bi) = (0u64, 0u32, 0usize);
3582        for o in out.iter_mut().take(gpr) {
3583            while have < 5 {
3584                acc |= (codes[bi] as u64) << have;
3585                bi += 1;
3586                have += 8;
3587            }
3588            *o = tab[(acc & 31) as usize];
3589            acc >>= 5;
3590            have -= 5;
3591        }
3592    }
3593}
3594
3595#[inline]
3596fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
3597    #[cfg(target_arch = "aarch64")]
3598    unsafe {
3599        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
3600    }
3601    #[cfg(target_arch = "x86_64")]
3602    unsafe {
3603        if vnni_tiles_enabled() {
3604            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
3605        }
3606        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
3607    }
3608    #[allow(unreachable_code)]
3609    {
3610        let mut acc = 0f32;
3611        for gi in 0..gpr {
3612            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3613            let s = scales[gi];
3614            let mut d = 0i32;
3615            for (k, &b) in tile.iter().enumerate() {
3616                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3617                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3618            }
3619            acc += d as f32 * s;
3620        }
3621        acc
3622    }
3623}
3624
3625/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
3626/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
3627#[cfg(target_arch = "aarch64")]
3628#[target_feature(enable = "neon,dotprod")]
3629unsafe fn dot_q4tp_row_sdot(
3630    nib: &[u8],
3631    r: usize,
3632    gpr: usize,
3633    xq: &[i8],
3634    scales: &[f32],
3635) -> f32 {
3636    // SAFETY: callers uphold slice-length contracts (16B tile per group,
3637    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
3638    unsafe {
3639        use core::arch::aarch64::*;
3640        use core::arch::asm;
3641        let lomask = vdupq_n_u8(0x0F);
3642        let eight = vdupq_n_s8(8);
3643        let mut acc = 0f32;
3644        for gi in 0..gpr {
3645            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3646            let s = *scales.get_unchecked(gi);
3647            let b = vld1q_u8(t);
3648            let lo = vandq_u8(b, lomask);
3649            let hi = vshrq_n_u8::<4>(b);
3650            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3651            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3652            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3653            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3654            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3655            asm!(
3656                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3657                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3658                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3659                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3660                options(pure, nomem, nostack),
3661            );
3662            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3663        }
3664        acc
3665    }
3666}
3667
3668#[cfg(target_arch = "x86_64")]
3669#[target_feature(enable = "avx2")]
3670unsafe fn dot_q4tp_row_avx2(
3671    nib: &[u8],
3672    r: usize,
3673    gpr: usize,
3674    xq: &[i8],
3675    scales: &[f32],
3676) -> f32 {
3677    // SAFETY: see dot_q4tp_row_sdot.
3678    unsafe {
3679        use core::arch::x86_64::*;
3680        let lomask = _mm_set1_epi8(0x0F);
3681        let eight = _mm256_set1_epi8(8);
3682        let ones = _mm256_set1_epi16(1);
3683        let mut acc = 0f32;
3684        for gi in 0..gpr {
3685            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3686            let s = *scales.get_unchecked(gi);
3687            let b = _mm_loadu_si128(t as *const __m128i);
3688            let lo = _mm_and_si128(b, lomask);
3689            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3690            let w = _mm256_sub_epi8(
3691                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3692                eight,
3693            );
3694            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3695            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3696            let d = _mm256_madd_epi16(p16, ones);
3697            let hi128 = _mm256_extracti128_si256::<1>(d);
3698            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3699            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3700            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3701            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3702        }
3703        acc
3704    }
3705}
3706
3707/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
3708/// 256-bit VL encoding is the one to use here).
3709#[cfg(target_arch = "x86_64")]
3710#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3711unsafe fn dot_q4tp_row_vnni(
3712    nib: &[u8],
3713    r: usize,
3714    gpr: usize,
3715    xq: &[i8],
3716    scales: &[f32],
3717) -> f32 {
3718    // SAFETY: see dot_q4tp_row_sdot.
3719    unsafe {
3720        use core::arch::x86_64::*;
3721        let lomask = _mm_set1_epi8(0x0F);
3722        let eight = _mm256_set1_epi8(8);
3723        let mut acc = 0f32;
3724        for gi in 0..gpr {
3725            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3726            let s = *scales.get_unchecked(gi);
3727            let b = _mm_loadu_si128(t as *const __m128i);
3728            let lo = _mm_and_si128(b, lomask);
3729            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3730            let w = _mm256_sub_epi8(
3731                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3732                eight,
3733            );
3734            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3735            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
3736        }
3737        acc
3738    }
3739}
3740
3741/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
3742/// accumulation shape as `q4t_row_exact`.
3743#[inline]
3744fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
3745    let mut acc = 0f32;
3746    for gi in 0..gpr {
3747        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3748        let s = scales[gi];
3749        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3750        let mut ga = 0f32;
3751        for (k, &b) in tile.iter().enumerate() {
3752            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3753                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3754        }
3755        acc += ga * s;
3756    }
3757    acc
3758}
3759
3760/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
3761/// activation outliers at full precision after the int8 pass.
3762#[inline]
3763fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
3764    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
3765    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
3766    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3767    ((n as i32 - 8) as f32, scales[gi])
3768}
3769
3770/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
3771fn q4tp_matvec(
3772    bytes: &[u8],
3773    x: &[f32],
3774    rows: usize,
3775    cols: usize,
3776    out: &mut [f32],
3777    pool: Option<&Pool>,
3778) {
3779    debug_assert_eq!(out.len(), rows);
3780    let gpr = cols / GROUP_SIZE;
3781    let v = Q4tpView::new(bytes, rows, cols);
3782    let out_addr = SendMut(out.as_mut_ptr());
3783    if a8w8_enabled() {
3784        let act = split_act(x);
3785        let run = |start: usize, end: usize| {
3786            // One scratch row of scales per worker, reused across its rows.
3787            let mut sc = vec![0f32; gpr];
3788            for r in start..end {
3789                v.scales_into(r, gpr, &mut sc);
3790                let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
3791                for &(j, xv) in &act.outliers {
3792                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3793                    acc += w * s * xv;
3794                }
3795                // SAFETY: disjoint row ranges per worker.
3796                unsafe { *out_addr.at(r) = acc };
3797            }
3798        };
3799        dispatch_rows(pool, rows, &run);
3800        return;
3801    }
3802    let run = |start: usize, end: usize| {
3803        let mut sc = vec![0f32; gpr];
3804        for r in start..end {
3805            v.scales_into(r, gpr, &mut sc);
3806            // SAFETY: disjoint row ranges per worker.
3807            unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
3808        }
3809    };
3810    dispatch_rows(pool, rows, &run);
3811}
3812
3813/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
3814/// row ladder are read once and spent on both activation streams.
3815#[allow(clippy::too_many_arguments)]
3816fn q4tp_matvec2(
3817    bytes: &[u8],
3818    x1: &[f32],
3819    x2: &[f32],
3820    rows: usize,
3821    cols: usize,
3822    o1: &mut [f32],
3823    o2: &mut [f32],
3824    pool: Option<&Pool>,
3825) {
3826    let gpr = cols / GROUP_SIZE;
3827    let v = Q4tpView::new(bytes, rows, cols);
3828    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
3829    let run = |start: usize, end: usize| {
3830        let mut sc = vec![0f32; gpr];
3831        for r in start..end {
3832            v.scales_into(r, gpr, &mut sc);
3833            // SAFETY: disjoint row ranges per worker.
3834            unsafe {
3835                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
3836                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
3837            }
3838        }
3839    };
3840    dispatch_rows(pool, rows, &run);
3841}
3842
3843/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
3844/// spent on four activation streams, which is where a prefill batch stops
3845/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
3846#[cfg(target_arch = "aarch64")]
3847#[target_feature(enable = "neon,dotprod")]
3848unsafe fn dot_q4tp_row_1x4_sdot(
3849    nib: &[u8],
3850    r: usize,
3851    gpr: usize,
3852    xs: [&[i8]; 4],
3853    scales: &[f32],
3854) -> [f32; 4] {
3855    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
3856    unsafe {
3857        use core::arch::aarch64::*;
3858        use core::arch::asm;
3859        let lomask = vdupq_n_u8(0x0F);
3860        let eight = vdupq_n_s8(8);
3861        // Named accumulators, NOT an array indexed by a loop variable: the
3862        // latter does not stay in registers (the same defect cost 2x in the
3863        // AVX2 q4t kernel and again in WGSL).
3864        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
3865        for gi in 0..gpr {
3866            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3867            let s = *scales.get_unchecked(gi);
3868            let bb = vld1q_u8(t);
3869            let lo = vandq_u8(bb, lomask);
3870            let hi = vshrq_n_u8::<4>(bb);
3871            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3872            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3873            let mut d = [0f32; 4];
3874            for (k, dk) in d.iter_mut().enumerate() {
3875                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
3876                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
3877                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3878                asm!(
3879                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3880                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3881                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3882                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3883                    options(pure, nomem, nostack),
3884                );
3885                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3886            }
3887            f0 += d[0];
3888            f1 += d[1];
3889            f2 += d[2];
3890            f3 += d[3];
3891        }
3892        [f0, f1, f2, f3]
3893    }
3894}
3895
3896/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
3897/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
3898/// the format was fine, the missing arms were the whole regression.
3899fn q4tp_matmat(
3900    bytes: &[u8],
3901    xs_all: &[f32],
3902    b: usize,
3903    rows: usize,
3904    cols: usize,
3905    out: &mut [f32],
3906    pool: Option<&Pool>,
3907) {
3908    debug_assert_eq!(out.len(), b * rows);
3909    let gpr = cols / GROUP_SIZE;
3910    let v = Q4tpView::new(bytes, rows, cols);
3911
3912    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
3913    #[cfg(target_os = "macos")]
3914    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3915        dequant_matmat_accel(
3916            &|r, dst| {
3917                let mut sc = [0f32; 32];
3918                let mut scv;
3919                let s: &[f32] = if gpr <= 32 {
3920                    v.scales_into(r, gpr, &mut sc);
3921                    &sc[..gpr]
3922                } else {
3923                    scv = vec![0f32; gpr];
3924                    v.scales_into(r, gpr, &mut scv);
3925                    &scv
3926                };
3927                for gi in 0..gpr {
3928                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3929                    for (k, &bb) in tile.iter().enumerate() {
3930                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
3931                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
3932                    }
3933                }
3934            },
3935            xs_all,
3936            b,
3937            rows,
3938            cols,
3939            out,
3940            pool,
3941        );
3942        return;
3943    }
3944
3945    let out_addr = SendMut(out.as_mut_ptr());
3946    if a8w8_enabled() {
3947        let acts: Vec<SplitAct> = (0..b)
3948            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
3949            .collect();
3950        let acts = &acts;
3951        #[cfg(target_arch = "aarch64")]
3952        let blocked_ok = sdot_enabled()
3953            && std::env::var("CMF_X86_BLOCKED")
3954                .map(|val| val != "0")
3955                .unwrap_or(true);
3956        #[cfg(not(target_arch = "aarch64"))]
3957        let blocked_ok = false;
3958        let run = |start: usize, end: usize| {
3959            let mut sc = vec![0f32; gpr];
3960            for r in start..end {
3961                v.scales_into(r, gpr, &mut sc);
3962                let mut bi = 0usize;
3963                #[cfg(target_arch = "aarch64")]
3964                if blocked_ok {
3965                    while bi + 4 <= acts.len() {
3966                        let xs = [
3967                            acts[bi].xq.as_slice(),
3968                            acts[bi + 1].xq.as_slice(),
3969                            acts[bi + 2].xq.as_slice(),
3970                            acts[bi + 3].xq.as_slice(),
3971                        ];
3972                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc) };
3973                        for k in 0..4 {
3974                            let act = &acts[bi + k];
3975                            let mut acc = d[k] * act.sx;
3976                            for &(j, xv) in &act.outliers {
3977                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3978                                acc += w * s * xv;
3979                            }
3980                            // SAFETY: disjoint (bi, r) cells per worker.
3981                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3982                        }
3983                        bi += 4;
3984                    }
3985                }
3986                let _ = blocked_ok;
3987                while bi < acts.len() {
3988                    let act = &acts[bi];
3989                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
3990                    for &(j, xv) in &act.outliers {
3991                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3992                        acc += w * s * xv;
3993                    }
3994                    // SAFETY: disjoint (bi, r) cells per worker range.
3995                    unsafe { *out_addr.at(bi * rows + r) = acc };
3996                    bi += 1;
3997                }
3998            }
3999        };
4000        dispatch_rows(pool, rows, &run);
4001        return;
4002    }
4003
4004    let run = |start: usize, end: usize| {
4005        let mut sc = vec![0f32; gpr];
4006        for r in start..end {
4007            v.scales_into(r, gpr, &mut sc);
4008            for bi in 0..b {
4009                let x = &xs_all[bi * cols..(bi + 1) * cols];
4010                // SAFETY: disjoint (bi, r) cells per worker range.
4011                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4012            }
4013        }
4014    };
4015    dispatch_rows(pool, rows, &run);
4016}
4017
4018/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
4019fn q4t_matvec(
4020    bytes: &[u8],
4021    x: &[f32],
4022    rows: usize,
4023    cols: usize,
4024    out: &mut [f32],
4025    pool: Option<&Pool>,
4026) {
4027    debug_assert_eq!(out.len(), rows);
4028    let gpr = cols / GROUP_SIZE;
4029    let out_addr = SendMut(out.as_mut_ptr());
4030    if a8w8_enabled() {
4031        let act = split_act(x);
4032        let run = move |start: usize, end: usize| {
4033            for r in start..end {
4034                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4035                for &(j, xv) in &act.outliers {
4036                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4037                    acc += w * s * xv;
4038                }
4039                // SAFETY: disjoint row ranges per worker.
4040                unsafe { *out_addr.at(r) = acc };
4041            }
4042        };
4043        dispatch_rows(pool, rows, &run);
4044        return;
4045    }
4046    let run = move |start: usize, end: usize| {
4047        for r in start..end {
4048            // SAFETY: disjoint row ranges per worker.
4049            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
4050        }
4051    };
4052    dispatch_rows(pool, rows, &run);
4053}
4054
4055/// Fused two-input q4_tiled matvec (weights read once per pair).
4056#[allow(clippy::too_many_arguments)]
4057fn q4t_matvec2(
4058    bytes: &[u8],
4059    x1: &[f32],
4060    x2: &[f32],
4061    rows: usize,
4062    cols: usize,
4063    o1: &mut [f32],
4064    o2: &mut [f32],
4065    pool: Option<&Pool>,
4066) {
4067    let gpr = cols / GROUP_SIZE;
4068    let p1 = SendMut(o1.as_mut_ptr());
4069    let p2 = SendMut(o2.as_mut_ptr());
4070    if a8w8_enabled() {
4071        let a1 = split_act(x1);
4072        let a2 = split_act(x2);
4073        let run = move |start: usize, end: usize| {
4074            for r in start..end {
4075                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
4076                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
4077                for &(j, xv) in &a1.outliers {
4078                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4079                    v1 += w * s * xv;
4080                }
4081                for &(j, xv) in &a2.outliers {
4082                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4083                    v2 += w * s * xv;
4084                }
4085                // SAFETY: disjoint row ranges per worker.
4086                unsafe {
4087                    *p1.at(r) = v1;
4088                    *p2.at(r) = v2;
4089                }
4090            }
4091        };
4092        dispatch_rows(pool, rows, &run);
4093        return;
4094    }
4095    let run = move |start: usize, end: usize| {
4096        for r in start..end {
4097            // SAFETY: disjoint row ranges per worker.
4098            unsafe {
4099                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
4100                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
4101            }
4102        }
4103    };
4104    dispatch_rows(pool, rows, &run);
4105}
4106
4107/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
4108#[allow(clippy::too_many_arguments)]
4109/// Prefill GEMM through Accelerate for group-quantized codecs: a
4110/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
4111/// each tile rides the AMX with one sgemm — the generic sibling of
4112/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
4113/// decode (b=1) never takes this path.
4114#[cfg(target_os = "macos")]
4115fn dequant_matmat_accel(
4116    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
4117    xs_all: &[f32],
4118    b: usize,
4119    rows: usize,
4120    cols: usize,
4121    out: &mut [f32],
4122    pool: Option<&Pool>,
4123) {
4124    const TR: usize = 2048;
4125    thread_local! {
4126        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4127    }
4128    WTILE.with(|wt| {
4129        let mut wtile = wt.borrow_mut();
4130        wtile.resize(TR * cols, 0.0);
4131        let mut r0 = 0usize;
4132        while r0 < rows {
4133            let tr = TR.min(rows - r0);
4134            let wt_addr = SendMut(wtile.as_mut_ptr());
4135            let run = |start: usize, end: usize| {
4136                for r in start..end {
4137                    // SAFETY: workers cover disjoint r ranges.
4138                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4139                    dequant_row(r0 + r, dst);
4140                }
4141            };
4142            dispatch_rows(pool, tr, &run);
4143            unsafe {
4144                accel_blas::cblas_sgemm(
4145                    101, // RowMajor
4146                    111, // NoTrans A
4147                    112, // Trans B
4148                    b as i32,
4149                    tr as i32,
4150                    cols as i32,
4151                    1.0,
4152                    xs_all.as_ptr(),
4153                    cols as i32,
4154                    wtile.as_ptr(),
4155                    cols as i32,
4156                    0.0,
4157                    out.as_mut_ptr().add(r0),
4158                    rows as i32,
4159                );
4160            }
4161            r0 += tr;
4162        }
4163    });
4164}
4165
4166fn q4t_matmat(
4167    bytes: &[u8],
4168    xs_all: &[f32],
4169    b: usize,
4170    rows: usize,
4171    cols: usize,
4172    out: &mut [f32],
4173    pool: Option<&Pool>,
4174) {
4175    debug_assert_eq!(out.len(), b * rows);
4176    let gpr = cols / GROUP_SIZE;
4177    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
4178    // the dequant-tile sgemm is an order above the SDOT row loop for
4179    // prefill shapes (imagegen DiT forwards are exactly this).
4180    #[cfg(target_os = "macos")]
4181    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4182        dequant_matmat_accel(
4183            &|r, dst| {
4184                for gi in 0..gpr {
4185                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4186                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4187                    for (k, &bb) in tile[2..].iter().enumerate() {
4188                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
4189                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
4190                    }
4191                }
4192            },
4193            xs_all,
4194            b,
4195            rows,
4196            cols,
4197            out,
4198            pool,
4199        );
4200        return;
4201    }
4202    let out_addr = SendMut(out.as_mut_ptr());
4203    if a8w8_enabled() {
4204        let acts: Vec<SplitAct> = (0..b)
4205            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4206            .collect();
4207        let acts = &acts;
4208        #[cfg(target_arch = "x86_64")]
4209        let blocked_ok = avx2_enabled()
4210            && std::env::var("CMF_X86_BLOCKED")
4211                .map(|v| v != "0")
4212                .unwrap_or(true);
4213        #[cfg(target_arch = "aarch64")]
4214        let blocked_ok = sdot_enabled()
4215            && std::env::var("CMF_X86_BLOCKED")
4216                .map(|v| v != "0")
4217                .unwrap_or(true);
4218        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
4219        let blocked_ok = false;
4220        let run = move |start: usize, end: usize| {
4221            for r in start..end {
4222                let mut bi = 0usize;
4223                #[cfg(target_arch = "aarch64")]
4224                if blocked_ok {
4225                    while bi + 4 <= acts.len() {
4226                        let xs = [
4227                            acts[bi].xq.as_slice(),
4228                            acts[bi + 1].xq.as_slice(),
4229                            acts[bi + 2].xq.as_slice(),
4230                            acts[bi + 3].xq.as_slice(),
4231                        ];
4232                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
4233                        for k in 0..4 {
4234                            let act = &acts[bi + k];
4235                            let mut acc = d[k] * act.sx;
4236                            for &(j, xv) in &act.outliers {
4237                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4238                                acc += w * sc * xv;
4239                            }
4240                            // SAFETY: disjoint (bi, r) cells per worker.
4241                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4242                        }
4243                        bi += 4;
4244                    }
4245                }
4246                #[cfg(target_arch = "x86_64")]
4247                if blocked_ok {
4248                    while bi + 4 <= acts.len() {
4249                        let xs = [
4250                            acts[bi].xq.as_slice(),
4251                            acts[bi + 1].xq.as_slice(),
4252                            acts[bi + 2].xq.as_slice(),
4253                            acts[bi + 3].xq.as_slice(),
4254                        ];
4255                        let d = unsafe {
4256                            if vnni_tiles_enabled() {
4257                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
4258                            } else {
4259                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
4260                            }
4261                        };
4262                        for k in 0..4 {
4263                            let act = &acts[bi + k];
4264                            let mut acc = d[k] * act.sx;
4265                            for &(j, xv) in &act.outliers {
4266                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4267                                acc += w * sc * xv;
4268                            }
4269                            // SAFETY: disjoint (bi, r) cells per worker.
4270                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4271                        }
4272                        bi += 4;
4273                    }
4274                }
4275                let _ = blocked_ok;
4276                while bi < acts.len() {
4277                    let act = &acts[bi];
4278                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4279                    for &(j, xv) in &act.outliers {
4280                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
4281                        acc += w * s * xv;
4282                    }
4283                    // SAFETY: disjoint (bi, r) cells per worker range.
4284                    unsafe { *out_addr.at(bi * rows + r) = acc };
4285                    bi += 1;
4286                }
4287            }
4288        };
4289        dispatch_rows(pool, rows, &run);
4290        return;
4291    }
4292    let run = move |start: usize, end: usize| {
4293        for r in start..end {
4294            for bi in 0..b {
4295                let x = &xs_all[bi * cols..(bi + 1) * cols];
4296                // SAFETY: disjoint (bi, r) cells per worker range.
4297                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
4298            }
4299        }
4300    };
4301    dispatch_rows(pool, rows, &run);
4302}
4303
4304// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
4305// 32-group tile. The kernel family mirrors q4_tiled: one sequential
4306// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
4307// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
4308
4309/// Per-32-group sums of the quantized activation — the ±1 identity's
4310/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
4311/// matvec and reused by every row.
4312fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
4313    (0..gpr)
4314        .map(|gi| {
4315            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
4316                .iter()
4317                .map(|&v| v as i32)
4318                .sum()
4319        })
4320        .collect()
4321}
4322
4323/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
4324/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
4325/// x86 pass).
4326#[inline]
4327#[allow(unreachable_code)]
4328/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
4329/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
4330/// masked activation sums through maddubs(1, x&mask), and
4331/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
4332#[cfg(target_arch = "x86_64")]
4333#[target_feature(enable = "avx2")]
4334unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4335    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4336    unsafe {
4337        use core::arch::x86_64::*;
4338        // Byte j of the mask must replicate bits-byte j/8.
4339        let expand = _mm256_setr_epi8(
4340            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,
4341            3, 3, 3,
4342        );
4343        let bitsel = _mm256_setr_epi8(
4344            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4345            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4346        );
4347        let ones8 = _mm256_set1_epi8(1);
4348        let ones16 = _mm256_set1_epi16(1);
4349        let mut acc = 0f32;
4350        for gi in 0..gpr {
4351            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4352            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4353            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4354            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4355            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4356            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4357            let sel = _mm256_and_si256(x, mask);
4358            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
4359            let p16 = _mm256_maddubs_epi16(ones8, sel);
4360            let d32 = _mm256_madd_epi16(p16, ones16);
4361            let hi128 = _mm256_extracti128_si256::<1>(d32);
4362            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4363            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4364            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4365            let msum = _mm_cvtsi128_si32(s32);
4366            // The and-select keeps x UN-negated (unlike ARM's −1-mask
4367            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
4368            let d = 2 * msum - gsum[gi];
4369            acc += d as f32 * s;
4370        }
4371        acc
4372    }
4373}
4374
4375/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
4376/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
4377#[cfg(target_arch = "x86_64")]
4378#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4379unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4380    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4381    unsafe {
4382        use core::arch::x86_64::*;
4383        let expand = _mm256_setr_epi8(
4384            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,
4385            3, 3, 3,
4386        );
4387        let bitsel = _mm256_setr_epi8(
4388            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4389            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4390        );
4391        let ones8 = _mm256_set1_epi8(1);
4392        let mut acc = 0f32;
4393        for gi in 0..gpr {
4394            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4395            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4396            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4397            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4398            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4399            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4400            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4401            let d = 2 * msum - gsum[gi];
4402            acc += d as f32 * s;
4403        }
4404        acc
4405    }
4406}
4407
4408/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
4409#[cfg(target_arch = "x86_64")]
4410#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4411unsafe fn dot_q1_row_1x4_vnni(
4412    bytes: &[u8],
4413    r: usize,
4414    gpr: usize,
4415    xs: [&[i8]; 4],
4416    gsums: [&[i32]; 4],
4417) -> [f32; 4] {
4418    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4419    unsafe {
4420        use core::arch::x86_64::*;
4421        let expand = _mm256_setr_epi8(
4422            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,
4423            3, 3, 3,
4424        );
4425        let bitsel = _mm256_setr_epi8(
4426            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4427            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4428        );
4429        let ones8 = _mm256_set1_epi8(1);
4430        let mut acc = [0f32; 4];
4431        for gi in 0..gpr {
4432            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4433            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4434            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4435            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4436            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4437            for (k, xq) in xs.iter().enumerate() {
4438                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4439                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4440                let d = 2 * msum - gsums[k][gi];
4441                acc[k] += d as f32 * s;
4442            }
4443        }
4444        acc
4445    }
4446}
4447
4448/// The blocked 1×4 flavor: the expanded bit mask serves four activation
4449/// streams per group (mask build once, four select+reduce chains).
4450#[cfg(target_arch = "x86_64")]
4451#[target_feature(enable = "avx2")]
4452unsafe fn dot_q1_row_1x4_avx2(
4453    bytes: &[u8],
4454    r: usize,
4455    gpr: usize,
4456    xs: [&[i8]; 4],
4457    gsums: [&[i32]; 4],
4458) -> [f32; 4] {
4459    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4460    unsafe {
4461        use core::arch::x86_64::*;
4462        let expand = _mm256_setr_epi8(
4463            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,
4464            3, 3, 3,
4465        );
4466        let bitsel = _mm256_setr_epi8(
4467            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4468            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4469        );
4470        let ones8 = _mm256_set1_epi8(1);
4471        let ones16 = _mm256_set1_epi16(1);
4472        let mut acc = [0f32; 4];
4473        for gi in 0..gpr {
4474            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4475            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4476            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4477            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4478            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4479            for (k, xq) in xs.iter().enumerate() {
4480                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4481                let sel = _mm256_and_si256(x, mask);
4482                let p16 = _mm256_maddubs_epi16(ones8, sel);
4483                let d32 = _mm256_madd_epi16(p16, ones16);
4484                let hi128 = _mm256_extracti128_si256::<1>(d32);
4485                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4486                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4487                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4488                let msum = _mm_cvtsi128_si32(s32);
4489                let d = 2 * msum - gsums[k][gi];
4490                acc[k] += d as f32 * s;
4491            }
4492        }
4493        acc
4494    }
4495}
4496
4497#[allow(unreachable_code)]
4498fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4499    #[cfg(target_arch = "aarch64")]
4500    unsafe {
4501        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
4502    }
4503    #[cfg(target_arch = "x86_64")]
4504    if avx2_enabled() {
4505        unsafe {
4506            if vnni_tiles_enabled() {
4507                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
4508            }
4509            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
4510        }
4511    }
4512    let _ = gsum;
4513    let mut acc = 0f32;
4514    for gi in 0..gpr {
4515        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4516        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4517        let mut d = 0i32;
4518        for (j, &b) in tile[2..].iter().enumerate() {
4519            for k in 0..8 {
4520                let w = ((b >> k) & 1) as i32 * 2 - 1;
4521                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
4522            }
4523        }
4524        acc += d as f32 * s;
4525    }
4526    acc
4527}
4528
4529/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
4530/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
4531/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
4532/// per-group activation sums shared across every row of the matvec.
4533/// Four tiles (128 weights) per iteration: integer dots reduce through
4534/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
4535/// fused f32 multiply-add. Integer math throughout — bit-identical to
4536/// the scalar ±1 reference.
4537#[cfg(target_arch = "aarch64")]
4538#[target_feature(enable = "neon,dotprod")]
4539unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4540    // SAFETY: callers uphold slice-length contracts (6B tile per group,
4541    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
4542    unsafe {
4543        use core::arch::aarch64::*;
4544        use core::arch::asm;
4545        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4546        let m = vld1q_u8(MASKS.as_ptr());
4547        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
4548        macro_rules! tile_dot {
4549            ($t:expr, $x:expr) => {{
4550                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
4551                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
4552                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4553                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4554                let x0 = vld1q_s8($x);
4555                let x1 = vld1q_s8($x.add(16));
4556                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4557                asm!(
4558                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4559                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4560                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4561                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4562                    options(pure, nomem, nostack),
4563                );
4564                vaddq_s32(a0, a1)
4565            }};
4566        }
4567        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
4568        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
4569        // bit-byte across 8 lanes for vtst, and the four scales gather
4570        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
4571        // 4 branchy software f16 conversions per 128 weights (the
4572        // measured load-port wall of this kernel) become 2 vector
4573        // loads + 9 table lookups. Integer math order is unchanged —
4574        // bit-identical results (FCVTL is exact on every f16).
4575        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4576        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4577        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4578        const IW11: [u8; 16] = [
4579            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4580        ];
4581        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4582        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4583        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4584        let isc = vld1_u8(ISC.as_ptr());
4585        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
4586        macro_rules! tile_dot_tbl {
4587            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
4588                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
4589                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
4590                let x0 = vld1q_s8($x);
4591                let x1 = vld1q_s8($x.add(16));
4592                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4593                asm!(
4594                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4595                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4596                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4597                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4598                    options(pure, nomem, nostack),
4599                );
4600                vaddq_s32(a0, a1)
4601            }};
4602        }
4603        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4604        let row_base = r * gpr * Q1_TILE;
4605        let abs_end = bytes.len();
4606        let xp = xq.as_ptr();
4607        let gp = gsum.as_ptr();
4608        let mut accv = vdupq_n_f32(0.0);
4609        let mut gi = 0;
4610        // The second pair load reads 4B past tile gi+3 — stay inside
4611        // the payload slice (only the file's final tiles fall back).
4612        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4613            let t0 = base.add(gi * Q1_TILE);
4614            let ld_a = vld1q_u8(t0);
4615            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4616            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
4617            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
4618            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
4619            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
4620            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
4621            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4622            let g = vld1q_s32(gp.add(gi));
4623            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4624            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4625            let scf: float32x4_t;
4626            asm!(
4627                "fcvtl {o:v}.4s, {i:v}.4h",
4628                o = out(vreg) scf, i = in(vreg) sc16,
4629                options(pure, nomem, nostack),
4630            );
4631            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
4632            gi += 4;
4633        }
4634        let mut acc = vaddvq_f32(accv);
4635        while gi < gpr {
4636            let t = base.add(gi * Q1_TILE);
4637            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4638            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
4639            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
4640            gi += 1;
4641        }
4642        acc
4643    }
4644}
4645
4646/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
4647/// activation streams (prefill amortization — the same idea as the
4648/// AVX2 twin; per stream the group order, fma order and tail match the
4649/// single-row kernel exactly, so batch == matvec bit-for-bit).
4650#[cfg(target_arch = "aarch64")]
4651#[target_feature(enable = "neon,dotprod")]
4652unsafe fn dot_q1_row_1x4_sdot(
4653    bytes: &[u8],
4654    r: usize,
4655    gpr: usize,
4656    xs: [&[i8]; 4],
4657    gs: [&[i32]; 4],
4658) -> [f32; 4] {
4659    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
4660    unsafe {
4661        use core::arch::aarch64::*;
4662        use core::arch::asm;
4663        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4664        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4665        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4666        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4667        const IW11: [u8; 16] = [
4668            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4669        ];
4670        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4671        let m = vld1q_u8(MASKS.as_ptr());
4672        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4673        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4674        let isc = vld1_u8(ISC.as_ptr());
4675        macro_rules! sdot2 {
4676            ($w0:expr, $w1:expr, $x:expr) => {{
4677                let x0 = vld1q_s8($x);
4678                let x1 = vld1q_s8($x.add(16));
4679                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4680                asm!(
4681                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4682                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4683                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4684                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
4685                    options(pure, nomem, nostack),
4686                );
4687                vaddq_s32(a0, a1)
4688            }};
4689        }
4690        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4691        let row_base = r * gpr * Q1_TILE;
4692        let abs_end = bytes.len();
4693        let mut accv = [vdupq_n_f32(0.0); 4];
4694        let mut gi = 0;
4695        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4696            let t0 = base.add(gi * Q1_TILE);
4697            let ld_a = vld1q_u8(t0);
4698            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4699            // Unpack ONCE — eight ±mask vectors serve all four streams.
4700            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
4701            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
4702            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
4703            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
4704            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
4705            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
4706            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
4707            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
4708            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4709            let scf: float32x4_t;
4710            asm!(
4711                "fcvtl {o:v}.4s, {i:v}.4h",
4712                o = out(vreg) scf, i = in(vreg) sc16,
4713                options(pure, nomem, nostack),
4714            );
4715            for k in 0..4 {
4716                let xp = xs[k].as_ptr();
4717                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
4718                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
4719                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
4720                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
4721                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4722                let g = vld1q_s32(gs[k].as_ptr().add(gi));
4723                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4724                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
4725            }
4726            gi += 4;
4727        }
4728        let mut acc = [
4729            vaddvq_f32(accv[0]),
4730            vaddvq_f32(accv[1]),
4731            vaddvq_f32(accv[2]),
4732            vaddvq_f32(accv[3]),
4733        ];
4734        while gi < gpr {
4735            let t = base.add(gi * Q1_TILE);
4736            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4737            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
4738            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
4739            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4740            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4741            for k in 0..4 {
4742                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
4743                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
4744            }
4745            gi += 1;
4746        }
4747        acc
4748    }
4749}
4750
4751/// (weight ±1, scale) of one q1 element — the exact outlier term.
4752#[inline]
4753fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4754    let gi = j / GROUP_SIZE;
4755    let k = j % GROUP_SIZE;
4756    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4757    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4758    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
4759    ((bit as i32 * 2 - 1) as f32, s)
4760}
4761
4762/// Exact scalar q1 row (CMF_SDOT=0 contract).
4763#[inline]
4764fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4765    let mut acc = 0f32;
4766    for gi in 0..gpr {
4767        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4768        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4769        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4770        let mut ga = 0f32;
4771        for (j, &b) in tile[2..].iter().enumerate() {
4772            for k in 0..8 {
4773                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
4774            }
4775        }
4776        acc += ga * s;
4777    }
4778    acc
4779}
4780
4781/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
4782/// extracted so multi-matrix jobs drive the same kernel).
4783#[allow(clippy::too_many_arguments)]
4784fn q1_range_a8w8(
4785    bytes: &[u8],
4786    gpr: usize,
4787    act: &SplitAct,
4788    gsum: &[i32],
4789    out: SendMut,
4790    start: usize,
4791    end: usize,
4792) {
4793    for r in start..end {
4794        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
4795        for &(j, xv) in &act.outliers {
4796            let (w, s) = q1_outlier(bytes, r, gpr, j);
4797            acc += w * s * xv;
4798        }
4799        // SAFETY: disjoint row ranges per worker.
4800        unsafe { *out.at(r) = acc };
4801    }
4802}
4803
4804/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
4805fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
4806    for r in start..end {
4807        // SAFETY: disjoint row ranges per worker.
4808        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
4809    }
4810}
4811
4812/// q1t per-row overlay locator. After the base (`base_len`) come
4813/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
4814/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
4815/// `(row_ptr offset, entries offset, present)`.
4816fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
4817    let entries = base_len + (rows + 1) * 4;
4818    (base_len, entries, entries <= bytes.len())
4819}
4820
4821/// Read `row_ptr[r]` from the overlay's prefix-sum table.
4822#[inline]
4823fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
4824    let o = rp_off + r * 4;
4825    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
4826}
4827
4828/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
4829/// decoding a q1t code is a table load, not the base-3 divide/modulo per
4830/// weight (division is ~20–40× the cost of a load). Built at compile time.
4831const SIGN5: [[f32; 5]; 256] = {
4832    let mut lut = [[0.0f32; 5]; 256];
4833    let pow3 = [1u16, 3, 9, 27, 81];
4834    let mut byte = 0usize;
4835    while byte < 256 {
4836        let mut i = 0usize;
4837        while i < 5 {
4838            let code = (byte as u16 / pow3[i]) % 3;
4839            lut[byte][i] = if code == 1 {
4840                1.0
4841            } else if code == 2 {
4842                -1.0
4843            } else {
4844                0.0
4845            };
4846            i += 1;
4847        }
4848        byte += 1;
4849    }
4850    lut
4851};
4852
4853/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
4854const SIGN5_I8: [[i8; 5]; 256] = {
4855    let mut lut = [[0i8; 5]; 256];
4856    let pow3 = [1u16, 3, 9, 27, 81];
4857    let mut byte = 0usize;
4858    while byte < 256 {
4859        let mut i = 0usize;
4860        while i < 5 {
4861            let code = (byte as u16 / pow3[i]) % 3;
4862            lut[byte][i] = if code == 1 {
4863                1
4864            } else if code == 2 {
4865                -1
4866            } else {
4867                0
4868            };
4869            i += 1;
4870        }
4871        byte += 1;
4872    }
4873    lut
4874};
4875
4876/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
4877/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
4878/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
4879/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
4880/// buffer is padded to 40). This is the decode/prefill hot inner op.
4881const SIGN5_U64: [u64; 256] = {
4882    let mut lut = [0u64; 256];
4883    let pow3 = [1u16, 3, 9, 27, 81];
4884    let mut byte = 0usize;
4885    while byte < 256 {
4886        let mut v = 0u64;
4887        let mut i = 0usize;
4888        while i < 5 {
4889            let code = (byte as u16 / pow3[i]) % 3;
4890            let s: u8 = if code == 1 {
4891                1
4892            } else if code == 2 {
4893                0xFF
4894            } else {
4895                0
4896            };
4897            v |= (s as u64) << (i * 8);
4898            i += 1;
4899        }
4900        lut[byte] = v;
4901        byte += 1;
4902    }
4903    lut
4904};
4905
4906/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
4907/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
4908/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
4909/// the overlay correction owns that column — no double counting.
4910#[inline]
4911fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
4912    const TILE: usize = cortiq_core::quant::Q1T_TILE;
4913    let off = (r * gpr + j / GROUP_SIZE) * TILE;
4914    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4915    let within = j % GROUP_SIZE;
4916    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
4917}
4918
4919/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
4920/// (integer accumulation is order-independent).
4921#[cfg(target_arch = "aarch64")]
4922#[target_feature(enable = "neon,dotprod")]
4923#[inline]
4924unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
4925    // SAFETY: caller guarantees 32 readable i8 at each pointer.
4926    unsafe {
4927        use core::arch::aarch64::*;
4928        use core::arch::asm;
4929        let w0 = vld1q_s8(w);
4930        let w1 = vld1q_s8(w.add(16));
4931        let x0 = vld1q_s8(x);
4932        let x1 = vld1q_s8(x.add(16));
4933        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4934        asm!(
4935            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4936            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4937            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4938            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4939            options(pure, nomem, nostack),
4940        );
4941        vaddvq_s32(vaddq_s32(a0, a1))
4942    }
4943}
4944
4945/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
4946/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
4947#[cfg(target_arch = "x86_64")]
4948#[target_feature(enable = "avx2")]
4949#[inline]
4950unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
4951    // SAFETY: caller guarantees 32 readable i8 at each pointer.
4952    unsafe {
4953        use core::arch::x86_64::*;
4954        let wv = _mm256_loadu_si256(w as *const __m256i);
4955        let xv = _mm256_loadu_si256(x as *const __m256i);
4956        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4957        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
4958        let hi128 = _mm256_extracti128_si256::<1>(d);
4959        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4960        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4961        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4962        _mm_cvtsi128_si32(s32)
4963    }
4964}
4965
4966/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
4967/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
4968/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
4969/// overwritten by the next; the final 6 padding bytes are unused by the dot.
4970#[inline]
4971fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
4972    debug_assert!(dst.len() >= 40);
4973    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
4974    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
4975    unsafe {
4976        let p = dst.as_mut_ptr();
4977        for bi in 0..7 {
4978            core::ptr::write_unaligned(
4979                p.add(bi * 5) as *mut u64,
4980                SIGN5_U64[*codes.add(bi) as usize],
4981            );
4982        }
4983    }
4984}
4985
4986/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
4987/// row's signs are unpacked once and dotted against every batch input).
4988/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
4989/// reachable; the scalar arm is a non-SIMD-arch fallback.
4990#[inline]
4991fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
4992    #[cfg(target_arch = "aarch64")]
4993    unsafe {
4994        return sdot32_i8(w, x);
4995    }
4996    #[cfg(target_arch = "x86_64")]
4997    unsafe {
4998        return i8dot32_avx2(w, x);
4999    }
5000    #[allow(unreachable_code)]
5001    unsafe {
5002        let mut s = 0i32;
5003        for k in 0..GROUP_SIZE {
5004            s += *w.add(k) as i32 * *x.add(k) as i32;
5005        }
5006        s
5007    }
5008}
5009
5010#[inline]
5011unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
5012    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
5013        (
5014            SIGN5_U64[*codes as usize],
5015            SIGN5_U64[*codes.add(1) as usize],
5016            SIGN5_U64[*codes.add(2) as usize],
5017            SIGN5_U64[*codes.add(3) as usize],
5018            SIGN5_U64[*codes.add(4) as usize],
5019            SIGN5_U64[*codes.add(5) as usize],
5020            SIGN5_U64[*codes.add(6) as usize],
5021        )
5022    };
5023
5024    let u0 = s0 | (s1 << 40);
5025    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
5026    let u2 = (s3 >> 8) | (s4 << 32);
5027    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
5028
5029    (u0, u1, u2, u3)
5030}
5031
5032/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
5033/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
5034/// ARM SDOT.
5035#[cfg(target_arch = "aarch64")]
5036#[target_feature(enable = "neon,dotprod")]
5037unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5038    use core::arch::aarch64::*;
5039    use core::arch::asm;
5040    unsafe {
5041        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5042        let mut acc = 0f32;
5043        let bytes_ptr = bytes.as_ptr();
5044        let xq_ptr = xq.as_ptr();
5045        let row_off = r * gpr * TILE;
5046
5047        let gpr2 = gpr & !1;
5048        let mut gi = 0;
5049        while gi < gpr2 {
5050            let off0 = row_off + gi * TILE;
5051            let off1 = off0 + TILE;
5052            let s0 = f16_to_f32(u16::from_le_bytes([
5053                *bytes_ptr.add(off0),
5054                *bytes_ptr.add(off0 + 1),
5055            ]));
5056            let s1 = f16_to_f32(u16::from_le_bytes([
5057                *bytes_ptr.add(off1),
5058                *bytes_ptr.add(off1 + 1),
5059            ]));
5060
5061            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5062            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5063
5064            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5065            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5066            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5067            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5068
5069            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5070            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5071            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
5072            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
5073
5074            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
5075            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5076            asm!(
5077                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
5078                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
5079                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
5080                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
5081                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
5082                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
5083                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
5084                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
5085                options(pure, nomem, nostack),
5086            );
5087            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
5088            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
5089            acc += d0 as f32 * s0 + d1 as f32 * s1;
5090            gi += 2;
5091        }
5092
5093        if gi < gpr {
5094            let off = row_off + gi * TILE;
5095            let s = f16_to_f32(u16::from_le_bytes([
5096                *bytes_ptr.add(off),
5097                *bytes_ptr.add(off + 1),
5098            ]));
5099            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5100            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5101            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5102            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5103            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5104            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5105            asm!(
5106                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5107                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5108                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5109                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5110                options(pure, nomem, nostack),
5111            );
5112            let d = vaddvq_s32(vaddq_s32(a0, a1));
5113            acc += d as f32 * s;
5114        }
5115        acc
5116    }
5117}
5118
5119/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
5120#[cfg(target_arch = "x86_64")]
5121#[target_feature(enable = "avx2")]
5122unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5123    use core::arch::x86_64::*;
5124    unsafe {
5125        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5126        let mut acc = 0f32;
5127        let bytes_ptr = bytes.as_ptr();
5128        let xq_ptr = xq.as_ptr();
5129        let row_off = r * gpr * TILE;
5130
5131        let ones = _mm256_set1_epi16(1);
5132        for gi in 0..gpr {
5133            let off = row_off + gi * TILE;
5134            let s = f16_to_f32(u16::from_le_bytes([
5135                *bytes_ptr.add(off),
5136                *bytes_ptr.add(off + 1),
5137            ]));
5138            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5139            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5140            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5141            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5142            let d256 = _mm256_madd_epi16(p16, ones);
5143            let d128 = _mm_add_epi32(
5144                _mm256_castsi256_si128(d256),
5145                _mm256_extracti128_si256(d256, 1),
5146            );
5147            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
5148            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
5149            acc += d32 as f32 * s;
5150        }
5151        acc
5152    }
5153}
5154
5155/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
5156#[cfg(target_arch = "x86_64")]
5157#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5158unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5159    use core::arch::x86_64::*;
5160    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
5161    unsafe {
5162        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5163        let mut acc = 0f32;
5164        let bytes_ptr = bytes.as_ptr();
5165        let xq_ptr = xq.as_ptr();
5166        let row_off = r * gpr * TILE;
5167        for gi in 0..gpr {
5168            let off = row_off + gi * TILE;
5169            let s = f16_to_f32(u16::from_le_bytes([
5170                *bytes_ptr.add(off),
5171                *bytes_ptr.add(off + 1),
5172            ]));
5173            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5174            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5175            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5176            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5177            acc += d as f32 * s;
5178        }
5179        acc
5180    }
5181}
5182
5183/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
5184/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
5185/// reachable.
5186#[inline]
5187fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5188    #[cfg(target_arch = "aarch64")]
5189    unsafe {
5190        return q1t_dot_row_sdot(bytes, r, gpr, xq);
5191    }
5192    #[cfg(target_arch = "x86_64")]
5193    unsafe {
5194        if vnni_tiles_enabled() {
5195            return q1t_dot_row_vnni(bytes, r, gpr, xq);
5196        }
5197        return q1t_dot_row_avx2(bytes, r, gpr, xq);
5198    }
5199    #[allow(unreachable_code)]
5200    {
5201        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5202        let mut acc = 0f32;
5203        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
5204        for gi in 0..gpr {
5205            let off = (r * gpr + gi) * TILE;
5206            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5207            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
5208            let mut d = 0i32;
5209            for k in 0..GROUP_SIZE {
5210                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
5211            }
5212            acc += d as f32 * s;
5213        }
5214        acc
5215    }
5216}
5217
5218/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
5219/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
5220/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
5221/// base contributes nothing there and this is a plain `value·x`, not
5222/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
5223/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
5224fn q1t_row_outlier_correction(
5225    bytes: &[u8],
5226    r: usize,
5227    rp_off: usize,
5228    entries_off: usize,
5229    has_ov: bool,
5230    x: &[f32],
5231) -> f32 {
5232    if !has_ov {
5233        return 0.0;
5234    }
5235    let (c0, c1) = (
5236        q1t_rowptr(bytes, rp_off, r),
5237        q1t_rowptr(bytes, rp_off, r + 1),
5238    );
5239    let mut corr = 0f32;
5240    for p in c0..c1 {
5241        let e = entries_off + p * 4;
5242        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5243        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5244        corr += val * x[col];
5245    }
5246    corr
5247}
5248
5249/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
5250/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
5251/// Used by the batched (prefill) path where the decode amortizes over the batch.
5252fn q1t_dequant_row(
5253    bytes: &[u8],
5254    r: usize,
5255    gpr: usize,
5256    rp_off: usize,
5257    entries_off: usize,
5258    has_ov: bool,
5259    buf: &mut [f32],
5260) {
5261    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5262    for g in 0..gpr {
5263        let off = (r * gpr + g) * TILE;
5264        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5265        let codes = &bytes[off + 2..off + TILE];
5266        let bc = g * GROUP_SIZE;
5267        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
5268        for bi in 0..6 {
5269            let lut = &SIGN5[codes[bi] as usize];
5270            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
5271            for i in 0..5 {
5272                d[i] = lut[i] * s;
5273            }
5274        }
5275        let lut = &SIGN5[codes[6] as usize];
5276        buf[bc + 30] = lut[0] * s;
5277        buf[bc + 31] = lut[1] * s;
5278    }
5279    if !has_ov {
5280        return;
5281    }
5282    let (c0, c1) = (
5283        q1t_rowptr(bytes, rp_off, r),
5284        q1t_rowptr(bytes, rp_off, r + 1),
5285    );
5286    for p in c0..c1 {
5287        let e = entries_off + p * 4;
5288        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5289        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5290    }
5291}
5292
5293/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
5294/// computes the ternary base; the overlay stays on the CPU — its entries are
5295/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
5296fn q1t_add_overlay(
5297    bytes: &[u8],
5298    x: &[f32],
5299    rows: usize,
5300    cols: usize,
5301    out: &mut [f32],
5302    pool: Option<&Pool>,
5303) {
5304    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5305    let gpr = cols / GROUP_SIZE;
5306    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5307    if !has_ov {
5308        return;
5309    }
5310    let out_addr = SendMut(out.as_mut_ptr());
5311    let run = move |start: usize, end: usize| {
5312        for r in start..end {
5313            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5314            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
5315            unsafe { *out_addr.at(r) += corr };
5316        }
5317    };
5318    dispatch_rows(pool, rows, &run);
5319}
5320
5321/// Q1T row range via the A8W8 int8 path — shared activation split,
5322/// per-row: base SDOT dot + outlier correction + overlay.
5323#[allow(clippy::too_many_arguments)]
5324fn q1t_range_a8w8(
5325    bytes: &[u8],
5326    gpr: usize,
5327    rp_off: usize,
5328    ent_off: usize,
5329    has_ov: bool,
5330    act: &SplitAct,
5331    x: &[f32],
5332    out: SendMut,
5333    start: usize,
5334    end: usize,
5335) {
5336    for r in start..end {
5337        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5338        for &(j, xv) in &act.outliers {
5339            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5340        }
5341        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5342        // SAFETY: disjoint row ranges per worker.
5343        unsafe { *out.at(r) = acc };
5344    }
5345}
5346
5347/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
5348/// dispatch when a8w8 is unavailable.
5349#[allow(clippy::too_many_arguments)]
5350fn q1t_range_f32_batch(
5351    bytes: &[u8],
5352    gpr: usize,
5353    rp_off: usize,
5354    ent_off: usize,
5355    has_ov: bool,
5356    x: &[f32],
5357    out: SendMut,
5358    start: usize,
5359    end: usize,
5360) {
5361    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5362    let mut sg = [0f32; GROUP_SIZE];
5363    for r in start..end {
5364        let mut acc = 0f32;
5365        for g in 0..gpr {
5366            let off = (r * gpr + g) * TILE;
5367            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5368            let codes = &bytes[off + 2..off + TILE];
5369            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5370            for bi in 0..6 {
5371                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5372            }
5373            let lut = &SIGN5[codes[6] as usize];
5374            sg[30] = lut[0];
5375            sg[31] = lut[1];
5376            let mut gsum = 0f32;
5377            for k in 0..GROUP_SIZE {
5378                gsum += sg[k] * xg[k];
5379            }
5380            acc += s * gsum;
5381        }
5382        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5383        // SAFETY: disjoint row ranges per worker.
5384        unsafe { *out.at(r) = acc };
5385    }
5386}
5387
5388/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
5389/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
5390/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
5391fn q1t_matvec(
5392    bytes: &[u8],
5393    x: &[f32],
5394    rows: usize,
5395    cols: usize,
5396    out: &mut [f32],
5397    pool: Option<&Pool>,
5398) {
5399    debug_assert_eq!(out.len(), rows);
5400    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5401    let gpr = cols / GROUP_SIZE;
5402    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5403    let out_addr = SendMut(out.as_mut_ptr());
5404    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
5405    // (`split_act`), activation outliers added back exactly in f32, weight
5406    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
5407    if a8w8_enabled() {
5408        let act = split_act(x);
5409        let act = &act;
5410        let run = move |start: usize, end: usize| {
5411            for r in start..end {
5412                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5413                for &(j, xv) in &act.outliers {
5414                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5415                }
5416                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5417                // SAFETY: disjoint row ranges per worker.
5418                unsafe { *out_addr.at(r) = acc };
5419            }
5420        };
5421        dispatch_rows(pool, rows, &run);
5422        return;
5423    }
5424    let run = move |start: usize, end: usize| {
5425        // Per-group signs, unpacked contiguously so the dot below is a clean
5426        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
5427        // 5-values-per-byte base-3 layout won't SIMD in place.
5428        let mut sg = [0f32; GROUP_SIZE];
5429        for r in start..end {
5430            let mut acc = 0f32;
5431            for g in 0..gpr {
5432                let off = (r * gpr + g) * TILE;
5433                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5434                let codes = &bytes[off + 2..off + TILE];
5435                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5436                for bi in 0..6 {
5437                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5438                }
5439                let lut = &SIGN5[codes[6] as usize];
5440                sg[30] = lut[0];
5441                sg[31] = lut[1];
5442                let mut gsum = 0f32;
5443                for k in 0..GROUP_SIZE {
5444                    gsum += sg[k] * xg[k];
5445                }
5446                acc += s * gsum;
5447            }
5448            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5449            unsafe { *out_addr.at(r) = acc };
5450        }
5451    };
5452    dispatch_rows(pool, rows, &run);
5453}
5454
5455/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
5456/// ternary codes serves BOTH activation streams (the unpack chain is
5457/// the dominant per-row cost — MTP verify pairs paid it twice). Per
5458/// stream the group order and f32 accumulation match the single-row
5459/// kernel exactly, so pair == 2×matvec bit-for-bit.
5460#[cfg(target_arch = "aarch64")]
5461#[target_feature(enable = "neon,dotprod")]
5462unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
5463    use core::arch::aarch64::*;
5464    use core::arch::asm;
5465    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
5466    unsafe {
5467        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5468        let bytes_ptr = bytes.as_ptr();
5469        let row_off = r * gpr * TILE;
5470        let xp = [xa.as_ptr(), xb.as_ptr()];
5471        let mut acc = [0f32; 2];
5472        macro_rules! sdot2 {
5473            ($w0:expr, $w1:expr, $x:expr) => {{
5474                let x0 = vld1q_s8($x);
5475                let x1 = vld1q_s8($x.add(16));
5476                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5477                asm!(
5478                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5479                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5480                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5481                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5482                    options(pure, nomem, nostack),
5483                );
5484                vaddvq_s32(vaddq_s32(a0, a1))
5485            }};
5486        }
5487        let gpr2 = gpr & !1;
5488        let mut gi = 0;
5489        while gi < gpr2 {
5490            let off0 = row_off + gi * TILE;
5491            let off1 = off0 + TILE;
5492            let s0 = f16_to_f32(u16::from_le_bytes([
5493                *bytes_ptr.add(off0),
5494                *bytes_ptr.add(off0 + 1),
5495            ]));
5496            let s1 = f16_to_f32(u16::from_le_bytes([
5497                *bytes_ptr.add(off1),
5498                *bytes_ptr.add(off1 + 1),
5499            ]));
5500            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5501            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5502            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5503            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5504            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5505            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5506            for k in 0..2 {
5507                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
5508                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
5509                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
5510            }
5511            gi += 2;
5512        }
5513        if gi < gpr {
5514            let off = row_off + gi * TILE;
5515            let s = f16_to_f32(u16::from_le_bytes([
5516                *bytes_ptr.add(off),
5517                *bytes_ptr.add(off + 1),
5518            ]));
5519            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5520            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5521            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5522            for k in 0..2 {
5523                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
5524                acc[k] += d as f32 * s;
5525            }
5526        }
5527        acc
5528    }
5529}
5530
5531/// Fused Q1T pair matvec: ONE pass over the rows serves both
5532/// activation streams — on ARM the ternary register unpack happens
5533/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
5534/// rides the row's L1-warm tile bytes. Per stream the math matches
5535/// `q1t_matvec` exactly.
5536fn q1t_matvec2(
5537    bytes: &[u8],
5538    x1: &[f32],
5539    x2: &[f32],
5540    rows: usize,
5541    cols: usize,
5542    o1: &mut [f32],
5543    o2: &mut [f32],
5544    pool: Option<&Pool>,
5545) {
5546    debug_assert_eq!(o1.len(), rows);
5547    debug_assert_eq!(o2.len(), rows);
5548    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5549    let gpr = cols / GROUP_SIZE;
5550    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5551    let out1 = SendMut(o1.as_mut_ptr());
5552    let out2 = SendMut(o2.as_mut_ptr());
5553    if a8w8_enabled() {
5554        let a1 = split_act(x1);
5555        let a2 = split_act(x2);
5556        let (a1, a2) = (&a1, &a2);
5557        let run = move |start: usize, end: usize| {
5558            for r in start..end {
5559                #[cfg(target_arch = "aarch64")]
5560                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
5561                // target features are present.
5562                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
5563                #[cfg(not(target_arch = "aarch64"))]
5564                let ds = [
5565                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
5566                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
5567                ];
5568                let mut acc1 = ds[0] * a1.sx;
5569                for &(j, xv) in &a1.outliers {
5570                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
5571                }
5572                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
5573                let mut acc2 = ds[1] * a2.sx;
5574                for &(j, xv) in &a2.outliers {
5575                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
5576                }
5577                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
5578                // SAFETY: disjoint row ranges per worker.
5579                unsafe {
5580                    *out1.at(r) = acc1;
5581                    *out2.at(r) = acc2;
5582                }
5583            }
5584        };
5585        dispatch_rows(pool, rows, &run);
5586        return;
5587    }
5588    let run = move |start: usize, end: usize| {
5589        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
5590        // dot both streams — same op order per stream as `q1t_matvec`.
5591        let mut sg = [0f32; GROUP_SIZE];
5592        for r in start..end {
5593            let mut acc1 = 0f32;
5594            let mut acc2 = 0f32;
5595            for g in 0..gpr {
5596                let off = (r * gpr + g) * TILE;
5597                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5598                let codes = &bytes[off + 2..off + TILE];
5599                for bi in 0..6 {
5600                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5601                }
5602                let lut = &SIGN5[codes[6] as usize];
5603                sg[30] = lut[0];
5604                sg[31] = lut[1];
5605                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5606                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5607                let mut gsum1 = 0f32;
5608                for k in 0..GROUP_SIZE {
5609                    gsum1 += sg[k] * xg1[k];
5610                }
5611                acc1 += s * gsum1;
5612                let mut gsum2 = 0f32;
5613                for k in 0..GROUP_SIZE {
5614                    gsum2 += sg[k] * xg2[k];
5615                }
5616                acc2 += s * gsum2;
5617            }
5618            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
5619            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
5620            // SAFETY: disjoint row ranges per worker.
5621            unsafe {
5622                *out1.at(r) = acc1;
5623                *out2.at(r) = acc2;
5624            }
5625        }
5626    };
5627    dispatch_rows(pool, rows, &run);
5628}
5629
5630/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
5631/// batch against it (amortizes the per-row decode).
5632fn q1t_matmat(
5633    bytes: &[u8],
5634    xs: &[f32],
5635    b: usize,
5636    rows: usize,
5637    cols: usize,
5638    out: &mut [f32],
5639    pool: Option<&Pool>,
5640) {
5641    debug_assert_eq!(out.len(), b * rows);
5642    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5643    let gpr = cols / GROUP_SIZE;
5644    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5645    let out_addr = SendMut(out.as_mut_ptr());
5646    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
5647    // each weight row's signs to i8 ONCE, then int8-dot against every input —
5648    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
5649    if a8w8_enabled() {
5650        let acts: Vec<SplitAct> = (0..b)
5651            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
5652            .collect();
5653        let acts = &acts;
5654        let run = move |start: usize, end: usize| {
5655            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
5656            let mut sc = vec![0f32; gpr]; // per-group scales
5657            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
5658            for r in start..end {
5659                for g in 0..gpr {
5660                    let off = (r * gpr + g) * TILE;
5661                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5662                    q1t_unpack_group_i8(
5663                        bytes.as_ptr().wrapping_add(off + 2),
5664                        &mut sg[g * GROUP_SIZE..],
5665                    );
5666                }
5667                for bi in 0..b {
5668                    let act = &acts[bi];
5669                    let mut isum = 0f32;
5670                    for g in 0..gpr {
5671                        let d = q1t_i8dot32(
5672                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
5673                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
5674                        );
5675                        isum += d as f32 * sc[g];
5676                    }
5677                    let mut acc = isum * act.sx;
5678                    for &(j, xv) in &act.outliers {
5679                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5680                    }
5681                    accs[bi] = acc;
5682                }
5683                // Overlay ONCE per row for the whole batch: read each (col, val)
5684                // from mmap a single time (was b× — the re-read dominated prefill)
5685                // and fan it out over the batch via the cached inputs.
5686                if has_ov {
5687                    let (c0, c1) = (
5688                        q1t_rowptr(bytes, rp_off, r),
5689                        q1t_rowptr(bytes, rp_off, r + 1),
5690                    );
5691                    for p in c0..c1 {
5692                        let e = ent_off + p * 4;
5693                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5694                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5695                        for bi in 0..b {
5696                            accs[bi] += val * xs[bi * cols + col];
5697                        }
5698                    }
5699                }
5700                for bi in 0..b {
5701                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
5702                }
5703            }
5704        };
5705        dispatch_rows(pool, rows, &run);
5706        return;
5707    }
5708    let run = move |start: usize, end: usize| {
5709        let mut buf = vec![0f32; cols];
5710        for r in start..end {
5711            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
5712            for bi in 0..b {
5713                let xr = &xs[bi * cols..(bi + 1) * cols];
5714                let mut acc = 0f32;
5715                for j in 0..cols {
5716                    acc += buf[j] * xr[j];
5717                }
5718                unsafe { *out_addr.at(bi * rows + r) = acc };
5719            }
5720        }
5721    };
5722    dispatch_rows(pool, rows, &run);
5723}
5724
5725fn q1_matvec(
5726    bytes: &[u8],
5727    x: &[f32],
5728    rows: usize,
5729    cols: usize,
5730    out: &mut [f32],
5731    pool: Option<&Pool>,
5732) {
5733    debug_assert_eq!(out.len(), rows);
5734    let gpr = cols / GROUP_SIZE;
5735    let out_addr = SendMut(out.as_mut_ptr());
5736    if a8w8_enabled() {
5737        let act = split_act(x);
5738        let gsum = q1_group_sums(&act.xq, gpr);
5739        let (act, gsum) = (&act, &gsum);
5740        let run = move |start: usize, end: usize| {
5741            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
5742        };
5743        dispatch_rows(pool, rows, &run);
5744        return;
5745    }
5746    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
5747    dispatch_rows(pool, rows, &run);
5748}
5749
5750/// Fused two-input q1 matvec (weights read once per pair).
5751#[allow(clippy::too_many_arguments)]
5752fn q1_matvec2(
5753    bytes: &[u8],
5754    x1: &[f32],
5755    x2: &[f32],
5756    rows: usize,
5757    cols: usize,
5758    o1: &mut [f32],
5759    o2: &mut [f32],
5760    pool: Option<&Pool>,
5761) {
5762    let gpr = cols / GROUP_SIZE;
5763    let p1 = SendMut(o1.as_mut_ptr());
5764    let p2 = SendMut(o2.as_mut_ptr());
5765    if a8w8_enabled() {
5766        let a1 = split_act(x1);
5767        let a2 = split_act(x2);
5768        let g1 = q1_group_sums(&a1.xq, gpr);
5769        let g2 = q1_group_sums(&a2.xq, gpr);
5770        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
5771        let run = move |start: usize, end: usize| {
5772            for r in start..end {
5773                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
5774                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
5775                for &(j, xv) in &a1.outliers {
5776                    let (w, s) = q1_outlier(bytes, r, gpr, j);
5777                    v1 += w * s * xv;
5778                }
5779                for &(j, xv) in &a2.outliers {
5780                    let (w, s) = q1_outlier(bytes, r, gpr, j);
5781                    v2 += w * s * xv;
5782                }
5783                // SAFETY: disjoint row ranges per worker.
5784                unsafe {
5785                    *p1.at(r) = v1;
5786                    *p2.at(r) = v2;
5787                }
5788            }
5789        };
5790        dispatch_rows(pool, rows, &run);
5791        return;
5792    }
5793    let run = move |start: usize, end: usize| {
5794        for r in start..end {
5795            // SAFETY: disjoint row ranges per worker.
5796            unsafe {
5797                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
5798                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
5799            }
5800        }
5801    };
5802    dispatch_rows(pool, rows, &run);
5803}
5804
5805/// Batched q1 matmat: each row's tiles stream once per microbatch.
5806#[allow(clippy::too_many_arguments)]
5807fn q1_matmat(
5808    bytes: &[u8],
5809    xs_all: &[f32],
5810    b: usize,
5811    rows: usize,
5812    cols: usize,
5813    out: &mut [f32],
5814    pool: Option<&Pool>,
5815) {
5816    debug_assert_eq!(out.len(), b * rows);
5817    let gpr = cols / GROUP_SIZE;
5818    let out_addr = SendMut(out.as_mut_ptr());
5819    if a8w8_enabled() {
5820        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
5821            .map(|bi| {
5822                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
5823                let gsum = q1_group_sums(&act.xq, gpr);
5824                (act, gsum)
5825            })
5826            .collect();
5827        let acts = &acts;
5828        #[cfg(target_arch = "x86_64")]
5829        let blocked_ok = avx2_enabled()
5830            && std::env::var("CMF_X86_BLOCKED")
5831                .map(|v| v != "0")
5832                .unwrap_or(true);
5833        #[cfg(target_arch = "aarch64")]
5834        let blocked_ok = sdot_enabled()
5835            && std::env::var("CMF_X86_BLOCKED")
5836                .map(|v| v != "0")
5837                .unwrap_or(true);
5838        let run = move |start: usize, end: usize| {
5839            for r in start..end {
5840                let mut bi = 0usize;
5841                // Blocked 1×4: the unpacked bit mask serves four
5842                // activation streams per group.
5843                #[cfg(target_arch = "aarch64")]
5844                if blocked_ok {
5845                    while bi + 4 <= acts.len() {
5846                        let xs = [
5847                            acts[bi].0.xq.as_slice(),
5848                            acts[bi + 1].0.xq.as_slice(),
5849                            acts[bi + 2].0.xq.as_slice(),
5850                            acts[bi + 3].0.xq.as_slice(),
5851                        ];
5852                        let gs = [
5853                            acts[bi].1.as_slice(),
5854                            acts[bi + 1].1.as_slice(),
5855                            acts[bi + 2].1.as_slice(),
5856                            acts[bi + 3].1.as_slice(),
5857                        ];
5858                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
5859                        for k in 0..4 {
5860                            let (act, _) = &acts[bi + k];
5861                            let mut acc = d[k] * act.sx;
5862                            for &(j, xv) in &act.outliers {
5863                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
5864                                acc += w * sc * xv;
5865                            }
5866                            // SAFETY: disjoint (bi, r) cells per worker.
5867                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5868                        }
5869                        bi += 4;
5870                    }
5871                }
5872                #[cfg(target_arch = "x86_64")]
5873                if blocked_ok {
5874                    while bi + 4 <= acts.len() {
5875                        let xs = [
5876                            acts[bi].0.xq.as_slice(),
5877                            acts[bi + 1].0.xq.as_slice(),
5878                            acts[bi + 2].0.xq.as_slice(),
5879                            acts[bi + 3].0.xq.as_slice(),
5880                        ];
5881                        let gs = [
5882                            acts[bi].1.as_slice(),
5883                            acts[bi + 1].1.as_slice(),
5884                            acts[bi + 2].1.as_slice(),
5885                            acts[bi + 3].1.as_slice(),
5886                        ];
5887                        let d = unsafe {
5888                            if vnni_tiles_enabled() {
5889                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
5890                            } else {
5891                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
5892                            }
5893                        };
5894                        for k in 0..4 {
5895                            let (act, _) = &acts[bi + k];
5896                            let mut acc = d[k] * act.sx;
5897                            for &(j, xv) in &act.outliers {
5898                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
5899                                acc += w * sc * xv;
5900                            }
5901                            // SAFETY: disjoint (bi, r) cells per worker.
5902                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5903                        }
5904                        bi += 4;
5905                    }
5906                }
5907                while bi < acts.len() {
5908                    let (act, gsum) = &acts[bi];
5909                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5910                    for &(j, xv) in &act.outliers {
5911                        let (w, s) = q1_outlier(bytes, r, gpr, j);
5912                        acc += w * s * xv;
5913                    }
5914                    // SAFETY: disjoint (bi, r) cells per worker range.
5915                    unsafe { *out_addr.at(bi * rows + r) = acc };
5916                    bi += 1;
5917                }
5918            }
5919        };
5920        dispatch_rows(pool, rows, &run);
5921        return;
5922    }
5923    let run = move |start: usize, end: usize| {
5924        for r in start..end {
5925            for bi in 0..b {
5926                let x = &xs_all[bi * cols..(bi + 1) * cols];
5927                // SAFETY: disjoint (bi, r) cells per worker range.
5928                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
5929            }
5930        }
5931    };
5932    dispatch_rows(pool, rows, &run);
5933}
5934
5935/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
5936/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
5937/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
5938/// 32-group, exact outlier correction — the same A8W8 contract as q8.
5939/// `CMF_SDOT=0` keeps the exact scalar path.
5940fn q4matvec(
5941    bytes: &[u8],
5942    x: &[f32],
5943    rows: usize,
5944    cols: usize,
5945    out: &mut [f32],
5946    pool: Option<&Pool>,
5947) {
5948    debug_assert_eq!(out.len(), rows);
5949    let (packed, scales) = q4_split(bytes, rows, cols);
5950    let gpr = cols / GROUP_SIZE;
5951    let out_addr = SendMut(out.as_mut_ptr());
5952
5953    if a8w8_enabled() {
5954        let act = split_act(x);
5955        let run = move |start: usize, end: usize| {
5956            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
5957        };
5958        dispatch_rows(pool, rows, &run);
5959        return;
5960    }
5961
5962    let run =
5963        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
5964    dispatch_rows(pool, rows, &run);
5965}
5966
5967/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
5968/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
5969#[inline]
5970#[allow(unreachable_code)]
5971/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
5972/// streams: the 32-byte weight chunk and its abs() load once per group,
5973/// the per-group f16 scale decodes once — four maddubs+reduce chains
5974/// instead of four full (load, abs, dot) rounds.
5975#[cfg(target_arch = "x86_64")]
5976#[target_feature(enable = "avx2")]
5977unsafe fn dot_q4b_row_1x4_avx2(
5978    buf: &[u8],
5979    scales: &[u8],
5980    g0: usize,
5981    gpr: usize,
5982    xs: [&[i8]; 4],
5983) -> [f32; 4] {
5984    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
5985    unsafe {
5986        use core::arch::x86_64::*;
5987        let ones = _mm256_set1_epi16(1);
5988        let mut acc = [0f32; 4];
5989        for gi in 0..gpr {
5990            let s = f16_to_f32(u16::from_le_bytes([
5991                scales[(g0 + gi) * 2],
5992                scales[(g0 + gi) * 2 + 1],
5993            ]));
5994            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5995            let aw = _mm256_abs_epi8(w);
5996            for (k, xq) in xs.iter().enumerate() {
5997                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5998                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5999                let d = _mm256_madd_epi16(p16, ones);
6000                let hi128 = _mm256_extracti128_si256::<1>(d);
6001                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6002                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6003                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6004                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
6005            }
6006        }
6007        acc
6008    }
6009}
6010
6011/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
6012#[cfg(target_arch = "x86_64")]
6013#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6014unsafe fn dot_q4b_row_1x4_vnni(
6015    buf: &[u8],
6016    scales: &[u8],
6017    g0: usize,
6018    gpr: usize,
6019    xs: [&[i8]; 4],
6020) -> [f32; 4] {
6021    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6022    unsafe {
6023        use core::arch::x86_64::*;
6024        let mut acc = [0f32; 4];
6025        for gi in 0..gpr {
6026            let s = f16_to_f32(u16::from_le_bytes([
6027                scales[(g0 + gi) * 2],
6028                scales[(g0 + gi) * 2 + 1],
6029            ]));
6030            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6031            let aw = _mm256_abs_epi8(w);
6032            for (k, xq) in xs.iter().enumerate() {
6033                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6034                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6035                acc[k] += d as f32 * s;
6036            }
6037        }
6038        acc
6039    }
6040}
6041
6042/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
6043/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
6044/// accumulation order (the q4_block flavor applies sx once at the end,
6045/// matching ITS single path; the two conventions are historical and
6046/// each blocked leg must mirror its own).
6047#[cfg(target_arch = "x86_64")]
6048#[target_feature(enable = "avx2")]
6049unsafe fn dot_q4b_row_1x4_sx_avx2(
6050    buf: &[u8],
6051    scales: &[u8],
6052    g0: usize,
6053    gpr: usize,
6054    xs: [&[i8]; 4],
6055    sxs: [f32; 4],
6056) -> [f32; 4] {
6057    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6058    unsafe {
6059        use core::arch::x86_64::*;
6060        let ones = _mm256_set1_epi16(1);
6061        let mut acc = [0f32; 4];
6062        for gi in 0..gpr {
6063            let s = f16_to_f32(u16::from_le_bytes([
6064                scales[(g0 + gi) * 2],
6065                scales[(g0 + gi) * 2 + 1],
6066            ]));
6067            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6068            let aw = _mm256_abs_epi8(w);
6069            for (k, xq) in xs.iter().enumerate() {
6070                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6071                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6072                let d = _mm256_madd_epi16(p16, ones);
6073                let hi128 = _mm256_extracti128_si256::<1>(d);
6074                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6075                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6076                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6077                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
6078            }
6079        }
6080        acc
6081    }
6082}
6083
6084/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
6085/// per-group `(d·sx)·s` fold mirrors the vbit single path).
6086#[cfg(target_arch = "x86_64")]
6087#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6088unsafe fn dot_q4b_row_1x4_sx_vnni(
6089    buf: &[u8],
6090    scales: &[u8],
6091    g0: usize,
6092    gpr: usize,
6093    xs: [&[i8]; 4],
6094    sxs: [f32; 4],
6095) -> [f32; 4] {
6096    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6097    unsafe {
6098        use core::arch::x86_64::*;
6099        let mut acc = [0f32; 4];
6100        for gi in 0..gpr {
6101            let s = f16_to_f32(u16::from_le_bytes([
6102                scales[(g0 + gi) * 2],
6103                scales[(g0 + gi) * 2 + 1],
6104            ]));
6105            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6106            let aw = _mm256_abs_epi8(w);
6107            for (k, xq) in xs.iter().enumerate() {
6108                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6109                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6110                acc[k] += (d as f32 * sxs[k]) * s;
6111            }
6112        }
6113        acc
6114    }
6115}
6116
6117#[allow(unreachable_code)]
6118fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6119    #[cfg(target_arch = "aarch64")]
6120    unsafe {
6121        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
6122    }
6123    #[cfg(target_arch = "x86_64")]
6124    unsafe {
6125        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
6126    }
6127    let mut acc = 0f32;
6128    for gi in 0..gpr {
6129        let g = g0 + gi;
6130        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6131        let mut d = 0i32;
6132        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6133            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6134                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6135        }
6136        acc += d as f32 * s;
6137    }
6138    acc
6139}
6140
6141/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
6142#[inline]
6143#[allow(unreachable_code)]
6144fn dot_q4_row_i8_2(
6145    packed: &[u8],
6146    scales: &[u8],
6147    g0: usize,
6148    gpr: usize,
6149    xq1: &[i8],
6150    xq2: &[i8],
6151) -> (f32, f32) {
6152    #[cfg(target_arch = "aarch64")]
6153    unsafe {
6154        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
6155    }
6156    #[cfg(target_arch = "x86_64")]
6157    unsafe {
6158        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
6159    }
6160    (
6161        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
6162        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
6163    )
6164}
6165
6166/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
6167/// multi-matrix jobs can drive it for several tensors in one dispatch).
6168#[allow(clippy::too_many_arguments)]
6169fn q4_range_a8w8(
6170    packed: &[u8],
6171    scales: &[u8],
6172    gpr: usize,
6173    cols: usize,
6174    act: &SplitAct,
6175    out: SendMut,
6176    start: usize,
6177    end: usize,
6178) {
6179    for r in start..end {
6180        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
6181        // xq is zeroed at outlier slots — add the exact terms.
6182        for &(j, xv) in &act.outliers {
6183            let flat = r * cols + j;
6184            let byte = packed[flat / 2];
6185            let nib = if flat & 1 == 0 {
6186                byte & 0x0F
6187            } else {
6188                byte >> 4
6189            };
6190            let s = f16_to_f32(u16::from_le_bytes([
6191                scales[(flat / GROUP_SIZE) * 2],
6192                scales[(flat / GROUP_SIZE) * 2 + 1],
6193            ]));
6194            acc += ((nib as i32 - 8) as f32) * s * xv;
6195        }
6196        // SAFETY: disjoint row ranges per worker.
6197        unsafe { *out.at(r) = acc };
6198    }
6199}
6200
6201/// Two-input q4 row range via the A8W8 int8 path — kernel body of
6202/// `q4matvec2`, extracted for pair multi-matrix jobs.
6203#[allow(clippy::too_many_arguments)]
6204fn q4_range2_a8w8(
6205    packed: &[u8],
6206    scales: &[u8],
6207    gpr: usize,
6208    cols: usize,
6209    a1: &SplitAct,
6210    a2: &SplitAct,
6211    p1: SendMut,
6212    p2: SendMut,
6213    start: usize,
6214    end: usize,
6215) {
6216    for r in start..end {
6217        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
6218        let mut acc1 = s1 * a1.sx;
6219        let mut acc2 = s2 * a2.sx;
6220        // xq is zeroed at outlier slots — add the exact terms.
6221        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
6222            for &(j, xv) in outliers {
6223                let flat = r * cols + j;
6224                let byte = packed[flat / 2];
6225                let nib = if flat & 1 == 0 {
6226                    byte & 0x0F
6227                } else {
6228                    byte >> 4
6229                };
6230                let s = f16_to_f32(u16::from_le_bytes([
6231                    scales[(flat / GROUP_SIZE) * 2],
6232                    scales[(flat / GROUP_SIZE) * 2 + 1],
6233                ]));
6234                *acc += ((nib as i32 - 8) as f32) * s * xv;
6235            }
6236        };
6237        fix(&a1.outliers, &mut acc1);
6238        fix(&a2.outliers, &mut acc2);
6239        // SAFETY: disjoint row ranges per worker.
6240        unsafe {
6241            *p1.at(r) = acc1;
6242            *p2.at(r) = acc2;
6243        }
6244    }
6245}
6246
6247/// Exact scalar q4 row range (same extraction, non-SDOT path).
6248fn q4_range_f32(
6249    packed: &[u8],
6250    scales: &[u8],
6251    gpr: usize,
6252    x: &[f32],
6253    out: SendMut,
6254    start: usize,
6255    end: usize,
6256) {
6257    for r in start..end {
6258        let mut acc = 0f32;
6259        for gi in 0..gpr {
6260            let g = r * gpr + gi;
6261            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6262            let pk = &packed[g * 16..(g + 1) * 16];
6263            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6264            let mut ga = 0f32;
6265            for (k, &b) in pk.iter().enumerate() {
6266                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6267                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6268            }
6269            acc += ga * s;
6270        }
6271        // SAFETY: disjoint row ranges per worker.
6272        unsafe { *out.at(r) = acc };
6273    }
6274}
6275
6276/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
6277/// dotted against both activations (was: two full matvecs — double
6278/// weight traffic). Per-lane math matches `q4matvec` exactly.
6279#[allow(clippy::too_many_arguments)]
6280fn q4matvec2(
6281    bytes: &[u8],
6282    x1: &[f32],
6283    x2: &[f32],
6284    rows: usize,
6285    cols: usize,
6286    o1: &mut [f32],
6287    o2: &mut [f32],
6288    pool: Option<&Pool>,
6289) {
6290    debug_assert_eq!(o1.len(), rows);
6291    debug_assert_eq!(o2.len(), rows);
6292    let (packed, scales) = q4_split(bytes, rows, cols);
6293    let gpr = cols / GROUP_SIZE;
6294
6295    if a8w8_enabled() {
6296        let a1 = split_act(x1);
6297        let a2 = split_act(x2);
6298        let p1 = SendMut(o1.as_mut_ptr());
6299        let p2 = SendMut(o2.as_mut_ptr());
6300        let run = move |start: usize, end: usize| {
6301            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
6302        };
6303        dispatch_rows(pool, rows, &run);
6304        return;
6305    }
6306
6307    let p1 = SendMut(o1.as_mut_ptr());
6308    let p2 = SendMut(o2.as_mut_ptr());
6309    let run = move |start: usize, end: usize| {
6310        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
6311    };
6312    dispatch_rows(pool, rows, &run);
6313}
6314
6315/// Two-input exact scalar q4 row range (same extraction).
6316#[allow(clippy::too_many_arguments)]
6317fn q4_range2_f32(
6318    packed: &[u8],
6319    scales: &[u8],
6320    gpr: usize,
6321    x1: &[f32],
6322    x2: &[f32],
6323    p1: SendMut,
6324    p2: SendMut,
6325    start: usize,
6326    end: usize,
6327) {
6328    for r in start..end {
6329        let (mut acc1, mut acc2) = (0f32, 0f32);
6330        for gi in 0..gpr {
6331            let g = r * gpr + gi;
6332            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6333            let pk = &packed[g * 16..(g + 1) * 16];
6334            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6335            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6336            let (mut g1, mut g2) = (0f32, 0f32);
6337            for (k, &b) in pk.iter().enumerate() {
6338                let wl = (b & 0x0F) as f32 - 8.0;
6339                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
6340                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
6341                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
6342            }
6343            acc1 += g1 * s;
6344            acc2 += g2 * s;
6345        }
6346        // SAFETY: disjoint row ranges per worker.
6347        unsafe {
6348            *p1.at(r) = acc1;
6349            *p2.at(r) = acc2;
6350        }
6351    }
6352}
6353
6354thread_local! {
6355    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
6356    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
6357    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
6358    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6359}
6360
6361/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
6362/// and dotted against ALL b activations (prefill used to fall back to b
6363/// full matvecs — b× weight traffic and b× nibble decode). Per-position
6364/// math matches `q4matvec` exactly: same group order, same accumulation.
6365/// `out` is row-major [b, rows] like `qmatmat`.
6366#[allow(clippy::too_many_arguments)]
6367fn q4matmat(
6368    bytes: &[u8],
6369    xs_all: &[f32],
6370    b: usize,
6371    rows: usize,
6372    cols: usize,
6373    out: &mut [f32],
6374    pool: Option<&Pool>,
6375) {
6376    debug_assert_eq!(xs_all.len(), b * cols);
6377    debug_assert_eq!(out.len(), b * rows);
6378    let (packed, scales) = q4_split(bytes, rows, cols);
6379    let gpr = cols / GROUP_SIZE;
6380    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6381
6382    if a8w8_enabled() {
6383        let acts: Vec<SplitAct> = (0..b)
6384            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6385            .collect();
6386        let acts = &acts;
6387        let out_addr = SendMut(out.as_mut_ptr());
6388        let run = move |start: usize, end: usize| {
6389            ROW_I8.with(|rb| {
6390                let mut buf = rb.borrow_mut();
6391                buf.resize(cols, 0);
6392                for r in start..end {
6393                    // Unpack the row's nibbles to centered i8 once
6394                    // (element 2k = low nibble, 2k+1 = high — flat order,
6395                    // same as dot_q4_row_sdot's zip).
6396                    for gi in 0..gpr {
6397                        let g = r * gpr + gi;
6398                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6399                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
6400                            buf[gi * GROUP_SIZE + k * 2 + 1] =
6401                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
6402                        }
6403                    }
6404                    let mut bi = 0usize;
6405                    #[cfg(target_arch = "x86_64")]
6406                    if avx2_enabled()
6407                        && std::env::var("CMF_X86_BLOCKED")
6408                            .map(|v| v != "0")
6409                            .unwrap_or(true)
6410                    {
6411                        while bi + 4 <= acts.len() {
6412                            let xs = [
6413                                acts[bi].xq.as_slice(),
6414                                acts[bi + 1].xq.as_slice(),
6415                                acts[bi + 2].xq.as_slice(),
6416                                acts[bi + 3].xq.as_slice(),
6417                            ];
6418                            let d = unsafe {
6419                                if vnni_tiles_enabled() {
6420                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
6421                                } else {
6422                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
6423                                }
6424                            };
6425                            for k in 0..4 {
6426                                let act = &acts[bi + k];
6427                                let mut acc = d[k] * act.sx;
6428                                for &(j, xv) in &act.outliers {
6429                                    acc += (buf[j] as i8) as f32
6430                                        * gscale((r * cols + j) / GROUP_SIZE)
6431                                        * xv;
6432                                }
6433                                // SAFETY: disjoint (bi, r) cells per worker.
6434                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6435                            }
6436                            bi += 4;
6437                        }
6438                    }
6439                    while bi < acts.len() {
6440                        let act = &acts[bi];
6441                        let mut acc = 0f32;
6442                        for gi in 0..gpr {
6443                            let d = dot_i8_i8(
6444                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6445                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6446                            );
6447                            acc += d as f32 * gscale(r * gpr + gi);
6448                        }
6449                        acc *= act.sx;
6450                        // xq is zeroed at outlier slots — exact terms.
6451                        for &(j, xv) in &act.outliers {
6452                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
6453                        }
6454                        // SAFETY: disjoint (bi, r) cells per worker row range.
6455                        unsafe { *out_addr.at(bi * rows + r) = acc };
6456                        bi += 1;
6457                    }
6458                }
6459            })
6460        };
6461        dispatch_rows(pool, rows, &run);
6462        return;
6463    }
6464
6465    let out_addr = SendMut(out.as_mut_ptr());
6466    let run = move |start: usize, end: usize| {
6467        ROW_F32.with(|rb| {
6468            let mut buf = rb.borrow_mut();
6469            buf.resize(cols, 0.0);
6470            for r in start..end {
6471                // Decode raw (nib − 8) values once; scales stay per-group
6472                // so the accumulation order matches q4matvec bit-for-bit.
6473                for gi in 0..gpr {
6474                    let g = r * gpr + gi;
6475                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6476                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
6477                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
6478                    }
6479                }
6480                for bi in 0..b {
6481                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6482                    let mut acc = 0f32;
6483                    for gi in 0..gpr {
6484                        let mut ga = 0f32;
6485                        // Pairwise (lo + hi) addition, matching
6486                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
6487                        // a flat one-per-element loop rounds differently
6488                        // and broke bit-parity on the scalar (x86) path.
6489                        for k in 0..GROUP_SIZE / 2 {
6490                            let e = gi * GROUP_SIZE + k * 2;
6491                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
6492                        }
6493                        acc += ga * gscale(r * gpr + gi);
6494                    }
6495                    // SAFETY: disjoint (bi, r) cells per worker row range.
6496                    unsafe { *out_addr.at(bi * rows + r) = acc };
6497                }
6498            }
6499        })
6500    };
6501    dispatch_rows(pool, rows, &run);
6502}
6503
6504/// Batched vbit matmat: each variable-bit row is decoded from the mmap
6505/// ONCE for the whole microbatch. Same per-position math as
6506/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
6507/// and the scalar path).
6508#[allow(clippy::too_many_arguments)]
6509fn vbitmatmat(
6510    bytes: &[u8],
6511    offsets: &[usize],
6512    xs_all: &[f32],
6513    b: usize,
6514    rows: usize,
6515    cols: usize,
6516    out: &mut [f32],
6517    pool: Option<&Pool>,
6518) {
6519    debug_assert_eq!(xs_all.len(), b * cols);
6520    debug_assert_eq!(out.len(), b * rows);
6521    debug_assert_eq!(offsets.len(), rows + 1);
6522    let ng = cols / GROUP_SIZE;
6523    let bits = &bytes[..rows];
6524    let sc_off = rows;
6525    let gscale = |r: usize, g: usize| {
6526        let so = (r * ng + g) * 2;
6527        f16_to_f32(u16::from_le_bytes([
6528            bytes[sc_off + so],
6529            bytes[sc_off + so + 1],
6530        ]))
6531    };
6532
6533    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
6534    let decode_f32 = |r: usize, dst: &mut [f32]| {
6535        let bw = bits[r] as usize;
6536        let l = ((1i32 << (bw - 1)) - 1) as f32;
6537        let data = &bytes[offsets[r]..offsets[r + 1]];
6538        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
6539        for d in dst.iter_mut() {
6540            while nbits < bw {
6541                acc = (acc << 8) | data[idx] as u64;
6542                idx += 1;
6543                nbits += 8;
6544            }
6545            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
6546            nbits -= bw;
6547            *d = u - l;
6548        }
6549    };
6550
6551    if a8w8_enabled() {
6552        let acts: Vec<SplitAct> = (0..b)
6553            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6554            .collect();
6555        let acts = &acts;
6556        let out_addr = SendMut(out.as_mut_ptr());
6557        let run = move |start: usize, end: usize| {
6558            for r in start..end {
6559                let bw = bits[r] as usize;
6560                if bw == 8 {
6561                    // u−L reaches 128 → no i8 path; decode once, exact
6562                    // f32 dots for every position (same as vbitmatvec).
6563                    ROW_F32.with(|rb| {
6564                        let mut buf = rb.borrow_mut();
6565                        buf.resize(cols, 0.0);
6566                        decode_f32(r, &mut buf);
6567                        for bi in 0..b {
6568                            let x = &xs_all[bi * cols..(bi + 1) * cols];
6569                            let mut dot = 0f32;
6570                            for g in 0..ng {
6571                                let mut gd = 0f32;
6572                                for k in 0..GROUP_SIZE {
6573                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6574                                }
6575                                dot += gd * gscale(r, g);
6576                            }
6577                            // SAFETY: disjoint (bi, r) cells per worker range.
6578                            unsafe { *out_addr.at(bi * rows + r) = dot };
6579                        }
6580                    });
6581                    continue;
6582                }
6583                let l = (1i32 << (bw - 1)) - 1;
6584                let data = &bytes[offsets[r]..offsets[r + 1]];
6585                ROW_I8.with(|rb| {
6586                    let mut buf = rb.borrow_mut();
6587                    buf.resize(cols, 0);
6588                    #[inline(always)]
6589                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
6590                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
6591                            let u = unpack8::<B>(&data[blk * B..]);
6592                            for k in 0..8 {
6593                                chunk[k] = (u[k] - l) as i8 as u8;
6594                            }
6595                        }
6596                    }
6597                    match bw {
6598                        3 => fill::<3>(data, l, &mut buf),
6599                        4 => vbit_fill4(data, &mut buf),
6600                        5 => fill::<5>(data, l, &mut buf),
6601                        6 => fill::<6>(data, l, &mut buf),
6602                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
6603                    }
6604                    let mut bi = 0usize;
6605                    // The vbit scale table shares q4_block's layout
6606                    // (contiguous f16 per (row·ng + g)), so the same
6607                    // blocked 1×4 kernel serves the decoded row.
6608                    #[cfg(target_arch = "x86_64")]
6609                    if avx2_enabled()
6610                        && std::env::var("CMF_X86_BLOCKED")
6611                            .map(|v| v != "0")
6612                            .unwrap_or(true)
6613                    {
6614                        while bi + 4 <= acts.len() {
6615                            let xs = [
6616                                acts[bi].xq.as_slice(),
6617                                acts[bi + 1].xq.as_slice(),
6618                                acts[bi + 2].xq.as_slice(),
6619                                acts[bi + 3].xq.as_slice(),
6620                            ];
6621                            let sxs = [
6622                                acts[bi].sx,
6623                                acts[bi + 1].sx,
6624                                acts[bi + 2].sx,
6625                                acts[bi + 3].sx,
6626                            ];
6627                            let d = unsafe {
6628                                if vnni_tiles_enabled() {
6629                                    dot_q4b_row_1x4_sx_vnni(
6630                                        &buf,
6631                                        &bytes[sc_off..],
6632                                        r * ng,
6633                                        ng,
6634                                        xs,
6635                                        sxs,
6636                                    )
6637                                } else {
6638                                    dot_q4b_row_1x4_sx_avx2(
6639                                        &buf,
6640                                        &bytes[sc_off..],
6641                                        r * ng,
6642                                        ng,
6643                                        xs,
6644                                        sxs,
6645                                    )
6646                                }
6647                            };
6648                            for k in 0..4 {
6649                                let act = &acts[bi + k];
6650                                let mut dot = d[k];
6651                                for &(j, xv) in &act.outliers {
6652                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6653                                }
6654                                // SAFETY: disjoint (bi, r) cells per worker.
6655                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
6656                            }
6657                            bi += 4;
6658                        }
6659                    }
6660                    while bi < acts.len() {
6661                        let act = &acts[bi];
6662                        let mut dot = 0f32;
6663                        for g in 0..ng {
6664                            let d = dot_i8_i8(
6665                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6666                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6667                            ) as f32
6668                                * act.sx;
6669                            dot += d * gscale(r, g);
6670                        }
6671                        for &(j, xv) in &act.outliers {
6672                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6673                        }
6674                        // SAFETY: disjoint (bi, r) cells per worker range.
6675                        unsafe { *out_addr.at(bi * rows + r) = dot };
6676                        bi += 1;
6677                    }
6678                });
6679            }
6680        };
6681        dispatch_rows(pool, rows, &run);
6682        return;
6683    }
6684
6685    let out_addr = SendMut(out.as_mut_ptr());
6686    let run = move |start: usize, end: usize| {
6687        ROW_F32.with(|rb| {
6688            let mut buf = rb.borrow_mut();
6689            buf.resize(cols, 0.0);
6690            for r in start..end {
6691                decode_f32(r, &mut buf);
6692                for bi in 0..b {
6693                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6694                    let mut dot = 0f32;
6695                    for g in 0..ng {
6696                        let mut gd = 0f32;
6697                        for k in 0..GROUP_SIZE {
6698                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6699                        }
6700                        dot += gd * gscale(r, g);
6701                    }
6702                    // SAFETY: disjoint (bi, r) cells per worker range.
6703                    unsafe { *out_addr.at(bi * rows + r) = dot };
6704                }
6705            }
6706        })
6707    };
6708    dispatch_rows(pool, rows, &run);
6709}
6710
6711/// Build a GPU batch job for a q8-family mapped tensor (primary
6712/// shard): prescaled input + directory coordinates. None → not
6713/// GPU-eligible, caller stays on the CPU.
6714pub(crate) fn gpu_batch_job<'a>(
6715    t: &'a QTensor,
6716    x: &[f32],
6717) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
6718    match t {
6719        QTensor::Mapped {
6720            model,
6721            idx,
6722            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
6723            rows,
6724            cols,
6725            row_scale,
6726            col_field,
6727            ..
6728        } => Some((
6729            model.clone(),
6730            crate::gpu::BatchJob {
6731                idx: *idx,
6732                rows: *rows,
6733                cols: *cols,
6734                row_scale,
6735                xs: prescale(x, col_field, *dt).into_owned(),
6736                q1: false,
6737            },
6738        )),
6739        // q1: raw f32 activations, tile-embedded scales.
6740        QTensor::Mapped {
6741            model,
6742            idx,
6743            dtype: TensorDtype::Q1,
6744            rows,
6745            cols,
6746            ..
6747        } => Some((
6748            model.clone(),
6749            crate::gpu::BatchJob {
6750                idx: *idx,
6751                rows: *rows,
6752                cols: *cols,
6753                row_scale: &[],
6754                xs: x.to_vec(),
6755                q1: true,
6756            },
6757        )),
6758        _ => None,
6759    }
6760}
6761
6762thread_local! {
6763    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6764    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6765}
6766
6767pub(crate) fn prescale<'a>(
6768    x: &'a [f32],
6769    col_field: &[f32],
6770    dtype: TensorDtype,
6771) -> std::borrow::Cow<'a, [f32]> {
6772    if dtype == TensorDtype::Q8_2f {
6773        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
6774    } else {
6775        std::borrow::Cow::Borrowed(x)
6776    }
6777}
6778
6779/// θ col-field fold for q8_2f activations. Borrowed pass-through for
6780/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
6781pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
6782    x: &[f32],
6783    col_field: &[f32],
6784    dtype: TensorDtype,
6785    buf_id: u8,
6786    f: F,
6787) -> R {
6788    if dtype == TensorDtype::Q8_2f {
6789        if buf_id == 1 {
6790            PRESCALE_BUF1.with(|b| {
6791                let mut buf = b.borrow_mut();
6792                buf.clear();
6793                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6794                f(&buf)
6795            })
6796        } else {
6797            PRESCALE_BUF2.with(|b| {
6798                let mut buf = b.borrow_mut();
6799                buf.clear();
6800                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6801                f(&buf)
6802            })
6803        }
6804    } else {
6805        f(x)
6806    }
6807}
6808
6809// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
6810
6811/// AVX2+FMA available? Default ON when the CPU supports both;
6812/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
6813#[cfg(target_arch = "x86_64")]
6814pub(crate) fn avx2_enabled() -> bool {
6815    use std::sync::OnceLock;
6816    static ON: OnceLock<bool> = OnceLock::new();
6817    *ON.get_or_init(|| {
6818        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
6819            && std::arch::is_x86_feature_detected!("avx2")
6820            && std::arch::is_x86_feature_detected!("fma")
6821    })
6822}
6823
6824/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
6825/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
6826/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
6827/// active either way, they are exact (regrouped sums only).
6828#[cfg(target_arch = "x86_64")]
6829fn avx2_a8w8_enabled() -> bool {
6830    use std::sync::OnceLock;
6831    static ON: OnceLock<bool> = OnceLock::new();
6832    *ON.get_or_init(|| {
6833        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
6834    })
6835}
6836
6837/// A8W8 quantized-activation path available on THIS machine? One
6838/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
6839/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
6840#[inline]
6841pub(crate) fn a8w8_enabled() -> bool {
6842    #[cfg(target_arch = "aarch64")]
6843    {
6844        sdot_enabled()
6845    }
6846    #[cfg(target_arch = "x86_64")]
6847    {
6848        avx2_a8w8_enabled()
6849    }
6850    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6851    {
6852        false
6853    }
6854}
6855
6856/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
6857/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
6858#[inline]
6859#[allow(unreachable_code)]
6860fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
6861    #[cfg(target_arch = "aarch64")]
6862    unsafe {
6863        return dot_i8_sdot(w, xq);
6864    }
6865    #[cfg(target_arch = "x86_64")]
6866    unsafe {
6867        if avx512vnni_enabled() {
6868            return dot_i8_i8_vnni(w, xq);
6869        }
6870        return dot_i8_i8_avx2(w, xq);
6871    }
6872    w.iter()
6873        .zip(xq)
6874        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
6875        .sum()
6876}
6877
6878/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
6879/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
6880/// `vpdpbusd` encoding.
6881#[cfg(target_arch = "x86_64")]
6882fn avx512vnni_enabled() -> bool {
6883    use std::sync::OnceLock;
6884    static ON: OnceLock<bool> = OnceLock::new();
6885    *ON.get_or_init(|| {
6886        std::env::var("CMF_AVX512")
6887            .map(|v| v != "0")
6888            .unwrap_or(true)
6889            && std::arch::is_x86_feature_detected!("avx512f")
6890            && std::arch::is_x86_feature_detected!("avx512bw")
6891            && std::arch::is_x86_feature_detected!("avx512vl")
6892            && std::arch::is_x86_feature_detected!("avx512vnni")
6893    })
6894}
6895
6896/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
6897/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
6898/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
6899/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
6900/// (+4%) — consistent, no leg regressed. The tile kernels keep a
6901/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
6902/// smaller than the long-dot q8 win (+13%), but it is real and free.
6903#[cfg(target_arch = "x86_64")]
6904fn vnni_tiles_enabled() -> bool {
6905    use std::sync::OnceLock;
6906    static ON: OnceLock<bool> = OnceLock::new();
6907    *ON.get_or_init(|| {
6908        std::env::var("CMF_VNNI_TILES")
6909            .map(|v| v != "0")
6910            .unwrap_or(true)
6911            && avx512vnni_enabled()
6912    })
6913}
6914
6915/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
6916/// plus the same horizontal reduce the AVX2 kernels use. Products are
6917/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
6918/// is bit-identical to the maddubs+madd pair it replaces.
6919#[cfg(target_arch = "x86_64")]
6920#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6921#[inline]
6922unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
6923    // SAFETY: pure register math.
6924    unsafe {
6925        use core::arch::x86_64::*;
6926        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
6927        let hi128 = _mm256_extracti128_si256::<1>(d);
6928        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6929        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6930        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6931        _mm_cvtsi128_si32(s32)
6932    }
6933}
6934
6935/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
6936/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
6937/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
6938/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
6939#[cfg(target_arch = "x86_64")]
6940#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6941unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
6942    // SAFETY: callers uphold slice-length contracts (see call sites).
6943    unsafe {
6944        use core::arch::x86_64::*;
6945        let n = w.len();
6946        let mut j = 0usize;
6947        let mut total: i32;
6948        // 4 independent accumulators: vpdpbusd is its own loop-carried
6949        // dependency (~5-cycle latency) — a single-acc loop runs
6950        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
6951        // on Granite Rapids.
6952        {
6953            #[inline(always)]
6954            unsafe fn step(
6955                w: *const u8,
6956                x: *const i8,
6957                acc: core::arch::x86_64::__m512i,
6958            ) -> core::arch::x86_64::__m512i {
6959                unsafe {
6960                    use core::arch::x86_64::*;
6961                    let wv = _mm512_loadu_si512(w as *const _);
6962                    let xv = _mm512_loadu_si512(x as *const _);
6963                    let aw = _mm512_abs_epi8(wv);
6964                    let neg = _mm512_movepi8_mask(wv);
6965                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
6966                    _mm512_dpbusd_epi32(acc, aw, sx)
6967                }
6968            }
6969            let (mut a0, mut a1, mut a2, mut a3) = (
6970                _mm512_setzero_si512(),
6971                _mm512_setzero_si512(),
6972                _mm512_setzero_si512(),
6973                _mm512_setzero_si512(),
6974            );
6975            while j + 256 <= n {
6976                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
6977                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
6978                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
6979                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
6980                j += 256;
6981            }
6982            while j + 64 <= n {
6983                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
6984                j += 64;
6985            }
6986            let s01 = _mm512_add_epi32(a0, a1);
6987            let s23 = _mm512_add_epi32(a2, a3);
6988            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
6989        }
6990        // 32-wide (q4/vbit groups are exactly 32 bytes).
6991        if j + 32 <= n {
6992            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
6993            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
6994            let d = _mm256_dpbusd_epi32(
6995                _mm256_setzero_si256(),
6996                _mm256_abs_epi8(wv),
6997                _mm256_sign_epi8(xv, wv),
6998            );
6999            let hi128 = _mm256_extracti128_si256::<1>(d);
7000            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7001            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7002            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7003            total += _mm_cvtsi128_si32(s32);
7004            j += 32;
7005        }
7006        while j < n {
7007            total += (w[j] as i8) as i32 * xq[j] as i32;
7008            j += 1;
7009        }
7010        total
7011    }
7012}
7013
7014/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
7015#[cfg(target_arch = "x86_64")]
7016#[target_feature(enable = "avx2,fma")]
7017unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
7018    // SAFETY: callers uphold slice-length contracts (see call sites).
7019    unsafe {
7020        use core::arch::x86_64::*;
7021        let n = x.len();
7022        let wp = w.as_ptr();
7023        let xp = x.as_ptr();
7024        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
7025        let mut j = 0usize;
7026        while j + 16 <= n {
7027            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
7028            let lo = _mm256_cvtepi8_epi32(wb);
7029            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
7030            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
7031            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
7032            j += 16;
7033        }
7034        let acc = _mm256_add_ps(a0, a1);
7035        let hi128 = _mm256_extractf128_ps::<1>(acc);
7036        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
7037        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
7038        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
7039        let mut sum = _mm_cvtss_f32(s32);
7040        while j < n {
7041            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
7042            j += 1;
7043        }
7044        sum
7045    }
7046}
7047
7048/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
7049/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
7050/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
7051/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
7052#[cfg(target_arch = "x86_64")]
7053#[target_feature(enable = "avx2")]
7054unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
7055    // SAFETY: callers uphold slice-length contracts (see call sites).
7056    unsafe {
7057        use core::arch::x86_64::*;
7058        let n = w.len();
7059        let ones = _mm256_set1_epi16(1);
7060        let mut acc = _mm256_setzero_si256();
7061        let mut j = 0usize;
7062        while j + 32 <= n {
7063            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7064            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7065            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7066            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
7067            j += 32;
7068        }
7069        let hi128 = _mm256_extracti128_si256::<1>(acc);
7070        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
7071        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7072        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7073        let mut s = _mm_cvtsi128_si32(s32);
7074        while j < n {
7075            s += (w[j] as i8) as i32 * xq[j] as i32;
7076            j += 1;
7077        }
7078        s
7079    }
7080}
7081
7082/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
7083/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
7084/// slice as a combined 2×8 register and meets two activation pairs.
7085#[cfg(target_arch = "aarch64")]
7086#[target_feature(enable = "neon,i8mm")]
7087unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7088    // SAFETY: callers uphold slice-length contracts.
7089    unsafe {
7090        use core::arch::aarch64::*;
7091        use core::arch::asm;
7092        let n = w0.len();
7093        let w0p = w0.as_ptr() as *const i8;
7094        let w1p = w1.as_ptr() as *const i8;
7095        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
7096        // same for x2/x3.
7097        let mut acc01 = vdupq_n_s32(0);
7098        let mut acc23 = vdupq_n_s32(0);
7099        let mut i = 0usize;
7100        while i + 8 <= n {
7101            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
7102            let xb01 = vcombine_s8(
7103                vld1_s8(xs[0].as_ptr().add(i)),
7104                vld1_s8(xs[1].as_ptr().add(i)),
7105            );
7106            let xb23 = vcombine_s8(
7107                vld1_s8(xs[2].as_ptr().add(i)),
7108                vld1_s8(xs[3].as_ptr().add(i)),
7109            );
7110            asm!(
7111                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
7112                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
7113                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
7114                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
7115                options(pure, nomem, nostack),
7116            );
7117            i += 8;
7118        }
7119        let mut out = [[0i32; 4]; 2];
7120        let a01: [i32; 4] = core::mem::transmute(acc01);
7121        let a23: [i32; 4] = core::mem::transmute(acc23);
7122        out[0][0] = a01[0];
7123        out[0][1] = a01[1];
7124        out[1][0] = a01[2];
7125        out[1][1] = a01[3];
7126        out[0][2] = a23[0];
7127        out[0][3] = a23[1];
7128        out[1][2] = a23[2];
7129        out[1][3] = a23[3];
7130        if i < n {
7131            for (k, x) in xs.iter().enumerate() {
7132                for j in i..n {
7133                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7134                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7135                }
7136            }
7137        }
7138        out
7139    }
7140}
7141
7142/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
7143/// registers across four activation streams, eight sdot accumulators.
7144/// (The per-row form re-read each W row once per activation.)
7145#[cfg(target_arch = "aarch64")]
7146#[target_feature(enable = "neon,dotprod")]
7147unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7148    // SAFETY: callers uphold slice-length contracts.
7149    unsafe {
7150        use core::arch::aarch64::*;
7151        use core::arch::asm;
7152        let n = w0.len();
7153        let w0p = w0.as_ptr() as *const i8;
7154        let w1p = w1.as_ptr() as *const i8;
7155        let mut acc = [[vdupq_n_s32(0); 4]; 2];
7156        let mut i = 0usize;
7157        while i + 16 <= n {
7158            let wv0 = vld1q_s8(w0p.add(i));
7159            let wv1 = vld1q_s8(w1p.add(i));
7160            for (k, x) in xs.iter().enumerate() {
7161                let xv = vld1q_s8(x.as_ptr().add(i));
7162                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
7163                asm!(
7164                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
7165                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
7166                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7167                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
7168                    options(pure, nomem, nostack),
7169                );
7170                acc[0][k] = a0;
7171                acc[1][k] = a1;
7172            }
7173            i += 16;
7174        }
7175        let mut out = [[0i32; 4]; 2];
7176        for r in 0..2 {
7177            for k in 0..4 {
7178                out[r][k] = vaddvq_s32(acc[r][k]);
7179            }
7180        }
7181        if i < n {
7182            for (k, x) in xs.iter().enumerate() {
7183                for j in i..n {
7184                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7185                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7186                }
7187            }
7188        }
7189        out
7190    }
7191}
7192
7193/// Blocked 2 weight rows × 4 activations for the prefill GEMM
7194/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
7195/// abs() live in registers across all four activation streams; the
7196/// sign-fixup is recomputed per pair (the price of the maddubs trick).
7197/// Returns raw i8·i8 dots; the caller applies scales and outliers.
7198#[cfg(target_arch = "x86_64")]
7199#[target_feature(enable = "avx2")]
7200unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7201    // SAFETY: callers uphold slice-length contracts.
7202    unsafe {
7203        use core::arch::x86_64::*;
7204        let n = w0.len();
7205        let ones = _mm256_set1_epi16(1);
7206        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
7207        let mut j = 0usize;
7208        while j + 32 <= n {
7209            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
7210            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
7211            let aw0 = _mm256_abs_epi8(wv0);
7212            let aw1 = _mm256_abs_epi8(wv1);
7213            for (k, x) in xs.iter().enumerate() {
7214                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
7215                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
7216                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
7217                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
7218                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
7219            }
7220            j += 32;
7221        }
7222        let mut out = [[0i32; 4]; 2];
7223        for r in 0..2 {
7224            for k in 0..4 {
7225                let a = acc[r][k];
7226                let hi128 = _mm256_extracti128_si256::<1>(a);
7227                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
7228                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7229                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7230                out[r][k] = _mm_cvtsi128_si32(s32);
7231            }
7232        }
7233        if j < n {
7234            for (k, x) in xs.iter().enumerate() {
7235                for i in j..n {
7236                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
7237                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
7238                }
7239            }
7240        }
7241        out
7242    }
7243}
7244
7245/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
7246/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
7247/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
7248/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
7249#[cfg(target_arch = "x86_64")]
7250#[inline]
7251fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
7252    let dot = if avx512vnni_enabled() && row.len() >= 64 {
7253        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
7254    } else {
7255        unsafe { dot_i8_i8_avx2(row, &act.xq) }
7256    };
7257    let mut acc = dot as f32 * act.sx;
7258    for &(j, xv) in &act.outliers {
7259        acc += (row[j] as i8) as f32 * xv;
7260    }
7261    acc
7262}
7263
7264/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
7265/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
7266/// a single-acc loop runs latency-bound, measured on Granite Rapids).
7267#[cfg(target_arch = "x86_64")]
7268#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7269unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7270    // SAFETY: callers uphold slice-length contracts (see call sites).
7271    unsafe {
7272        use core::arch::x86_64::*;
7273        let n = w.len();
7274        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
7275        #[inline(always)]
7276        unsafe fn step(
7277            w: *const u8,
7278            x: *const i8,
7279            flip: core::arch::x86_64::__m512i,
7280            acc: core::arch::x86_64::__m512i,
7281        ) -> core::arch::x86_64::__m512i {
7282            unsafe {
7283                use core::arch::x86_64::*;
7284                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
7285                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
7286            }
7287        }
7288        let (mut a0, mut a1, mut a2, mut a3) = (
7289            _mm512_setzero_si512(),
7290            _mm512_setzero_si512(),
7291            _mm512_setzero_si512(),
7292            _mm512_setzero_si512(),
7293        );
7294        let mut j = 0usize;
7295        while j + 256 <= n {
7296            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7297            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
7298            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
7299            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
7300            j += 256;
7301        }
7302        while j + 64 <= n {
7303            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7304            j += 64;
7305        }
7306        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
7307            _mm512_add_epi32(a0, a1),
7308            _mm512_add_epi32(a2, a3),
7309        ));
7310        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
7311        while j < n {
7312            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
7313            j += 1;
7314        }
7315        total
7316    }
7317}
7318
7319/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
7320/// writer's flat order, same as the NEON vzip pair), maddubs against
7321/// the pre-quantized activation group, × the group's f16 scale. Pair
7322/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
7323/// `dot_q4_row_sdot`.
7324#[cfg(target_arch = "x86_64")]
7325#[target_feature(enable = "avx2")]
7326unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7327    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7328    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
7329    unsafe {
7330        use core::arch::x86_64::*;
7331        let lomask = _mm_set1_epi8(0x0F);
7332        let eight = _mm256_set1_epi8(8);
7333        let ones = _mm256_set1_epi16(1);
7334        let mut acc = 0f32;
7335        for gi in 0..gpr {
7336            let g = g0 + gi;
7337            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7338            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7339            let lo = _mm_and_si128(b, lomask);
7340            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7341            let w = _mm256_sub_epi8(
7342                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7343                eight,
7344            );
7345            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7346            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
7347            let d = _mm256_madd_epi16(p16, ones);
7348            let hi128 = _mm256_extracti128_si256::<1>(d);
7349            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7350            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7351            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7352            acc += _mm_cvtsi128_si32(s32) as f32 * s;
7353        }
7354        acc
7355    }
7356}
7357
7358/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
7359/// both activations dotted against the same centered i8 register.
7360#[cfg(target_arch = "x86_64")]
7361#[target_feature(enable = "avx2")]
7362unsafe fn dot_q4_row_avx2_2(
7363    packed: &[u8],
7364    scales: &[u8],
7365    g0: usize,
7366    gpr: usize,
7367    xq1: &[i8],
7368    xq2: &[i8],
7369) -> (f32, f32) {
7370    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
7371    unsafe {
7372        use core::arch::x86_64::*;
7373        let lomask = _mm_set1_epi8(0x0F);
7374        let eight = _mm256_set1_epi8(8);
7375        let ones = _mm256_set1_epi16(1);
7376        let (mut acc1, mut acc2) = (0f32, 0f32);
7377        #[inline(always)]
7378        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
7379            unsafe {
7380                use core::arch::x86_64::*;
7381                let hi128 = _mm256_extracti128_si256::<1>(d);
7382                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7383                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7384                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7385                _mm_cvtsi128_si32(s32)
7386            }
7387        }
7388        for gi in 0..gpr {
7389            let g = g0 + gi;
7390            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7391            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7392            let lo = _mm_and_si128(b, lomask);
7393            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7394            let w = _mm256_sub_epi8(
7395                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7396                eight,
7397            );
7398            let aw = _mm256_abs_epi8(w);
7399            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7400            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7401            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
7402            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
7403            acc1 += hsum(d1) as f32 * s;
7404            acc2 += hsum(d2) as f32 * s;
7405        }
7406        (acc1, acc2)
7407    }
7408}
7409
7410/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
7411#[cfg(target_arch = "x86_64")]
7412fn q8_range_avx2(
7413    q: &[u8],
7414    row_scale: &[f32],
7415    act: &SplitAct,
7416    cols: usize,
7417    out_addr: SendMut,
7418    start: usize,
7419    end: usize,
7420) {
7421    for o in start..end {
7422        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7423        // SAFETY: disjoint row ranges per worker.
7424        unsafe { *out_addr.at(o) = v };
7425    }
7426}
7427
7428/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
7429#[cfg(target_arch = "x86_64")]
7430#[allow(clippy::too_many_arguments)]
7431fn q8_range2_avx2(
7432    q: &[u8],
7433    row_scale: &[f32],
7434    a1: &SplitAct,
7435    a2: &SplitAct,
7436    cols: usize,
7437    p1: SendMut,
7438    p2: SendMut,
7439    start: usize,
7440    end: usize,
7441) {
7442    for o in start..end {
7443        let row = &q[o * cols..(o + 1) * cols];
7444        // SAFETY: disjoint row ranges per worker.
7445        unsafe {
7446            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
7447            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
7448        }
7449    }
7450}
7451
7452// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
7453
7454/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
7455/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
7456/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
7457/// accumulator dependency chain swamp the MAC advantage, and Apple's
7458/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
7459/// field trials on Cortex-A710/X-class parts with two pipes, where the
7460/// balance may differ; a pre-interleaved weight layout (repack infra)
7461/// is the known path if it ever earns its keep.
7462#[cfg(target_arch = "aarch64")]
7463fn i8mm_enabled() -> bool {
7464    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7465    *ON.get_or_init(|| {
7466        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
7467            && std::arch::is_aarch64_feature_detected!("i8mm")
7468    })
7469}
7470
7471/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
7472/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
7473/// (On non-ARM release builds only the test tolerance switch calls it.)
7474#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
7475fn sdot_enabled() -> bool {
7476    use std::sync::OnceLock;
7477    static ON: OnceLock<bool> = OnceLock::new();
7478    *ON.get_or_init(|| {
7479        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
7480        if !want {
7481            return false;
7482        }
7483
7484        #[cfg(target_arch = "aarch64")]
7485        {
7486            if std::arch::is_aarch64_feature_detected!("dotprod") {
7487                return true;
7488            }
7489            #[cfg(target_os = "android")]
7490            {
7491                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
7492                    if cpuinfo.lines().any(|l| {
7493                        (l.starts_with("Features") || l.starts_with("features"))
7494                            && l.contains("asimddp")
7495                    }) {
7496                        return true;
7497                    }
7498                }
7499            }
7500            false
7501        }
7502        #[cfg(not(target_arch = "aarch64"))]
7503        {
7504            false
7505        }
7506    })
7507}
7508
7509/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
7510/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
7511/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
7512/// matvec, shared by all rows/workers.
7513struct SplitAct {
7514    xq: Vec<i8>,
7515    sx: f32,
7516    outliers: Vec<(usize, f32)>,
7517    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
7518    /// `−128·Σx`); one i32 per split, computed once per matvec.
7519    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
7520    xsum: i32,
7521}
7522
7523thread_local! {
7524    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
7525    /// and its hidden-size allocation was steady-state heap churn.
7526    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
7527        const { std::cell::RefCell::new(Vec::new()) };
7528}
7529
7530impl Drop for SplitAct {
7531    fn drop(&mut self) {
7532        let buf = std::mem::take(&mut self.xq);
7533        if buf.capacity() > 0 {
7534            XQ_FREE.with(|f| {
7535                let mut f = f.borrow_mut();
7536                if f.len() < 16 {
7537                    f.push(buf);
7538                }
7539            });
7540        }
7541    }
7542}
7543
7544fn split_act(x: &[f32]) -> SplitAct {
7545    let n = x.len();
7546    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
7547    let thr = 8.0 * rms;
7548    // One pass: collect outliers and the bulk absmax (outliers excluded —
7549    // identical to the old zero-then-fold over a copied buffer, minus the
7550    // full-vector copy).
7551    let mut outliers: Vec<(usize, f32)> = Vec::new();
7552    let mut amax = 0f32;
7553    for (j, &v) in x.iter().enumerate() {
7554        let a = v.abs();
7555        if a > thr {
7556            outliers.push((j, v));
7557        } else if a > amax {
7558            amax = a;
7559        }
7560    }
7561    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7562    let inv = 1.0 / sx;
7563    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7564    xq.clear();
7565    xq.reserve(n);
7566    if outliers.is_empty() {
7567        xq.extend(
7568            x.iter()
7569                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
7570        );
7571    } else {
7572        // Outlier slots quantize to 0 (their exact term is added later).
7573        xq.extend(x.iter().map(|&v| {
7574            if v.abs() > thr {
7575                0
7576            } else {
7577                (v * inv).round().clamp(-127.0, 127.0) as i8
7578            }
7579        }));
7580    }
7581    let xsum = xq.iter().map(|&v| v as i32).sum();
7582    SplitAct {
7583        xq,
7584        sx,
7585        outliers,
7586        xsum,
7587    }
7588}
7589
7590fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
7591    let n = x.len();
7592    let rms = (x
7593        .iter()
7594        .zip(col)
7595        .map(|(&a, &c)| {
7596            let v = a * c;
7597            (v * v) as f64
7598        })
7599        .sum::<f64>()
7600        / n.max(1) as f64)
7601        .sqrt() as f32;
7602    let thr = 8.0 * rms;
7603
7604    let mut outliers = Vec::new();
7605    let mut amax = 0f32;
7606    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
7607        let v = a * c;
7608        let s = v.abs();
7609        if s > thr {
7610            outliers.push((j, v));
7611        } else if s > amax {
7612            amax = s;
7613        }
7614    }
7615
7616    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7617    let inv = 1.0 / sx;
7618    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7619    xq.clear();
7620    xq.reserve(n);
7621    if outliers.is_empty() {
7622        xq.extend(
7623            x.iter()
7624                .zip(col)
7625                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
7626        );
7627    } else {
7628        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
7629            let v = a * c;
7630            if v.abs() > thr {
7631                0
7632            } else {
7633                (v * inv).round().clamp(-127.0, 127.0) as i8
7634            }
7635        }));
7636    }
7637    let xsum = xq.iter().map(|&v| v as i32).sum();
7638    SplitAct {
7639        xq,
7640        sx,
7641        outliers,
7642        xsum,
7643    }
7644}
7645
7646/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
7647/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
7648#[cfg(target_arch = "aarch64")]
7649#[target_feature(enable = "neon,dotprod")]
7650unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
7651    // SAFETY: callers uphold slice-length contracts (see call sites).
7652    unsafe {
7653        use core::arch::aarch64::*;
7654        use core::arch::asm;
7655        let wp = w.as_ptr() as *const i8;
7656        let n = w.len();
7657        let (mut a0, mut a1, mut a2, mut a3) = (
7658            vdupq_n_s32(0),
7659            vdupq_n_s32(0),
7660            vdupq_n_s32(0),
7661            vdupq_n_s32(0),
7662        );
7663        let mut i = 0;
7664        while i + 64 <= n {
7665            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7666            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
7667            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
7668            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
7669            asm!(
7670                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7671                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7672                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
7673                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
7674                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7675                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7676                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
7677                options(pure, nomem, nostack),
7678            );
7679            i += 64;
7680        }
7681        while i + 16 <= n {
7682            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7683            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
7684                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
7685            i += 16;
7686        }
7687        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
7688        while i < n {
7689            s += (*wp.add(i)) as i32 * xq[i] as i32;
7690            i += 1;
7691        }
7692        s
7693    }
7694}
7695
7696/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
7697/// loaded once and reused, 4 independent accumulators hide sdot latency
7698/// (port of vmfcore `dot_i8_sdot_4rows`).
7699#[cfg(target_arch = "aarch64")]
7700#[target_feature(enable = "neon,dotprod")]
7701unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
7702    // SAFETY: callers uphold slice-length contracts (see call sites).
7703    unsafe {
7704        use core::arch::aarch64::*;
7705        use core::arch::asm;
7706        let n = xq.len();
7707        let px = xq.as_ptr();
7708        let (p0, p1, p2, p3) = (
7709            w0.as_ptr() as *const i8,
7710            w1.as_ptr() as *const i8,
7711            w2.as_ptr() as *const i8,
7712            w3.as_ptr() as *const i8,
7713        );
7714        let (mut a0, mut a1, mut a2, mut a3) = (
7715            vdupq_n_s32(0),
7716            vdupq_n_s32(0),
7717            vdupq_n_s32(0),
7718            vdupq_n_s32(0),
7719        );
7720        let mut i = 0;
7721        while i + 16 <= n {
7722            let x = vld1q_s8(px.add(i));
7723            let v0 = vld1q_s8(p0.add(i));
7724            let v1 = vld1q_s8(p1.add(i));
7725            let v2 = vld1q_s8(p2.add(i));
7726            let v3 = vld1q_s8(p3.add(i));
7727            asm!(
7728                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7729                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7730                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7731                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7732                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7733                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7734                options(pure, nomem, nostack),
7735            );
7736            i += 16;
7737        }
7738        let mut r = [
7739            vaddvq_s32(a0),
7740            vaddvq_s32(a1),
7741            vaddvq_s32(a2),
7742            vaddvq_s32(a3),
7743        ];
7744        while i < n {
7745            let xi = *px.add(i) as i32;
7746            r[0] += (*p0.add(i)) as i32 * xi;
7747            r[1] += (*p1.add(i)) as i32 * xi;
7748            r[2] += (*p2.add(i)) as i32 * xi;
7749            r[3] += (*p3.add(i)) as i32 * xi;
7750            i += 1;
7751        }
7752        r
7753    }
7754}
7755
7756/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
7757/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
7758/// line plus the shared activation chunk — a single sequential weight
7759/// stream per worker. Per-row accumulation is the same one-accumulator
7760/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
7761/// are bit-identical to the mmap-layout kernel.
7762#[cfg(target_arch = "aarch64")]
7763#[target_feature(enable = "neon,dotprod")]
7764unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
7765    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
7766    // n % 16 == 0 — guaranteed by the repack gate).
7767    unsafe {
7768        use core::arch::aarch64::*;
7769        use core::arch::asm;
7770        let n = xq.len();
7771        let px = xq.as_ptr();
7772        let pg = g.as_ptr() as *const i8;
7773        let (mut a0, mut a1, mut a2, mut a3) = (
7774            vdupq_n_s32(0),
7775            vdupq_n_s32(0),
7776            vdupq_n_s32(0),
7777            vdupq_n_s32(0),
7778        );
7779        let mut i = 0;
7780        while i + 16 <= n {
7781            let x = vld1q_s8(px.add(i));
7782            let base = pg.add(4 * i);
7783            let v0 = vld1q_s8(base);
7784            let v1 = vld1q_s8(base.add(16));
7785            let v2 = vld1q_s8(base.add(32));
7786            let v3 = vld1q_s8(base.add(48));
7787            asm!(
7788                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7789                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7790                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7791                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7792                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7793                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7794                options(pure, nomem, nostack),
7795            );
7796            i += 16;
7797        }
7798        [
7799            vaddvq_s32(a0),
7800            vaddvq_s32(a1),
7801            vaddvq_s32(a2),
7802            vaddvq_s32(a3),
7803        ]
7804    }
7805}
7806
7807/// One q8 row range via SDOT (4-row blocks + tail) — the body of
7808/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
7809/// SAME kernel for several tensors under one pool dispatch. `rep` — the
7810/// load-time interleaved repack (empty = mmap layout only); rows outside
7811/// full 4-row groups always come from the mmap layout.
7812#[cfg(target_arch = "aarch64")]
7813fn q8_range_sdot(
7814    q: &[u8],
7815    rep: &[u8],
7816    row_scale: &[f32],
7817    act: &SplitAct,
7818    cols: usize,
7819    out_addr: SendMut,
7820    start: usize,
7821    end: usize,
7822) {
7823    let mut o = start;
7824    // Leading rows to the group boundary (repack path only): the pool
7825    // splits row ranges arbitrarily, groups are absolute.
7826    if !rep.is_empty() {
7827        while o < end && o % 4 != 0 {
7828            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7829            unsafe { *out_addr.at(o) = v };
7830            o += 1;
7831        }
7832    }
7833    while o + 4 <= end {
7834        let r = if rep.is_empty() {
7835            unsafe {
7836                dot_i8_sdot_4rows(
7837                    &q[o * cols..(o + 1) * cols],
7838                    &q[(o + 1) * cols..(o + 2) * cols],
7839                    &q[(o + 2) * cols..(o + 3) * cols],
7840                    &q[(o + 3) * cols..(o + 4) * cols],
7841                    &act.xq,
7842                )
7843            }
7844        } else {
7845            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
7846        };
7847        for k in 0..4 {
7848            let mut acc = r[k] as f32 * act.sx;
7849            for &(j, xv) in &act.outliers {
7850                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
7851            }
7852            // SAFETY: disjoint row ranges per worker.
7853            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
7854        }
7855        o += 4;
7856    }
7857    while o < end {
7858        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7859        unsafe { *out_addr.at(o) = v };
7860        o += 1;
7861    }
7862}
7863
7864/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
7865/// for the fused pair multi-matrix job (`matvec2_many`).
7866#[cfg(target_arch = "aarch64")]
7867#[allow(clippy::too_many_arguments)]
7868fn q8_range2_sdot(
7869    q: &[u8],
7870    row_scale: &[f32],
7871    a1: &SplitAct,
7872    a2: &SplitAct,
7873    cols: usize,
7874    p1: SendMut,
7875    p2: SendMut,
7876    start: usize,
7877    end: usize,
7878) {
7879    for o in start..end {
7880        let row = &q[o * cols..(o + 1) * cols];
7881        // SAFETY: disjoint row ranges per worker.
7882        unsafe {
7883            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
7884            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
7885        }
7886    }
7887}
7888
7889/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
7890#[allow(clippy::too_many_arguments)]
7891fn q8_range2_f32(
7892    q: &[u8],
7893    row_scale: &[f32],
7894    x1: &[f32],
7895    x2: &[f32],
7896    cols: usize,
7897    p1: SendMut,
7898    p2: SendMut,
7899    start: usize,
7900    end: usize,
7901) {
7902    for o in start..end {
7903        let row = &q[o * cols..(o + 1) * cols];
7904        // SAFETY: disjoint row ranges per worker.
7905        unsafe {
7906            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
7907            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
7908        }
7909    }
7910}
7911
7912/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
7913fn q8_range_f32(
7914    q: &[u8],
7915    row_scale: &[f32],
7916    xs: &[f32],
7917    cols: usize,
7918    out_addr: SendMut,
7919    start: usize,
7920    end: usize,
7921) {
7922    for o in start..end {
7923        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
7924        // SAFETY: disjoint row ranges per worker.
7925        unsafe { *out_addr.at(o) = v };
7926    }
7927}
7928
7929/// SDOT row dot with exact outlier correction:
7930/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
7931#[cfg(target_arch = "aarch64")]
7932#[inline]
7933fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
7934    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
7935    for &(j, xv) in &act.outliers {
7936        acc += (row[j] as i8) as f32 * xv;
7937    }
7938    acc
7939}
7940
7941/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
7942/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
7943/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
7944/// the caller multiplies by the activation scale and adds the exact
7945/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
7946/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
7947/// → zip(lo,hi) restores flat order.
7948#[cfg(target_arch = "aarch64")]
7949#[target_feature(enable = "neon,dotprod")]
7950unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7951    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7952    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
7953    unsafe {
7954        use core::arch::aarch64::*;
7955        use core::arch::asm;
7956        let lomask = vdupq_n_u8(0x0F);
7957        let eight = vdupq_n_s8(8);
7958        let mut acc = 0f32;
7959        for gi in 0..gpr {
7960            let g = g0 + gi;
7961            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7962            let b = vld1q_u8(packed.as_ptr().add(g * 16));
7963            let lo = vandq_u8(b, lomask);
7964            let hi = vshrq_n_u8::<4>(b);
7965            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7966            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7967            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
7968            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
7969            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7970            asm!(
7971                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7972                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7973                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7974                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7975                options(pure, nomem, nostack),
7976            );
7977            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7978        }
7979        acc
7980    }
7981}
7982
7983/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
7984/// part) happens ONCE per group; both pre-quantized activations are
7985/// dotted against the same centered i8 registers. Per-lane math matches
7986/// `dot_q4_row_sdot` exactly.
7987#[cfg(target_arch = "aarch64")]
7988#[target_feature(enable = "neon,dotprod")]
7989unsafe fn dot_q4_row_sdot2(
7990    packed: &[u8],
7991    scales: &[u8],
7992    g0: usize,
7993    gpr: usize,
7994    xq1: &[i8],
7995    xq2: &[i8],
7996) -> (f32, f32) {
7997    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7998    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
7999    unsafe {
8000        use core::arch::aarch64::*;
8001        use core::arch::asm;
8002        let lomask = vdupq_n_u8(0x0F);
8003        let eight = vdupq_n_s8(8);
8004        let (mut acc1, mut acc2) = (0f32, 0f32);
8005        for gi in 0..gpr {
8006            let g = g0 + gi;
8007            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8008            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8009            let lo = vandq_u8(b, lomask);
8010            let hi = vshrq_n_u8::<4>(b);
8011            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8012            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8013            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
8014            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
8015            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
8016            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
8017            let (mut a0, mut a1, mut b0, mut b1) = (
8018                vdupq_n_s32(0),
8019                vdupq_n_s32(0),
8020                vdupq_n_s32(0),
8021                vdupq_n_s32(0),
8022            );
8023            asm!(
8024                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
8025                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
8026                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
8027                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
8028                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8029                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
8030                e0 = in(vreg) e0, e1 = in(vreg) e1,
8031                x10 = in(vreg) x10, x11 = in(vreg) x11,
8032                x20 = in(vreg) x20, x21 = in(vreg) x21,
8033                options(pure, nomem, nostack),
8034            );
8035            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8036            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
8037        }
8038        (acc1, acc2)
8039    }
8040}
8041
8042// ───────────────────── fused int8 kernels ─────────────────────
8043
8044/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
8045/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
8046#[inline]
8047pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
8048    #[cfg(target_arch = "aarch64")]
8049    unsafe {
8050        return axpy_i8_f32_neon(acc, row, w);
8051    }
8052    #[cfg(target_arch = "x86_64")]
8053    if avx2_enabled() {
8054        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
8055    }
8056    #[allow(unreachable_code)]
8057    {
8058        for (a, &b) in acc.iter_mut().zip(row) {
8059            *a += w * b as f32;
8060        }
8061    }
8062}
8063
8064/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
8065#[cfg(target_arch = "x86_64")]
8066#[target_feature(enable = "avx2,fma")]
8067unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
8068    // SAFETY: callers uphold slice-length contracts (see call sites).
8069    unsafe {
8070        use core::arch::x86_64::*;
8071        let n = acc.len().min(row.len());
8072        let ap = acc.as_mut_ptr();
8073        let rp = row.as_ptr();
8074        let wv = _mm256_set1_ps(w);
8075        let mut j = 0usize;
8076        while j + 16 <= n {
8077            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
8078            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
8079            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
8080            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
8081            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
8082            _mm256_storeu_ps(ap.add(j), v0);
8083            _mm256_storeu_ps(ap.add(j + 8), v1);
8084            j += 16;
8085        }
8086        while j < n {
8087            *ap.add(j) += w * (*rp.add(j)) as f32;
8088            j += 1;
8089        }
8090    }
8091}
8092
8093#[cfg(target_arch = "aarch64")]
8094#[target_feature(enable = "neon")]
8095unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
8096    // SAFETY: callers uphold slice-length contracts (see call sites).
8097    unsafe {
8098        use core::arch::aarch64::*;
8099        let n = acc.len().min(row.len());
8100        let ap = acc.as_mut_ptr();
8101        let rp = row.as_ptr();
8102        let wv = vdupq_n_f32(w);
8103        let mut j = 0usize;
8104        while j + 16 <= n {
8105            let rb = vld1q_s8(rp.add(j));
8106            let lo = vmovl_s8(vget_low_s8(rb));
8107            let hi = vmovl_s8(vget_high_s8(rb));
8108            for (off, half) in [(0, lo), (8, hi)] {
8109                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
8110                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
8111                let o = j + off;
8112                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
8113                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
8114            }
8115            j += 16;
8116        }
8117        while j < n {
8118            *ap.add(j) += w * (*rp.add(j)) as f32;
8119            j += 1;
8120        }
8121    }
8122}
8123
8124/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
8125/// ≈9× scalar), scalar elsewhere.
8126#[inline]
8127pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
8128    #[cfg(target_arch = "aarch64")]
8129    unsafe {
8130        return dot_i8_f32_neon(w, x);
8131    }
8132    #[cfg(target_arch = "x86_64")]
8133    if avx2_enabled() {
8134        return unsafe { dot_i8_f32_avx2(w, x) };
8135    }
8136    #[allow(unreachable_code)]
8137    {
8138        let mut sum = 0.0f32;
8139        for (j, &b) in w.iter().enumerate() {
8140            sum += (b as i8) as f32 * x[j];
8141        }
8142        sum
8143    }
8144}
8145
8146/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
8147/// folded into the product (no prescaled copy of x). NEON on aarch64,
8148/// scalar elsewhere. Used by the active-neuron path `row_dot`.
8149#[inline]
8150fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8151    #[cfg(target_arch = "aarch64")]
8152    unsafe {
8153        return dot_i8_col_f32_neon(w, x, col);
8154    }
8155    #[allow(unreachable_code)]
8156    {
8157        let mut sum = 0.0f32;
8158        for (j, &b) in w.iter().enumerate() {
8159            sum += (b as i8) as f32 * x[j] * col[j];
8160        }
8161        sum
8162    }
8163}
8164
8165#[cfg(target_arch = "aarch64")]
8166#[target_feature(enable = "neon")]
8167unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8168    // SAFETY: callers uphold slice-length contracts (see call sites).
8169    unsafe {
8170        use core::arch::aarch64::*;
8171        let n = x.len();
8172        let wp = w.as_ptr() as *const i8;
8173        let xp = x.as_ptr();
8174        let cp = col.as_ptr();
8175        let (mut a0, mut a1, mut a2, mut a3) = (
8176            vdupq_n_f32(0.0),
8177            vdupq_n_f32(0.0),
8178            vdupq_n_f32(0.0),
8179            vdupq_n_f32(0.0),
8180        );
8181        let mut j = 0usize;
8182        while j + 16 <= n {
8183            let wb = vld1q_s8(wp.add(j));
8184            let lo = vmovl_s8(vget_low_s8(wb));
8185            let hi = vmovl_s8(vget_high_s8(wb));
8186            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8187            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8188            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8189            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8190            a0 = vfmaq_f32(
8191                a0,
8192                w0,
8193                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
8194            );
8195            a1 = vfmaq_f32(
8196                a1,
8197                w1,
8198                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
8199            );
8200            a2 = vfmaq_f32(
8201                a2,
8202                w2,
8203                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
8204            );
8205            a3 = vfmaq_f32(
8206                a3,
8207                w3,
8208                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
8209            );
8210            j += 16;
8211        }
8212        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8213        while j < n {
8214            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
8215            j += 1;
8216        }
8217        sum
8218    }
8219}
8220
8221#[cfg(target_arch = "aarch64")]
8222#[target_feature(enable = "neon")]
8223unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
8224    // SAFETY: callers uphold slice-length contracts (see call sites).
8225    unsafe {
8226        use core::arch::aarch64::*;
8227        let n = x.len();
8228        let wp = w.as_ptr() as *const i8;
8229        let xp = x.as_ptr();
8230        let (mut a0, mut a1, mut a2, mut a3) = (
8231            vdupq_n_f32(0.0),
8232            vdupq_n_f32(0.0),
8233            vdupq_n_f32(0.0),
8234            vdupq_n_f32(0.0),
8235        );
8236        let mut j = 0usize;
8237        while j + 16 <= n {
8238            let wb = vld1q_s8(wp.add(j));
8239            let lo = vmovl_s8(vget_low_s8(wb));
8240            let hi = vmovl_s8(vget_high_s8(wb));
8241            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8242            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8243            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8244            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8245            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
8246            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
8247            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
8248            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
8249            j += 16;
8250        }
8251        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8252        while j < n {
8253            sum += (*wp.add(j)) as f32 * *xp.add(j);
8254            j += 1;
8255        }
8256        sum
8257    }
8258}
8259
8260#[allow(clippy::too_many_arguments)]
8261fn qmatvec(
8262    q: &[u8],
8263    rep: &[u8],
8264    row_scale: &[f32],
8265    x: &[f32],
8266    col_field: &[f32],
8267    dtype: TensorDtype,
8268    rows: usize,
8269    cols: usize,
8270    out: &mut [f32],
8271    pool: Option<&Pool>,
8272) {
8273    debug_assert_eq!(out.len(), rows);
8274    #[cfg(not(target_arch = "aarch64"))]
8275    let _ = rep;
8276
8277    #[cfg(target_arch = "aarch64")]
8278    if sdot_enabled() {
8279        let act = if dtype == TensorDtype::Q8_2f {
8280            split_act_q8_2f(x, col_field)
8281        } else {
8282            split_act(x)
8283        };
8284        let out_addr = SendMut(out.as_mut_ptr());
8285        let run_range = |start: usize, end: usize| {
8286            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
8287        };
8288        match pool {
8289            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8290            _ => run_range(0, rows),
8291        }
8292        return;
8293    }
8294    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
8295    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
8296    #[cfg(target_arch = "x86_64")]
8297    if avx2_a8w8_enabled() {
8298        let act = if dtype == TensorDtype::Q8_2f {
8299            split_act_q8_2f(x, col_field)
8300        } else {
8301            split_act(x)
8302        };
8303        let out_addr = SendMut(out.as_mut_ptr());
8304        let run_range = |start: usize, end: usize| {
8305            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
8306        };
8307        match pool {
8308            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8309            _ => run_range(0, rows),
8310        }
8311        return;
8312    }
8313
8314    prescale_with(x, col_field, dtype, 1, |xs| {
8315        let out_addr = SendMut(out.as_mut_ptr());
8316        let run_range = move |start: usize, end: usize| {
8317            for o in start..end {
8318                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8319                // SAFETY: disjoint row ranges per worker.
8320                unsafe { *out_addr.at(o) = v };
8321            }
8322        };
8323        match pool {
8324            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8325            _ => run_range(0, rows),
8326        }
8327    });
8328}
8329
8330#[allow(clippy::too_many_arguments)]
8331fn qmatvec2(
8332    q: &[u8],
8333    row_scale: &[f32],
8334    x1: &[f32],
8335    x2: &[f32],
8336    col_field: &[f32],
8337    dtype: TensorDtype,
8338    rows: usize,
8339    cols: usize,
8340    o1: &mut [f32],
8341    o2: &mut [f32],
8342    pool: Option<&Pool>,
8343) {
8344    #[cfg(target_arch = "aarch64")]
8345    if sdot_enabled() {
8346        let a1s = if dtype == TensorDtype::Q8_2f {
8347            split_act_q8_2f(x1, col_field)
8348        } else {
8349            split_act(x1)
8350        };
8351        let a2s = if dtype == TensorDtype::Q8_2f {
8352            split_act_q8_2f(x2, col_field)
8353        } else {
8354            split_act(x2)
8355        };
8356        let p1 = SendMut(o1.as_mut_ptr());
8357        let p2 = SendMut(o2.as_mut_ptr());
8358        let run_range = |start: usize, end: usize| {
8359            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8360        };
8361        match pool {
8362            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8363            _ => run_range(0, rows),
8364        }
8365        return;
8366    }
8367    #[cfg(target_arch = "x86_64")]
8368    if avx2_a8w8_enabled() {
8369        let a1s = if dtype == TensorDtype::Q8_2f {
8370            split_act_q8_2f(x1, col_field)
8371        } else {
8372            split_act(x1)
8373        };
8374        let a2s = if dtype == TensorDtype::Q8_2f {
8375            split_act_q8_2f(x2, col_field)
8376        } else {
8377            split_act(x2)
8378        };
8379        let p1 = SendMut(o1.as_mut_ptr());
8380        let p2 = SendMut(o2.as_mut_ptr());
8381        let run_range = |start: usize, end: usize| {
8382            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8383        };
8384        match pool {
8385            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8386            _ => run_range(0, rows),
8387        }
8388        return;
8389    }
8390
8391    prescale_with(x1, col_field, dtype, 1, |x1s| {
8392        prescale_with(x2, col_field, dtype, 2, |x2s| {
8393            let p1 = SendMut(o1.as_mut_ptr());
8394            let p2 = SendMut(o2.as_mut_ptr());
8395            let run_range = move |start: usize, end: usize| {
8396                for o in start..end {
8397                    let row = &q[o * cols..(o + 1) * cols];
8398                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
8399                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
8400                    // SAFETY: disjoint row ranges per worker.
8401                    unsafe {
8402                        *p1.at(o) = s1;
8403                        *p2.at(o) = s2;
8404                    }
8405                }
8406            };
8407            match pool {
8408                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8409                _ => run_range(0, rows),
8410            }
8411        });
8412    });
8413}
8414
8415#[derive(Clone, Copy)]
8416struct SendMut(*mut f32);
8417unsafe impl Send for SendMut {}
8418unsafe impl Sync for SendMut {}
8419
8420impl SendMut {
8421    #[inline]
8422    fn at(self, i: usize) -> *mut f32 {
8423        unsafe { self.0.add(i) }
8424    }
8425}
8426
8427#[cfg(test)]
8428mod tests {
8429    use super::*;
8430
8431    #[test]
8432    fn f32_matvec_matches_matvec_rows_bitexact() {
8433        let (rows, cols) = (300, 40);
8434        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
8435        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
8436        let qt = QTensor::from_f32(w.clone(), rows, cols);
8437
8438        let mut a = vec![0.0f32; rows];
8439        matvec_rows(None, &w, &x, &mut a);
8440        let mut b = vec![0.0f32; rows];
8441        qt.matvec(&x, &mut b, None);
8442        assert_eq!(a, b);
8443    }
8444
8445    #[test]
8446    fn sdot_kernel_exact_on_grid() {
8447        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
8448        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
8449        // exact f32 dot to float rounding. This isolates kernel
8450        // correctness from quantization noise.
8451        eprintln!("sdot_enabled = {}", sdot_enabled());
8452        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
8453        let w: Vec<u8> = (0..rows * cols)
8454            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8455            .collect();
8456        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
8457        let x: Vec<f32> = (0..cols)
8458            .map(|i| match i % 3 {
8459                0 => 1.0,
8460                1 => -1.0,
8461                _ => 0.0,
8462            })
8463            .collect();
8464        let mut a = vec![0.0f32; rows];
8465        qmatvec(
8466            &w,
8467            &[],
8468            &scales,
8469            &x,
8470            &[],
8471            TensorDtype::Q8Row,
8472            rows,
8473            cols,
8474            &mut a,
8475            None,
8476        );
8477        for o in 0..rows {
8478            let mut acc = 0.0f32;
8479            for j in 0..cols {
8480                acc += (w[o * cols + j] as i8) as f32 * x[j];
8481            }
8482            let expect = acc * scales[o];
8483            assert!(
8484                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8485                "row {o}: {} vs {expect}",
8486                a[o]
8487            );
8488        }
8489    }
8490
8491    #[test]
8492    fn q1_tbl_fast_path_matches_reference() {
8493        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
8494        // row's final 4-tile window trips the 4B-overread guard (the
8495        // payload ends exactly at the last tile) — both paths must
8496        // agree with the dequant reference.
8497        let (rows, cols) = (5, 256);
8498        let gpr = cols / GROUP_SIZE;
8499        let mut bytes = Vec::new();
8500        for t in 0..rows * gpr {
8501            let s = 0.007 + (t % 11) as f32 * 0.004;
8502            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8503            for j in 0..4 {
8504                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
8505            }
8506        }
8507        let x: Vec<f32> = (0..cols)
8508            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
8509            .collect();
8510        let mut w = vec![0.0f32; rows * cols];
8511        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8512        let mut got = vec![0.0f32; rows];
8513        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8514        for o in 0..rows {
8515            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8516            assert!(
8517                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8518                "row {o}: {} vs {expect}",
8519                got[o]
8520            );
8521        }
8522        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
8523        // single-matvec path bit-for-bit.
8524        let b = 5usize;
8525        let mut xs_all = Vec::new();
8526        for bi in 0..b {
8527            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
8528        }
8529        let mut mm = vec![0.0f32; b * rows];
8530        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
8531        for bi in 0..b {
8532            let mut single = vec![0.0f32; rows];
8533            q1_matvec(
8534                &bytes,
8535                &xs_all[bi * cols..(bi + 1) * cols],
8536                rows,
8537                cols,
8538                &mut single,
8539                None,
8540            );
8541            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
8542        }
8543    }
8544
8545    #[test]
8546    fn q1_kernels_match_exact_reference() {
8547        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
8548        let (rows, cols) = (7, 96);
8549        let gpr = cols / GROUP_SIZE;
8550        let mut bytes = Vec::new();
8551        for t in 0..rows * gpr {
8552            let s = 0.01 + (t % 13) as f32 * 0.003;
8553            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8554            for j in 0..4 {
8555                bytes.push(((t * 31 + j * 97) % 251) as u8);
8556            }
8557        }
8558        // On-grid activations (±1, amax 1) → the SDOT path is exact.
8559        let x: Vec<f32> = (0..cols)
8560            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
8561            .collect();
8562        // Reference through the core dequant.
8563        let mut w = vec![0.0f32; rows * cols];
8564        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8565        let mut expect = vec![0.0f32; rows];
8566        for o in 0..rows {
8567            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8568        }
8569        let mut got = vec![0.0f32; rows];
8570        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8571        for o in 0..rows {
8572            assert!(
8573                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
8574                "row {o}: {} vs {}",
8575                got[o],
8576                expect[o]
8577            );
8578        }
8579        // Pair and batch paths agree with the single path.
8580        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
8581        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
8582        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
8583        assert_eq!(a1, got);
8584        let mut xs = x.clone();
8585        xs.extend_from_slice(&x2);
8586        let mut mm = vec![0.0f32; 2 * rows];
8587        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
8588        assert_eq!(&mm[..rows], got.as_slice());
8589        assert_eq!(&mm[rows..], a2.as_slice());
8590    }
8591
8592    #[test]
8593    fn repack_is_bit_identical() {
8594        // The interleaved-repack kernel must produce EXACTLY the same
8595        // bits as the mmap-layout kernel: integer accumulation is order-
8596        // exact, the f32 epilogue is identical. Odd rows exercise the
8597        // tail; direct range calls exercise unaligned pool splits.
8598        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
8599        let w: Vec<u8> = (0..rows * cols)
8600            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
8601            .collect();
8602        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
8603        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
8604        let rep = q8_repack_layout(&w, rows, cols);
8605        // Group interleave round-trips.
8606        for g in 0..rows / 4 {
8607            for c in 0..cols / 16 {
8608                for lane in 0..4 {
8609                    assert_eq!(
8610                        &rep[g * 4 * cols + c * 64 + lane * 16
8611                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
8612                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
8613                    );
8614                }
8615            }
8616        }
8617        let mut a = vec![0.0f32; rows];
8618        qmatvec(
8619            &w,
8620            &[],
8621            &scales,
8622            &x,
8623            &[],
8624            TensorDtype::Q8Row,
8625            rows,
8626            cols,
8627            &mut a,
8628            None,
8629        );
8630        let mut b = vec![0.0f32; rows];
8631        qmatvec(
8632            &w,
8633            &rep,
8634            &scales,
8635            &x,
8636            &[],
8637            TensorDtype::Q8Row,
8638            rows,
8639            cols,
8640            &mut b,
8641            None,
8642        );
8643        assert_eq!(a, b, "full-range repack output diverged");
8644
8645        #[cfg(target_arch = "aarch64")]
8646        if sdot_enabled() {
8647            // Unaligned range split (pool workers get arbitrary bounds).
8648            let act = split_act(&x);
8649            let mut c1 = vec![0.0f32; rows];
8650            let mut c2 = vec![0.0f32; rows];
8651            q8_range_sdot(
8652                &w,
8653                &[],
8654                &scales,
8655                &act,
8656                cols,
8657                SendMut(c1.as_mut_ptr()),
8658                3,
8659                rows - 2,
8660            );
8661            q8_range_sdot(
8662                &w,
8663                &rep,
8664                &scales,
8665                &act,
8666                cols,
8667                SendMut(c2.as_mut_ptr()),
8668                3,
8669                rows - 2,
8670            );
8671            assert_eq!(c1, c2, "unaligned-range repack output diverged");
8672        }
8673    }
8674
8675    #[test]
8676    fn sdot_a8w8_noise_is_bounded() {
8677        // Off-grid activations: A8 quantization noise must stay small in
8678        // relative L2 over the whole output (realistic accuracy contract;
8679        // vmfcore measured argmax-identical decode on real models).
8680        let (rows, cols) = (16, 512);
8681        let w: Vec<u8> = (0..rows * cols)
8682            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8683            .collect();
8684        let scales = vec![0.01f32; rows];
8685        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
8686        let mut a = vec![0.0f32; rows];
8687        qmatvec(
8688            &w,
8689            &[],
8690            &scales,
8691            &x,
8692            &[],
8693            TensorDtype::Q8Row,
8694            rows,
8695            cols,
8696            &mut a,
8697            None,
8698        );
8699        let (mut num, mut den) = (0f64, 0f64);
8700        for o in 0..rows {
8701            let mut acc = 0.0f32;
8702            for j in 0..cols {
8703                acc += (w[o * cols + j] as i8) as f32 * x[j];
8704            }
8705            let expect = acc * scales[o];
8706            num += ((a[o] - expect) as f64).powi(2);
8707            den += (expect as f64).powi(2);
8708        }
8709        let rel = (num / den.max(1e-12)).sqrt();
8710        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
8711    }
8712
8713    #[test]
8714    fn i8_dot_neon_matches_scalar() {
8715        let n = 100;
8716        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
8717        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
8718        let mut scalar = 0.0f32;
8719        for j in 0..n {
8720            scalar += (w[j] as i8) as f32 * x[j];
8721        }
8722        let fast = dot_i8_f32(&w, &x);
8723        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
8724    }
8725
8726    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
8727    #[test]
8728    fn vbitmatvec_matches_full_dequant() {
8729        let (rows, cols) = (6, 64);
8730        let ng = cols / GROUP_SIZE;
8731        // Hand-craft: bits per row, f16 scales, packed rows.
8732        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8733        let mut bytes = bits.clone();
8734        for g in 0..rows * ng {
8735            let s = 0.02 + 0.001 * g as f32;
8736            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8737        }
8738        for r in 0..rows {
8739            let b = bits[r] as usize;
8740            let (mut acc, mut nb) = (0u64, 0usize);
8741            let mut rowbytes = Vec::new();
8742            for i in 0..cols {
8743                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8744                acc = (acc << b) | v;
8745                nb += b;
8746                while nb >= 8 {
8747                    nb -= 8;
8748                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
8749                }
8750            }
8751            if nb > 0 {
8752                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8753            }
8754            bytes.extend_from_slice(&rowbytes);
8755        }
8756        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8757
8758        let mut reference = vec![0f32; rows * cols];
8759        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
8760        let mut expect = vec![0f32; rows];
8761        for r in 0..rows {
8762            expect[r] = reference[r * cols..(r + 1) * cols]
8763                .iter()
8764                .zip(&x)
8765                .map(|(w, xv)| w * xv)
8766                .sum();
8767        }
8768        let mut got = vec![0f32; rows];
8769        let offsets = vbit_row_offsets(&bytes, rows, cols);
8770        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
8771        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
8772        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
8773        // the golden-parity gate).
8774        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
8775        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
8776        for r in 0..rows {
8777            assert!(
8778                (got[r] - expect[r]).abs() < tol * scale,
8779                "row {r}: {} vs {}",
8780                got[r],
8781                expect[r]
8782            );
8783        }
8784    }
8785
8786    /// Fused q4 matvec must match the reference full-dequant + dense
8787    /// matvec bit-for-bit in structure (same f32 math, group order).
8788    /// vbit matmat: the blocked 1×4 leg must match the per-row path
8789    /// (paired env toggle; larger shape so both code paths engage).
8790    #[test]
8791    #[cfg(target_arch = "x86_64")]
8792    fn vbit_matmat_blocked_matches_per_row() {
8793        let (rows, cols, b) = (64usize, 128usize, 9usize);
8794        let ng = cols / GROUP_SIZE;
8795        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
8796        let mut bytes = bits.clone();
8797        for g in 0..rows * ng {
8798            let sc = 0.02 + 0.0005 * g as f32;
8799            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8800        }
8801        for r in 0..rows {
8802            let bw = bits[r] as usize;
8803            let (mut acc, mut nb) = (0u64, 0usize);
8804            let mut rowbytes = Vec::new();
8805            for i in 0..cols {
8806                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
8807                acc = (acc << bw) | v;
8808                nb += bw;
8809                while nb >= 8 {
8810                    nb -= 8;
8811                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
8812                }
8813            }
8814            if nb > 0 {
8815                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8816            }
8817            bytes.extend_from_slice(&rowbytes);
8818        }
8819        let x: Vec<f32> = (0..b * cols)
8820            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8821            .collect();
8822        let offsets = vbit_row_offsets(&bytes, rows, cols);
8823        let mut y_a = vec![0f32; b * rows];
8824        let mut y_b = vec![0f32; b * rows];
8825        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8826        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
8827        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8828        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
8829        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8830        let max_d = y_a
8831            .iter()
8832            .zip(&y_b)
8833            .map(|(p, q)| (p - q).abs())
8834            .fold(0.0f32, f32::max);
8835        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
8836    }
8837
8838    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
8839    /// per-row path exactly: same nibble unpack, same group order,
8840    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
8841    /// two full 1×4 blocks plus a remainder through the single-row
8842    /// kernel. (Both paths produce identical output, so the shared
8843    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
8844    /// the verdict — worst case both sides take the same path.)
8845    #[test]
8846    fn q4t_matmat_blocked_matches_per_row() {
8847        let (rows, cols, b) = (16usize, 64usize, 9usize);
8848        let gpr = cols / GROUP_SIZE;
8849        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
8850        for r in 0..rows {
8851            for g in 0..gpr {
8852                let t = (r * gpr + g) * Q4_TILE;
8853                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
8854                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8855                for k in 0..16 {
8856                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8857                }
8858            }
8859        }
8860        let x: Vec<f32> = (0..b * cols)
8861            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8862            .collect();
8863        let mut y_blk = vec![0f32; b * rows];
8864        let mut y_row = vec![0f32; b * rows];
8865        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8866        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
8867        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8868        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
8869        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8870        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
8871    }
8872
8873    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
8874    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
8875    /// order differs — tight tolerance.
8876    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
8877    /// span varies row to row, so the codes actually exercise the full 0..31
8878    /// range rather than clustering on one rung.
8879    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
8880        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
8881        let gpr = cols / GROUP_SIZE;
8882        let stride = q4tp_code_stride(gpr);
8883        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
8884        let mut b = vec![0u8; codes_off + rows * stride];
8885        for r in 0..rows {
8886            for g in 0..gpr {
8887                let t = (r * gpr + g) * Q4TP_NIB;
8888                for k in 0..16 {
8889                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8890                }
8891            }
8892            let lo = -6.0 - 0.03 * (r % 17) as f32;
8893            let step = 0.01 + 0.004 * (r % 11) as f32;
8894            let p = params_off + r * 4;
8895            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
8896            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
8897            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
8898            for g in 0..gpr {
8899                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
8900            }
8901        }
8902        b
8903    }
8904
8905    /// The same weights re-expressed as q4_tiled, so the proven kernel can
8906    /// be the reference: each tile stores the ladder scale its code selects.
8907    /// Only the f16 rounding of that scale separates the two payloads.
8908    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
8909        let gpr = cols / GROUP_SIZE;
8910        let v = Q4tpView::new(bytes, rows, cols);
8911        let mut out = vec![0u8; rows * gpr * Q4_TILE];
8912        let mut sc = vec![0f32; gpr];
8913        for r in 0..rows {
8914            v.scales_into(r, gpr, &mut sc);
8915            for g in 0..gpr {
8916                let t = (r * gpr + g) * Q4_TILE;
8917                let s = sc[g];
8918                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8919                let src = (r * gpr + g) * Q4TP_NIB;
8920                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
8921            }
8922        }
8923        out
8924    }
8925
8926    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
8927    /// rounding — that scalar routine is the format's definition, and the
8928    /// kernels re-derive the scale from the ladder independently. Call the
8929    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
8930    /// so routing through it would test the other path by accident.
8931    #[test]
8932    fn q4tp_exact_path_matches_dequant_reference() {
8933        let (rows, cols) = (256usize, 512usize);
8934        let gpr = cols / GROUP_SIZE;
8935        let bytes = synth_q4tp(rows, cols);
8936        let mut w = vec![0f32; rows * cols];
8937        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
8938
8939        let x: Vec<f32> = (0..cols)
8940            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8941            .collect();
8942        let v = Q4tpView::new(&bytes, rows, cols);
8943        let mut sc = vec![0f32; gpr];
8944        for r in 0..rows {
8945            v.scales_into(r, gpr, &mut sc);
8946            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
8947            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
8948            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
8949            // the meaningful yardstick is the summed magnitude, not the result:
8950            // against the result any reordering of a 512-term f32 sum "fails".
8951            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
8952            assert!(
8953                (got - want).abs() <= 1e-5 * mag,
8954                "row {r}: kernel {got} vs dequant {want}"
8955            );
8956        }
8957    }
8958
8959    /// The int8 (a8w8) path can't be checked against an f32 reference — the
8960    /// activation quantization dominates. Check it against the q4t kernel it
8961    /// was ported from instead, on payloads holding the same weights: that
8962    /// isolates exactly what the port could break (16 B stride, ladder
8963    /// lookup, nibble unpack) from what it deliberately shares.
8964    #[test]
8965    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
8966        let (rows, cols) = (256usize, 512usize);
8967        let bytes = synth_q4tp(rows, cols);
8968        let twin = q4tp_as_q4t(&bytes, rows, cols);
8969        let x: Vec<f32> = (0..cols)
8970            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8971            .collect();
8972
8973        let mut got = vec![0f32; rows];
8974        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
8975        let mut want = vec![0f32; rows];
8976        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
8977
8978        // Scale is f16 in the twin and f32 here, so allow that rounding on
8979        // top of the summed magnitude (same cancellation argument as above).
8980        let mut w = vec![0f32; rows * cols];
8981        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
8982        for r in 0..rows {
8983            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
8984            assert!(
8985                (got[r] - want[r]).abs() <= 1e-3 * mag,
8986                "row {r}: q4tp {} vs q4t {}",
8987                got[r],
8988                want[r]
8989            );
8990        }
8991    }
8992
8993    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
8994    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
8995    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
8996    /// code and its four accumulators are exactly what tends to go wrong.
8997    #[test]
8998    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
8999        let (rows, cols, b) = (256usize, 512usize, 5usize);
9000        let bytes = synth_q4tp(rows, cols);
9001        let twin = q4tp_as_q4t(&bytes, rows, cols);
9002        let xs: Vec<f32> = (0..b * cols)
9003            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9004            .collect();
9005
9006        let mut got = vec![0f32; b * rows];
9007        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
9008        let mut want = vec![0f32; b * rows];
9009        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
9010
9011        let mut w = vec![0f32; rows * cols];
9012        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9013        for t in 0..b {
9014            for r in 0..rows {
9015                let mag: f32 = (0..cols)
9016                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
9017                    .sum();
9018                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
9019                assert!(
9020                    (g - wa).abs() <= 1e-3 * mag,
9021                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
9022                );
9023            }
9024        }
9025    }
9026
9027    #[test]
9028    fn q4tp_matvec2_matches_the_single_stream_kernel() {
9029        let (rows, cols) = (128usize, 256usize);
9030        let gpr = cols / GROUP_SIZE;
9031        let bytes = synth_q4tp(rows, cols);
9032        let xs: Vec<f32> = (0..2 * cols)
9033            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9034            .collect();
9035
9036        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
9037        q4tp_matvec2(
9038            &bytes,
9039            &xs[..cols],
9040            &xs[cols..],
9041            rows,
9042            cols,
9043            &mut o1,
9044            &mut o2,
9045            None,
9046        );
9047
9048        // matvec2 takes the exact path for both streams, so the single-row
9049        // kernel is an exact reference — no tolerance for path differences.
9050        let v = Q4tpView::new(&bytes, rows, cols);
9051        let mut sc = vec![0f32; gpr];
9052        for r in 0..rows {
9053            v.scales_into(r, gpr, &mut sc);
9054            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
9055            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
9056        }
9057    }
9058
9059    /// q4tp must not COST speed — it exists to save bytes, and a format that
9060    /// trades 7% of a file for a slower model is a bad trade. This guard is
9061    /// here because correctness tests happily passed while `q4tp_matmat` was
9062    /// missing its int8 and Accelerate arms and the model ran 5x slower.
9063    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
9064    /// aligned than q4t's 18 B, which pays for the scale indirection).
9065    #[test]
9066    fn q4tp_matvec_keeps_pace_with_q4t() {
9067        let (rows, cols) = (4096usize, 3072usize);
9068        let bytes = synth_q4tp(rows, cols);
9069        let twin = q4tp_as_q4t(&bytes, rows, cols);
9070        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
9071        let mut o = vec![0f32; rows];
9072        let n = 12;
9073        let mut best = (f64::MAX, f64::MAX);
9074        // Interleaved A/B, minimum statistic: this machine throttles, and a
9075        // mean over a thermal ramp reliably indicts whichever ran second.
9076        for _ in 0..3 {
9077            let t0 = std::time::Instant::now();
9078            for _ in 0..n {
9079                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
9080            }
9081            best.0 = best.0.min(t0.elapsed().as_secs_f64());
9082            let t0 = std::time::Instant::now();
9083            for _ in 0..n {
9084                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
9085            }
9086            best.1 = best.1.min(t0.elapsed().as_secs_f64());
9087        }
9088        let ratio = best.1 / best.0;
9089        println!("q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x", best.0 * 1e3 / n as f64, best.1 * 1e3 / n as f64);
9090        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
9091    }
9092
9093    #[cfg(target_os = "macos")]
9094    #[test]
9095    fn q4t_matmat_accel_matches_dequant_reference() {
9096        if !accel_gemm_enabled() {
9097            return; // CMF_ACCEL=0
9098        }
9099        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
9100        let gpr = cols / GROUP_SIZE;
9101        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9102        for r in 0..rows {
9103            for g in 0..gpr {
9104                let t = (r * gpr + g) * Q4_TILE;
9105                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
9106                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9107                for k in 0..16 {
9108                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9109                }
9110            }
9111        }
9112        let x: Vec<f32> = (0..b * cols)
9113            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9114            .collect();
9115        let mut got = vec![0f32; b * rows];
9116        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
9117        // Brute-force reference off the same tiles.
9118        let mut w = vec![0f32; rows * cols];
9119        for r in 0..rows {
9120            for g in 0..gpr {
9121                let t = (r * gpr + g) * Q4_TILE;
9122                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
9123                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
9124                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
9125                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
9126                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
9127                }
9128            }
9129        }
9130        for bi in 0..b {
9131            for r in 0..rows {
9132                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
9133                let d = (got[bi * rows + r] - want).abs();
9134                assert!(
9135                    d <= want.abs().max(1.0) * 1e-4,
9136                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
9137                    got[bi * rows + r]
9138                );
9139            }
9140        }
9141    }
9142
9143    #[test]
9144    fn q4matvec_matches_full_dequant() {
9145        let (rows, cols) = (8, 64);
9146        let groups = rows * cols / GROUP_SIZE;
9147        // Hand-craft a q4_block blob: nibbles then f16 scales.
9148        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9149        for i in 0..groups * 16 {
9150            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9151        }
9152        for g in 0..groups {
9153            let s = 0.01 + 0.003 * g as f32;
9154            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9155        }
9156        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9157
9158        let mut reference = vec![0.0f32; rows * cols];
9159        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9160        let mut expect = vec![0.0f32; rows];
9161        for r in 0..rows {
9162            expect[r] = reference[r * cols..(r + 1) * cols]
9163                .iter()
9164                .zip(&x)
9165                .map(|(w, xv)| w * xv)
9166                .sum();
9167        }
9168
9169        let mut got = vec![0.0f32; rows];
9170        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9171        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9172        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
9173        // in the golden-parity gate).
9174        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9175        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9176        for r in 0..rows {
9177            assert!(
9178                (got[r] - expect[r]).abs() < tol * scale,
9179                "row {r}: {} vs {}",
9180                got[r],
9181                expect[r]
9182            );
9183        }
9184    }
9185
9186    /// Fused two-input vbit matvec must equal two single matvecs exactly
9187    /// (same per-lane accumulation order on both scalar and SDOT paths).
9188    #[test]
9189    fn vbitmatvec2_equals_two_singles() {
9190        let (rows, cols) = (6, 64);
9191        let ng = cols / GROUP_SIZE;
9192        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9193        let mut bytes = bits.clone();
9194        for g in 0..rows * ng {
9195            let s = 0.02 + 0.001 * g as f32;
9196            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9197        }
9198        for r in 0..rows {
9199            let b = bits[r] as usize;
9200            let (mut acc, mut nb) = (0u64, 0usize);
9201            let mut rowbytes = Vec::new();
9202            for i in 0..cols {
9203                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9204                acc = (acc << b) | v;
9205                nb += b;
9206                while nb >= 8 {
9207                    nb -= 8;
9208                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9209                }
9210            }
9211            if nb > 0 {
9212                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9213            }
9214            bytes.extend_from_slice(&rowbytes);
9215        }
9216        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9217        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
9218        let offsets = vbit_row_offsets(&bytes, rows, cols);
9219
9220        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9221        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
9222        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
9223        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9224        vbitmatvec2(
9225            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
9226        );
9227        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
9228        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
9229    }
9230
9231    /// Fused two-input q4 matvec must equal two single matvecs exactly.
9232    #[test]
9233    fn q4matvec2_equals_two_singles() {
9234        let (rows, cols) = (8, 128);
9235        let groups = rows * cols / GROUP_SIZE;
9236        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9237        for i in 0..groups * 16 {
9238            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9239        }
9240        for g in 0..groups {
9241            let s = 0.01 + 0.003 * g as f32;
9242            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9243        }
9244        // Include an outlier channel so the SDOT correction path is
9245        // exercised in the pair kernel too.
9246        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9247        x1[9] = 250.0;
9248        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9249
9250        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9251        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
9252        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
9253        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9254        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
9255        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
9256        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
9257    }
9258
9259    /// Multi-matrix job must equal separate matvecs exactly — same
9260    /// kernels, only the dispatch is fused.
9261    #[test]
9262    fn matvec_many_equals_separate_matvecs() {
9263        use crate::pool::Pool;
9264        let (r1, r2, cols) = (300, 200, 64);
9265        let mk = |salt: usize, rows: usize| {
9266            QTensor::from_f32(
9267                (0..rows * cols)
9268                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
9269                    .collect(),
9270                rows,
9271                cols,
9272            )
9273        };
9274        let (a, b) = (mk(1, r1), mk(5, r2));
9275        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
9276        let pool = Pool::new(3);
9277
9278        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
9279        a.matvec(&x, &mut ea, Some(&pool));
9280        b.matvec(&x, &mut eb, Some(&pool));
9281        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
9282        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
9283        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
9284        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
9285    }
9286
9287    /// Batched q4/vbit matmat must equal per-position matvec calls
9288    /// exactly (the fallback it replaced) — same kernels, same order.
9289    #[test]
9290    fn batched_matmat_equals_per_position_matvec() {
9291        let (rows, cols, b) = (8, 64, 5);
9292        // q4 blob.
9293        let groups = rows * cols / GROUP_SIZE;
9294        let mut q4 = Vec::new();
9295        for i in 0..groups * 16 {
9296            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9297        }
9298        for g in 0..groups {
9299            q4.extend_from_slice(
9300                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9301            );
9302        }
9303        // vbit blob (mixed widths incl. 8).
9304        let ng = cols / GROUP_SIZE;
9305        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
9306        let mut vb = bits.clone();
9307        for g in 0..rows * ng {
9308            vb.extend_from_slice(
9309                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
9310            );
9311        }
9312        for r in 0..rows {
9313            let bw = bits[r] as usize;
9314            let (mut acc, mut nb) = (0u64, 0usize);
9315            let mut rowbytes = Vec::new();
9316            for i in 0..cols {
9317                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9318                acc = (acc << bw) | v;
9319                nb += bw;
9320                while nb >= 8 {
9321                    nb -= 8;
9322                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9323                }
9324            }
9325            if nb > 0 {
9326                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9327            }
9328            vb.extend_from_slice(&rowbytes);
9329        }
9330        let offsets = vbit_row_offsets(&vb, rows, cols);
9331
9332        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9333
9334        // q4: batch vs singles.
9335        let mut got = vec![0f32; b * rows];
9336        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
9337        for bi in 0..b {
9338            let mut expect = vec![0f32; rows];
9339            q4matvec(
9340                &q4,
9341                &xs[bi * cols..(bi + 1) * cols],
9342                rows,
9343                cols,
9344                &mut expect,
9345                None,
9346            );
9347            assert_eq!(
9348                &got[bi * rows..(bi + 1) * rows],
9349                &expect[..],
9350                "q4 batch pos {bi}"
9351            );
9352        }
9353
9354        // vbit: batch vs singles.
9355        let mut got = vec![0f32; b * rows];
9356        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
9357        for bi in 0..b {
9358            let mut expect = vec![0f32; rows];
9359            vbitmatvec(
9360                &vb,
9361                &offsets,
9362                &xs[bi * cols..(bi + 1) * cols],
9363                rows,
9364                cols,
9365                &mut expect,
9366                None,
9367            );
9368            assert_eq!(
9369                &got[bi * rows..(bi + 1) * rows],
9370                &expect[..],
9371                "vbit batch pos {bi}"
9372            );
9373        }
9374    }
9375
9376    /// q4_tiled kernels must produce BIT-identical outputs to the q4
9377    /// split kernels on the same values (same ints, same order — only
9378    /// the byte placement differs).
9379    #[test]
9380    fn q4_tiled_matches_q4_block_bitexact() {
9381        let (rows, cols, b) = (8usize, 128usize, 3usize);
9382        let groups = rows * cols / GROUP_SIZE;
9383        let mut split = Vec::with_capacity(groups * 18);
9384        for i in 0..groups * 16 {
9385            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9386        }
9387        for g in 0..groups {
9388            split.extend_from_slice(
9389                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9390            );
9391        }
9392        // Re-tile: [scale][nibbles] per group.
9393        let (packed, scales) = split.split_at(groups * 16);
9394        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
9395        for g in 0..groups {
9396            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
9397            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
9398        }
9399
9400        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9401        x1[9] = 250.0; // exercise the outlier path
9402        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9403
9404        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
9405        q4matvec(&split, &x1, rows, cols, &mut a, None);
9406        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
9407        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
9408
9409        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9410        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
9411        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
9412        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
9413        assert_eq!(a1, t1);
9414        assert_eq!(a2, t2);
9415
9416        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9417        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
9418        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
9419        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
9420        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
9421    }
9422
9423    /// q4 SDOT outlier correction: a single huge activation channel
9424    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
9425    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
9426    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
9427    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
9428    /// can never qualify (8² = n).
9429    #[test]
9430    fn q4matvec_sdot_outlier_exact() {
9431        let (rows, cols) = (4, 128);
9432        let groups = rows * cols / GROUP_SIZE;
9433        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9434        for i in 0..groups * 16 {
9435            bytes.push(((i * 11 + 5) % 256) as u8);
9436        }
9437        for g in 0..groups {
9438            let s = 0.02 + 0.002 * g as f32;
9439            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9440        }
9441        let mut x: Vec<f32> = (0..cols)
9442            .map(|i| match i % 3 {
9443                0 => 1.0,
9444                1 => -1.0,
9445                _ => 0.0,
9446            })
9447            .collect();
9448        x[17] = 300.0; // ≫ 8·rms → outlier channel
9449
9450        let mut reference = vec![0.0f32; rows * cols];
9451        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9452        let mut expect = vec![0.0f32; rows];
9453        for r in 0..rows {
9454            expect[r] = reference[r * cols..(r + 1) * cols]
9455                .iter()
9456                .zip(&x)
9457                .map(|(w, xv)| w * xv)
9458                .sum();
9459        }
9460        let mut got = vec![0.0f32; rows];
9461        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9462        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9463        for r in 0..rows {
9464            assert!(
9465                (got[r] - expect[r]).abs() < 2e-3 * scale,
9466                "row {r}: {} vs {} (outlier term must be exact)",
9467                got[r],
9468                expect[r]
9469            );
9470        }
9471    }
9472
9473    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
9474    /// including the ternary zero level and the binary-searched outlier
9475    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
9476    #[test]
9477    fn q1t_matvec_matches_reference() {
9478        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
9479        let (rows, cols) = (3usize, 64usize); // gpr = 2
9480        let gpr = cols / GROUP_SIZE;
9481        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
9482        // Overlay (must be sorted by flat index): a few spikes across rows.
9483        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
9484        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
9485        let mut bytes = Vec::new();
9486        for r in 0..rows {
9487            for g in 0..gpr {
9488                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
9489                let mut c = [0u8; 7];
9490                for k in 0..GROUP_SIZE {
9491                    // Encoder invariant: code 0 at outlier positions.
9492                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
9493                        0
9494                    } else {
9495                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
9496                    };
9497                    cortiq_core::quant::q1t_pack(&mut c, k, code);
9498                }
9499                bytes.extend_from_slice(&c);
9500            }
9501        }
9502        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
9503        // row (outliers are sorted by flat index → already grouped by row).
9504        let mut row_ptr = vec![0u32; rows + 1];
9505        for &(idx, _) in &outliers {
9506            row_ptr[idx as usize / cols + 1] += 1;
9507        }
9508        for r in 0..rows {
9509            row_ptr[r + 1] += row_ptr[r];
9510        }
9511        for &p in &row_ptr {
9512            bytes.extend_from_slice(&p.to_le_bytes());
9513        }
9514        for &(idx, v) in &outliers {
9515            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
9516            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
9517        }
9518
9519        let mut refw = vec![0f32; rows * cols];
9520        dequant_q1t(&bytes, rows, cols, &mut refw);
9521        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
9522        // x exactly and matches the f32 reference (same trick as the q1 test).
9523        let x: Vec<f32> = (0..cols)
9524            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9525            .collect();
9526        let mut expect = vec![0f32; rows];
9527        for r in 0..rows {
9528            let mut a = 0.0f32;
9529            for j in 0..cols {
9530                a += refw[r * cols + j] * x[j];
9531            }
9532            expect[r] = a;
9533        }
9534        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
9535        let mut got = vec![0f32; rows];
9536        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
9537        for r in 0..rows {
9538            assert!(
9539                (got[r] - expect[r]).abs() < tol(expect[r]),
9540                "row {r}: {} vs {}",
9541                got[r],
9542                expect[r]
9543            );
9544        }
9545        // matmat (b=2, f32 decode path) must agree too.
9546        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
9547        let mut gm = vec![0f32; 2 * rows];
9548        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
9549        for r in 0..rows {
9550            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
9551            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
9552        }
9553        // Fused pair (q1t_matvec2) must equal two single matvecs
9554        // bit-for-bit: same unpack, same group order, same f32
9555        // accumulation per stream. Distinct x2 exercises both lanes.
9556        let xb: Vec<f32> = (0..cols)
9557            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
9558            .collect();
9559        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9560        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
9561        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
9562        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9563        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
9564        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
9565        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
9566    }
9567
9568    /// Pair == 2×matvec with an ODD group count (the kernel's tail
9569    /// group) and no overlay section.
9570    #[test]
9571    fn q1t_matvec2_odd_gpr_matches_singles() {
9572        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
9573        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
9574        let gpr = cols / GROUP_SIZE;
9575        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
9576        for r in 0..rows {
9577            for g in 0..gpr {
9578                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
9579                let mut c = [0u8; 7];
9580                for k in 0..GROUP_SIZE {
9581                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
9582                }
9583                bytes.extend_from_slice(&c);
9584            }
9585        }
9586        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
9587        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
9588        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9589        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9590        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
9591        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9592        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9593        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
9594        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
9595    }
9596
9597    // Speed A/B: fused pair (one unpack, two streams) vs two single
9598    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
9599    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
9600    #[test]
9601    #[ignore]
9602    fn q1t_matvec2_speed() {
9603        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
9604        use std::time::Instant;
9605        let (rows, cols) = (8192usize, 4096usize);
9606        let gpr = cols / GROUP_SIZE;
9607        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
9608        for r in 0..rows {
9609            for g in 0..gpr {
9610                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
9611                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
9612                let mut c = [0u8; 7];
9613                for k in 0..GROUP_SIZE {
9614                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
9615                }
9616                bytes.extend_from_slice(&c);
9617            }
9618        }
9619        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
9620        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
9621        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9622        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9623        // Warm both paths once.
9624        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9625        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9626        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
9627        for _ in 0..8 {
9628            let t0 = Instant::now();
9629            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9630            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
9631            let t1 = Instant::now();
9632            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9633            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
9634            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
9635        }
9636        assert_eq!(p1, s1);
9637        assert_eq!(p2, s2);
9638        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
9639    }
9640
9641    // Speed A/B: the base-3-division decode (what the packing commit left in
9642    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
9643    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
9644    #[test]
9645    #[ignore]
9646    fn q1t_matvec_speed() {
9647        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
9648        use std::time::Instant;
9649        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
9650        let gpr = cols / GROUP_SIZE;
9651        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
9652        for r in 0..rows {
9653            for g in 0..gpr {
9654                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
9655                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
9656                let mut c = [0u8; 7];
9657                for k in 0..GROUP_SIZE {
9658                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
9659                }
9660                bytes.extend_from_slice(&c);
9661            }
9662        }
9663        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
9664        let mut row_ptr = vec![0u32; rows + 1];
9665        let mut idx = 0usize;
9666        while idx < n {
9667            row_ptr[idx / cols + 1] += 1;
9668            idx += stride;
9669        }
9670        for r in 0..rows {
9671            row_ptr[r + 1] += row_ptr[r];
9672        }
9673        for &p in &row_ptr {
9674            bytes.extend_from_slice(&p.to_le_bytes());
9675        }
9676        let mut idx = 0usize;
9677        while idx < n {
9678            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
9679            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
9680            idx += stride;
9681        }
9682        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
9683        // reference (the A/B is a timing check; values must still agree).
9684        let x: Vec<f32> = (0..cols)
9685            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9686            .collect();
9687        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
9688
9689        // "before": base-3 division decode into a buffer, then dot.
9690        let slow = |out: &mut [f32]| {
9691            let mut buf = vec![0f32; cols];
9692            for r in 0..rows {
9693                for g in 0..gpr {
9694                    let off = (r * gpr + g) * Q1T_TILE;
9695                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9696                    let codes = &bytes[off + 2..off + Q1T_TILE];
9697                    for k in 0..GROUP_SIZE {
9698                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
9699                            1 => s,
9700                            2 => -s,
9701                            _ => 0.0,
9702                        };
9703                    }
9704                }
9705                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
9706                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
9707            }
9708        };
9709        let iters = 5;
9710        let mut a = vec![0f32; rows];
9711        slow(&mut a); // warm
9712        let t = Instant::now();
9713        for _ in 0..iters {
9714            slow(&mut a);
9715        }
9716        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
9717
9718        let mut b = vec![0f32; rows];
9719        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
9720        let t = Instant::now();
9721        for _ in 0..iters {
9722            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
9723        }
9724        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
9725
9726        for r in 0..rows {
9727            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
9728        }
9729        println!(
9730            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
9731            slow_ms / fast_ms
9732        );
9733    }
9734}
cortiq_engine/qtensor.rs

cortiq_engine/
qtensor.rs
