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    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
406    /// its kernels by which of the two answers.
407    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
408        match self {
409            Self::Mapped {
410                model,
411                idx,
412                dtype: TensorDtype::Q4TiledP,
413                ..
414            } => Some((model, *idx)),
415            _ => None,
416        }
417    }
418
419    pub fn cols(&self) -> usize {
420        match self {
421            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
422        }
423    }
424
425    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
426    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
427    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
428        match self {
429            Self::Mapped {
430                model,
431                idx,
432                dtype: TensorDtype::Q1,
433                ..
434            } => Some((model, *idx)),
435            _ => None,
436        }
437    }
438
439    /// (model, idx, kind, row_scale) for a graph-capable mapped weight. kind:
440    /// 0=q8_row (per-row scales), 1=q1, 2=q4_tiled, 3=q1t (tile-embedded, no
441    /// rs). None for dtypes the token graph does not handle (q8_2f/q4_block/vbit).
442    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
443        match self {
444            Self::Mapped {
445                model,
446                idx,
447                dtype: TensorDtype::Q8Row,
448                row_scale,
449                ..
450            } => Some((model, *idx, 0, row_scale.as_slice())),
451            Self::Mapped {
452                model,
453                idx,
454                dtype: TensorDtype::Q1,
455                ..
456            } => Some((model, *idx, 1, &[])),
457            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
458            // the wgpu token graph fed 18B interleaved tiles to the
459            // split-layout q4b kernel — garbage output on q4t models
460            // (caught by an end-to-end answer check on real Vulkan).
461            Self::Mapped {
462                model,
463                idx,
464                dtype: TensorDtype::Q4Tiled,
465                ..
466            } => Some((model, *idx, 5, &[])),
467            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
468            // and feeding them to the q4t kernel is exactly the mistake that
469            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
470            Self::Mapped {
471                model,
472                idx,
473                dtype: TensorDtype::Q4TiledP,
474                ..
475            } => Some((model, *idx, 6, &[])),
476            Self::Mapped {
477                model,
478                idx,
479                dtype: TensorDtype::Q4Block,
480                ..
481            } => Some((model, *idx, 2, &[])),
482            Self::Mapped {
483                model,
484                idx,
485                dtype: TensorDtype::Q1T,
486                ..
487            } => Some((model, *idx, 3, &[])),
488            _ => None,
489        }
490    }
491
492    /// Dense f32 view — only for owned tensors. Masked/sparse execution
493    /// paths require it; quantized weights don't support masks yet.
494    pub fn as_f32(&self) -> Option<&[f32]> {
495        match self {
496            Self::F32 { data, .. } => Some(data),
497            Self::Mapped { .. } => None,
498        }
499    }
500
501    fn quant_bytes(&self) -> &[u8] {
502        match self {
503            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
504            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
505        }
506    }
507
508    /// Dequantize one row into `dst` (embedding lookup).
509    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
510        let cols = self.cols();
511        debug_assert_eq!(dst.len(), cols);
512        match self {
513            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
514            Self::Mapped {
515                dtype,
516                row_scale,
517                col_field,
518                vbit_offsets,
519                ..
520            } => {
521                if *dtype == TensorDtype::Q4Tiled {
522                    let bytes = self.quant_bytes();
523                    let gpr = cols / GROUP_SIZE;
524                    for gi in 0..gpr {
525                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
526                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
527                        for (k, &b) in tile[2..].iter().enumerate() {
528                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
529                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
530                        }
531                    }
532                    return;
533                }
534                if *dtype == TensorDtype::Q4TiledP {
535                    let bytes = self.quant_bytes();
536                    let gpr = cols / GROUP_SIZE;
537                    let v = Q4tpView::new(bytes, self.rows(), cols);
538                    let mut sc = vec![0f32; gpr];
539                    v.scales_into(r, gpr, &mut sc);
540                    for gi in 0..gpr {
541                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
542                        let s = sc[gi];
543                        for (k, &b) in tile.iter().enumerate() {
544                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
545                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
546                        }
547                    }
548                    return;
549                }
550                if *dtype == TensorDtype::Q4Block {
551                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
552                    let gpr = cols / GROUP_SIZE;
553                    for gi in 0..gpr {
554                        let g = r * gpr + gi;
555                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
556                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
557                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
558                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
559                        }
560                    }
561                    return;
562                }
563                if *dtype == TensorDtype::Q1 {
564                    let bytes = self.quant_bytes();
565                    let gpr = cols / GROUP_SIZE;
566                    for gi in 0..gpr {
567                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
568                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
569                        for (j, &b) in tile[2..].iter().enumerate() {
570                            for k in 0..8 {
571                                dst[gi * GROUP_SIZE + j * 8 + k] =
572                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
573                            }
574                        }
575                    }
576                    return;
577                }
578                if *dtype == TensorDtype::Q1T {
579                    let bytes = self.quant_bytes();
580                    let gpr = cols / GROUP_SIZE;
581                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
582                    for gi in 0..gpr {
583                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
584                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
585                            bytes[off],
586                            bytes[off + 1],
587                        ]));
588                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
589                        for k in 0..GROUP_SIZE {
590                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
591                            {
592                                1 => s,
593                                2 => -s,
594                                _ => 0.0,
595                            };
596                        }
597                    }
598                    // Overlay
599                    let rows = self.rows();
600                    let entries = base_len + (rows + 1) * 4;
601                    if entries <= bytes.len() {
602                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
603                        let r0 = u32::from_le_bytes([
604                            ptrs[r * 4],
605                            ptrs[r * 4 + 1],
606                            ptrs[r * 4 + 2],
607                            ptrs[r * 4 + 3],
608                        ]) as usize;
609                        let r1 = u32::from_le_bytes([
610                            ptrs[(r + 1) * 4],
611                            ptrs[(r + 1) * 4 + 1],
612                            ptrs[(r + 1) * 4 + 2],
613                            ptrs[(r + 1) * 4 + 3],
614                        ]) as usize;
615                        let off = entries + r0 * 4;
616                        for i in 0..r1 - r0 {
617                            let item = &bytes[off + i * 4..off + i * 4 + 4];
618                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
619                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
620                                item[2], item[3],
621                            ]));
622                            if c < cols {
623                                dst[c] = v;
624                            }
625                        }
626                    }
627                    return;
628                }
629                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
630                    let bytes = self.quant_bytes();
631                    let rows = self.rows();
632                    let ng = cols / GROUP_SIZE;
633                    let bits = &bytes[..rows];
634                    let sc_off = rows;
635                    // Precomputed at load — embedding lookup used to scan
636                    // the bit-widths of every preceding row (O(token_id)).
637                    let off = vbit_offsets[r];
638                    let b = bits[r] as usize;
639                    let l = ((1usize << (b - 1)) - 1) as f32;
640                    let data = &bytes[off..];
641                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
642                    for (i, d) in dst.iter_mut().enumerate() {
643                        while nbits < b {
644                            acc = (acc << 8) | data[idx] as u64;
645                            idx += 1;
646                            nbits += 8;
647                        }
648                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
649                        nbits -= b;
650                        let so = (r * ng + i / GROUP_SIZE) * 2;
651                        let sv = f16_to_f32(u16::from_le_bytes([
652                            bytes[sc_off + so],
653                            bytes[sc_off + so + 1],
654                        ]));
655                        *d = (u - l) * sv;
656                    }
657                    return;
658                }
659                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
660                let s = row_scale[r];
661                match dtype {
662                    TensorDtype::Q8Row => {
663                        for (d, &b) in dst.iter_mut().zip(q) {
664                            *d = (b as i8) as f32 * s;
665                        }
666                    }
667                    TensorDtype::Q8_2f => {
668                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
669                            *d = (b as i8) as f32 * s * col_field[i];
670                        }
671                    }
672                    _ => unreachable!(),
673                }
674            }
675        }
676    }
677
678    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
679    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
680    /// false for group-packed q4/vbit (column access would unpack whole
681    /// groups — sparse execution falls back to f32 for those).
682    pub fn sparse_col_ok(&self) -> bool {
683        match self {
684            Self::F32 { .. } => true,
685            Self::Mapped { dtype, .. } => {
686                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
687            }
688        }
689    }
690
691    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
692    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
693    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
694    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
695        let inter = self.cols();
696        let hidden = self.rows();
697        debug_assert_eq!(out.len(), hidden);
698        match self {
699            Self::F32 { data, .. } => {
700                for (k, o) in out.iter_mut().enumerate() {
701                    *o += w * data[k * inter + c];
702                }
703            }
704            Self::Mapped {
705                dtype,
706                row_scale,
707                col_field,
708                ..
709            } => {
710                let q = self.quant_bytes();
711                let colf = if *dtype == TensorDtype::Q8_2f {
712                    col_field[c]
713                } else {
714                    1.0
715                };
716                let wc = w * colf;
717                for (k, o) in out.iter_mut().enumerate() {
718                    let b = q[k * inter + c] as i8 as f32;
719                    *o += wc * b * row_scale[k];
720                }
721            }
722        }
723    }
724
725    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
726    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
727    /// into `scratch` first (rare for active-FFN weights).
728    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
729        let cols = self.cols();
730        match self {
731            Self::F32 { data, .. } => {
732                let row = &data[r * cols..(r + 1) * cols];
733                row.iter().zip(x).map(|(w, v)| w * v).sum()
734            }
735            Self::Mapped {
736                dtype,
737                row_scale,
738                col_field,
739                ..
740            } => match dtype {
741                TensorDtype::Q8Row => {
742                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
743                    dot_i8_f32(q, x) * row_scale[r]
744                }
745                TensorDtype::Q8_2f => {
746                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
747                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
748                }
749                _ => {
750                    self.row_f32(r, scratch);
751                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
752                }
753            },
754        }
755    }
756
757    /// `out = W · x` (row-major). F32 delegates to the historical
758    /// bit-exact path; Mapped runs the fused int8 kernel.
759    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
760        match self {
761            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
762            Self::Mapped {
763                model,
764                idx,
765                dtype,
766                rows,
767                cols,
768                row_scale,
769                col_field,
770                vbit_offsets,
771                repack,
772            } => {
773                let _ = (model, idx);
774                if *dtype == TensorDtype::Q4Block {
775                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
776                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
777                    // the winner; Metal returns false → the CPU kernel below.
778                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
779                        let t0 = std::time::Instant::now();
780                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
781                            crate::gpu::ProbeArm::Gpu => {
782                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
783                                    crate::gpu::probe_record(
784                                        crate::gpu::OpClass::Matvec,
785                                        true,
786                                        t0.elapsed(),
787                                    );
788                                    return;
789                                }
790                            }
791                            crate::gpu::ProbeArm::CpuTimed => {
792                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
793                                crate::gpu::probe_record(
794                                    crate::gpu::OpClass::Matvec,
795                                    false,
796                                    t0.elapsed(),
797                                );
798                                return;
799                            }
800                            crate::gpu::ProbeArm::Cpu => {}
801                        }
802                    }
803                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
804                    return;
805                }
806                if *dtype == TensorDtype::Q4Tiled {
807                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
808                    return;
809                }
810                if *dtype == TensorDtype::Q4TiledP {
811                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
812                    return;
813                }
814                if *dtype == TensorDtype::Q1 {
815                    // GPU route for large q1 matvecs (out_proj / lm_head
816                    // class): the CPU q1 kernel is load-port-bound at
817                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
818                    // probe measures both arms and keeps the winner.
819                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
820                        let t0 = std::time::Instant::now();
821                        let arm = if crate::gpu::q1_force() {
822                            crate::gpu::ProbeArm::Gpu
823                        } else {
824                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
825                        };
826                        match arm {
827                            crate::gpu::ProbeArm::Gpu => {
828                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
829                                    crate::gpu::probe_record(
830                                        crate::gpu::OpClass::Matvec,
831                                        true,
832                                        t0.elapsed(),
833                                    );
834                                    return;
835                                }
836                            }
837                            crate::gpu::ProbeArm::CpuTimed => {
838                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
839                                crate::gpu::probe_record(
840                                    crate::gpu::OpClass::Matvec,
841                                    false,
842                                    t0.elapsed(),
843                                );
844                                return;
845                            }
846                            crate::gpu::ProbeArm::Cpu => {}
847                        }
848                    }
849                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
850                    return;
851                }
852                if *dtype == TensorDtype::Q1T {
853                    // GPU route for large q1t matvecs: the ternary BASE dot runs
854                    // on the GPU (load-port-bound on CPU, like q1), then the
855                    // sparse overlay is added on the CPU. Probe keeps the winner.
856                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
857                        let t0 = std::time::Instant::now();
858                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
859                            crate::gpu::ProbeArm::Gpu => {
860                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
861                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
862                                    crate::gpu::probe_record(
863                                        crate::gpu::OpClass::Matvec,
864                                        true,
865                                        t0.elapsed(),
866                                    );
867                                    return;
868                                }
869                            }
870                            crate::gpu::ProbeArm::CpuTimed => {
871                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
872                                crate::gpu::probe_record(
873                                    crate::gpu::OpClass::Matvec,
874                                    false,
875                                    t0.elapsed(),
876                                );
877                                return;
878                            }
879                            crate::gpu::ProbeArm::Cpu => {}
880                        }
881                    }
882                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
883                    return;
884                }
885                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
886                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
887                    return;
888                }
889                let xs = prescale(x, col_field, *dtype);
890                // D5: large q8 matrices (lm_head-class) — hybrid
891                // CPU∥GPU: split the rows, both sides compute
892                // SIMULTANEOUSLY (same math, shared prescale).
893                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
894                if *rows >= crate::gpu::min_rows()
895                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
896                    && std::env::var("CMF_GPU_LMHEAD")
897                        .map(|v| v != "0")
898                        .unwrap_or(true)
899                    && crate::gpu::enabled_here()
900                {
901                    // Runtime probe: alternate the hybrid against the
902                    // pure-CPU matvec, keep whichever is faster HERE.
903                    let t0 = std::time::Instant::now();
904                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
905                        crate::gpu::ProbeArm::Gpu => {}
906                        crate::gpu::ProbeArm::CpuTimed => {
907                            qmatvec(
908                                self.quant_bytes(),
909                                repack,
910                                row_scale,
911                                x,
912                                col_field,
913                                *dtype,
914                                *rows,
915                                *cols,
916                                out,
917                                pool,
918                            );
919                            crate::gpu::probe_record(
920                                crate::gpu::OpClass::Matvec,
921                                false,
922                                t0.elapsed(),
923                            );
924                            return;
925                        }
926                        crate::gpu::ProbeArm::Cpu => {
927                            qmatvec(
928                                self.quant_bytes(),
929                                repack,
930                                row_scale,
931                                x,
932                                col_field,
933                                *dtype,
934                                *rows,
935                                *cols,
936                                out,
937                                pool,
938                            );
939                            return;
940                        }
941                    }
942                    let frac = std::env::var("CMF_GPU_SPLIT")
943                        .ok()
944                        .and_then(|v| v.parse::<f32>().ok())
945                        .unwrap_or(0.5)
946                        .clamp(0.0, 1.0);
947                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
948                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
949                    let bytes = self.quant_bytes();
950                    let ok = std::thread::scope(|sc| {
951                        let g = sc.spawn(|| {
952                            crate::gpu::q8_matvec_range(
953                                model,
954                                *idx,
955                                cpu_rows,
956                                &row_scale[cpu_rows..],
957                                &xs,
958                                *rows - cpu_rows,
959                                *cols,
960                                out_gpu,
961                            )
962                        });
963                        if cpu_rows > 0 {
964                            // Repack prefix covers the full groups of the
965                            // CPU half (the split starts at row 0).
966                            let rep_cpu = if repack.is_empty() {
967                                &[][..]
968                            } else {
969                                &repack[..(cpu_rows / 4) * 4 * *cols]
970                            };
971                            qmatvec(
972                                &bytes[..cpu_rows * *cols],
973                                rep_cpu,
974                                &row_scale[..cpu_rows],
975                                x,
976                                col_field,
977                                *dtype,
978                                cpu_rows,
979                                *cols,
980                                out_cpu,
981                                pool,
982                            );
983                        }
984                        g.join().unwrap_or(false)
985                    });
986                    if ok {
987                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
988                        return;
989                    }
990                    // GPU failed — CPU finishes its half (rows rebased —
991                    // group offsets don't line up, mmap layout only).
992                    qmatvec(
993                        &bytes[cpu_rows * *cols..(*rows) * *cols],
994                        &[],
995                        &row_scale[cpu_rows..],
996                        x,
997                        col_field,
998                        *dtype,
999                        *rows - cpu_rows,
1000                        *cols,
1001                        out_gpu,
1002                        pool,
1003                    );
1004                    return;
1005                }
1006                qmatvec(
1007                    self.quant_bytes(),
1008                    repack,
1009                    row_scale,
1010                    x,
1011                    col_field,
1012                    *dtype,
1013                    *rows,
1014                    *cols,
1015                    out,
1016                    pool,
1017                );
1018            }
1019        }
1020    }
1021
1022    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1023    pub fn matvec2(
1024        &self,
1025        x1: &[f32],
1026        x2: &[f32],
1027        o1: &mut [f32],
1028        o2: &mut [f32],
1029        pool: Option<&Pool>,
1030    ) {
1031        match self {
1032            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1033            Self::Mapped {
1034                dtype,
1035                rows,
1036                cols,
1037                row_scale,
1038                col_field,
1039                vbit_offsets,
1040                ..
1041            } => {
1042                if *dtype == TensorDtype::Q4Block {
1043                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1044                    return;
1045                }
1046                if *dtype == TensorDtype::Q4Tiled {
1047                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1048                    return;
1049                }
1050                if *dtype == TensorDtype::Q4TiledP {
1051                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1052                    return;
1053                }
1054                if *dtype == TensorDtype::Q1 {
1055                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1056                    return;
1057                }
1058                if *dtype == TensorDtype::Q1T {
1059                    // Fused ternary pair: one row pass, the register
1060                    // unpack shared across both streams on ARM. (Q1T
1061                    // lacks a row_scale array — scales live inline in
1062                    // the tiles — so it must not fall through to the
1063                    // q8 qmatvec2 below.)
1064                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1065                    return;
1066                }
1067                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1068                    vbitmatvec2(
1069                        self.quant_bytes(),
1070                        vbit_offsets,
1071                        x1,
1072                        x2,
1073                        *rows,
1074                        *cols,
1075                        o1,
1076                        o2,
1077                        pool,
1078                    );
1079                    return;
1080                }
1081                qmatvec2(
1082                    self.quant_bytes(),
1083                    row_scale,
1084                    x1,
1085                    x2,
1086                    col_field,
1087                    *dtype,
1088                    *rows,
1089                    *cols,
1090                    o1,
1091                    o2,
1092                    pool,
1093                );
1094            }
1095        }
1096    }
1097}
1098
1099impl QTensor {
1100    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1101    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1102    /// to b matvec calls (same dot kernels in the same order); the win —
1103    /// the weight row streams from DRAM once per batch, not b times.
1104    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1105        let cols = self.cols();
1106        let rows = self.rows();
1107        debug_assert_eq!(xs_all.len(), b * cols);
1108        debug_assert_eq!(out.len(), b * rows);
1109        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1110        // Mapped tensors carry a directory name; the check is a relaxed
1111        // atomic load, free when not calibrating.
1112        if crate::gptq_capture::capturing() {
1113            if let Self::Mapped { model, idx, .. } = self {
1114                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1115            }
1116        }
1117        match self {
1118            Self::F32 { data, .. } => {
1119                let out_addr = SendMut(out.as_mut_ptr());
1120                let run = |start: usize, end: usize| {
1121                    for o in start..end {
1122                        let row = &data[o * cols..(o + 1) * cols];
1123                        for bi in 0..b {
1124                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1125                            let mut acc = 0f32;
1126                            for j in 0..cols {
1127                                acc += row[j] * x[j];
1128                            }
1129                            unsafe { *out_addr.at(bi * rows + o) = acc };
1130                        }
1131                    }
1132                };
1133                dispatch_rows(pool, rows, &run);
1134            }
1135            Self::Mapped {
1136                dtype,
1137                row_scale,
1138                col_field,
1139                vbit_offsets,
1140                ..
1141            } => {
1142                if *dtype == TensorDtype::Q4Block {
1143                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1144                    return;
1145                }
1146                if *dtype == TensorDtype::Q4TiledP {
1147                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1148                    // device); the probe keeps whichever beats the CPU arm.
1149                    // Narrow (prompt-encode) and wide (DiT) batches probe
1150                    // as separate classes — the regimes have opposite
1151                    // winners and one shared verdict locked the wrong arm.
1152                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1153                    // (a fair-condition op is ≤~100 ms even at 1024px)
1154                    // means the device is contended by another process
1155                    // (e.g. a simulator) — verdicts are per-process, so
1156                    // without the bail the whole render crawls behind
1157                    // someone else's queue.
1158                    if b >= 32
1159                        && b * rows * cols >= 128_000_000
1160                        && cols % 32 == 0
1161                        && !crate::gpu::mm_killed()
1162                        && crate::gpu::enabled_here()
1163                    {
1164                        let class = if b >= 128 {
1165                            crate::gpu::OpClass::MatmatWide
1166                        } else {
1167                            crate::gpu::OpClass::Matmat
1168                        };
1169                        if let Self::Mapped { model, idx, .. } = self {
1170                            let t0 = std::time::Instant::now();
1171                            match crate::gpu::probe_arm(class) {
1172                                crate::gpu::ProbeArm::Gpu => {
1173                                    if crate::gpu::q4tp_matmat(
1174                                        model, *idx, xs_all, b, rows, cols, out,
1175                                    ) {
1176                                        let el = t0.elapsed();
1177                                        // Work-proportional budget: ~8× the
1178                                        // fair-device estimate (+20 ms slack).
1179                                        // An absolute cap missed the worst
1180                                        // case — contended ops sit at
1181                                        // 100–240 ms each and still bury a
1182                                        // render whose fair op is 3–9 ms.
1183                                        // Cold ops (first PSO build, buffer
1184                                        // alloc) are exempt: a one-off
1185                                        // ~50 ms compile is not contention.
1186                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1187                                        let budget = std::time::Duration::from_secs_f64(
1188                                            flops / 1.5e12 * 8.0 + 0.020,
1189                                        );
1190                                        if el > budget && !crate::gpu::probe_was_cold() {
1191                                            tracing::warn!(
1192                                                "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1193                                                 device contended, CPU for the rest of the process"
1194                                            );
1195                                            crate::gpu::mm_kill();
1196                                        }
1197                                        crate::gpu::probe_record(class, true, el);
1198                                        return;
1199                                    }
1200                                }
1201                                crate::gpu::ProbeArm::CpuTimed => {
1202                                    q4tp_matmat(
1203                                        self.quant_bytes(),
1204                                        xs_all,
1205                                        b,
1206                                        rows,
1207                                        cols,
1208                                        out,
1209                                        pool,
1210                                    );
1211                                    crate::gpu::probe_record(class, false, t0.elapsed());
1212                                    return;
1213                                }
1214                                crate::gpu::ProbeArm::Cpu => {}
1215                            }
1216                        }
1217                    }
1218                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1219                    return;
1220                }
1221                if *dtype == TensorDtype::Q4Tiled {
1222                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1223                    // device); the probe keeps whichever beats the CPU arm.
1224                    // Narrow (prompt-encode) and wide (DiT) batches probe
1225                    // as separate classes — the regimes have opposite
1226                    // winners and one shared verdict locked the wrong arm.
1227                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1228                    // (a fair-condition op is ≤~100 ms even at 1024px)
1229                    // means the device is contended by another process
1230                    // (e.g. a simulator) — verdicts are per-process, so
1231                    // without the bail the whole render crawls behind
1232                    // someone else's queue.
1233                    if b >= 32
1234                        && b * rows * cols >= 128_000_000
1235                        && cols % 32 == 0
1236                        && !crate::gpu::mm_killed()
1237                        && crate::gpu::enabled_here()
1238                    {
1239                        let class = if b >= 128 {
1240                            crate::gpu::OpClass::MatmatWide
1241                        } else {
1242                            crate::gpu::OpClass::Matmat
1243                        };
1244                        if let Self::Mapped { model, idx, .. } = self {
1245                            let t0 = std::time::Instant::now();
1246                            match crate::gpu::probe_arm(class) {
1247                                crate::gpu::ProbeArm::Gpu => {
1248                                    if crate::gpu::q4t_matmat(
1249                                        model, *idx, xs_all, b, rows, cols, out,
1250                                    ) {
1251                                        let el = t0.elapsed();
1252                                        // Work-proportional budget: ~8× the
1253                                        // fair-device estimate (+20 ms slack).
1254                                        // An absolute cap missed the worst
1255                                        // case — contended ops sit at
1256                                        // 100–240 ms each and still bury a
1257                                        // render whose fair op is 3–9 ms.
1258                                        // Cold ops (first PSO build, buffer
1259                                        // alloc) are exempt: a one-off
1260                                        // ~50 ms compile is not contention.
1261                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1262                                        let budget = std::time::Duration::from_secs_f64(
1263                                            flops / 1.5e12 * 8.0 + 0.020,
1264                                        );
1265                                        if el > budget && !crate::gpu::probe_was_cold() {
1266                                            tracing::warn!(
1267                                                "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1268                                                 device contended, CPU for the rest of the process"
1269                                            );
1270                                            crate::gpu::mm_kill();
1271                                        }
1272                                        crate::gpu::probe_record(class, true, el);
1273                                        return;
1274                                    }
1275                                }
1276                                crate::gpu::ProbeArm::CpuTimed => {
1277                                    q4t_matmat(
1278                                        self.quant_bytes(),
1279                                        xs_all,
1280                                        b,
1281                                        rows,
1282                                        cols,
1283                                        out,
1284                                        pool,
1285                                    );
1286                                    crate::gpu::probe_record(class, false, t0.elapsed());
1287                                    return;
1288                                }
1289                                crate::gpu::ProbeArm::Cpu => {}
1290                            }
1291                        }
1292                    }
1293                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1294                    return;
1295                }
1296                if *dtype == TensorDtype::Q1 {
1297                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1298                    // device); the probe keeps whichever beats the CPU matmat.
1299                    if b >= 32
1300                        && b * rows * cols >= 128_000_000
1301                        && cols % 64 == 0
1302                        && crate::gpu::enabled_here()
1303                    {
1304                        if let Self::Mapped { model, idx, .. } = self {
1305                            let t0 = std::time::Instant::now();
1306                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1307                                crate::gpu::ProbeArm::Gpu => {
1308                                    if crate::gpu::q1_matmat(
1309                                        model, *idx, xs_all, b, rows, cols, out,
1310                                    ) {
1311                                        crate::gpu::probe_record(
1312                                            crate::gpu::OpClass::Matmat,
1313                                            true,
1314                                            t0.elapsed(),
1315                                        );
1316                                        return;
1317                                    }
1318                                }
1319                                crate::gpu::ProbeArm::CpuTimed => {
1320                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1321                                    crate::gpu::probe_record(
1322                                        crate::gpu::OpClass::Matmat,
1323                                        false,
1324                                        t0.elapsed(),
1325                                    );
1326                                    return;
1327                                }
1328                                crate::gpu::ProbeArm::Cpu => {}
1329                            }
1330                        }
1331                    }
1332                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1333                    return;
1334                }
1335                if *dtype == TensorDtype::Q1T {
1336                    // GPU batched GEMM for wide prefill (base + overlay on the
1337                    // device); probe keeps the winner vs the CPU matmat.
1338                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1339                        if let Self::Mapped { model, idx, .. } = self {
1340                            let t0 = std::time::Instant::now();
1341                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1342                                crate::gpu::ProbeArm::Gpu => {
1343                                    if crate::gpu::q1t_matmat(
1344                                        model, *idx, xs_all, b, rows, cols, out,
1345                                    ) {
1346                                        crate::gpu::probe_record(
1347                                            crate::gpu::OpClass::Matmat,
1348                                            true,
1349                                            t0.elapsed(),
1350                                        );
1351                                        return;
1352                                    }
1353                                }
1354                                crate::gpu::ProbeArm::CpuTimed => {
1355                                    q1t_matmat(
1356                                        self.quant_bytes(),
1357                                        xs_all,
1358                                        b,
1359                                        rows,
1360                                        cols,
1361                                        out,
1362                                        pool,
1363                                    );
1364                                    crate::gpu::probe_record(
1365                                        crate::gpu::OpClass::Matmat,
1366                                        false,
1367                                        t0.elapsed(),
1368                                    );
1369                                    return;
1370                                }
1371                                crate::gpu::ProbeArm::Cpu => {}
1372                            }
1373                        }
1374                    }
1375                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1376                    return;
1377                }
1378                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1379                    vbitmatmat(
1380                        self.quant_bytes(),
1381                        vbit_offsets,
1382                        xs_all,
1383                        b,
1384                        rows,
1385                        cols,
1386                        out,
1387                        pool,
1388                    );
1389                    return;
1390                }
1391                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1392                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1393                    .collect();
1394                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1395                // work volume: submission carries b×rows×cols MACs).
1396                // Runtime probe: the naive GEMM shader + sync readback
1397                // lose to the CPU GEMM on slow driver stacks — alternate
1398                // both arms and keep the winner.
1399                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1400                    if let Self::Mapped { model, idx, .. } = self {
1401                        let t0 = std::time::Instant::now();
1402                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1403                            crate::gpu::ProbeArm::Gpu
1404                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1405                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1406                            {
1407                                // Cold weights during probing: the upload
1408                                // has started, the count runs on the CPU —
1409                                // the GPU arm samples on the next touch.
1410                                let q = self.quant_bytes();
1411                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1412                                return;
1413                            }
1414                            crate::gpu::ProbeArm::Gpu => {
1415                                let flat: Vec<f32> =
1416                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1417                                if crate::gpu::q8_matmat(
1418                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1419                                ) {
1420                                    crate::gpu::probe_record(
1421                                        crate::gpu::OpClass::Matmat,
1422                                        true,
1423                                        t0.elapsed(),
1424                                    );
1425                                    return;
1426                                }
1427                            }
1428                            crate::gpu::ProbeArm::CpuTimed => {
1429                                let q = self.quant_bytes();
1430                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1431                                crate::gpu::probe_record(
1432                                    crate::gpu::OpClass::Matmat,
1433                                    false,
1434                                    t0.elapsed(),
1435                                );
1436                                return;
1437                            }
1438                            crate::gpu::ProbeArm::Cpu => {}
1439                        }
1440                    }
1441                }
1442                let q = self.quant_bytes();
1443                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1444            }
1445        }
1446    }
1447}
1448
1449impl QTensor {
1450    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1451    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1452    /// barrier instead of N. Per-row math is the exact same kernel as
1453    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1454    /// Falls back to N sequential matvecs when the set is not a uniform
1455    /// q8-family/F32 group or there is no pool.
1456    pub fn matvec_many<const N: usize>(
1457        ts: [&QTensor; N],
1458        x: &[f32],
1459        mut outs: [&mut [f32]; N],
1460        pool: Option<&Pool>,
1461    ) {
1462        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1463        let uniform_q8 = ts.iter().all(|t| {
1464            matches!(
1465                t,
1466                Self::Mapped {
1467                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1468                    ..
1469                }
1470            )
1471        });
1472        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1473        let uniform_q4 = ts.iter().all(|t| {
1474            matches!(
1475                t,
1476                Self::Mapped {
1477                    dtype: TensorDtype::Q4Block,
1478                    ..
1479                }
1480            )
1481        });
1482        let uniform_vbit = ts.iter().all(|t| {
1483            matches!(
1484                t,
1485                Self::Mapped {
1486                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1487                    ..
1488                }
1489            )
1490        });
1491        let uniform_q1 = ts.iter().all(|t| {
1492            matches!(
1493                t,
1494                Self::Mapped {
1495                    dtype: TensorDtype::Q1,
1496                    ..
1497                }
1498            )
1499        });
1500        let uniform_q1t = ts.iter().all(|t| {
1501            matches!(
1502                t,
1503                Self::Mapped {
1504                    dtype: TensorDtype::Q1T,
1505                    ..
1506                }
1507            )
1508        });
1509        let Some(pool) = pool else {
1510            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1511                t.matvec(x, o, None);
1512            }
1513            return;
1514        };
1515        if total_rows < 256
1516            || !(uniform_q8
1517                || uniform_f32
1518                || uniform_q4
1519                || uniform_vbit
1520                || uniform_q1
1521                || uniform_q1t)
1522        {
1523            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1524                t.matvec(x, o, Some(pool));
1525            }
1526            return;
1527        }
1528
1529        if uniform_q1 {
1530            // One shared activation split + group sums (q1 has no col
1531            // field; the same input feeds every tensor).
1532            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1533            if a8w8_enabled() {
1534                let act = split_act(x);
1535                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1536                let (act, gsum) = (&act, &gsum);
1537                let closures: [_; N] = std::array::from_fn(|i| {
1538                    let (bytes, gpr, out) =
1539                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1540                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1541                });
1542                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1543                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1544                pool.run_many(&parts);
1545            } else {
1546                let closures: [_; N] = std::array::from_fn(|i| {
1547                    let (bytes, gpr, out) =
1548                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1549                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1550                });
1551                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1552                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1553                pool.run_many(&parts);
1554            }
1555            return;
1556        }
1557
1558        if uniform_q1t {
1559            // Q1T batched: one shared activation split + overlay decode,
1560            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1561            // and N−1 redundant split_act calls per layer).
1562            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1563            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1564            if a8w8_enabled() {
1565                let act = split_act(x);
1566                let act = &act;
1567                let x_ref = x;
1568                let closures: [_; N] = std::array::from_fn(|i| {
1569                    let bytes = ts[i].quant_bytes();
1570                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1571                    let gpr = cols / GROUP_SIZE;
1572                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1573                    let out = outs_addr[i];
1574                    move |s: usize, e: usize| {
1575                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1576                    }
1577                });
1578                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1579                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1580                pool.run_many(&parts);
1581            } else {
1582                let x_ref = x;
1583                let closures: [_; N] = std::array::from_fn(|i| {
1584                    let bytes = ts[i].quant_bytes();
1585                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1586                    let gpr = cols / GROUP_SIZE;
1587                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1588                    let out = outs_addr[i];
1589                    move |s: usize, e: usize| {
1590                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1591                    }
1592                });
1593                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1594                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1595                pool.run_many(&parts);
1596            }
1597            return;
1598        }
1599
1600        if uniform_q4 || uniform_vbit {
1601            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1602            // q4/vbit share one activation split — no per-tensor col field.
1603            if a8w8_enabled() {
1604                let act = split_act(x);
1605                let act = &act;
1606                if uniform_q4 {
1607                    let closures: [_; N] = std::array::from_fn(|i| {
1608                        let (packed, scales) =
1609                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1610                        let (gpr, cols, out) =
1611                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1612                        move |s: usize, e: usize| {
1613                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1614                        }
1615                    });
1616                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1617                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1618                    pool.run_many(&parts);
1619                } else {
1620                    let closures: [_; N] = std::array::from_fn(|i| {
1621                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1622                            unreachable!()
1623                        };
1624                        let (bytes, rows, cols, out) = (
1625                            ts[i].quant_bytes(),
1626                            ts[i].rows(),
1627                            ts[i].cols(),
1628                            outs_addr[i],
1629                        );
1630                        move |s: usize, e: usize| {
1631                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1632                        }
1633                    });
1634                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1635                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1636                    pool.run_many(&parts);
1637                }
1638                return;
1639            }
1640            if uniform_q4 {
1641                let closures: [_; N] = std::array::from_fn(|i| {
1642                    let (packed, scales) =
1643                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1644                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1645                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1646                });
1647                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1648                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1649                pool.run_many(&parts);
1650            } else {
1651                let closures: [_; N] = std::array::from_fn(|i| {
1652                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1653                        unreachable!()
1654                    };
1655                    let (bytes, rows, cols, out) = (
1656                        ts[i].quant_bytes(),
1657                        ts[i].rows(),
1658                        ts[i].cols(),
1659                        outs_addr[i],
1660                    );
1661                    move |s: usize, e: usize| {
1662                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1663                    }
1664                });
1665                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1666                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1667                pool.run_many(&parts);
1668            }
1669            return;
1670        }
1671
1672        if uniform_f32 {
1673            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1674            let closures: [_; N] = std::array::from_fn(|i| {
1675                let Self::F32 { data, cols, .. } = ts[i] else {
1676                    unreachable!()
1677                };
1678                let out = outs_addr[i];
1679                move |start: usize, end: usize| {
1680                    for o in start..end {
1681                        let row = &data[o * cols..(o + 1) * cols];
1682                        let mut sum = 0.0f32;
1683                        for j in 0..*cols {
1684                            sum += row[j] * x[j];
1685                        }
1686                        // SAFETY: disjoint (tensor, row) cells per worker.
1687                        unsafe { *out.at(o) = sum };
1688                    }
1689                }
1690            });
1691            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1692                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1693            pool.run_many(&parts);
1694            return;
1695        }
1696
1697        // Uniform q8-family: per-tensor prescale (q8_2f col fields
1698        // differ per tensor) + the shared range kernels.
1699        struct Ctx<'a> {
1700            bytes: &'a [u8],
1701            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1702            rep: &'a [u8],
1703            row_scale: &'a [f32],
1704            cols: usize,
1705            xs: std::borrow::Cow<'a, [f32]>,
1706        }
1707        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1708            let Self::Mapped {
1709                dtype,
1710                cols,
1711                row_scale,
1712                col_field,
1713                repack,
1714                ..
1715            } = ts[i]
1716            else {
1717                unreachable!()
1718            };
1719            Ctx {
1720                bytes: ts[i].quant_bytes(),
1721                rep: repack,
1722                row_scale,
1723                cols: *cols,
1724                xs: prescale(x, col_field, *dtype),
1725            }
1726        });
1727        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1728        #[cfg(target_arch = "aarch64")]
1729        if sdot_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_sdot(c.bytes, c.rep, 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        #[cfg(target_arch = "x86_64")]
1743        if avx2_a8w8_enabled() {
1744            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1745            let closures: [_; N] = std::array::from_fn(|i| {
1746                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1747                move |start: usize, end: usize| {
1748                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1749                }
1750            });
1751            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1752                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1753            pool.run_many(&parts);
1754            return;
1755        }
1756        let closures: [_; N] = std::array::from_fn(|i| {
1757            let (c, out) = (&ctxs[i], outs_addr[i]);
1758            move |start: usize, end: usize| {
1759                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1760            }
1761        });
1762        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1763            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1764        pool.run_many(&parts);
1765    }
1766}
1767
1768impl QTensor {
1769    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
1770    /// single pool dispatch — the MTP/pair decode path publishes one job
1771    /// for Q/K/V (and one for gate+up) instead of one per tensor.
1772    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
1773    #[allow(clippy::needless_range_loop)]
1774    pub fn matvec2_many<const N: usize>(
1775        ts: [&QTensor; N],
1776        x1: &[f32],
1777        x2: &[f32],
1778        mut o1s: [&mut [f32]; N],
1779        mut o2s: [&mut [f32]; N],
1780        pool: Option<&Pool>,
1781    ) {
1782        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1783        let uniform_q8 = ts.iter().all(|t| {
1784            matches!(
1785                t,
1786                Self::Mapped {
1787                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1788                    ..
1789                }
1790            )
1791        });
1792        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1793        let uniform_q4 = ts.iter().all(|t| {
1794            matches!(
1795                t,
1796                Self::Mapped {
1797                    dtype: TensorDtype::Q4Block,
1798                    ..
1799                }
1800            )
1801        });
1802        let uniform_vbit = ts.iter().all(|t| {
1803            matches!(
1804                t,
1805                Self::Mapped {
1806                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1807                    ..
1808                }
1809            )
1810        });
1811        let fusable = pool.is_some()
1812            && total_rows >= 256
1813            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1814        if !fusable {
1815            for i in 0..N {
1816                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1817            }
1818            return;
1819        }
1820        let pool = pool.unwrap();
1821
1822        if uniform_q4 || uniform_vbit {
1823            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1824            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1825            // q4/vbit share activation splits — no per-tensor col field.
1826            if a8w8_enabled() {
1827                let a1 = split_act(x1);
1828                let a2 = split_act(x2);
1829                let (a1, a2) = (&a1, &a2);
1830                if uniform_q4 {
1831                    let closures: [_; N] = std::array::from_fn(|i| {
1832                        let (packed, scales) =
1833                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1834                        let (gpr, cols, o1, o2) =
1835                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1836                        move |s: usize, e: usize| {
1837                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1838                        }
1839                    });
1840                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1841                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1842                    pool.run_many(&parts);
1843                } else {
1844                    let closures: [_; N] = std::array::from_fn(|i| {
1845                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1846                            unreachable!()
1847                        };
1848                        let (bytes, rows, cols, o1, o2) = (
1849                            ts[i].quant_bytes(),
1850                            ts[i].rows(),
1851                            ts[i].cols(),
1852                            p1[i],
1853                            p2[i],
1854                        );
1855                        move |s: usize, e: usize| {
1856                            vbit_range2_a8w8(
1857                                bytes,
1858                                vbit_offsets,
1859                                x1,
1860                                x2,
1861                                a1,
1862                                a2,
1863                                rows,
1864                                cols,
1865                                o1,
1866                                o2,
1867                                s,
1868                                e,
1869                            )
1870                        }
1871                    });
1872                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1873                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1874                    pool.run_many(&parts);
1875                }
1876                return;
1877            }
1878            if uniform_q4 {
1879                let closures: [_; N] = std::array::from_fn(|i| {
1880                    let (packed, scales) =
1881                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1882                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1883                    move |s: usize, e: usize| {
1884                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1885                    }
1886                });
1887                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1888                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1889                pool.run_many(&parts);
1890            } else {
1891                let closures: [_; N] = std::array::from_fn(|i| {
1892                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1893                        unreachable!()
1894                    };
1895                    let (bytes, rows, cols, o1, o2) = (
1896                        ts[i].quant_bytes(),
1897                        ts[i].rows(),
1898                        ts[i].cols(),
1899                        p1[i],
1900                        p2[i],
1901                    );
1902                    move |s: usize, e: usize| {
1903                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1904                    }
1905                });
1906                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1907                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1908                pool.run_many(&parts);
1909            }
1910            return;
1911        }
1912
1913        if uniform_f32 {
1914            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1915            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1916            let closures: [_; N] = std::array::from_fn(|i| {
1917                let Self::F32 { data, cols, .. } = ts[i] else {
1918                    unreachable!()
1919                };
1920                let (o1, o2) = (p1[i], p2[i]);
1921                move |start: usize, end: usize| {
1922                    for o in start..end {
1923                        let row = &data[o * cols..(o + 1) * cols];
1924                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
1925                        for j in 0..*cols {
1926                            s1 += row[j] * x1[j];
1927                            s2 += row[j] * x2[j];
1928                        }
1929                        // SAFETY: disjoint (tensor, row) cells per worker.
1930                        unsafe {
1931                            *o1.at(o) = s1;
1932                            *o2.at(o) = s2;
1933                        }
1934                    }
1935                }
1936            });
1937            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1938                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1939            pool.run_many(&parts);
1940            return;
1941        }
1942
1943        struct Ctx<'a> {
1944            bytes: &'a [u8],
1945            row_scale: &'a [f32],
1946            cols: usize,
1947            xs1: std::borrow::Cow<'a, [f32]>,
1948            xs2: std::borrow::Cow<'a, [f32]>,
1949        }
1950        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1951            let Self::Mapped {
1952                dtype,
1953                cols,
1954                row_scale,
1955                col_field,
1956                ..
1957            } = ts[i]
1958            else {
1959                unreachable!()
1960            };
1961            Ctx {
1962                bytes: ts[i].quant_bytes(),
1963                row_scale,
1964                cols: *cols,
1965                xs1: prescale(x1, col_field, *dtype),
1966                xs2: prescale(x2, col_field, *dtype),
1967            }
1968        });
1969        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1970        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1971        #[cfg(target_arch = "aarch64")]
1972        if sdot_enabled() {
1973            let acts: [(SplitAct, SplitAct); N] =
1974                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1975            let closures: [_; N] = std::array::from_fn(|i| {
1976                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1977                move |start: usize, end: usize| {
1978                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1979                }
1980            });
1981            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1982                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1983            pool.run_many(&parts);
1984            return;
1985        }
1986        #[cfg(target_arch = "x86_64")]
1987        if avx2_a8w8_enabled() {
1988            let acts: [(SplitAct, SplitAct); N] =
1989                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1990            let closures: [_; N] = std::array::from_fn(|i| {
1991                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1992                move |start: usize, end: usize| {
1993                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1994                }
1995            });
1996            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1997                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1998            pool.run_many(&parts);
1999            return;
2000        }
2001        let closures: [_; N] = std::array::from_fn(|i| {
2002            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2003            move |start: usize, end: usize| {
2004                q8_range2_f32(
2005                    c.bytes,
2006                    c.row_scale,
2007                    &c.xs1,
2008                    &c.xs2,
2009                    c.cols,
2010                    o1,
2011                    o2,
2012                    start,
2013                    end,
2014                )
2015            }
2016        });
2017        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2018            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2019        pool.run_many(&parts);
2020    }
2021
2022    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2023    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2024    /// no intermediate g/u buffers, no separate silu pass. Falls back
2025    /// (returns false) for unsupported dtype combos.
2026    pub fn matvec_silu_mul(
2027        gate: &QTensor,
2028        up: &QTensor,
2029        x: &[f32],
2030        out: &mut [f32],
2031        pool: Option<&Pool>,
2032    ) -> bool {
2033        let inter = gate.rows();
2034        debug_assert_eq!(up.rows(), inter);
2035        debug_assert_eq!(out.len(), inter);
2036        debug_assert_eq!(gate.cols(), up.cols());
2037        if !a8w8_enabled() {
2038            return false;
2039        }
2040        let act = split_act(x);
2041        let act = &act;
2042        let x_ref = x;
2043        let out_addr = SendMut(out.as_mut_ptr());
2044
2045        match (gate, up) {
2046            // Q4Block gate + Q4Block up (most common mobile q4 models)
2047            (
2048                Self::Mapped {
2049                    dtype: TensorDtype::Q4Block,
2050                    ..
2051                },
2052                Self::Mapped {
2053                    dtype: TensorDtype::Q4Block,
2054                    ..
2055                },
2056            ) => {
2057                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2058                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2059                let gpr = gate.cols() / GROUP_SIZE;
2060                let cols = gate.cols();
2061                let run = move |start: usize, end: usize| {
2062                    for r in start..end {
2063                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2064                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2065                        for &(j, xv) in &act.outliers {
2066                            let flat = r * cols + j;
2067                            let gb = gp[flat / 2];
2068                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2069                            let gsc = f16_to_f32(u16::from_le_bytes([
2070                                gs[(flat / GROUP_SIZE) * 2],
2071                                gs[(flat / GROUP_SIZE) * 2 + 1],
2072                            ]));
2073                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2074                            let ub = up_p[flat / 2];
2075                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2076                            let usc = f16_to_f32(u16::from_le_bytes([
2077                                up_s[(flat / GROUP_SIZE) * 2],
2078                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2079                            ]));
2080                            uv += ((un as i32 - 8) as f32) * usc * xv;
2081                        }
2082                        let silu_g = gv / (1.0 + (-gv).exp());
2083                        // SAFETY: disjoint row ranges per worker.
2084                        unsafe { *out_addr.at(r) = silu_g * uv };
2085                    }
2086                };
2087                dispatch_rows(pool, inter, &run);
2088                true
2089            }
2090            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2091            // streams sequential, silu·mul fused (same per-row math as
2092            // `q4t_matvec`).
2093            (
2094                Self::Mapped {
2095                    dtype: TensorDtype::Q4Tiled,
2096                    ..
2097                },
2098                Self::Mapped {
2099                    dtype: TensorDtype::Q4Tiled,
2100                    ..
2101                },
2102            ) => {
2103                let g_bytes = gate.quant_bytes();
2104                let u_bytes = up.quant_bytes();
2105                let gpr = gate.cols() / GROUP_SIZE;
2106                let run = move |start: usize, end: usize| {
2107                    for r in start..end {
2108                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2109                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2110                        for &(j, xv) in &act.outliers {
2111                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2112                            gv += w * s * xv;
2113                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2114                            uv += w * s * xv;
2115                        }
2116                        let silu_g = gv / (1.0 + (-gv).exp());
2117                        // SAFETY: disjoint row ranges per worker.
2118                        unsafe { *out_addr.at(r) = silu_g * uv };
2119                    }
2120                };
2121                dispatch_rows(pool, inter, &run);
2122                true
2123            }
2124            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2125            // each row's two ladders built once and spent on both streams.
2126            (
2127                Self::Mapped {
2128                    dtype: TensorDtype::Q4TiledP,
2129                    ..
2130                },
2131                Self::Mapped {
2132                    dtype: TensorDtype::Q4TiledP,
2133                    ..
2134                },
2135            ) => {
2136                let cols = gate.cols();
2137                let gpr = cols / GROUP_SIZE;
2138                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2139                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2140                let run = |start: usize, end: usize| {
2141                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2142                    for r in start..end {
2143                        gv_view.scales_into(r, gpr, &mut gsc);
2144                        uv_view.scales_into(r, gpr, &mut usc);
2145                        let mut gv =
2146                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2147                        let mut uv =
2148                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2149                        for &(j, xv) in &act.outliers {
2150                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2151                            gv += w * s * xv;
2152                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2153                            uv += w * s * xv;
2154                        }
2155                        let silu_g = gv / (1.0 + (-gv).exp());
2156                        // SAFETY: disjoint row ranges per worker.
2157                        unsafe { *out_addr.at(r) = silu_g * uv };
2158                    }
2159                };
2160                dispatch_rows(pool, inter, &run);
2161                true
2162            }
2163            // Q1T gate + Q1T up
2164            (
2165                Self::Mapped {
2166                    dtype: TensorDtype::Q1T,
2167                    ..
2168                },
2169                Self::Mapped {
2170                    dtype: TensorDtype::Q1T,
2171                    ..
2172                },
2173            ) => {
2174                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2175                let g_bytes = gate.quant_bytes();
2176                let u_bytes = up.quant_bytes();
2177                let gpr = gate.cols() / GROUP_SIZE;
2178                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2179                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2180                let run = move |start: usize, end: usize| {
2181                    for r in start..end {
2182                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2183                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2184                        for &(j, xv) in &act.outliers {
2185                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2186                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2187                        }
2188                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2189                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2190                        let silu_g = gv / (1.0 + (-gv).exp());
2191                        // SAFETY: disjoint row ranges per worker.
2192                        unsafe { *out_addr.at(r) = silu_g * uv };
2193                    }
2194                };
2195                dispatch_rows(pool, inter, &run);
2196                true
2197            }
2198            _ => false,
2199        }
2200    }
2201}
2202
2203/// Batched q8 kernel: same math as qmatvec, the row makes a single
2204/// pass from memory for the whole batch.
2205/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2206/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2207#[cfg(target_os = "macos")]
2208mod accel_blas {
2209    #[link(name = "Accelerate", kind = "framework")]
2210    unsafe extern "C" {
2211        pub fn cblas_sgemm(
2212            order: i32,
2213            trans_a: i32,
2214            trans_b: i32,
2215            m: i32,
2216            n: i32,
2217            k: i32,
2218            alpha: f32,
2219            a: *const f32,
2220            lda: i32,
2221            b: *const f32,
2222            ldb: i32,
2223            beta: f32,
2224            c: *mut f32,
2225            ldc: i32,
2226        );
2227    }
2228}
2229
2230#[cfg(target_os = "macos")]
2231pub(crate) fn accel_gemm_enabled() -> bool {
2232    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2233    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2234}
2235
2236/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2237/// same entry point, so the batched-attention path opens on mobile.
2238#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2239pub(crate) fn accel_gemm_enabled() -> bool {
2240    true
2241}
2242
2243/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2244/// micro-kernel with A broadcast against B panels — the mobile stand-in
2245/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2246/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2247/// k = head_dim or context), and the goal is removing the per-position
2248/// quadratic wall, not peak GEMM.
2249#[cfg(target_arch = "aarch64")]
2250#[allow(clippy::too_many_arguments)]
2251pub(crate) fn neon_gemm_rm(
2252    m: usize,
2253    n: usize,
2254    k: usize,
2255    alpha: f32,
2256    a: &[f32],
2257    lda: usize,
2258    b_mat: &[f32],
2259    ldb: usize,
2260    b_rows_are_n: bool,
2261    c: &mut [f32],
2262    ldc: usize,
2263) {
2264    debug_assert!(a.len() >= (m - 1) * lda + k);
2265    debug_assert!(c.len() >= (m - 1) * ldc + n);
2266    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2267    unsafe {
2268        use core::arch::aarch64::*;
2269        let mut i = 0usize;
2270        while i < m {
2271            let mi = (m - i).min(4);
2272            let mut j = 0usize;
2273            while j < n {
2274                let nj = (n - j).min(8);
2275                if mi == 4 && nj == 8 {
2276                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2277                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2278                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2279                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2280                    for p in 0..k {
2281                        let (b0, b1) = if b_rows_are_n {
2282                            // B is [n, k]: column p of Bᵀ = element p of
2283                            // eight consecutive B rows — gathered.
2284                            let base = b_mat.as_ptr().add(j * ldb + p);
2285                            let g = |o: usize| *base.add(o * ldb);
2286                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2287                        } else {
2288                            let base = b_mat.as_ptr().add(p * ldb + j);
2289                            (
2290                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2291                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2292                            )
2293                        };
2294                        let bv0 = vld1q_f32(b0.as_ptr());
2295                        let bv1 = vld1q_f32(b1.as_ptr());
2296                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2297                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2298                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2299                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2300                        c0a = vfmaq_f32(c0a, a0, bv0);
2301                        c0b = vfmaq_f32(c0b, a0, bv1);
2302                        c1a = vfmaq_f32(c1a, a1, bv0);
2303                        c1b = vfmaq_f32(c1b, a1, bv1);
2304                        c2a = vfmaq_f32(c2a, a2, bv0);
2305                        c2b = vfmaq_f32(c2b, a2, bv1);
2306                        c3a = vfmaq_f32(c3a, a3, bv0);
2307                        c3b = vfmaq_f32(c3b, a3, bv1);
2308                    }
2309                    let al = vdupq_n_f32(alpha);
2310                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2311                        .iter()
2312                        .enumerate()
2313                    {
2314                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2315                        vst1q_f32(dst, vmulq_f32(*ca, al));
2316                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2317                    }
2318                } else {
2319                    for r in 0..mi {
2320                        for q in 0..nj {
2321                            let mut acc = 0f32;
2322                            for p in 0..k {
2323                                let bv = if b_rows_are_n {
2324                                    b_mat[(j + q) * ldb + p]
2325                                } else {
2326                                    b_mat[p * ldb + j + q]
2327                                };
2328                                acc += a[(i + r) * lda + p] * bv;
2329                            }
2330                            c[(i + r) * ldc + j + q] = acc * alpha;
2331                        }
2332                    }
2333                }
2334                j += nj;
2335            }
2336            i += mi;
2337        }
2338    }
2339}
2340
2341/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2342#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2343#[allow(clippy::too_many_arguments)]
2344pub(crate) fn sgemm_rm(
2345    m: usize,
2346    n: usize,
2347    k: usize,
2348    alpha: f32,
2349    a: &[f32],
2350    lda: usize,
2351    b_mat: &[f32],
2352    ldb: usize,
2353    b_rows_are_n: bool,
2354    c: &mut [f32],
2355    ldc: usize,
2356) {
2357    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2358}
2359
2360/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
2361/// per-layer projection and applies it to every expert; a naive triple loop
2362/// would turn a two-minute job into half an hour).
2363#[allow(clippy::too_many_arguments)]
2364pub fn sgemm_public(
2365    m: usize,
2366    n: usize,
2367    k: usize,
2368    alpha: f32,
2369    a: &[f32],
2370    lda: usize,
2371    b_mat: &[f32],
2372    ldb: usize,
2373    b_rows_are_n: bool,
2374    c: &mut [f32],
2375    ldc: usize,
2376) {
2377    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
2378    {
2379        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2380    }
2381    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
2382    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
2383    // this, so correctness matters and throughput does not — a triple loop is
2384    // the honest fallback rather than a reason to make the tool macOS-only.
2385    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
2386    {
2387        for i in 0..m {
2388            for j in 0..n {
2389                let mut acc = 0f32;
2390                for p in 0..k {
2391                    let bv = if b_rows_are_n {
2392                        b_mat[j * ldb + p]
2393                    } else {
2394                        b_mat[p * ldb + j]
2395                    };
2396                    acc += a[i * lda + p] * bv;
2397                }
2398                c[i * ldc + j] = alpha * acc;
2399            }
2400        }
2401    }
2402}
2403
2404/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
2405/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
2406#[cfg(target_os = "macos")]
2407#[allow(clippy::too_many_arguments)]
2408pub(crate) fn sgemm_rm(
2409    m: usize,
2410    n: usize,
2411    k: usize,
2412    alpha: f32,
2413    a: &[f32],
2414    lda: usize,
2415    b_mat: &[f32],
2416    ldb: usize,
2417    b_rows_are_n: bool,
2418    c: &mut [f32],
2419    ldc: usize,
2420) {
2421    debug_assert!(a.len() >= (m - 1) * lda + k);
2422    debug_assert!(c.len() >= (m - 1) * ldc + n);
2423    // Test hook: route the attention GEMMs through the portable NEON
2424    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
2425    // measured without a phone in the loop. (Intel macOS has no NEON —
2426    // the hook is a no-op there, Accelerate continues below.)
2427    #[cfg(target_arch = "aarch64")]
2428    if std::env::var("CMF_FORCE_NEON_GEMM")
2429        .map(|v| v == "1")
2430        .unwrap_or(false)
2431    {
2432        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2433    }
2434    unsafe {
2435        accel_blas::cblas_sgemm(
2436            101, // RowMajor
2437            111, // NoTrans A
2438            if b_rows_are_n { 112 } else { 111 },
2439            m as i32,
2440            n as i32,
2441            k as i32,
2442            alpha,
2443            a.as_ptr(),
2444            lda as i32,
2445            b_mat.as_ptr(),
2446            ldb as i32,
2447            0.0,
2448            c.as_mut_ptr(),
2449            ldc as i32,
2450        );
2451    }
2452}
2453
2454/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
2455/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
2456/// on the AMX with one row-major sgemm. Tiles live in cache, weights
2457/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
2458/// logits shift within f32 rounding — tolerance-class, like every
2459/// reduction-order change; decode (M=1) never takes this path.
2460#[cfg(target_os = "macos")]
2461fn qmatmat_accel(
2462    q: &[u8],
2463    row_scale: &[f32],
2464    pre: &[std::borrow::Cow<'_, [f32]>],
2465    rows: usize,
2466    cols: usize,
2467    out: &mut [f32],
2468    pool: Option<&Pool>,
2469) {
2470    // NOTE: double-buffering the dequant against the sgemm (a scoped
2471    // thread driving the pool on tile k+1 while the caller multiplies
2472    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
2473    // multithreaded, and the dequant workers just steal its cores.
2474    const TR: usize = 2048;
2475    let b = pre.len();
2476    thread_local! {
2477        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2478        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2479    }
2480    XPANEL.with(|xp| {
2481        WTILE.with(|wt| {
2482            let mut xpanel = xp.borrow_mut();
2483            xpanel.clear();
2484            for x in pre {
2485                xpanel.extend_from_slice(x);
2486            }
2487            let mut wtile = wt.borrow_mut();
2488            wtile.resize(TR * cols, 0.0);
2489            let mut r0 = 0usize;
2490            while r0 < rows {
2491                let tr = TR.min(rows - r0);
2492                // Dequant the tile (scale folded) — pool-parallel.
2493                let wt_addr = SendMut(wtile.as_mut_ptr());
2494                let run = |start: usize, end: usize| {
2495                    for r in start..end {
2496                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2497                        let s = row_scale[r0 + r];
2498                        // SAFETY: workers cover disjoint r ranges.
2499                        let dst =
2500                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2501                        for (d, &v) in dst.iter_mut().zip(row) {
2502                            *d = (v as i8) as f32 * s;
2503                        }
2504                    }
2505                };
2506                dispatch_rows(pool, tr, &run);
2507                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
2508                unsafe {
2509                    accel_blas::cblas_sgemm(
2510                        101, // RowMajor
2511                        111, // NoTrans A
2512                        112, // Trans B
2513                        b as i32,
2514                        tr as i32,
2515                        cols as i32,
2516                        1.0,
2517                        xpanel.as_ptr(),
2518                        cols as i32,
2519                        wtile.as_ptr(),
2520                        cols as i32,
2521                        0.0,
2522                        out.as_mut_ptr().add(r0),
2523                        rows as i32,
2524                    );
2525                }
2526                r0 += tr;
2527            }
2528        })
2529    });
2530}
2531
2532fn qmatmat(
2533    q: &[u8],
2534    row_scale: &[f32],
2535    pre: &[std::borrow::Cow<'_, [f32]>],
2536    rows: usize,
2537    cols: usize,
2538    out: &mut [f32],
2539    pool: Option<&Pool>,
2540) {
2541    let b = pre.len();
2542    debug_assert_eq!(out.len(), b * rows);
2543    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
2544    // SDOT loop below peaks near the CPU's dot throughput, an order
2545    // below the matrix units. Small tensors and tiny test models stay
2546    // on the exact integer path.
2547    #[cfg(target_os = "macos")]
2548    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2549        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2550        return;
2551    }
2552    #[cfg(target_arch = "aarch64")]
2553    if sdot_enabled() {
2554        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2555        let out_addr = SendMut(out.as_mut_ptr());
2556        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
2557        // path IS the ARM prefill GEMM off Apple silicon).
2558        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2559            .map(|v| v != "0")
2560            .unwrap_or(true);
2561        let use_i8mm = i8mm_enabled();
2562        if blocked_ok {
2563            let run = |start: usize, end: usize| {
2564                let mut o = start;
2565                while o < end {
2566                    if o + 2 <= end {
2567                        let r0 = &q[o * cols..(o + 1) * cols];
2568                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2569                        let mut bi = 0usize;
2570                        while bi + 4 <= acts.len() {
2571                            let xs = [
2572                                acts[bi].xq.as_slice(),
2573                                acts[bi + 1].xq.as_slice(),
2574                                acts[bi + 2].xq.as_slice(),
2575                                acts[bi + 3].xq.as_slice(),
2576                            ];
2577                            let d = if use_i8mm {
2578                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2579                            } else {
2580                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2581                            };
2582                            for (r, row) in [r0, r1].into_iter().enumerate() {
2583                                for k in 0..4 {
2584                                    let act = &acts[bi + k];
2585                                    let mut v = d[r][k] as f32 * act.sx;
2586                                    for &(j, xv) in &act.outliers {
2587                                        v += (row[j] as i8) as f32 * xv;
2588                                    }
2589                                    unsafe {
2590                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2591                                    };
2592                                }
2593                            }
2594                            bi += 4;
2595                        }
2596                        while bi < acts.len() {
2597                            for (r, row) in [r0, r1].into_iter().enumerate() {
2598                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2599                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2600                            }
2601                            bi += 1;
2602                        }
2603                        o += 2;
2604                    } else {
2605                        let row = &q[o * cols..(o + 1) * cols];
2606                        for (bi, act) in acts.iter().enumerate() {
2607                            let v = row_dot_sdot(row, act) * row_scale[o];
2608                            unsafe { *out_addr.at(bi * rows + o) = v };
2609                        }
2610                        o += 1;
2611                    }
2612                }
2613            };
2614            dispatch_rows(pool, rows, &run);
2615            return;
2616        }
2617        let run = |start: usize, end: usize| {
2618            for o in start..end {
2619                let row = &q[o * cols..(o + 1) * cols];
2620                for (bi, act) in acts.iter().enumerate() {
2621                    let v = row_dot_sdot(row, act) * row_scale[o];
2622                    unsafe { *out_addr.at(bi * rows + o) = v };
2623                }
2624            }
2625        };
2626        dispatch_rows(pool, rows, &run);
2627        return;
2628    }
2629    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
2630    // (roadmap P0: two weight rows' abs() stay in registers across four
2631    // activation streams); VNNI machines keep the per-row bias-trick
2632    // dot, which is already throughput-bound there.
2633    #[cfg(target_arch = "x86_64")]
2634    if avx2_a8w8_enabled() {
2635        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2636        let out_addr = SendMut(out.as_mut_ptr());
2637        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
2638        // A/B on noisy shared-vCPU hosts).
2639        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2640            .map(|v| v != "0")
2641            .unwrap_or(true);
2642        if !avx512vnni_enabled() && blocked_ok {
2643            let run = |start: usize, end: usize| {
2644                let mut o = start;
2645                while o < end {
2646                    if o + 2 <= end {
2647                        let r0 = &q[o * cols..(o + 1) * cols];
2648                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2649                        let mut bi = 0usize;
2650                        while bi + 4 <= acts.len() {
2651                            let xs = [
2652                                acts[bi].xq.as_slice(),
2653                                acts[bi + 1].xq.as_slice(),
2654                                acts[bi + 2].xq.as_slice(),
2655                                acts[bi + 3].xq.as_slice(),
2656                            ];
2657                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
2658                            for (r, row) in [r0, r1].into_iter().enumerate() {
2659                                for k in 0..4 {
2660                                    let act = &acts[bi + k];
2661                                    let mut v = d[r][k] as f32 * act.sx;
2662                                    for &(j, xv) in &act.outliers {
2663                                        v += (row[j] as i8) as f32 * xv;
2664                                    }
2665                                    unsafe {
2666                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2667                                    };
2668                                }
2669                            }
2670                            bi += 4;
2671                        }
2672                        while bi < acts.len() {
2673                            for (r, row) in [r0, r1].into_iter().enumerate() {
2674                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
2675                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2676                            }
2677                            bi += 1;
2678                        }
2679                        o += 2;
2680                    } else {
2681                        let row = &q[o * cols..(o + 1) * cols];
2682                        for (bi, act) in acts.iter().enumerate() {
2683                            let v = row_dot_avx2(row, act) * row_scale[o];
2684                            unsafe { *out_addr.at(bi * rows + o) = v };
2685                        }
2686                        o += 1;
2687                    }
2688                }
2689            };
2690            dispatch_rows(pool, rows, &run);
2691            return;
2692        }
2693        let run = |start: usize, end: usize| {
2694            for o in start..end {
2695                let row = &q[o * cols..(o + 1) * cols];
2696                for (bi, act) in acts.iter().enumerate() {
2697                    let v = row_dot_avx2(row, act) * row_scale[o];
2698                    unsafe { *out_addr.at(bi * rows + o) = v };
2699                }
2700            }
2701        };
2702        dispatch_rows(pool, rows, &run);
2703        return;
2704    }
2705    let out_addr = SendMut(out.as_mut_ptr());
2706    let run = |start: usize, end: usize| {
2707        for o in start..end {
2708            let row = &q[o * cols..(o + 1) * cols];
2709            for (bi, x) in pre.iter().enumerate() {
2710                let mut acc = 0f32;
2711                for j in 0..cols {
2712                    acc += (row[j] as i8) as f32 * x[j];
2713                }
2714                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
2715            }
2716        }
2717    };
2718    dispatch_rows(pool, rows, &run);
2719}
2720
2721/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
2722/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
2723fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
2724    match pool {
2725        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
2726        _ => run(0, rows),
2727    }
2728}
2729
2730/// Split a q4_block blob into (packed nibbles, f16 group scales).
2731fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
2732    let groups = rows * cols / GROUP_SIZE;
2733    bytes.split_at(groups * 16)
2734}
2735
2736/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
2737/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
2738/// vbit packs MSB-first, so the HIGH nibble is the even element
2739/// (opposite of q4_block's lo-first interleave). Centering is u-7.
2740#[inline]
2741fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
2742    #[cfg(target_arch = "aarch64")]
2743    unsafe {
2744        return vbit_fill4_neon(data, buf);
2745    }
2746    #[cfg(target_arch = "x86_64")]
2747    if avx2_enabled() {
2748        return unsafe { vbit_fill4_avx2(data, buf) };
2749    }
2750    #[allow(unreachable_code)]
2751    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2752        let u = unpack8::<4>(&data[blk * 4..]);
2753        for k in 0..8 {
2754            chunk[k] = (u[k] - 7) as i8 as u8;
2755        }
2756    }
2757}
2758
2759#[cfg(target_arch = "aarch64")]
2760#[target_feature(enable = "neon")]
2761unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
2762    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
2763    // buf.len()/2 packed bytes (validated at load).
2764    unsafe {
2765        use core::arch::aarch64::*;
2766        let n = buf.len();
2767        let mask = vdupq_n_u8(0x0F);
2768        let seven = vdupq_n_s8(7);
2769        let mut g = 0usize;
2770        while g * 32 + 32 <= n {
2771            let b = vld1q_u8(data.as_ptr().add(g * 16));
2772            let hi = vshrq_n_u8::<4>(b);
2773            let lo = vandq_u8(b, mask);
2774            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
2775            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
2776            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
2777            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
2778            g += 1;
2779        }
2780    }
2781}
2782
2783#[cfg(target_arch = "x86_64")]
2784#[target_feature(enable = "avx2")]
2785unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2786    // SAFETY: see vbit_fill4_neon.
2787    unsafe {
2788        use core::arch::x86_64::*;
2789        let n = buf.len();
2790        let mask = _mm_set1_epi8(0x0F);
2791        let seven = _mm256_set1_epi8(7);
2792        let mut g = 0usize;
2793        while g * 32 + 32 <= n {
2794            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2795            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2796            let lo = _mm_and_si128(b, mask);
2797            let z = _mm256_sub_epi8(
2798                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2799                seven,
2800            );
2801            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2802            g += 1;
2803        }
2804    }
2805}
2806
2807/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
2808/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
2809/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
2810/// into 4 such blocks.
2811#[inline(always)]
2812fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2813    let mut acc = 0u64;
2814    for i in 0..B {
2815        acc = (acc << 8) | data[i] as u64;
2816    }
2817    let mask = (1u64 << B) - 1;
2818    let mut out = [0i32; 8];
2819    for (k, o) in out.iter_mut().enumerate() {
2820        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2821    }
2822    out
2823}
2824
2825/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
2826/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
2827/// MSB-first, byte-padded]. Row data offsets are precomputed at load
2828/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
2829/// overhead on every matvec.
2830#[allow(clippy::too_many_arguments)]
2831fn vbitmatvec(
2832    bytes: &[u8],
2833    offsets: &[usize],
2834    x: &[f32],
2835    rows: usize,
2836    cols: usize,
2837    out: &mut [f32],
2838    pool: Option<&Pool>,
2839) {
2840    debug_assert_eq!(out.len(), rows);
2841    debug_assert_eq!(offsets.len(), rows + 1);
2842
2843    // SDOT path: unpack the row to centered i8 once, then per-group
2844    // int8 dot against the quantized activations — same A8W8 contract
2845    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
2846    if a8w8_enabled() {
2847        let act = split_act(x);
2848        let out_addr = SendMut(out.as_mut_ptr());
2849        let run = move |start: usize, end: usize| {
2850            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2851        };
2852        dispatch_rows(pool, rows, &run);
2853        return;
2854    }
2855
2856    let out_addr = SendMut(out.as_mut_ptr());
2857    let run = move |start: usize, end: usize| {
2858        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2859    };
2860    dispatch_rows(pool, rows, &run);
2861}
2862
2863/// One vbit row range via the A8W8 int8 path — kernel body of
2864/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
2865/// several tensors in one dispatch (b=8 rows go exact f32).
2866#[allow(clippy::too_many_arguments)]
2867fn vbit_range_a8w8(
2868    bytes: &[u8],
2869    offsets: &[usize],
2870    x: &[f32],
2871    act: &SplitAct,
2872    rows: usize,
2873    cols: usize,
2874    out: SendMut,
2875    start: usize,
2876    end: usize,
2877) {
2878    let ng = cols / GROUP_SIZE;
2879    let bits = &bytes[..rows];
2880    let sc_off = rows;
2881    let row_dot = |r: usize| -> f32 {
2882        let b = bits[r] as usize;
2883        let l = (1i32 << (b - 1)) - 1;
2884        let mask = (1u64 << b) - 1;
2885        let data = &bytes[offsets[r]..offsets[r + 1]];
2886        if b == 8 {
2887            // u−L reaches 128 → does not fit i8; exact f32 path.
2888            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2889            let mut dot = 0f32;
2890            for g in 0..ng {
2891                let so = (r * ng + g) * 2;
2892                let sgf = f16_to_f32(u16::from_le_bytes([
2893                    bytes[sc_off + so],
2894                    bytes[sc_off + so + 1],
2895                ]));
2896                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2897                let mut gd = 0f32;
2898                for &xv in xg.iter() {
2899                    if nbits < 8 {
2900                        acc = (acc << 8) | data[idx] as u64;
2901                        idx += 1;
2902                        nbits += 8;
2903                    }
2904                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2905                    nbits -= 8;
2906                    gd += (u - l) as f32 * xv;
2907                }
2908                dot += gd * sgf;
2909            }
2910            return dot;
2911        }
2912        // Per-worker scratch: this closure runs for every row of the
2913        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
2914        // row was measurable pure overhead.
2915        thread_local! {
2916            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2917                const { std::cell::RefCell::new(Vec::new()) };
2918        }
2919        #[inline(always)]
2920        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2921            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2922                let u = unpack8::<B>(&data[blk * B..]);
2923                for k in 0..8 {
2924                    chunk[k] = (u[k] - l) as i8 as u8;
2925                }
2926            }
2927        }
2928        let _ = mask;
2929        VBIT_SCRATCH.with(|scratch| {
2930            let mut buf = scratch.borrow_mut();
2931            buf.resize(cols, 0);
2932            match b {
2933                3 => fill::<3>(data, l, &mut buf),
2934                4 => vbit_fill4(data, &mut buf),
2935                5 => fill::<5>(data, l, &mut buf),
2936                6 => fill::<6>(data, l, &mut buf),
2937                _ => unreachable!(),
2938            }
2939            let mut dot = 0f32;
2940            for g in 0..ng {
2941                let so = (r * ng + g) * 2;
2942                let s = f16_to_f32(u16::from_le_bytes([
2943                    bytes[sc_off + so],
2944                    bytes[sc_off + so + 1],
2945                ]));
2946                let d = dot_i8_i8(
2947                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2948                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2949                ) as f32
2950                    * act.sx;
2951                dot += d * s;
2952            }
2953            for &(j, xv) in &act.outliers {
2954                let so = (r * ng + j / GROUP_SIZE) * 2;
2955                let s = f16_to_f32(u16::from_le_bytes([
2956                    bytes[sc_off + so],
2957                    bytes[sc_off + so + 1],
2958                ]));
2959                // xq is zeroed at outlier slots — add the exact term.
2960                dot += (buf[j] as i8) as f32 * s * xv;
2961            }
2962            dot
2963        })
2964    };
2965    for r in start..end {
2966        // SAFETY: disjoint row ranges per worker.
2967        unsafe { *out.at(r) = row_dot(r) };
2968    }
2969}
2970
2971/// Exact scalar vbit row range (same extraction, non-SDOT path).
2972#[allow(clippy::too_many_arguments)]
2973fn vbit_range_f32(
2974    bytes: &[u8],
2975    offsets: &[usize],
2976    x: &[f32],
2977    rows: usize,
2978    cols: usize,
2979    out: SendMut,
2980    start: usize,
2981    end: usize,
2982) {
2983    let ng = cols / GROUP_SIZE;
2984    let bits = &bytes[..rows];
2985    let sc_off = rows;
2986    // Per-bit-width specialized inner loops: the compiler unrolls the
2987    // constant shifts (the generic bit-buffer loop was branch-bound —
2988    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
2989    #[inline(always)]
2990    fn dot_row<const B: usize>(
2991        data: &[u8],
2992        bytes: &[u8],
2993        sc_off: usize,
2994        r: usize,
2995        ng: usize,
2996        x: &[f32],
2997    ) -> f32 {
2998        let l = ((1i32 << (B - 1)) - 1) as f32;
2999        let gbytes = GROUP_SIZE * B / 8;
3000        let mut dot = 0f32;
3001        for g in 0..ng {
3002            let so = (r * ng + g) * 2;
3003            let s = f16_to_f32(u16::from_le_bytes([
3004                bytes[sc_off + so],
3005                bytes[sc_off + so + 1],
3006            ]));
3007            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3008            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3009            let mut gd = 0f32;
3010            for blk in 0..GROUP_SIZE / 8 {
3011                let u = unpack8::<B>(&gd0[blk * B..]);
3012                let xb = &xg[blk * 8..blk * 8 + 8];
3013                for k in 0..8 {
3014                    gd += (u[k] as f32 - l) * xb[k];
3015                }
3016            }
3017            dot += gd * s;
3018        }
3019        dot
3020    }
3021    for r in start..end {
3022        let data = &bytes[offsets[r]..offsets[r + 1]];
3023        let v = match bits[r] {
3024            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3025            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3026            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3027            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3028            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3029            b => unreachable!("vbit bit-width {b} (validated at load)"),
3030        };
3031        // SAFETY: disjoint row ranges per worker.
3032        unsafe { *out.at(r) = v };
3033    }
3034}
3035
3036/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3037/// and dotted against BOTH activations (MTP verify / pair prefill used
3038/// to run two full matvecs — double weight traffic and double unpack).
3039/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3040#[allow(clippy::too_many_arguments)]
3041fn vbitmatvec2(
3042    bytes: &[u8],
3043    offsets: &[usize],
3044    x1: &[f32],
3045    x2: &[f32],
3046    rows: usize,
3047    cols: usize,
3048    o1: &mut [f32],
3049    o2: &mut [f32],
3050    pool: Option<&Pool>,
3051) {
3052    debug_assert_eq!(o1.len(), rows);
3053    debug_assert_eq!(o2.len(), rows);
3054
3055    if a8w8_enabled() {
3056        let a1 = split_act(x1);
3057        let a2 = split_act(x2);
3058        let p1 = SendMut(o1.as_mut_ptr());
3059        let p2 = SendMut(o2.as_mut_ptr());
3060        let run = move |start: usize, end: usize| {
3061            vbit_range2_a8w8(
3062                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3063            )
3064        };
3065        dispatch_rows(pool, rows, &run);
3066        return;
3067    }
3068
3069    let p1 = SendMut(o1.as_mut_ptr());
3070    let p2 = SendMut(o2.as_mut_ptr());
3071    let run = move |start: usize, end: usize| {
3072        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3073    };
3074    dispatch_rows(pool, rows, &run);
3075}
3076
3077/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3078/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3079/// exact f32 for both lanes, bits streamed once).
3080#[allow(clippy::too_many_arguments)]
3081fn vbit_range2_a8w8(
3082    bytes: &[u8],
3083    offsets: &[usize],
3084    x1: &[f32],
3085    x2: &[f32],
3086    a1: &SplitAct,
3087    a2: &SplitAct,
3088    rows: usize,
3089    cols: usize,
3090    p1: SendMut,
3091    p2: SendMut,
3092    start: usize,
3093    end: usize,
3094) {
3095    let ng = cols / GROUP_SIZE;
3096    let bits = &bytes[..rows];
3097    let sc_off = rows;
3098    let row_dots = |r: usize| -> (f32, f32) {
3099        let b = bits[r] as usize;
3100        let l = (1i32 << (b - 1)) - 1;
3101        let data = &bytes[offsets[r]..offsets[r + 1]];
3102        if b == 8 {
3103            // u−L reaches 128 → does not fit i8; exact f32 path,
3104            // bits still streamed once for both lanes.
3105            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3106            let (mut d1, mut d2) = (0f32, 0f32);
3107            for g in 0..ng {
3108                let so = (r * ng + g) * 2;
3109                let sgf = f16_to_f32(u16::from_le_bytes([
3110                    bytes[sc_off + so],
3111                    bytes[sc_off + so + 1],
3112                ]));
3113                let (mut g1, mut g2) = (0f32, 0f32);
3114                for k in 0..GROUP_SIZE {
3115                    if nbits < 8 {
3116                        acc = (acc << 8) | data[idx] as u64;
3117                        idx += 1;
3118                        nbits += 8;
3119                    }
3120                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3121                    nbits -= 8;
3122                    let w = (u - l) as f32;
3123                    g1 += w * x1[g * GROUP_SIZE + k];
3124                    g2 += w * x2[g * GROUP_SIZE + k];
3125                }
3126                d1 += g1 * sgf;
3127                d2 += g2 * sgf;
3128            }
3129            return (d1, d2);
3130        }
3131        thread_local! {
3132            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3133                const { std::cell::RefCell::new(Vec::new()) };
3134        }
3135        #[inline(always)]
3136        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3137            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3138                let u = unpack8::<B>(&data[blk * B..]);
3139                for k in 0..8 {
3140                    chunk[k] = (u[k] - l) as i8 as u8;
3141                }
3142            }
3143        }
3144        VBIT_SCRATCH2.with(|scratch| {
3145            let mut buf = scratch.borrow_mut();
3146            buf.resize(cols, 0);
3147            match b {
3148                3 => fill::<3>(data, l, &mut buf),
3149                4 => vbit_fill4(data, &mut buf),
3150                5 => fill::<5>(data, l, &mut buf),
3151                6 => fill::<6>(data, l, &mut buf),
3152                _ => unreachable!(),
3153            }
3154            let (mut d1, mut d2) = (0f32, 0f32);
3155            for g in 0..ng {
3156                let so = (r * ng + g) * 2;
3157                let s = f16_to_f32(u16::from_le_bytes([
3158                    bytes[sc_off + so],
3159                    bytes[sc_off + so + 1],
3160                ]));
3161                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3162                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3163                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3164                d1 += v1 * s;
3165                d2 += v2 * s;
3166            }
3167            for &(j, xv) in &a1.outliers {
3168                let so = (r * ng + j / GROUP_SIZE) * 2;
3169                let s = f16_to_f32(u16::from_le_bytes([
3170                    bytes[sc_off + so],
3171                    bytes[sc_off + so + 1],
3172                ]));
3173                d1 += (buf[j] as i8) as f32 * s * xv;
3174            }
3175            for &(j, xv) in &a2.outliers {
3176                let so = (r * ng + j / GROUP_SIZE) * 2;
3177                let s = f16_to_f32(u16::from_le_bytes([
3178                    bytes[sc_off + so],
3179                    bytes[sc_off + so + 1],
3180                ]));
3181                d2 += (buf[j] as i8) as f32 * s * xv;
3182            }
3183            (d1, d2)
3184        })
3185    };
3186    for r in start..end {
3187        let (v1, v2) = row_dots(r);
3188        // SAFETY: disjoint row ranges per worker.
3189        unsafe {
3190            *p1.at(r) = v1;
3191            *p2.at(r) = v2;
3192        }
3193    }
3194}
3195
3196/// Two-input exact scalar vbit row range (same extraction) —
3197/// per-bit-width specialized, two accumulators per row; per-lane
3198/// accumulation order matches `vbitmatvec` exactly.
3199#[allow(clippy::too_many_arguments)]
3200fn vbit_range2_f32(
3201    bytes: &[u8],
3202    offsets: &[usize],
3203    x1: &[f32],
3204    x2: &[f32],
3205    rows: usize,
3206    cols: usize,
3207    p1: SendMut,
3208    p2: SendMut,
3209    start: usize,
3210    end: usize,
3211) {
3212    let ng = cols / GROUP_SIZE;
3213    let bits = &bytes[..rows];
3214    let sc_off = rows;
3215    #[inline(always)]
3216    #[allow(clippy::too_many_arguments)]
3217    fn dot_row2<const B: usize>(
3218        data: &[u8],
3219        bytes: &[u8],
3220        sc_off: usize,
3221        r: usize,
3222        ng: usize,
3223        x1: &[f32],
3224        x2: &[f32],
3225    ) -> (f32, f32) {
3226        let l = ((1i32 << (B - 1)) - 1) as f32;
3227        let gbytes = GROUP_SIZE * B / 8;
3228        let (mut d1, mut d2) = (0f32, 0f32);
3229        for g in 0..ng {
3230            let so = (r * ng + g) * 2;
3231            let s = f16_to_f32(u16::from_le_bytes([
3232                bytes[sc_off + so],
3233                bytes[sc_off + so + 1],
3234            ]));
3235            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3236            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3237            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3238            let (mut g1, mut g2) = (0f32, 0f32);
3239            for blk in 0..GROUP_SIZE / 8 {
3240                let u = unpack8::<B>(&gd0[blk * B..]);
3241                for k in 0..8 {
3242                    let w = u[k] as f32 - l;
3243                    g1 += w * x1g[blk * 8 + k];
3244                    g2 += w * x2g[blk * 8 + k];
3245                }
3246            }
3247            d1 += g1 * s;
3248            d2 += g2 * s;
3249        }
3250        (d1, d2)
3251    }
3252    for r in start..end {
3253        let data = &bytes[offsets[r]..offsets[r + 1]];
3254        let (v1, v2) = match bits[r] {
3255            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3256            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3257            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3258            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3259            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3260            b => unreachable!("vbit bit-width {b} (validated at load)"),
3261        };
3262        // SAFETY: disjoint row ranges per worker.
3263        unsafe {
3264            *p1.at(r) = v1;
3265            *p2.at(r) = v2;
3266        }
3267    }
3268}
3269
3270// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3271
3272/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3273/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3274/// distant streams of the split layout. Values/order identical to the
3275/// split kernels.
3276#[inline]
3277#[allow(unreachable_code)]
3278fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3279    #[cfg(target_arch = "aarch64")]
3280    unsafe {
3281        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3282    }
3283    #[cfg(target_arch = "x86_64")]
3284    unsafe {
3285        if vnni_tiles_enabled() {
3286            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3287        }
3288        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3289    }
3290    let mut acc = 0f32;
3291    for gi in 0..gpr {
3292        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3293        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3294        let mut d = 0i32;
3295        for (k, &b) in tile[2..].iter().enumerate() {
3296            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3297                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3298        }
3299        acc += d as f32 * s;
3300    }
3301    acc
3302}
3303
3304#[cfg(target_arch = "aarch64")]
3305#[target_feature(enable = "neon,dotprod")]
3306unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3307    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3308    // xq.len() == gpr·GROUP_SIZE).
3309    unsafe {
3310        use core::arch::aarch64::*;
3311        use core::arch::asm;
3312        let lomask = vdupq_n_u8(0x0F);
3313        let eight = vdupq_n_s8(8);
3314        let mut acc = 0f32;
3315        for gi in 0..gpr {
3316            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3317            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3318            let b = vld1q_u8(t.add(2));
3319            let lo = vandq_u8(b, lomask);
3320            let hi = vshrq_n_u8::<4>(b);
3321            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3322            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3323            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3324            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3325            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3326            asm!(
3327                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3328                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3329                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3330                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3331                options(pure, nomem, nostack),
3332            );
3333            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3334        }
3335        acc
3336    }
3337}
3338
3339#[cfg(target_arch = "x86_64")]
3340#[target_feature(enable = "avx2")]
3341unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3342    // SAFETY: see dot_q4t_row_sdot.
3343    unsafe {
3344        use core::arch::x86_64::*;
3345        let lomask = _mm_set1_epi8(0x0F);
3346        let eight = _mm256_set1_epi8(8);
3347        let ones = _mm256_set1_epi16(1);
3348        let mut acc = 0f32;
3349        for gi in 0..gpr {
3350            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3351            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3352            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3353            let lo = _mm_and_si128(b, lomask);
3354            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3355            let w = _mm256_sub_epi8(
3356                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3357                eight,
3358            );
3359            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3360            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3361            let d = _mm256_madd_epi16(p16, ones);
3362            let hi128 = _mm256_extracti128_si256::<1>(d);
3363            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3364            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3365            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3366            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3367        }
3368        acc
3369    }
3370}
3371
3372/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
3373/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
3374/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
3375#[cfg(target_arch = "x86_64")]
3376#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3377unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3378    // SAFETY: see dot_q4t_row_sdot.
3379    unsafe {
3380        use core::arch::x86_64::*;
3381        let lomask = _mm_set1_epi8(0x0F);
3382        let eight = _mm256_set1_epi8(8);
3383        let mut acc = 0f32;
3384        for gi in 0..gpr {
3385            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3386            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3387            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3388            let lo = _mm_and_si128(b, lomask);
3389            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3390            let w = _mm256_sub_epi8(
3391                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3392                eight,
3393            );
3394            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3395            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3396            acc += d as f32 * s;
3397        }
3398        acc
3399    }
3400}
3401
3402/// One q4_tiled row against FOUR activation streams: the nibble unpack
3403/// and abs() happen once per group instead of once per (group,
3404/// activation) — the unpack is the dominant per-element cost of the
3405/// tiled format (roadmap P0 portable blocking, q4t leg).
3406#[cfg(target_arch = "x86_64")]
3407// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
3408// to a libm call per lane — measured 2x slower than the reduction this
3409// kernel replaces. The runtime gate (`avx2_enabled`) already requires
3410// both features, so declaring it here is safe.
3411#[target_feature(enable = "avx2,fma")]
3412unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3413    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3414    unsafe {
3415        use core::arch::x86_64::*;
3416        let lomask = _mm_set1_epi8(0x0F);
3417        let eight = _mm256_set1_epi8(8);
3418        let ones = _mm256_set1_epi16(1);
3419        // One f32 accumulator VECTOR per activation, reduced once at the
3420        // end. Folding each group's i32 lanes to a scalar inside the loop
3421        // costs an extracti128 + three shift/add + a movd — a cross-lane
3422        // dependency chain per (group, activation), 288 of them per row at
3423        // cols=2304. The per-group scale is what forces a float
3424        // accumulator; it does not force a horizontal sum.
3425        //
3426        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
3427        // indexed by a loop variable LLVM keeps them in memory and every
3428        // group pays four 32-byte loads and stores. That alone made this
3429        // kernel 2x SLOWER than the per-group reduction it replaces
3430        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
3431        let mut f0 = _mm256_setzero_ps();
3432        let mut f1 = _mm256_setzero_ps();
3433        let mut f2 = _mm256_setzero_ps();
3434        let mut f3 = _mm256_setzero_ps();
3435        for gi in 0..gpr {
3436            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3437            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3438            let sv = _mm256_set1_ps(s);
3439            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3440            let lo = _mm_and_si128(bb, lomask);
3441            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3442            let w = _mm256_sub_epi8(
3443                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3444                eight,
3445            );
3446            let aw = _mm256_abs_epi8(w);
3447            let off = gi * GROUP_SIZE;
3448            let dot = |xq: &[i8]| {
3449                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3450                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3451                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3452            };
3453            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3454            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3455            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3456            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3457        }
3458        [
3459            hsum256_ps(f0),
3460            hsum256_ps(f1),
3461            hsum256_ps(f2),
3462            hsum256_ps(f3),
3463        ]
3464    }
3465}
3466
3467/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
3468/// blocked kernels pay, once per row instead of once per group.
3469#[cfg(target_arch = "x86_64")]
3470#[target_feature(enable = "avx2")]
3471#[inline]
3472unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3473    // SAFETY: pure register arithmetic on the caller's vector.
3474    unsafe {
3475        use core::arch::x86_64::*;
3476        let hi = _mm256_extractf128_ps::<1>(v);
3477        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3478        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3479        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3480        _mm_cvtss_f32(s)
3481    }
3482}
3483
3484/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
3485#[cfg(target_arch = "x86_64")]
3486#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3487unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3488    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3489    unsafe {
3490        use core::arch::x86_64::*;
3491        let lomask = _mm_set1_epi8(0x0F);
3492        let eight = _mm256_set1_epi8(8);
3493        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
3494        // one cross-lane reduction per row, not per (group, activation).
3495        let mut f0 = _mm256_setzero_ps();
3496        let mut f1 = _mm256_setzero_ps();
3497        let mut f2 = _mm256_setzero_ps();
3498        let mut f3 = _mm256_setzero_ps();
3499        for gi in 0..gpr {
3500            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3501            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3502            let sv = _mm256_set1_ps(s);
3503            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3504            let lo = _mm_and_si128(bb, lomask);
3505            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3506            let w = _mm256_sub_epi8(
3507                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3508                eight,
3509            );
3510            let aw = _mm256_abs_epi8(w);
3511            let off = gi * GROUP_SIZE;
3512            let dot = |xq: &[i8]| {
3513                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3514                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3515                    _mm256_setzero_si256(),
3516                    aw,
3517                    _mm256_sign_epi8(x, w),
3518                ))
3519            };
3520            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3521            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3522            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3523            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3524        }
3525        let acc = [
3526            hsum256_ps(f0),
3527            hsum256_ps(f1),
3528            hsum256_ps(f2),
3529            hsum256_ps(f3),
3530        ];
3531        acc
3532    }
3533}
3534
3535/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
3536/// serves FOUR activation streams. Per stream the group order and f32
3537/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
3538/// bit-for-bit.
3539#[cfg(target_arch = "aarch64")]
3540#[target_feature(enable = "neon,dotprod")]
3541unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3542    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3543    unsafe {
3544        use core::arch::aarch64::*;
3545        use core::arch::asm;
3546        let lomask = vdupq_n_u8(0x0F);
3547        let eight = vdupq_n_s8(8);
3548        let mut acc = [0f32; 4];
3549        for gi in 0..gpr {
3550            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3551            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3552            let b = vld1q_u8(t.add(2));
3553            let lo = vandq_u8(b, lomask);
3554            let hi = vshrq_n_u8::<4>(b);
3555            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3556            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3557            for (k, xq) in xs.iter().enumerate() {
3558                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3559                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3560                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3561                asm!(
3562                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3563                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3564                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3565                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3566                    options(pure, nomem, nostack),
3567                );
3568                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3569            }
3570        }
3571        acc
3572    }
3573}
3574
3575/// Exact-term correction for A8W8 outliers on a tiled row.
3576#[inline]
3577fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3578    let gi = j / GROUP_SIZE;
3579    let k = j % GROUP_SIZE;
3580    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3581    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3582    let byte = tile[2 + k / 2];
3583    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3584    ((nib as i32 - 8) as f32, s)
3585}
3586
3587/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
3588/// accumulation shape as `q4_range_f32`.
3589#[inline]
3590fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3591    let mut acc = 0f32;
3592    for gi in 0..gpr {
3593        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3594        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3595        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3596        let mut ga = 0f32;
3597        for (k, &b) in tile[2..].iter().enumerate() {
3598            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3599                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3600        }
3601        acc += ga * s;
3602    }
3603    acc
3604}
3605
3606/// Split view of a `q4tp` payload. The three planes are resolved once per
3607/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
3608/// the row loop would put a division on the hot path for nothing.
3609struct Q4tpView<'a> {
3610    nib: &'a [u8],
3611    params: &'a [u8],
3612    codes: &'a [u8],
3613    stride: usize,
3614}
3615
3616impl<'a> Q4tpView<'a> {
3617    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
3618        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
3619        Self {
3620            nib: &bytes[..params_off],
3621            params: &bytes[params_off..codes_off],
3622            codes: &bytes[codes_off..],
3623            stride,
3624        }
3625    }
3626
3627    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
3628    ///
3629    /// Doing this once per row — rather than decoding a 5-bit code inside the
3630    /// tile loop — is what makes the format free at runtime. Random access to
3631    /// a packed 5-bit field costs a division, two bounds checks and a branch;
3632    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
3633    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
3634    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
3635    #[inline]
3636    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
3637        let tab = q4tp_ladder(self.params, r);
3638        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
3639        let (mut acc, mut have, mut bi) = (0u64, 0u32, 0usize);
3640        for o in out.iter_mut().take(gpr) {
3641            while have < 5 {
3642                acc |= (codes[bi] as u64) << have;
3643                bi += 1;
3644                have += 8;
3645            }
3646            *o = tab[(acc & 31) as usize];
3647            acc >>= 5;
3648            have -= 5;
3649        }
3650    }
3651}
3652
3653#[inline]
3654fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
3655    #[cfg(target_arch = "aarch64")]
3656    unsafe {
3657        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
3658    }
3659    #[cfg(target_arch = "x86_64")]
3660    unsafe {
3661        if vnni_tiles_enabled() {
3662            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
3663        }
3664        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
3665    }
3666    #[allow(unreachable_code)]
3667    {
3668        let mut acc = 0f32;
3669        for gi in 0..gpr {
3670            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3671            let s = scales[gi];
3672            let mut d = 0i32;
3673            for (k, &b) in tile.iter().enumerate() {
3674                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3675                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3676            }
3677            acc += d as f32 * s;
3678        }
3679        acc
3680    }
3681}
3682
3683/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
3684/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
3685#[cfg(target_arch = "aarch64")]
3686#[target_feature(enable = "neon,dotprod")]
3687unsafe fn dot_q4tp_row_sdot(
3688    nib: &[u8],
3689    r: usize,
3690    gpr: usize,
3691    xq: &[i8],
3692    scales: &[f32],
3693) -> f32 {
3694    // SAFETY: callers uphold slice-length contracts (16B tile per group,
3695    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
3696    unsafe {
3697        use core::arch::aarch64::*;
3698        use core::arch::asm;
3699        let lomask = vdupq_n_u8(0x0F);
3700        let eight = vdupq_n_s8(8);
3701        let mut acc = 0f32;
3702        for gi in 0..gpr {
3703            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3704            let s = *scales.get_unchecked(gi);
3705            let b = vld1q_u8(t);
3706            let lo = vandq_u8(b, lomask);
3707            let hi = vshrq_n_u8::<4>(b);
3708            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3709            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3710            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3711            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3712            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3713            asm!(
3714                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3715                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3716                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3717                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3718                options(pure, nomem, nostack),
3719            );
3720            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3721        }
3722        acc
3723    }
3724}
3725
3726#[cfg(target_arch = "x86_64")]
3727#[target_feature(enable = "avx2")]
3728unsafe fn dot_q4tp_row_avx2(
3729    nib: &[u8],
3730    r: usize,
3731    gpr: usize,
3732    xq: &[i8],
3733    scales: &[f32],
3734) -> f32 {
3735    // SAFETY: see dot_q4tp_row_sdot.
3736    unsafe {
3737        use core::arch::x86_64::*;
3738        let lomask = _mm_set1_epi8(0x0F);
3739        let eight = _mm256_set1_epi8(8);
3740        let ones = _mm256_set1_epi16(1);
3741        let mut acc = 0f32;
3742        for gi in 0..gpr {
3743            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3744            let s = *scales.get_unchecked(gi);
3745            let b = _mm_loadu_si128(t as *const __m128i);
3746            let lo = _mm_and_si128(b, lomask);
3747            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3748            let w = _mm256_sub_epi8(
3749                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3750                eight,
3751            );
3752            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3753            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3754            let d = _mm256_madd_epi16(p16, ones);
3755            let hi128 = _mm256_extracti128_si256::<1>(d);
3756            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3757            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3758            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3759            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3760        }
3761        acc
3762    }
3763}
3764
3765/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
3766/// 256-bit VL encoding is the one to use here).
3767#[cfg(target_arch = "x86_64")]
3768#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3769unsafe fn dot_q4tp_row_vnni(
3770    nib: &[u8],
3771    r: usize,
3772    gpr: usize,
3773    xq: &[i8],
3774    scales: &[f32],
3775) -> f32 {
3776    // SAFETY: see dot_q4tp_row_sdot.
3777    unsafe {
3778        use core::arch::x86_64::*;
3779        let lomask = _mm_set1_epi8(0x0F);
3780        let eight = _mm256_set1_epi8(8);
3781        let mut acc = 0f32;
3782        for gi in 0..gpr {
3783            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3784            let s = *scales.get_unchecked(gi);
3785            let b = _mm_loadu_si128(t as *const __m128i);
3786            let lo = _mm_and_si128(b, lomask);
3787            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3788            let w = _mm256_sub_epi8(
3789                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3790                eight,
3791            );
3792            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3793            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
3794        }
3795        acc
3796    }
3797}
3798
3799/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
3800/// accumulation shape as `q4t_row_exact`.
3801#[inline]
3802fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
3803    let mut acc = 0f32;
3804    for gi in 0..gpr {
3805        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3806        let s = scales[gi];
3807        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3808        let mut ga = 0f32;
3809        for (k, &b) in tile.iter().enumerate() {
3810            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3811                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3812        }
3813        acc += ga * s;
3814    }
3815    acc
3816}
3817
3818/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
3819/// activation outliers at full precision after the int8 pass.
3820#[inline]
3821fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
3822    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
3823    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
3824    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3825    ((n as i32 - 8) as f32, scales[gi])
3826}
3827
3828/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
3829fn q4tp_matvec(
3830    bytes: &[u8],
3831    x: &[f32],
3832    rows: usize,
3833    cols: usize,
3834    out: &mut [f32],
3835    pool: Option<&Pool>,
3836) {
3837    debug_assert_eq!(out.len(), rows);
3838    let gpr = cols / GROUP_SIZE;
3839    let v = Q4tpView::new(bytes, rows, cols);
3840    let out_addr = SendMut(out.as_mut_ptr());
3841    if a8w8_enabled() {
3842        let act = split_act(x);
3843        let run = |start: usize, end: usize| {
3844            // One scratch row of scales per worker, reused across its rows.
3845            let mut sc = vec![0f32; gpr];
3846            for r in start..end {
3847                v.scales_into(r, gpr, &mut sc);
3848                let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
3849                for &(j, xv) in &act.outliers {
3850                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3851                    acc += w * s * xv;
3852                }
3853                // SAFETY: disjoint row ranges per worker.
3854                unsafe { *out_addr.at(r) = acc };
3855            }
3856        };
3857        dispatch_rows(pool, rows, &run);
3858        return;
3859    }
3860    let run = |start: usize, end: usize| {
3861        let mut sc = vec![0f32; gpr];
3862        for r in start..end {
3863            v.scales_into(r, gpr, &mut sc);
3864            // SAFETY: disjoint row ranges per worker.
3865            unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
3866        }
3867    };
3868    dispatch_rows(pool, rows, &run);
3869}
3870
3871/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
3872/// row ladder are read once and spent on both activation streams.
3873#[allow(clippy::too_many_arguments)]
3874fn q4tp_matvec2(
3875    bytes: &[u8],
3876    x1: &[f32],
3877    x2: &[f32],
3878    rows: usize,
3879    cols: usize,
3880    o1: &mut [f32],
3881    o2: &mut [f32],
3882    pool: Option<&Pool>,
3883) {
3884    let gpr = cols / GROUP_SIZE;
3885    let v = Q4tpView::new(bytes, rows, cols);
3886    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
3887    let run = |start: usize, end: usize| {
3888        let mut sc = vec![0f32; gpr];
3889        for r in start..end {
3890            v.scales_into(r, gpr, &mut sc);
3891            // SAFETY: disjoint row ranges per worker.
3892            unsafe {
3893                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
3894                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
3895            }
3896        }
3897    };
3898    dispatch_rows(pool, rows, &run);
3899}
3900
3901/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
3902/// spent on four activation streams, which is where a prefill batch stops
3903/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
3904#[cfg(target_arch = "aarch64")]
3905#[target_feature(enable = "neon,dotprod")]
3906unsafe fn dot_q4tp_row_1x4_sdot(
3907    nib: &[u8],
3908    r: usize,
3909    gpr: usize,
3910    xs: [&[i8]; 4],
3911    scales: &[f32],
3912) -> [f32; 4] {
3913    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
3914    unsafe {
3915        use core::arch::aarch64::*;
3916        use core::arch::asm;
3917        let lomask = vdupq_n_u8(0x0F);
3918        let eight = vdupq_n_s8(8);
3919        // Named accumulators, NOT an array indexed by a loop variable: the
3920        // latter does not stay in registers (the same defect cost 2x in the
3921        // AVX2 q4t kernel and again in WGSL).
3922        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
3923        for gi in 0..gpr {
3924            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3925            let s = *scales.get_unchecked(gi);
3926            let bb = vld1q_u8(t);
3927            let lo = vandq_u8(bb, lomask);
3928            let hi = vshrq_n_u8::<4>(bb);
3929            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3930            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3931            let mut d = [0f32; 4];
3932            for (k, dk) in d.iter_mut().enumerate() {
3933                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
3934                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
3935                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3936                asm!(
3937                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3938                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3939                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3940                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3941                    options(pure, nomem, nostack),
3942                );
3943                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3944            }
3945            f0 += d[0];
3946            f1 += d[1];
3947            f2 += d[2];
3948            f3 += d[3];
3949        }
3950        [f0, f1, f2, f3]
3951    }
3952}
3953
3954/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
3955/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
3956/// the format was fine, the missing arms were the whole regression.
3957fn q4tp_matmat(
3958    bytes: &[u8],
3959    xs_all: &[f32],
3960    b: usize,
3961    rows: usize,
3962    cols: usize,
3963    out: &mut [f32],
3964    pool: Option<&Pool>,
3965) {
3966    debug_assert_eq!(out.len(), b * rows);
3967    let gpr = cols / GROUP_SIZE;
3968    let v = Q4tpView::new(bytes, rows, cols);
3969
3970    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
3971    #[cfg(target_os = "macos")]
3972    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3973        dequant_matmat_accel(
3974            &|r, dst| {
3975                let mut sc = [0f32; 32];
3976                let mut scv;
3977                let s: &[f32] = if gpr <= 32 {
3978                    v.scales_into(r, gpr, &mut sc);
3979                    &sc[..gpr]
3980                } else {
3981                    scv = vec![0f32; gpr];
3982                    v.scales_into(r, gpr, &mut scv);
3983                    &scv
3984                };
3985                for gi in 0..gpr {
3986                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3987                    for (k, &bb) in tile.iter().enumerate() {
3988                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
3989                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
3990                    }
3991                }
3992            },
3993            xs_all,
3994            b,
3995            rows,
3996            cols,
3997            out,
3998            pool,
3999        );
4000        return;
4001    }
4002
4003    let out_addr = SendMut(out.as_mut_ptr());
4004    if a8w8_enabled() {
4005        let acts: Vec<SplitAct> = (0..b)
4006            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4007            .collect();
4008        let acts = &acts;
4009        #[cfg(target_arch = "aarch64")]
4010        let blocked_ok = sdot_enabled()
4011            && std::env::var("CMF_X86_BLOCKED")
4012                .map(|val| val != "0")
4013                .unwrap_or(true);
4014        #[cfg(not(target_arch = "aarch64"))]
4015        let blocked_ok = false;
4016        let run = |start: usize, end: usize| {
4017            let mut sc = vec![0f32; gpr];
4018            for r in start..end {
4019                v.scales_into(r, gpr, &mut sc);
4020                let mut bi = 0usize;
4021                #[cfg(target_arch = "aarch64")]
4022                if blocked_ok {
4023                    while bi + 4 <= acts.len() {
4024                        let xs = [
4025                            acts[bi].xq.as_slice(),
4026                            acts[bi + 1].xq.as_slice(),
4027                            acts[bi + 2].xq.as_slice(),
4028                            acts[bi + 3].xq.as_slice(),
4029                        ];
4030                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc) };
4031                        for k in 0..4 {
4032                            let act = &acts[bi + k];
4033                            let mut acc = d[k] * act.sx;
4034                            for &(j, xv) in &act.outliers {
4035                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4036                                acc += w * s * xv;
4037                            }
4038                            // SAFETY: disjoint (bi, r) cells per worker.
4039                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4040                        }
4041                        bi += 4;
4042                    }
4043                }
4044                let _ = blocked_ok;
4045                while bi < acts.len() {
4046                    let act = &acts[bi];
4047                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4048                    for &(j, xv) in &act.outliers {
4049                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4050                        acc += w * s * xv;
4051                    }
4052                    // SAFETY: disjoint (bi, r) cells per worker range.
4053                    unsafe { *out_addr.at(bi * rows + r) = acc };
4054                    bi += 1;
4055                }
4056            }
4057        };
4058        dispatch_rows(pool, rows, &run);
4059        return;
4060    }
4061
4062    let run = |start: usize, end: usize| {
4063        let mut sc = vec![0f32; gpr];
4064        for r in start..end {
4065            v.scales_into(r, gpr, &mut sc);
4066            for bi in 0..b {
4067                let x = &xs_all[bi * cols..(bi + 1) * cols];
4068                // SAFETY: disjoint (bi, r) cells per worker range.
4069                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4070            }
4071        }
4072    };
4073    dispatch_rows(pool, rows, &run);
4074}
4075
4076/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
4077fn q4t_matvec(
4078    bytes: &[u8],
4079    x: &[f32],
4080    rows: usize,
4081    cols: usize,
4082    out: &mut [f32],
4083    pool: Option<&Pool>,
4084) {
4085    debug_assert_eq!(out.len(), rows);
4086    let gpr = cols / GROUP_SIZE;
4087    let out_addr = SendMut(out.as_mut_ptr());
4088    if a8w8_enabled() {
4089        let act = split_act(x);
4090        let run = move |start: usize, end: usize| {
4091            for r in start..end {
4092                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4093                for &(j, xv) in &act.outliers {
4094                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4095                    acc += w * s * xv;
4096                }
4097                // SAFETY: disjoint row ranges per worker.
4098                unsafe { *out_addr.at(r) = acc };
4099            }
4100        };
4101        dispatch_rows(pool, rows, &run);
4102        return;
4103    }
4104    let run = move |start: usize, end: usize| {
4105        for r in start..end {
4106            // SAFETY: disjoint row ranges per worker.
4107            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
4108        }
4109    };
4110    dispatch_rows(pool, rows, &run);
4111}
4112
4113/// Fused two-input q4_tiled matvec (weights read once per pair).
4114#[allow(clippy::too_many_arguments)]
4115fn q4t_matvec2(
4116    bytes: &[u8],
4117    x1: &[f32],
4118    x2: &[f32],
4119    rows: usize,
4120    cols: usize,
4121    o1: &mut [f32],
4122    o2: &mut [f32],
4123    pool: Option<&Pool>,
4124) {
4125    let gpr = cols / GROUP_SIZE;
4126    let p1 = SendMut(o1.as_mut_ptr());
4127    let p2 = SendMut(o2.as_mut_ptr());
4128    if a8w8_enabled() {
4129        let a1 = split_act(x1);
4130        let a2 = split_act(x2);
4131        let run = move |start: usize, end: usize| {
4132            for r in start..end {
4133                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
4134                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
4135                for &(j, xv) in &a1.outliers {
4136                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4137                    v1 += w * s * xv;
4138                }
4139                for &(j, xv) in &a2.outliers {
4140                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4141                    v2 += w * s * xv;
4142                }
4143                // SAFETY: disjoint row ranges per worker.
4144                unsafe {
4145                    *p1.at(r) = v1;
4146                    *p2.at(r) = v2;
4147                }
4148            }
4149        };
4150        dispatch_rows(pool, rows, &run);
4151        return;
4152    }
4153    let run = move |start: usize, end: usize| {
4154        for r in start..end {
4155            // SAFETY: disjoint row ranges per worker.
4156            unsafe {
4157                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
4158                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
4159            }
4160        }
4161    };
4162    dispatch_rows(pool, rows, &run);
4163}
4164
4165/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
4166#[allow(clippy::too_many_arguments)]
4167/// Prefill GEMM through Accelerate for group-quantized codecs: a
4168/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
4169/// each tile rides the AMX with one sgemm — the generic sibling of
4170/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
4171/// decode (b=1) never takes this path.
4172#[cfg(target_os = "macos")]
4173fn dequant_matmat_accel(
4174    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
4175    xs_all: &[f32],
4176    b: usize,
4177    rows: usize,
4178    cols: usize,
4179    out: &mut [f32],
4180    pool: Option<&Pool>,
4181) {
4182    const TR: usize = 2048;
4183    thread_local! {
4184        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4185    }
4186    WTILE.with(|wt| {
4187        let mut wtile = wt.borrow_mut();
4188        wtile.resize(TR * cols, 0.0);
4189        let mut r0 = 0usize;
4190        while r0 < rows {
4191            let tr = TR.min(rows - r0);
4192            let wt_addr = SendMut(wtile.as_mut_ptr());
4193            let run = |start: usize, end: usize| {
4194                for r in start..end {
4195                    // SAFETY: workers cover disjoint r ranges.
4196                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4197                    dequant_row(r0 + r, dst);
4198                }
4199            };
4200            dispatch_rows(pool, tr, &run);
4201            unsafe {
4202                accel_blas::cblas_sgemm(
4203                    101, // RowMajor
4204                    111, // NoTrans A
4205                    112, // Trans B
4206                    b as i32,
4207                    tr as i32,
4208                    cols as i32,
4209                    1.0,
4210                    xs_all.as_ptr(),
4211                    cols as i32,
4212                    wtile.as_ptr(),
4213                    cols as i32,
4214                    0.0,
4215                    out.as_mut_ptr().add(r0),
4216                    rows as i32,
4217                );
4218            }
4219            r0 += tr;
4220        }
4221    });
4222}
4223
4224fn q4t_matmat(
4225    bytes: &[u8],
4226    xs_all: &[f32],
4227    b: usize,
4228    rows: usize,
4229    cols: usize,
4230    out: &mut [f32],
4231    pool: Option<&Pool>,
4232) {
4233    debug_assert_eq!(out.len(), b * rows);
4234    let gpr = cols / GROUP_SIZE;
4235    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
4236    // the dequant-tile sgemm is an order above the SDOT row loop for
4237    // prefill shapes (imagegen DiT forwards are exactly this).
4238    #[cfg(target_os = "macos")]
4239    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4240        dequant_matmat_accel(
4241            &|r, dst| {
4242                for gi in 0..gpr {
4243                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4244                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4245                    for (k, &bb) in tile[2..].iter().enumerate() {
4246                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
4247                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
4248                    }
4249                }
4250            },
4251            xs_all,
4252            b,
4253            rows,
4254            cols,
4255            out,
4256            pool,
4257        );
4258        return;
4259    }
4260    let out_addr = SendMut(out.as_mut_ptr());
4261    if a8w8_enabled() {
4262        let acts: Vec<SplitAct> = (0..b)
4263            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4264            .collect();
4265        let acts = &acts;
4266        #[cfg(target_arch = "x86_64")]
4267        let blocked_ok = avx2_enabled()
4268            && std::env::var("CMF_X86_BLOCKED")
4269                .map(|v| v != "0")
4270                .unwrap_or(true);
4271        #[cfg(target_arch = "aarch64")]
4272        let blocked_ok = sdot_enabled()
4273            && std::env::var("CMF_X86_BLOCKED")
4274                .map(|v| v != "0")
4275                .unwrap_or(true);
4276        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
4277        let blocked_ok = false;
4278        let run = move |start: usize, end: usize| {
4279            for r in start..end {
4280                let mut bi = 0usize;
4281                #[cfg(target_arch = "aarch64")]
4282                if blocked_ok {
4283                    while bi + 4 <= acts.len() {
4284                        let xs = [
4285                            acts[bi].xq.as_slice(),
4286                            acts[bi + 1].xq.as_slice(),
4287                            acts[bi + 2].xq.as_slice(),
4288                            acts[bi + 3].xq.as_slice(),
4289                        ];
4290                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
4291                        for k in 0..4 {
4292                            let act = &acts[bi + k];
4293                            let mut acc = d[k] * act.sx;
4294                            for &(j, xv) in &act.outliers {
4295                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4296                                acc += w * sc * xv;
4297                            }
4298                            // SAFETY: disjoint (bi, r) cells per worker.
4299                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4300                        }
4301                        bi += 4;
4302                    }
4303                }
4304                #[cfg(target_arch = "x86_64")]
4305                if blocked_ok {
4306                    while bi + 4 <= acts.len() {
4307                        let xs = [
4308                            acts[bi].xq.as_slice(),
4309                            acts[bi + 1].xq.as_slice(),
4310                            acts[bi + 2].xq.as_slice(),
4311                            acts[bi + 3].xq.as_slice(),
4312                        ];
4313                        let d = unsafe {
4314                            if vnni_tiles_enabled() {
4315                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
4316                            } else {
4317                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
4318                            }
4319                        };
4320                        for k in 0..4 {
4321                            let act = &acts[bi + k];
4322                            let mut acc = d[k] * act.sx;
4323                            for &(j, xv) in &act.outliers {
4324                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4325                                acc += w * sc * xv;
4326                            }
4327                            // SAFETY: disjoint (bi, r) cells per worker.
4328                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4329                        }
4330                        bi += 4;
4331                    }
4332                }
4333                let _ = blocked_ok;
4334                while bi < acts.len() {
4335                    let act = &acts[bi];
4336                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4337                    for &(j, xv) in &act.outliers {
4338                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
4339                        acc += w * s * xv;
4340                    }
4341                    // SAFETY: disjoint (bi, r) cells per worker range.
4342                    unsafe { *out_addr.at(bi * rows + r) = acc };
4343                    bi += 1;
4344                }
4345            }
4346        };
4347        dispatch_rows(pool, rows, &run);
4348        return;
4349    }
4350    let run = move |start: usize, end: usize| {
4351        for r in start..end {
4352            for bi in 0..b {
4353                let x = &xs_all[bi * cols..(bi + 1) * cols];
4354                // SAFETY: disjoint (bi, r) cells per worker range.
4355                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
4356            }
4357        }
4358    };
4359    dispatch_rows(pool, rows, &run);
4360}
4361
4362// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
4363// 32-group tile. The kernel family mirrors q4_tiled: one sequential
4364// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
4365// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
4366
4367/// Per-32-group sums of the quantized activation — the ±1 identity's
4368/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
4369/// matvec and reused by every row.
4370fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
4371    (0..gpr)
4372        .map(|gi| {
4373            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
4374                .iter()
4375                .map(|&v| v as i32)
4376                .sum()
4377        })
4378        .collect()
4379}
4380
4381/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
4382/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
4383/// x86 pass).
4384#[inline]
4385#[allow(unreachable_code)]
4386/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
4387/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
4388/// masked activation sums through maddubs(1, x&mask), and
4389/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
4390#[cfg(target_arch = "x86_64")]
4391#[target_feature(enable = "avx2")]
4392unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4393    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4394    unsafe {
4395        use core::arch::x86_64::*;
4396        // Byte j of the mask must replicate bits-byte j/8.
4397        let expand = _mm256_setr_epi8(
4398            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,
4399            3, 3, 3,
4400        );
4401        let bitsel = _mm256_setr_epi8(
4402            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4403            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4404        );
4405        let ones8 = _mm256_set1_epi8(1);
4406        let ones16 = _mm256_set1_epi16(1);
4407        let mut acc = 0f32;
4408        for gi in 0..gpr {
4409            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4410            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4411            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4412            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4413            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4414            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4415            let sel = _mm256_and_si256(x, mask);
4416            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
4417            let p16 = _mm256_maddubs_epi16(ones8, sel);
4418            let d32 = _mm256_madd_epi16(p16, ones16);
4419            let hi128 = _mm256_extracti128_si256::<1>(d32);
4420            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4421            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4422            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4423            let msum = _mm_cvtsi128_si32(s32);
4424            // The and-select keeps x UN-negated (unlike ARM's −1-mask
4425            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
4426            let d = 2 * msum - gsum[gi];
4427            acc += d as f32 * s;
4428        }
4429        acc
4430    }
4431}
4432
4433/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
4434/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
4435#[cfg(target_arch = "x86_64")]
4436#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4437unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4438    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4439    unsafe {
4440        use core::arch::x86_64::*;
4441        let expand = _mm256_setr_epi8(
4442            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,
4443            3, 3, 3,
4444        );
4445        let bitsel = _mm256_setr_epi8(
4446            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4447            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4448        );
4449        let ones8 = _mm256_set1_epi8(1);
4450        let mut acc = 0f32;
4451        for gi in 0..gpr {
4452            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4453            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4454            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4455            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4456            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4457            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4458            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4459            let d = 2 * msum - gsum[gi];
4460            acc += d as f32 * s;
4461        }
4462        acc
4463    }
4464}
4465
4466/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
4467#[cfg(target_arch = "x86_64")]
4468#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4469unsafe fn dot_q1_row_1x4_vnni(
4470    bytes: &[u8],
4471    r: usize,
4472    gpr: usize,
4473    xs: [&[i8]; 4],
4474    gsums: [&[i32]; 4],
4475) -> [f32; 4] {
4476    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4477    unsafe {
4478        use core::arch::x86_64::*;
4479        let expand = _mm256_setr_epi8(
4480            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,
4481            3, 3, 3,
4482        );
4483        let bitsel = _mm256_setr_epi8(
4484            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4485            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4486        );
4487        let ones8 = _mm256_set1_epi8(1);
4488        let mut acc = [0f32; 4];
4489        for gi in 0..gpr {
4490            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4491            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4492            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4493            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4494            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4495            for (k, xq) in xs.iter().enumerate() {
4496                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4497                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4498                let d = 2 * msum - gsums[k][gi];
4499                acc[k] += d as f32 * s;
4500            }
4501        }
4502        acc
4503    }
4504}
4505
4506/// The blocked 1×4 flavor: the expanded bit mask serves four activation
4507/// streams per group (mask build once, four select+reduce chains).
4508#[cfg(target_arch = "x86_64")]
4509#[target_feature(enable = "avx2")]
4510unsafe fn dot_q1_row_1x4_avx2(
4511    bytes: &[u8],
4512    r: usize,
4513    gpr: usize,
4514    xs: [&[i8]; 4],
4515    gsums: [&[i32]; 4],
4516) -> [f32; 4] {
4517    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4518    unsafe {
4519        use core::arch::x86_64::*;
4520        let expand = _mm256_setr_epi8(
4521            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,
4522            3, 3, 3,
4523        );
4524        let bitsel = _mm256_setr_epi8(
4525            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4526            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4527        );
4528        let ones8 = _mm256_set1_epi8(1);
4529        let ones16 = _mm256_set1_epi16(1);
4530        let mut acc = [0f32; 4];
4531        for gi in 0..gpr {
4532            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4533            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4534            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4535            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4536            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4537            for (k, xq) in xs.iter().enumerate() {
4538                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4539                let sel = _mm256_and_si256(x, mask);
4540                let p16 = _mm256_maddubs_epi16(ones8, sel);
4541                let d32 = _mm256_madd_epi16(p16, ones16);
4542                let hi128 = _mm256_extracti128_si256::<1>(d32);
4543                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4544                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4545                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4546                let msum = _mm_cvtsi128_si32(s32);
4547                let d = 2 * msum - gsums[k][gi];
4548                acc[k] += d as f32 * s;
4549            }
4550        }
4551        acc
4552    }
4553}
4554
4555#[allow(unreachable_code)]
4556fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4557    #[cfg(target_arch = "aarch64")]
4558    unsafe {
4559        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
4560    }
4561    #[cfg(target_arch = "x86_64")]
4562    if avx2_enabled() {
4563        unsafe {
4564            if vnni_tiles_enabled() {
4565                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
4566            }
4567            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
4568        }
4569    }
4570    let _ = gsum;
4571    let mut acc = 0f32;
4572    for gi in 0..gpr {
4573        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4574        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4575        let mut d = 0i32;
4576        for (j, &b) in tile[2..].iter().enumerate() {
4577            for k in 0..8 {
4578                let w = ((b >> k) & 1) as i32 * 2 - 1;
4579                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
4580            }
4581        }
4582        acc += d as f32 * s;
4583    }
4584    acc
4585}
4586
4587/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
4588/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
4589/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
4590/// per-group activation sums shared across every row of the matvec.
4591/// Four tiles (128 weights) per iteration: integer dots reduce through
4592/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
4593/// fused f32 multiply-add. Integer math throughout — bit-identical to
4594/// the scalar ±1 reference.
4595#[cfg(target_arch = "aarch64")]
4596#[target_feature(enable = "neon,dotprod")]
4597unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4598    // SAFETY: callers uphold slice-length contracts (6B tile per group,
4599    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
4600    unsafe {
4601        use core::arch::aarch64::*;
4602        use core::arch::asm;
4603        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4604        let m = vld1q_u8(MASKS.as_ptr());
4605        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
4606        macro_rules! tile_dot {
4607            ($t:expr, $x:expr) => {{
4608                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
4609                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
4610                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4611                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4612                let x0 = vld1q_s8($x);
4613                let x1 = vld1q_s8($x.add(16));
4614                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4615                asm!(
4616                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4617                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4618                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4619                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4620                    options(pure, nomem, nostack),
4621                );
4622                vaddq_s32(a0, a1)
4623            }};
4624        }
4625        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
4626        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
4627        // bit-byte across 8 lanes for vtst, and the four scales gather
4628        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
4629        // 4 branchy software f16 conversions per 128 weights (the
4630        // measured load-port wall of this kernel) become 2 vector
4631        // loads + 9 table lookups. Integer math order is unchanged —
4632        // bit-identical results (FCVTL is exact on every f16).
4633        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4634        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4635        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4636        const IW11: [u8; 16] = [
4637            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4638        ];
4639        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4640        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4641        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4642        let isc = vld1_u8(ISC.as_ptr());
4643        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
4644        macro_rules! tile_dot_tbl {
4645            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
4646                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
4647                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
4648                let x0 = vld1q_s8($x);
4649                let x1 = vld1q_s8($x.add(16));
4650                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4651                asm!(
4652                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4653                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4654                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4655                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4656                    options(pure, nomem, nostack),
4657                );
4658                vaddq_s32(a0, a1)
4659            }};
4660        }
4661        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4662        let row_base = r * gpr * Q1_TILE;
4663        let abs_end = bytes.len();
4664        let xp = xq.as_ptr();
4665        let gp = gsum.as_ptr();
4666        let mut accv = vdupq_n_f32(0.0);
4667        let mut gi = 0;
4668        // The second pair load reads 4B past tile gi+3 — stay inside
4669        // the payload slice (only the file's final tiles fall back).
4670        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4671            let t0 = base.add(gi * Q1_TILE);
4672            let ld_a = vld1q_u8(t0);
4673            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4674            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
4675            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
4676            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
4677            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
4678            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
4679            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4680            let g = vld1q_s32(gp.add(gi));
4681            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4682            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4683            let scf: float32x4_t;
4684            asm!(
4685                "fcvtl {o:v}.4s, {i:v}.4h",
4686                o = out(vreg) scf, i = in(vreg) sc16,
4687                options(pure, nomem, nostack),
4688            );
4689            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
4690            gi += 4;
4691        }
4692        let mut acc = vaddvq_f32(accv);
4693        while gi < gpr {
4694            let t = base.add(gi * Q1_TILE);
4695            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4696            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
4697            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
4698            gi += 1;
4699        }
4700        acc
4701    }
4702}
4703
4704/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
4705/// activation streams (prefill amortization — the same idea as the
4706/// AVX2 twin; per stream the group order, fma order and tail match the
4707/// single-row kernel exactly, so batch == matvec bit-for-bit).
4708#[cfg(target_arch = "aarch64")]
4709#[target_feature(enable = "neon,dotprod")]
4710unsafe fn dot_q1_row_1x4_sdot(
4711    bytes: &[u8],
4712    r: usize,
4713    gpr: usize,
4714    xs: [&[i8]; 4],
4715    gs: [&[i32]; 4],
4716) -> [f32; 4] {
4717    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
4718    unsafe {
4719        use core::arch::aarch64::*;
4720        use core::arch::asm;
4721        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4722        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4723        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4724        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4725        const IW11: [u8; 16] = [
4726            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4727        ];
4728        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4729        let m = vld1q_u8(MASKS.as_ptr());
4730        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4731        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4732        let isc = vld1_u8(ISC.as_ptr());
4733        macro_rules! sdot2 {
4734            ($w0:expr, $w1:expr, $x:expr) => {{
4735                let x0 = vld1q_s8($x);
4736                let x1 = vld1q_s8($x.add(16));
4737                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4738                asm!(
4739                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4740                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4741                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4742                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
4743                    options(pure, nomem, nostack),
4744                );
4745                vaddq_s32(a0, a1)
4746            }};
4747        }
4748        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4749        let row_base = r * gpr * Q1_TILE;
4750        let abs_end = bytes.len();
4751        let mut accv = [vdupq_n_f32(0.0); 4];
4752        let mut gi = 0;
4753        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4754            let t0 = base.add(gi * Q1_TILE);
4755            let ld_a = vld1q_u8(t0);
4756            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4757            // Unpack ONCE — eight ±mask vectors serve all four streams.
4758            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
4759            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
4760            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
4761            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
4762            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
4763            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
4764            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
4765            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
4766            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4767            let scf: float32x4_t;
4768            asm!(
4769                "fcvtl {o:v}.4s, {i:v}.4h",
4770                o = out(vreg) scf, i = in(vreg) sc16,
4771                options(pure, nomem, nostack),
4772            );
4773            for k in 0..4 {
4774                let xp = xs[k].as_ptr();
4775                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
4776                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
4777                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
4778                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
4779                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4780                let g = vld1q_s32(gs[k].as_ptr().add(gi));
4781                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4782                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
4783            }
4784            gi += 4;
4785        }
4786        let mut acc = [
4787            vaddvq_f32(accv[0]),
4788            vaddvq_f32(accv[1]),
4789            vaddvq_f32(accv[2]),
4790            vaddvq_f32(accv[3]),
4791        ];
4792        while gi < gpr {
4793            let t = base.add(gi * Q1_TILE);
4794            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4795            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
4796            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
4797            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4798            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4799            for k in 0..4 {
4800                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
4801                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
4802            }
4803            gi += 1;
4804        }
4805        acc
4806    }
4807}
4808
4809/// (weight ±1, scale) of one q1 element — the exact outlier term.
4810#[inline]
4811fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4812    let gi = j / GROUP_SIZE;
4813    let k = j % GROUP_SIZE;
4814    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4815    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4816    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
4817    ((bit as i32 * 2 - 1) as f32, s)
4818}
4819
4820/// Exact scalar q1 row (CMF_SDOT=0 contract).
4821#[inline]
4822fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4823    let mut acc = 0f32;
4824    for gi in 0..gpr {
4825        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4826        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4827        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4828        let mut ga = 0f32;
4829        for (j, &b) in tile[2..].iter().enumerate() {
4830            for k in 0..8 {
4831                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
4832            }
4833        }
4834        acc += ga * s;
4835    }
4836    acc
4837}
4838
4839/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
4840/// extracted so multi-matrix jobs drive the same kernel).
4841#[allow(clippy::too_many_arguments)]
4842fn q1_range_a8w8(
4843    bytes: &[u8],
4844    gpr: usize,
4845    act: &SplitAct,
4846    gsum: &[i32],
4847    out: SendMut,
4848    start: usize,
4849    end: usize,
4850) {
4851    for r in start..end {
4852        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
4853        for &(j, xv) in &act.outliers {
4854            let (w, s) = q1_outlier(bytes, r, gpr, j);
4855            acc += w * s * xv;
4856        }
4857        // SAFETY: disjoint row ranges per worker.
4858        unsafe { *out.at(r) = acc };
4859    }
4860}
4861
4862/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
4863fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
4864    for r in start..end {
4865        // SAFETY: disjoint row ranges per worker.
4866        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
4867    }
4868}
4869
4870/// q1t per-row overlay locator. After the base (`base_len`) come
4871/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
4872/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
4873/// `(row_ptr offset, entries offset, present)`.
4874fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
4875    let entries = base_len + (rows + 1) * 4;
4876    (base_len, entries, entries <= bytes.len())
4877}
4878
4879/// Read `row_ptr[r]` from the overlay's prefix-sum table.
4880#[inline]
4881fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
4882    let o = rp_off + r * 4;
4883    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
4884}
4885
4886/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
4887/// decoding a q1t code is a table load, not the base-3 divide/modulo per
4888/// weight (division is ~20–40× the cost of a load). Built at compile time.
4889const SIGN5: [[f32; 5]; 256] = {
4890    let mut lut = [[0.0f32; 5]; 256];
4891    let pow3 = [1u16, 3, 9, 27, 81];
4892    let mut byte = 0usize;
4893    while byte < 256 {
4894        let mut i = 0usize;
4895        while i < 5 {
4896            let code = (byte as u16 / pow3[i]) % 3;
4897            lut[byte][i] = if code == 1 {
4898                1.0
4899            } else if code == 2 {
4900                -1.0
4901            } else {
4902                0.0
4903            };
4904            i += 1;
4905        }
4906        byte += 1;
4907    }
4908    lut
4909};
4910
4911/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
4912const SIGN5_I8: [[i8; 5]; 256] = {
4913    let mut lut = [[0i8; 5]; 256];
4914    let pow3 = [1u16, 3, 9, 27, 81];
4915    let mut byte = 0usize;
4916    while byte < 256 {
4917        let mut i = 0usize;
4918        while i < 5 {
4919            let code = (byte as u16 / pow3[i]) % 3;
4920            lut[byte][i] = if code == 1 {
4921                1
4922            } else if code == 2 {
4923                -1
4924            } else {
4925                0
4926            };
4927            i += 1;
4928        }
4929        byte += 1;
4930    }
4931    lut
4932};
4933
4934/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
4935/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
4936/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
4937/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
4938/// buffer is padded to 40). This is the decode/prefill hot inner op.
4939const SIGN5_U64: [u64; 256] = {
4940    let mut lut = [0u64; 256];
4941    let pow3 = [1u16, 3, 9, 27, 81];
4942    let mut byte = 0usize;
4943    while byte < 256 {
4944        let mut v = 0u64;
4945        let mut i = 0usize;
4946        while i < 5 {
4947            let code = (byte as u16 / pow3[i]) % 3;
4948            let s: u8 = if code == 1 {
4949                1
4950            } else if code == 2 {
4951                0xFF
4952            } else {
4953                0
4954            };
4955            v |= (s as u64) << (i * 8);
4956            i += 1;
4957        }
4958        lut[byte] = v;
4959        byte += 1;
4960    }
4961    lut
4962};
4963
4964/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
4965/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
4966/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
4967/// the overlay correction owns that column — no double counting.
4968#[inline]
4969fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
4970    const TILE: usize = cortiq_core::quant::Q1T_TILE;
4971    let off = (r * gpr + j / GROUP_SIZE) * TILE;
4972    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4973    let within = j % GROUP_SIZE;
4974    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
4975}
4976
4977/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
4978/// (integer accumulation is order-independent).
4979#[cfg(target_arch = "aarch64")]
4980#[target_feature(enable = "neon,dotprod")]
4981#[inline]
4982unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
4983    // SAFETY: caller guarantees 32 readable i8 at each pointer.
4984    unsafe {
4985        use core::arch::aarch64::*;
4986        use core::arch::asm;
4987        let w0 = vld1q_s8(w);
4988        let w1 = vld1q_s8(w.add(16));
4989        let x0 = vld1q_s8(x);
4990        let x1 = vld1q_s8(x.add(16));
4991        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4992        asm!(
4993            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4994            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4995            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4996            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4997            options(pure, nomem, nostack),
4998        );
4999        vaddvq_s32(vaddq_s32(a0, a1))
5000    }
5001}
5002
5003/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
5004/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
5005#[cfg(target_arch = "x86_64")]
5006#[target_feature(enable = "avx2")]
5007#[inline]
5008unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
5009    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5010    unsafe {
5011        use core::arch::x86_64::*;
5012        let wv = _mm256_loadu_si256(w as *const __m256i);
5013        let xv = _mm256_loadu_si256(x as *const __m256i);
5014        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5015        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
5016        let hi128 = _mm256_extracti128_si256::<1>(d);
5017        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5018        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5019        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5020        _mm_cvtsi128_si32(s32)
5021    }
5022}
5023
5024/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
5025/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
5026/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
5027/// overwritten by the next; the final 6 padding bytes are unused by the dot.
5028#[inline]
5029fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
5030    debug_assert!(dst.len() >= 40);
5031    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
5032    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
5033    unsafe {
5034        let p = dst.as_mut_ptr();
5035        for bi in 0..7 {
5036            core::ptr::write_unaligned(
5037                p.add(bi * 5) as *mut u64,
5038                SIGN5_U64[*codes.add(bi) as usize],
5039            );
5040        }
5041    }
5042}
5043
5044/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
5045/// row's signs are unpacked once and dotted against every batch input).
5046/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
5047/// reachable; the scalar arm is a non-SIMD-arch fallback.
5048#[inline]
5049fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
5050    #[cfg(target_arch = "aarch64")]
5051    unsafe {
5052        return sdot32_i8(w, x);
5053    }
5054    #[cfg(target_arch = "x86_64")]
5055    unsafe {
5056        return i8dot32_avx2(w, x);
5057    }
5058    #[allow(unreachable_code)]
5059    unsafe {
5060        let mut s = 0i32;
5061        for k in 0..GROUP_SIZE {
5062            s += *w.add(k) as i32 * *x.add(k) as i32;
5063        }
5064        s
5065    }
5066}
5067
5068#[inline]
5069unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
5070    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
5071        (
5072            SIGN5_U64[*codes as usize],
5073            SIGN5_U64[*codes.add(1) as usize],
5074            SIGN5_U64[*codes.add(2) as usize],
5075            SIGN5_U64[*codes.add(3) as usize],
5076            SIGN5_U64[*codes.add(4) as usize],
5077            SIGN5_U64[*codes.add(5) as usize],
5078            SIGN5_U64[*codes.add(6) as usize],
5079        )
5080    };
5081
5082    let u0 = s0 | (s1 << 40);
5083    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
5084    let u2 = (s3 >> 8) | (s4 << 32);
5085    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
5086
5087    (u0, u1, u2, u3)
5088}
5089
5090/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
5091/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
5092/// ARM SDOT.
5093#[cfg(target_arch = "aarch64")]
5094#[target_feature(enable = "neon,dotprod")]
5095unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5096    use core::arch::aarch64::*;
5097    use core::arch::asm;
5098    unsafe {
5099        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5100        let mut acc = 0f32;
5101        let bytes_ptr = bytes.as_ptr();
5102        let xq_ptr = xq.as_ptr();
5103        let row_off = r * gpr * TILE;
5104
5105        let gpr2 = gpr & !1;
5106        let mut gi = 0;
5107        while gi < gpr2 {
5108            let off0 = row_off + gi * TILE;
5109            let off1 = off0 + TILE;
5110            let s0 = f16_to_f32(u16::from_le_bytes([
5111                *bytes_ptr.add(off0),
5112                *bytes_ptr.add(off0 + 1),
5113            ]));
5114            let s1 = f16_to_f32(u16::from_le_bytes([
5115                *bytes_ptr.add(off1),
5116                *bytes_ptr.add(off1 + 1),
5117            ]));
5118
5119            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5120            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5121
5122            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5123            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5124            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5125            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5126
5127            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5128            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5129            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
5130            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
5131
5132            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
5133            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5134            asm!(
5135                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
5136                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
5137                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
5138                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
5139                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
5140                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
5141                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
5142                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
5143                options(pure, nomem, nostack),
5144            );
5145            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
5146            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
5147            acc += d0 as f32 * s0 + d1 as f32 * s1;
5148            gi += 2;
5149        }
5150
5151        if gi < gpr {
5152            let off = row_off + gi * TILE;
5153            let s = f16_to_f32(u16::from_le_bytes([
5154                *bytes_ptr.add(off),
5155                *bytes_ptr.add(off + 1),
5156            ]));
5157            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5158            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5159            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5160            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5161            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5162            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5163            asm!(
5164                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5165                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5166                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5167                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5168                options(pure, nomem, nostack),
5169            );
5170            let d = vaddvq_s32(vaddq_s32(a0, a1));
5171            acc += d as f32 * s;
5172        }
5173        acc
5174    }
5175}
5176
5177/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
5178#[cfg(target_arch = "x86_64")]
5179#[target_feature(enable = "avx2")]
5180unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5181    use core::arch::x86_64::*;
5182    unsafe {
5183        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5184        let mut acc = 0f32;
5185        let bytes_ptr = bytes.as_ptr();
5186        let xq_ptr = xq.as_ptr();
5187        let row_off = r * gpr * TILE;
5188
5189        let ones = _mm256_set1_epi16(1);
5190        for gi in 0..gpr {
5191            let off = row_off + gi * TILE;
5192            let s = f16_to_f32(u16::from_le_bytes([
5193                *bytes_ptr.add(off),
5194                *bytes_ptr.add(off + 1),
5195            ]));
5196            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5197            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5198            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5199            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5200            let d256 = _mm256_madd_epi16(p16, ones);
5201            let d128 = _mm_add_epi32(
5202                _mm256_castsi256_si128(d256),
5203                _mm256_extracti128_si256(d256, 1),
5204            );
5205            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
5206            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
5207            acc += d32 as f32 * s;
5208        }
5209        acc
5210    }
5211}
5212
5213/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
5214#[cfg(target_arch = "x86_64")]
5215#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5216unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5217    use core::arch::x86_64::*;
5218    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
5219    unsafe {
5220        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5221        let mut acc = 0f32;
5222        let bytes_ptr = bytes.as_ptr();
5223        let xq_ptr = xq.as_ptr();
5224        let row_off = r * gpr * TILE;
5225        for gi in 0..gpr {
5226            let off = row_off + gi * TILE;
5227            let s = f16_to_f32(u16::from_le_bytes([
5228                *bytes_ptr.add(off),
5229                *bytes_ptr.add(off + 1),
5230            ]));
5231            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5232            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5233            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5234            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5235            acc += d as f32 * s;
5236        }
5237        acc
5238    }
5239}
5240
5241/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
5242/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
5243/// reachable.
5244#[inline]
5245fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5246    #[cfg(target_arch = "aarch64")]
5247    unsafe {
5248        return q1t_dot_row_sdot(bytes, r, gpr, xq);
5249    }
5250    #[cfg(target_arch = "x86_64")]
5251    unsafe {
5252        if vnni_tiles_enabled() {
5253            return q1t_dot_row_vnni(bytes, r, gpr, xq);
5254        }
5255        return q1t_dot_row_avx2(bytes, r, gpr, xq);
5256    }
5257    #[allow(unreachable_code)]
5258    {
5259        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5260        let mut acc = 0f32;
5261        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
5262        for gi in 0..gpr {
5263            let off = (r * gpr + gi) * TILE;
5264            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5265            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
5266            let mut d = 0i32;
5267            for k in 0..GROUP_SIZE {
5268                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
5269            }
5270            acc += d as f32 * s;
5271        }
5272        acc
5273    }
5274}
5275
5276/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
5277/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
5278/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
5279/// base contributes nothing there and this is a plain `value·x`, not
5280/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
5281/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
5282fn q1t_row_outlier_correction(
5283    bytes: &[u8],
5284    r: usize,
5285    rp_off: usize,
5286    entries_off: usize,
5287    has_ov: bool,
5288    x: &[f32],
5289) -> f32 {
5290    if !has_ov {
5291        return 0.0;
5292    }
5293    let (c0, c1) = (
5294        q1t_rowptr(bytes, rp_off, r),
5295        q1t_rowptr(bytes, rp_off, r + 1),
5296    );
5297    let mut corr = 0f32;
5298    for p in c0..c1 {
5299        let e = entries_off + p * 4;
5300        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5301        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5302        corr += val * x[col];
5303    }
5304    corr
5305}
5306
5307/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
5308/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
5309/// Used by the batched (prefill) path where the decode amortizes over the batch.
5310fn q1t_dequant_row(
5311    bytes: &[u8],
5312    r: usize,
5313    gpr: usize,
5314    rp_off: usize,
5315    entries_off: usize,
5316    has_ov: bool,
5317    buf: &mut [f32],
5318) {
5319    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5320    for g in 0..gpr {
5321        let off = (r * gpr + g) * TILE;
5322        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5323        let codes = &bytes[off + 2..off + TILE];
5324        let bc = g * GROUP_SIZE;
5325        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
5326        for bi in 0..6 {
5327            let lut = &SIGN5[codes[bi] as usize];
5328            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
5329            for i in 0..5 {
5330                d[i] = lut[i] * s;
5331            }
5332        }
5333        let lut = &SIGN5[codes[6] as usize];
5334        buf[bc + 30] = lut[0] * s;
5335        buf[bc + 31] = lut[1] * s;
5336    }
5337    if !has_ov {
5338        return;
5339    }
5340    let (c0, c1) = (
5341        q1t_rowptr(bytes, rp_off, r),
5342        q1t_rowptr(bytes, rp_off, r + 1),
5343    );
5344    for p in c0..c1 {
5345        let e = entries_off + p * 4;
5346        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5347        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5348    }
5349}
5350
5351/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
5352/// computes the ternary base; the overlay stays on the CPU — its entries are
5353/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
5354fn q1t_add_overlay(
5355    bytes: &[u8],
5356    x: &[f32],
5357    rows: usize,
5358    cols: usize,
5359    out: &mut [f32],
5360    pool: Option<&Pool>,
5361) {
5362    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5363    let gpr = cols / GROUP_SIZE;
5364    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5365    if !has_ov {
5366        return;
5367    }
5368    let out_addr = SendMut(out.as_mut_ptr());
5369    let run = move |start: usize, end: usize| {
5370        for r in start..end {
5371            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5372            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
5373            unsafe { *out_addr.at(r) += corr };
5374        }
5375    };
5376    dispatch_rows(pool, rows, &run);
5377}
5378
5379/// Q1T row range via the A8W8 int8 path — shared activation split,
5380/// per-row: base SDOT dot + outlier correction + overlay.
5381#[allow(clippy::too_many_arguments)]
5382fn q1t_range_a8w8(
5383    bytes: &[u8],
5384    gpr: usize,
5385    rp_off: usize,
5386    ent_off: usize,
5387    has_ov: bool,
5388    act: &SplitAct,
5389    x: &[f32],
5390    out: SendMut,
5391    start: usize,
5392    end: usize,
5393) {
5394    for r in start..end {
5395        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5396        for &(j, xv) in &act.outliers {
5397            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5398        }
5399        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5400        // SAFETY: disjoint row ranges per worker.
5401        unsafe { *out.at(r) = acc };
5402    }
5403}
5404
5405/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
5406/// dispatch when a8w8 is unavailable.
5407#[allow(clippy::too_many_arguments)]
5408fn q1t_range_f32_batch(
5409    bytes: &[u8],
5410    gpr: usize,
5411    rp_off: usize,
5412    ent_off: usize,
5413    has_ov: bool,
5414    x: &[f32],
5415    out: SendMut,
5416    start: usize,
5417    end: usize,
5418) {
5419    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5420    let mut sg = [0f32; GROUP_SIZE];
5421    for r in start..end {
5422        let mut acc = 0f32;
5423        for g in 0..gpr {
5424            let off = (r * gpr + g) * TILE;
5425            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5426            let codes = &bytes[off + 2..off + TILE];
5427            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5428            for bi in 0..6 {
5429                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5430            }
5431            let lut = &SIGN5[codes[6] as usize];
5432            sg[30] = lut[0];
5433            sg[31] = lut[1];
5434            let mut gsum = 0f32;
5435            for k in 0..GROUP_SIZE {
5436                gsum += sg[k] * xg[k];
5437            }
5438            acc += s * gsum;
5439        }
5440        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5441        // SAFETY: disjoint row ranges per worker.
5442        unsafe { *out.at(r) = acc };
5443    }
5444}
5445
5446/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
5447/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
5448/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
5449fn q1t_matvec(
5450    bytes: &[u8],
5451    x: &[f32],
5452    rows: usize,
5453    cols: usize,
5454    out: &mut [f32],
5455    pool: Option<&Pool>,
5456) {
5457    debug_assert_eq!(out.len(), rows);
5458    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5459    let gpr = cols / GROUP_SIZE;
5460    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5461    let out_addr = SendMut(out.as_mut_ptr());
5462    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
5463    // (`split_act`), activation outliers added back exactly in f32, weight
5464    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
5465    if a8w8_enabled() {
5466        let act = split_act(x);
5467        let act = &act;
5468        let run = move |start: usize, end: usize| {
5469            for r in start..end {
5470                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5471                for &(j, xv) in &act.outliers {
5472                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5473                }
5474                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5475                // SAFETY: disjoint row ranges per worker.
5476                unsafe { *out_addr.at(r) = acc };
5477            }
5478        };
5479        dispatch_rows(pool, rows, &run);
5480        return;
5481    }
5482    let run = move |start: usize, end: usize| {
5483        // Per-group signs, unpacked contiguously so the dot below is a clean
5484        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
5485        // 5-values-per-byte base-3 layout won't SIMD in place.
5486        let mut sg = [0f32; GROUP_SIZE];
5487        for r in start..end {
5488            let mut acc = 0f32;
5489            for g in 0..gpr {
5490                let off = (r * gpr + g) * TILE;
5491                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5492                let codes = &bytes[off + 2..off + TILE];
5493                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5494                for bi in 0..6 {
5495                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5496                }
5497                let lut = &SIGN5[codes[6] as usize];
5498                sg[30] = lut[0];
5499                sg[31] = lut[1];
5500                let mut gsum = 0f32;
5501                for k in 0..GROUP_SIZE {
5502                    gsum += sg[k] * xg[k];
5503                }
5504                acc += s * gsum;
5505            }
5506            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5507            unsafe { *out_addr.at(r) = acc };
5508        }
5509    };
5510    dispatch_rows(pool, rows, &run);
5511}
5512
5513/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
5514/// ternary codes serves BOTH activation streams (the unpack chain is
5515/// the dominant per-row cost — MTP verify pairs paid it twice). Per
5516/// stream the group order and f32 accumulation match the single-row
5517/// kernel exactly, so pair == 2×matvec bit-for-bit.
5518#[cfg(target_arch = "aarch64")]
5519#[target_feature(enable = "neon,dotprod")]
5520unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
5521    use core::arch::aarch64::*;
5522    use core::arch::asm;
5523    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
5524    unsafe {
5525        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5526        let bytes_ptr = bytes.as_ptr();
5527        let row_off = r * gpr * TILE;
5528        let xp = [xa.as_ptr(), xb.as_ptr()];
5529        let mut acc = [0f32; 2];
5530        macro_rules! sdot2 {
5531            ($w0:expr, $w1:expr, $x:expr) => {{
5532                let x0 = vld1q_s8($x);
5533                let x1 = vld1q_s8($x.add(16));
5534                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5535                asm!(
5536                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5537                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5538                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5539                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5540                    options(pure, nomem, nostack),
5541                );
5542                vaddvq_s32(vaddq_s32(a0, a1))
5543            }};
5544        }
5545        let gpr2 = gpr & !1;
5546        let mut gi = 0;
5547        while gi < gpr2 {
5548            let off0 = row_off + gi * TILE;
5549            let off1 = off0 + TILE;
5550            let s0 = f16_to_f32(u16::from_le_bytes([
5551                *bytes_ptr.add(off0),
5552                *bytes_ptr.add(off0 + 1),
5553            ]));
5554            let s1 = f16_to_f32(u16::from_le_bytes([
5555                *bytes_ptr.add(off1),
5556                *bytes_ptr.add(off1 + 1),
5557            ]));
5558            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5559            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5560            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5561            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5562            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5563            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5564            for k in 0..2 {
5565                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
5566                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
5567                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
5568            }
5569            gi += 2;
5570        }
5571        if gi < gpr {
5572            let off = row_off + gi * TILE;
5573            let s = f16_to_f32(u16::from_le_bytes([
5574                *bytes_ptr.add(off),
5575                *bytes_ptr.add(off + 1),
5576            ]));
5577            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5578            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5579            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5580            for k in 0..2 {
5581                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
5582                acc[k] += d as f32 * s;
5583            }
5584        }
5585        acc
5586    }
5587}
5588
5589/// Fused Q1T pair matvec: ONE pass over the rows serves both
5590/// activation streams — on ARM the ternary register unpack happens
5591/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
5592/// rides the row's L1-warm tile bytes. Per stream the math matches
5593/// `q1t_matvec` exactly.
5594fn q1t_matvec2(
5595    bytes: &[u8],
5596    x1: &[f32],
5597    x2: &[f32],
5598    rows: usize,
5599    cols: usize,
5600    o1: &mut [f32],
5601    o2: &mut [f32],
5602    pool: Option<&Pool>,
5603) {
5604    debug_assert_eq!(o1.len(), rows);
5605    debug_assert_eq!(o2.len(), rows);
5606    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5607    let gpr = cols / GROUP_SIZE;
5608    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5609    let out1 = SendMut(o1.as_mut_ptr());
5610    let out2 = SendMut(o2.as_mut_ptr());
5611    if a8w8_enabled() {
5612        let a1 = split_act(x1);
5613        let a2 = split_act(x2);
5614        let (a1, a2) = (&a1, &a2);
5615        let run = move |start: usize, end: usize| {
5616            for r in start..end {
5617                #[cfg(target_arch = "aarch64")]
5618                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
5619                // target features are present.
5620                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
5621                #[cfg(not(target_arch = "aarch64"))]
5622                let ds = [
5623                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
5624                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
5625                ];
5626                let mut acc1 = ds[0] * a1.sx;
5627                for &(j, xv) in &a1.outliers {
5628                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
5629                }
5630                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
5631                let mut acc2 = ds[1] * a2.sx;
5632                for &(j, xv) in &a2.outliers {
5633                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
5634                }
5635                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
5636                // SAFETY: disjoint row ranges per worker.
5637                unsafe {
5638                    *out1.at(r) = acc1;
5639                    *out2.at(r) = acc2;
5640                }
5641            }
5642        };
5643        dispatch_rows(pool, rows, &run);
5644        return;
5645    }
5646    let run = move |start: usize, end: usize| {
5647        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
5648        // dot both streams — same op order per stream as `q1t_matvec`.
5649        let mut sg = [0f32; GROUP_SIZE];
5650        for r in start..end {
5651            let mut acc1 = 0f32;
5652            let mut acc2 = 0f32;
5653            for g in 0..gpr {
5654                let off = (r * gpr + g) * TILE;
5655                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5656                let codes = &bytes[off + 2..off + TILE];
5657                for bi in 0..6 {
5658                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5659                }
5660                let lut = &SIGN5[codes[6] as usize];
5661                sg[30] = lut[0];
5662                sg[31] = lut[1];
5663                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5664                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5665                let mut gsum1 = 0f32;
5666                for k in 0..GROUP_SIZE {
5667                    gsum1 += sg[k] * xg1[k];
5668                }
5669                acc1 += s * gsum1;
5670                let mut gsum2 = 0f32;
5671                for k in 0..GROUP_SIZE {
5672                    gsum2 += sg[k] * xg2[k];
5673                }
5674                acc2 += s * gsum2;
5675            }
5676            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
5677            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
5678            // SAFETY: disjoint row ranges per worker.
5679            unsafe {
5680                *out1.at(r) = acc1;
5681                *out2.at(r) = acc2;
5682            }
5683        }
5684    };
5685    dispatch_rows(pool, rows, &run);
5686}
5687
5688/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
5689/// batch against it (amortizes the per-row decode).
5690fn q1t_matmat(
5691    bytes: &[u8],
5692    xs: &[f32],
5693    b: usize,
5694    rows: usize,
5695    cols: usize,
5696    out: &mut [f32],
5697    pool: Option<&Pool>,
5698) {
5699    debug_assert_eq!(out.len(), b * rows);
5700    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5701    let gpr = cols / GROUP_SIZE;
5702    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5703    let out_addr = SendMut(out.as_mut_ptr());
5704    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
5705    // each weight row's signs to i8 ONCE, then int8-dot against every input —
5706    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
5707    if a8w8_enabled() {
5708        let acts: Vec<SplitAct> = (0..b)
5709            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
5710            .collect();
5711        let acts = &acts;
5712        let run = move |start: usize, end: usize| {
5713            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
5714            let mut sc = vec![0f32; gpr]; // per-group scales
5715            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
5716            for r in start..end {
5717                for g in 0..gpr {
5718                    let off = (r * gpr + g) * TILE;
5719                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5720                    q1t_unpack_group_i8(
5721                        bytes.as_ptr().wrapping_add(off + 2),
5722                        &mut sg[g * GROUP_SIZE..],
5723                    );
5724                }
5725                for bi in 0..b {
5726                    let act = &acts[bi];
5727                    let mut isum = 0f32;
5728                    for g in 0..gpr {
5729                        let d = q1t_i8dot32(
5730                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
5731                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
5732                        );
5733                        isum += d as f32 * sc[g];
5734                    }
5735                    let mut acc = isum * act.sx;
5736                    for &(j, xv) in &act.outliers {
5737                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5738                    }
5739                    accs[bi] = acc;
5740                }
5741                // Overlay ONCE per row for the whole batch: read each (col, val)
5742                // from mmap a single time (was b× — the re-read dominated prefill)
5743                // and fan it out over the batch via the cached inputs.
5744                if has_ov {
5745                    let (c0, c1) = (
5746                        q1t_rowptr(bytes, rp_off, r),
5747                        q1t_rowptr(bytes, rp_off, r + 1),
5748                    );
5749                    for p in c0..c1 {
5750                        let e = ent_off + p * 4;
5751                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5752                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5753                        for bi in 0..b {
5754                            accs[bi] += val * xs[bi * cols + col];
5755                        }
5756                    }
5757                }
5758                for bi in 0..b {
5759                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
5760                }
5761            }
5762        };
5763        dispatch_rows(pool, rows, &run);
5764        return;
5765    }
5766    let run = move |start: usize, end: usize| {
5767        let mut buf = vec![0f32; cols];
5768        for r in start..end {
5769            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
5770            for bi in 0..b {
5771                let xr = &xs[bi * cols..(bi + 1) * cols];
5772                let mut acc = 0f32;
5773                for j in 0..cols {
5774                    acc += buf[j] * xr[j];
5775                }
5776                unsafe { *out_addr.at(bi * rows + r) = acc };
5777            }
5778        }
5779    };
5780    dispatch_rows(pool, rows, &run);
5781}
5782
5783fn q1_matvec(
5784    bytes: &[u8],
5785    x: &[f32],
5786    rows: usize,
5787    cols: usize,
5788    out: &mut [f32],
5789    pool: Option<&Pool>,
5790) {
5791    debug_assert_eq!(out.len(), rows);
5792    let gpr = cols / GROUP_SIZE;
5793    let out_addr = SendMut(out.as_mut_ptr());
5794    if a8w8_enabled() {
5795        let act = split_act(x);
5796        let gsum = q1_group_sums(&act.xq, gpr);
5797        let (act, gsum) = (&act, &gsum);
5798        let run = move |start: usize, end: usize| {
5799            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
5800        };
5801        dispatch_rows(pool, rows, &run);
5802        return;
5803    }
5804    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
5805    dispatch_rows(pool, rows, &run);
5806}
5807
5808/// Fused two-input q1 matvec (weights read once per pair).
5809#[allow(clippy::too_many_arguments)]
5810fn q1_matvec2(
5811    bytes: &[u8],
5812    x1: &[f32],
5813    x2: &[f32],
5814    rows: usize,
5815    cols: usize,
5816    o1: &mut [f32],
5817    o2: &mut [f32],
5818    pool: Option<&Pool>,
5819) {
5820    let gpr = cols / GROUP_SIZE;
5821    let p1 = SendMut(o1.as_mut_ptr());
5822    let p2 = SendMut(o2.as_mut_ptr());
5823    if a8w8_enabled() {
5824        let a1 = split_act(x1);
5825        let a2 = split_act(x2);
5826        let g1 = q1_group_sums(&a1.xq, gpr);
5827        let g2 = q1_group_sums(&a2.xq, gpr);
5828        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
5829        let run = move |start: usize, end: usize| {
5830            for r in start..end {
5831                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
5832                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
5833                for &(j, xv) in &a1.outliers {
5834                    let (w, s) = q1_outlier(bytes, r, gpr, j);
5835                    v1 += w * s * xv;
5836                }
5837                for &(j, xv) in &a2.outliers {
5838                    let (w, s) = q1_outlier(bytes, r, gpr, j);
5839                    v2 += w * s * xv;
5840                }
5841                // SAFETY: disjoint row ranges per worker.
5842                unsafe {
5843                    *p1.at(r) = v1;
5844                    *p2.at(r) = v2;
5845                }
5846            }
5847        };
5848        dispatch_rows(pool, rows, &run);
5849        return;
5850    }
5851    let run = move |start: usize, end: usize| {
5852        for r in start..end {
5853            // SAFETY: disjoint row ranges per worker.
5854            unsafe {
5855                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
5856                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
5857            }
5858        }
5859    };
5860    dispatch_rows(pool, rows, &run);
5861}
5862
5863/// Batched q1 matmat: each row's tiles stream once per microbatch.
5864#[allow(clippy::too_many_arguments)]
5865fn q1_matmat(
5866    bytes: &[u8],
5867    xs_all: &[f32],
5868    b: usize,
5869    rows: usize,
5870    cols: usize,
5871    out: &mut [f32],
5872    pool: Option<&Pool>,
5873) {
5874    debug_assert_eq!(out.len(), b * rows);
5875    let gpr = cols / GROUP_SIZE;
5876    let out_addr = SendMut(out.as_mut_ptr());
5877    if a8w8_enabled() {
5878        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
5879            .map(|bi| {
5880                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
5881                let gsum = q1_group_sums(&act.xq, gpr);
5882                (act, gsum)
5883            })
5884            .collect();
5885        let acts = &acts;
5886        #[cfg(target_arch = "x86_64")]
5887        let blocked_ok = avx2_enabled()
5888            && std::env::var("CMF_X86_BLOCKED")
5889                .map(|v| v != "0")
5890                .unwrap_or(true);
5891        #[cfg(target_arch = "aarch64")]
5892        let blocked_ok = sdot_enabled()
5893            && std::env::var("CMF_X86_BLOCKED")
5894                .map(|v| v != "0")
5895                .unwrap_or(true);
5896        let run = move |start: usize, end: usize| {
5897            for r in start..end {
5898                let mut bi = 0usize;
5899                // Blocked 1×4: the unpacked bit mask serves four
5900                // activation streams per group.
5901                #[cfg(target_arch = "aarch64")]
5902                if blocked_ok {
5903                    while bi + 4 <= acts.len() {
5904                        let xs = [
5905                            acts[bi].0.xq.as_slice(),
5906                            acts[bi + 1].0.xq.as_slice(),
5907                            acts[bi + 2].0.xq.as_slice(),
5908                            acts[bi + 3].0.xq.as_slice(),
5909                        ];
5910                        let gs = [
5911                            acts[bi].1.as_slice(),
5912                            acts[bi + 1].1.as_slice(),
5913                            acts[bi + 2].1.as_slice(),
5914                            acts[bi + 3].1.as_slice(),
5915                        ];
5916                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
5917                        for k in 0..4 {
5918                            let (act, _) = &acts[bi + k];
5919                            let mut acc = d[k] * act.sx;
5920                            for &(j, xv) in &act.outliers {
5921                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
5922                                acc += w * sc * xv;
5923                            }
5924                            // SAFETY: disjoint (bi, r) cells per worker.
5925                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5926                        }
5927                        bi += 4;
5928                    }
5929                }
5930                #[cfg(target_arch = "x86_64")]
5931                if blocked_ok {
5932                    while bi + 4 <= acts.len() {
5933                        let xs = [
5934                            acts[bi].0.xq.as_slice(),
5935                            acts[bi + 1].0.xq.as_slice(),
5936                            acts[bi + 2].0.xq.as_slice(),
5937                            acts[bi + 3].0.xq.as_slice(),
5938                        ];
5939                        let gs = [
5940                            acts[bi].1.as_slice(),
5941                            acts[bi + 1].1.as_slice(),
5942                            acts[bi + 2].1.as_slice(),
5943                            acts[bi + 3].1.as_slice(),
5944                        ];
5945                        let d = unsafe {
5946                            if vnni_tiles_enabled() {
5947                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
5948                            } else {
5949                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
5950                            }
5951                        };
5952                        for k in 0..4 {
5953                            let (act, _) = &acts[bi + k];
5954                            let mut acc = d[k] * act.sx;
5955                            for &(j, xv) in &act.outliers {
5956                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
5957                                acc += w * sc * xv;
5958                            }
5959                            // SAFETY: disjoint (bi, r) cells per worker.
5960                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5961                        }
5962                        bi += 4;
5963                    }
5964                }
5965                while bi < acts.len() {
5966                    let (act, gsum) = &acts[bi];
5967                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5968                    for &(j, xv) in &act.outliers {
5969                        let (w, s) = q1_outlier(bytes, r, gpr, j);
5970                        acc += w * s * xv;
5971                    }
5972                    // SAFETY: disjoint (bi, r) cells per worker range.
5973                    unsafe { *out_addr.at(bi * rows + r) = acc };
5974                    bi += 1;
5975                }
5976            }
5977        };
5978        dispatch_rows(pool, rows, &run);
5979        return;
5980    }
5981    let run = move |start: usize, end: usize| {
5982        for r in start..end {
5983            for bi in 0..b {
5984                let x = &xs_all[bi * cols..(bi + 1) * cols];
5985                // SAFETY: disjoint (bi, r) cells per worker range.
5986                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
5987            }
5988        }
5989    };
5990    dispatch_rows(pool, rows, &run);
5991}
5992
5993/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
5994/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
5995/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
5996/// 32-group, exact outlier correction — the same A8W8 contract as q8.
5997/// `CMF_SDOT=0` keeps the exact scalar path.
5998fn q4matvec(
5999    bytes: &[u8],
6000    x: &[f32],
6001    rows: usize,
6002    cols: usize,
6003    out: &mut [f32],
6004    pool: Option<&Pool>,
6005) {
6006    debug_assert_eq!(out.len(), rows);
6007    let (packed, scales) = q4_split(bytes, rows, cols);
6008    let gpr = cols / GROUP_SIZE;
6009    let out_addr = SendMut(out.as_mut_ptr());
6010
6011    if a8w8_enabled() {
6012        let act = split_act(x);
6013        let run = move |start: usize, end: usize| {
6014            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
6015        };
6016        dispatch_rows(pool, rows, &run);
6017        return;
6018    }
6019
6020    let run =
6021        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
6022    dispatch_rows(pool, rows, &run);
6023}
6024
6025/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
6026/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
6027#[inline]
6028#[allow(unreachable_code)]
6029/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
6030/// streams: the 32-byte weight chunk and its abs() load once per group,
6031/// the per-group f16 scale decodes once — four maddubs+reduce chains
6032/// instead of four full (load, abs, dot) rounds.
6033#[cfg(target_arch = "x86_64")]
6034#[target_feature(enable = "avx2")]
6035unsafe fn dot_q4b_row_1x4_avx2(
6036    buf: &[u8],
6037    scales: &[u8],
6038    g0: usize,
6039    gpr: usize,
6040    xs: [&[i8]; 4],
6041) -> [f32; 4] {
6042    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6043    unsafe {
6044        use core::arch::x86_64::*;
6045        let ones = _mm256_set1_epi16(1);
6046        let mut acc = [0f32; 4];
6047        for gi in 0..gpr {
6048            let s = f16_to_f32(u16::from_le_bytes([
6049                scales[(g0 + gi) * 2],
6050                scales[(g0 + gi) * 2 + 1],
6051            ]));
6052            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6053            let aw = _mm256_abs_epi8(w);
6054            for (k, xq) in xs.iter().enumerate() {
6055                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6056                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6057                let d = _mm256_madd_epi16(p16, ones);
6058                let hi128 = _mm256_extracti128_si256::<1>(d);
6059                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6060                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6061                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6062                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
6063            }
6064        }
6065        acc
6066    }
6067}
6068
6069/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
6070#[cfg(target_arch = "x86_64")]
6071#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6072unsafe fn dot_q4b_row_1x4_vnni(
6073    buf: &[u8],
6074    scales: &[u8],
6075    g0: usize,
6076    gpr: usize,
6077    xs: [&[i8]; 4],
6078) -> [f32; 4] {
6079    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6080    unsafe {
6081        use core::arch::x86_64::*;
6082        let mut acc = [0f32; 4];
6083        for gi in 0..gpr {
6084            let s = f16_to_f32(u16::from_le_bytes([
6085                scales[(g0 + gi) * 2],
6086                scales[(g0 + gi) * 2 + 1],
6087            ]));
6088            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6089            let aw = _mm256_abs_epi8(w);
6090            for (k, xq) in xs.iter().enumerate() {
6091                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6092                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6093                acc[k] += d as f32 * s;
6094            }
6095        }
6096        acc
6097    }
6098}
6099
6100/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
6101/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
6102/// accumulation order (the q4_block flavor applies sx once at the end,
6103/// matching ITS single path; the two conventions are historical and
6104/// each blocked leg must mirror its own).
6105#[cfg(target_arch = "x86_64")]
6106#[target_feature(enable = "avx2")]
6107unsafe fn dot_q4b_row_1x4_sx_avx2(
6108    buf: &[u8],
6109    scales: &[u8],
6110    g0: usize,
6111    gpr: usize,
6112    xs: [&[i8]; 4],
6113    sxs: [f32; 4],
6114) -> [f32; 4] {
6115    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6116    unsafe {
6117        use core::arch::x86_64::*;
6118        let ones = _mm256_set1_epi16(1);
6119        let mut acc = [0f32; 4];
6120        for gi in 0..gpr {
6121            let s = f16_to_f32(u16::from_le_bytes([
6122                scales[(g0 + gi) * 2],
6123                scales[(g0 + gi) * 2 + 1],
6124            ]));
6125            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6126            let aw = _mm256_abs_epi8(w);
6127            for (k, xq) in xs.iter().enumerate() {
6128                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6129                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6130                let d = _mm256_madd_epi16(p16, ones);
6131                let hi128 = _mm256_extracti128_si256::<1>(d);
6132                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6133                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6134                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6135                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
6136            }
6137        }
6138        acc
6139    }
6140}
6141
6142/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
6143/// per-group `(d·sx)·s` fold mirrors the vbit single path).
6144#[cfg(target_arch = "x86_64")]
6145#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6146unsafe fn dot_q4b_row_1x4_sx_vnni(
6147    buf: &[u8],
6148    scales: &[u8],
6149    g0: usize,
6150    gpr: usize,
6151    xs: [&[i8]; 4],
6152    sxs: [f32; 4],
6153) -> [f32; 4] {
6154    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6155    unsafe {
6156        use core::arch::x86_64::*;
6157        let mut acc = [0f32; 4];
6158        for gi in 0..gpr {
6159            let s = f16_to_f32(u16::from_le_bytes([
6160                scales[(g0 + gi) * 2],
6161                scales[(g0 + gi) * 2 + 1],
6162            ]));
6163            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6164            let aw = _mm256_abs_epi8(w);
6165            for (k, xq) in xs.iter().enumerate() {
6166                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6167                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6168                acc[k] += (d as f32 * sxs[k]) * s;
6169            }
6170        }
6171        acc
6172    }
6173}
6174
6175#[allow(unreachable_code)]
6176fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6177    #[cfg(target_arch = "aarch64")]
6178    unsafe {
6179        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
6180    }
6181    #[cfg(target_arch = "x86_64")]
6182    unsafe {
6183        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
6184    }
6185    let mut acc = 0f32;
6186    for gi in 0..gpr {
6187        let g = g0 + gi;
6188        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6189        let mut d = 0i32;
6190        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6191            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6192                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6193        }
6194        acc += d as f32 * s;
6195    }
6196    acc
6197}
6198
6199/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
6200#[inline]
6201#[allow(unreachable_code)]
6202fn dot_q4_row_i8_2(
6203    packed: &[u8],
6204    scales: &[u8],
6205    g0: usize,
6206    gpr: usize,
6207    xq1: &[i8],
6208    xq2: &[i8],
6209) -> (f32, f32) {
6210    #[cfg(target_arch = "aarch64")]
6211    unsafe {
6212        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
6213    }
6214    #[cfg(target_arch = "x86_64")]
6215    unsafe {
6216        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
6217    }
6218    (
6219        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
6220        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
6221    )
6222}
6223
6224/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
6225/// multi-matrix jobs can drive it for several tensors in one dispatch).
6226#[allow(clippy::too_many_arguments)]
6227fn q4_range_a8w8(
6228    packed: &[u8],
6229    scales: &[u8],
6230    gpr: usize,
6231    cols: usize,
6232    act: &SplitAct,
6233    out: SendMut,
6234    start: usize,
6235    end: usize,
6236) {
6237    for r in start..end {
6238        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
6239        // xq is zeroed at outlier slots — add the exact terms.
6240        for &(j, xv) in &act.outliers {
6241            let flat = r * cols + j;
6242            let byte = packed[flat / 2];
6243            let nib = if flat & 1 == 0 {
6244                byte & 0x0F
6245            } else {
6246                byte >> 4
6247            };
6248            let s = f16_to_f32(u16::from_le_bytes([
6249                scales[(flat / GROUP_SIZE) * 2],
6250                scales[(flat / GROUP_SIZE) * 2 + 1],
6251            ]));
6252            acc += ((nib as i32 - 8) as f32) * s * xv;
6253        }
6254        // SAFETY: disjoint row ranges per worker.
6255        unsafe { *out.at(r) = acc };
6256    }
6257}
6258
6259/// Two-input q4 row range via the A8W8 int8 path — kernel body of
6260/// `q4matvec2`, extracted for pair multi-matrix jobs.
6261#[allow(clippy::too_many_arguments)]
6262fn q4_range2_a8w8(
6263    packed: &[u8],
6264    scales: &[u8],
6265    gpr: usize,
6266    cols: usize,
6267    a1: &SplitAct,
6268    a2: &SplitAct,
6269    p1: SendMut,
6270    p2: SendMut,
6271    start: usize,
6272    end: usize,
6273) {
6274    for r in start..end {
6275        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
6276        let mut acc1 = s1 * a1.sx;
6277        let mut acc2 = s2 * a2.sx;
6278        // xq is zeroed at outlier slots — add the exact terms.
6279        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
6280            for &(j, xv) in outliers {
6281                let flat = r * cols + j;
6282                let byte = packed[flat / 2];
6283                let nib = if flat & 1 == 0 {
6284                    byte & 0x0F
6285                } else {
6286                    byte >> 4
6287                };
6288                let s = f16_to_f32(u16::from_le_bytes([
6289                    scales[(flat / GROUP_SIZE) * 2],
6290                    scales[(flat / GROUP_SIZE) * 2 + 1],
6291                ]));
6292                *acc += ((nib as i32 - 8) as f32) * s * xv;
6293            }
6294        };
6295        fix(&a1.outliers, &mut acc1);
6296        fix(&a2.outliers, &mut acc2);
6297        // SAFETY: disjoint row ranges per worker.
6298        unsafe {
6299            *p1.at(r) = acc1;
6300            *p2.at(r) = acc2;
6301        }
6302    }
6303}
6304
6305/// Exact scalar q4 row range (same extraction, non-SDOT path).
6306fn q4_range_f32(
6307    packed: &[u8],
6308    scales: &[u8],
6309    gpr: usize,
6310    x: &[f32],
6311    out: SendMut,
6312    start: usize,
6313    end: usize,
6314) {
6315    for r in start..end {
6316        let mut acc = 0f32;
6317        for gi in 0..gpr {
6318            let g = r * gpr + gi;
6319            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6320            let pk = &packed[g * 16..(g + 1) * 16];
6321            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6322            let mut ga = 0f32;
6323            for (k, &b) in pk.iter().enumerate() {
6324                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6325                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6326            }
6327            acc += ga * s;
6328        }
6329        // SAFETY: disjoint row ranges per worker.
6330        unsafe { *out.at(r) = acc };
6331    }
6332}
6333
6334/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
6335/// dotted against both activations (was: two full matvecs — double
6336/// weight traffic). Per-lane math matches `q4matvec` exactly.
6337#[allow(clippy::too_many_arguments)]
6338fn q4matvec2(
6339    bytes: &[u8],
6340    x1: &[f32],
6341    x2: &[f32],
6342    rows: usize,
6343    cols: usize,
6344    o1: &mut [f32],
6345    o2: &mut [f32],
6346    pool: Option<&Pool>,
6347) {
6348    debug_assert_eq!(o1.len(), rows);
6349    debug_assert_eq!(o2.len(), rows);
6350    let (packed, scales) = q4_split(bytes, rows, cols);
6351    let gpr = cols / GROUP_SIZE;
6352
6353    if a8w8_enabled() {
6354        let a1 = split_act(x1);
6355        let a2 = split_act(x2);
6356        let p1 = SendMut(o1.as_mut_ptr());
6357        let p2 = SendMut(o2.as_mut_ptr());
6358        let run = move |start: usize, end: usize| {
6359            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
6360        };
6361        dispatch_rows(pool, rows, &run);
6362        return;
6363    }
6364
6365    let p1 = SendMut(o1.as_mut_ptr());
6366    let p2 = SendMut(o2.as_mut_ptr());
6367    let run = move |start: usize, end: usize| {
6368        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
6369    };
6370    dispatch_rows(pool, rows, &run);
6371}
6372
6373/// Two-input exact scalar q4 row range (same extraction).
6374#[allow(clippy::too_many_arguments)]
6375fn q4_range2_f32(
6376    packed: &[u8],
6377    scales: &[u8],
6378    gpr: usize,
6379    x1: &[f32],
6380    x2: &[f32],
6381    p1: SendMut,
6382    p2: SendMut,
6383    start: usize,
6384    end: usize,
6385) {
6386    for r in start..end {
6387        let (mut acc1, mut acc2) = (0f32, 0f32);
6388        for gi in 0..gpr {
6389            let g = r * gpr + gi;
6390            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6391            let pk = &packed[g * 16..(g + 1) * 16];
6392            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6393            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6394            let (mut g1, mut g2) = (0f32, 0f32);
6395            for (k, &b) in pk.iter().enumerate() {
6396                let wl = (b & 0x0F) as f32 - 8.0;
6397                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
6398                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
6399                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
6400            }
6401            acc1 += g1 * s;
6402            acc2 += g2 * s;
6403        }
6404        // SAFETY: disjoint row ranges per worker.
6405        unsafe {
6406            *p1.at(r) = acc1;
6407            *p2.at(r) = acc2;
6408        }
6409    }
6410}
6411
6412thread_local! {
6413    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
6414    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
6415    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
6416    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6417}
6418
6419/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
6420/// and dotted against ALL b activations (prefill used to fall back to b
6421/// full matvecs — b× weight traffic and b× nibble decode). Per-position
6422/// math matches `q4matvec` exactly: same group order, same accumulation.
6423/// `out` is row-major [b, rows] like `qmatmat`.
6424#[allow(clippy::too_many_arguments)]
6425fn q4matmat(
6426    bytes: &[u8],
6427    xs_all: &[f32],
6428    b: usize,
6429    rows: usize,
6430    cols: usize,
6431    out: &mut [f32],
6432    pool: Option<&Pool>,
6433) {
6434    debug_assert_eq!(xs_all.len(), b * cols);
6435    debug_assert_eq!(out.len(), b * rows);
6436    let (packed, scales) = q4_split(bytes, rows, cols);
6437    let gpr = cols / GROUP_SIZE;
6438    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6439
6440    if a8w8_enabled() {
6441        let acts: Vec<SplitAct> = (0..b)
6442            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6443            .collect();
6444        let acts = &acts;
6445        let out_addr = SendMut(out.as_mut_ptr());
6446        let run = move |start: usize, end: usize| {
6447            ROW_I8.with(|rb| {
6448                let mut buf = rb.borrow_mut();
6449                buf.resize(cols, 0);
6450                for r in start..end {
6451                    // Unpack the row's nibbles to centered i8 once
6452                    // (element 2k = low nibble, 2k+1 = high — flat order,
6453                    // same as dot_q4_row_sdot's zip).
6454                    for gi in 0..gpr {
6455                        let g = r * gpr + gi;
6456                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6457                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
6458                            buf[gi * GROUP_SIZE + k * 2 + 1] =
6459                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
6460                        }
6461                    }
6462                    let mut bi = 0usize;
6463                    #[cfg(target_arch = "x86_64")]
6464                    if avx2_enabled()
6465                        && std::env::var("CMF_X86_BLOCKED")
6466                            .map(|v| v != "0")
6467                            .unwrap_or(true)
6468                    {
6469                        while bi + 4 <= acts.len() {
6470                            let xs = [
6471                                acts[bi].xq.as_slice(),
6472                                acts[bi + 1].xq.as_slice(),
6473                                acts[bi + 2].xq.as_slice(),
6474                                acts[bi + 3].xq.as_slice(),
6475                            ];
6476                            let d = unsafe {
6477                                if vnni_tiles_enabled() {
6478                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
6479                                } else {
6480                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
6481                                }
6482                            };
6483                            for k in 0..4 {
6484                                let act = &acts[bi + k];
6485                                let mut acc = d[k] * act.sx;
6486                                for &(j, xv) in &act.outliers {
6487                                    acc += (buf[j] as i8) as f32
6488                                        * gscale((r * cols + j) / GROUP_SIZE)
6489                                        * xv;
6490                                }
6491                                // SAFETY: disjoint (bi, r) cells per worker.
6492                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6493                            }
6494                            bi += 4;
6495                        }
6496                    }
6497                    while bi < acts.len() {
6498                        let act = &acts[bi];
6499                        let mut acc = 0f32;
6500                        for gi in 0..gpr {
6501                            let d = dot_i8_i8(
6502                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6503                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6504                            );
6505                            acc += d as f32 * gscale(r * gpr + gi);
6506                        }
6507                        acc *= act.sx;
6508                        // xq is zeroed at outlier slots — exact terms.
6509                        for &(j, xv) in &act.outliers {
6510                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
6511                        }
6512                        // SAFETY: disjoint (bi, r) cells per worker row range.
6513                        unsafe { *out_addr.at(bi * rows + r) = acc };
6514                        bi += 1;
6515                    }
6516                }
6517            })
6518        };
6519        dispatch_rows(pool, rows, &run);
6520        return;
6521    }
6522
6523    let out_addr = SendMut(out.as_mut_ptr());
6524    let run = move |start: usize, end: usize| {
6525        ROW_F32.with(|rb| {
6526            let mut buf = rb.borrow_mut();
6527            buf.resize(cols, 0.0);
6528            for r in start..end {
6529                // Decode raw (nib − 8) values once; scales stay per-group
6530                // so the accumulation order matches q4matvec bit-for-bit.
6531                for gi in 0..gpr {
6532                    let g = r * gpr + gi;
6533                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6534                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
6535                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
6536                    }
6537                }
6538                for bi in 0..b {
6539                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6540                    let mut acc = 0f32;
6541                    for gi in 0..gpr {
6542                        let mut ga = 0f32;
6543                        // Pairwise (lo + hi) addition, matching
6544                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
6545                        // a flat one-per-element loop rounds differently
6546                        // and broke bit-parity on the scalar (x86) path.
6547                        for k in 0..GROUP_SIZE / 2 {
6548                            let e = gi * GROUP_SIZE + k * 2;
6549                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
6550                        }
6551                        acc += ga * gscale(r * gpr + gi);
6552                    }
6553                    // SAFETY: disjoint (bi, r) cells per worker row range.
6554                    unsafe { *out_addr.at(bi * rows + r) = acc };
6555                }
6556            }
6557        })
6558    };
6559    dispatch_rows(pool, rows, &run);
6560}
6561
6562/// Batched vbit matmat: each variable-bit row is decoded from the mmap
6563/// ONCE for the whole microbatch. Same per-position math as
6564/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
6565/// and the scalar path).
6566#[allow(clippy::too_many_arguments)]
6567fn vbitmatmat(
6568    bytes: &[u8],
6569    offsets: &[usize],
6570    xs_all: &[f32],
6571    b: usize,
6572    rows: usize,
6573    cols: usize,
6574    out: &mut [f32],
6575    pool: Option<&Pool>,
6576) {
6577    debug_assert_eq!(xs_all.len(), b * cols);
6578    debug_assert_eq!(out.len(), b * rows);
6579    debug_assert_eq!(offsets.len(), rows + 1);
6580    let ng = cols / GROUP_SIZE;
6581    let bits = &bytes[..rows];
6582    let sc_off = rows;
6583    let gscale = |r: usize, g: usize| {
6584        let so = (r * ng + g) * 2;
6585        f16_to_f32(u16::from_le_bytes([
6586            bytes[sc_off + so],
6587            bytes[sc_off + so + 1],
6588        ]))
6589    };
6590
6591    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
6592    let decode_f32 = |r: usize, dst: &mut [f32]| {
6593        let bw = bits[r] as usize;
6594        let l = ((1i32 << (bw - 1)) - 1) as f32;
6595        let data = &bytes[offsets[r]..offsets[r + 1]];
6596        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
6597        for d in dst.iter_mut() {
6598            while nbits < bw {
6599                acc = (acc << 8) | data[idx] as u64;
6600                idx += 1;
6601                nbits += 8;
6602            }
6603            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
6604            nbits -= bw;
6605            *d = u - l;
6606        }
6607    };
6608
6609    if a8w8_enabled() {
6610        let acts: Vec<SplitAct> = (0..b)
6611            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6612            .collect();
6613        let acts = &acts;
6614        let out_addr = SendMut(out.as_mut_ptr());
6615        let run = move |start: usize, end: usize| {
6616            for r in start..end {
6617                let bw = bits[r] as usize;
6618                if bw == 8 {
6619                    // u−L reaches 128 → no i8 path; decode once, exact
6620                    // f32 dots for every position (same as vbitmatvec).
6621                    ROW_F32.with(|rb| {
6622                        let mut buf = rb.borrow_mut();
6623                        buf.resize(cols, 0.0);
6624                        decode_f32(r, &mut buf);
6625                        for bi in 0..b {
6626                            let x = &xs_all[bi * cols..(bi + 1) * cols];
6627                            let mut dot = 0f32;
6628                            for g in 0..ng {
6629                                let mut gd = 0f32;
6630                                for k in 0..GROUP_SIZE {
6631                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6632                                }
6633                                dot += gd * gscale(r, g);
6634                            }
6635                            // SAFETY: disjoint (bi, r) cells per worker range.
6636                            unsafe { *out_addr.at(bi * rows + r) = dot };
6637                        }
6638                    });
6639                    continue;
6640                }
6641                let l = (1i32 << (bw - 1)) - 1;
6642                let data = &bytes[offsets[r]..offsets[r + 1]];
6643                ROW_I8.with(|rb| {
6644                    let mut buf = rb.borrow_mut();
6645                    buf.resize(cols, 0);
6646                    #[inline(always)]
6647                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
6648                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
6649                            let u = unpack8::<B>(&data[blk * B..]);
6650                            for k in 0..8 {
6651                                chunk[k] = (u[k] - l) as i8 as u8;
6652                            }
6653                        }
6654                    }
6655                    match bw {
6656                        3 => fill::<3>(data, l, &mut buf),
6657                        4 => vbit_fill4(data, &mut buf),
6658                        5 => fill::<5>(data, l, &mut buf),
6659                        6 => fill::<6>(data, l, &mut buf),
6660                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
6661                    }
6662                    let mut bi = 0usize;
6663                    // The vbit scale table shares q4_block's layout
6664                    // (contiguous f16 per (row·ng + g)), so the same
6665                    // blocked 1×4 kernel serves the decoded row.
6666                    #[cfg(target_arch = "x86_64")]
6667                    if avx2_enabled()
6668                        && std::env::var("CMF_X86_BLOCKED")
6669                            .map(|v| v != "0")
6670                            .unwrap_or(true)
6671                    {
6672                        while bi + 4 <= acts.len() {
6673                            let xs = [
6674                                acts[bi].xq.as_slice(),
6675                                acts[bi + 1].xq.as_slice(),
6676                                acts[bi + 2].xq.as_slice(),
6677                                acts[bi + 3].xq.as_slice(),
6678                            ];
6679                            let sxs = [
6680                                acts[bi].sx,
6681                                acts[bi + 1].sx,
6682                                acts[bi + 2].sx,
6683                                acts[bi + 3].sx,
6684                            ];
6685                            let d = unsafe {
6686                                if vnni_tiles_enabled() {
6687                                    dot_q4b_row_1x4_sx_vnni(
6688                                        &buf,
6689                                        &bytes[sc_off..],
6690                                        r * ng,
6691                                        ng,
6692                                        xs,
6693                                        sxs,
6694                                    )
6695                                } else {
6696                                    dot_q4b_row_1x4_sx_avx2(
6697                                        &buf,
6698                                        &bytes[sc_off..],
6699                                        r * ng,
6700                                        ng,
6701                                        xs,
6702                                        sxs,
6703                                    )
6704                                }
6705                            };
6706                            for k in 0..4 {
6707                                let act = &acts[bi + k];
6708                                let mut dot = d[k];
6709                                for &(j, xv) in &act.outliers {
6710                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6711                                }
6712                                // SAFETY: disjoint (bi, r) cells per worker.
6713                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
6714                            }
6715                            bi += 4;
6716                        }
6717                    }
6718                    while bi < acts.len() {
6719                        let act = &acts[bi];
6720                        let mut dot = 0f32;
6721                        for g in 0..ng {
6722                            let d = dot_i8_i8(
6723                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6724                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6725                            ) as f32
6726                                * act.sx;
6727                            dot += d * gscale(r, g);
6728                        }
6729                        for &(j, xv) in &act.outliers {
6730                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6731                        }
6732                        // SAFETY: disjoint (bi, r) cells per worker range.
6733                        unsafe { *out_addr.at(bi * rows + r) = dot };
6734                        bi += 1;
6735                    }
6736                });
6737            }
6738        };
6739        dispatch_rows(pool, rows, &run);
6740        return;
6741    }
6742
6743    let out_addr = SendMut(out.as_mut_ptr());
6744    let run = move |start: usize, end: usize| {
6745        ROW_F32.with(|rb| {
6746            let mut buf = rb.borrow_mut();
6747            buf.resize(cols, 0.0);
6748            for r in start..end {
6749                decode_f32(r, &mut buf);
6750                for bi in 0..b {
6751                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6752                    let mut dot = 0f32;
6753                    for g in 0..ng {
6754                        let mut gd = 0f32;
6755                        for k in 0..GROUP_SIZE {
6756                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6757                        }
6758                        dot += gd * gscale(r, g);
6759                    }
6760                    // SAFETY: disjoint (bi, r) cells per worker range.
6761                    unsafe { *out_addr.at(bi * rows + r) = dot };
6762                }
6763            }
6764        })
6765    };
6766    dispatch_rows(pool, rows, &run);
6767}
6768
6769/// Build a GPU batch job for a q8-family mapped tensor (primary
6770/// shard): prescaled input + directory coordinates. None → not
6771/// GPU-eligible, caller stays on the CPU.
6772pub(crate) fn gpu_batch_job<'a>(
6773    t: &'a QTensor,
6774    x: &[f32],
6775) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
6776    match t {
6777        QTensor::Mapped {
6778            model,
6779            idx,
6780            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
6781            rows,
6782            cols,
6783            row_scale,
6784            col_field,
6785            ..
6786        } => Some((
6787            model.clone(),
6788            crate::gpu::BatchJob {
6789                idx: *idx,
6790                rows: *rows,
6791                cols: *cols,
6792                row_scale,
6793                xs: prescale(x, col_field, *dt).into_owned(),
6794                q1: false,
6795            },
6796        )),
6797        // q1: raw f32 activations, tile-embedded scales.
6798        QTensor::Mapped {
6799            model,
6800            idx,
6801            dtype: TensorDtype::Q1,
6802            rows,
6803            cols,
6804            ..
6805        } => Some((
6806            model.clone(),
6807            crate::gpu::BatchJob {
6808                idx: *idx,
6809                rows: *rows,
6810                cols: *cols,
6811                row_scale: &[],
6812                xs: x.to_vec(),
6813                q1: true,
6814            },
6815        )),
6816        _ => None,
6817    }
6818}
6819
6820thread_local! {
6821    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6822    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6823}
6824
6825pub(crate) fn prescale<'a>(
6826    x: &'a [f32],
6827    col_field: &[f32],
6828    dtype: TensorDtype,
6829) -> std::borrow::Cow<'a, [f32]> {
6830    if dtype == TensorDtype::Q8_2f {
6831        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
6832    } else {
6833        std::borrow::Cow::Borrowed(x)
6834    }
6835}
6836
6837/// θ col-field fold for q8_2f activations. Borrowed pass-through for
6838/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
6839pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
6840    x: &[f32],
6841    col_field: &[f32],
6842    dtype: TensorDtype,
6843    buf_id: u8,
6844    f: F,
6845) -> R {
6846    if dtype == TensorDtype::Q8_2f {
6847        if buf_id == 1 {
6848            PRESCALE_BUF1.with(|b| {
6849                let mut buf = b.borrow_mut();
6850                buf.clear();
6851                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6852                f(&buf)
6853            })
6854        } else {
6855            PRESCALE_BUF2.with(|b| {
6856                let mut buf = b.borrow_mut();
6857                buf.clear();
6858                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6859                f(&buf)
6860            })
6861        }
6862    } else {
6863        f(x)
6864    }
6865}
6866
6867// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
6868
6869/// AVX2+FMA available? Default ON when the CPU supports both;
6870/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
6871#[cfg(target_arch = "x86_64")]
6872pub(crate) fn avx2_enabled() -> bool {
6873    use std::sync::OnceLock;
6874    static ON: OnceLock<bool> = OnceLock::new();
6875    *ON.get_or_init(|| {
6876        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
6877            && std::arch::is_x86_feature_detected!("avx2")
6878            && std::arch::is_x86_feature_detected!("fma")
6879    })
6880}
6881
6882/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
6883/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
6884/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
6885/// active either way, they are exact (regrouped sums only).
6886#[cfg(target_arch = "x86_64")]
6887fn avx2_a8w8_enabled() -> bool {
6888    use std::sync::OnceLock;
6889    static ON: OnceLock<bool> = OnceLock::new();
6890    *ON.get_or_init(|| {
6891        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
6892    })
6893}
6894
6895/// A8W8 quantized-activation path available on THIS machine? One
6896/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
6897/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
6898#[inline]
6899pub(crate) fn a8w8_enabled() -> bool {
6900    #[cfg(target_arch = "aarch64")]
6901    {
6902        sdot_enabled()
6903    }
6904    #[cfg(target_arch = "x86_64")]
6905    {
6906        avx2_a8w8_enabled()
6907    }
6908    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6909    {
6910        false
6911    }
6912}
6913
6914/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
6915/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
6916#[inline]
6917#[allow(unreachable_code)]
6918fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
6919    #[cfg(target_arch = "aarch64")]
6920    unsafe {
6921        return dot_i8_sdot(w, xq);
6922    }
6923    #[cfg(target_arch = "x86_64")]
6924    unsafe {
6925        if avx512vnni_enabled() {
6926            return dot_i8_i8_vnni(w, xq);
6927        }
6928        return dot_i8_i8_avx2(w, xq);
6929    }
6930    w.iter()
6931        .zip(xq)
6932        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
6933        .sum()
6934}
6935
6936/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
6937/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
6938/// `vpdpbusd` encoding.
6939#[cfg(target_arch = "x86_64")]
6940fn avx512vnni_enabled() -> bool {
6941    use std::sync::OnceLock;
6942    static ON: OnceLock<bool> = OnceLock::new();
6943    *ON.get_or_init(|| {
6944        std::env::var("CMF_AVX512")
6945            .map(|v| v != "0")
6946            .unwrap_or(true)
6947            && std::arch::is_x86_feature_detected!("avx512f")
6948            && std::arch::is_x86_feature_detected!("avx512bw")
6949            && std::arch::is_x86_feature_detected!("avx512vl")
6950            && std::arch::is_x86_feature_detected!("avx512vnni")
6951    })
6952}
6953
6954/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
6955/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
6956/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
6957/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
6958/// (+4%) — consistent, no leg regressed. The tile kernels keep a
6959/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
6960/// smaller than the long-dot q8 win (+13%), but it is real and free.
6961#[cfg(target_arch = "x86_64")]
6962fn vnni_tiles_enabled() -> bool {
6963    use std::sync::OnceLock;
6964    static ON: OnceLock<bool> = OnceLock::new();
6965    *ON.get_or_init(|| {
6966        std::env::var("CMF_VNNI_TILES")
6967            .map(|v| v != "0")
6968            .unwrap_or(true)
6969            && avx512vnni_enabled()
6970    })
6971}
6972
6973/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
6974/// plus the same horizontal reduce the AVX2 kernels use. Products are
6975/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
6976/// is bit-identical to the maddubs+madd pair it replaces.
6977#[cfg(target_arch = "x86_64")]
6978#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6979#[inline]
6980unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
6981    // SAFETY: pure register math.
6982    unsafe {
6983        use core::arch::x86_64::*;
6984        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
6985        let hi128 = _mm256_extracti128_si256::<1>(d);
6986        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6987        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6988        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6989        _mm_cvtsi128_si32(s32)
6990    }
6991}
6992
6993/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
6994/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
6995/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
6996/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
6997#[cfg(target_arch = "x86_64")]
6998#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6999unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7000    // SAFETY: callers uphold slice-length contracts (see call sites).
7001    unsafe {
7002        use core::arch::x86_64::*;
7003        let n = w.len();
7004        let mut j = 0usize;
7005        let mut total: i32;
7006        // 4 independent accumulators: vpdpbusd is its own loop-carried
7007        // dependency (~5-cycle latency) — a single-acc loop runs
7008        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
7009        // on Granite Rapids.
7010        {
7011            #[inline(always)]
7012            unsafe fn step(
7013                w: *const u8,
7014                x: *const i8,
7015                acc: core::arch::x86_64::__m512i,
7016            ) -> core::arch::x86_64::__m512i {
7017                unsafe {
7018                    use core::arch::x86_64::*;
7019                    let wv = _mm512_loadu_si512(w as *const _);
7020                    let xv = _mm512_loadu_si512(x as *const _);
7021                    let aw = _mm512_abs_epi8(wv);
7022                    let neg = _mm512_movepi8_mask(wv);
7023                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
7024                    _mm512_dpbusd_epi32(acc, aw, sx)
7025                }
7026            }
7027            let (mut a0, mut a1, mut a2, mut a3) = (
7028                _mm512_setzero_si512(),
7029                _mm512_setzero_si512(),
7030                _mm512_setzero_si512(),
7031                _mm512_setzero_si512(),
7032            );
7033            while j + 256 <= n {
7034                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7035                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
7036                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
7037                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
7038                j += 256;
7039            }
7040            while j + 64 <= n {
7041                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7042                j += 64;
7043            }
7044            let s01 = _mm512_add_epi32(a0, a1);
7045            let s23 = _mm512_add_epi32(a2, a3);
7046            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
7047        }
7048        // 32-wide (q4/vbit groups are exactly 32 bytes).
7049        if j + 32 <= n {
7050            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7051            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7052            let d = _mm256_dpbusd_epi32(
7053                _mm256_setzero_si256(),
7054                _mm256_abs_epi8(wv),
7055                _mm256_sign_epi8(xv, wv),
7056            );
7057            let hi128 = _mm256_extracti128_si256::<1>(d);
7058            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7059            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7060            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7061            total += _mm_cvtsi128_si32(s32);
7062            j += 32;
7063        }
7064        while j < n {
7065            total += (w[j] as i8) as i32 * xq[j] as i32;
7066            j += 1;
7067        }
7068        total
7069    }
7070}
7071
7072/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
7073#[cfg(target_arch = "x86_64")]
7074#[target_feature(enable = "avx2,fma")]
7075unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
7076    // SAFETY: callers uphold slice-length contracts (see call sites).
7077    unsafe {
7078        use core::arch::x86_64::*;
7079        let n = x.len();
7080        let wp = w.as_ptr();
7081        let xp = x.as_ptr();
7082        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
7083        let mut j = 0usize;
7084        while j + 16 <= n {
7085            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
7086            let lo = _mm256_cvtepi8_epi32(wb);
7087            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
7088            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
7089            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
7090            j += 16;
7091        }
7092        let acc = _mm256_add_ps(a0, a1);
7093        let hi128 = _mm256_extractf128_ps::<1>(acc);
7094        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
7095        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
7096        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
7097        let mut sum = _mm_cvtss_f32(s32);
7098        while j < n {
7099            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
7100            j += 1;
7101        }
7102        sum
7103    }
7104}
7105
7106/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
7107/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
7108/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
7109/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
7110#[cfg(target_arch = "x86_64")]
7111#[target_feature(enable = "avx2")]
7112unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
7113    // SAFETY: callers uphold slice-length contracts (see call sites).
7114    unsafe {
7115        use core::arch::x86_64::*;
7116        let n = w.len();
7117        let ones = _mm256_set1_epi16(1);
7118        let mut acc = _mm256_setzero_si256();
7119        let mut j = 0usize;
7120        while j + 32 <= n {
7121            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7122            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7123            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7124            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
7125            j += 32;
7126        }
7127        let hi128 = _mm256_extracti128_si256::<1>(acc);
7128        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
7129        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7130        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7131        let mut s = _mm_cvtsi128_si32(s32);
7132        while j < n {
7133            s += (w[j] as i8) as i32 * xq[j] as i32;
7134            j += 1;
7135        }
7136        s
7137    }
7138}
7139
7140/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
7141/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
7142/// slice as a combined 2×8 register and meets two activation pairs.
7143#[cfg(target_arch = "aarch64")]
7144#[target_feature(enable = "neon,i8mm")]
7145unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7146    // SAFETY: callers uphold slice-length contracts.
7147    unsafe {
7148        use core::arch::aarch64::*;
7149        use core::arch::asm;
7150        let n = w0.len();
7151        let w0p = w0.as_ptr() as *const i8;
7152        let w1p = w1.as_ptr() as *const i8;
7153        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
7154        // same for x2/x3.
7155        let mut acc01 = vdupq_n_s32(0);
7156        let mut acc23 = vdupq_n_s32(0);
7157        let mut i = 0usize;
7158        while i + 8 <= n {
7159            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
7160            let xb01 = vcombine_s8(
7161                vld1_s8(xs[0].as_ptr().add(i)),
7162                vld1_s8(xs[1].as_ptr().add(i)),
7163            );
7164            let xb23 = vcombine_s8(
7165                vld1_s8(xs[2].as_ptr().add(i)),
7166                vld1_s8(xs[3].as_ptr().add(i)),
7167            );
7168            asm!(
7169                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
7170                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
7171                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
7172                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
7173                options(pure, nomem, nostack),
7174            );
7175            i += 8;
7176        }
7177        let mut out = [[0i32; 4]; 2];
7178        let a01: [i32; 4] = core::mem::transmute(acc01);
7179        let a23: [i32; 4] = core::mem::transmute(acc23);
7180        out[0][0] = a01[0];
7181        out[0][1] = a01[1];
7182        out[1][0] = a01[2];
7183        out[1][1] = a01[3];
7184        out[0][2] = a23[0];
7185        out[0][3] = a23[1];
7186        out[1][2] = a23[2];
7187        out[1][3] = a23[3];
7188        if i < n {
7189            for (k, x) in xs.iter().enumerate() {
7190                for j in i..n {
7191                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7192                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7193                }
7194            }
7195        }
7196        out
7197    }
7198}
7199
7200/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
7201/// registers across four activation streams, eight sdot accumulators.
7202/// (The per-row form re-read each W row once per activation.)
7203#[cfg(target_arch = "aarch64")]
7204#[target_feature(enable = "neon,dotprod")]
7205unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7206    // SAFETY: callers uphold slice-length contracts.
7207    unsafe {
7208        use core::arch::aarch64::*;
7209        use core::arch::asm;
7210        let n = w0.len();
7211        let w0p = w0.as_ptr() as *const i8;
7212        let w1p = w1.as_ptr() as *const i8;
7213        let mut acc = [[vdupq_n_s32(0); 4]; 2];
7214        let mut i = 0usize;
7215        while i + 16 <= n {
7216            let wv0 = vld1q_s8(w0p.add(i));
7217            let wv1 = vld1q_s8(w1p.add(i));
7218            for (k, x) in xs.iter().enumerate() {
7219                let xv = vld1q_s8(x.as_ptr().add(i));
7220                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
7221                asm!(
7222                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
7223                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
7224                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7225                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
7226                    options(pure, nomem, nostack),
7227                );
7228                acc[0][k] = a0;
7229                acc[1][k] = a1;
7230            }
7231            i += 16;
7232        }
7233        let mut out = [[0i32; 4]; 2];
7234        for r in 0..2 {
7235            for k in 0..4 {
7236                out[r][k] = vaddvq_s32(acc[r][k]);
7237            }
7238        }
7239        if i < n {
7240            for (k, x) in xs.iter().enumerate() {
7241                for j in i..n {
7242                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7243                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7244                }
7245            }
7246        }
7247        out
7248    }
7249}
7250
7251/// Blocked 2 weight rows × 4 activations for the prefill GEMM
7252/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
7253/// abs() live in registers across all four activation streams; the
7254/// sign-fixup is recomputed per pair (the price of the maddubs trick).
7255/// Returns raw i8·i8 dots; the caller applies scales and outliers.
7256#[cfg(target_arch = "x86_64")]
7257#[target_feature(enable = "avx2")]
7258unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7259    // SAFETY: callers uphold slice-length contracts.
7260    unsafe {
7261        use core::arch::x86_64::*;
7262        let n = w0.len();
7263        let ones = _mm256_set1_epi16(1);
7264        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
7265        let mut j = 0usize;
7266        while j + 32 <= n {
7267            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
7268            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
7269            let aw0 = _mm256_abs_epi8(wv0);
7270            let aw1 = _mm256_abs_epi8(wv1);
7271            for (k, x) in xs.iter().enumerate() {
7272                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
7273                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
7274                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
7275                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
7276                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
7277            }
7278            j += 32;
7279        }
7280        let mut out = [[0i32; 4]; 2];
7281        for r in 0..2 {
7282            for k in 0..4 {
7283                let a = acc[r][k];
7284                let hi128 = _mm256_extracti128_si256::<1>(a);
7285                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
7286                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7287                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7288                out[r][k] = _mm_cvtsi128_si32(s32);
7289            }
7290        }
7291        if j < n {
7292            for (k, x) in xs.iter().enumerate() {
7293                for i in j..n {
7294                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
7295                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
7296                }
7297            }
7298        }
7299        out
7300    }
7301}
7302
7303/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
7304/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
7305/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
7306/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
7307#[cfg(target_arch = "x86_64")]
7308#[inline]
7309fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
7310    let dot = if avx512vnni_enabled() && row.len() >= 64 {
7311        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
7312    } else {
7313        unsafe { dot_i8_i8_avx2(row, &act.xq) }
7314    };
7315    let mut acc = dot as f32 * act.sx;
7316    for &(j, xv) in &act.outliers {
7317        acc += (row[j] as i8) as f32 * xv;
7318    }
7319    acc
7320}
7321
7322/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
7323/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
7324/// a single-acc loop runs latency-bound, measured on Granite Rapids).
7325#[cfg(target_arch = "x86_64")]
7326#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7327unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7328    // SAFETY: callers uphold slice-length contracts (see call sites).
7329    unsafe {
7330        use core::arch::x86_64::*;
7331        let n = w.len();
7332        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
7333        #[inline(always)]
7334        unsafe fn step(
7335            w: *const u8,
7336            x: *const i8,
7337            flip: core::arch::x86_64::__m512i,
7338            acc: core::arch::x86_64::__m512i,
7339        ) -> core::arch::x86_64::__m512i {
7340            unsafe {
7341                use core::arch::x86_64::*;
7342                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
7343                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
7344            }
7345        }
7346        let (mut a0, mut a1, mut a2, mut a3) = (
7347            _mm512_setzero_si512(),
7348            _mm512_setzero_si512(),
7349            _mm512_setzero_si512(),
7350            _mm512_setzero_si512(),
7351        );
7352        let mut j = 0usize;
7353        while j + 256 <= n {
7354            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7355            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
7356            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
7357            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
7358            j += 256;
7359        }
7360        while j + 64 <= n {
7361            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7362            j += 64;
7363        }
7364        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
7365            _mm512_add_epi32(a0, a1),
7366            _mm512_add_epi32(a2, a3),
7367        ));
7368        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
7369        while j < n {
7370            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
7371            j += 1;
7372        }
7373        total
7374    }
7375}
7376
7377/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
7378/// writer's flat order, same as the NEON vzip pair), maddubs against
7379/// the pre-quantized activation group, × the group's f16 scale. Pair
7380/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
7381/// `dot_q4_row_sdot`.
7382#[cfg(target_arch = "x86_64")]
7383#[target_feature(enable = "avx2")]
7384unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7385    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7386    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
7387    unsafe {
7388        use core::arch::x86_64::*;
7389        let lomask = _mm_set1_epi8(0x0F);
7390        let eight = _mm256_set1_epi8(8);
7391        let ones = _mm256_set1_epi16(1);
7392        let mut acc = 0f32;
7393        for gi in 0..gpr {
7394            let g = g0 + gi;
7395            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7396            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7397            let lo = _mm_and_si128(b, lomask);
7398            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7399            let w = _mm256_sub_epi8(
7400                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7401                eight,
7402            );
7403            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7404            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
7405            let d = _mm256_madd_epi16(p16, ones);
7406            let hi128 = _mm256_extracti128_si256::<1>(d);
7407            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7408            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7409            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7410            acc += _mm_cvtsi128_si32(s32) as f32 * s;
7411        }
7412        acc
7413    }
7414}
7415
7416/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
7417/// both activations dotted against the same centered i8 register.
7418#[cfg(target_arch = "x86_64")]
7419#[target_feature(enable = "avx2")]
7420unsafe fn dot_q4_row_avx2_2(
7421    packed: &[u8],
7422    scales: &[u8],
7423    g0: usize,
7424    gpr: usize,
7425    xq1: &[i8],
7426    xq2: &[i8],
7427) -> (f32, f32) {
7428    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
7429    unsafe {
7430        use core::arch::x86_64::*;
7431        let lomask = _mm_set1_epi8(0x0F);
7432        let eight = _mm256_set1_epi8(8);
7433        let ones = _mm256_set1_epi16(1);
7434        let (mut acc1, mut acc2) = (0f32, 0f32);
7435        #[inline(always)]
7436        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
7437            unsafe {
7438                use core::arch::x86_64::*;
7439                let hi128 = _mm256_extracti128_si256::<1>(d);
7440                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7441                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7442                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7443                _mm_cvtsi128_si32(s32)
7444            }
7445        }
7446        for gi in 0..gpr {
7447            let g = g0 + gi;
7448            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7449            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7450            let lo = _mm_and_si128(b, lomask);
7451            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7452            let w = _mm256_sub_epi8(
7453                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7454                eight,
7455            );
7456            let aw = _mm256_abs_epi8(w);
7457            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7458            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7459            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
7460            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
7461            acc1 += hsum(d1) as f32 * s;
7462            acc2 += hsum(d2) as f32 * s;
7463        }
7464        (acc1, acc2)
7465    }
7466}
7467
7468/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
7469#[cfg(target_arch = "x86_64")]
7470fn q8_range_avx2(
7471    q: &[u8],
7472    row_scale: &[f32],
7473    act: &SplitAct,
7474    cols: usize,
7475    out_addr: SendMut,
7476    start: usize,
7477    end: usize,
7478) {
7479    for o in start..end {
7480        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7481        // SAFETY: disjoint row ranges per worker.
7482        unsafe { *out_addr.at(o) = v };
7483    }
7484}
7485
7486/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
7487#[cfg(target_arch = "x86_64")]
7488#[allow(clippy::too_many_arguments)]
7489fn q8_range2_avx2(
7490    q: &[u8],
7491    row_scale: &[f32],
7492    a1: &SplitAct,
7493    a2: &SplitAct,
7494    cols: usize,
7495    p1: SendMut,
7496    p2: SendMut,
7497    start: usize,
7498    end: usize,
7499) {
7500    for o in start..end {
7501        let row = &q[o * cols..(o + 1) * cols];
7502        // SAFETY: disjoint row ranges per worker.
7503        unsafe {
7504            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
7505            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
7506        }
7507    }
7508}
7509
7510// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
7511
7512/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
7513/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
7514/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
7515/// accumulator dependency chain swamp the MAC advantage, and Apple's
7516/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
7517/// field trials on Cortex-A710/X-class parts with two pipes, where the
7518/// balance may differ; a pre-interleaved weight layout (repack infra)
7519/// is the known path if it ever earns its keep.
7520#[cfg(target_arch = "aarch64")]
7521fn i8mm_enabled() -> bool {
7522    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7523    *ON.get_or_init(|| {
7524        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
7525            && std::arch::is_aarch64_feature_detected!("i8mm")
7526    })
7527}
7528
7529/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
7530/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
7531/// (On non-ARM release builds only the test tolerance switch calls it.)
7532#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
7533fn sdot_enabled() -> bool {
7534    use std::sync::OnceLock;
7535    static ON: OnceLock<bool> = OnceLock::new();
7536    *ON.get_or_init(|| {
7537        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
7538        if !want {
7539            return false;
7540        }
7541
7542        #[cfg(target_arch = "aarch64")]
7543        {
7544            if std::arch::is_aarch64_feature_detected!("dotprod") {
7545                return true;
7546            }
7547            #[cfg(target_os = "android")]
7548            {
7549                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
7550                    if cpuinfo.lines().any(|l| {
7551                        (l.starts_with("Features") || l.starts_with("features"))
7552                            && l.contains("asimddp")
7553                    }) {
7554                        return true;
7555                    }
7556                }
7557            }
7558            false
7559        }
7560        #[cfg(not(target_arch = "aarch64"))]
7561        {
7562            false
7563        }
7564    })
7565}
7566
7567/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
7568/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
7569/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
7570/// matvec, shared by all rows/workers.
7571struct SplitAct {
7572    xq: Vec<i8>,
7573    sx: f32,
7574    outliers: Vec<(usize, f32)>,
7575    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
7576    /// `−128·Σx`); one i32 per split, computed once per matvec.
7577    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
7578    xsum: i32,
7579}
7580
7581thread_local! {
7582    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
7583    /// and its hidden-size allocation was steady-state heap churn.
7584    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
7585        const { std::cell::RefCell::new(Vec::new()) };
7586}
7587
7588impl Drop for SplitAct {
7589    fn drop(&mut self) {
7590        let buf = std::mem::take(&mut self.xq);
7591        if buf.capacity() > 0 {
7592            XQ_FREE.with(|f| {
7593                let mut f = f.borrow_mut();
7594                if f.len() < 16 {
7595                    f.push(buf);
7596                }
7597            });
7598        }
7599    }
7600}
7601
7602fn split_act(x: &[f32]) -> SplitAct {
7603    let n = x.len();
7604    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
7605    let thr = 8.0 * rms;
7606    // One pass: collect outliers and the bulk absmax (outliers excluded —
7607    // identical to the old zero-then-fold over a copied buffer, minus the
7608    // full-vector copy).
7609    let mut outliers: Vec<(usize, f32)> = Vec::new();
7610    let mut amax = 0f32;
7611    for (j, &v) in x.iter().enumerate() {
7612        let a = v.abs();
7613        if a > thr {
7614            outliers.push((j, v));
7615        } else if a > amax {
7616            amax = a;
7617        }
7618    }
7619    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7620    let inv = 1.0 / sx;
7621    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7622    xq.clear();
7623    xq.reserve(n);
7624    if outliers.is_empty() {
7625        xq.extend(
7626            x.iter()
7627                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
7628        );
7629    } else {
7630        // Outlier slots quantize to 0 (their exact term is added later).
7631        xq.extend(x.iter().map(|&v| {
7632            if v.abs() > thr {
7633                0
7634            } else {
7635                (v * inv).round().clamp(-127.0, 127.0) as i8
7636            }
7637        }));
7638    }
7639    let xsum = xq.iter().map(|&v| v as i32).sum();
7640    SplitAct {
7641        xq,
7642        sx,
7643        outliers,
7644        xsum,
7645    }
7646}
7647
7648fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
7649    let n = x.len();
7650    let rms = (x
7651        .iter()
7652        .zip(col)
7653        .map(|(&a, &c)| {
7654            let v = a * c;
7655            (v * v) as f64
7656        })
7657        .sum::<f64>()
7658        / n.max(1) as f64)
7659        .sqrt() as f32;
7660    let thr = 8.0 * rms;
7661
7662    let mut outliers = Vec::new();
7663    let mut amax = 0f32;
7664    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
7665        let v = a * c;
7666        let s = v.abs();
7667        if s > thr {
7668            outliers.push((j, v));
7669        } else if s > amax {
7670            amax = s;
7671        }
7672    }
7673
7674    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7675    let inv = 1.0 / sx;
7676    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7677    xq.clear();
7678    xq.reserve(n);
7679    if outliers.is_empty() {
7680        xq.extend(
7681            x.iter()
7682                .zip(col)
7683                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
7684        );
7685    } else {
7686        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
7687            let v = a * c;
7688            if v.abs() > thr {
7689                0
7690            } else {
7691                (v * inv).round().clamp(-127.0, 127.0) as i8
7692            }
7693        }));
7694    }
7695    let xsum = xq.iter().map(|&v| v as i32).sum();
7696    SplitAct {
7697        xq,
7698        sx,
7699        outliers,
7700        xsum,
7701    }
7702}
7703
7704/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
7705/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
7706#[cfg(target_arch = "aarch64")]
7707#[target_feature(enable = "neon,dotprod")]
7708unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
7709    // SAFETY: callers uphold slice-length contracts (see call sites).
7710    unsafe {
7711        use core::arch::aarch64::*;
7712        use core::arch::asm;
7713        let wp = w.as_ptr() as *const i8;
7714        let n = w.len();
7715        let (mut a0, mut a1, mut a2, mut a3) = (
7716            vdupq_n_s32(0),
7717            vdupq_n_s32(0),
7718            vdupq_n_s32(0),
7719            vdupq_n_s32(0),
7720        );
7721        let mut i = 0;
7722        while i + 64 <= n {
7723            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7724            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
7725            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
7726            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
7727            asm!(
7728                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7729                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7730                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
7731                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
7732                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7733                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7734                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
7735                options(pure, nomem, nostack),
7736            );
7737            i += 64;
7738        }
7739        while i + 16 <= n {
7740            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7741            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
7742                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
7743            i += 16;
7744        }
7745        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
7746        while i < n {
7747            s += (*wp.add(i)) as i32 * xq[i] as i32;
7748            i += 1;
7749        }
7750        s
7751    }
7752}
7753
7754/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
7755/// loaded once and reused, 4 independent accumulators hide sdot latency
7756/// (port of vmfcore `dot_i8_sdot_4rows`).
7757#[cfg(target_arch = "aarch64")]
7758#[target_feature(enable = "neon,dotprod")]
7759unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
7760    // SAFETY: callers uphold slice-length contracts (see call sites).
7761    unsafe {
7762        use core::arch::aarch64::*;
7763        use core::arch::asm;
7764        let n = xq.len();
7765        let px = xq.as_ptr();
7766        let (p0, p1, p2, p3) = (
7767            w0.as_ptr() as *const i8,
7768            w1.as_ptr() as *const i8,
7769            w2.as_ptr() as *const i8,
7770            w3.as_ptr() as *const i8,
7771        );
7772        let (mut a0, mut a1, mut a2, mut a3) = (
7773            vdupq_n_s32(0),
7774            vdupq_n_s32(0),
7775            vdupq_n_s32(0),
7776            vdupq_n_s32(0),
7777        );
7778        let mut i = 0;
7779        while i + 16 <= n {
7780            let x = vld1q_s8(px.add(i));
7781            let v0 = vld1q_s8(p0.add(i));
7782            let v1 = vld1q_s8(p1.add(i));
7783            let v2 = vld1q_s8(p2.add(i));
7784            let v3 = vld1q_s8(p3.add(i));
7785            asm!(
7786                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7787                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7788                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7789                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7790                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7791                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7792                options(pure, nomem, nostack),
7793            );
7794            i += 16;
7795        }
7796        let mut r = [
7797            vaddvq_s32(a0),
7798            vaddvq_s32(a1),
7799            vaddvq_s32(a2),
7800            vaddvq_s32(a3),
7801        ];
7802        while i < n {
7803            let xi = *px.add(i) as i32;
7804            r[0] += (*p0.add(i)) as i32 * xi;
7805            r[1] += (*p1.add(i)) as i32 * xi;
7806            r[2] += (*p2.add(i)) as i32 * xi;
7807            r[3] += (*p3.add(i)) as i32 * xi;
7808            i += 1;
7809        }
7810        r
7811    }
7812}
7813
7814/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
7815/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
7816/// line plus the shared activation chunk — a single sequential weight
7817/// stream per worker. Per-row accumulation is the same one-accumulator
7818/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
7819/// are bit-identical to the mmap-layout kernel.
7820#[cfg(target_arch = "aarch64")]
7821#[target_feature(enable = "neon,dotprod")]
7822unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
7823    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
7824    // n % 16 == 0 — guaranteed by the repack gate).
7825    unsafe {
7826        use core::arch::aarch64::*;
7827        use core::arch::asm;
7828        let n = xq.len();
7829        let px = xq.as_ptr();
7830        let pg = g.as_ptr() as *const i8;
7831        let (mut a0, mut a1, mut a2, mut a3) = (
7832            vdupq_n_s32(0),
7833            vdupq_n_s32(0),
7834            vdupq_n_s32(0),
7835            vdupq_n_s32(0),
7836        );
7837        let mut i = 0;
7838        while i + 16 <= n {
7839            let x = vld1q_s8(px.add(i));
7840            let base = pg.add(4 * i);
7841            let v0 = vld1q_s8(base);
7842            let v1 = vld1q_s8(base.add(16));
7843            let v2 = vld1q_s8(base.add(32));
7844            let v3 = vld1q_s8(base.add(48));
7845            asm!(
7846                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7847                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7848                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7849                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7850                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7851                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7852                options(pure, nomem, nostack),
7853            );
7854            i += 16;
7855        }
7856        [
7857            vaddvq_s32(a0),
7858            vaddvq_s32(a1),
7859            vaddvq_s32(a2),
7860            vaddvq_s32(a3),
7861        ]
7862    }
7863}
7864
7865/// One q8 row range via SDOT (4-row blocks + tail) — the body of
7866/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
7867/// SAME kernel for several tensors under one pool dispatch. `rep` — the
7868/// load-time interleaved repack (empty = mmap layout only); rows outside
7869/// full 4-row groups always come from the mmap layout.
7870#[cfg(target_arch = "aarch64")]
7871fn q8_range_sdot(
7872    q: &[u8],
7873    rep: &[u8],
7874    row_scale: &[f32],
7875    act: &SplitAct,
7876    cols: usize,
7877    out_addr: SendMut,
7878    start: usize,
7879    end: usize,
7880) {
7881    let mut o = start;
7882    // Leading rows to the group boundary (repack path only): the pool
7883    // splits row ranges arbitrarily, groups are absolute.
7884    if !rep.is_empty() {
7885        while o < end && o % 4 != 0 {
7886            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7887            unsafe { *out_addr.at(o) = v };
7888            o += 1;
7889        }
7890    }
7891    while o + 4 <= end {
7892        let r = if rep.is_empty() {
7893            unsafe {
7894                dot_i8_sdot_4rows(
7895                    &q[o * cols..(o + 1) * cols],
7896                    &q[(o + 1) * cols..(o + 2) * cols],
7897                    &q[(o + 2) * cols..(o + 3) * cols],
7898                    &q[(o + 3) * cols..(o + 4) * cols],
7899                    &act.xq,
7900                )
7901            }
7902        } else {
7903            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
7904        };
7905        for k in 0..4 {
7906            let mut acc = r[k] as f32 * act.sx;
7907            for &(j, xv) in &act.outliers {
7908                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
7909            }
7910            // SAFETY: disjoint row ranges per worker.
7911            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
7912        }
7913        o += 4;
7914    }
7915    while o < end {
7916        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7917        unsafe { *out_addr.at(o) = v };
7918        o += 1;
7919    }
7920}
7921
7922/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
7923/// for the fused pair multi-matrix job (`matvec2_many`).
7924#[cfg(target_arch = "aarch64")]
7925#[allow(clippy::too_many_arguments)]
7926fn q8_range2_sdot(
7927    q: &[u8],
7928    row_scale: &[f32],
7929    a1: &SplitAct,
7930    a2: &SplitAct,
7931    cols: usize,
7932    p1: SendMut,
7933    p2: SendMut,
7934    start: usize,
7935    end: usize,
7936) {
7937    for o in start..end {
7938        let row = &q[o * cols..(o + 1) * cols];
7939        // SAFETY: disjoint row ranges per worker.
7940        unsafe {
7941            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
7942            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
7943        }
7944    }
7945}
7946
7947/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
7948#[allow(clippy::too_many_arguments)]
7949fn q8_range2_f32(
7950    q: &[u8],
7951    row_scale: &[f32],
7952    x1: &[f32],
7953    x2: &[f32],
7954    cols: usize,
7955    p1: SendMut,
7956    p2: SendMut,
7957    start: usize,
7958    end: usize,
7959) {
7960    for o in start..end {
7961        let row = &q[o * cols..(o + 1) * cols];
7962        // SAFETY: disjoint row ranges per worker.
7963        unsafe {
7964            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
7965            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
7966        }
7967    }
7968}
7969
7970/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
7971fn q8_range_f32(
7972    q: &[u8],
7973    row_scale: &[f32],
7974    xs: &[f32],
7975    cols: usize,
7976    out_addr: SendMut,
7977    start: usize,
7978    end: usize,
7979) {
7980    for o in start..end {
7981        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
7982        // SAFETY: disjoint row ranges per worker.
7983        unsafe { *out_addr.at(o) = v };
7984    }
7985}
7986
7987/// SDOT row dot with exact outlier correction:
7988/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
7989#[cfg(target_arch = "aarch64")]
7990#[inline]
7991fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
7992    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
7993    for &(j, xv) in &act.outliers {
7994        acc += (row[j] as i8) as f32 * xv;
7995    }
7996    acc
7997}
7998
7999/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
8000/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
8001/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
8002/// the caller multiplies by the activation scale and adds the exact
8003/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
8004/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
8005/// → zip(lo,hi) restores flat order.
8006#[cfg(target_arch = "aarch64")]
8007#[target_feature(enable = "neon,dotprod")]
8008unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8009    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8010    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8011    unsafe {
8012        use core::arch::aarch64::*;
8013        use core::arch::asm;
8014        let lomask = vdupq_n_u8(0x0F);
8015        let eight = vdupq_n_s8(8);
8016        let mut acc = 0f32;
8017        for gi in 0..gpr {
8018            let g = g0 + gi;
8019            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8020            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8021            let lo = vandq_u8(b, lomask);
8022            let hi = vshrq_n_u8::<4>(b);
8023            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8024            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8025            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
8026            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
8027            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8028            asm!(
8029                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8030                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8031                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8032                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8033                options(pure, nomem, nostack),
8034            );
8035            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8036        }
8037        acc
8038    }
8039}
8040
8041/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
8042/// part) happens ONCE per group; both pre-quantized activations are
8043/// dotted against the same centered i8 registers. Per-lane math matches
8044/// `dot_q4_row_sdot` exactly.
8045#[cfg(target_arch = "aarch64")]
8046#[target_feature(enable = "neon,dotprod")]
8047unsafe fn dot_q4_row_sdot2(
8048    packed: &[u8],
8049    scales: &[u8],
8050    g0: usize,
8051    gpr: usize,
8052    xq1: &[i8],
8053    xq2: &[i8],
8054) -> (f32, f32) {
8055    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8056    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
8057    unsafe {
8058        use core::arch::aarch64::*;
8059        use core::arch::asm;
8060        let lomask = vdupq_n_u8(0x0F);
8061        let eight = vdupq_n_s8(8);
8062        let (mut acc1, mut acc2) = (0f32, 0f32);
8063        for gi in 0..gpr {
8064            let g = g0 + gi;
8065            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8066            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8067            let lo = vandq_u8(b, lomask);
8068            let hi = vshrq_n_u8::<4>(b);
8069            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8070            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8071            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
8072            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
8073            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
8074            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
8075            let (mut a0, mut a1, mut b0, mut b1) = (
8076                vdupq_n_s32(0),
8077                vdupq_n_s32(0),
8078                vdupq_n_s32(0),
8079                vdupq_n_s32(0),
8080            );
8081            asm!(
8082                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
8083                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
8084                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
8085                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
8086                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8087                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
8088                e0 = in(vreg) e0, e1 = in(vreg) e1,
8089                x10 = in(vreg) x10, x11 = in(vreg) x11,
8090                x20 = in(vreg) x20, x21 = in(vreg) x21,
8091                options(pure, nomem, nostack),
8092            );
8093            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8094            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
8095        }
8096        (acc1, acc2)
8097    }
8098}
8099
8100// ───────────────────── fused int8 kernels ─────────────────────
8101
8102/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
8103/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
8104#[inline]
8105pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
8106    #[cfg(target_arch = "aarch64")]
8107    unsafe {
8108        return axpy_i8_f32_neon(acc, row, w);
8109    }
8110    #[cfg(target_arch = "x86_64")]
8111    if avx2_enabled() {
8112        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
8113    }
8114    #[allow(unreachable_code)]
8115    {
8116        for (a, &b) in acc.iter_mut().zip(row) {
8117            *a += w * b as f32;
8118        }
8119    }
8120}
8121
8122/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
8123#[cfg(target_arch = "x86_64")]
8124#[target_feature(enable = "avx2,fma")]
8125unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
8126    // SAFETY: callers uphold slice-length contracts (see call sites).
8127    unsafe {
8128        use core::arch::x86_64::*;
8129        let n = acc.len().min(row.len());
8130        let ap = acc.as_mut_ptr();
8131        let rp = row.as_ptr();
8132        let wv = _mm256_set1_ps(w);
8133        let mut j = 0usize;
8134        while j + 16 <= n {
8135            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
8136            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
8137            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
8138            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
8139            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
8140            _mm256_storeu_ps(ap.add(j), v0);
8141            _mm256_storeu_ps(ap.add(j + 8), v1);
8142            j += 16;
8143        }
8144        while j < n {
8145            *ap.add(j) += w * (*rp.add(j)) as f32;
8146            j += 1;
8147        }
8148    }
8149}
8150
8151#[cfg(target_arch = "aarch64")]
8152#[target_feature(enable = "neon")]
8153unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
8154    // SAFETY: callers uphold slice-length contracts (see call sites).
8155    unsafe {
8156        use core::arch::aarch64::*;
8157        let n = acc.len().min(row.len());
8158        let ap = acc.as_mut_ptr();
8159        let rp = row.as_ptr();
8160        let wv = vdupq_n_f32(w);
8161        let mut j = 0usize;
8162        while j + 16 <= n {
8163            let rb = vld1q_s8(rp.add(j));
8164            let lo = vmovl_s8(vget_low_s8(rb));
8165            let hi = vmovl_s8(vget_high_s8(rb));
8166            for (off, half) in [(0, lo), (8, hi)] {
8167                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
8168                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
8169                let o = j + off;
8170                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
8171                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
8172            }
8173            j += 16;
8174        }
8175        while j < n {
8176            *ap.add(j) += w * (*rp.add(j)) as f32;
8177            j += 1;
8178        }
8179    }
8180}
8181
8182/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
8183/// ≈9× scalar), scalar elsewhere.
8184#[inline]
8185pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
8186    #[cfg(target_arch = "aarch64")]
8187    unsafe {
8188        return dot_i8_f32_neon(w, x);
8189    }
8190    #[cfg(target_arch = "x86_64")]
8191    if avx2_enabled() {
8192        return unsafe { dot_i8_f32_avx2(w, x) };
8193    }
8194    #[allow(unreachable_code)]
8195    {
8196        let mut sum = 0.0f32;
8197        for (j, &b) in w.iter().enumerate() {
8198            sum += (b as i8) as f32 * x[j];
8199        }
8200        sum
8201    }
8202}
8203
8204/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
8205/// folded into the product (no prescaled copy of x). NEON on aarch64,
8206/// scalar elsewhere. Used by the active-neuron path `row_dot`.
8207#[inline]
8208fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8209    #[cfg(target_arch = "aarch64")]
8210    unsafe {
8211        return dot_i8_col_f32_neon(w, x, col);
8212    }
8213    #[allow(unreachable_code)]
8214    {
8215        let mut sum = 0.0f32;
8216        for (j, &b) in w.iter().enumerate() {
8217            sum += (b as i8) as f32 * x[j] * col[j];
8218        }
8219        sum
8220    }
8221}
8222
8223#[cfg(target_arch = "aarch64")]
8224#[target_feature(enable = "neon")]
8225unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8226    // SAFETY: callers uphold slice-length contracts (see call sites).
8227    unsafe {
8228        use core::arch::aarch64::*;
8229        let n = x.len();
8230        let wp = w.as_ptr() as *const i8;
8231        let xp = x.as_ptr();
8232        let cp = col.as_ptr();
8233        let (mut a0, mut a1, mut a2, mut a3) = (
8234            vdupq_n_f32(0.0),
8235            vdupq_n_f32(0.0),
8236            vdupq_n_f32(0.0),
8237            vdupq_n_f32(0.0),
8238        );
8239        let mut j = 0usize;
8240        while j + 16 <= n {
8241            let wb = vld1q_s8(wp.add(j));
8242            let lo = vmovl_s8(vget_low_s8(wb));
8243            let hi = vmovl_s8(vget_high_s8(wb));
8244            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8245            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8246            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8247            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8248            a0 = vfmaq_f32(
8249                a0,
8250                w0,
8251                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
8252            );
8253            a1 = vfmaq_f32(
8254                a1,
8255                w1,
8256                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
8257            );
8258            a2 = vfmaq_f32(
8259                a2,
8260                w2,
8261                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
8262            );
8263            a3 = vfmaq_f32(
8264                a3,
8265                w3,
8266                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
8267            );
8268            j += 16;
8269        }
8270        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8271        while j < n {
8272            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
8273            j += 1;
8274        }
8275        sum
8276    }
8277}
8278
8279#[cfg(target_arch = "aarch64")]
8280#[target_feature(enable = "neon")]
8281unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
8282    // SAFETY: callers uphold slice-length contracts (see call sites).
8283    unsafe {
8284        use core::arch::aarch64::*;
8285        let n = x.len();
8286        let wp = w.as_ptr() as *const i8;
8287        let xp = x.as_ptr();
8288        let (mut a0, mut a1, mut a2, mut a3) = (
8289            vdupq_n_f32(0.0),
8290            vdupq_n_f32(0.0),
8291            vdupq_n_f32(0.0),
8292            vdupq_n_f32(0.0),
8293        );
8294        let mut j = 0usize;
8295        while j + 16 <= n {
8296            let wb = vld1q_s8(wp.add(j));
8297            let lo = vmovl_s8(vget_low_s8(wb));
8298            let hi = vmovl_s8(vget_high_s8(wb));
8299            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8300            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8301            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8302            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8303            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
8304            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
8305            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
8306            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
8307            j += 16;
8308        }
8309        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8310        while j < n {
8311            sum += (*wp.add(j)) as f32 * *xp.add(j);
8312            j += 1;
8313        }
8314        sum
8315    }
8316}
8317
8318#[allow(clippy::too_many_arguments)]
8319fn qmatvec(
8320    q: &[u8],
8321    rep: &[u8],
8322    row_scale: &[f32],
8323    x: &[f32],
8324    col_field: &[f32],
8325    dtype: TensorDtype,
8326    rows: usize,
8327    cols: usize,
8328    out: &mut [f32],
8329    pool: Option<&Pool>,
8330) {
8331    debug_assert_eq!(out.len(), rows);
8332    #[cfg(not(target_arch = "aarch64"))]
8333    let _ = rep;
8334
8335    #[cfg(target_arch = "aarch64")]
8336    if sdot_enabled() {
8337        let act = if dtype == TensorDtype::Q8_2f {
8338            split_act_q8_2f(x, col_field)
8339        } else {
8340            split_act(x)
8341        };
8342        let out_addr = SendMut(out.as_mut_ptr());
8343        let run_range = |start: usize, end: usize| {
8344            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
8345        };
8346        match pool {
8347            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8348            _ => run_range(0, rows),
8349        }
8350        return;
8351    }
8352    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
8353    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
8354    #[cfg(target_arch = "x86_64")]
8355    if avx2_a8w8_enabled() {
8356        let act = if dtype == TensorDtype::Q8_2f {
8357            split_act_q8_2f(x, col_field)
8358        } else {
8359            split_act(x)
8360        };
8361        let out_addr = SendMut(out.as_mut_ptr());
8362        let run_range = |start: usize, end: usize| {
8363            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
8364        };
8365        match pool {
8366            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8367            _ => run_range(0, rows),
8368        }
8369        return;
8370    }
8371
8372    prescale_with(x, col_field, dtype, 1, |xs| {
8373        let out_addr = SendMut(out.as_mut_ptr());
8374        let run_range = move |start: usize, end: usize| {
8375            for o in start..end {
8376                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8377                // SAFETY: disjoint row ranges per worker.
8378                unsafe { *out_addr.at(o) = v };
8379            }
8380        };
8381        match pool {
8382            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8383            _ => run_range(0, rows),
8384        }
8385    });
8386}
8387
8388#[allow(clippy::too_many_arguments)]
8389fn qmatvec2(
8390    q: &[u8],
8391    row_scale: &[f32],
8392    x1: &[f32],
8393    x2: &[f32],
8394    col_field: &[f32],
8395    dtype: TensorDtype,
8396    rows: usize,
8397    cols: usize,
8398    o1: &mut [f32],
8399    o2: &mut [f32],
8400    pool: Option<&Pool>,
8401) {
8402    #[cfg(target_arch = "aarch64")]
8403    if sdot_enabled() {
8404        let a1s = if dtype == TensorDtype::Q8_2f {
8405            split_act_q8_2f(x1, col_field)
8406        } else {
8407            split_act(x1)
8408        };
8409        let a2s = if dtype == TensorDtype::Q8_2f {
8410            split_act_q8_2f(x2, col_field)
8411        } else {
8412            split_act(x2)
8413        };
8414        let p1 = SendMut(o1.as_mut_ptr());
8415        let p2 = SendMut(o2.as_mut_ptr());
8416        let run_range = |start: usize, end: usize| {
8417            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8418        };
8419        match pool {
8420            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8421            _ => run_range(0, rows),
8422        }
8423        return;
8424    }
8425    #[cfg(target_arch = "x86_64")]
8426    if avx2_a8w8_enabled() {
8427        let a1s = if dtype == TensorDtype::Q8_2f {
8428            split_act_q8_2f(x1, col_field)
8429        } else {
8430            split_act(x1)
8431        };
8432        let a2s = if dtype == TensorDtype::Q8_2f {
8433            split_act_q8_2f(x2, col_field)
8434        } else {
8435            split_act(x2)
8436        };
8437        let p1 = SendMut(o1.as_mut_ptr());
8438        let p2 = SendMut(o2.as_mut_ptr());
8439        let run_range = |start: usize, end: usize| {
8440            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8441        };
8442        match pool {
8443            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8444            _ => run_range(0, rows),
8445        }
8446        return;
8447    }
8448
8449    prescale_with(x1, col_field, dtype, 1, |x1s| {
8450        prescale_with(x2, col_field, dtype, 2, |x2s| {
8451            let p1 = SendMut(o1.as_mut_ptr());
8452            let p2 = SendMut(o2.as_mut_ptr());
8453            let run_range = move |start: usize, end: usize| {
8454                for o in start..end {
8455                    let row = &q[o * cols..(o + 1) * cols];
8456                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
8457                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
8458                    // SAFETY: disjoint row ranges per worker.
8459                    unsafe {
8460                        *p1.at(o) = s1;
8461                        *p2.at(o) = s2;
8462                    }
8463                }
8464            };
8465            match pool {
8466                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8467                _ => run_range(0, rows),
8468            }
8469        });
8470    });
8471}
8472
8473#[derive(Clone, Copy)]
8474struct SendMut(*mut f32);
8475unsafe impl Send for SendMut {}
8476unsafe impl Sync for SendMut {}
8477
8478impl SendMut {
8479    #[inline]
8480    fn at(self, i: usize) -> *mut f32 {
8481        unsafe { self.0.add(i) }
8482    }
8483}
8484
8485#[cfg(test)]
8486mod tests {
8487    use super::*;
8488
8489    #[test]
8490    fn f32_matvec_matches_matvec_rows_bitexact() {
8491        let (rows, cols) = (300, 40);
8492        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
8493        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
8494        let qt = QTensor::from_f32(w.clone(), rows, cols);
8495
8496        let mut a = vec![0.0f32; rows];
8497        matvec_rows(None, &w, &x, &mut a);
8498        let mut b = vec![0.0f32; rows];
8499        qt.matvec(&x, &mut b, None);
8500        assert_eq!(a, b);
8501    }
8502
8503    #[test]
8504    fn sdot_kernel_exact_on_grid() {
8505        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
8506        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
8507        // exact f32 dot to float rounding. This isolates kernel
8508        // correctness from quantization noise.
8509        eprintln!("sdot_enabled = {}", sdot_enabled());
8510        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
8511        let w: Vec<u8> = (0..rows * cols)
8512            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8513            .collect();
8514        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
8515        let x: Vec<f32> = (0..cols)
8516            .map(|i| match i % 3 {
8517                0 => 1.0,
8518                1 => -1.0,
8519                _ => 0.0,
8520            })
8521            .collect();
8522        let mut a = vec![0.0f32; rows];
8523        qmatvec(
8524            &w,
8525            &[],
8526            &scales,
8527            &x,
8528            &[],
8529            TensorDtype::Q8Row,
8530            rows,
8531            cols,
8532            &mut a,
8533            None,
8534        );
8535        for o in 0..rows {
8536            let mut acc = 0.0f32;
8537            for j in 0..cols {
8538                acc += (w[o * cols + j] as i8) as f32 * x[j];
8539            }
8540            let expect = acc * scales[o];
8541            assert!(
8542                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8543                "row {o}: {} vs {expect}",
8544                a[o]
8545            );
8546        }
8547    }
8548
8549    #[test]
8550    fn q1_tbl_fast_path_matches_reference() {
8551        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
8552        // row's final 4-tile window trips the 4B-overread guard (the
8553        // payload ends exactly at the last tile) — both paths must
8554        // agree with the dequant reference.
8555        let (rows, cols) = (5, 256);
8556        let gpr = cols / GROUP_SIZE;
8557        let mut bytes = Vec::new();
8558        for t in 0..rows * gpr {
8559            let s = 0.007 + (t % 11) as f32 * 0.004;
8560            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8561            for j in 0..4 {
8562                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
8563            }
8564        }
8565        let x: Vec<f32> = (0..cols)
8566            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
8567            .collect();
8568        let mut w = vec![0.0f32; rows * cols];
8569        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8570        let mut got = vec![0.0f32; rows];
8571        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8572        for o in 0..rows {
8573            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8574            assert!(
8575                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8576                "row {o}: {} vs {expect}",
8577                got[o]
8578            );
8579        }
8580        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
8581        // single-matvec path bit-for-bit.
8582        let b = 5usize;
8583        let mut xs_all = Vec::new();
8584        for bi in 0..b {
8585            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
8586        }
8587        let mut mm = vec![0.0f32; b * rows];
8588        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
8589        for bi in 0..b {
8590            let mut single = vec![0.0f32; rows];
8591            q1_matvec(
8592                &bytes,
8593                &xs_all[bi * cols..(bi + 1) * cols],
8594                rows,
8595                cols,
8596                &mut single,
8597                None,
8598            );
8599            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
8600        }
8601    }
8602
8603    #[test]
8604    fn q1_kernels_match_exact_reference() {
8605        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
8606        let (rows, cols) = (7, 96);
8607        let gpr = cols / GROUP_SIZE;
8608        let mut bytes = Vec::new();
8609        for t in 0..rows * gpr {
8610            let s = 0.01 + (t % 13) as f32 * 0.003;
8611            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8612            for j in 0..4 {
8613                bytes.push(((t * 31 + j * 97) % 251) as u8);
8614            }
8615        }
8616        // On-grid activations (±1, amax 1) → the SDOT path is exact.
8617        let x: Vec<f32> = (0..cols)
8618            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
8619            .collect();
8620        // Reference through the core dequant.
8621        let mut w = vec![0.0f32; rows * cols];
8622        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8623        let mut expect = vec![0.0f32; rows];
8624        for o in 0..rows {
8625            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8626        }
8627        let mut got = vec![0.0f32; rows];
8628        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8629        for o in 0..rows {
8630            assert!(
8631                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
8632                "row {o}: {} vs {}",
8633                got[o],
8634                expect[o]
8635            );
8636        }
8637        // Pair and batch paths agree with the single path.
8638        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
8639        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
8640        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
8641        assert_eq!(a1, got);
8642        let mut xs = x.clone();
8643        xs.extend_from_slice(&x2);
8644        let mut mm = vec![0.0f32; 2 * rows];
8645        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
8646        assert_eq!(&mm[..rows], got.as_slice());
8647        assert_eq!(&mm[rows..], a2.as_slice());
8648    }
8649
8650    #[test]
8651    fn repack_is_bit_identical() {
8652        // The interleaved-repack kernel must produce EXACTLY the same
8653        // bits as the mmap-layout kernel: integer accumulation is order-
8654        // exact, the f32 epilogue is identical. Odd rows exercise the
8655        // tail; direct range calls exercise unaligned pool splits.
8656        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
8657        let w: Vec<u8> = (0..rows * cols)
8658            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
8659            .collect();
8660        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
8661        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
8662        let rep = q8_repack_layout(&w, rows, cols);
8663        // Group interleave round-trips.
8664        for g in 0..rows / 4 {
8665            for c in 0..cols / 16 {
8666                for lane in 0..4 {
8667                    assert_eq!(
8668                        &rep[g * 4 * cols + c * 64 + lane * 16
8669                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
8670                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
8671                    );
8672                }
8673            }
8674        }
8675        let mut a = vec![0.0f32; rows];
8676        qmatvec(
8677            &w,
8678            &[],
8679            &scales,
8680            &x,
8681            &[],
8682            TensorDtype::Q8Row,
8683            rows,
8684            cols,
8685            &mut a,
8686            None,
8687        );
8688        let mut b = vec![0.0f32; rows];
8689        qmatvec(
8690            &w,
8691            &rep,
8692            &scales,
8693            &x,
8694            &[],
8695            TensorDtype::Q8Row,
8696            rows,
8697            cols,
8698            &mut b,
8699            None,
8700        );
8701        assert_eq!(a, b, "full-range repack output diverged");
8702
8703        #[cfg(target_arch = "aarch64")]
8704        if sdot_enabled() {
8705            // Unaligned range split (pool workers get arbitrary bounds).
8706            let act = split_act(&x);
8707            let mut c1 = vec![0.0f32; rows];
8708            let mut c2 = vec![0.0f32; rows];
8709            q8_range_sdot(
8710                &w,
8711                &[],
8712                &scales,
8713                &act,
8714                cols,
8715                SendMut(c1.as_mut_ptr()),
8716                3,
8717                rows - 2,
8718            );
8719            q8_range_sdot(
8720                &w,
8721                &rep,
8722                &scales,
8723                &act,
8724                cols,
8725                SendMut(c2.as_mut_ptr()),
8726                3,
8727                rows - 2,
8728            );
8729            assert_eq!(c1, c2, "unaligned-range repack output diverged");
8730        }
8731    }
8732
8733    #[test]
8734    fn sdot_a8w8_noise_is_bounded() {
8735        // Off-grid activations: A8 quantization noise must stay small in
8736        // relative L2 over the whole output (realistic accuracy contract;
8737        // vmfcore measured argmax-identical decode on real models).
8738        let (rows, cols) = (16, 512);
8739        let w: Vec<u8> = (0..rows * cols)
8740            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8741            .collect();
8742        let scales = vec![0.01f32; rows];
8743        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
8744        let mut a = vec![0.0f32; rows];
8745        qmatvec(
8746            &w,
8747            &[],
8748            &scales,
8749            &x,
8750            &[],
8751            TensorDtype::Q8Row,
8752            rows,
8753            cols,
8754            &mut a,
8755            None,
8756        );
8757        let (mut num, mut den) = (0f64, 0f64);
8758        for o in 0..rows {
8759            let mut acc = 0.0f32;
8760            for j in 0..cols {
8761                acc += (w[o * cols + j] as i8) as f32 * x[j];
8762            }
8763            let expect = acc * scales[o];
8764            num += ((a[o] - expect) as f64).powi(2);
8765            den += (expect as f64).powi(2);
8766        }
8767        let rel = (num / den.max(1e-12)).sqrt();
8768        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
8769    }
8770
8771    #[test]
8772    fn i8_dot_neon_matches_scalar() {
8773        let n = 100;
8774        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
8775        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
8776        let mut scalar = 0.0f32;
8777        for j in 0..n {
8778            scalar += (w[j] as i8) as f32 * x[j];
8779        }
8780        let fast = dot_i8_f32(&w, &x);
8781        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
8782    }
8783
8784    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
8785    #[test]
8786    fn vbitmatvec_matches_full_dequant() {
8787        let (rows, cols) = (6, 64);
8788        let ng = cols / GROUP_SIZE;
8789        // Hand-craft: bits per row, f16 scales, packed rows.
8790        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8791        let mut bytes = bits.clone();
8792        for g in 0..rows * ng {
8793            let s = 0.02 + 0.001 * g as f32;
8794            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8795        }
8796        for r in 0..rows {
8797            let b = bits[r] as usize;
8798            let (mut acc, mut nb) = (0u64, 0usize);
8799            let mut rowbytes = Vec::new();
8800            for i in 0..cols {
8801                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8802                acc = (acc << b) | v;
8803                nb += b;
8804                while nb >= 8 {
8805                    nb -= 8;
8806                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
8807                }
8808            }
8809            if nb > 0 {
8810                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8811            }
8812            bytes.extend_from_slice(&rowbytes);
8813        }
8814        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8815
8816        let mut reference = vec![0f32; rows * cols];
8817        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
8818        let mut expect = vec![0f32; rows];
8819        for r in 0..rows {
8820            expect[r] = reference[r * cols..(r + 1) * cols]
8821                .iter()
8822                .zip(&x)
8823                .map(|(w, xv)| w * xv)
8824                .sum();
8825        }
8826        let mut got = vec![0f32; rows];
8827        let offsets = vbit_row_offsets(&bytes, rows, cols);
8828        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
8829        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
8830        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
8831        // the golden-parity gate).
8832        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
8833        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
8834        for r in 0..rows {
8835            assert!(
8836                (got[r] - expect[r]).abs() < tol * scale,
8837                "row {r}: {} vs {}",
8838                got[r],
8839                expect[r]
8840            );
8841        }
8842    }
8843
8844    /// Fused q4 matvec must match the reference full-dequant + dense
8845    /// matvec bit-for-bit in structure (same f32 math, group order).
8846    /// vbit matmat: the blocked 1×4 leg must match the per-row path
8847    /// (paired env toggle; larger shape so both code paths engage).
8848    #[test]
8849    #[cfg(target_arch = "x86_64")]
8850    fn vbit_matmat_blocked_matches_per_row() {
8851        let (rows, cols, b) = (64usize, 128usize, 9usize);
8852        let ng = cols / GROUP_SIZE;
8853        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
8854        let mut bytes = bits.clone();
8855        for g in 0..rows * ng {
8856            let sc = 0.02 + 0.0005 * g as f32;
8857            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8858        }
8859        for r in 0..rows {
8860            let bw = bits[r] as usize;
8861            let (mut acc, mut nb) = (0u64, 0usize);
8862            let mut rowbytes = Vec::new();
8863            for i in 0..cols {
8864                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
8865                acc = (acc << bw) | v;
8866                nb += bw;
8867                while nb >= 8 {
8868                    nb -= 8;
8869                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
8870                }
8871            }
8872            if nb > 0 {
8873                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8874            }
8875            bytes.extend_from_slice(&rowbytes);
8876        }
8877        let x: Vec<f32> = (0..b * cols)
8878            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8879            .collect();
8880        let offsets = vbit_row_offsets(&bytes, rows, cols);
8881        let mut y_a = vec![0f32; b * rows];
8882        let mut y_b = vec![0f32; b * rows];
8883        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8884        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
8885        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8886        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
8887        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8888        let max_d = y_a
8889            .iter()
8890            .zip(&y_b)
8891            .map(|(p, q)| (p - q).abs())
8892            .fold(0.0f32, f32::max);
8893        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
8894    }
8895
8896    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
8897    /// per-row path exactly: same nibble unpack, same group order,
8898    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
8899    /// two full 1×4 blocks plus a remainder through the single-row
8900    /// kernel. (Both paths produce identical output, so the shared
8901    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
8902    /// the verdict — worst case both sides take the same path.)
8903    #[test]
8904    fn q4t_matmat_blocked_matches_per_row() {
8905        let (rows, cols, b) = (16usize, 64usize, 9usize);
8906        let gpr = cols / GROUP_SIZE;
8907        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
8908        for r in 0..rows {
8909            for g in 0..gpr {
8910                let t = (r * gpr + g) * Q4_TILE;
8911                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
8912                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8913                for k in 0..16 {
8914                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8915                }
8916            }
8917        }
8918        let x: Vec<f32> = (0..b * cols)
8919            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8920            .collect();
8921        let mut y_blk = vec![0f32; b * rows];
8922        let mut y_row = vec![0f32; b * rows];
8923        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8924        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
8925        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8926        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
8927        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8928        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
8929    }
8930
8931    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
8932    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
8933    /// order differs — tight tolerance.
8934    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
8935    /// span varies row to row, so the codes actually exercise the full 0..31
8936    /// range rather than clustering on one rung.
8937    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
8938        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
8939        let gpr = cols / GROUP_SIZE;
8940        let stride = q4tp_code_stride(gpr);
8941        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
8942        let mut b = vec![0u8; codes_off + rows * stride];
8943        for r in 0..rows {
8944            for g in 0..gpr {
8945                let t = (r * gpr + g) * Q4TP_NIB;
8946                for k in 0..16 {
8947                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8948                }
8949            }
8950            let lo = -6.0 - 0.03 * (r % 17) as f32;
8951            let step = 0.01 + 0.004 * (r % 11) as f32;
8952            let p = params_off + r * 4;
8953            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
8954            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
8955            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
8956            for g in 0..gpr {
8957                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
8958            }
8959        }
8960        b
8961    }
8962
8963    /// The same weights re-expressed as q4_tiled, so the proven kernel can
8964    /// be the reference: each tile stores the ladder scale its code selects.
8965    /// Only the f16 rounding of that scale separates the two payloads.
8966    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
8967        let gpr = cols / GROUP_SIZE;
8968        let v = Q4tpView::new(bytes, rows, cols);
8969        let mut out = vec![0u8; rows * gpr * Q4_TILE];
8970        let mut sc = vec![0f32; gpr];
8971        for r in 0..rows {
8972            v.scales_into(r, gpr, &mut sc);
8973            for g in 0..gpr {
8974                let t = (r * gpr + g) * Q4_TILE;
8975                let s = sc[g];
8976                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8977                let src = (r * gpr + g) * Q4TP_NIB;
8978                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
8979            }
8980        }
8981        out
8982    }
8983
8984    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
8985    /// rounding — that scalar routine is the format's definition, and the
8986    /// kernels re-derive the scale from the ladder independently. Call the
8987    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
8988    /// so routing through it would test the other path by accident.
8989    #[test]
8990    fn q4tp_exact_path_matches_dequant_reference() {
8991        let (rows, cols) = (256usize, 512usize);
8992        let gpr = cols / GROUP_SIZE;
8993        let bytes = synth_q4tp(rows, cols);
8994        let mut w = vec![0f32; rows * cols];
8995        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
8996
8997        let x: Vec<f32> = (0..cols)
8998            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8999            .collect();
9000        let v = Q4tpView::new(&bytes, rows, cols);
9001        let mut sc = vec![0f32; gpr];
9002        for r in 0..rows {
9003            v.scales_into(r, gpr, &mut sc);
9004            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
9005            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
9006            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
9007            // the meaningful yardstick is the summed magnitude, not the result:
9008            // against the result any reordering of a 512-term f32 sum "fails".
9009            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9010            assert!(
9011                (got - want).abs() <= 1e-5 * mag,
9012                "row {r}: kernel {got} vs dequant {want}"
9013            );
9014        }
9015    }
9016
9017    /// The int8 (a8w8) path can't be checked against an f32 reference — the
9018    /// activation quantization dominates. Check it against the q4t kernel it
9019    /// was ported from instead, on payloads holding the same weights: that
9020    /// isolates exactly what the port could break (16 B stride, ladder
9021    /// lookup, nibble unpack) from what it deliberately shares.
9022    #[test]
9023    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
9024        let (rows, cols) = (256usize, 512usize);
9025        let bytes = synth_q4tp(rows, cols);
9026        let twin = q4tp_as_q4t(&bytes, rows, cols);
9027        let x: Vec<f32> = (0..cols)
9028            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9029            .collect();
9030
9031        let mut got = vec![0f32; rows];
9032        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
9033        let mut want = vec![0f32; rows];
9034        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
9035
9036        // Scale is f16 in the twin and f32 here, so allow that rounding on
9037        // top of the summed magnitude (same cancellation argument as above).
9038        let mut w = vec![0f32; rows * cols];
9039        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9040        for r in 0..rows {
9041            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9042            assert!(
9043                (got[r] - want[r]).abs() <= 1e-3 * mag,
9044                "row {r}: q4tp {} vs q4t {}",
9045                got[r],
9046                want[r]
9047            );
9048        }
9049    }
9050
9051    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
9052    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
9053    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
9054    /// code and its four accumulators are exactly what tends to go wrong.
9055    #[test]
9056    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
9057        let (rows, cols, b) = (256usize, 512usize, 5usize);
9058        let bytes = synth_q4tp(rows, cols);
9059        let twin = q4tp_as_q4t(&bytes, rows, cols);
9060        let xs: Vec<f32> = (0..b * cols)
9061            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9062            .collect();
9063
9064        let mut got = vec![0f32; b * rows];
9065        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
9066        let mut want = vec![0f32; b * rows];
9067        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
9068
9069        let mut w = vec![0f32; rows * cols];
9070        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9071        for t in 0..b {
9072            for r in 0..rows {
9073                let mag: f32 = (0..cols)
9074                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
9075                    .sum();
9076                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
9077                assert!(
9078                    (g - wa).abs() <= 1e-3 * mag,
9079                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
9080                );
9081            }
9082        }
9083    }
9084
9085    #[test]
9086    fn q4tp_matvec2_matches_the_single_stream_kernel() {
9087        let (rows, cols) = (128usize, 256usize);
9088        let gpr = cols / GROUP_SIZE;
9089        let bytes = synth_q4tp(rows, cols);
9090        let xs: Vec<f32> = (0..2 * cols)
9091            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9092            .collect();
9093
9094        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
9095        q4tp_matvec2(
9096            &bytes,
9097            &xs[..cols],
9098            &xs[cols..],
9099            rows,
9100            cols,
9101            &mut o1,
9102            &mut o2,
9103            None,
9104        );
9105
9106        // matvec2 takes the exact path for both streams, so the single-row
9107        // kernel is an exact reference — no tolerance for path differences.
9108        let v = Q4tpView::new(&bytes, rows, cols);
9109        let mut sc = vec![0f32; gpr];
9110        for r in 0..rows {
9111            v.scales_into(r, gpr, &mut sc);
9112            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
9113            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
9114        }
9115    }
9116
9117    /// q4tp must not COST speed — it exists to save bytes, and a format that
9118    /// trades 7% of a file for a slower model is a bad trade. This guard is
9119    /// here because correctness tests happily passed while `q4tp_matmat` was
9120    /// missing its int8 and Accelerate arms and the model ran 5x slower.
9121    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
9122    /// aligned than q4t's 18 B, which pays for the scale indirection).
9123    #[test]
9124    fn q4tp_matvec_keeps_pace_with_q4t() {
9125        let (rows, cols) = (4096usize, 3072usize);
9126        let bytes = synth_q4tp(rows, cols);
9127        let twin = q4tp_as_q4t(&bytes, rows, cols);
9128        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
9129        let mut o = vec![0f32; rows];
9130        let n = 12;
9131        let mut best = (f64::MAX, f64::MAX);
9132        // Interleaved A/B, minimum statistic: this machine throttles, and a
9133        // mean over a thermal ramp reliably indicts whichever ran second.
9134        for _ in 0..3 {
9135            let t0 = std::time::Instant::now();
9136            for _ in 0..n {
9137                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
9138            }
9139            best.0 = best.0.min(t0.elapsed().as_secs_f64());
9140            let t0 = std::time::Instant::now();
9141            for _ in 0..n {
9142                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
9143            }
9144            best.1 = best.1.min(t0.elapsed().as_secs_f64());
9145        }
9146        let ratio = best.1 / best.0;
9147        println!("q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x", best.0 * 1e3 / n as f64, best.1 * 1e3 / n as f64);
9148        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
9149    }
9150
9151    #[cfg(target_os = "macos")]
9152    #[test]
9153    fn q4t_matmat_accel_matches_dequant_reference() {
9154        if !accel_gemm_enabled() {
9155            return; // CMF_ACCEL=0
9156        }
9157        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
9158        let gpr = cols / GROUP_SIZE;
9159        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9160        for r in 0..rows {
9161            for g in 0..gpr {
9162                let t = (r * gpr + g) * Q4_TILE;
9163                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
9164                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9165                for k in 0..16 {
9166                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9167                }
9168            }
9169        }
9170        let x: Vec<f32> = (0..b * cols)
9171            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9172            .collect();
9173        let mut got = vec![0f32; b * rows];
9174        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
9175        // Brute-force reference off the same tiles.
9176        let mut w = vec![0f32; rows * cols];
9177        for r in 0..rows {
9178            for g in 0..gpr {
9179                let t = (r * gpr + g) * Q4_TILE;
9180                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
9181                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
9182                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
9183                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
9184                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
9185                }
9186            }
9187        }
9188        for bi in 0..b {
9189            for r in 0..rows {
9190                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
9191                let d = (got[bi * rows + r] - want).abs();
9192                assert!(
9193                    d <= want.abs().max(1.0) * 1e-4,
9194                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
9195                    got[bi * rows + r]
9196                );
9197            }
9198        }
9199    }
9200
9201    #[test]
9202    fn q4matvec_matches_full_dequant() {
9203        let (rows, cols) = (8, 64);
9204        let groups = rows * cols / GROUP_SIZE;
9205        // Hand-craft a q4_block blob: nibbles then f16 scales.
9206        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9207        for i in 0..groups * 16 {
9208            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9209        }
9210        for g in 0..groups {
9211            let s = 0.01 + 0.003 * g as f32;
9212            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9213        }
9214        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9215
9216        let mut reference = vec![0.0f32; rows * cols];
9217        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9218        let mut expect = vec![0.0f32; rows];
9219        for r in 0..rows {
9220            expect[r] = reference[r * cols..(r + 1) * cols]
9221                .iter()
9222                .zip(&x)
9223                .map(|(w, xv)| w * xv)
9224                .sum();
9225        }
9226
9227        let mut got = vec![0.0f32; rows];
9228        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9229        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9230        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
9231        // in the golden-parity gate).
9232        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9233        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9234        for r in 0..rows {
9235            assert!(
9236                (got[r] - expect[r]).abs() < tol * scale,
9237                "row {r}: {} vs {}",
9238                got[r],
9239                expect[r]
9240            );
9241        }
9242    }
9243
9244    /// Fused two-input vbit matvec must equal two single matvecs exactly
9245    /// (same per-lane accumulation order on both scalar and SDOT paths).
9246    #[test]
9247    fn vbitmatvec2_equals_two_singles() {
9248        let (rows, cols) = (6, 64);
9249        let ng = cols / GROUP_SIZE;
9250        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9251        let mut bytes = bits.clone();
9252        for g in 0..rows * ng {
9253            let s = 0.02 + 0.001 * g as f32;
9254            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9255        }
9256        for r in 0..rows {
9257            let b = bits[r] as usize;
9258            let (mut acc, mut nb) = (0u64, 0usize);
9259            let mut rowbytes = Vec::new();
9260            for i in 0..cols {
9261                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9262                acc = (acc << b) | v;
9263                nb += b;
9264                while nb >= 8 {
9265                    nb -= 8;
9266                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9267                }
9268            }
9269            if nb > 0 {
9270                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9271            }
9272            bytes.extend_from_slice(&rowbytes);
9273        }
9274        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9275        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
9276        let offsets = vbit_row_offsets(&bytes, rows, cols);
9277
9278        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9279        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
9280        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
9281        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9282        vbitmatvec2(
9283            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
9284        );
9285        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
9286        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
9287    }
9288
9289    /// Fused two-input q4 matvec must equal two single matvecs exactly.
9290    #[test]
9291    fn q4matvec2_equals_two_singles() {
9292        let (rows, cols) = (8, 128);
9293        let groups = rows * cols / GROUP_SIZE;
9294        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9295        for i in 0..groups * 16 {
9296            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9297        }
9298        for g in 0..groups {
9299            let s = 0.01 + 0.003 * g as f32;
9300            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9301        }
9302        // Include an outlier channel so the SDOT correction path is
9303        // exercised in the pair kernel too.
9304        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9305        x1[9] = 250.0;
9306        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9307
9308        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9309        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
9310        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
9311        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9312        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
9313        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
9314        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
9315    }
9316
9317    /// Multi-matrix job must equal separate matvecs exactly — same
9318    /// kernels, only the dispatch is fused.
9319    #[test]
9320    fn matvec_many_equals_separate_matvecs() {
9321        use crate::pool::Pool;
9322        let (r1, r2, cols) = (300, 200, 64);
9323        let mk = |salt: usize, rows: usize| {
9324            QTensor::from_f32(
9325                (0..rows * cols)
9326                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
9327                    .collect(),
9328                rows,
9329                cols,
9330            )
9331        };
9332        let (a, b) = (mk(1, r1), mk(5, r2));
9333        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
9334        let pool = Pool::new(3);
9335
9336        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
9337        a.matvec(&x, &mut ea, Some(&pool));
9338        b.matvec(&x, &mut eb, Some(&pool));
9339        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
9340        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
9341        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
9342        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
9343    }
9344
9345    /// Batched q4/vbit matmat must equal per-position matvec calls
9346    /// exactly (the fallback it replaced) — same kernels, same order.
9347    #[test]
9348    fn batched_matmat_equals_per_position_matvec() {
9349        let (rows, cols, b) = (8, 64, 5);
9350        // q4 blob.
9351        let groups = rows * cols / GROUP_SIZE;
9352        let mut q4 = Vec::new();
9353        for i in 0..groups * 16 {
9354            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9355        }
9356        for g in 0..groups {
9357            q4.extend_from_slice(
9358                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9359            );
9360        }
9361        // vbit blob (mixed widths incl. 8).
9362        let ng = cols / GROUP_SIZE;
9363        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
9364        let mut vb = bits.clone();
9365        for g in 0..rows * ng {
9366            vb.extend_from_slice(
9367                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
9368            );
9369        }
9370        for r in 0..rows {
9371            let bw = bits[r] as usize;
9372            let (mut acc, mut nb) = (0u64, 0usize);
9373            let mut rowbytes = Vec::new();
9374            for i in 0..cols {
9375                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9376                acc = (acc << bw) | v;
9377                nb += bw;
9378                while nb >= 8 {
9379                    nb -= 8;
9380                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9381                }
9382            }
9383            if nb > 0 {
9384                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9385            }
9386            vb.extend_from_slice(&rowbytes);
9387        }
9388        let offsets = vbit_row_offsets(&vb, rows, cols);
9389
9390        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9391
9392        // q4: batch vs singles.
9393        let mut got = vec![0f32; b * rows];
9394        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
9395        for bi in 0..b {
9396            let mut expect = vec![0f32; rows];
9397            q4matvec(
9398                &q4,
9399                &xs[bi * cols..(bi + 1) * cols],
9400                rows,
9401                cols,
9402                &mut expect,
9403                None,
9404            );
9405            assert_eq!(
9406                &got[bi * rows..(bi + 1) * rows],
9407                &expect[..],
9408                "q4 batch pos {bi}"
9409            );
9410        }
9411
9412        // vbit: batch vs singles.
9413        let mut got = vec![0f32; b * rows];
9414        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
9415        for bi in 0..b {
9416            let mut expect = vec![0f32; rows];
9417            vbitmatvec(
9418                &vb,
9419                &offsets,
9420                &xs[bi * cols..(bi + 1) * cols],
9421                rows,
9422                cols,
9423                &mut expect,
9424                None,
9425            );
9426            assert_eq!(
9427                &got[bi * rows..(bi + 1) * rows],
9428                &expect[..],
9429                "vbit batch pos {bi}"
9430            );
9431        }
9432    }
9433
9434    /// q4_tiled kernels must produce BIT-identical outputs to the q4
9435    /// split kernels on the same values (same ints, same order — only
9436    /// the byte placement differs).
9437    #[test]
9438    fn q4_tiled_matches_q4_block_bitexact() {
9439        let (rows, cols, b) = (8usize, 128usize, 3usize);
9440        let groups = rows * cols / GROUP_SIZE;
9441        let mut split = Vec::with_capacity(groups * 18);
9442        for i in 0..groups * 16 {
9443            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9444        }
9445        for g in 0..groups {
9446            split.extend_from_slice(
9447                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9448            );
9449        }
9450        // Re-tile: [scale][nibbles] per group.
9451        let (packed, scales) = split.split_at(groups * 16);
9452        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
9453        for g in 0..groups {
9454            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
9455            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
9456        }
9457
9458        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9459        x1[9] = 250.0; // exercise the outlier path
9460        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9461
9462        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
9463        q4matvec(&split, &x1, rows, cols, &mut a, None);
9464        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
9465        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
9466
9467        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9468        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
9469        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
9470        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
9471        assert_eq!(a1, t1);
9472        assert_eq!(a2, t2);
9473
9474        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9475        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
9476        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
9477        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
9478        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
9479    }
9480
9481    /// q4 SDOT outlier correction: a single huge activation channel
9482    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
9483    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
9484    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
9485    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
9486    /// can never qualify (8² = n).
9487    #[test]
9488    fn q4matvec_sdot_outlier_exact() {
9489        let (rows, cols) = (4, 128);
9490        let groups = rows * cols / GROUP_SIZE;
9491        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9492        for i in 0..groups * 16 {
9493            bytes.push(((i * 11 + 5) % 256) as u8);
9494        }
9495        for g in 0..groups {
9496            let s = 0.02 + 0.002 * g as f32;
9497            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9498        }
9499        let mut x: Vec<f32> = (0..cols)
9500            .map(|i| match i % 3 {
9501                0 => 1.0,
9502                1 => -1.0,
9503                _ => 0.0,
9504            })
9505            .collect();
9506        x[17] = 300.0; // ≫ 8·rms → outlier channel
9507
9508        let mut reference = vec![0.0f32; rows * cols];
9509        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9510        let mut expect = vec![0.0f32; rows];
9511        for r in 0..rows {
9512            expect[r] = reference[r * cols..(r + 1) * cols]
9513                .iter()
9514                .zip(&x)
9515                .map(|(w, xv)| w * xv)
9516                .sum();
9517        }
9518        let mut got = vec![0.0f32; rows];
9519        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9520        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9521        for r in 0..rows {
9522            assert!(
9523                (got[r] - expect[r]).abs() < 2e-3 * scale,
9524                "row {r}: {} vs {} (outlier term must be exact)",
9525                got[r],
9526                expect[r]
9527            );
9528        }
9529    }
9530
9531    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
9532    /// including the ternary zero level and the binary-searched outlier
9533    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
9534    #[test]
9535    fn q1t_matvec_matches_reference() {
9536        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
9537        let (rows, cols) = (3usize, 64usize); // gpr = 2
9538        let gpr = cols / GROUP_SIZE;
9539        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
9540        // Overlay (must be sorted by flat index): a few spikes across rows.
9541        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
9542        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
9543        let mut bytes = Vec::new();
9544        for r in 0..rows {
9545            for g in 0..gpr {
9546                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
9547                let mut c = [0u8; 7];
9548                for k in 0..GROUP_SIZE {
9549                    // Encoder invariant: code 0 at outlier positions.
9550                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
9551                        0
9552                    } else {
9553                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
9554                    };
9555                    cortiq_core::quant::q1t_pack(&mut c, k, code);
9556                }
9557                bytes.extend_from_slice(&c);
9558            }
9559        }
9560        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
9561        // row (outliers are sorted by flat index → already grouped by row).
9562        let mut row_ptr = vec![0u32; rows + 1];
9563        for &(idx, _) in &outliers {
9564            row_ptr[idx as usize / cols + 1] += 1;
9565        }
9566        for r in 0..rows {
9567            row_ptr[r + 1] += row_ptr[r];
9568        }
9569        for &p in &row_ptr {
9570            bytes.extend_from_slice(&p.to_le_bytes());
9571        }
9572        for &(idx, v) in &outliers {
9573            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
9574            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
9575        }
9576
9577        let mut refw = vec![0f32; rows * cols];
9578        dequant_q1t(&bytes, rows, cols, &mut refw);
9579        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
9580        // x exactly and matches the f32 reference (same trick as the q1 test).
9581        let x: Vec<f32> = (0..cols)
9582            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9583            .collect();
9584        let mut expect = vec![0f32; rows];
9585        for r in 0..rows {
9586            let mut a = 0.0f32;
9587            for j in 0..cols {
9588                a += refw[r * cols + j] * x[j];
9589            }
9590            expect[r] = a;
9591        }
9592        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
9593        let mut got = vec![0f32; rows];
9594        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
9595        for r in 0..rows {
9596            assert!(
9597                (got[r] - expect[r]).abs() < tol(expect[r]),
9598                "row {r}: {} vs {}",
9599                got[r],
9600                expect[r]
9601            );
9602        }
9603        // matmat (b=2, f32 decode path) must agree too.
9604        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
9605        let mut gm = vec![0f32; 2 * rows];
9606        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
9607        for r in 0..rows {
9608            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
9609            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
9610        }
9611        // Fused pair (q1t_matvec2) must equal two single matvecs
9612        // bit-for-bit: same unpack, same group order, same f32
9613        // accumulation per stream. Distinct x2 exercises both lanes.
9614        let xb: Vec<f32> = (0..cols)
9615            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
9616            .collect();
9617        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9618        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
9619        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
9620        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9621        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
9622        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
9623        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
9624    }
9625
9626    /// Pair == 2×matvec with an ODD group count (the kernel's tail
9627    /// group) and no overlay section.
9628    #[test]
9629    fn q1t_matvec2_odd_gpr_matches_singles() {
9630        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
9631        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
9632        let gpr = cols / GROUP_SIZE;
9633        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
9634        for r in 0..rows {
9635            for g in 0..gpr {
9636                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
9637                let mut c = [0u8; 7];
9638                for k in 0..GROUP_SIZE {
9639                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
9640                }
9641                bytes.extend_from_slice(&c);
9642            }
9643        }
9644        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
9645        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
9646        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9647        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9648        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
9649        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9650        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9651        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
9652        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
9653    }
9654
9655    // Speed A/B: fused pair (one unpack, two streams) vs two single
9656    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
9657    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
9658    #[test]
9659    #[ignore]
9660    fn q1t_matvec2_speed() {
9661        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
9662        use std::time::Instant;
9663        let (rows, cols) = (8192usize, 4096usize);
9664        let gpr = cols / GROUP_SIZE;
9665        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
9666        for r in 0..rows {
9667            for g in 0..gpr {
9668                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
9669                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
9670                let mut c = [0u8; 7];
9671                for k in 0..GROUP_SIZE {
9672                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
9673                }
9674                bytes.extend_from_slice(&c);
9675            }
9676        }
9677        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
9678        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
9679        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9680        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9681        // Warm both paths once.
9682        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9683        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9684        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
9685        for _ in 0..8 {
9686            let t0 = Instant::now();
9687            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9688            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
9689            let t1 = Instant::now();
9690            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9691            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
9692            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
9693        }
9694        assert_eq!(p1, s1);
9695        assert_eq!(p2, s2);
9696        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
9697    }
9698
9699    // Speed A/B: the base-3-division decode (what the packing commit left in
9700    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
9701    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
9702    #[test]
9703    #[ignore]
9704    fn q1t_matvec_speed() {
9705        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
9706        use std::time::Instant;
9707        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
9708        let gpr = cols / GROUP_SIZE;
9709        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
9710        for r in 0..rows {
9711            for g in 0..gpr {
9712                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
9713                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
9714                let mut c = [0u8; 7];
9715                for k in 0..GROUP_SIZE {
9716                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
9717                }
9718                bytes.extend_from_slice(&c);
9719            }
9720        }
9721        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
9722        let mut row_ptr = vec![0u32; rows + 1];
9723        let mut idx = 0usize;
9724        while idx < n {
9725            row_ptr[idx / cols + 1] += 1;
9726            idx += stride;
9727        }
9728        for r in 0..rows {
9729            row_ptr[r + 1] += row_ptr[r];
9730        }
9731        for &p in &row_ptr {
9732            bytes.extend_from_slice(&p.to_le_bytes());
9733        }
9734        let mut idx = 0usize;
9735        while idx < n {
9736            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
9737            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
9738            idx += stride;
9739        }
9740        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
9741        // reference (the A/B is a timing check; values must still agree).
9742        let x: Vec<f32> = (0..cols)
9743            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9744            .collect();
9745        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
9746
9747        // "before": base-3 division decode into a buffer, then dot.
9748        let slow = |out: &mut [f32]| {
9749            let mut buf = vec![0f32; cols];
9750            for r in 0..rows {
9751                for g in 0..gpr {
9752                    let off = (r * gpr + g) * Q1T_TILE;
9753                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9754                    let codes = &bytes[off + 2..off + Q1T_TILE];
9755                    for k in 0..GROUP_SIZE {
9756                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
9757                            1 => s,
9758                            2 => -s,
9759                            _ => 0.0,
9760                        };
9761                    }
9762                }
9763                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
9764                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
9765            }
9766        };
9767        let iters = 5;
9768        let mut a = vec![0f32; rows];
9769        slow(&mut a); // warm
9770        let t = Instant::now();
9771        for _ in 0..iters {
9772            slow(&mut a);
9773        }
9774        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
9775
9776        let mut b = vec![0f32; rows];
9777        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
9778        let t = Instant::now();
9779        for _ in 0..iters {
9780            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
9781        }
9782        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
9783
9784        for r in 0..rows {
9785            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
9786        }
9787        println!(
9788            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
9789            slow_ms / fast_ms
9790        );
9791    }
9792}
cortiq_engine/qtensor.rs

cortiq_engine/
qtensor.rs
