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    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2203    ///
2204    /// The per-expert path pays a pool barrier per expert per stage: at 9
2205    /// experts over 40 layers that is ~720 barriers a token, and a decode
2206    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2207    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2208    /// every expert's rows end-to-end in one virtual row space collapses
2209    /// the stage to a single dispatch. The per-row body is the
2210    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2211    ///
2212    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2213    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2214    /// per-expert path.
2215    pub fn moe_gate_up_many(
2216        pairs: &[(&QTensor, &QTensor)],
2217        x: &[f32],
2218        outs: &mut [Vec<f32>],
2219        pool: Option<&Pool>,
2220    ) -> bool {
2221        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2222            return false;
2223        }
2224        let inter = pairs[0].0.rows();
2225        let cols = pairs[0].0.cols();
2226        if cols % GROUP_SIZE != 0 {
2227            return false;
2228        }
2229        let gpr = cols / GROUP_SIZE;
2230        let mut views = Vec::with_capacity(pairs.len() * 2);
2231        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2232            let both_q4tp = matches!(
2233                g,
2234                Self::Mapped {
2235                    dtype: TensorDtype::Q4TiledP,
2236                    ..
2237                }
2238            ) && matches!(
2239                u,
2240                Self::Mapped {
2241                    dtype: TensorDtype::Q4TiledP,
2242                    ..
2243                }
2244            );
2245            if !both_q4tp
2246                || g.rows() != inter
2247                || u.rows() != inter
2248                || g.cols() != cols
2249                || u.cols() != cols
2250                || o.len() != inter
2251            {
2252                return false;
2253            }
2254            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
2255            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
2256        }
2257        let act = split_act(x);
2258        let act = &act;
2259        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2260        let (views, ptrs) = (&views, &ptrs);
2261        let run = |start: usize, end: usize| {
2262            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2263            for flat in start..end {
2264                let (e, r) = (flat / inter, flat % inter);
2265                let gv_view = &views[e * 2];
2266                let uv_view = &views[e * 2 + 1];
2267                gv_view.scales_into(r, gpr, &mut gsc);
2268                uv_view.scales_into(r, gpr, &mut usc);
2269                let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2270                let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2271                for &(j, xv) in &act.outliers {
2272                    let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2273                    gv += w * s * xv;
2274                    let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2275                    uv += w * s * xv;
2276                }
2277                let silu_g = gv / (1.0 + (-gv).exp());
2278                // SAFETY: one worker owns each (expert, row) pair.
2279                unsafe { *ptrs[e].at(r) = silu_g * uv };
2280            }
2281        };
2282        dispatch_rows(pool, pairs.len() * inter, &run);
2283        true
2284    }
2285
2286    /// Every routed expert's down projection, weighted and summed into
2287    /// `out`, under ONE pool dispatch.
2288    ///
2289    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2290    /// by a single worker, so the experts are summed in the caller's order
2291    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2292    /// performs, hence bit-identical. Partitioning by expert instead would
2293    /// race on the shared accumulator.
2294    pub fn moe_down_many(
2295        downs: &[&QTensor],
2296        gs: &[Vec<f32>],
2297        weights: &[f32],
2298        out: &mut [f32],
2299        pool: Option<&Pool>,
2300    ) -> bool {
2301        if downs.is_empty()
2302            || downs.len() != gs.len()
2303            || downs.len() != weights.len()
2304            || !a8w8_enabled()
2305        {
2306            return false;
2307        }
2308        let rows = out.len();
2309        let cols = downs[0].cols();
2310        if cols % GROUP_SIZE != 0 {
2311            return false;
2312        }
2313        let gpr = cols / GROUP_SIZE;
2314        let mut views = Vec::with_capacity(downs.len());
2315        for (d, g) in downs.iter().zip(gs.iter()) {
2316            if !matches!(
2317                d,
2318                Self::Mapped {
2319                    dtype: TensorDtype::Q4TiledP,
2320                    ..
2321                }
2322            ) || d.rows() != rows
2323                || d.cols() != cols
2324                || g.len() != cols
2325            {
2326                return false;
2327            }
2328            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2329        }
2330        // One int8 split per expert — the activation vectors differ.
2331        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2332        // Partitioned by OUTPUT row, with the experts folded inside: each
2333        // row is owned by one worker, so they are summed in the caller's
2334        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2335        // loop produces. Partitioning by expert instead would either race
2336        // on the accumulator or need a scratch plane and a second pass;
2337        // measured, that variant was a wash, so this keeps the simpler
2338        // shape.
2339        let out_addr = SendMut(out.as_mut_ptr());
2340        let (views, acts, weights) = (&views, &acts, &weights);
2341        let run = |start: usize, end: usize| {
2342            let mut sc = vec![0f32; gpr];
2343            for r in start..end {
2344                let mut acc = 0f32;
2345                for (e, v) in views.iter().enumerate() {
2346                    v.scales_into(r, gpr, &mut sc);
2347                    let a = &acts[e];
2348                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2349                    for &(j, xv) in &a.outliers {
2350                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2351                        d += w * s * xv;
2352                    }
2353                    acc += weights[e] * d;
2354                }
2355                // SAFETY: disjoint row ranges per worker.
2356                unsafe { *out_addr.at(r) = acc };
2357            }
2358        };
2359        dispatch_rows(pool, rows, &run);
2360        true
2361    }
2362}
2363
2364/// Batched q8 kernel: same math as qmatvec, the row makes a single
2365/// pass from memory for the whole batch.
2366/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2367/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2368#[cfg(target_os = "macos")]
2369mod accel_blas {
2370    #[link(name = "Accelerate", kind = "framework")]
2371    unsafe extern "C" {
2372        pub fn cblas_sgemm(
2373            order: i32,
2374            trans_a: i32,
2375            trans_b: i32,
2376            m: i32,
2377            n: i32,
2378            k: i32,
2379            alpha: f32,
2380            a: *const f32,
2381            lda: i32,
2382            b: *const f32,
2383            ldb: i32,
2384            beta: f32,
2385            c: *mut f32,
2386            ldc: i32,
2387        );
2388    }
2389}
2390
2391#[cfg(target_os = "macos")]
2392pub(crate) fn accel_gemm_enabled() -> bool {
2393    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2394    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2395}
2396
2397/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2398/// same entry point, so the batched-attention path opens on mobile.
2399#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2400pub(crate) fn accel_gemm_enabled() -> bool {
2401    true
2402}
2403
2404/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2405/// micro-kernel with A broadcast against B panels — the mobile stand-in
2406/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2407/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2408/// k = head_dim or context), and the goal is removing the per-position
2409/// quadratic wall, not peak GEMM.
2410#[cfg(target_arch = "aarch64")]
2411#[allow(clippy::too_many_arguments)]
2412pub(crate) fn neon_gemm_rm(
2413    m: usize,
2414    n: usize,
2415    k: usize,
2416    alpha: f32,
2417    a: &[f32],
2418    lda: usize,
2419    b_mat: &[f32],
2420    ldb: usize,
2421    b_rows_are_n: bool,
2422    c: &mut [f32],
2423    ldc: usize,
2424) {
2425    debug_assert!(a.len() >= (m - 1) * lda + k);
2426    debug_assert!(c.len() >= (m - 1) * ldc + n);
2427    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2428    unsafe {
2429        use core::arch::aarch64::*;
2430        let mut i = 0usize;
2431        while i < m {
2432            let mi = (m - i).min(4);
2433            let mut j = 0usize;
2434            while j < n {
2435                let nj = (n - j).min(8);
2436                if mi == 4 && nj == 8 {
2437                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2438                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2439                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2440                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2441                    for p in 0..k {
2442                        let (b0, b1) = if b_rows_are_n {
2443                            // B is [n, k]: column p of Bᵀ = element p of
2444                            // eight consecutive B rows — gathered.
2445                            let base = b_mat.as_ptr().add(j * ldb + p);
2446                            let g = |o: usize| *base.add(o * ldb);
2447                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2448                        } else {
2449                            let base = b_mat.as_ptr().add(p * ldb + j);
2450                            (
2451                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2452                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2453                            )
2454                        };
2455                        let bv0 = vld1q_f32(b0.as_ptr());
2456                        let bv1 = vld1q_f32(b1.as_ptr());
2457                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2458                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2459                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2460                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2461                        c0a = vfmaq_f32(c0a, a0, bv0);
2462                        c0b = vfmaq_f32(c0b, a0, bv1);
2463                        c1a = vfmaq_f32(c1a, a1, bv0);
2464                        c1b = vfmaq_f32(c1b, a1, bv1);
2465                        c2a = vfmaq_f32(c2a, a2, bv0);
2466                        c2b = vfmaq_f32(c2b, a2, bv1);
2467                        c3a = vfmaq_f32(c3a, a3, bv0);
2468                        c3b = vfmaq_f32(c3b, a3, bv1);
2469                    }
2470                    let al = vdupq_n_f32(alpha);
2471                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2472                        .iter()
2473                        .enumerate()
2474                    {
2475                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2476                        vst1q_f32(dst, vmulq_f32(*ca, al));
2477                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2478                    }
2479                } else {
2480                    for r in 0..mi {
2481                        for q in 0..nj {
2482                            let mut acc = 0f32;
2483                            for p in 0..k {
2484                                let bv = if b_rows_are_n {
2485                                    b_mat[(j + q) * ldb + p]
2486                                } else {
2487                                    b_mat[p * ldb + j + q]
2488                                };
2489                                acc += a[(i + r) * lda + p] * bv;
2490                            }
2491                            c[(i + r) * ldc + j + q] = acc * alpha;
2492                        }
2493                    }
2494                }
2495                j += nj;
2496            }
2497            i += mi;
2498        }
2499    }
2500}
2501
2502/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2503#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2504#[allow(clippy::too_many_arguments)]
2505pub(crate) fn sgemm_rm(
2506    m: usize,
2507    n: usize,
2508    k: usize,
2509    alpha: f32,
2510    a: &[f32],
2511    lda: usize,
2512    b_mat: &[f32],
2513    ldb: usize,
2514    b_rows_are_n: bool,
2515    c: &mut [f32],
2516    ldc: usize,
2517) {
2518    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2519}
2520
2521/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
2522/// per-layer projection and applies it to every expert; a naive triple loop
2523/// would turn a two-minute job into half an hour).
2524#[allow(clippy::too_many_arguments)]
2525pub fn sgemm_public(
2526    m: usize,
2527    n: usize,
2528    k: usize,
2529    alpha: f32,
2530    a: &[f32],
2531    lda: usize,
2532    b_mat: &[f32],
2533    ldb: usize,
2534    b_rows_are_n: bool,
2535    c: &mut [f32],
2536    ldc: usize,
2537) {
2538    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
2539    {
2540        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2541    }
2542    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
2543    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
2544    // this, so correctness matters and throughput does not — a triple loop is
2545    // the honest fallback rather than a reason to make the tool macOS-only.
2546    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
2547    {
2548        for i in 0..m {
2549            for j in 0..n {
2550                let mut acc = 0f32;
2551                for p in 0..k {
2552                    let bv = if b_rows_are_n {
2553                        b_mat[j * ldb + p]
2554                    } else {
2555                        b_mat[p * ldb + j]
2556                    };
2557                    acc += a[i * lda + p] * bv;
2558                }
2559                c[i * ldc + j] = alpha * acc;
2560            }
2561        }
2562    }
2563}
2564
2565/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
2566/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
2567#[cfg(target_os = "macos")]
2568#[allow(clippy::too_many_arguments)]
2569pub(crate) fn sgemm_rm(
2570    m: usize,
2571    n: usize,
2572    k: usize,
2573    alpha: f32,
2574    a: &[f32],
2575    lda: usize,
2576    b_mat: &[f32],
2577    ldb: usize,
2578    b_rows_are_n: bool,
2579    c: &mut [f32],
2580    ldc: usize,
2581) {
2582    debug_assert!(a.len() >= (m - 1) * lda + k);
2583    debug_assert!(c.len() >= (m - 1) * ldc + n);
2584    // Test hook: route the attention GEMMs through the portable NEON
2585    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
2586    // measured without a phone in the loop. (Intel macOS has no NEON —
2587    // the hook is a no-op there, Accelerate continues below.)
2588    #[cfg(target_arch = "aarch64")]
2589    if std::env::var("CMF_FORCE_NEON_GEMM")
2590        .map(|v| v == "1")
2591        .unwrap_or(false)
2592    {
2593        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2594    }
2595    unsafe {
2596        accel_blas::cblas_sgemm(
2597            101, // RowMajor
2598            111, // NoTrans A
2599            if b_rows_are_n { 112 } else { 111 },
2600            m as i32,
2601            n as i32,
2602            k as i32,
2603            alpha,
2604            a.as_ptr(),
2605            lda as i32,
2606            b_mat.as_ptr(),
2607            ldb as i32,
2608            0.0,
2609            c.as_mut_ptr(),
2610            ldc as i32,
2611        );
2612    }
2613}
2614
2615/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
2616/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
2617/// on the AMX with one row-major sgemm. Tiles live in cache, weights
2618/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
2619/// logits shift within f32 rounding — tolerance-class, like every
2620/// reduction-order change; decode (M=1) never takes this path.
2621#[cfg(target_os = "macos")]
2622fn qmatmat_accel(
2623    q: &[u8],
2624    row_scale: &[f32],
2625    pre: &[std::borrow::Cow<'_, [f32]>],
2626    rows: usize,
2627    cols: usize,
2628    out: &mut [f32],
2629    pool: Option<&Pool>,
2630) {
2631    // NOTE: double-buffering the dequant against the sgemm (a scoped
2632    // thread driving the pool on tile k+1 while the caller multiplies
2633    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
2634    // multithreaded, and the dequant workers just steal its cores.
2635    const TR: usize = 2048;
2636    let b = pre.len();
2637    thread_local! {
2638        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2639        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2640    }
2641    XPANEL.with(|xp| {
2642        WTILE.with(|wt| {
2643            let mut xpanel = xp.borrow_mut();
2644            xpanel.clear();
2645            for x in pre {
2646                xpanel.extend_from_slice(x);
2647            }
2648            let mut wtile = wt.borrow_mut();
2649            wtile.resize(TR * cols, 0.0);
2650            let mut r0 = 0usize;
2651            while r0 < rows {
2652                let tr = TR.min(rows - r0);
2653                // Dequant the tile (scale folded) — pool-parallel.
2654                let wt_addr = SendMut(wtile.as_mut_ptr());
2655                let run = |start: usize, end: usize| {
2656                    for r in start..end {
2657                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2658                        let s = row_scale[r0 + r];
2659                        // SAFETY: workers cover disjoint r ranges.
2660                        let dst =
2661                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2662                        for (d, &v) in dst.iter_mut().zip(row) {
2663                            *d = (v as i8) as f32 * s;
2664                        }
2665                    }
2666                };
2667                dispatch_rows(pool, tr, &run);
2668                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
2669                unsafe {
2670                    accel_blas::cblas_sgemm(
2671                        101, // RowMajor
2672                        111, // NoTrans A
2673                        112, // Trans B
2674                        b as i32,
2675                        tr as i32,
2676                        cols as i32,
2677                        1.0,
2678                        xpanel.as_ptr(),
2679                        cols as i32,
2680                        wtile.as_ptr(),
2681                        cols as i32,
2682                        0.0,
2683                        out.as_mut_ptr().add(r0),
2684                        rows as i32,
2685                    );
2686                }
2687                r0 += tr;
2688            }
2689        })
2690    });
2691}
2692
2693fn qmatmat(
2694    q: &[u8],
2695    row_scale: &[f32],
2696    pre: &[std::borrow::Cow<'_, [f32]>],
2697    rows: usize,
2698    cols: usize,
2699    out: &mut [f32],
2700    pool: Option<&Pool>,
2701) {
2702    let b = pre.len();
2703    debug_assert_eq!(out.len(), b * rows);
2704    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
2705    // SDOT loop below peaks near the CPU's dot throughput, an order
2706    // below the matrix units. Small tensors and tiny test models stay
2707    // on the exact integer path.
2708    #[cfg(target_os = "macos")]
2709    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2710        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2711        return;
2712    }
2713    #[cfg(target_arch = "aarch64")]
2714    if sdot_enabled() {
2715        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2716        let out_addr = SendMut(out.as_mut_ptr());
2717        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
2718        // path IS the ARM prefill GEMM off Apple silicon).
2719        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2720            .map(|v| v != "0")
2721            .unwrap_or(true);
2722        let use_i8mm = i8mm_enabled();
2723        if blocked_ok {
2724            let run = |start: usize, end: usize| {
2725                let mut o = start;
2726                while o < end {
2727                    if o + 2 <= end {
2728                        let r0 = &q[o * cols..(o + 1) * cols];
2729                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2730                        let mut bi = 0usize;
2731                        while bi + 4 <= acts.len() {
2732                            let xs = [
2733                                acts[bi].xq.as_slice(),
2734                                acts[bi + 1].xq.as_slice(),
2735                                acts[bi + 2].xq.as_slice(),
2736                                acts[bi + 3].xq.as_slice(),
2737                            ];
2738                            let d = if use_i8mm {
2739                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2740                            } else {
2741                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2742                            };
2743                            for (r, row) in [r0, r1].into_iter().enumerate() {
2744                                for k in 0..4 {
2745                                    let act = &acts[bi + k];
2746                                    let mut v = d[r][k] as f32 * act.sx;
2747                                    for &(j, xv) in &act.outliers {
2748                                        v += (row[j] as i8) as f32 * xv;
2749                                    }
2750                                    unsafe {
2751                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2752                                    };
2753                                }
2754                            }
2755                            bi += 4;
2756                        }
2757                        while bi < acts.len() {
2758                            for (r, row) in [r0, r1].into_iter().enumerate() {
2759                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2760                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2761                            }
2762                            bi += 1;
2763                        }
2764                        o += 2;
2765                    } else {
2766                        let row = &q[o * cols..(o + 1) * cols];
2767                        for (bi, act) in acts.iter().enumerate() {
2768                            let v = row_dot_sdot(row, act) * row_scale[o];
2769                            unsafe { *out_addr.at(bi * rows + o) = v };
2770                        }
2771                        o += 1;
2772                    }
2773                }
2774            };
2775            dispatch_rows(pool, rows, &run);
2776            return;
2777        }
2778        let run = |start: usize, end: usize| {
2779            for o in start..end {
2780                let row = &q[o * cols..(o + 1) * cols];
2781                for (bi, act) in acts.iter().enumerate() {
2782                    let v = row_dot_sdot(row, act) * row_scale[o];
2783                    unsafe { *out_addr.at(bi * rows + o) = v };
2784                }
2785            }
2786        };
2787        dispatch_rows(pool, rows, &run);
2788        return;
2789    }
2790    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
2791    // (roadmap P0: two weight rows' abs() stay in registers across four
2792    // activation streams); VNNI machines keep the per-row bias-trick
2793    // dot, which is already throughput-bound there.
2794    #[cfg(target_arch = "x86_64")]
2795    if avx2_a8w8_enabled() {
2796        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2797        let out_addr = SendMut(out.as_mut_ptr());
2798        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
2799        // A/B on noisy shared-vCPU hosts).
2800        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2801            .map(|v| v != "0")
2802            .unwrap_or(true);
2803        if !avx512vnni_enabled() && blocked_ok {
2804            let run = |start: usize, end: usize| {
2805                let mut o = start;
2806                while o < end {
2807                    if o + 2 <= end {
2808                        let r0 = &q[o * cols..(o + 1) * cols];
2809                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2810                        let mut bi = 0usize;
2811                        while bi + 4 <= acts.len() {
2812                            let xs = [
2813                                acts[bi].xq.as_slice(),
2814                                acts[bi + 1].xq.as_slice(),
2815                                acts[bi + 2].xq.as_slice(),
2816                                acts[bi + 3].xq.as_slice(),
2817                            ];
2818                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
2819                            for (r, row) in [r0, r1].into_iter().enumerate() {
2820                                for k in 0..4 {
2821                                    let act = &acts[bi + k];
2822                                    let mut v = d[r][k] as f32 * act.sx;
2823                                    for &(j, xv) in &act.outliers {
2824                                        v += (row[j] as i8) as f32 * xv;
2825                                    }
2826                                    unsafe {
2827                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2828                                    };
2829                                }
2830                            }
2831                            bi += 4;
2832                        }
2833                        while bi < acts.len() {
2834                            for (r, row) in [r0, r1].into_iter().enumerate() {
2835                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
2836                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2837                            }
2838                            bi += 1;
2839                        }
2840                        o += 2;
2841                    } else {
2842                        let row = &q[o * cols..(o + 1) * cols];
2843                        for (bi, act) in acts.iter().enumerate() {
2844                            let v = row_dot_avx2(row, act) * row_scale[o];
2845                            unsafe { *out_addr.at(bi * rows + o) = v };
2846                        }
2847                        o += 1;
2848                    }
2849                }
2850            };
2851            dispatch_rows(pool, rows, &run);
2852            return;
2853        }
2854        let run = |start: usize, end: usize| {
2855            for o in start..end {
2856                let row = &q[o * cols..(o + 1) * cols];
2857                for (bi, act) in acts.iter().enumerate() {
2858                    let v = row_dot_avx2(row, act) * row_scale[o];
2859                    unsafe { *out_addr.at(bi * rows + o) = v };
2860                }
2861            }
2862        };
2863        dispatch_rows(pool, rows, &run);
2864        return;
2865    }
2866    let out_addr = SendMut(out.as_mut_ptr());
2867    let run = |start: usize, end: usize| {
2868        for o in start..end {
2869            let row = &q[o * cols..(o + 1) * cols];
2870            for (bi, x) in pre.iter().enumerate() {
2871                let mut acc = 0f32;
2872                for j in 0..cols {
2873                    acc += (row[j] as i8) as f32 * x[j];
2874                }
2875                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
2876            }
2877        }
2878    };
2879    dispatch_rows(pool, rows, &run);
2880}
2881
2882/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
2883/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
2884fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
2885    match pool {
2886        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
2887        _ => run(0, rows),
2888    }
2889}
2890
2891/// Split a q4_block blob into (packed nibbles, f16 group scales).
2892fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
2893    let groups = rows * cols / GROUP_SIZE;
2894    bytes.split_at(groups * 16)
2895}
2896
2897/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
2898/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
2899/// vbit packs MSB-first, so the HIGH nibble is the even element
2900/// (opposite of q4_block's lo-first interleave). Centering is u-7.
2901#[inline]
2902fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
2903    #[cfg(target_arch = "aarch64")]
2904    unsafe {
2905        return vbit_fill4_neon(data, buf);
2906    }
2907    #[cfg(target_arch = "x86_64")]
2908    if avx2_enabled() {
2909        return unsafe { vbit_fill4_avx2(data, buf) };
2910    }
2911    #[allow(unreachable_code)]
2912    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2913        let u = unpack8::<4>(&data[blk * 4..]);
2914        for k in 0..8 {
2915            chunk[k] = (u[k] - 7) as i8 as u8;
2916        }
2917    }
2918}
2919
2920#[cfg(target_arch = "aarch64")]
2921#[target_feature(enable = "neon")]
2922unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
2923    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
2924    // buf.len()/2 packed bytes (validated at load).
2925    unsafe {
2926        use core::arch::aarch64::*;
2927        let n = buf.len();
2928        let mask = vdupq_n_u8(0x0F);
2929        let seven = vdupq_n_s8(7);
2930        let mut g = 0usize;
2931        while g * 32 + 32 <= n {
2932            let b = vld1q_u8(data.as_ptr().add(g * 16));
2933            let hi = vshrq_n_u8::<4>(b);
2934            let lo = vandq_u8(b, mask);
2935            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
2936            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
2937            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
2938            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
2939            g += 1;
2940        }
2941    }
2942}
2943
2944#[cfg(target_arch = "x86_64")]
2945#[target_feature(enable = "avx2")]
2946unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2947    // SAFETY: see vbit_fill4_neon.
2948    unsafe {
2949        use core::arch::x86_64::*;
2950        let n = buf.len();
2951        let mask = _mm_set1_epi8(0x0F);
2952        let seven = _mm256_set1_epi8(7);
2953        let mut g = 0usize;
2954        while g * 32 + 32 <= n {
2955            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2956            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2957            let lo = _mm_and_si128(b, mask);
2958            let z = _mm256_sub_epi8(
2959                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2960                seven,
2961            );
2962            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2963            g += 1;
2964        }
2965    }
2966}
2967
2968/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
2969/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
2970/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
2971/// into 4 such blocks.
2972#[inline(always)]
2973fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2974    let mut acc = 0u64;
2975    for i in 0..B {
2976        acc = (acc << 8) | data[i] as u64;
2977    }
2978    let mask = (1u64 << B) - 1;
2979    let mut out = [0i32; 8];
2980    for (k, o) in out.iter_mut().enumerate() {
2981        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2982    }
2983    out
2984}
2985
2986/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
2987/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
2988/// MSB-first, byte-padded]. Row data offsets are precomputed at load
2989/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
2990/// overhead on every matvec.
2991#[allow(clippy::too_many_arguments)]
2992fn vbitmatvec(
2993    bytes: &[u8],
2994    offsets: &[usize],
2995    x: &[f32],
2996    rows: usize,
2997    cols: usize,
2998    out: &mut [f32],
2999    pool: Option<&Pool>,
3000) {
3001    debug_assert_eq!(out.len(), rows);
3002    debug_assert_eq!(offsets.len(), rows + 1);
3003
3004    // SDOT path: unpack the row to centered i8 once, then per-group
3005    // int8 dot against the quantized activations — same A8W8 contract
3006    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3007    if a8w8_enabled() {
3008        let act = split_act(x);
3009        let out_addr = SendMut(out.as_mut_ptr());
3010        let run = move |start: usize, end: usize| {
3011            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3012        };
3013        dispatch_rows(pool, rows, &run);
3014        return;
3015    }
3016
3017    let out_addr = SendMut(out.as_mut_ptr());
3018    let run = move |start: usize, end: usize| {
3019        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3020    };
3021    dispatch_rows(pool, rows, &run);
3022}
3023
3024/// One vbit row range via the A8W8 int8 path — kernel body of
3025/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3026/// several tensors in one dispatch (b=8 rows go exact f32).
3027#[allow(clippy::too_many_arguments)]
3028fn vbit_range_a8w8(
3029    bytes: &[u8],
3030    offsets: &[usize],
3031    x: &[f32],
3032    act: &SplitAct,
3033    rows: usize,
3034    cols: usize,
3035    out: SendMut,
3036    start: usize,
3037    end: usize,
3038) {
3039    let ng = cols / GROUP_SIZE;
3040    let bits = &bytes[..rows];
3041    let sc_off = rows;
3042    let row_dot = |r: usize| -> f32 {
3043        let b = bits[r] as usize;
3044        let l = (1i32 << (b - 1)) - 1;
3045        let mask = (1u64 << b) - 1;
3046        let data = &bytes[offsets[r]..offsets[r + 1]];
3047        if b == 8 {
3048            // u−L reaches 128 → does not fit i8; exact f32 path.
3049            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3050            let mut dot = 0f32;
3051            for g in 0..ng {
3052                let so = (r * ng + g) * 2;
3053                let sgf = f16_to_f32(u16::from_le_bytes([
3054                    bytes[sc_off + so],
3055                    bytes[sc_off + so + 1],
3056                ]));
3057                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3058                let mut gd = 0f32;
3059                for &xv in xg.iter() {
3060                    if nbits < 8 {
3061                        acc = (acc << 8) | data[idx] as u64;
3062                        idx += 1;
3063                        nbits += 8;
3064                    }
3065                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3066                    nbits -= 8;
3067                    gd += (u - l) as f32 * xv;
3068                }
3069                dot += gd * sgf;
3070            }
3071            return dot;
3072        }
3073        // Per-worker scratch: this closure runs for every row of the
3074        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3075        // row was measurable pure overhead.
3076        thread_local! {
3077            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3078                const { std::cell::RefCell::new(Vec::new()) };
3079        }
3080        #[inline(always)]
3081        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3082            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3083                let u = unpack8::<B>(&data[blk * B..]);
3084                for k in 0..8 {
3085                    chunk[k] = (u[k] - l) as i8 as u8;
3086                }
3087            }
3088        }
3089        let _ = mask;
3090        VBIT_SCRATCH.with(|scratch| {
3091            let mut buf = scratch.borrow_mut();
3092            buf.resize(cols, 0);
3093            match b {
3094                3 => fill::<3>(data, l, &mut buf),
3095                4 => vbit_fill4(data, &mut buf),
3096                5 => fill::<5>(data, l, &mut buf),
3097                6 => fill::<6>(data, l, &mut buf),
3098                _ => unreachable!(),
3099            }
3100            let mut dot = 0f32;
3101            for g in 0..ng {
3102                let so = (r * ng + g) * 2;
3103                let s = f16_to_f32(u16::from_le_bytes([
3104                    bytes[sc_off + so],
3105                    bytes[sc_off + so + 1],
3106                ]));
3107                let d = dot_i8_i8(
3108                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3109                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3110                ) as f32
3111                    * act.sx;
3112                dot += d * s;
3113            }
3114            for &(j, xv) in &act.outliers {
3115                let so = (r * ng + j / GROUP_SIZE) * 2;
3116                let s = f16_to_f32(u16::from_le_bytes([
3117                    bytes[sc_off + so],
3118                    bytes[sc_off + so + 1],
3119                ]));
3120                // xq is zeroed at outlier slots — add the exact term.
3121                dot += (buf[j] as i8) as f32 * s * xv;
3122            }
3123            dot
3124        })
3125    };
3126    for r in start..end {
3127        // SAFETY: disjoint row ranges per worker.
3128        unsafe { *out.at(r) = row_dot(r) };
3129    }
3130}
3131
3132/// Exact scalar vbit row range (same extraction, non-SDOT path).
3133#[allow(clippy::too_many_arguments)]
3134fn vbit_range_f32(
3135    bytes: &[u8],
3136    offsets: &[usize],
3137    x: &[f32],
3138    rows: usize,
3139    cols: usize,
3140    out: SendMut,
3141    start: usize,
3142    end: usize,
3143) {
3144    let ng = cols / GROUP_SIZE;
3145    let bits = &bytes[..rows];
3146    let sc_off = rows;
3147    // Per-bit-width specialized inner loops: the compiler unrolls the
3148    // constant shifts (the generic bit-buffer loop was branch-bound —
3149    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3150    #[inline(always)]
3151    fn dot_row<const B: usize>(
3152        data: &[u8],
3153        bytes: &[u8],
3154        sc_off: usize,
3155        r: usize,
3156        ng: usize,
3157        x: &[f32],
3158    ) -> f32 {
3159        let l = ((1i32 << (B - 1)) - 1) as f32;
3160        let gbytes = GROUP_SIZE * B / 8;
3161        let mut dot = 0f32;
3162        for g in 0..ng {
3163            let so = (r * ng + g) * 2;
3164            let s = f16_to_f32(u16::from_le_bytes([
3165                bytes[sc_off + so],
3166                bytes[sc_off + so + 1],
3167            ]));
3168            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3169            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3170            let mut gd = 0f32;
3171            for blk in 0..GROUP_SIZE / 8 {
3172                let u = unpack8::<B>(&gd0[blk * B..]);
3173                let xb = &xg[blk * 8..blk * 8 + 8];
3174                for k in 0..8 {
3175                    gd += (u[k] as f32 - l) * xb[k];
3176                }
3177            }
3178            dot += gd * s;
3179        }
3180        dot
3181    }
3182    for r in start..end {
3183        let data = &bytes[offsets[r]..offsets[r + 1]];
3184        let v = match bits[r] {
3185            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3186            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3187            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3188            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3189            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3190            b => unreachable!("vbit bit-width {b} (validated at load)"),
3191        };
3192        // SAFETY: disjoint row ranges per worker.
3193        unsafe { *out.at(r) = v };
3194    }
3195}
3196
3197/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3198/// and dotted against BOTH activations (MTP verify / pair prefill used
3199/// to run two full matvecs — double weight traffic and double unpack).
3200/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3201#[allow(clippy::too_many_arguments)]
3202fn vbitmatvec2(
3203    bytes: &[u8],
3204    offsets: &[usize],
3205    x1: &[f32],
3206    x2: &[f32],
3207    rows: usize,
3208    cols: usize,
3209    o1: &mut [f32],
3210    o2: &mut [f32],
3211    pool: Option<&Pool>,
3212) {
3213    debug_assert_eq!(o1.len(), rows);
3214    debug_assert_eq!(o2.len(), rows);
3215
3216    if a8w8_enabled() {
3217        let a1 = split_act(x1);
3218        let a2 = split_act(x2);
3219        let p1 = SendMut(o1.as_mut_ptr());
3220        let p2 = SendMut(o2.as_mut_ptr());
3221        let run = move |start: usize, end: usize| {
3222            vbit_range2_a8w8(
3223                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3224            )
3225        };
3226        dispatch_rows(pool, rows, &run);
3227        return;
3228    }
3229
3230    let p1 = SendMut(o1.as_mut_ptr());
3231    let p2 = SendMut(o2.as_mut_ptr());
3232    let run = move |start: usize, end: usize| {
3233        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3234    };
3235    dispatch_rows(pool, rows, &run);
3236}
3237
3238/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3239/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3240/// exact f32 for both lanes, bits streamed once).
3241#[allow(clippy::too_many_arguments)]
3242fn vbit_range2_a8w8(
3243    bytes: &[u8],
3244    offsets: &[usize],
3245    x1: &[f32],
3246    x2: &[f32],
3247    a1: &SplitAct,
3248    a2: &SplitAct,
3249    rows: usize,
3250    cols: usize,
3251    p1: SendMut,
3252    p2: SendMut,
3253    start: usize,
3254    end: usize,
3255) {
3256    let ng = cols / GROUP_SIZE;
3257    let bits = &bytes[..rows];
3258    let sc_off = rows;
3259    let row_dots = |r: usize| -> (f32, f32) {
3260        let b = bits[r] as usize;
3261        let l = (1i32 << (b - 1)) - 1;
3262        let data = &bytes[offsets[r]..offsets[r + 1]];
3263        if b == 8 {
3264            // u−L reaches 128 → does not fit i8; exact f32 path,
3265            // bits still streamed once for both lanes.
3266            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3267            let (mut d1, mut d2) = (0f32, 0f32);
3268            for g in 0..ng {
3269                let so = (r * ng + g) * 2;
3270                let sgf = f16_to_f32(u16::from_le_bytes([
3271                    bytes[sc_off + so],
3272                    bytes[sc_off + so + 1],
3273                ]));
3274                let (mut g1, mut g2) = (0f32, 0f32);
3275                for k in 0..GROUP_SIZE {
3276                    if nbits < 8 {
3277                        acc = (acc << 8) | data[idx] as u64;
3278                        idx += 1;
3279                        nbits += 8;
3280                    }
3281                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3282                    nbits -= 8;
3283                    let w = (u - l) as f32;
3284                    g1 += w * x1[g * GROUP_SIZE + k];
3285                    g2 += w * x2[g * GROUP_SIZE + k];
3286                }
3287                d1 += g1 * sgf;
3288                d2 += g2 * sgf;
3289            }
3290            return (d1, d2);
3291        }
3292        thread_local! {
3293            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3294                const { std::cell::RefCell::new(Vec::new()) };
3295        }
3296        #[inline(always)]
3297        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3298            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3299                let u = unpack8::<B>(&data[blk * B..]);
3300                for k in 0..8 {
3301                    chunk[k] = (u[k] - l) as i8 as u8;
3302                }
3303            }
3304        }
3305        VBIT_SCRATCH2.with(|scratch| {
3306            let mut buf = scratch.borrow_mut();
3307            buf.resize(cols, 0);
3308            match b {
3309                3 => fill::<3>(data, l, &mut buf),
3310                4 => vbit_fill4(data, &mut buf),
3311                5 => fill::<5>(data, l, &mut buf),
3312                6 => fill::<6>(data, l, &mut buf),
3313                _ => unreachable!(),
3314            }
3315            let (mut d1, mut d2) = (0f32, 0f32);
3316            for g in 0..ng {
3317                let so = (r * ng + g) * 2;
3318                let s = f16_to_f32(u16::from_le_bytes([
3319                    bytes[sc_off + so],
3320                    bytes[sc_off + so + 1],
3321                ]));
3322                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3323                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3324                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3325                d1 += v1 * s;
3326                d2 += v2 * s;
3327            }
3328            for &(j, xv) in &a1.outliers {
3329                let so = (r * ng + j / GROUP_SIZE) * 2;
3330                let s = f16_to_f32(u16::from_le_bytes([
3331                    bytes[sc_off + so],
3332                    bytes[sc_off + so + 1],
3333                ]));
3334                d1 += (buf[j] as i8) as f32 * s * xv;
3335            }
3336            for &(j, xv) in &a2.outliers {
3337                let so = (r * ng + j / GROUP_SIZE) * 2;
3338                let s = f16_to_f32(u16::from_le_bytes([
3339                    bytes[sc_off + so],
3340                    bytes[sc_off + so + 1],
3341                ]));
3342                d2 += (buf[j] as i8) as f32 * s * xv;
3343            }
3344            (d1, d2)
3345        })
3346    };
3347    for r in start..end {
3348        let (v1, v2) = row_dots(r);
3349        // SAFETY: disjoint row ranges per worker.
3350        unsafe {
3351            *p1.at(r) = v1;
3352            *p2.at(r) = v2;
3353        }
3354    }
3355}
3356
3357/// Two-input exact scalar vbit row range (same extraction) —
3358/// per-bit-width specialized, two accumulators per row; per-lane
3359/// accumulation order matches `vbitmatvec` exactly.
3360#[allow(clippy::too_many_arguments)]
3361fn vbit_range2_f32(
3362    bytes: &[u8],
3363    offsets: &[usize],
3364    x1: &[f32],
3365    x2: &[f32],
3366    rows: usize,
3367    cols: usize,
3368    p1: SendMut,
3369    p2: SendMut,
3370    start: usize,
3371    end: usize,
3372) {
3373    let ng = cols / GROUP_SIZE;
3374    let bits = &bytes[..rows];
3375    let sc_off = rows;
3376    #[inline(always)]
3377    #[allow(clippy::too_many_arguments)]
3378    fn dot_row2<const B: usize>(
3379        data: &[u8],
3380        bytes: &[u8],
3381        sc_off: usize,
3382        r: usize,
3383        ng: usize,
3384        x1: &[f32],
3385        x2: &[f32],
3386    ) -> (f32, f32) {
3387        let l = ((1i32 << (B - 1)) - 1) as f32;
3388        let gbytes = GROUP_SIZE * B / 8;
3389        let (mut d1, mut d2) = (0f32, 0f32);
3390        for g in 0..ng {
3391            let so = (r * ng + g) * 2;
3392            let s = f16_to_f32(u16::from_le_bytes([
3393                bytes[sc_off + so],
3394                bytes[sc_off + so + 1],
3395            ]));
3396            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3397            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3398            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3399            let (mut g1, mut g2) = (0f32, 0f32);
3400            for blk in 0..GROUP_SIZE / 8 {
3401                let u = unpack8::<B>(&gd0[blk * B..]);
3402                for k in 0..8 {
3403                    let w = u[k] as f32 - l;
3404                    g1 += w * x1g[blk * 8 + k];
3405                    g2 += w * x2g[blk * 8 + k];
3406                }
3407            }
3408            d1 += g1 * s;
3409            d2 += g2 * s;
3410        }
3411        (d1, d2)
3412    }
3413    for r in start..end {
3414        let data = &bytes[offsets[r]..offsets[r + 1]];
3415        let (v1, v2) = match bits[r] {
3416            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3417            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3418            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3419            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3420            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3421            b => unreachable!("vbit bit-width {b} (validated at load)"),
3422        };
3423        // SAFETY: disjoint row ranges per worker.
3424        unsafe {
3425            *p1.at(r) = v1;
3426            *p2.at(r) = v2;
3427        }
3428    }
3429}
3430
3431// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3432
3433/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3434/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3435/// distant streams of the split layout. Values/order identical to the
3436/// split kernels.
3437#[inline]
3438#[allow(unreachable_code)]
3439fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3440    #[cfg(target_arch = "aarch64")]
3441    unsafe {
3442        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3443    }
3444    #[cfg(target_arch = "x86_64")]
3445    unsafe {
3446        if vnni_tiles_enabled() {
3447            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3448        }
3449        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3450    }
3451    let mut acc = 0f32;
3452    for gi in 0..gpr {
3453        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3454        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3455        let mut d = 0i32;
3456        for (k, &b) in tile[2..].iter().enumerate() {
3457            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3458                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3459        }
3460        acc += d as f32 * s;
3461    }
3462    acc
3463}
3464
3465#[cfg(target_arch = "aarch64")]
3466#[target_feature(enable = "neon,dotprod")]
3467unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3468    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3469    // xq.len() == gpr·GROUP_SIZE).
3470    unsafe {
3471        use core::arch::aarch64::*;
3472        use core::arch::asm;
3473        let lomask = vdupq_n_u8(0x0F);
3474        let eight = vdupq_n_s8(8);
3475        let mut acc = 0f32;
3476        for gi in 0..gpr {
3477            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3478            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3479            let b = vld1q_u8(t.add(2));
3480            let lo = vandq_u8(b, lomask);
3481            let hi = vshrq_n_u8::<4>(b);
3482            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3483            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3484            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3485            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3486            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3487            asm!(
3488                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3489                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3490                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3491                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3492                options(pure, nomem, nostack),
3493            );
3494            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3495        }
3496        acc
3497    }
3498}
3499
3500#[cfg(target_arch = "x86_64")]
3501#[target_feature(enable = "avx2")]
3502unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3503    // SAFETY: see dot_q4t_row_sdot.
3504    unsafe {
3505        use core::arch::x86_64::*;
3506        let lomask = _mm_set1_epi8(0x0F);
3507        let eight = _mm256_set1_epi8(8);
3508        let ones = _mm256_set1_epi16(1);
3509        let mut acc = 0f32;
3510        for gi in 0..gpr {
3511            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3512            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3513            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3514            let lo = _mm_and_si128(b, lomask);
3515            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3516            let w = _mm256_sub_epi8(
3517                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3518                eight,
3519            );
3520            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3521            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3522            let d = _mm256_madd_epi16(p16, ones);
3523            let hi128 = _mm256_extracti128_si256::<1>(d);
3524            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3525            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3526            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3527            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3528        }
3529        acc
3530    }
3531}
3532
3533/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
3534/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
3535/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
3536#[cfg(target_arch = "x86_64")]
3537#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3538unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3539    // SAFETY: see dot_q4t_row_sdot.
3540    unsafe {
3541        use core::arch::x86_64::*;
3542        let lomask = _mm_set1_epi8(0x0F);
3543        let eight = _mm256_set1_epi8(8);
3544        let mut acc = 0f32;
3545        for gi in 0..gpr {
3546            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3547            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3548            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3549            let lo = _mm_and_si128(b, lomask);
3550            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3551            let w = _mm256_sub_epi8(
3552                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3553                eight,
3554            );
3555            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3556            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3557            acc += d as f32 * s;
3558        }
3559        acc
3560    }
3561}
3562
3563/// One q4_tiled row against FOUR activation streams: the nibble unpack
3564/// and abs() happen once per group instead of once per (group,
3565/// activation) — the unpack is the dominant per-element cost of the
3566/// tiled format (roadmap P0 portable blocking, q4t leg).
3567#[cfg(target_arch = "x86_64")]
3568// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
3569// to a libm call per lane — measured 2x slower than the reduction this
3570// kernel replaces. The runtime gate (`avx2_enabled`) already requires
3571// both features, so declaring it here is safe.
3572#[target_feature(enable = "avx2,fma")]
3573unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3574    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3575    unsafe {
3576        use core::arch::x86_64::*;
3577        let lomask = _mm_set1_epi8(0x0F);
3578        let eight = _mm256_set1_epi8(8);
3579        let ones = _mm256_set1_epi16(1);
3580        // One f32 accumulator VECTOR per activation, reduced once at the
3581        // end. Folding each group's i32 lanes to a scalar inside the loop
3582        // costs an extracti128 + three shift/add + a movd — a cross-lane
3583        // dependency chain per (group, activation), 288 of them per row at
3584        // cols=2304. The per-group scale is what forces a float
3585        // accumulator; it does not force a horizontal sum.
3586        //
3587        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
3588        // indexed by a loop variable LLVM keeps them in memory and every
3589        // group pays four 32-byte loads and stores. That alone made this
3590        // kernel 2x SLOWER than the per-group reduction it replaces
3591        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
3592        let mut f0 = _mm256_setzero_ps();
3593        let mut f1 = _mm256_setzero_ps();
3594        let mut f2 = _mm256_setzero_ps();
3595        let mut f3 = _mm256_setzero_ps();
3596        for gi in 0..gpr {
3597            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3598            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3599            let sv = _mm256_set1_ps(s);
3600            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3601            let lo = _mm_and_si128(bb, lomask);
3602            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3603            let w = _mm256_sub_epi8(
3604                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3605                eight,
3606            );
3607            let aw = _mm256_abs_epi8(w);
3608            let off = gi * GROUP_SIZE;
3609            let dot = |xq: &[i8]| {
3610                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3611                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3612                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3613            };
3614            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3615            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3616            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3617            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3618        }
3619        [
3620            hsum256_ps(f0),
3621            hsum256_ps(f1),
3622            hsum256_ps(f2),
3623            hsum256_ps(f3),
3624        ]
3625    }
3626}
3627
3628/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
3629/// blocked kernels pay, once per row instead of once per group.
3630#[cfg(target_arch = "x86_64")]
3631#[target_feature(enable = "avx2")]
3632#[inline]
3633unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3634    // SAFETY: pure register arithmetic on the caller's vector.
3635    unsafe {
3636        use core::arch::x86_64::*;
3637        let hi = _mm256_extractf128_ps::<1>(v);
3638        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3639        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3640        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3641        _mm_cvtss_f32(s)
3642    }
3643}
3644
3645/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
3646#[cfg(target_arch = "x86_64")]
3647#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3648unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3649    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3650    unsafe {
3651        use core::arch::x86_64::*;
3652        let lomask = _mm_set1_epi8(0x0F);
3653        let eight = _mm256_set1_epi8(8);
3654        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
3655        // one cross-lane reduction per row, not per (group, activation).
3656        let mut f0 = _mm256_setzero_ps();
3657        let mut f1 = _mm256_setzero_ps();
3658        let mut f2 = _mm256_setzero_ps();
3659        let mut f3 = _mm256_setzero_ps();
3660        for gi in 0..gpr {
3661            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3662            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3663            let sv = _mm256_set1_ps(s);
3664            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3665            let lo = _mm_and_si128(bb, lomask);
3666            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3667            let w = _mm256_sub_epi8(
3668                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3669                eight,
3670            );
3671            let aw = _mm256_abs_epi8(w);
3672            let off = gi * GROUP_SIZE;
3673            let dot = |xq: &[i8]| {
3674                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3675                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3676                    _mm256_setzero_si256(),
3677                    aw,
3678                    _mm256_sign_epi8(x, w),
3679                ))
3680            };
3681            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3682            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3683            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3684            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3685        }
3686        let acc = [
3687            hsum256_ps(f0),
3688            hsum256_ps(f1),
3689            hsum256_ps(f2),
3690            hsum256_ps(f3),
3691        ];
3692        acc
3693    }
3694}
3695
3696/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
3697/// serves FOUR activation streams. Per stream the group order and f32
3698/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
3699/// bit-for-bit.
3700#[cfg(target_arch = "aarch64")]
3701#[target_feature(enable = "neon,dotprod")]
3702unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3703    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3704    unsafe {
3705        use core::arch::aarch64::*;
3706        use core::arch::asm;
3707        let lomask = vdupq_n_u8(0x0F);
3708        let eight = vdupq_n_s8(8);
3709        let mut acc = [0f32; 4];
3710        for gi in 0..gpr {
3711            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3712            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3713            let b = vld1q_u8(t.add(2));
3714            let lo = vandq_u8(b, lomask);
3715            let hi = vshrq_n_u8::<4>(b);
3716            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3717            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3718            for (k, xq) in xs.iter().enumerate() {
3719                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3720                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3721                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3722                asm!(
3723                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3724                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3725                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3726                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3727                    options(pure, nomem, nostack),
3728                );
3729                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3730            }
3731        }
3732        acc
3733    }
3734}
3735
3736/// Exact-term correction for A8W8 outliers on a tiled row.
3737#[inline]
3738fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3739    let gi = j / GROUP_SIZE;
3740    let k = j % GROUP_SIZE;
3741    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3742    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3743    let byte = tile[2 + k / 2];
3744    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3745    ((nib as i32 - 8) as f32, s)
3746}
3747
3748/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
3749/// accumulation shape as `q4_range_f32`.
3750#[inline]
3751fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3752    let mut acc = 0f32;
3753    for gi in 0..gpr {
3754        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3755        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3756        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3757        let mut ga = 0f32;
3758        for (k, &b) in tile[2..].iter().enumerate() {
3759            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3760                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3761        }
3762        acc += ga * s;
3763    }
3764    acc
3765}
3766
3767/// Split view of a `q4tp` payload. The three planes are resolved once per
3768/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
3769/// the row loop would put a division on the hot path for nothing.
3770struct Q4tpView<'a> {
3771    nib: &'a [u8],
3772    params: &'a [u8],
3773    codes: &'a [u8],
3774    stride: usize,
3775}
3776
3777impl<'a> Q4tpView<'a> {
3778    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
3779        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
3780        Self {
3781            nib: &bytes[..params_off],
3782            params: &bytes[params_off..codes_off],
3783            codes: &bytes[codes_off..],
3784            stride,
3785        }
3786    }
3787
3788    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
3789    ///
3790    /// Doing this once per row — rather than decoding a 5-bit code inside the
3791    /// tile loop — is what makes the format free at runtime. Random access to
3792    /// a packed 5-bit field costs a division, two bounds checks and a branch;
3793    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
3794    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
3795    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
3796    /// Eight 5-bit codes are exactly five bytes, so a whole group of
3797    /// eight decodes from one little-endian word at fixed shifts. The
3798    /// bit-accumulator this replaces carried a data-dependent `while
3799    /// have < 5` refill whose branch sat in the innermost loop of every
3800    /// q4tp row; a decode profile put this function above the dot
3801    /// products it feeds. Same bitstream, same codes — just no branch
3802    /// and eight independent extractions.
3803    #[inline]
3804    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
3805        let tab = q4tp_ladder(self.params, r);
3806        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
3807        let out = &mut out[..gpr];
3808        let mut chunks = out.chunks_exact_mut(8);
3809        let mut ci = 0usize;
3810        for c in &mut chunks {
3811            let w = u64::from(codes[ci])
3812                | u64::from(codes[ci + 1]) << 8
3813                | u64::from(codes[ci + 2]) << 16
3814                | u64::from(codes[ci + 3]) << 24
3815                | u64::from(codes[ci + 4]) << 32;
3816            for (k, o) in c.iter_mut().enumerate() {
3817                *o = tab[((w >> (5 * k)) & 31) as usize];
3818            }
3819            ci += 5;
3820        }
3821        // Fewer than eight codes left: the shared total accessor, which
3822        // tolerates a 5-bit field whose spill byte is past the stride.
3823        let tail = &codes[ci..];
3824        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
3825            *o = tab[q4tp_code(tail, k)];
3826        }
3827    }
3828}
3829
3830#[inline]
3831fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
3832    #[cfg(target_arch = "aarch64")]
3833    unsafe {
3834        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
3835    }
3836    #[cfg(target_arch = "x86_64")]
3837    unsafe {
3838        if vnni_tiles_enabled() {
3839            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
3840        }
3841        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
3842    }
3843    #[allow(unreachable_code)]
3844    {
3845        let mut acc = 0f32;
3846        for gi in 0..gpr {
3847            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3848            let s = scales[gi];
3849            let mut d = 0i32;
3850            for (k, &b) in tile.iter().enumerate() {
3851                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3852                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3853            }
3854            acc += d as f32 * s;
3855        }
3856        acc
3857    }
3858}
3859
3860/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
3861/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
3862#[cfg(target_arch = "aarch64")]
3863#[target_feature(enable = "neon,dotprod")]
3864unsafe fn dot_q4tp_row_sdot(
3865    nib: &[u8],
3866    r: usize,
3867    gpr: usize,
3868    xq: &[i8],
3869    scales: &[f32],
3870) -> f32 {
3871    // SAFETY: callers uphold slice-length contracts (16B tile per group,
3872    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
3873    unsafe {
3874        use core::arch::aarch64::*;
3875        use core::arch::asm;
3876        let lomask = vdupq_n_u8(0x0F);
3877        let eight = vdupq_n_s8(8);
3878        let mut acc = 0f32;
3879        for gi in 0..gpr {
3880            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3881            let s = *scales.get_unchecked(gi);
3882            let b = vld1q_u8(t);
3883            let lo = vandq_u8(b, lomask);
3884            let hi = vshrq_n_u8::<4>(b);
3885            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3886            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3887            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3888            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3889            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3890            asm!(
3891                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3892                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3893                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3894                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3895                options(pure, nomem, nostack),
3896            );
3897            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3898        }
3899        acc
3900    }
3901}
3902
3903#[cfg(target_arch = "x86_64")]
3904#[target_feature(enable = "avx2")]
3905unsafe fn dot_q4tp_row_avx2(
3906    nib: &[u8],
3907    r: usize,
3908    gpr: usize,
3909    xq: &[i8],
3910    scales: &[f32],
3911) -> f32 {
3912    // SAFETY: see dot_q4tp_row_sdot.
3913    unsafe {
3914        use core::arch::x86_64::*;
3915        let lomask = _mm_set1_epi8(0x0F);
3916        let eight = _mm256_set1_epi8(8);
3917        let ones = _mm256_set1_epi16(1);
3918        let mut acc = 0f32;
3919        for gi in 0..gpr {
3920            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3921            let s = *scales.get_unchecked(gi);
3922            let b = _mm_loadu_si128(t as *const __m128i);
3923            let lo = _mm_and_si128(b, lomask);
3924            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3925            let w = _mm256_sub_epi8(
3926                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3927                eight,
3928            );
3929            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3930            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3931            let d = _mm256_madd_epi16(p16, ones);
3932            let hi128 = _mm256_extracti128_si256::<1>(d);
3933            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3934            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3935            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3936            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3937        }
3938        acc
3939    }
3940}
3941
3942/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
3943/// 256-bit VL encoding is the one to use here).
3944#[cfg(target_arch = "x86_64")]
3945#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3946unsafe fn dot_q4tp_row_vnni(
3947    nib: &[u8],
3948    r: usize,
3949    gpr: usize,
3950    xq: &[i8],
3951    scales: &[f32],
3952) -> f32 {
3953    // SAFETY: see dot_q4tp_row_sdot.
3954    unsafe {
3955        use core::arch::x86_64::*;
3956        let lomask = _mm_set1_epi8(0x0F);
3957        let eight = _mm256_set1_epi8(8);
3958        let mut acc = 0f32;
3959        for gi in 0..gpr {
3960            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3961            let s = *scales.get_unchecked(gi);
3962            let b = _mm_loadu_si128(t as *const __m128i);
3963            let lo = _mm_and_si128(b, lomask);
3964            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3965            let w = _mm256_sub_epi8(
3966                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3967                eight,
3968            );
3969            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3970            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
3971        }
3972        acc
3973    }
3974}
3975
3976/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
3977/// accumulation shape as `q4t_row_exact`.
3978#[inline]
3979fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
3980    let mut acc = 0f32;
3981    for gi in 0..gpr {
3982        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3983        let s = scales[gi];
3984        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3985        let mut ga = 0f32;
3986        for (k, &b) in tile.iter().enumerate() {
3987            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3988                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3989        }
3990        acc += ga * s;
3991    }
3992    acc
3993}
3994
3995/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
3996/// activation outliers at full precision after the int8 pass.
3997#[inline]
3998fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
3999    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4000    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4001    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4002    ((n as i32 - 8) as f32, scales[gi])
4003}
4004
4005/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4006fn q4tp_matvec(
4007    bytes: &[u8],
4008    x: &[f32],
4009    rows: usize,
4010    cols: usize,
4011    out: &mut [f32],
4012    pool: Option<&Pool>,
4013) {
4014    debug_assert_eq!(out.len(), rows);
4015    let gpr = cols / GROUP_SIZE;
4016    let v = Q4tpView::new(bytes, rows, cols);
4017    let out_addr = SendMut(out.as_mut_ptr());
4018    if a8w8_enabled() {
4019        let act = split_act(x);
4020        let run = |start: usize, end: usize| {
4021            // One scratch row of scales per worker, reused across its rows.
4022            let mut sc = vec![0f32; gpr];
4023            for r in start..end {
4024                v.scales_into(r, gpr, &mut sc);
4025                let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4026                for &(j, xv) in &act.outliers {
4027                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4028                    acc += w * s * xv;
4029                }
4030                // SAFETY: disjoint row ranges per worker.
4031                unsafe { *out_addr.at(r) = acc };
4032            }
4033        };
4034        dispatch_rows(pool, rows, &run);
4035        return;
4036    }
4037    let run = |start: usize, end: usize| {
4038        let mut sc = vec![0f32; gpr];
4039        for r in start..end {
4040            v.scales_into(r, gpr, &mut sc);
4041            // SAFETY: disjoint row ranges per worker.
4042            unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4043        }
4044    };
4045    dispatch_rows(pool, rows, &run);
4046}
4047
4048/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4049/// row ladder are read once and spent on both activation streams.
4050#[allow(clippy::too_many_arguments)]
4051fn q4tp_matvec2(
4052    bytes: &[u8],
4053    x1: &[f32],
4054    x2: &[f32],
4055    rows: usize,
4056    cols: usize,
4057    o1: &mut [f32],
4058    o2: &mut [f32],
4059    pool: Option<&Pool>,
4060) {
4061    let gpr = cols / GROUP_SIZE;
4062    let v = Q4tpView::new(bytes, rows, cols);
4063    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4064    let run = |start: usize, end: usize| {
4065        let mut sc = vec![0f32; gpr];
4066        for r in start..end {
4067            v.scales_into(r, gpr, &mut sc);
4068            // SAFETY: disjoint row ranges per worker.
4069            unsafe {
4070                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4071                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4072            }
4073        }
4074    };
4075    dispatch_rows(pool, rows, &run);
4076}
4077
4078/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
4079/// spent on four activation streams, which is where a prefill batch stops
4080/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
4081#[cfg(target_arch = "aarch64")]
4082#[target_feature(enable = "neon,dotprod")]
4083unsafe fn dot_q4tp_row_1x4_sdot(
4084    nib: &[u8],
4085    r: usize,
4086    gpr: usize,
4087    xs: [&[i8]; 4],
4088    scales: &[f32],
4089) -> [f32; 4] {
4090    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
4091    unsafe {
4092        use core::arch::aarch64::*;
4093        use core::arch::asm;
4094        let lomask = vdupq_n_u8(0x0F);
4095        let eight = vdupq_n_s8(8);
4096        // Named accumulators, NOT an array indexed by a loop variable: the
4097        // latter does not stay in registers (the same defect cost 2x in the
4098        // AVX2 q4t kernel and again in WGSL).
4099        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4100        for gi in 0..gpr {
4101            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4102            let s = *scales.get_unchecked(gi);
4103            let bb = vld1q_u8(t);
4104            let lo = vandq_u8(bb, lomask);
4105            let hi = vshrq_n_u8::<4>(bb);
4106            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4107            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4108            let mut d = [0f32; 4];
4109            for (k, dk) in d.iter_mut().enumerate() {
4110                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4111                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4112                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4113                asm!(
4114                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4115                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4116                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4117                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4118                    options(pure, nomem, nostack),
4119                );
4120                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4121            }
4122            f0 += d[0];
4123            f1 += d[1];
4124            f2 += d[2];
4125            f3 += d[3];
4126        }
4127        [f0, f1, f2, f3]
4128    }
4129}
4130
4131/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
4132/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
4133/// the format was fine, the missing arms were the whole regression.
4134fn q4tp_matmat(
4135    bytes: &[u8],
4136    xs_all: &[f32],
4137    b: usize,
4138    rows: usize,
4139    cols: usize,
4140    out: &mut [f32],
4141    pool: Option<&Pool>,
4142) {
4143    debug_assert_eq!(out.len(), b * rows);
4144    let gpr = cols / GROUP_SIZE;
4145    let v = Q4tpView::new(bytes, rows, cols);
4146
4147    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
4148    #[cfg(target_os = "macos")]
4149    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4150        dequant_matmat_accel(
4151            &|r, dst| {
4152                let mut sc = [0f32; 32];
4153                let mut scv;
4154                let s: &[f32] = if gpr <= 32 {
4155                    v.scales_into(r, gpr, &mut sc);
4156                    &sc[..gpr]
4157                } else {
4158                    scv = vec![0f32; gpr];
4159                    v.scales_into(r, gpr, &mut scv);
4160                    &scv
4161                };
4162                for gi in 0..gpr {
4163                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4164                    for (k, &bb) in tile.iter().enumerate() {
4165                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
4166                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
4167                    }
4168                }
4169            },
4170            xs_all,
4171            b,
4172            rows,
4173            cols,
4174            out,
4175            pool,
4176        );
4177        return;
4178    }
4179
4180    let out_addr = SendMut(out.as_mut_ptr());
4181    if a8w8_enabled() {
4182        let acts: Vec<SplitAct> = (0..b)
4183            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4184            .collect();
4185        let acts = &acts;
4186        #[cfg(target_arch = "aarch64")]
4187        let blocked_ok = sdot_enabled()
4188            && std::env::var("CMF_X86_BLOCKED")
4189                .map(|val| val != "0")
4190                .unwrap_or(true);
4191        #[cfg(not(target_arch = "aarch64"))]
4192        let blocked_ok = false;
4193        let run = |start: usize, end: usize| {
4194            let mut sc = vec![0f32; gpr];
4195            for r in start..end {
4196                v.scales_into(r, gpr, &mut sc);
4197                let mut bi = 0usize;
4198                #[cfg(target_arch = "aarch64")]
4199                if blocked_ok {
4200                    while bi + 4 <= acts.len() {
4201                        let xs = [
4202                            acts[bi].xq.as_slice(),
4203                            acts[bi + 1].xq.as_slice(),
4204                            acts[bi + 2].xq.as_slice(),
4205                            acts[bi + 3].xq.as_slice(),
4206                        ];
4207                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc) };
4208                        for k in 0..4 {
4209                            let act = &acts[bi + k];
4210                            let mut acc = d[k] * act.sx;
4211                            for &(j, xv) in &act.outliers {
4212                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4213                                acc += w * s * xv;
4214                            }
4215                            // SAFETY: disjoint (bi, r) cells per worker.
4216                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4217                        }
4218                        bi += 4;
4219                    }
4220                }
4221                let _ = blocked_ok;
4222                while bi < acts.len() {
4223                    let act = &acts[bi];
4224                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4225                    for &(j, xv) in &act.outliers {
4226                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4227                        acc += w * s * xv;
4228                    }
4229                    // SAFETY: disjoint (bi, r) cells per worker range.
4230                    unsafe { *out_addr.at(bi * rows + r) = acc };
4231                    bi += 1;
4232                }
4233            }
4234        };
4235        dispatch_rows(pool, rows, &run);
4236        return;
4237    }
4238
4239    let run = |start: usize, end: usize| {
4240        let mut sc = vec![0f32; gpr];
4241        for r in start..end {
4242            v.scales_into(r, gpr, &mut sc);
4243            for bi in 0..b {
4244                let x = &xs_all[bi * cols..(bi + 1) * cols];
4245                // SAFETY: disjoint (bi, r) cells per worker range.
4246                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4247            }
4248        }
4249    };
4250    dispatch_rows(pool, rows, &run);
4251}
4252
4253/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
4254fn q4t_matvec(
4255    bytes: &[u8],
4256    x: &[f32],
4257    rows: usize,
4258    cols: usize,
4259    out: &mut [f32],
4260    pool: Option<&Pool>,
4261) {
4262    debug_assert_eq!(out.len(), rows);
4263    let gpr = cols / GROUP_SIZE;
4264    let out_addr = SendMut(out.as_mut_ptr());
4265    if a8w8_enabled() {
4266        let act = split_act(x);
4267        let run = move |start: usize, end: usize| {
4268            for r in start..end {
4269                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4270                for &(j, xv) in &act.outliers {
4271                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4272                    acc += w * s * xv;
4273                }
4274                // SAFETY: disjoint row ranges per worker.
4275                unsafe { *out_addr.at(r) = acc };
4276            }
4277        };
4278        dispatch_rows(pool, rows, &run);
4279        return;
4280    }
4281    let run = move |start: usize, end: usize| {
4282        for r in start..end {
4283            // SAFETY: disjoint row ranges per worker.
4284            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
4285        }
4286    };
4287    dispatch_rows(pool, rows, &run);
4288}
4289
4290/// Fused two-input q4_tiled matvec (weights read once per pair).
4291#[allow(clippy::too_many_arguments)]
4292fn q4t_matvec2(
4293    bytes: &[u8],
4294    x1: &[f32],
4295    x2: &[f32],
4296    rows: usize,
4297    cols: usize,
4298    o1: &mut [f32],
4299    o2: &mut [f32],
4300    pool: Option<&Pool>,
4301) {
4302    let gpr = cols / GROUP_SIZE;
4303    let p1 = SendMut(o1.as_mut_ptr());
4304    let p2 = SendMut(o2.as_mut_ptr());
4305    if a8w8_enabled() {
4306        let a1 = split_act(x1);
4307        let a2 = split_act(x2);
4308        let run = move |start: usize, end: usize| {
4309            for r in start..end {
4310                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
4311                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
4312                for &(j, xv) in &a1.outliers {
4313                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4314                    v1 += w * s * xv;
4315                }
4316                for &(j, xv) in &a2.outliers {
4317                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4318                    v2 += w * s * xv;
4319                }
4320                // SAFETY: disjoint row ranges per worker.
4321                unsafe {
4322                    *p1.at(r) = v1;
4323                    *p2.at(r) = v2;
4324                }
4325            }
4326        };
4327        dispatch_rows(pool, rows, &run);
4328        return;
4329    }
4330    let run = move |start: usize, end: usize| {
4331        for r in start..end {
4332            // SAFETY: disjoint row ranges per worker.
4333            unsafe {
4334                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
4335                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
4336            }
4337        }
4338    };
4339    dispatch_rows(pool, rows, &run);
4340}
4341
4342/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
4343#[allow(clippy::too_many_arguments)]
4344/// Prefill GEMM through Accelerate for group-quantized codecs: a
4345/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
4346/// each tile rides the AMX with one sgemm — the generic sibling of
4347/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
4348/// decode (b=1) never takes this path.
4349#[cfg(target_os = "macos")]
4350fn dequant_matmat_accel(
4351    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
4352    xs_all: &[f32],
4353    b: usize,
4354    rows: usize,
4355    cols: usize,
4356    out: &mut [f32],
4357    pool: Option<&Pool>,
4358) {
4359    const TR: usize = 2048;
4360    thread_local! {
4361        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4362    }
4363    WTILE.with(|wt| {
4364        let mut wtile = wt.borrow_mut();
4365        wtile.resize(TR * cols, 0.0);
4366        let mut r0 = 0usize;
4367        while r0 < rows {
4368            let tr = TR.min(rows - r0);
4369            let wt_addr = SendMut(wtile.as_mut_ptr());
4370            let run = |start: usize, end: usize| {
4371                for r in start..end {
4372                    // SAFETY: workers cover disjoint r ranges.
4373                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4374                    dequant_row(r0 + r, dst);
4375                }
4376            };
4377            dispatch_rows(pool, tr, &run);
4378            unsafe {
4379                accel_blas::cblas_sgemm(
4380                    101, // RowMajor
4381                    111, // NoTrans A
4382                    112, // Trans B
4383                    b as i32,
4384                    tr as i32,
4385                    cols as i32,
4386                    1.0,
4387                    xs_all.as_ptr(),
4388                    cols as i32,
4389                    wtile.as_ptr(),
4390                    cols as i32,
4391                    0.0,
4392                    out.as_mut_ptr().add(r0),
4393                    rows as i32,
4394                );
4395            }
4396            r0 += tr;
4397        }
4398    });
4399}
4400
4401fn q4t_matmat(
4402    bytes: &[u8],
4403    xs_all: &[f32],
4404    b: usize,
4405    rows: usize,
4406    cols: usize,
4407    out: &mut [f32],
4408    pool: Option<&Pool>,
4409) {
4410    debug_assert_eq!(out.len(), b * rows);
4411    let gpr = cols / GROUP_SIZE;
4412    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
4413    // the dequant-tile sgemm is an order above the SDOT row loop for
4414    // prefill shapes (imagegen DiT forwards are exactly this).
4415    #[cfg(target_os = "macos")]
4416    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4417        dequant_matmat_accel(
4418            &|r, dst| {
4419                for gi in 0..gpr {
4420                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4421                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4422                    for (k, &bb) in tile[2..].iter().enumerate() {
4423                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
4424                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
4425                    }
4426                }
4427            },
4428            xs_all,
4429            b,
4430            rows,
4431            cols,
4432            out,
4433            pool,
4434        );
4435        return;
4436    }
4437    let out_addr = SendMut(out.as_mut_ptr());
4438    if a8w8_enabled() {
4439        let acts: Vec<SplitAct> = (0..b)
4440            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4441            .collect();
4442        let acts = &acts;
4443        #[cfg(target_arch = "x86_64")]
4444        let blocked_ok = avx2_enabled()
4445            && std::env::var("CMF_X86_BLOCKED")
4446                .map(|v| v != "0")
4447                .unwrap_or(true);
4448        #[cfg(target_arch = "aarch64")]
4449        let blocked_ok = sdot_enabled()
4450            && std::env::var("CMF_X86_BLOCKED")
4451                .map(|v| v != "0")
4452                .unwrap_or(true);
4453        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
4454        let blocked_ok = false;
4455        let run = move |start: usize, end: usize| {
4456            for r in start..end {
4457                let mut bi = 0usize;
4458                #[cfg(target_arch = "aarch64")]
4459                if blocked_ok {
4460                    while bi + 4 <= acts.len() {
4461                        let xs = [
4462                            acts[bi].xq.as_slice(),
4463                            acts[bi + 1].xq.as_slice(),
4464                            acts[bi + 2].xq.as_slice(),
4465                            acts[bi + 3].xq.as_slice(),
4466                        ];
4467                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
4468                        for k in 0..4 {
4469                            let act = &acts[bi + k];
4470                            let mut acc = d[k] * act.sx;
4471                            for &(j, xv) in &act.outliers {
4472                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4473                                acc += w * sc * xv;
4474                            }
4475                            // SAFETY: disjoint (bi, r) cells per worker.
4476                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4477                        }
4478                        bi += 4;
4479                    }
4480                }
4481                #[cfg(target_arch = "x86_64")]
4482                if blocked_ok {
4483                    while bi + 4 <= acts.len() {
4484                        let xs = [
4485                            acts[bi].xq.as_slice(),
4486                            acts[bi + 1].xq.as_slice(),
4487                            acts[bi + 2].xq.as_slice(),
4488                            acts[bi + 3].xq.as_slice(),
4489                        ];
4490                        let d = unsafe {
4491                            if vnni_tiles_enabled() {
4492                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
4493                            } else {
4494                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
4495                            }
4496                        };
4497                        for k in 0..4 {
4498                            let act = &acts[bi + k];
4499                            let mut acc = d[k] * act.sx;
4500                            for &(j, xv) in &act.outliers {
4501                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4502                                acc += w * sc * xv;
4503                            }
4504                            // SAFETY: disjoint (bi, r) cells per worker.
4505                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4506                        }
4507                        bi += 4;
4508                    }
4509                }
4510                let _ = blocked_ok;
4511                while bi < acts.len() {
4512                    let act = &acts[bi];
4513                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4514                    for &(j, xv) in &act.outliers {
4515                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
4516                        acc += w * s * xv;
4517                    }
4518                    // SAFETY: disjoint (bi, r) cells per worker range.
4519                    unsafe { *out_addr.at(bi * rows + r) = acc };
4520                    bi += 1;
4521                }
4522            }
4523        };
4524        dispatch_rows(pool, rows, &run);
4525        return;
4526    }
4527    let run = move |start: usize, end: usize| {
4528        for r in start..end {
4529            for bi in 0..b {
4530                let x = &xs_all[bi * cols..(bi + 1) * cols];
4531                // SAFETY: disjoint (bi, r) cells per worker range.
4532                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
4533            }
4534        }
4535    };
4536    dispatch_rows(pool, rows, &run);
4537}
4538
4539// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
4540// 32-group tile. The kernel family mirrors q4_tiled: one sequential
4541// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
4542// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
4543
4544/// Per-32-group sums of the quantized activation — the ±1 identity's
4545/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
4546/// matvec and reused by every row.
4547fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
4548    (0..gpr)
4549        .map(|gi| {
4550            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
4551                .iter()
4552                .map(|&v| v as i32)
4553                .sum()
4554        })
4555        .collect()
4556}
4557
4558/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
4559/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
4560/// x86 pass).
4561#[inline]
4562#[allow(unreachable_code)]
4563/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
4564/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
4565/// masked activation sums through maddubs(1, x&mask), and
4566/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
4567#[cfg(target_arch = "x86_64")]
4568#[target_feature(enable = "avx2")]
4569unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4570    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4571    unsafe {
4572        use core::arch::x86_64::*;
4573        // Byte j of the mask must replicate bits-byte j/8.
4574        let expand = _mm256_setr_epi8(
4575            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,
4576            3, 3, 3,
4577        );
4578        let bitsel = _mm256_setr_epi8(
4579            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4580            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4581        );
4582        let ones8 = _mm256_set1_epi8(1);
4583        let ones16 = _mm256_set1_epi16(1);
4584        let mut acc = 0f32;
4585        for gi in 0..gpr {
4586            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4587            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4588            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4589            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4590            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4591            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4592            let sel = _mm256_and_si256(x, mask);
4593            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
4594            let p16 = _mm256_maddubs_epi16(ones8, sel);
4595            let d32 = _mm256_madd_epi16(p16, ones16);
4596            let hi128 = _mm256_extracti128_si256::<1>(d32);
4597            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4598            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4599            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4600            let msum = _mm_cvtsi128_si32(s32);
4601            // The and-select keeps x UN-negated (unlike ARM's −1-mask
4602            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
4603            let d = 2 * msum - gsum[gi];
4604            acc += d as f32 * s;
4605        }
4606        acc
4607    }
4608}
4609
4610/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
4611/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
4612#[cfg(target_arch = "x86_64")]
4613#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4614unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4615    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4616    unsafe {
4617        use core::arch::x86_64::*;
4618        let expand = _mm256_setr_epi8(
4619            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,
4620            3, 3, 3,
4621        );
4622        let bitsel = _mm256_setr_epi8(
4623            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4624            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4625        );
4626        let ones8 = _mm256_set1_epi8(1);
4627        let mut acc = 0f32;
4628        for gi in 0..gpr {
4629            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4630            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4631            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4632            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4633            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4634            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4635            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4636            let d = 2 * msum - gsum[gi];
4637            acc += d as f32 * s;
4638        }
4639        acc
4640    }
4641}
4642
4643/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
4644#[cfg(target_arch = "x86_64")]
4645#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4646unsafe fn dot_q1_row_1x4_vnni(
4647    bytes: &[u8],
4648    r: usize,
4649    gpr: usize,
4650    xs: [&[i8]; 4],
4651    gsums: [&[i32]; 4],
4652) -> [f32; 4] {
4653    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4654    unsafe {
4655        use core::arch::x86_64::*;
4656        let expand = _mm256_setr_epi8(
4657            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,
4658            3, 3, 3,
4659        );
4660        let bitsel = _mm256_setr_epi8(
4661            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4662            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4663        );
4664        let ones8 = _mm256_set1_epi8(1);
4665        let mut acc = [0f32; 4];
4666        for gi in 0..gpr {
4667            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4668            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4669            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4670            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4671            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4672            for (k, xq) in xs.iter().enumerate() {
4673                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4674                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4675                let d = 2 * msum - gsums[k][gi];
4676                acc[k] += d as f32 * s;
4677            }
4678        }
4679        acc
4680    }
4681}
4682
4683/// The blocked 1×4 flavor: the expanded bit mask serves four activation
4684/// streams per group (mask build once, four select+reduce chains).
4685#[cfg(target_arch = "x86_64")]
4686#[target_feature(enable = "avx2")]
4687unsafe fn dot_q1_row_1x4_avx2(
4688    bytes: &[u8],
4689    r: usize,
4690    gpr: usize,
4691    xs: [&[i8]; 4],
4692    gsums: [&[i32]; 4],
4693) -> [f32; 4] {
4694    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4695    unsafe {
4696        use core::arch::x86_64::*;
4697        let expand = _mm256_setr_epi8(
4698            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,
4699            3, 3, 3,
4700        );
4701        let bitsel = _mm256_setr_epi8(
4702            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4703            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4704        );
4705        let ones8 = _mm256_set1_epi8(1);
4706        let ones16 = _mm256_set1_epi16(1);
4707        let mut acc = [0f32; 4];
4708        for gi in 0..gpr {
4709            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4710            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4711            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4712            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4713            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4714            for (k, xq) in xs.iter().enumerate() {
4715                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4716                let sel = _mm256_and_si256(x, mask);
4717                let p16 = _mm256_maddubs_epi16(ones8, sel);
4718                let d32 = _mm256_madd_epi16(p16, ones16);
4719                let hi128 = _mm256_extracti128_si256::<1>(d32);
4720                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4721                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4722                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4723                let msum = _mm_cvtsi128_si32(s32);
4724                let d = 2 * msum - gsums[k][gi];
4725                acc[k] += d as f32 * s;
4726            }
4727        }
4728        acc
4729    }
4730}
4731
4732#[allow(unreachable_code)]
4733fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4734    #[cfg(target_arch = "aarch64")]
4735    unsafe {
4736        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
4737    }
4738    #[cfg(target_arch = "x86_64")]
4739    if avx2_enabled() {
4740        unsafe {
4741            if vnni_tiles_enabled() {
4742                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
4743            }
4744            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
4745        }
4746    }
4747    let _ = gsum;
4748    let mut acc = 0f32;
4749    for gi in 0..gpr {
4750        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4751        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4752        let mut d = 0i32;
4753        for (j, &b) in tile[2..].iter().enumerate() {
4754            for k in 0..8 {
4755                let w = ((b >> k) & 1) as i32 * 2 - 1;
4756                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
4757            }
4758        }
4759        acc += d as f32 * s;
4760    }
4761    acc
4762}
4763
4764/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
4765/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
4766/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
4767/// per-group activation sums shared across every row of the matvec.
4768/// Four tiles (128 weights) per iteration: integer dots reduce through
4769/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
4770/// fused f32 multiply-add. Integer math throughout — bit-identical to
4771/// the scalar ±1 reference.
4772#[cfg(target_arch = "aarch64")]
4773#[target_feature(enable = "neon,dotprod")]
4774unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4775    // SAFETY: callers uphold slice-length contracts (6B tile per group,
4776    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
4777    unsafe {
4778        use core::arch::aarch64::*;
4779        use core::arch::asm;
4780        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4781        let m = vld1q_u8(MASKS.as_ptr());
4782        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
4783        macro_rules! tile_dot {
4784            ($t:expr, $x:expr) => {{
4785                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
4786                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
4787                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4788                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4789                let x0 = vld1q_s8($x);
4790                let x1 = vld1q_s8($x.add(16));
4791                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4792                asm!(
4793                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4794                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4795                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4796                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4797                    options(pure, nomem, nostack),
4798                );
4799                vaddq_s32(a0, a1)
4800            }};
4801        }
4802        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
4803        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
4804        // bit-byte across 8 lanes for vtst, and the four scales gather
4805        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
4806        // 4 branchy software f16 conversions per 128 weights (the
4807        // measured load-port wall of this kernel) become 2 vector
4808        // loads + 9 table lookups. Integer math order is unchanged —
4809        // bit-identical results (FCVTL is exact on every f16).
4810        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4811        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4812        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4813        const IW11: [u8; 16] = [
4814            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4815        ];
4816        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4817        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4818        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4819        let isc = vld1_u8(ISC.as_ptr());
4820        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
4821        macro_rules! tile_dot_tbl {
4822            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
4823                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
4824                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
4825                let x0 = vld1q_s8($x);
4826                let x1 = vld1q_s8($x.add(16));
4827                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4828                asm!(
4829                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4830                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4831                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4832                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4833                    options(pure, nomem, nostack),
4834                );
4835                vaddq_s32(a0, a1)
4836            }};
4837        }
4838        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4839        let row_base = r * gpr * Q1_TILE;
4840        let abs_end = bytes.len();
4841        let xp = xq.as_ptr();
4842        let gp = gsum.as_ptr();
4843        let mut accv = vdupq_n_f32(0.0);
4844        let mut gi = 0;
4845        // The second pair load reads 4B past tile gi+3 — stay inside
4846        // the payload slice (only the file's final tiles fall back).
4847        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4848            let t0 = base.add(gi * Q1_TILE);
4849            let ld_a = vld1q_u8(t0);
4850            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4851            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
4852            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
4853            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
4854            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
4855            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
4856            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4857            let g = vld1q_s32(gp.add(gi));
4858            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4859            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4860            let scf: float32x4_t;
4861            asm!(
4862                "fcvtl {o:v}.4s, {i:v}.4h",
4863                o = out(vreg) scf, i = in(vreg) sc16,
4864                options(pure, nomem, nostack),
4865            );
4866            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
4867            gi += 4;
4868        }
4869        let mut acc = vaddvq_f32(accv);
4870        while gi < gpr {
4871            let t = base.add(gi * Q1_TILE);
4872            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4873            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
4874            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
4875            gi += 1;
4876        }
4877        acc
4878    }
4879}
4880
4881/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
4882/// activation streams (prefill amortization — the same idea as the
4883/// AVX2 twin; per stream the group order, fma order and tail match the
4884/// single-row kernel exactly, so batch == matvec bit-for-bit).
4885#[cfg(target_arch = "aarch64")]
4886#[target_feature(enable = "neon,dotprod")]
4887unsafe fn dot_q1_row_1x4_sdot(
4888    bytes: &[u8],
4889    r: usize,
4890    gpr: usize,
4891    xs: [&[i8]; 4],
4892    gs: [&[i32]; 4],
4893) -> [f32; 4] {
4894    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
4895    unsafe {
4896        use core::arch::aarch64::*;
4897        use core::arch::asm;
4898        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4899        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4900        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4901        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4902        const IW11: [u8; 16] = [
4903            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4904        ];
4905        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4906        let m = vld1q_u8(MASKS.as_ptr());
4907        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4908        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4909        let isc = vld1_u8(ISC.as_ptr());
4910        macro_rules! sdot2 {
4911            ($w0:expr, $w1:expr, $x:expr) => {{
4912                let x0 = vld1q_s8($x);
4913                let x1 = vld1q_s8($x.add(16));
4914                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4915                asm!(
4916                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4917                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4918                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4919                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
4920                    options(pure, nomem, nostack),
4921                );
4922                vaddq_s32(a0, a1)
4923            }};
4924        }
4925        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4926        let row_base = r * gpr * Q1_TILE;
4927        let abs_end = bytes.len();
4928        let mut accv = [vdupq_n_f32(0.0); 4];
4929        let mut gi = 0;
4930        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4931            let t0 = base.add(gi * Q1_TILE);
4932            let ld_a = vld1q_u8(t0);
4933            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4934            // Unpack ONCE — eight ±mask vectors serve all four streams.
4935            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
4936            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
4937            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
4938            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
4939            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
4940            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
4941            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
4942            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
4943            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4944            let scf: float32x4_t;
4945            asm!(
4946                "fcvtl {o:v}.4s, {i:v}.4h",
4947                o = out(vreg) scf, i = in(vreg) sc16,
4948                options(pure, nomem, nostack),
4949            );
4950            for k in 0..4 {
4951                let xp = xs[k].as_ptr();
4952                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
4953                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
4954                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
4955                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
4956                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4957                let g = vld1q_s32(gs[k].as_ptr().add(gi));
4958                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4959                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
4960            }
4961            gi += 4;
4962        }
4963        let mut acc = [
4964            vaddvq_f32(accv[0]),
4965            vaddvq_f32(accv[1]),
4966            vaddvq_f32(accv[2]),
4967            vaddvq_f32(accv[3]),
4968        ];
4969        while gi < gpr {
4970            let t = base.add(gi * Q1_TILE);
4971            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4972            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
4973            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
4974            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4975            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4976            for k in 0..4 {
4977                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
4978                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
4979            }
4980            gi += 1;
4981        }
4982        acc
4983    }
4984}
4985
4986/// (weight ±1, scale) of one q1 element — the exact outlier term.
4987#[inline]
4988fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4989    let gi = j / GROUP_SIZE;
4990    let k = j % GROUP_SIZE;
4991    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4992    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4993    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
4994    ((bit as i32 * 2 - 1) as f32, s)
4995}
4996
4997/// Exact scalar q1 row (CMF_SDOT=0 contract).
4998#[inline]
4999fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5000    let mut acc = 0f32;
5001    for gi in 0..gpr {
5002        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5003        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5004        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5005        let mut ga = 0f32;
5006        for (j, &b) in tile[2..].iter().enumerate() {
5007            for k in 0..8 {
5008                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
5009            }
5010        }
5011        acc += ga * s;
5012    }
5013    acc
5014}
5015
5016/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
5017/// extracted so multi-matrix jobs drive the same kernel).
5018#[allow(clippy::too_many_arguments)]
5019fn q1_range_a8w8(
5020    bytes: &[u8],
5021    gpr: usize,
5022    act: &SplitAct,
5023    gsum: &[i32],
5024    out: SendMut,
5025    start: usize,
5026    end: usize,
5027) {
5028    for r in start..end {
5029        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5030        for &(j, xv) in &act.outliers {
5031            let (w, s) = q1_outlier(bytes, r, gpr, j);
5032            acc += w * s * xv;
5033        }
5034        // SAFETY: disjoint row ranges per worker.
5035        unsafe { *out.at(r) = acc };
5036    }
5037}
5038
5039/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
5040fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
5041    for r in start..end {
5042        // SAFETY: disjoint row ranges per worker.
5043        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
5044    }
5045}
5046
5047/// q1t per-row overlay locator. After the base (`base_len`) come
5048/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
5049/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
5050/// `(row_ptr offset, entries offset, present)`.
5051fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
5052    let entries = base_len + (rows + 1) * 4;
5053    (base_len, entries, entries <= bytes.len())
5054}
5055
5056/// Read `row_ptr[r]` from the overlay's prefix-sum table.
5057#[inline]
5058fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
5059    let o = rp_off + r * 4;
5060    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
5061}
5062
5063/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
5064/// decoding a q1t code is a table load, not the base-3 divide/modulo per
5065/// weight (division is ~20–40× the cost of a load). Built at compile time.
5066const SIGN5: [[f32; 5]; 256] = {
5067    let mut lut = [[0.0f32; 5]; 256];
5068    let pow3 = [1u16, 3, 9, 27, 81];
5069    let mut byte = 0usize;
5070    while byte < 256 {
5071        let mut i = 0usize;
5072        while i < 5 {
5073            let code = (byte as u16 / pow3[i]) % 3;
5074            lut[byte][i] = if code == 1 {
5075                1.0
5076            } else if code == 2 {
5077                -1.0
5078            } else {
5079                0.0
5080            };
5081            i += 1;
5082        }
5083        byte += 1;
5084    }
5085    lut
5086};
5087
5088/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
5089const SIGN5_I8: [[i8; 5]; 256] = {
5090    let mut lut = [[0i8; 5]; 256];
5091    let pow3 = [1u16, 3, 9, 27, 81];
5092    let mut byte = 0usize;
5093    while byte < 256 {
5094        let mut i = 0usize;
5095        while i < 5 {
5096            let code = (byte as u16 / pow3[i]) % 3;
5097            lut[byte][i] = if code == 1 {
5098                1
5099            } else if code == 2 {
5100                -1
5101            } else {
5102                0
5103            };
5104            i += 1;
5105        }
5106        byte += 1;
5107    }
5108    lut
5109};
5110
5111/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
5112/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
5113/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
5114/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
5115/// buffer is padded to 40). This is the decode/prefill hot inner op.
5116const SIGN5_U64: [u64; 256] = {
5117    let mut lut = [0u64; 256];
5118    let pow3 = [1u16, 3, 9, 27, 81];
5119    let mut byte = 0usize;
5120    while byte < 256 {
5121        let mut v = 0u64;
5122        let mut i = 0usize;
5123        while i < 5 {
5124            let code = (byte as u16 / pow3[i]) % 3;
5125            let s: u8 = if code == 1 {
5126                1
5127            } else if code == 2 {
5128                0xFF
5129            } else {
5130                0
5131            };
5132            v |= (s as u64) << (i * 8);
5133            i += 1;
5134        }
5135        lut[byte] = v;
5136        byte += 1;
5137    }
5138    lut
5139};
5140
5141/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
5142/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
5143/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
5144/// the overlay correction owns that column — no double counting.
5145#[inline]
5146fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
5147    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5148    let off = (r * gpr + j / GROUP_SIZE) * TILE;
5149    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5150    let within = j % GROUP_SIZE;
5151    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
5152}
5153
5154/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
5155/// (integer accumulation is order-independent).
5156#[cfg(target_arch = "aarch64")]
5157#[target_feature(enable = "neon,dotprod")]
5158#[inline]
5159unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
5160    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5161    unsafe {
5162        use core::arch::aarch64::*;
5163        use core::arch::asm;
5164        let w0 = vld1q_s8(w);
5165        let w1 = vld1q_s8(w.add(16));
5166        let x0 = vld1q_s8(x);
5167        let x1 = vld1q_s8(x.add(16));
5168        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5169        asm!(
5170            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5171            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5172            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5173            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5174            options(pure, nomem, nostack),
5175        );
5176        vaddvq_s32(vaddq_s32(a0, a1))
5177    }
5178}
5179
5180/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
5181/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
5182#[cfg(target_arch = "x86_64")]
5183#[target_feature(enable = "avx2")]
5184#[inline]
5185unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
5186    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5187    unsafe {
5188        use core::arch::x86_64::*;
5189        let wv = _mm256_loadu_si256(w as *const __m256i);
5190        let xv = _mm256_loadu_si256(x as *const __m256i);
5191        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5192        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
5193        let hi128 = _mm256_extracti128_si256::<1>(d);
5194        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5195        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5196        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5197        _mm_cvtsi128_si32(s32)
5198    }
5199}
5200
5201/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
5202/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
5203/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
5204/// overwritten by the next; the final 6 padding bytes are unused by the dot.
5205#[inline]
5206fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
5207    debug_assert!(dst.len() >= 40);
5208    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
5209    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
5210    unsafe {
5211        let p = dst.as_mut_ptr();
5212        for bi in 0..7 {
5213            core::ptr::write_unaligned(
5214                p.add(bi * 5) as *mut u64,
5215                SIGN5_U64[*codes.add(bi) as usize],
5216            );
5217        }
5218    }
5219}
5220
5221/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
5222/// row's signs are unpacked once and dotted against every batch input).
5223/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
5224/// reachable; the scalar arm is a non-SIMD-arch fallback.
5225#[inline]
5226fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
5227    #[cfg(target_arch = "aarch64")]
5228    unsafe {
5229        return sdot32_i8(w, x);
5230    }
5231    #[cfg(target_arch = "x86_64")]
5232    unsafe {
5233        return i8dot32_avx2(w, x);
5234    }
5235    #[allow(unreachable_code)]
5236    unsafe {
5237        let mut s = 0i32;
5238        for k in 0..GROUP_SIZE {
5239            s += *w.add(k) as i32 * *x.add(k) as i32;
5240        }
5241        s
5242    }
5243}
5244
5245#[inline]
5246unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
5247    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
5248        (
5249            SIGN5_U64[*codes as usize],
5250            SIGN5_U64[*codes.add(1) as usize],
5251            SIGN5_U64[*codes.add(2) as usize],
5252            SIGN5_U64[*codes.add(3) as usize],
5253            SIGN5_U64[*codes.add(4) as usize],
5254            SIGN5_U64[*codes.add(5) as usize],
5255            SIGN5_U64[*codes.add(6) as usize],
5256        )
5257    };
5258
5259    let u0 = s0 | (s1 << 40);
5260    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
5261    let u2 = (s3 >> 8) | (s4 << 32);
5262    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
5263
5264    (u0, u1, u2, u3)
5265}
5266
5267/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
5268/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
5269/// ARM SDOT.
5270#[cfg(target_arch = "aarch64")]
5271#[target_feature(enable = "neon,dotprod")]
5272unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5273    use core::arch::aarch64::*;
5274    use core::arch::asm;
5275    unsafe {
5276        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5277        let mut acc = 0f32;
5278        let bytes_ptr = bytes.as_ptr();
5279        let xq_ptr = xq.as_ptr();
5280        let row_off = r * gpr * TILE;
5281
5282        let gpr2 = gpr & !1;
5283        let mut gi = 0;
5284        while gi < gpr2 {
5285            let off0 = row_off + gi * TILE;
5286            let off1 = off0 + TILE;
5287            let s0 = f16_to_f32(u16::from_le_bytes([
5288                *bytes_ptr.add(off0),
5289                *bytes_ptr.add(off0 + 1),
5290            ]));
5291            let s1 = f16_to_f32(u16::from_le_bytes([
5292                *bytes_ptr.add(off1),
5293                *bytes_ptr.add(off1 + 1),
5294            ]));
5295
5296            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5297            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5298
5299            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5300            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5301            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5302            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5303
5304            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5305            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5306            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
5307            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
5308
5309            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
5310            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5311            asm!(
5312                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
5313                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
5314                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
5315                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
5316                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
5317                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
5318                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
5319                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
5320                options(pure, nomem, nostack),
5321            );
5322            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
5323            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
5324            acc += d0 as f32 * s0 + d1 as f32 * s1;
5325            gi += 2;
5326        }
5327
5328        if gi < gpr {
5329            let off = row_off + gi * TILE;
5330            let s = f16_to_f32(u16::from_le_bytes([
5331                *bytes_ptr.add(off),
5332                *bytes_ptr.add(off + 1),
5333            ]));
5334            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5335            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5336            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5337            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5338            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5339            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5340            asm!(
5341                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5342                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5343                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5344                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5345                options(pure, nomem, nostack),
5346            );
5347            let d = vaddvq_s32(vaddq_s32(a0, a1));
5348            acc += d as f32 * s;
5349        }
5350        acc
5351    }
5352}
5353
5354/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
5355#[cfg(target_arch = "x86_64")]
5356#[target_feature(enable = "avx2")]
5357unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5358    use core::arch::x86_64::*;
5359    unsafe {
5360        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5361        let mut acc = 0f32;
5362        let bytes_ptr = bytes.as_ptr();
5363        let xq_ptr = xq.as_ptr();
5364        let row_off = r * gpr * TILE;
5365
5366        let ones = _mm256_set1_epi16(1);
5367        for gi in 0..gpr {
5368            let off = row_off + gi * TILE;
5369            let s = f16_to_f32(u16::from_le_bytes([
5370                *bytes_ptr.add(off),
5371                *bytes_ptr.add(off + 1),
5372            ]));
5373            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5374            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5375            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5376            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5377            let d256 = _mm256_madd_epi16(p16, ones);
5378            let d128 = _mm_add_epi32(
5379                _mm256_castsi256_si128(d256),
5380                _mm256_extracti128_si256(d256, 1),
5381            );
5382            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
5383            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
5384            acc += d32 as f32 * s;
5385        }
5386        acc
5387    }
5388}
5389
5390/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
5391#[cfg(target_arch = "x86_64")]
5392#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5393unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5394    use core::arch::x86_64::*;
5395    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
5396    unsafe {
5397        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5398        let mut acc = 0f32;
5399        let bytes_ptr = bytes.as_ptr();
5400        let xq_ptr = xq.as_ptr();
5401        let row_off = r * gpr * TILE;
5402        for gi in 0..gpr {
5403            let off = row_off + gi * TILE;
5404            let s = f16_to_f32(u16::from_le_bytes([
5405                *bytes_ptr.add(off),
5406                *bytes_ptr.add(off + 1),
5407            ]));
5408            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5409            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5410            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5411            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5412            acc += d as f32 * s;
5413        }
5414        acc
5415    }
5416}
5417
5418/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
5419/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
5420/// reachable.
5421#[inline]
5422fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5423    #[cfg(target_arch = "aarch64")]
5424    unsafe {
5425        return q1t_dot_row_sdot(bytes, r, gpr, xq);
5426    }
5427    #[cfg(target_arch = "x86_64")]
5428    unsafe {
5429        if vnni_tiles_enabled() {
5430            return q1t_dot_row_vnni(bytes, r, gpr, xq);
5431        }
5432        return q1t_dot_row_avx2(bytes, r, gpr, xq);
5433    }
5434    #[allow(unreachable_code)]
5435    {
5436        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5437        let mut acc = 0f32;
5438        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
5439        for gi in 0..gpr {
5440            let off = (r * gpr + gi) * TILE;
5441            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5442            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
5443            let mut d = 0i32;
5444            for k in 0..GROUP_SIZE {
5445                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
5446            }
5447            acc += d as f32 * s;
5448        }
5449        acc
5450    }
5451}
5452
5453/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
5454/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
5455/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
5456/// base contributes nothing there and this is a plain `value·x`, not
5457/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
5458/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
5459fn q1t_row_outlier_correction(
5460    bytes: &[u8],
5461    r: usize,
5462    rp_off: usize,
5463    entries_off: usize,
5464    has_ov: bool,
5465    x: &[f32],
5466) -> f32 {
5467    if !has_ov {
5468        return 0.0;
5469    }
5470    let (c0, c1) = (
5471        q1t_rowptr(bytes, rp_off, r),
5472        q1t_rowptr(bytes, rp_off, r + 1),
5473    );
5474    let mut corr = 0f32;
5475    for p in c0..c1 {
5476        let e = entries_off + p * 4;
5477        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5478        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5479        corr += val * x[col];
5480    }
5481    corr
5482}
5483
5484/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
5485/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
5486/// Used by the batched (prefill) path where the decode amortizes over the batch.
5487fn q1t_dequant_row(
5488    bytes: &[u8],
5489    r: usize,
5490    gpr: usize,
5491    rp_off: usize,
5492    entries_off: usize,
5493    has_ov: bool,
5494    buf: &mut [f32],
5495) {
5496    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5497    for g in 0..gpr {
5498        let off = (r * gpr + g) * TILE;
5499        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5500        let codes = &bytes[off + 2..off + TILE];
5501        let bc = g * GROUP_SIZE;
5502        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
5503        for bi in 0..6 {
5504            let lut = &SIGN5[codes[bi] as usize];
5505            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
5506            for i in 0..5 {
5507                d[i] = lut[i] * s;
5508            }
5509        }
5510        let lut = &SIGN5[codes[6] as usize];
5511        buf[bc + 30] = lut[0] * s;
5512        buf[bc + 31] = lut[1] * s;
5513    }
5514    if !has_ov {
5515        return;
5516    }
5517    let (c0, c1) = (
5518        q1t_rowptr(bytes, rp_off, r),
5519        q1t_rowptr(bytes, rp_off, r + 1),
5520    );
5521    for p in c0..c1 {
5522        let e = entries_off + p * 4;
5523        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5524        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5525    }
5526}
5527
5528/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
5529/// computes the ternary base; the overlay stays on the CPU — its entries are
5530/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
5531fn q1t_add_overlay(
5532    bytes: &[u8],
5533    x: &[f32],
5534    rows: usize,
5535    cols: usize,
5536    out: &mut [f32],
5537    pool: Option<&Pool>,
5538) {
5539    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5540    let gpr = cols / GROUP_SIZE;
5541    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5542    if !has_ov {
5543        return;
5544    }
5545    let out_addr = SendMut(out.as_mut_ptr());
5546    let run = move |start: usize, end: usize| {
5547        for r in start..end {
5548            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5549            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
5550            unsafe { *out_addr.at(r) += corr };
5551        }
5552    };
5553    dispatch_rows(pool, rows, &run);
5554}
5555
5556/// Q1T row range via the A8W8 int8 path — shared activation split,
5557/// per-row: base SDOT dot + outlier correction + overlay.
5558#[allow(clippy::too_many_arguments)]
5559fn q1t_range_a8w8(
5560    bytes: &[u8],
5561    gpr: usize,
5562    rp_off: usize,
5563    ent_off: usize,
5564    has_ov: bool,
5565    act: &SplitAct,
5566    x: &[f32],
5567    out: SendMut,
5568    start: usize,
5569    end: usize,
5570) {
5571    for r in start..end {
5572        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5573        for &(j, xv) in &act.outliers {
5574            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5575        }
5576        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5577        // SAFETY: disjoint row ranges per worker.
5578        unsafe { *out.at(r) = acc };
5579    }
5580}
5581
5582/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
5583/// dispatch when a8w8 is unavailable.
5584#[allow(clippy::too_many_arguments)]
5585fn q1t_range_f32_batch(
5586    bytes: &[u8],
5587    gpr: usize,
5588    rp_off: usize,
5589    ent_off: usize,
5590    has_ov: bool,
5591    x: &[f32],
5592    out: SendMut,
5593    start: usize,
5594    end: usize,
5595) {
5596    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5597    let mut sg = [0f32; GROUP_SIZE];
5598    for r in start..end {
5599        let mut acc = 0f32;
5600        for g in 0..gpr {
5601            let off = (r * gpr + g) * TILE;
5602            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5603            let codes = &bytes[off + 2..off + TILE];
5604            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5605            for bi in 0..6 {
5606                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5607            }
5608            let lut = &SIGN5[codes[6] as usize];
5609            sg[30] = lut[0];
5610            sg[31] = lut[1];
5611            let mut gsum = 0f32;
5612            for k in 0..GROUP_SIZE {
5613                gsum += sg[k] * xg[k];
5614            }
5615            acc += s * gsum;
5616        }
5617        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5618        // SAFETY: disjoint row ranges per worker.
5619        unsafe { *out.at(r) = acc };
5620    }
5621}
5622
5623/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
5624/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
5625/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
5626fn q1t_matvec(
5627    bytes: &[u8],
5628    x: &[f32],
5629    rows: usize,
5630    cols: usize,
5631    out: &mut [f32],
5632    pool: Option<&Pool>,
5633) {
5634    debug_assert_eq!(out.len(), rows);
5635    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5636    let gpr = cols / GROUP_SIZE;
5637    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5638    let out_addr = SendMut(out.as_mut_ptr());
5639    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
5640    // (`split_act`), activation outliers added back exactly in f32, weight
5641    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
5642    if a8w8_enabled() {
5643        let act = split_act(x);
5644        let act = &act;
5645        let run = move |start: usize, end: usize| {
5646            for r in start..end {
5647                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5648                for &(j, xv) in &act.outliers {
5649                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5650                }
5651                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5652                // SAFETY: disjoint row ranges per worker.
5653                unsafe { *out_addr.at(r) = acc };
5654            }
5655        };
5656        dispatch_rows(pool, rows, &run);
5657        return;
5658    }
5659    let run = move |start: usize, end: usize| {
5660        // Per-group signs, unpacked contiguously so the dot below is a clean
5661        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
5662        // 5-values-per-byte base-3 layout won't SIMD in place.
5663        let mut sg = [0f32; GROUP_SIZE];
5664        for r in start..end {
5665            let mut acc = 0f32;
5666            for g in 0..gpr {
5667                let off = (r * gpr + g) * TILE;
5668                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5669                let codes = &bytes[off + 2..off + TILE];
5670                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5671                for bi in 0..6 {
5672                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5673                }
5674                let lut = &SIGN5[codes[6] as usize];
5675                sg[30] = lut[0];
5676                sg[31] = lut[1];
5677                let mut gsum = 0f32;
5678                for k in 0..GROUP_SIZE {
5679                    gsum += sg[k] * xg[k];
5680                }
5681                acc += s * gsum;
5682            }
5683            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5684            unsafe { *out_addr.at(r) = acc };
5685        }
5686    };
5687    dispatch_rows(pool, rows, &run);
5688}
5689
5690/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
5691/// ternary codes serves BOTH activation streams (the unpack chain is
5692/// the dominant per-row cost — MTP verify pairs paid it twice). Per
5693/// stream the group order and f32 accumulation match the single-row
5694/// kernel exactly, so pair == 2×matvec bit-for-bit.
5695#[cfg(target_arch = "aarch64")]
5696#[target_feature(enable = "neon,dotprod")]
5697unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
5698    use core::arch::aarch64::*;
5699    use core::arch::asm;
5700    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
5701    unsafe {
5702        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5703        let bytes_ptr = bytes.as_ptr();
5704        let row_off = r * gpr * TILE;
5705        let xp = [xa.as_ptr(), xb.as_ptr()];
5706        let mut acc = [0f32; 2];
5707        macro_rules! sdot2 {
5708            ($w0:expr, $w1:expr, $x:expr) => {{
5709                let x0 = vld1q_s8($x);
5710                let x1 = vld1q_s8($x.add(16));
5711                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5712                asm!(
5713                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5714                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5715                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5716                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5717                    options(pure, nomem, nostack),
5718                );
5719                vaddvq_s32(vaddq_s32(a0, a1))
5720            }};
5721        }
5722        let gpr2 = gpr & !1;
5723        let mut gi = 0;
5724        while gi < gpr2 {
5725            let off0 = row_off + gi * TILE;
5726            let off1 = off0 + TILE;
5727            let s0 = f16_to_f32(u16::from_le_bytes([
5728                *bytes_ptr.add(off0),
5729                *bytes_ptr.add(off0 + 1),
5730            ]));
5731            let s1 = f16_to_f32(u16::from_le_bytes([
5732                *bytes_ptr.add(off1),
5733                *bytes_ptr.add(off1 + 1),
5734            ]));
5735            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5736            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5737            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5738            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5739            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5740            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5741            for k in 0..2 {
5742                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
5743                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
5744                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
5745            }
5746            gi += 2;
5747        }
5748        if gi < gpr {
5749            let off = row_off + gi * TILE;
5750            let s = f16_to_f32(u16::from_le_bytes([
5751                *bytes_ptr.add(off),
5752                *bytes_ptr.add(off + 1),
5753            ]));
5754            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5755            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5756            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5757            for k in 0..2 {
5758                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
5759                acc[k] += d as f32 * s;
5760            }
5761        }
5762        acc
5763    }
5764}
5765
5766/// Fused Q1T pair matvec: ONE pass over the rows serves both
5767/// activation streams — on ARM the ternary register unpack happens
5768/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
5769/// rides the row's L1-warm tile bytes. Per stream the math matches
5770/// `q1t_matvec` exactly.
5771fn q1t_matvec2(
5772    bytes: &[u8],
5773    x1: &[f32],
5774    x2: &[f32],
5775    rows: usize,
5776    cols: usize,
5777    o1: &mut [f32],
5778    o2: &mut [f32],
5779    pool: Option<&Pool>,
5780) {
5781    debug_assert_eq!(o1.len(), rows);
5782    debug_assert_eq!(o2.len(), rows);
5783    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5784    let gpr = cols / GROUP_SIZE;
5785    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5786    let out1 = SendMut(o1.as_mut_ptr());
5787    let out2 = SendMut(o2.as_mut_ptr());
5788    if a8w8_enabled() {
5789        let a1 = split_act(x1);
5790        let a2 = split_act(x2);
5791        let (a1, a2) = (&a1, &a2);
5792        let run = move |start: usize, end: usize| {
5793            for r in start..end {
5794                #[cfg(target_arch = "aarch64")]
5795                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
5796                // target features are present.
5797                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
5798                #[cfg(not(target_arch = "aarch64"))]
5799                let ds = [
5800                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
5801                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
5802                ];
5803                let mut acc1 = ds[0] * a1.sx;
5804                for &(j, xv) in &a1.outliers {
5805                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
5806                }
5807                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
5808                let mut acc2 = ds[1] * a2.sx;
5809                for &(j, xv) in &a2.outliers {
5810                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
5811                }
5812                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
5813                // SAFETY: disjoint row ranges per worker.
5814                unsafe {
5815                    *out1.at(r) = acc1;
5816                    *out2.at(r) = acc2;
5817                }
5818            }
5819        };
5820        dispatch_rows(pool, rows, &run);
5821        return;
5822    }
5823    let run = move |start: usize, end: usize| {
5824        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
5825        // dot both streams — same op order per stream as `q1t_matvec`.
5826        let mut sg = [0f32; GROUP_SIZE];
5827        for r in start..end {
5828            let mut acc1 = 0f32;
5829            let mut acc2 = 0f32;
5830            for g in 0..gpr {
5831                let off = (r * gpr + g) * TILE;
5832                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5833                let codes = &bytes[off + 2..off + TILE];
5834                for bi in 0..6 {
5835                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5836                }
5837                let lut = &SIGN5[codes[6] as usize];
5838                sg[30] = lut[0];
5839                sg[31] = lut[1];
5840                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5841                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5842                let mut gsum1 = 0f32;
5843                for k in 0..GROUP_SIZE {
5844                    gsum1 += sg[k] * xg1[k];
5845                }
5846                acc1 += s * gsum1;
5847                let mut gsum2 = 0f32;
5848                for k in 0..GROUP_SIZE {
5849                    gsum2 += sg[k] * xg2[k];
5850                }
5851                acc2 += s * gsum2;
5852            }
5853            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
5854            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
5855            // SAFETY: disjoint row ranges per worker.
5856            unsafe {
5857                *out1.at(r) = acc1;
5858                *out2.at(r) = acc2;
5859            }
5860        }
5861    };
5862    dispatch_rows(pool, rows, &run);
5863}
5864
5865/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
5866/// batch against it (amortizes the per-row decode).
5867fn q1t_matmat(
5868    bytes: &[u8],
5869    xs: &[f32],
5870    b: usize,
5871    rows: usize,
5872    cols: usize,
5873    out: &mut [f32],
5874    pool: Option<&Pool>,
5875) {
5876    debug_assert_eq!(out.len(), b * rows);
5877    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5878    let gpr = cols / GROUP_SIZE;
5879    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5880    let out_addr = SendMut(out.as_mut_ptr());
5881    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
5882    // each weight row's signs to i8 ONCE, then int8-dot against every input —
5883    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
5884    if a8w8_enabled() {
5885        let acts: Vec<SplitAct> = (0..b)
5886            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
5887            .collect();
5888        let acts = &acts;
5889        let run = move |start: usize, end: usize| {
5890            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
5891            let mut sc = vec![0f32; gpr]; // per-group scales
5892            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
5893            for r in start..end {
5894                for g in 0..gpr {
5895                    let off = (r * gpr + g) * TILE;
5896                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5897                    q1t_unpack_group_i8(
5898                        bytes.as_ptr().wrapping_add(off + 2),
5899                        &mut sg[g * GROUP_SIZE..],
5900                    );
5901                }
5902                for bi in 0..b {
5903                    let act = &acts[bi];
5904                    let mut isum = 0f32;
5905                    for g in 0..gpr {
5906                        let d = q1t_i8dot32(
5907                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
5908                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
5909                        );
5910                        isum += d as f32 * sc[g];
5911                    }
5912                    let mut acc = isum * act.sx;
5913                    for &(j, xv) in &act.outliers {
5914                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5915                    }
5916                    accs[bi] = acc;
5917                }
5918                // Overlay ONCE per row for the whole batch: read each (col, val)
5919                // from mmap a single time (was b× — the re-read dominated prefill)
5920                // and fan it out over the batch via the cached inputs.
5921                if has_ov {
5922                    let (c0, c1) = (
5923                        q1t_rowptr(bytes, rp_off, r),
5924                        q1t_rowptr(bytes, rp_off, r + 1),
5925                    );
5926                    for p in c0..c1 {
5927                        let e = ent_off + p * 4;
5928                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5929                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5930                        for bi in 0..b {
5931                            accs[bi] += val * xs[bi * cols + col];
5932                        }
5933                    }
5934                }
5935                for bi in 0..b {
5936                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
5937                }
5938            }
5939        };
5940        dispatch_rows(pool, rows, &run);
5941        return;
5942    }
5943    let run = move |start: usize, end: usize| {
5944        let mut buf = vec![0f32; cols];
5945        for r in start..end {
5946            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
5947            for bi in 0..b {
5948                let xr = &xs[bi * cols..(bi + 1) * cols];
5949                let mut acc = 0f32;
5950                for j in 0..cols {
5951                    acc += buf[j] * xr[j];
5952                }
5953                unsafe { *out_addr.at(bi * rows + r) = acc };
5954            }
5955        }
5956    };
5957    dispatch_rows(pool, rows, &run);
5958}
5959
5960fn q1_matvec(
5961    bytes: &[u8],
5962    x: &[f32],
5963    rows: usize,
5964    cols: usize,
5965    out: &mut [f32],
5966    pool: Option<&Pool>,
5967) {
5968    debug_assert_eq!(out.len(), rows);
5969    let gpr = cols / GROUP_SIZE;
5970    let out_addr = SendMut(out.as_mut_ptr());
5971    if a8w8_enabled() {
5972        let act = split_act(x);
5973        let gsum = q1_group_sums(&act.xq, gpr);
5974        let (act, gsum) = (&act, &gsum);
5975        let run = move |start: usize, end: usize| {
5976            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
5977        };
5978        dispatch_rows(pool, rows, &run);
5979        return;
5980    }
5981    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
5982    dispatch_rows(pool, rows, &run);
5983}
5984
5985/// Fused two-input q1 matvec (weights read once per pair).
5986#[allow(clippy::too_many_arguments)]
5987fn q1_matvec2(
5988    bytes: &[u8],
5989    x1: &[f32],
5990    x2: &[f32],
5991    rows: usize,
5992    cols: usize,
5993    o1: &mut [f32],
5994    o2: &mut [f32],
5995    pool: Option<&Pool>,
5996) {
5997    let gpr = cols / GROUP_SIZE;
5998    let p1 = SendMut(o1.as_mut_ptr());
5999    let p2 = SendMut(o2.as_mut_ptr());
6000    if a8w8_enabled() {
6001        let a1 = split_act(x1);
6002        let a2 = split_act(x2);
6003        let g1 = q1_group_sums(&a1.xq, gpr);
6004        let g2 = q1_group_sums(&a2.xq, gpr);
6005        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
6006        let run = move |start: usize, end: usize| {
6007            for r in start..end {
6008                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
6009                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
6010                for &(j, xv) in &a1.outliers {
6011                    let (w, s) = q1_outlier(bytes, r, gpr, j);
6012                    v1 += w * s * xv;
6013                }
6014                for &(j, xv) in &a2.outliers {
6015                    let (w, s) = q1_outlier(bytes, r, gpr, j);
6016                    v2 += w * s * xv;
6017                }
6018                // SAFETY: disjoint row ranges per worker.
6019                unsafe {
6020                    *p1.at(r) = v1;
6021                    *p2.at(r) = v2;
6022                }
6023            }
6024        };
6025        dispatch_rows(pool, rows, &run);
6026        return;
6027    }
6028    let run = move |start: usize, end: usize| {
6029        for r in start..end {
6030            // SAFETY: disjoint row ranges per worker.
6031            unsafe {
6032                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
6033                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
6034            }
6035        }
6036    };
6037    dispatch_rows(pool, rows, &run);
6038}
6039
6040/// Batched q1 matmat: each row's tiles stream once per microbatch.
6041#[allow(clippy::too_many_arguments)]
6042fn q1_matmat(
6043    bytes: &[u8],
6044    xs_all: &[f32],
6045    b: usize,
6046    rows: usize,
6047    cols: usize,
6048    out: &mut [f32],
6049    pool: Option<&Pool>,
6050) {
6051    debug_assert_eq!(out.len(), b * rows);
6052    let gpr = cols / GROUP_SIZE;
6053    let out_addr = SendMut(out.as_mut_ptr());
6054    if a8w8_enabled() {
6055        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
6056            .map(|bi| {
6057                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
6058                let gsum = q1_group_sums(&act.xq, gpr);
6059                (act, gsum)
6060            })
6061            .collect();
6062        let acts = &acts;
6063        #[cfg(target_arch = "x86_64")]
6064        let blocked_ok = avx2_enabled()
6065            && std::env::var("CMF_X86_BLOCKED")
6066                .map(|v| v != "0")
6067                .unwrap_or(true);
6068        #[cfg(target_arch = "aarch64")]
6069        let blocked_ok = sdot_enabled()
6070            && std::env::var("CMF_X86_BLOCKED")
6071                .map(|v| v != "0")
6072                .unwrap_or(true);
6073        let run = move |start: usize, end: usize| {
6074            for r in start..end {
6075                let mut bi = 0usize;
6076                // Blocked 1×4: the unpacked bit mask serves four
6077                // activation streams per group.
6078                #[cfg(target_arch = "aarch64")]
6079                if blocked_ok {
6080                    while bi + 4 <= acts.len() {
6081                        let xs = [
6082                            acts[bi].0.xq.as_slice(),
6083                            acts[bi + 1].0.xq.as_slice(),
6084                            acts[bi + 2].0.xq.as_slice(),
6085                            acts[bi + 3].0.xq.as_slice(),
6086                        ];
6087                        let gs = [
6088                            acts[bi].1.as_slice(),
6089                            acts[bi + 1].1.as_slice(),
6090                            acts[bi + 2].1.as_slice(),
6091                            acts[bi + 3].1.as_slice(),
6092                        ];
6093                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
6094                        for k in 0..4 {
6095                            let (act, _) = &acts[bi + k];
6096                            let mut acc = d[k] * act.sx;
6097                            for &(j, xv) in &act.outliers {
6098                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
6099                                acc += w * sc * xv;
6100                            }
6101                            // SAFETY: disjoint (bi, r) cells per worker.
6102                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6103                        }
6104                        bi += 4;
6105                    }
6106                }
6107                #[cfg(target_arch = "x86_64")]
6108                if blocked_ok {
6109                    while bi + 4 <= acts.len() {
6110                        let xs = [
6111                            acts[bi].0.xq.as_slice(),
6112                            acts[bi + 1].0.xq.as_slice(),
6113                            acts[bi + 2].0.xq.as_slice(),
6114                            acts[bi + 3].0.xq.as_slice(),
6115                        ];
6116                        let gs = [
6117                            acts[bi].1.as_slice(),
6118                            acts[bi + 1].1.as_slice(),
6119                            acts[bi + 2].1.as_slice(),
6120                            acts[bi + 3].1.as_slice(),
6121                        ];
6122                        let d = unsafe {
6123                            if vnni_tiles_enabled() {
6124                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
6125                            } else {
6126                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
6127                            }
6128                        };
6129                        for k in 0..4 {
6130                            let (act, _) = &acts[bi + k];
6131                            let mut acc = d[k] * act.sx;
6132                            for &(j, xv) in &act.outliers {
6133                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
6134                                acc += w * sc * xv;
6135                            }
6136                            // SAFETY: disjoint (bi, r) cells per worker.
6137                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6138                        }
6139                        bi += 4;
6140                    }
6141                }
6142                while bi < acts.len() {
6143                    let (act, gsum) = &acts[bi];
6144                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6145                    for &(j, xv) in &act.outliers {
6146                        let (w, s) = q1_outlier(bytes, r, gpr, j);
6147                        acc += w * s * xv;
6148                    }
6149                    // SAFETY: disjoint (bi, r) cells per worker range.
6150                    unsafe { *out_addr.at(bi * rows + r) = acc };
6151                    bi += 1;
6152                }
6153            }
6154        };
6155        dispatch_rows(pool, rows, &run);
6156        return;
6157    }
6158    let run = move |start: usize, end: usize| {
6159        for r in start..end {
6160            for bi in 0..b {
6161                let x = &xs_all[bi * cols..(bi + 1) * cols];
6162                // SAFETY: disjoint (bi, r) cells per worker range.
6163                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
6164            }
6165        }
6166    };
6167    dispatch_rows(pool, rows, &run);
6168}
6169
6170/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
6171/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
6172/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
6173/// 32-group, exact outlier correction — the same A8W8 contract as q8.
6174/// `CMF_SDOT=0` keeps the exact scalar path.
6175fn q4matvec(
6176    bytes: &[u8],
6177    x: &[f32],
6178    rows: usize,
6179    cols: usize,
6180    out: &mut [f32],
6181    pool: Option<&Pool>,
6182) {
6183    debug_assert_eq!(out.len(), rows);
6184    let (packed, scales) = q4_split(bytes, rows, cols);
6185    let gpr = cols / GROUP_SIZE;
6186    let out_addr = SendMut(out.as_mut_ptr());
6187
6188    if a8w8_enabled() {
6189        let act = split_act(x);
6190        let run = move |start: usize, end: usize| {
6191            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
6192        };
6193        dispatch_rows(pool, rows, &run);
6194        return;
6195    }
6196
6197    let run =
6198        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
6199    dispatch_rows(pool, rows, &run);
6200}
6201
6202/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
6203/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
6204#[inline]
6205#[allow(unreachable_code)]
6206/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
6207/// streams: the 32-byte weight chunk and its abs() load once per group,
6208/// the per-group f16 scale decodes once — four maddubs+reduce chains
6209/// instead of four full (load, abs, dot) rounds.
6210#[cfg(target_arch = "x86_64")]
6211#[target_feature(enable = "avx2")]
6212unsafe fn dot_q4b_row_1x4_avx2(
6213    buf: &[u8],
6214    scales: &[u8],
6215    g0: usize,
6216    gpr: usize,
6217    xs: [&[i8]; 4],
6218) -> [f32; 4] {
6219    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6220    unsafe {
6221        use core::arch::x86_64::*;
6222        let ones = _mm256_set1_epi16(1);
6223        let mut acc = [0f32; 4];
6224        for gi in 0..gpr {
6225            let s = f16_to_f32(u16::from_le_bytes([
6226                scales[(g0 + gi) * 2],
6227                scales[(g0 + gi) * 2 + 1],
6228            ]));
6229            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6230            let aw = _mm256_abs_epi8(w);
6231            for (k, xq) in xs.iter().enumerate() {
6232                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6233                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6234                let d = _mm256_madd_epi16(p16, ones);
6235                let hi128 = _mm256_extracti128_si256::<1>(d);
6236                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6237                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6238                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6239                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
6240            }
6241        }
6242        acc
6243    }
6244}
6245
6246/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
6247#[cfg(target_arch = "x86_64")]
6248#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6249unsafe fn dot_q4b_row_1x4_vnni(
6250    buf: &[u8],
6251    scales: &[u8],
6252    g0: usize,
6253    gpr: usize,
6254    xs: [&[i8]; 4],
6255) -> [f32; 4] {
6256    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6257    unsafe {
6258        use core::arch::x86_64::*;
6259        let mut acc = [0f32; 4];
6260        for gi in 0..gpr {
6261            let s = f16_to_f32(u16::from_le_bytes([
6262                scales[(g0 + gi) * 2],
6263                scales[(g0 + gi) * 2 + 1],
6264            ]));
6265            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6266            let aw = _mm256_abs_epi8(w);
6267            for (k, xq) in xs.iter().enumerate() {
6268                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6269                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6270                acc[k] += d as f32 * s;
6271            }
6272        }
6273        acc
6274    }
6275}
6276
6277/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
6278/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
6279/// accumulation order (the q4_block flavor applies sx once at the end,
6280/// matching ITS single path; the two conventions are historical and
6281/// each blocked leg must mirror its own).
6282#[cfg(target_arch = "x86_64")]
6283#[target_feature(enable = "avx2")]
6284unsafe fn dot_q4b_row_1x4_sx_avx2(
6285    buf: &[u8],
6286    scales: &[u8],
6287    g0: usize,
6288    gpr: usize,
6289    xs: [&[i8]; 4],
6290    sxs: [f32; 4],
6291) -> [f32; 4] {
6292    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6293    unsafe {
6294        use core::arch::x86_64::*;
6295        let ones = _mm256_set1_epi16(1);
6296        let mut acc = [0f32; 4];
6297        for gi in 0..gpr {
6298            let s = f16_to_f32(u16::from_le_bytes([
6299                scales[(g0 + gi) * 2],
6300                scales[(g0 + gi) * 2 + 1],
6301            ]));
6302            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6303            let aw = _mm256_abs_epi8(w);
6304            for (k, xq) in xs.iter().enumerate() {
6305                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6306                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6307                let d = _mm256_madd_epi16(p16, ones);
6308                let hi128 = _mm256_extracti128_si256::<1>(d);
6309                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6310                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6311                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6312                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
6313            }
6314        }
6315        acc
6316    }
6317}
6318
6319/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
6320/// per-group `(d·sx)·s` fold mirrors the vbit single path).
6321#[cfg(target_arch = "x86_64")]
6322#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6323unsafe fn dot_q4b_row_1x4_sx_vnni(
6324    buf: &[u8],
6325    scales: &[u8],
6326    g0: usize,
6327    gpr: usize,
6328    xs: [&[i8]; 4],
6329    sxs: [f32; 4],
6330) -> [f32; 4] {
6331    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6332    unsafe {
6333        use core::arch::x86_64::*;
6334        let mut acc = [0f32; 4];
6335        for gi in 0..gpr {
6336            let s = f16_to_f32(u16::from_le_bytes([
6337                scales[(g0 + gi) * 2],
6338                scales[(g0 + gi) * 2 + 1],
6339            ]));
6340            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6341            let aw = _mm256_abs_epi8(w);
6342            for (k, xq) in xs.iter().enumerate() {
6343                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6344                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6345                acc[k] += (d as f32 * sxs[k]) * s;
6346            }
6347        }
6348        acc
6349    }
6350}
6351
6352#[allow(unreachable_code)]
6353fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6354    #[cfg(target_arch = "aarch64")]
6355    unsafe {
6356        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
6357    }
6358    #[cfg(target_arch = "x86_64")]
6359    unsafe {
6360        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
6361    }
6362    let mut acc = 0f32;
6363    for gi in 0..gpr {
6364        let g = g0 + gi;
6365        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6366        let mut d = 0i32;
6367        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6368            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6369                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6370        }
6371        acc += d as f32 * s;
6372    }
6373    acc
6374}
6375
6376/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
6377#[inline]
6378#[allow(unreachable_code)]
6379fn dot_q4_row_i8_2(
6380    packed: &[u8],
6381    scales: &[u8],
6382    g0: usize,
6383    gpr: usize,
6384    xq1: &[i8],
6385    xq2: &[i8],
6386) -> (f32, f32) {
6387    #[cfg(target_arch = "aarch64")]
6388    unsafe {
6389        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
6390    }
6391    #[cfg(target_arch = "x86_64")]
6392    unsafe {
6393        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
6394    }
6395    (
6396        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
6397        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
6398    )
6399}
6400
6401/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
6402/// multi-matrix jobs can drive it for several tensors in one dispatch).
6403#[allow(clippy::too_many_arguments)]
6404fn q4_range_a8w8(
6405    packed: &[u8],
6406    scales: &[u8],
6407    gpr: usize,
6408    cols: usize,
6409    act: &SplitAct,
6410    out: SendMut,
6411    start: usize,
6412    end: usize,
6413) {
6414    for r in start..end {
6415        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
6416        // xq is zeroed at outlier slots — add the exact terms.
6417        for &(j, xv) in &act.outliers {
6418            let flat = r * cols + j;
6419            let byte = packed[flat / 2];
6420            let nib = if flat & 1 == 0 {
6421                byte & 0x0F
6422            } else {
6423                byte >> 4
6424            };
6425            let s = f16_to_f32(u16::from_le_bytes([
6426                scales[(flat / GROUP_SIZE) * 2],
6427                scales[(flat / GROUP_SIZE) * 2 + 1],
6428            ]));
6429            acc += ((nib as i32 - 8) as f32) * s * xv;
6430        }
6431        // SAFETY: disjoint row ranges per worker.
6432        unsafe { *out.at(r) = acc };
6433    }
6434}
6435
6436/// Two-input q4 row range via the A8W8 int8 path — kernel body of
6437/// `q4matvec2`, extracted for pair multi-matrix jobs.
6438#[allow(clippy::too_many_arguments)]
6439fn q4_range2_a8w8(
6440    packed: &[u8],
6441    scales: &[u8],
6442    gpr: usize,
6443    cols: usize,
6444    a1: &SplitAct,
6445    a2: &SplitAct,
6446    p1: SendMut,
6447    p2: SendMut,
6448    start: usize,
6449    end: usize,
6450) {
6451    for r in start..end {
6452        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
6453        let mut acc1 = s1 * a1.sx;
6454        let mut acc2 = s2 * a2.sx;
6455        // xq is zeroed at outlier slots — add the exact terms.
6456        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
6457            for &(j, xv) in outliers {
6458                let flat = r * cols + j;
6459                let byte = packed[flat / 2];
6460                let nib = if flat & 1 == 0 {
6461                    byte & 0x0F
6462                } else {
6463                    byte >> 4
6464                };
6465                let s = f16_to_f32(u16::from_le_bytes([
6466                    scales[(flat / GROUP_SIZE) * 2],
6467                    scales[(flat / GROUP_SIZE) * 2 + 1],
6468                ]));
6469                *acc += ((nib as i32 - 8) as f32) * s * xv;
6470            }
6471        };
6472        fix(&a1.outliers, &mut acc1);
6473        fix(&a2.outliers, &mut acc2);
6474        // SAFETY: disjoint row ranges per worker.
6475        unsafe {
6476            *p1.at(r) = acc1;
6477            *p2.at(r) = acc2;
6478        }
6479    }
6480}
6481
6482/// Exact scalar q4 row range (same extraction, non-SDOT path).
6483fn q4_range_f32(
6484    packed: &[u8],
6485    scales: &[u8],
6486    gpr: usize,
6487    x: &[f32],
6488    out: SendMut,
6489    start: usize,
6490    end: usize,
6491) {
6492    for r in start..end {
6493        let mut acc = 0f32;
6494        for gi in 0..gpr {
6495            let g = r * gpr + gi;
6496            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6497            let pk = &packed[g * 16..(g + 1) * 16];
6498            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6499            let mut ga = 0f32;
6500            for (k, &b) in pk.iter().enumerate() {
6501                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6502                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6503            }
6504            acc += ga * s;
6505        }
6506        // SAFETY: disjoint row ranges per worker.
6507        unsafe { *out.at(r) = acc };
6508    }
6509}
6510
6511/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
6512/// dotted against both activations (was: two full matvecs — double
6513/// weight traffic). Per-lane math matches `q4matvec` exactly.
6514#[allow(clippy::too_many_arguments)]
6515fn q4matvec2(
6516    bytes: &[u8],
6517    x1: &[f32],
6518    x2: &[f32],
6519    rows: usize,
6520    cols: usize,
6521    o1: &mut [f32],
6522    o2: &mut [f32],
6523    pool: Option<&Pool>,
6524) {
6525    debug_assert_eq!(o1.len(), rows);
6526    debug_assert_eq!(o2.len(), rows);
6527    let (packed, scales) = q4_split(bytes, rows, cols);
6528    let gpr = cols / GROUP_SIZE;
6529
6530    if a8w8_enabled() {
6531        let a1 = split_act(x1);
6532        let a2 = split_act(x2);
6533        let p1 = SendMut(o1.as_mut_ptr());
6534        let p2 = SendMut(o2.as_mut_ptr());
6535        let run = move |start: usize, end: usize| {
6536            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
6537        };
6538        dispatch_rows(pool, rows, &run);
6539        return;
6540    }
6541
6542    let p1 = SendMut(o1.as_mut_ptr());
6543    let p2 = SendMut(o2.as_mut_ptr());
6544    let run = move |start: usize, end: usize| {
6545        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
6546    };
6547    dispatch_rows(pool, rows, &run);
6548}
6549
6550/// Two-input exact scalar q4 row range (same extraction).
6551#[allow(clippy::too_many_arguments)]
6552fn q4_range2_f32(
6553    packed: &[u8],
6554    scales: &[u8],
6555    gpr: usize,
6556    x1: &[f32],
6557    x2: &[f32],
6558    p1: SendMut,
6559    p2: SendMut,
6560    start: usize,
6561    end: usize,
6562) {
6563    for r in start..end {
6564        let (mut acc1, mut acc2) = (0f32, 0f32);
6565        for gi in 0..gpr {
6566            let g = r * gpr + gi;
6567            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6568            let pk = &packed[g * 16..(g + 1) * 16];
6569            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6570            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6571            let (mut g1, mut g2) = (0f32, 0f32);
6572            for (k, &b) in pk.iter().enumerate() {
6573                let wl = (b & 0x0F) as f32 - 8.0;
6574                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
6575                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
6576                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
6577            }
6578            acc1 += g1 * s;
6579            acc2 += g2 * s;
6580        }
6581        // SAFETY: disjoint row ranges per worker.
6582        unsafe {
6583            *p1.at(r) = acc1;
6584            *p2.at(r) = acc2;
6585        }
6586    }
6587}
6588
6589thread_local! {
6590    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
6591    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
6592    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
6593    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6594}
6595
6596/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
6597/// and dotted against ALL b activations (prefill used to fall back to b
6598/// full matvecs — b× weight traffic and b× nibble decode). Per-position
6599/// math matches `q4matvec` exactly: same group order, same accumulation.
6600/// `out` is row-major [b, rows] like `qmatmat`.
6601#[allow(clippy::too_many_arguments)]
6602fn q4matmat(
6603    bytes: &[u8],
6604    xs_all: &[f32],
6605    b: usize,
6606    rows: usize,
6607    cols: usize,
6608    out: &mut [f32],
6609    pool: Option<&Pool>,
6610) {
6611    debug_assert_eq!(xs_all.len(), b * cols);
6612    debug_assert_eq!(out.len(), b * rows);
6613    let (packed, scales) = q4_split(bytes, rows, cols);
6614    let gpr = cols / GROUP_SIZE;
6615    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6616
6617    if a8w8_enabled() {
6618        let acts: Vec<SplitAct> = (0..b)
6619            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6620            .collect();
6621        let acts = &acts;
6622        let out_addr = SendMut(out.as_mut_ptr());
6623        let run = move |start: usize, end: usize| {
6624            ROW_I8.with(|rb| {
6625                let mut buf = rb.borrow_mut();
6626                buf.resize(cols, 0);
6627                for r in start..end {
6628                    // Unpack the row's nibbles to centered i8 once
6629                    // (element 2k = low nibble, 2k+1 = high — flat order,
6630                    // same as dot_q4_row_sdot's zip).
6631                    for gi in 0..gpr {
6632                        let g = r * gpr + gi;
6633                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6634                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
6635                            buf[gi * GROUP_SIZE + k * 2 + 1] =
6636                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
6637                        }
6638                    }
6639                    let mut bi = 0usize;
6640                    #[cfg(target_arch = "x86_64")]
6641                    if avx2_enabled()
6642                        && std::env::var("CMF_X86_BLOCKED")
6643                            .map(|v| v != "0")
6644                            .unwrap_or(true)
6645                    {
6646                        while bi + 4 <= acts.len() {
6647                            let xs = [
6648                                acts[bi].xq.as_slice(),
6649                                acts[bi + 1].xq.as_slice(),
6650                                acts[bi + 2].xq.as_slice(),
6651                                acts[bi + 3].xq.as_slice(),
6652                            ];
6653                            let d = unsafe {
6654                                if vnni_tiles_enabled() {
6655                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
6656                                } else {
6657                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
6658                                }
6659                            };
6660                            for k in 0..4 {
6661                                let act = &acts[bi + k];
6662                                let mut acc = d[k] * act.sx;
6663                                for &(j, xv) in &act.outliers {
6664                                    acc += (buf[j] as i8) as f32
6665                                        * gscale((r * cols + j) / GROUP_SIZE)
6666                                        * xv;
6667                                }
6668                                // SAFETY: disjoint (bi, r) cells per worker.
6669                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6670                            }
6671                            bi += 4;
6672                        }
6673                    }
6674                    while bi < acts.len() {
6675                        let act = &acts[bi];
6676                        let mut acc = 0f32;
6677                        for gi in 0..gpr {
6678                            let d = dot_i8_i8(
6679                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6680                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6681                            );
6682                            acc += d as f32 * gscale(r * gpr + gi);
6683                        }
6684                        acc *= act.sx;
6685                        // xq is zeroed at outlier slots — exact terms.
6686                        for &(j, xv) in &act.outliers {
6687                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
6688                        }
6689                        // SAFETY: disjoint (bi, r) cells per worker row range.
6690                        unsafe { *out_addr.at(bi * rows + r) = acc };
6691                        bi += 1;
6692                    }
6693                }
6694            })
6695        };
6696        dispatch_rows(pool, rows, &run);
6697        return;
6698    }
6699
6700    let out_addr = SendMut(out.as_mut_ptr());
6701    let run = move |start: usize, end: usize| {
6702        ROW_F32.with(|rb| {
6703            let mut buf = rb.borrow_mut();
6704            buf.resize(cols, 0.0);
6705            for r in start..end {
6706                // Decode raw (nib − 8) values once; scales stay per-group
6707                // so the accumulation order matches q4matvec bit-for-bit.
6708                for gi in 0..gpr {
6709                    let g = r * gpr + gi;
6710                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6711                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
6712                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
6713                    }
6714                }
6715                for bi in 0..b {
6716                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6717                    let mut acc = 0f32;
6718                    for gi in 0..gpr {
6719                        let mut ga = 0f32;
6720                        // Pairwise (lo + hi) addition, matching
6721                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
6722                        // a flat one-per-element loop rounds differently
6723                        // and broke bit-parity on the scalar (x86) path.
6724                        for k in 0..GROUP_SIZE / 2 {
6725                            let e = gi * GROUP_SIZE + k * 2;
6726                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
6727                        }
6728                        acc += ga * gscale(r * gpr + gi);
6729                    }
6730                    // SAFETY: disjoint (bi, r) cells per worker row range.
6731                    unsafe { *out_addr.at(bi * rows + r) = acc };
6732                }
6733            }
6734        })
6735    };
6736    dispatch_rows(pool, rows, &run);
6737}
6738
6739/// Batched vbit matmat: each variable-bit row is decoded from the mmap
6740/// ONCE for the whole microbatch. Same per-position math as
6741/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
6742/// and the scalar path).
6743#[allow(clippy::too_many_arguments)]
6744fn vbitmatmat(
6745    bytes: &[u8],
6746    offsets: &[usize],
6747    xs_all: &[f32],
6748    b: usize,
6749    rows: usize,
6750    cols: usize,
6751    out: &mut [f32],
6752    pool: Option<&Pool>,
6753) {
6754    debug_assert_eq!(xs_all.len(), b * cols);
6755    debug_assert_eq!(out.len(), b * rows);
6756    debug_assert_eq!(offsets.len(), rows + 1);
6757    let ng = cols / GROUP_SIZE;
6758    let bits = &bytes[..rows];
6759    let sc_off = rows;
6760    let gscale = |r: usize, g: usize| {
6761        let so = (r * ng + g) * 2;
6762        f16_to_f32(u16::from_le_bytes([
6763            bytes[sc_off + so],
6764            bytes[sc_off + so + 1],
6765        ]))
6766    };
6767
6768    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
6769    let decode_f32 = |r: usize, dst: &mut [f32]| {
6770        let bw = bits[r] as usize;
6771        let l = ((1i32 << (bw - 1)) - 1) as f32;
6772        let data = &bytes[offsets[r]..offsets[r + 1]];
6773        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
6774        for d in dst.iter_mut() {
6775            while nbits < bw {
6776                acc = (acc << 8) | data[idx] as u64;
6777                idx += 1;
6778                nbits += 8;
6779            }
6780            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
6781            nbits -= bw;
6782            *d = u - l;
6783        }
6784    };
6785
6786    if a8w8_enabled() {
6787        let acts: Vec<SplitAct> = (0..b)
6788            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6789            .collect();
6790        let acts = &acts;
6791        let out_addr = SendMut(out.as_mut_ptr());
6792        let run = move |start: usize, end: usize| {
6793            for r in start..end {
6794                let bw = bits[r] as usize;
6795                if bw == 8 {
6796                    // u−L reaches 128 → no i8 path; decode once, exact
6797                    // f32 dots for every position (same as vbitmatvec).
6798                    ROW_F32.with(|rb| {
6799                        let mut buf = rb.borrow_mut();
6800                        buf.resize(cols, 0.0);
6801                        decode_f32(r, &mut buf);
6802                        for bi in 0..b {
6803                            let x = &xs_all[bi * cols..(bi + 1) * cols];
6804                            let mut dot = 0f32;
6805                            for g in 0..ng {
6806                                let mut gd = 0f32;
6807                                for k in 0..GROUP_SIZE {
6808                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6809                                }
6810                                dot += gd * gscale(r, g);
6811                            }
6812                            // SAFETY: disjoint (bi, r) cells per worker range.
6813                            unsafe { *out_addr.at(bi * rows + r) = dot };
6814                        }
6815                    });
6816                    continue;
6817                }
6818                let l = (1i32 << (bw - 1)) - 1;
6819                let data = &bytes[offsets[r]..offsets[r + 1]];
6820                ROW_I8.with(|rb| {
6821                    let mut buf = rb.borrow_mut();
6822                    buf.resize(cols, 0);
6823                    #[inline(always)]
6824                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
6825                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
6826                            let u = unpack8::<B>(&data[blk * B..]);
6827                            for k in 0..8 {
6828                                chunk[k] = (u[k] - l) as i8 as u8;
6829                            }
6830                        }
6831                    }
6832                    match bw {
6833                        3 => fill::<3>(data, l, &mut buf),
6834                        4 => vbit_fill4(data, &mut buf),
6835                        5 => fill::<5>(data, l, &mut buf),
6836                        6 => fill::<6>(data, l, &mut buf),
6837                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
6838                    }
6839                    let mut bi = 0usize;
6840                    // The vbit scale table shares q4_block's layout
6841                    // (contiguous f16 per (row·ng + g)), so the same
6842                    // blocked 1×4 kernel serves the decoded row.
6843                    #[cfg(target_arch = "x86_64")]
6844                    if avx2_enabled()
6845                        && std::env::var("CMF_X86_BLOCKED")
6846                            .map(|v| v != "0")
6847                            .unwrap_or(true)
6848                    {
6849                        while bi + 4 <= acts.len() {
6850                            let xs = [
6851                                acts[bi].xq.as_slice(),
6852                                acts[bi + 1].xq.as_slice(),
6853                                acts[bi + 2].xq.as_slice(),
6854                                acts[bi + 3].xq.as_slice(),
6855                            ];
6856                            let sxs = [
6857                                acts[bi].sx,
6858                                acts[bi + 1].sx,
6859                                acts[bi + 2].sx,
6860                                acts[bi + 3].sx,
6861                            ];
6862                            let d = unsafe {
6863                                if vnni_tiles_enabled() {
6864                                    dot_q4b_row_1x4_sx_vnni(
6865                                        &buf,
6866                                        &bytes[sc_off..],
6867                                        r * ng,
6868                                        ng,
6869                                        xs,
6870                                        sxs,
6871                                    )
6872                                } else {
6873                                    dot_q4b_row_1x4_sx_avx2(
6874                                        &buf,
6875                                        &bytes[sc_off..],
6876                                        r * ng,
6877                                        ng,
6878                                        xs,
6879                                        sxs,
6880                                    )
6881                                }
6882                            };
6883                            for k in 0..4 {
6884                                let act = &acts[bi + k];
6885                                let mut dot = d[k];
6886                                for &(j, xv) in &act.outliers {
6887                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6888                                }
6889                                // SAFETY: disjoint (bi, r) cells per worker.
6890                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
6891                            }
6892                            bi += 4;
6893                        }
6894                    }
6895                    while bi < acts.len() {
6896                        let act = &acts[bi];
6897                        let mut dot = 0f32;
6898                        for g in 0..ng {
6899                            let d = dot_i8_i8(
6900                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6901                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6902                            ) as f32
6903                                * act.sx;
6904                            dot += d * gscale(r, g);
6905                        }
6906                        for &(j, xv) in &act.outliers {
6907                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6908                        }
6909                        // SAFETY: disjoint (bi, r) cells per worker range.
6910                        unsafe { *out_addr.at(bi * rows + r) = dot };
6911                        bi += 1;
6912                    }
6913                });
6914            }
6915        };
6916        dispatch_rows(pool, rows, &run);
6917        return;
6918    }
6919
6920    let out_addr = SendMut(out.as_mut_ptr());
6921    let run = move |start: usize, end: usize| {
6922        ROW_F32.with(|rb| {
6923            let mut buf = rb.borrow_mut();
6924            buf.resize(cols, 0.0);
6925            for r in start..end {
6926                decode_f32(r, &mut buf);
6927                for bi in 0..b {
6928                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6929                    let mut dot = 0f32;
6930                    for g in 0..ng {
6931                        let mut gd = 0f32;
6932                        for k in 0..GROUP_SIZE {
6933                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6934                        }
6935                        dot += gd * gscale(r, g);
6936                    }
6937                    // SAFETY: disjoint (bi, r) cells per worker range.
6938                    unsafe { *out_addr.at(bi * rows + r) = dot };
6939                }
6940            }
6941        })
6942    };
6943    dispatch_rows(pool, rows, &run);
6944}
6945
6946/// Build a GPU batch job for a q8-family mapped tensor (primary
6947/// shard): prescaled input + directory coordinates. None → not
6948/// GPU-eligible, caller stays on the CPU.
6949pub(crate) fn gpu_batch_job<'a>(
6950    t: &'a QTensor,
6951    x: &[f32],
6952) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
6953    match t {
6954        QTensor::Mapped {
6955            model,
6956            idx,
6957            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
6958            rows,
6959            cols,
6960            row_scale,
6961            col_field,
6962            ..
6963        } => Some((
6964            model.clone(),
6965            crate::gpu::BatchJob {
6966                idx: *idx,
6967                rows: *rows,
6968                cols: *cols,
6969                row_scale,
6970                xs: prescale(x, col_field, *dt).into_owned(),
6971                layout: crate::gpu::BatchLayout::Q8,
6972            },
6973        )),
6974        // q1: raw f32 activations, tile-embedded scales.
6975        QTensor::Mapped {
6976            model,
6977            idx,
6978            dtype: TensorDtype::Q1,
6979            rows,
6980            cols,
6981            ..
6982        } => Some((
6983            model.clone(),
6984            crate::gpu::BatchJob {
6985                idx: *idx,
6986                rows: *rows,
6987                cols: *cols,
6988                row_scale: &[],
6989                xs: x.to_vec(),
6990                layout: crate::gpu::BatchLayout::Q1,
6991            },
6992        )),
6993        _ => None,
6994    }
6995}
6996
6997thread_local! {
6998    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6999    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7000}
7001
7002pub(crate) fn prescale<'a>(
7003    x: &'a [f32],
7004    col_field: &[f32],
7005    dtype: TensorDtype,
7006) -> std::borrow::Cow<'a, [f32]> {
7007    if dtype == TensorDtype::Q8_2f {
7008        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
7009    } else {
7010        std::borrow::Cow::Borrowed(x)
7011    }
7012}
7013
7014/// θ col-field fold for q8_2f activations. Borrowed pass-through for
7015/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
7016pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
7017    x: &[f32],
7018    col_field: &[f32],
7019    dtype: TensorDtype,
7020    buf_id: u8,
7021    f: F,
7022) -> R {
7023    if dtype == TensorDtype::Q8_2f {
7024        if buf_id == 1 {
7025            PRESCALE_BUF1.with(|b| {
7026                let mut buf = b.borrow_mut();
7027                buf.clear();
7028                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7029                f(&buf)
7030            })
7031        } else {
7032            PRESCALE_BUF2.with(|b| {
7033                let mut buf = b.borrow_mut();
7034                buf.clear();
7035                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7036                f(&buf)
7037            })
7038        }
7039    } else {
7040        f(x)
7041    }
7042}
7043
7044// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
7045
7046/// AVX2+FMA available? Default ON when the CPU supports both;
7047/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
7048#[cfg(target_arch = "x86_64")]
7049pub(crate) fn avx2_enabled() -> bool {
7050    use std::sync::OnceLock;
7051    static ON: OnceLock<bool> = OnceLock::new();
7052    *ON.get_or_init(|| {
7053        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
7054            && std::arch::is_x86_feature_detected!("avx2")
7055            && std::arch::is_x86_feature_detected!("fma")
7056    })
7057}
7058
7059/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
7060/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
7061/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
7062/// active either way, they are exact (regrouped sums only).
7063#[cfg(target_arch = "x86_64")]
7064fn avx2_a8w8_enabled() -> bool {
7065    use std::sync::OnceLock;
7066    static ON: OnceLock<bool> = OnceLock::new();
7067    *ON.get_or_init(|| {
7068        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
7069    })
7070}
7071
7072/// A8W8 quantized-activation path available on THIS machine? One
7073/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
7074/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
7075#[inline]
7076pub(crate) fn a8w8_enabled() -> bool {
7077    #[cfg(target_arch = "aarch64")]
7078    {
7079        sdot_enabled()
7080    }
7081    #[cfg(target_arch = "x86_64")]
7082    {
7083        avx2_a8w8_enabled()
7084    }
7085    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7086    {
7087        false
7088    }
7089}
7090
7091/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
7092/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
7093#[inline]
7094#[allow(unreachable_code)]
7095fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
7096    #[cfg(target_arch = "aarch64")]
7097    unsafe {
7098        return dot_i8_sdot(w, xq);
7099    }
7100    #[cfg(target_arch = "x86_64")]
7101    unsafe {
7102        if avx512vnni_enabled() {
7103            return dot_i8_i8_vnni(w, xq);
7104        }
7105        return dot_i8_i8_avx2(w, xq);
7106    }
7107    w.iter()
7108        .zip(xq)
7109        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
7110        .sum()
7111}
7112
7113/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
7114/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
7115/// `vpdpbusd` encoding.
7116#[cfg(target_arch = "x86_64")]
7117fn avx512vnni_enabled() -> bool {
7118    use std::sync::OnceLock;
7119    static ON: OnceLock<bool> = OnceLock::new();
7120    *ON.get_or_init(|| {
7121        std::env::var("CMF_AVX512")
7122            .map(|v| v != "0")
7123            .unwrap_or(true)
7124            && std::arch::is_x86_feature_detected!("avx512f")
7125            && std::arch::is_x86_feature_detected!("avx512bw")
7126            && std::arch::is_x86_feature_detected!("avx512vl")
7127            && std::arch::is_x86_feature_detected!("avx512vnni")
7128    })
7129}
7130
7131/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
7132/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
7133/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
7134/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
7135/// (+4%) — consistent, no leg regressed. The tile kernels keep a
7136/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
7137/// smaller than the long-dot q8 win (+13%), but it is real and free.
7138#[cfg(target_arch = "x86_64")]
7139fn vnni_tiles_enabled() -> bool {
7140    use std::sync::OnceLock;
7141    static ON: OnceLock<bool> = OnceLock::new();
7142    *ON.get_or_init(|| {
7143        std::env::var("CMF_VNNI_TILES")
7144            .map(|v| v != "0")
7145            .unwrap_or(true)
7146            && avx512vnni_enabled()
7147    })
7148}
7149
7150/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
7151/// plus the same horizontal reduce the AVX2 kernels use. Products are
7152/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
7153/// is bit-identical to the maddubs+madd pair it replaces.
7154#[cfg(target_arch = "x86_64")]
7155#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7156#[inline]
7157unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
7158    // SAFETY: pure register math.
7159    unsafe {
7160        use core::arch::x86_64::*;
7161        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
7162        let hi128 = _mm256_extracti128_si256::<1>(d);
7163        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7164        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7165        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7166        _mm_cvtsi128_si32(s32)
7167    }
7168}
7169
7170/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
7171/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
7172/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
7173/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
7174#[cfg(target_arch = "x86_64")]
7175#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7176unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7177    // SAFETY: callers uphold slice-length contracts (see call sites).
7178    unsafe {
7179        use core::arch::x86_64::*;
7180        let n = w.len();
7181        let mut j = 0usize;
7182        let mut total: i32;
7183        // 4 independent accumulators: vpdpbusd is its own loop-carried
7184        // dependency (~5-cycle latency) — a single-acc loop runs
7185        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
7186        // on Granite Rapids.
7187        {
7188            #[inline(always)]
7189            unsafe fn step(
7190                w: *const u8,
7191                x: *const i8,
7192                acc: core::arch::x86_64::__m512i,
7193            ) -> core::arch::x86_64::__m512i {
7194                unsafe {
7195                    use core::arch::x86_64::*;
7196                    let wv = _mm512_loadu_si512(w as *const _);
7197                    let xv = _mm512_loadu_si512(x as *const _);
7198                    let aw = _mm512_abs_epi8(wv);
7199                    let neg = _mm512_movepi8_mask(wv);
7200                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
7201                    _mm512_dpbusd_epi32(acc, aw, sx)
7202                }
7203            }
7204            let (mut a0, mut a1, mut a2, mut a3) = (
7205                _mm512_setzero_si512(),
7206                _mm512_setzero_si512(),
7207                _mm512_setzero_si512(),
7208                _mm512_setzero_si512(),
7209            );
7210            while j + 256 <= n {
7211                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7212                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
7213                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
7214                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
7215                j += 256;
7216            }
7217            while j + 64 <= n {
7218                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7219                j += 64;
7220            }
7221            let s01 = _mm512_add_epi32(a0, a1);
7222            let s23 = _mm512_add_epi32(a2, a3);
7223            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
7224        }
7225        // 32-wide (q4/vbit groups are exactly 32 bytes).
7226        if j + 32 <= n {
7227            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7228            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7229            let d = _mm256_dpbusd_epi32(
7230                _mm256_setzero_si256(),
7231                _mm256_abs_epi8(wv),
7232                _mm256_sign_epi8(xv, wv),
7233            );
7234            let hi128 = _mm256_extracti128_si256::<1>(d);
7235            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7236            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7237            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7238            total += _mm_cvtsi128_si32(s32);
7239            j += 32;
7240        }
7241        while j < n {
7242            total += (w[j] as i8) as i32 * xq[j] as i32;
7243            j += 1;
7244        }
7245        total
7246    }
7247}
7248
7249/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
7250#[cfg(target_arch = "x86_64")]
7251#[target_feature(enable = "avx2,fma")]
7252unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
7253    // SAFETY: callers uphold slice-length contracts (see call sites).
7254    unsafe {
7255        use core::arch::x86_64::*;
7256        let n = x.len();
7257        let wp = w.as_ptr();
7258        let xp = x.as_ptr();
7259        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
7260        let mut j = 0usize;
7261        while j + 16 <= n {
7262            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
7263            let lo = _mm256_cvtepi8_epi32(wb);
7264            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
7265            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
7266            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
7267            j += 16;
7268        }
7269        let acc = _mm256_add_ps(a0, a1);
7270        let hi128 = _mm256_extractf128_ps::<1>(acc);
7271        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
7272        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
7273        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
7274        let mut sum = _mm_cvtss_f32(s32);
7275        while j < n {
7276            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
7277            j += 1;
7278        }
7279        sum
7280    }
7281}
7282
7283/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
7284/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
7285/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
7286/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
7287#[cfg(target_arch = "x86_64")]
7288#[target_feature(enable = "avx2")]
7289unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
7290    // SAFETY: callers uphold slice-length contracts (see call sites).
7291    unsafe {
7292        use core::arch::x86_64::*;
7293        let n = w.len();
7294        let ones = _mm256_set1_epi16(1);
7295        let mut acc = _mm256_setzero_si256();
7296        let mut j = 0usize;
7297        while j + 32 <= n {
7298            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7299            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7300            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7301            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
7302            j += 32;
7303        }
7304        let hi128 = _mm256_extracti128_si256::<1>(acc);
7305        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
7306        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7307        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7308        let mut s = _mm_cvtsi128_si32(s32);
7309        while j < n {
7310            s += (w[j] as i8) as i32 * xq[j] as i32;
7311            j += 1;
7312        }
7313        s
7314    }
7315}
7316
7317/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
7318/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
7319/// slice as a combined 2×8 register and meets two activation pairs.
7320#[cfg(target_arch = "aarch64")]
7321#[target_feature(enable = "neon,i8mm")]
7322unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7323    // SAFETY: callers uphold slice-length contracts.
7324    unsafe {
7325        use core::arch::aarch64::*;
7326        use core::arch::asm;
7327        let n = w0.len();
7328        let w0p = w0.as_ptr() as *const i8;
7329        let w1p = w1.as_ptr() as *const i8;
7330        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
7331        // same for x2/x3.
7332        let mut acc01 = vdupq_n_s32(0);
7333        let mut acc23 = vdupq_n_s32(0);
7334        let mut i = 0usize;
7335        while i + 8 <= n {
7336            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
7337            let xb01 = vcombine_s8(
7338                vld1_s8(xs[0].as_ptr().add(i)),
7339                vld1_s8(xs[1].as_ptr().add(i)),
7340            );
7341            let xb23 = vcombine_s8(
7342                vld1_s8(xs[2].as_ptr().add(i)),
7343                vld1_s8(xs[3].as_ptr().add(i)),
7344            );
7345            asm!(
7346                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
7347                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
7348                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
7349                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
7350                options(pure, nomem, nostack),
7351            );
7352            i += 8;
7353        }
7354        let mut out = [[0i32; 4]; 2];
7355        let a01: [i32; 4] = core::mem::transmute(acc01);
7356        let a23: [i32; 4] = core::mem::transmute(acc23);
7357        out[0][0] = a01[0];
7358        out[0][1] = a01[1];
7359        out[1][0] = a01[2];
7360        out[1][1] = a01[3];
7361        out[0][2] = a23[0];
7362        out[0][3] = a23[1];
7363        out[1][2] = a23[2];
7364        out[1][3] = a23[3];
7365        if i < n {
7366            for (k, x) in xs.iter().enumerate() {
7367                for j in i..n {
7368                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7369                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7370                }
7371            }
7372        }
7373        out
7374    }
7375}
7376
7377/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
7378/// registers across four activation streams, eight sdot accumulators.
7379/// (The per-row form re-read each W row once per activation.)
7380#[cfg(target_arch = "aarch64")]
7381#[target_feature(enable = "neon,dotprod")]
7382unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7383    // SAFETY: callers uphold slice-length contracts.
7384    unsafe {
7385        use core::arch::aarch64::*;
7386        use core::arch::asm;
7387        let n = w0.len();
7388        let w0p = w0.as_ptr() as *const i8;
7389        let w1p = w1.as_ptr() as *const i8;
7390        let mut acc = [[vdupq_n_s32(0); 4]; 2];
7391        let mut i = 0usize;
7392        while i + 16 <= n {
7393            let wv0 = vld1q_s8(w0p.add(i));
7394            let wv1 = vld1q_s8(w1p.add(i));
7395            for (k, x) in xs.iter().enumerate() {
7396                let xv = vld1q_s8(x.as_ptr().add(i));
7397                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
7398                asm!(
7399                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
7400                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
7401                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7402                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
7403                    options(pure, nomem, nostack),
7404                );
7405                acc[0][k] = a0;
7406                acc[1][k] = a1;
7407            }
7408            i += 16;
7409        }
7410        let mut out = [[0i32; 4]; 2];
7411        for r in 0..2 {
7412            for k in 0..4 {
7413                out[r][k] = vaddvq_s32(acc[r][k]);
7414            }
7415        }
7416        if i < n {
7417            for (k, x) in xs.iter().enumerate() {
7418                for j in i..n {
7419                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7420                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7421                }
7422            }
7423        }
7424        out
7425    }
7426}
7427
7428/// Blocked 2 weight rows × 4 activations for the prefill GEMM
7429/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
7430/// abs() live in registers across all four activation streams; the
7431/// sign-fixup is recomputed per pair (the price of the maddubs trick).
7432/// Returns raw i8·i8 dots; the caller applies scales and outliers.
7433#[cfg(target_arch = "x86_64")]
7434#[target_feature(enable = "avx2")]
7435unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7436    // SAFETY: callers uphold slice-length contracts.
7437    unsafe {
7438        use core::arch::x86_64::*;
7439        let n = w0.len();
7440        let ones = _mm256_set1_epi16(1);
7441        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
7442        let mut j = 0usize;
7443        while j + 32 <= n {
7444            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
7445            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
7446            let aw0 = _mm256_abs_epi8(wv0);
7447            let aw1 = _mm256_abs_epi8(wv1);
7448            for (k, x) in xs.iter().enumerate() {
7449                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
7450                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
7451                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
7452                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
7453                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
7454            }
7455            j += 32;
7456        }
7457        let mut out = [[0i32; 4]; 2];
7458        for r in 0..2 {
7459            for k in 0..4 {
7460                let a = acc[r][k];
7461                let hi128 = _mm256_extracti128_si256::<1>(a);
7462                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
7463                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7464                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7465                out[r][k] = _mm_cvtsi128_si32(s32);
7466            }
7467        }
7468        if j < n {
7469            for (k, x) in xs.iter().enumerate() {
7470                for i in j..n {
7471                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
7472                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
7473                }
7474            }
7475        }
7476        out
7477    }
7478}
7479
7480/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
7481/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
7482/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
7483/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
7484#[cfg(target_arch = "x86_64")]
7485#[inline]
7486fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
7487    let dot = if avx512vnni_enabled() && row.len() >= 64 {
7488        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
7489    } else {
7490        unsafe { dot_i8_i8_avx2(row, &act.xq) }
7491    };
7492    let mut acc = dot as f32 * act.sx;
7493    for &(j, xv) in &act.outliers {
7494        acc += (row[j] as i8) as f32 * xv;
7495    }
7496    acc
7497}
7498
7499/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
7500/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
7501/// a single-acc loop runs latency-bound, measured on Granite Rapids).
7502#[cfg(target_arch = "x86_64")]
7503#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7504unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7505    // SAFETY: callers uphold slice-length contracts (see call sites).
7506    unsafe {
7507        use core::arch::x86_64::*;
7508        let n = w.len();
7509        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
7510        #[inline(always)]
7511        unsafe fn step(
7512            w: *const u8,
7513            x: *const i8,
7514            flip: core::arch::x86_64::__m512i,
7515            acc: core::arch::x86_64::__m512i,
7516        ) -> core::arch::x86_64::__m512i {
7517            unsafe {
7518                use core::arch::x86_64::*;
7519                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
7520                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
7521            }
7522        }
7523        let (mut a0, mut a1, mut a2, mut a3) = (
7524            _mm512_setzero_si512(),
7525            _mm512_setzero_si512(),
7526            _mm512_setzero_si512(),
7527            _mm512_setzero_si512(),
7528        );
7529        let mut j = 0usize;
7530        while j + 256 <= n {
7531            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7532            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
7533            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
7534            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
7535            j += 256;
7536        }
7537        while j + 64 <= n {
7538            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7539            j += 64;
7540        }
7541        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
7542            _mm512_add_epi32(a0, a1),
7543            _mm512_add_epi32(a2, a3),
7544        ));
7545        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
7546        while j < n {
7547            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
7548            j += 1;
7549        }
7550        total
7551    }
7552}
7553
7554/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
7555/// writer's flat order, same as the NEON vzip pair), maddubs against
7556/// the pre-quantized activation group, × the group's f16 scale. Pair
7557/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
7558/// `dot_q4_row_sdot`.
7559#[cfg(target_arch = "x86_64")]
7560#[target_feature(enable = "avx2")]
7561unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7562    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7563    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
7564    unsafe {
7565        use core::arch::x86_64::*;
7566        let lomask = _mm_set1_epi8(0x0F);
7567        let eight = _mm256_set1_epi8(8);
7568        let ones = _mm256_set1_epi16(1);
7569        let mut acc = 0f32;
7570        for gi in 0..gpr {
7571            let g = g0 + gi;
7572            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7573            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7574            let lo = _mm_and_si128(b, lomask);
7575            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7576            let w = _mm256_sub_epi8(
7577                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7578                eight,
7579            );
7580            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7581            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
7582            let d = _mm256_madd_epi16(p16, ones);
7583            let hi128 = _mm256_extracti128_si256::<1>(d);
7584            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7585            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7586            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7587            acc += _mm_cvtsi128_si32(s32) as f32 * s;
7588        }
7589        acc
7590    }
7591}
7592
7593/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
7594/// both activations dotted against the same centered i8 register.
7595#[cfg(target_arch = "x86_64")]
7596#[target_feature(enable = "avx2")]
7597unsafe fn dot_q4_row_avx2_2(
7598    packed: &[u8],
7599    scales: &[u8],
7600    g0: usize,
7601    gpr: usize,
7602    xq1: &[i8],
7603    xq2: &[i8],
7604) -> (f32, f32) {
7605    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
7606    unsafe {
7607        use core::arch::x86_64::*;
7608        let lomask = _mm_set1_epi8(0x0F);
7609        let eight = _mm256_set1_epi8(8);
7610        let ones = _mm256_set1_epi16(1);
7611        let (mut acc1, mut acc2) = (0f32, 0f32);
7612        #[inline(always)]
7613        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
7614            unsafe {
7615                use core::arch::x86_64::*;
7616                let hi128 = _mm256_extracti128_si256::<1>(d);
7617                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7618                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7619                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7620                _mm_cvtsi128_si32(s32)
7621            }
7622        }
7623        for gi in 0..gpr {
7624            let g = g0 + gi;
7625            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7626            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7627            let lo = _mm_and_si128(b, lomask);
7628            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7629            let w = _mm256_sub_epi8(
7630                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7631                eight,
7632            );
7633            let aw = _mm256_abs_epi8(w);
7634            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7635            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7636            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
7637            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
7638            acc1 += hsum(d1) as f32 * s;
7639            acc2 += hsum(d2) as f32 * s;
7640        }
7641        (acc1, acc2)
7642    }
7643}
7644
7645/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
7646#[cfg(target_arch = "x86_64")]
7647fn q8_range_avx2(
7648    q: &[u8],
7649    row_scale: &[f32],
7650    act: &SplitAct,
7651    cols: usize,
7652    out_addr: SendMut,
7653    start: usize,
7654    end: usize,
7655) {
7656    for o in start..end {
7657        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7658        // SAFETY: disjoint row ranges per worker.
7659        unsafe { *out_addr.at(o) = v };
7660    }
7661}
7662
7663/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
7664#[cfg(target_arch = "x86_64")]
7665#[allow(clippy::too_many_arguments)]
7666fn q8_range2_avx2(
7667    q: &[u8],
7668    row_scale: &[f32],
7669    a1: &SplitAct,
7670    a2: &SplitAct,
7671    cols: usize,
7672    p1: SendMut,
7673    p2: SendMut,
7674    start: usize,
7675    end: usize,
7676) {
7677    for o in start..end {
7678        let row = &q[o * cols..(o + 1) * cols];
7679        // SAFETY: disjoint row ranges per worker.
7680        unsafe {
7681            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
7682            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
7683        }
7684    }
7685}
7686
7687// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
7688
7689/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
7690/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
7691/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
7692/// accumulator dependency chain swamp the MAC advantage, and Apple's
7693/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
7694/// field trials on Cortex-A710/X-class parts with two pipes, where the
7695/// balance may differ; a pre-interleaved weight layout (repack infra)
7696/// is the known path if it ever earns its keep.
7697#[cfg(target_arch = "aarch64")]
7698fn i8mm_enabled() -> bool {
7699    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7700    *ON.get_or_init(|| {
7701        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
7702            && std::arch::is_aarch64_feature_detected!("i8mm")
7703    })
7704}
7705
7706/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
7707/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
7708/// (On non-ARM release builds only the test tolerance switch calls it.)
7709#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
7710fn sdot_enabled() -> bool {
7711    use std::sync::OnceLock;
7712    static ON: OnceLock<bool> = OnceLock::new();
7713    *ON.get_or_init(|| {
7714        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
7715        if !want {
7716            return false;
7717        }
7718
7719        #[cfg(target_arch = "aarch64")]
7720        {
7721            if std::arch::is_aarch64_feature_detected!("dotprod") {
7722                return true;
7723            }
7724            #[cfg(target_os = "android")]
7725            {
7726                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
7727                    if cpuinfo.lines().any(|l| {
7728                        (l.starts_with("Features") || l.starts_with("features"))
7729                            && l.contains("asimddp")
7730                    }) {
7731                        return true;
7732                    }
7733                }
7734            }
7735            false
7736        }
7737        #[cfg(not(target_arch = "aarch64"))]
7738        {
7739            false
7740        }
7741    })
7742}
7743
7744/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
7745/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
7746/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
7747/// matvec, shared by all rows/workers.
7748struct SplitAct {
7749    xq: Vec<i8>,
7750    sx: f32,
7751    outliers: Vec<(usize, f32)>,
7752    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
7753    /// `−128·Σx`); one i32 per split, computed once per matvec.
7754    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
7755    xsum: i32,
7756}
7757
7758thread_local! {
7759    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
7760    /// and its hidden-size allocation was steady-state heap churn.
7761    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
7762        const { std::cell::RefCell::new(Vec::new()) };
7763}
7764
7765impl Drop for SplitAct {
7766    fn drop(&mut self) {
7767        let buf = std::mem::take(&mut self.xq);
7768        if buf.capacity() > 0 {
7769            XQ_FREE.with(|f| {
7770                let mut f = f.borrow_mut();
7771                if f.len() < 16 {
7772                    f.push(buf);
7773                }
7774            });
7775        }
7776    }
7777}
7778
7779fn split_act(x: &[f32]) -> SplitAct {
7780    let n = x.len();
7781    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
7782    let thr = 8.0 * rms;
7783    // One pass: collect outliers and the bulk absmax (outliers excluded —
7784    // identical to the old zero-then-fold over a copied buffer, minus the
7785    // full-vector copy).
7786    let mut outliers: Vec<(usize, f32)> = Vec::new();
7787    let mut amax = 0f32;
7788    for (j, &v) in x.iter().enumerate() {
7789        let a = v.abs();
7790        if a > thr {
7791            outliers.push((j, v));
7792        } else if a > amax {
7793            amax = a;
7794        }
7795    }
7796    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7797    let inv = 1.0 / sx;
7798    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7799    xq.clear();
7800    xq.reserve(n);
7801    if outliers.is_empty() {
7802        xq.extend(
7803            x.iter()
7804                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
7805        );
7806    } else {
7807        // Outlier slots quantize to 0 (their exact term is added later).
7808        xq.extend(x.iter().map(|&v| {
7809            if v.abs() > thr {
7810                0
7811            } else {
7812                (v * inv).round().clamp(-127.0, 127.0) as i8
7813            }
7814        }));
7815    }
7816    let xsum = xq.iter().map(|&v| v as i32).sum();
7817    SplitAct {
7818        xq,
7819        sx,
7820        outliers,
7821        xsum,
7822    }
7823}
7824
7825fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
7826    let n = x.len();
7827    let rms = (x
7828        .iter()
7829        .zip(col)
7830        .map(|(&a, &c)| {
7831            let v = a * c;
7832            (v * v) as f64
7833        })
7834        .sum::<f64>()
7835        / n.max(1) as f64)
7836        .sqrt() as f32;
7837    let thr = 8.0 * rms;
7838
7839    let mut outliers = Vec::new();
7840    let mut amax = 0f32;
7841    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
7842        let v = a * c;
7843        let s = v.abs();
7844        if s > thr {
7845            outliers.push((j, v));
7846        } else if s > amax {
7847            amax = s;
7848        }
7849    }
7850
7851    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7852    let inv = 1.0 / sx;
7853    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7854    xq.clear();
7855    xq.reserve(n);
7856    if outliers.is_empty() {
7857        xq.extend(
7858            x.iter()
7859                .zip(col)
7860                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
7861        );
7862    } else {
7863        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
7864            let v = a * c;
7865            if v.abs() > thr {
7866                0
7867            } else {
7868                (v * inv).round().clamp(-127.0, 127.0) as i8
7869            }
7870        }));
7871    }
7872    let xsum = xq.iter().map(|&v| v as i32).sum();
7873    SplitAct {
7874        xq,
7875        sx,
7876        outliers,
7877        xsum,
7878    }
7879}
7880
7881/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
7882/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
7883#[cfg(target_arch = "aarch64")]
7884#[target_feature(enable = "neon,dotprod")]
7885unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
7886    // SAFETY: callers uphold slice-length contracts (see call sites).
7887    unsafe {
7888        use core::arch::aarch64::*;
7889        use core::arch::asm;
7890        let wp = w.as_ptr() as *const i8;
7891        let n = w.len();
7892        let (mut a0, mut a1, mut a2, mut a3) = (
7893            vdupq_n_s32(0),
7894            vdupq_n_s32(0),
7895            vdupq_n_s32(0),
7896            vdupq_n_s32(0),
7897        );
7898        let mut i = 0;
7899        while i + 64 <= n {
7900            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7901            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
7902            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
7903            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
7904            asm!(
7905                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7906                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7907                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
7908                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
7909                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7910                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7911                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
7912                options(pure, nomem, nostack),
7913            );
7914            i += 64;
7915        }
7916        while i + 16 <= n {
7917            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7918            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
7919                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
7920            i += 16;
7921        }
7922        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
7923        while i < n {
7924            s += (*wp.add(i)) as i32 * xq[i] as i32;
7925            i += 1;
7926        }
7927        s
7928    }
7929}
7930
7931/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
7932/// loaded once and reused, 4 independent accumulators hide sdot latency
7933/// (port of vmfcore `dot_i8_sdot_4rows`).
7934#[cfg(target_arch = "aarch64")]
7935#[target_feature(enable = "neon,dotprod")]
7936unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
7937    // SAFETY: callers uphold slice-length contracts (see call sites).
7938    unsafe {
7939        use core::arch::aarch64::*;
7940        use core::arch::asm;
7941        let n = xq.len();
7942        let px = xq.as_ptr();
7943        let (p0, p1, p2, p3) = (
7944            w0.as_ptr() as *const i8,
7945            w1.as_ptr() as *const i8,
7946            w2.as_ptr() as *const i8,
7947            w3.as_ptr() as *const i8,
7948        );
7949        let (mut a0, mut a1, mut a2, mut a3) = (
7950            vdupq_n_s32(0),
7951            vdupq_n_s32(0),
7952            vdupq_n_s32(0),
7953            vdupq_n_s32(0),
7954        );
7955        let mut i = 0;
7956        while i + 16 <= n {
7957            let x = vld1q_s8(px.add(i));
7958            let v0 = vld1q_s8(p0.add(i));
7959            let v1 = vld1q_s8(p1.add(i));
7960            let v2 = vld1q_s8(p2.add(i));
7961            let v3 = vld1q_s8(p3.add(i));
7962            asm!(
7963                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7964                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7965                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7966                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7967                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7968                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7969                options(pure, nomem, nostack),
7970            );
7971            i += 16;
7972        }
7973        let mut r = [
7974            vaddvq_s32(a0),
7975            vaddvq_s32(a1),
7976            vaddvq_s32(a2),
7977            vaddvq_s32(a3),
7978        ];
7979        while i < n {
7980            let xi = *px.add(i) as i32;
7981            r[0] += (*p0.add(i)) as i32 * xi;
7982            r[1] += (*p1.add(i)) as i32 * xi;
7983            r[2] += (*p2.add(i)) as i32 * xi;
7984            r[3] += (*p3.add(i)) as i32 * xi;
7985            i += 1;
7986        }
7987        r
7988    }
7989}
7990
7991/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
7992/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
7993/// line plus the shared activation chunk — a single sequential weight
7994/// stream per worker. Per-row accumulation is the same one-accumulator
7995/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
7996/// are bit-identical to the mmap-layout kernel.
7997#[cfg(target_arch = "aarch64")]
7998#[target_feature(enable = "neon,dotprod")]
7999unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
8000    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
8001    // n % 16 == 0 — guaranteed by the repack gate).
8002    unsafe {
8003        use core::arch::aarch64::*;
8004        use core::arch::asm;
8005        let n = xq.len();
8006        let px = xq.as_ptr();
8007        let pg = g.as_ptr() as *const i8;
8008        let (mut a0, mut a1, mut a2, mut a3) = (
8009            vdupq_n_s32(0),
8010            vdupq_n_s32(0),
8011            vdupq_n_s32(0),
8012            vdupq_n_s32(0),
8013        );
8014        let mut i = 0;
8015        while i + 16 <= n {
8016            let x = vld1q_s8(px.add(i));
8017            let base = pg.add(4 * i);
8018            let v0 = vld1q_s8(base);
8019            let v1 = vld1q_s8(base.add(16));
8020            let v2 = vld1q_s8(base.add(32));
8021            let v3 = vld1q_s8(base.add(48));
8022            asm!(
8023                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8024                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8025                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8026                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8027                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8028                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8029                options(pure, nomem, nostack),
8030            );
8031            i += 16;
8032        }
8033        [
8034            vaddvq_s32(a0),
8035            vaddvq_s32(a1),
8036            vaddvq_s32(a2),
8037            vaddvq_s32(a3),
8038        ]
8039    }
8040}
8041
8042/// One q8 row range via SDOT (4-row blocks + tail) — the body of
8043/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
8044/// SAME kernel for several tensors under one pool dispatch. `rep` — the
8045/// load-time interleaved repack (empty = mmap layout only); rows outside
8046/// full 4-row groups always come from the mmap layout.
8047#[cfg(target_arch = "aarch64")]
8048fn q8_range_sdot(
8049    q: &[u8],
8050    rep: &[u8],
8051    row_scale: &[f32],
8052    act: &SplitAct,
8053    cols: usize,
8054    out_addr: SendMut,
8055    start: usize,
8056    end: usize,
8057) {
8058    let mut o = start;
8059    // Leading rows to the group boundary (repack path only): the pool
8060    // splits row ranges arbitrarily, groups are absolute.
8061    if !rep.is_empty() {
8062        while o < end && o % 4 != 0 {
8063            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8064            unsafe { *out_addr.at(o) = v };
8065            o += 1;
8066        }
8067    }
8068    while o + 4 <= end {
8069        let r = if rep.is_empty() {
8070            unsafe {
8071                dot_i8_sdot_4rows(
8072                    &q[o * cols..(o + 1) * cols],
8073                    &q[(o + 1) * cols..(o + 2) * cols],
8074                    &q[(o + 2) * cols..(o + 3) * cols],
8075                    &q[(o + 3) * cols..(o + 4) * cols],
8076                    &act.xq,
8077                )
8078            }
8079        } else {
8080            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
8081        };
8082        for k in 0..4 {
8083            let mut acc = r[k] as f32 * act.sx;
8084            for &(j, xv) in &act.outliers {
8085                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
8086            }
8087            // SAFETY: disjoint row ranges per worker.
8088            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
8089        }
8090        o += 4;
8091    }
8092    while o < end {
8093        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8094        unsafe { *out_addr.at(o) = v };
8095        o += 1;
8096    }
8097}
8098
8099/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
8100/// for the fused pair multi-matrix job (`matvec2_many`).
8101#[cfg(target_arch = "aarch64")]
8102#[allow(clippy::too_many_arguments)]
8103fn q8_range2_sdot(
8104    q: &[u8],
8105    row_scale: &[f32],
8106    a1: &SplitAct,
8107    a2: &SplitAct,
8108    cols: usize,
8109    p1: SendMut,
8110    p2: SendMut,
8111    start: usize,
8112    end: usize,
8113) {
8114    for o in start..end {
8115        let row = &q[o * cols..(o + 1) * cols];
8116        // SAFETY: disjoint row ranges per worker.
8117        unsafe {
8118            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
8119            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
8120        }
8121    }
8122}
8123
8124/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
8125#[allow(clippy::too_many_arguments)]
8126fn q8_range2_f32(
8127    q: &[u8],
8128    row_scale: &[f32],
8129    x1: &[f32],
8130    x2: &[f32],
8131    cols: usize,
8132    p1: SendMut,
8133    p2: SendMut,
8134    start: usize,
8135    end: usize,
8136) {
8137    for o in start..end {
8138        let row = &q[o * cols..(o + 1) * cols];
8139        // SAFETY: disjoint row ranges per worker.
8140        unsafe {
8141            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
8142            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
8143        }
8144    }
8145}
8146
8147/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
8148fn q8_range_f32(
8149    q: &[u8],
8150    row_scale: &[f32],
8151    xs: &[f32],
8152    cols: usize,
8153    out_addr: SendMut,
8154    start: usize,
8155    end: usize,
8156) {
8157    for o in start..end {
8158        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8159        // SAFETY: disjoint row ranges per worker.
8160        unsafe { *out_addr.at(o) = v };
8161    }
8162}
8163
8164/// SDOT row dot with exact outlier correction:
8165/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
8166#[cfg(target_arch = "aarch64")]
8167#[inline]
8168fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
8169    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
8170    for &(j, xv) in &act.outliers {
8171        acc += (row[j] as i8) as f32 * xv;
8172    }
8173    acc
8174}
8175
8176/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
8177/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
8178/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
8179/// the caller multiplies by the activation scale and adds the exact
8180/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
8181/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
8182/// → zip(lo,hi) restores flat order.
8183#[cfg(target_arch = "aarch64")]
8184#[target_feature(enable = "neon,dotprod")]
8185unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8186    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8187    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8188    unsafe {
8189        use core::arch::aarch64::*;
8190        use core::arch::asm;
8191        let lomask = vdupq_n_u8(0x0F);
8192        let eight = vdupq_n_s8(8);
8193        let mut acc = 0f32;
8194        for gi in 0..gpr {
8195            let g = g0 + gi;
8196            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8197            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8198            let lo = vandq_u8(b, lomask);
8199            let hi = vshrq_n_u8::<4>(b);
8200            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8201            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8202            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
8203            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
8204            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8205            asm!(
8206                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8207                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8208                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8209                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8210                options(pure, nomem, nostack),
8211            );
8212            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8213        }
8214        acc
8215    }
8216}
8217
8218/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
8219/// part) happens ONCE per group; both pre-quantized activations are
8220/// dotted against the same centered i8 registers. Per-lane math matches
8221/// `dot_q4_row_sdot` exactly.
8222#[cfg(target_arch = "aarch64")]
8223#[target_feature(enable = "neon,dotprod")]
8224unsafe fn dot_q4_row_sdot2(
8225    packed: &[u8],
8226    scales: &[u8],
8227    g0: usize,
8228    gpr: usize,
8229    xq1: &[i8],
8230    xq2: &[i8],
8231) -> (f32, f32) {
8232    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8233    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
8234    unsafe {
8235        use core::arch::aarch64::*;
8236        use core::arch::asm;
8237        let lomask = vdupq_n_u8(0x0F);
8238        let eight = vdupq_n_s8(8);
8239        let (mut acc1, mut acc2) = (0f32, 0f32);
8240        for gi in 0..gpr {
8241            let g = g0 + gi;
8242            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8243            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8244            let lo = vandq_u8(b, lomask);
8245            let hi = vshrq_n_u8::<4>(b);
8246            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8247            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8248            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
8249            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
8250            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
8251            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
8252            let (mut a0, mut a1, mut b0, mut b1) = (
8253                vdupq_n_s32(0),
8254                vdupq_n_s32(0),
8255                vdupq_n_s32(0),
8256                vdupq_n_s32(0),
8257            );
8258            asm!(
8259                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
8260                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
8261                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
8262                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
8263                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8264                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
8265                e0 = in(vreg) e0, e1 = in(vreg) e1,
8266                x10 = in(vreg) x10, x11 = in(vreg) x11,
8267                x20 = in(vreg) x20, x21 = in(vreg) x21,
8268                options(pure, nomem, nostack),
8269            );
8270            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8271            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
8272        }
8273        (acc1, acc2)
8274    }
8275}
8276
8277// ───────────────────── fused int8 kernels ─────────────────────
8278
8279/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
8280/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
8281#[inline]
8282pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
8283    #[cfg(target_arch = "aarch64")]
8284    unsafe {
8285        return axpy_i8_f32_neon(acc, row, w);
8286    }
8287    #[cfg(target_arch = "x86_64")]
8288    if avx2_enabled() {
8289        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
8290    }
8291    #[allow(unreachable_code)]
8292    {
8293        for (a, &b) in acc.iter_mut().zip(row) {
8294            *a += w * b as f32;
8295        }
8296    }
8297}
8298
8299/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
8300#[cfg(target_arch = "x86_64")]
8301#[target_feature(enable = "avx2,fma")]
8302unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
8303    // SAFETY: callers uphold slice-length contracts (see call sites).
8304    unsafe {
8305        use core::arch::x86_64::*;
8306        let n = acc.len().min(row.len());
8307        let ap = acc.as_mut_ptr();
8308        let rp = row.as_ptr();
8309        let wv = _mm256_set1_ps(w);
8310        let mut j = 0usize;
8311        while j + 16 <= n {
8312            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
8313            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
8314            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
8315            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
8316            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
8317            _mm256_storeu_ps(ap.add(j), v0);
8318            _mm256_storeu_ps(ap.add(j + 8), v1);
8319            j += 16;
8320        }
8321        while j < n {
8322            *ap.add(j) += w * (*rp.add(j)) as f32;
8323            j += 1;
8324        }
8325    }
8326}
8327
8328#[cfg(target_arch = "aarch64")]
8329#[target_feature(enable = "neon")]
8330unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
8331    // SAFETY: callers uphold slice-length contracts (see call sites).
8332    unsafe {
8333        use core::arch::aarch64::*;
8334        let n = acc.len().min(row.len());
8335        let ap = acc.as_mut_ptr();
8336        let rp = row.as_ptr();
8337        let wv = vdupq_n_f32(w);
8338        let mut j = 0usize;
8339        while j + 16 <= n {
8340            let rb = vld1q_s8(rp.add(j));
8341            let lo = vmovl_s8(vget_low_s8(rb));
8342            let hi = vmovl_s8(vget_high_s8(rb));
8343            for (off, half) in [(0, lo), (8, hi)] {
8344                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
8345                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
8346                let o = j + off;
8347                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
8348                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
8349            }
8350            j += 16;
8351        }
8352        while j < n {
8353            *ap.add(j) += w * (*rp.add(j)) as f32;
8354            j += 1;
8355        }
8356    }
8357}
8358
8359/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
8360/// ≈9× scalar), scalar elsewhere.
8361#[inline]
8362pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
8363    #[cfg(target_arch = "aarch64")]
8364    unsafe {
8365        return dot_i8_f32_neon(w, x);
8366    }
8367    #[cfg(target_arch = "x86_64")]
8368    if avx2_enabled() {
8369        return unsafe { dot_i8_f32_avx2(w, x) };
8370    }
8371    #[allow(unreachable_code)]
8372    {
8373        let mut sum = 0.0f32;
8374        for (j, &b) in w.iter().enumerate() {
8375            sum += (b as i8) as f32 * x[j];
8376        }
8377        sum
8378    }
8379}
8380
8381/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
8382/// folded into the product (no prescaled copy of x). NEON on aarch64,
8383/// scalar elsewhere. Used by the active-neuron path `row_dot`.
8384#[inline]
8385fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8386    #[cfg(target_arch = "aarch64")]
8387    unsafe {
8388        return dot_i8_col_f32_neon(w, x, col);
8389    }
8390    #[allow(unreachable_code)]
8391    {
8392        let mut sum = 0.0f32;
8393        for (j, &b) in w.iter().enumerate() {
8394            sum += (b as i8) as f32 * x[j] * col[j];
8395        }
8396        sum
8397    }
8398}
8399
8400#[cfg(target_arch = "aarch64")]
8401#[target_feature(enable = "neon")]
8402unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8403    // SAFETY: callers uphold slice-length contracts (see call sites).
8404    unsafe {
8405        use core::arch::aarch64::*;
8406        let n = x.len();
8407        let wp = w.as_ptr() as *const i8;
8408        let xp = x.as_ptr();
8409        let cp = col.as_ptr();
8410        let (mut a0, mut a1, mut a2, mut a3) = (
8411            vdupq_n_f32(0.0),
8412            vdupq_n_f32(0.0),
8413            vdupq_n_f32(0.0),
8414            vdupq_n_f32(0.0),
8415        );
8416        let mut j = 0usize;
8417        while j + 16 <= n {
8418            let wb = vld1q_s8(wp.add(j));
8419            let lo = vmovl_s8(vget_low_s8(wb));
8420            let hi = vmovl_s8(vget_high_s8(wb));
8421            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8422            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8423            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8424            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8425            a0 = vfmaq_f32(
8426                a0,
8427                w0,
8428                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
8429            );
8430            a1 = vfmaq_f32(
8431                a1,
8432                w1,
8433                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
8434            );
8435            a2 = vfmaq_f32(
8436                a2,
8437                w2,
8438                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
8439            );
8440            a3 = vfmaq_f32(
8441                a3,
8442                w3,
8443                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
8444            );
8445            j += 16;
8446        }
8447        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8448        while j < n {
8449            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
8450            j += 1;
8451        }
8452        sum
8453    }
8454}
8455
8456#[cfg(target_arch = "aarch64")]
8457#[target_feature(enable = "neon")]
8458unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
8459    // SAFETY: callers uphold slice-length contracts (see call sites).
8460    unsafe {
8461        use core::arch::aarch64::*;
8462        let n = x.len();
8463        let wp = w.as_ptr() as *const i8;
8464        let xp = x.as_ptr();
8465        let (mut a0, mut a1, mut a2, mut a3) = (
8466            vdupq_n_f32(0.0),
8467            vdupq_n_f32(0.0),
8468            vdupq_n_f32(0.0),
8469            vdupq_n_f32(0.0),
8470        );
8471        let mut j = 0usize;
8472        while j + 16 <= n {
8473            let wb = vld1q_s8(wp.add(j));
8474            let lo = vmovl_s8(vget_low_s8(wb));
8475            let hi = vmovl_s8(vget_high_s8(wb));
8476            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8477            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8478            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8479            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8480            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
8481            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
8482            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
8483            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
8484            j += 16;
8485        }
8486        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8487        while j < n {
8488            sum += (*wp.add(j)) as f32 * *xp.add(j);
8489            j += 1;
8490        }
8491        sum
8492    }
8493}
8494
8495#[allow(clippy::too_many_arguments)]
8496fn qmatvec(
8497    q: &[u8],
8498    rep: &[u8],
8499    row_scale: &[f32],
8500    x: &[f32],
8501    col_field: &[f32],
8502    dtype: TensorDtype,
8503    rows: usize,
8504    cols: usize,
8505    out: &mut [f32],
8506    pool: Option<&Pool>,
8507) {
8508    debug_assert_eq!(out.len(), rows);
8509    #[cfg(not(target_arch = "aarch64"))]
8510    let _ = rep;
8511
8512    #[cfg(target_arch = "aarch64")]
8513    if sdot_enabled() {
8514        let act = if dtype == TensorDtype::Q8_2f {
8515            split_act_q8_2f(x, col_field)
8516        } else {
8517            split_act(x)
8518        };
8519        let out_addr = SendMut(out.as_mut_ptr());
8520        let run_range = |start: usize, end: usize| {
8521            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
8522        };
8523        match pool {
8524            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8525            _ => run_range(0, rows),
8526        }
8527        return;
8528    }
8529    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
8530    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
8531    #[cfg(target_arch = "x86_64")]
8532    if avx2_a8w8_enabled() {
8533        let act = if dtype == TensorDtype::Q8_2f {
8534            split_act_q8_2f(x, col_field)
8535        } else {
8536            split_act(x)
8537        };
8538        let out_addr = SendMut(out.as_mut_ptr());
8539        let run_range = |start: usize, end: usize| {
8540            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
8541        };
8542        match pool {
8543            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8544            _ => run_range(0, rows),
8545        }
8546        return;
8547    }
8548
8549    prescale_with(x, col_field, dtype, 1, |xs| {
8550        let out_addr = SendMut(out.as_mut_ptr());
8551        let run_range = move |start: usize, end: usize| {
8552            for o in start..end {
8553                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8554                // SAFETY: disjoint row ranges per worker.
8555                unsafe { *out_addr.at(o) = v };
8556            }
8557        };
8558        match pool {
8559            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8560            _ => run_range(0, rows),
8561        }
8562    });
8563}
8564
8565#[allow(clippy::too_many_arguments)]
8566fn qmatvec2(
8567    q: &[u8],
8568    row_scale: &[f32],
8569    x1: &[f32],
8570    x2: &[f32],
8571    col_field: &[f32],
8572    dtype: TensorDtype,
8573    rows: usize,
8574    cols: usize,
8575    o1: &mut [f32],
8576    o2: &mut [f32],
8577    pool: Option<&Pool>,
8578) {
8579    #[cfg(target_arch = "aarch64")]
8580    if sdot_enabled() {
8581        let a1s = if dtype == TensorDtype::Q8_2f {
8582            split_act_q8_2f(x1, col_field)
8583        } else {
8584            split_act(x1)
8585        };
8586        let a2s = if dtype == TensorDtype::Q8_2f {
8587            split_act_q8_2f(x2, col_field)
8588        } else {
8589            split_act(x2)
8590        };
8591        let p1 = SendMut(o1.as_mut_ptr());
8592        let p2 = SendMut(o2.as_mut_ptr());
8593        let run_range = |start: usize, end: usize| {
8594            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8595        };
8596        match pool {
8597            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8598            _ => run_range(0, rows),
8599        }
8600        return;
8601    }
8602    #[cfg(target_arch = "x86_64")]
8603    if avx2_a8w8_enabled() {
8604        let a1s = if dtype == TensorDtype::Q8_2f {
8605            split_act_q8_2f(x1, col_field)
8606        } else {
8607            split_act(x1)
8608        };
8609        let a2s = if dtype == TensorDtype::Q8_2f {
8610            split_act_q8_2f(x2, col_field)
8611        } else {
8612            split_act(x2)
8613        };
8614        let p1 = SendMut(o1.as_mut_ptr());
8615        let p2 = SendMut(o2.as_mut_ptr());
8616        let run_range = |start: usize, end: usize| {
8617            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8618        };
8619        match pool {
8620            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8621            _ => run_range(0, rows),
8622        }
8623        return;
8624    }
8625
8626    prescale_with(x1, col_field, dtype, 1, |x1s| {
8627        prescale_with(x2, col_field, dtype, 2, |x2s| {
8628            let p1 = SendMut(o1.as_mut_ptr());
8629            let p2 = SendMut(o2.as_mut_ptr());
8630            let run_range = move |start: usize, end: usize| {
8631                for o in start..end {
8632                    let row = &q[o * cols..(o + 1) * cols];
8633                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
8634                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
8635                    // SAFETY: disjoint row ranges per worker.
8636                    unsafe {
8637                        *p1.at(o) = s1;
8638                        *p2.at(o) = s2;
8639                    }
8640                }
8641            };
8642            match pool {
8643                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8644                _ => run_range(0, rows),
8645            }
8646        });
8647    });
8648}
8649
8650#[derive(Clone, Copy)]
8651struct SendMut(*mut f32);
8652unsafe impl Send for SendMut {}
8653unsafe impl Sync for SendMut {}
8654
8655impl SendMut {
8656    #[inline]
8657    fn at(self, i: usize) -> *mut f32 {
8658        unsafe { self.0.add(i) }
8659    }
8660}
8661
8662#[cfg(test)]
8663mod tests {
8664    use super::*;
8665
8666    #[test]
8667    fn f32_matvec_matches_matvec_rows_bitexact() {
8668        let (rows, cols) = (300, 40);
8669        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
8670        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
8671        let qt = QTensor::from_f32(w.clone(), rows, cols);
8672
8673        let mut a = vec![0.0f32; rows];
8674        matvec_rows(None, &w, &x, &mut a);
8675        let mut b = vec![0.0f32; rows];
8676        qt.matvec(&x, &mut b, None);
8677        assert_eq!(a, b);
8678    }
8679
8680    #[test]
8681    fn sdot_kernel_exact_on_grid() {
8682        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
8683        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
8684        // exact f32 dot to float rounding. This isolates kernel
8685        // correctness from quantization noise.
8686        eprintln!("sdot_enabled = {}", sdot_enabled());
8687        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
8688        let w: Vec<u8> = (0..rows * cols)
8689            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8690            .collect();
8691        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
8692        let x: Vec<f32> = (0..cols)
8693            .map(|i| match i % 3 {
8694                0 => 1.0,
8695                1 => -1.0,
8696                _ => 0.0,
8697            })
8698            .collect();
8699        let mut a = vec![0.0f32; rows];
8700        qmatvec(
8701            &w,
8702            &[],
8703            &scales,
8704            &x,
8705            &[],
8706            TensorDtype::Q8Row,
8707            rows,
8708            cols,
8709            &mut a,
8710            None,
8711        );
8712        for o in 0..rows {
8713            let mut acc = 0.0f32;
8714            for j in 0..cols {
8715                acc += (w[o * cols + j] as i8) as f32 * x[j];
8716            }
8717            let expect = acc * scales[o];
8718            assert!(
8719                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8720                "row {o}: {} vs {expect}",
8721                a[o]
8722            );
8723        }
8724    }
8725
8726    #[test]
8727    fn q1_tbl_fast_path_matches_reference() {
8728        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
8729        // row's final 4-tile window trips the 4B-overread guard (the
8730        // payload ends exactly at the last tile) — both paths must
8731        // agree with the dequant reference.
8732        let (rows, cols) = (5, 256);
8733        let gpr = cols / GROUP_SIZE;
8734        let mut bytes = Vec::new();
8735        for t in 0..rows * gpr {
8736            let s = 0.007 + (t % 11) as f32 * 0.004;
8737            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8738            for j in 0..4 {
8739                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
8740            }
8741        }
8742        let x: Vec<f32> = (0..cols)
8743            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
8744            .collect();
8745        let mut w = vec![0.0f32; rows * cols];
8746        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8747        let mut got = vec![0.0f32; rows];
8748        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8749        for o in 0..rows {
8750            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8751            assert!(
8752                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8753                "row {o}: {} vs {expect}",
8754                got[o]
8755            );
8756        }
8757        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
8758        // single-matvec path bit-for-bit.
8759        let b = 5usize;
8760        let mut xs_all = Vec::new();
8761        for bi in 0..b {
8762            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
8763        }
8764        let mut mm = vec![0.0f32; b * rows];
8765        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
8766        for bi in 0..b {
8767            let mut single = vec![0.0f32; rows];
8768            q1_matvec(
8769                &bytes,
8770                &xs_all[bi * cols..(bi + 1) * cols],
8771                rows,
8772                cols,
8773                &mut single,
8774                None,
8775            );
8776            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
8777        }
8778    }
8779
8780    #[test]
8781    fn q1_kernels_match_exact_reference() {
8782        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
8783        let (rows, cols) = (7, 96);
8784        let gpr = cols / GROUP_SIZE;
8785        let mut bytes = Vec::new();
8786        for t in 0..rows * gpr {
8787            let s = 0.01 + (t % 13) as f32 * 0.003;
8788            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8789            for j in 0..4 {
8790                bytes.push(((t * 31 + j * 97) % 251) as u8);
8791            }
8792        }
8793        // On-grid activations (±1, amax 1) → the SDOT path is exact.
8794        let x: Vec<f32> = (0..cols)
8795            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
8796            .collect();
8797        // Reference through the core dequant.
8798        let mut w = vec![0.0f32; rows * cols];
8799        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8800        let mut expect = vec![0.0f32; rows];
8801        for o in 0..rows {
8802            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8803        }
8804        let mut got = vec![0.0f32; rows];
8805        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8806        for o in 0..rows {
8807            assert!(
8808                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
8809                "row {o}: {} vs {}",
8810                got[o],
8811                expect[o]
8812            );
8813        }
8814        // Pair and batch paths agree with the single path.
8815        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
8816        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
8817        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
8818        assert_eq!(a1, got);
8819        let mut xs = x.clone();
8820        xs.extend_from_slice(&x2);
8821        let mut mm = vec![0.0f32; 2 * rows];
8822        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
8823        assert_eq!(&mm[..rows], got.as_slice());
8824        assert_eq!(&mm[rows..], a2.as_slice());
8825    }
8826
8827    #[test]
8828    fn repack_is_bit_identical() {
8829        // The interleaved-repack kernel must produce EXACTLY the same
8830        // bits as the mmap-layout kernel: integer accumulation is order-
8831        // exact, the f32 epilogue is identical. Odd rows exercise the
8832        // tail; direct range calls exercise unaligned pool splits.
8833        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
8834        let w: Vec<u8> = (0..rows * cols)
8835            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
8836            .collect();
8837        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
8838        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
8839        let rep = q8_repack_layout(&w, rows, cols);
8840        // Group interleave round-trips.
8841        for g in 0..rows / 4 {
8842            for c in 0..cols / 16 {
8843                for lane in 0..4 {
8844                    assert_eq!(
8845                        &rep[g * 4 * cols + c * 64 + lane * 16
8846                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
8847                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
8848                    );
8849                }
8850            }
8851        }
8852        let mut a = vec![0.0f32; rows];
8853        qmatvec(
8854            &w,
8855            &[],
8856            &scales,
8857            &x,
8858            &[],
8859            TensorDtype::Q8Row,
8860            rows,
8861            cols,
8862            &mut a,
8863            None,
8864        );
8865        let mut b = vec![0.0f32; rows];
8866        qmatvec(
8867            &w,
8868            &rep,
8869            &scales,
8870            &x,
8871            &[],
8872            TensorDtype::Q8Row,
8873            rows,
8874            cols,
8875            &mut b,
8876            None,
8877        );
8878        assert_eq!(a, b, "full-range repack output diverged");
8879
8880        #[cfg(target_arch = "aarch64")]
8881        if sdot_enabled() {
8882            // Unaligned range split (pool workers get arbitrary bounds).
8883            let act = split_act(&x);
8884            let mut c1 = vec![0.0f32; rows];
8885            let mut c2 = vec![0.0f32; rows];
8886            q8_range_sdot(
8887                &w,
8888                &[],
8889                &scales,
8890                &act,
8891                cols,
8892                SendMut(c1.as_mut_ptr()),
8893                3,
8894                rows - 2,
8895            );
8896            q8_range_sdot(
8897                &w,
8898                &rep,
8899                &scales,
8900                &act,
8901                cols,
8902                SendMut(c2.as_mut_ptr()),
8903                3,
8904                rows - 2,
8905            );
8906            assert_eq!(c1, c2, "unaligned-range repack output diverged");
8907        }
8908    }
8909
8910    #[test]
8911    fn sdot_a8w8_noise_is_bounded() {
8912        // Off-grid activations: A8 quantization noise must stay small in
8913        // relative L2 over the whole output (realistic accuracy contract;
8914        // vmfcore measured argmax-identical decode on real models).
8915        let (rows, cols) = (16, 512);
8916        let w: Vec<u8> = (0..rows * cols)
8917            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8918            .collect();
8919        let scales = vec![0.01f32; rows];
8920        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
8921        let mut a = vec![0.0f32; rows];
8922        qmatvec(
8923            &w,
8924            &[],
8925            &scales,
8926            &x,
8927            &[],
8928            TensorDtype::Q8Row,
8929            rows,
8930            cols,
8931            &mut a,
8932            None,
8933        );
8934        let (mut num, mut den) = (0f64, 0f64);
8935        for o in 0..rows {
8936            let mut acc = 0.0f32;
8937            for j in 0..cols {
8938                acc += (w[o * cols + j] as i8) as f32 * x[j];
8939            }
8940            let expect = acc * scales[o];
8941            num += ((a[o] - expect) as f64).powi(2);
8942            den += (expect as f64).powi(2);
8943        }
8944        let rel = (num / den.max(1e-12)).sqrt();
8945        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
8946    }
8947
8948    #[test]
8949    fn i8_dot_neon_matches_scalar() {
8950        let n = 100;
8951        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
8952        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
8953        let mut scalar = 0.0f32;
8954        for j in 0..n {
8955            scalar += (w[j] as i8) as f32 * x[j];
8956        }
8957        let fast = dot_i8_f32(&w, &x);
8958        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
8959    }
8960
8961    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
8962    #[test]
8963    fn vbitmatvec_matches_full_dequant() {
8964        let (rows, cols) = (6, 64);
8965        let ng = cols / GROUP_SIZE;
8966        // Hand-craft: bits per row, f16 scales, packed rows.
8967        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8968        let mut bytes = bits.clone();
8969        for g in 0..rows * ng {
8970            let s = 0.02 + 0.001 * g as f32;
8971            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8972        }
8973        for r in 0..rows {
8974            let b = bits[r] as usize;
8975            let (mut acc, mut nb) = (0u64, 0usize);
8976            let mut rowbytes = Vec::new();
8977            for i in 0..cols {
8978                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8979                acc = (acc << b) | v;
8980                nb += b;
8981                while nb >= 8 {
8982                    nb -= 8;
8983                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
8984                }
8985            }
8986            if nb > 0 {
8987                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8988            }
8989            bytes.extend_from_slice(&rowbytes);
8990        }
8991        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8992
8993        let mut reference = vec![0f32; rows * cols];
8994        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
8995        let mut expect = vec![0f32; rows];
8996        for r in 0..rows {
8997            expect[r] = reference[r * cols..(r + 1) * cols]
8998                .iter()
8999                .zip(&x)
9000                .map(|(w, xv)| w * xv)
9001                .sum();
9002        }
9003        let mut got = vec![0f32; rows];
9004        let offsets = vbit_row_offsets(&bytes, rows, cols);
9005        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
9006        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9007        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
9008        // the golden-parity gate).
9009        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9010        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
9011        for r in 0..rows {
9012            assert!(
9013                (got[r] - expect[r]).abs() < tol * scale,
9014                "row {r}: {} vs {}",
9015                got[r],
9016                expect[r]
9017            );
9018        }
9019    }
9020
9021    /// Fused q4 matvec must match the reference full-dequant + dense
9022    /// matvec bit-for-bit in structure (same f32 math, group order).
9023    /// vbit matmat: the blocked 1×4 leg must match the per-row path
9024    /// (paired env toggle; larger shape so both code paths engage).
9025    #[test]
9026    #[cfg(target_arch = "x86_64")]
9027    fn vbit_matmat_blocked_matches_per_row() {
9028        let (rows, cols, b) = (64usize, 128usize, 9usize);
9029        let ng = cols / GROUP_SIZE;
9030        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
9031        let mut bytes = bits.clone();
9032        for g in 0..rows * ng {
9033            let sc = 0.02 + 0.0005 * g as f32;
9034            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9035        }
9036        for r in 0..rows {
9037            let bw = bits[r] as usize;
9038            let (mut acc, mut nb) = (0u64, 0usize);
9039            let mut rowbytes = Vec::new();
9040            for i in 0..cols {
9041                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9042                acc = (acc << bw) | v;
9043                nb += bw;
9044                while nb >= 8 {
9045                    nb -= 8;
9046                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9047                }
9048            }
9049            if nb > 0 {
9050                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9051            }
9052            bytes.extend_from_slice(&rowbytes);
9053        }
9054        let x: Vec<f32> = (0..b * cols)
9055            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9056            .collect();
9057        let offsets = vbit_row_offsets(&bytes, rows, cols);
9058        let mut y_a = vec![0f32; b * rows];
9059        let mut y_b = vec![0f32; b * rows];
9060        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9061        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
9062        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9063        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
9064        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9065        let max_d = y_a
9066            .iter()
9067            .zip(&y_b)
9068            .map(|(p, q)| (p - q).abs())
9069            .fold(0.0f32, f32::max);
9070        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
9071    }
9072
9073    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
9074    /// per-row path exactly: same nibble unpack, same group order,
9075    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
9076    /// two full 1×4 blocks plus a remainder through the single-row
9077    /// kernel. (Both paths produce identical output, so the shared
9078    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
9079    /// the verdict — worst case both sides take the same path.)
9080    #[test]
9081    fn q4t_matmat_blocked_matches_per_row() {
9082        let (rows, cols, b) = (16usize, 64usize, 9usize);
9083        let gpr = cols / GROUP_SIZE;
9084        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9085        for r in 0..rows {
9086            for g in 0..gpr {
9087                let t = (r * gpr + g) * Q4_TILE;
9088                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
9089                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9090                for k in 0..16 {
9091                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9092                }
9093            }
9094        }
9095        let x: Vec<f32> = (0..b * cols)
9096            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9097            .collect();
9098        let mut y_blk = vec![0f32; b * rows];
9099        let mut y_row = vec![0f32; b * rows];
9100        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9101        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
9102        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9103        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
9104        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9105        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
9106    }
9107
9108    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
9109    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
9110    /// order differs — tight tolerance.
9111    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
9112    /// span varies row to row, so the codes actually exercise the full 0..31
9113    /// range rather than clustering on one rung.
9114    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
9115        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
9116        let gpr = cols / GROUP_SIZE;
9117        let stride = q4tp_code_stride(gpr);
9118        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
9119        let mut b = vec![0u8; codes_off + rows * stride];
9120        for r in 0..rows {
9121            for g in 0..gpr {
9122                let t = (r * gpr + g) * Q4TP_NIB;
9123                for k in 0..16 {
9124                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9125                }
9126            }
9127            let lo = -6.0 - 0.03 * (r % 17) as f32;
9128            let step = 0.01 + 0.004 * (r % 11) as f32;
9129            let p = params_off + r * 4;
9130            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
9131            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
9132            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
9133            for g in 0..gpr {
9134                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
9135            }
9136        }
9137        b
9138    }
9139
9140    /// The same weights re-expressed as q4_tiled, so the proven kernel can
9141    /// be the reference: each tile stores the ladder scale its code selects.
9142    /// Only the f16 rounding of that scale separates the two payloads.
9143    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
9144        let gpr = cols / GROUP_SIZE;
9145        let v = Q4tpView::new(bytes, rows, cols);
9146        let mut out = vec![0u8; rows * gpr * Q4_TILE];
9147        let mut sc = vec![0f32; gpr];
9148        for r in 0..rows {
9149            v.scales_into(r, gpr, &mut sc);
9150            for g in 0..gpr {
9151                let t = (r * gpr + g) * Q4_TILE;
9152                let s = sc[g];
9153                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9154                let src = (r * gpr + g) * Q4TP_NIB;
9155                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
9156            }
9157        }
9158        out
9159    }
9160
9161    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
9162    /// rounding — that scalar routine is the format's definition, and the
9163    /// kernels re-derive the scale from the ladder independently. Call the
9164    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
9165    /// so routing through it would test the other path by accident.
9166    #[test]
9167    fn q4tp_exact_path_matches_dequant_reference() {
9168        let (rows, cols) = (256usize, 512usize);
9169        let gpr = cols / GROUP_SIZE;
9170        let bytes = synth_q4tp(rows, cols);
9171        let mut w = vec![0f32; rows * cols];
9172        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9173
9174        let x: Vec<f32> = (0..cols)
9175            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9176            .collect();
9177        let v = Q4tpView::new(&bytes, rows, cols);
9178        let mut sc = vec![0f32; gpr];
9179        for r in 0..rows {
9180            v.scales_into(r, gpr, &mut sc);
9181            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
9182            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
9183            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
9184            // the meaningful yardstick is the summed magnitude, not the result:
9185            // against the result any reordering of a 512-term f32 sum "fails".
9186            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9187            assert!(
9188                (got - want).abs() <= 1e-5 * mag,
9189                "row {r}: kernel {got} vs dequant {want}"
9190            );
9191        }
9192    }
9193
9194    /// The int8 (a8w8) path can't be checked against an f32 reference — the
9195    /// activation quantization dominates. Check it against the q4t kernel it
9196    /// was ported from instead, on payloads holding the same weights: that
9197    /// isolates exactly what the port could break (16 B stride, ladder
9198    /// lookup, nibble unpack) from what it deliberately shares.
9199    #[test]
9200    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
9201        let (rows, cols) = (256usize, 512usize);
9202        let bytes = synth_q4tp(rows, cols);
9203        let twin = q4tp_as_q4t(&bytes, rows, cols);
9204        let x: Vec<f32> = (0..cols)
9205            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9206            .collect();
9207
9208        let mut got = vec![0f32; rows];
9209        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
9210        let mut want = vec![0f32; rows];
9211        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
9212
9213        // Scale is f16 in the twin and f32 here, so allow that rounding on
9214        // top of the summed magnitude (same cancellation argument as above).
9215        let mut w = vec![0f32; rows * cols];
9216        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9217        for r in 0..rows {
9218            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9219            assert!(
9220                (got[r] - want[r]).abs() <= 1e-3 * mag,
9221                "row {r}: q4tp {} vs q4t {}",
9222                got[r],
9223                want[r]
9224            );
9225        }
9226    }
9227
9228    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
9229    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
9230    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
9231    /// code and its four accumulators are exactly what tends to go wrong.
9232    #[test]
9233    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
9234        let (rows, cols, b) = (256usize, 512usize, 5usize);
9235        let bytes = synth_q4tp(rows, cols);
9236        let twin = q4tp_as_q4t(&bytes, rows, cols);
9237        let xs: Vec<f32> = (0..b * cols)
9238            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9239            .collect();
9240
9241        let mut got = vec![0f32; b * rows];
9242        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
9243        let mut want = vec![0f32; b * rows];
9244        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
9245
9246        let mut w = vec![0f32; rows * cols];
9247        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9248        for t in 0..b {
9249            for r in 0..rows {
9250                let mag: f32 = (0..cols)
9251                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
9252                    .sum();
9253                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
9254                assert!(
9255                    (g - wa).abs() <= 1e-3 * mag,
9256                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
9257                );
9258            }
9259        }
9260    }
9261
9262    #[test]
9263    fn q4tp_matvec2_matches_the_single_stream_kernel() {
9264        let (rows, cols) = (128usize, 256usize);
9265        let gpr = cols / GROUP_SIZE;
9266        let bytes = synth_q4tp(rows, cols);
9267        let xs: Vec<f32> = (0..2 * cols)
9268            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9269            .collect();
9270
9271        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
9272        q4tp_matvec2(
9273            &bytes,
9274            &xs[..cols],
9275            &xs[cols..],
9276            rows,
9277            cols,
9278            &mut o1,
9279            &mut o2,
9280            None,
9281        );
9282
9283        // matvec2 takes the exact path for both streams, so the single-row
9284        // kernel is an exact reference — no tolerance for path differences.
9285        let v = Q4tpView::new(&bytes, rows, cols);
9286        let mut sc = vec![0f32; gpr];
9287        for r in 0..rows {
9288            v.scales_into(r, gpr, &mut sc);
9289            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
9290            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
9291        }
9292    }
9293
9294    /// q4tp must not COST speed — it exists to save bytes, and a format that
9295    /// trades 7% of a file for a slower model is a bad trade. This guard is
9296    /// here because correctness tests happily passed while `q4tp_matmat` was
9297    /// missing its int8 and Accelerate arms and the model ran 5x slower.
9298    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
9299    /// aligned than q4t's 18 B, which pays for the scale indirection).
9300    #[test]
9301    fn q4tp_matvec_keeps_pace_with_q4t() {
9302        let (rows, cols) = (4096usize, 3072usize);
9303        let bytes = synth_q4tp(rows, cols);
9304        let twin = q4tp_as_q4t(&bytes, rows, cols);
9305        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
9306        let mut o = vec![0f32; rows];
9307        let n = 12;
9308        let mut best = (f64::MAX, f64::MAX);
9309        // Interleaved A/B, minimum statistic: this machine throttles, and a
9310        // mean over a thermal ramp reliably indicts whichever ran second.
9311        for _ in 0..3 {
9312            let t0 = std::time::Instant::now();
9313            for _ in 0..n {
9314                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
9315            }
9316            best.0 = best.0.min(t0.elapsed().as_secs_f64());
9317            let t0 = std::time::Instant::now();
9318            for _ in 0..n {
9319                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
9320            }
9321            best.1 = best.1.min(t0.elapsed().as_secs_f64());
9322        }
9323        let ratio = best.1 / best.0;
9324        println!("q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x", best.0 * 1e3 / n as f64, best.1 * 1e3 / n as f64);
9325        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
9326    }
9327
9328    #[cfg(target_os = "macos")]
9329    #[test]
9330    fn q4t_matmat_accel_matches_dequant_reference() {
9331        if !accel_gemm_enabled() {
9332            return; // CMF_ACCEL=0
9333        }
9334        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
9335        let gpr = cols / GROUP_SIZE;
9336        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9337        for r in 0..rows {
9338            for g in 0..gpr {
9339                let t = (r * gpr + g) * Q4_TILE;
9340                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
9341                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9342                for k in 0..16 {
9343                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9344                }
9345            }
9346        }
9347        let x: Vec<f32> = (0..b * cols)
9348            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9349            .collect();
9350        let mut got = vec![0f32; b * rows];
9351        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
9352        // Brute-force reference off the same tiles.
9353        let mut w = vec![0f32; rows * cols];
9354        for r in 0..rows {
9355            for g in 0..gpr {
9356                let t = (r * gpr + g) * Q4_TILE;
9357                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
9358                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
9359                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
9360                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
9361                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
9362                }
9363            }
9364        }
9365        for bi in 0..b {
9366            for r in 0..rows {
9367                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
9368                let d = (got[bi * rows + r] - want).abs();
9369                assert!(
9370                    d <= want.abs().max(1.0) * 1e-4,
9371                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
9372                    got[bi * rows + r]
9373                );
9374            }
9375        }
9376    }
9377
9378    #[test]
9379    fn q4matvec_matches_full_dequant() {
9380        let (rows, cols) = (8, 64);
9381        let groups = rows * cols / GROUP_SIZE;
9382        // Hand-craft a q4_block blob: nibbles then f16 scales.
9383        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9384        for i in 0..groups * 16 {
9385            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9386        }
9387        for g in 0..groups {
9388            let s = 0.01 + 0.003 * g as f32;
9389            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9390        }
9391        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9392
9393        let mut reference = vec![0.0f32; rows * cols];
9394        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9395        let mut expect = vec![0.0f32; rows];
9396        for r in 0..rows {
9397            expect[r] = reference[r * cols..(r + 1) * cols]
9398                .iter()
9399                .zip(&x)
9400                .map(|(w, xv)| w * xv)
9401                .sum();
9402        }
9403
9404        let mut got = vec![0.0f32; rows];
9405        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9406        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9407        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
9408        // in the golden-parity gate).
9409        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9410        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9411        for r in 0..rows {
9412            assert!(
9413                (got[r] - expect[r]).abs() < tol * scale,
9414                "row {r}: {} vs {}",
9415                got[r],
9416                expect[r]
9417            );
9418        }
9419    }
9420
9421    /// Fused two-input vbit matvec must equal two single matvecs exactly
9422    /// (same per-lane accumulation order on both scalar and SDOT paths).
9423    #[test]
9424    fn vbitmatvec2_equals_two_singles() {
9425        let (rows, cols) = (6, 64);
9426        let ng = cols / GROUP_SIZE;
9427        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9428        let mut bytes = bits.clone();
9429        for g in 0..rows * ng {
9430            let s = 0.02 + 0.001 * g as f32;
9431            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9432        }
9433        for r in 0..rows {
9434            let b = bits[r] as usize;
9435            let (mut acc, mut nb) = (0u64, 0usize);
9436            let mut rowbytes = Vec::new();
9437            for i in 0..cols {
9438                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9439                acc = (acc << b) | v;
9440                nb += b;
9441                while nb >= 8 {
9442                    nb -= 8;
9443                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9444                }
9445            }
9446            if nb > 0 {
9447                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9448            }
9449            bytes.extend_from_slice(&rowbytes);
9450        }
9451        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9452        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
9453        let offsets = vbit_row_offsets(&bytes, rows, cols);
9454
9455        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9456        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
9457        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
9458        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9459        vbitmatvec2(
9460            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
9461        );
9462        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
9463        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
9464    }
9465
9466    /// Fused two-input q4 matvec must equal two single matvecs exactly.
9467    #[test]
9468    fn q4matvec2_equals_two_singles() {
9469        let (rows, cols) = (8, 128);
9470        let groups = rows * cols / GROUP_SIZE;
9471        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9472        for i in 0..groups * 16 {
9473            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9474        }
9475        for g in 0..groups {
9476            let s = 0.01 + 0.003 * g as f32;
9477            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9478        }
9479        // Include an outlier channel so the SDOT correction path is
9480        // exercised in the pair kernel too.
9481        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9482        x1[9] = 250.0;
9483        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9484
9485        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9486        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
9487        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
9488        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9489        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
9490        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
9491        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
9492    }
9493
9494    /// Multi-matrix job must equal separate matvecs exactly — same
9495    /// kernels, only the dispatch is fused.
9496    #[test]
9497    fn matvec_many_equals_separate_matvecs() {
9498        use crate::pool::Pool;
9499        let (r1, r2, cols) = (300, 200, 64);
9500        let mk = |salt: usize, rows: usize| {
9501            QTensor::from_f32(
9502                (0..rows * cols)
9503                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
9504                    .collect(),
9505                rows,
9506                cols,
9507            )
9508        };
9509        let (a, b) = (mk(1, r1), mk(5, r2));
9510        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
9511        let pool = Pool::new(3);
9512
9513        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
9514        a.matvec(&x, &mut ea, Some(&pool));
9515        b.matvec(&x, &mut eb, Some(&pool));
9516        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
9517        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
9518        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
9519        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
9520    }
9521
9522    /// Batched q4/vbit matmat must equal per-position matvec calls
9523    /// exactly (the fallback it replaced) — same kernels, same order.
9524    #[test]
9525    fn batched_matmat_equals_per_position_matvec() {
9526        let (rows, cols, b) = (8, 64, 5);
9527        // q4 blob.
9528        let groups = rows * cols / GROUP_SIZE;
9529        let mut q4 = Vec::new();
9530        for i in 0..groups * 16 {
9531            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9532        }
9533        for g in 0..groups {
9534            q4.extend_from_slice(
9535                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9536            );
9537        }
9538        // vbit blob (mixed widths incl. 8).
9539        let ng = cols / GROUP_SIZE;
9540        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
9541        let mut vb = bits.clone();
9542        for g in 0..rows * ng {
9543            vb.extend_from_slice(
9544                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
9545            );
9546        }
9547        for r in 0..rows {
9548            let bw = bits[r] as usize;
9549            let (mut acc, mut nb) = (0u64, 0usize);
9550            let mut rowbytes = Vec::new();
9551            for i in 0..cols {
9552                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9553                acc = (acc << bw) | v;
9554                nb += bw;
9555                while nb >= 8 {
9556                    nb -= 8;
9557                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9558                }
9559            }
9560            if nb > 0 {
9561                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9562            }
9563            vb.extend_from_slice(&rowbytes);
9564        }
9565        let offsets = vbit_row_offsets(&vb, rows, cols);
9566
9567        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9568
9569        // q4: batch vs singles.
9570        let mut got = vec![0f32; b * rows];
9571        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
9572        for bi in 0..b {
9573            let mut expect = vec![0f32; rows];
9574            q4matvec(
9575                &q4,
9576                &xs[bi * cols..(bi + 1) * cols],
9577                rows,
9578                cols,
9579                &mut expect,
9580                None,
9581            );
9582            assert_eq!(
9583                &got[bi * rows..(bi + 1) * rows],
9584                &expect[..],
9585                "q4 batch pos {bi}"
9586            );
9587        }
9588
9589        // vbit: batch vs singles.
9590        let mut got = vec![0f32; b * rows];
9591        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
9592        for bi in 0..b {
9593            let mut expect = vec![0f32; rows];
9594            vbitmatvec(
9595                &vb,
9596                &offsets,
9597                &xs[bi * cols..(bi + 1) * cols],
9598                rows,
9599                cols,
9600                &mut expect,
9601                None,
9602            );
9603            assert_eq!(
9604                &got[bi * rows..(bi + 1) * rows],
9605                &expect[..],
9606                "vbit batch pos {bi}"
9607            );
9608        }
9609    }
9610
9611    /// q4_tiled kernels must produce BIT-identical outputs to the q4
9612    /// split kernels on the same values (same ints, same order — only
9613    /// the byte placement differs).
9614    #[test]
9615    fn q4_tiled_matches_q4_block_bitexact() {
9616        let (rows, cols, b) = (8usize, 128usize, 3usize);
9617        let groups = rows * cols / GROUP_SIZE;
9618        let mut split = Vec::with_capacity(groups * 18);
9619        for i in 0..groups * 16 {
9620            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9621        }
9622        for g in 0..groups {
9623            split.extend_from_slice(
9624                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9625            );
9626        }
9627        // Re-tile: [scale][nibbles] per group.
9628        let (packed, scales) = split.split_at(groups * 16);
9629        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
9630        for g in 0..groups {
9631            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
9632            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
9633        }
9634
9635        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9636        x1[9] = 250.0; // exercise the outlier path
9637        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9638
9639        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
9640        q4matvec(&split, &x1, rows, cols, &mut a, None);
9641        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
9642        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
9643
9644        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9645        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
9646        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
9647        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
9648        assert_eq!(a1, t1);
9649        assert_eq!(a2, t2);
9650
9651        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9652        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
9653        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
9654        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
9655        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
9656    }
9657
9658    /// q4 SDOT outlier correction: a single huge activation channel
9659    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
9660    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
9661    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
9662    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
9663    /// can never qualify (8² = n).
9664    #[test]
9665    fn q4matvec_sdot_outlier_exact() {
9666        let (rows, cols) = (4, 128);
9667        let groups = rows * cols / GROUP_SIZE;
9668        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9669        for i in 0..groups * 16 {
9670            bytes.push(((i * 11 + 5) % 256) as u8);
9671        }
9672        for g in 0..groups {
9673            let s = 0.02 + 0.002 * g as f32;
9674            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9675        }
9676        let mut x: Vec<f32> = (0..cols)
9677            .map(|i| match i % 3 {
9678                0 => 1.0,
9679                1 => -1.0,
9680                _ => 0.0,
9681            })
9682            .collect();
9683        x[17] = 300.0; // ≫ 8·rms → outlier channel
9684
9685        let mut reference = vec![0.0f32; rows * cols];
9686        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9687        let mut expect = vec![0.0f32; rows];
9688        for r in 0..rows {
9689            expect[r] = reference[r * cols..(r + 1) * cols]
9690                .iter()
9691                .zip(&x)
9692                .map(|(w, xv)| w * xv)
9693                .sum();
9694        }
9695        let mut got = vec![0.0f32; rows];
9696        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9697        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9698        for r in 0..rows {
9699            assert!(
9700                (got[r] - expect[r]).abs() < 2e-3 * scale,
9701                "row {r}: {} vs {} (outlier term must be exact)",
9702                got[r],
9703                expect[r]
9704            );
9705        }
9706    }
9707
9708    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
9709    /// including the ternary zero level and the binary-searched outlier
9710    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
9711    #[test]
9712    fn q1t_matvec_matches_reference() {
9713        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
9714        let (rows, cols) = (3usize, 64usize); // gpr = 2
9715        let gpr = cols / GROUP_SIZE;
9716        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
9717        // Overlay (must be sorted by flat index): a few spikes across rows.
9718        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
9719        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
9720        let mut bytes = Vec::new();
9721        for r in 0..rows {
9722            for g in 0..gpr {
9723                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
9724                let mut c = [0u8; 7];
9725                for k in 0..GROUP_SIZE {
9726                    // Encoder invariant: code 0 at outlier positions.
9727                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
9728                        0
9729                    } else {
9730                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
9731                    };
9732                    cortiq_core::quant::q1t_pack(&mut c, k, code);
9733                }
9734                bytes.extend_from_slice(&c);
9735            }
9736        }
9737        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
9738        // row (outliers are sorted by flat index → already grouped by row).
9739        let mut row_ptr = vec![0u32; rows + 1];
9740        for &(idx, _) in &outliers {
9741            row_ptr[idx as usize / cols + 1] += 1;
9742        }
9743        for r in 0..rows {
9744            row_ptr[r + 1] += row_ptr[r];
9745        }
9746        for &p in &row_ptr {
9747            bytes.extend_from_slice(&p.to_le_bytes());
9748        }
9749        for &(idx, v) in &outliers {
9750            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
9751            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
9752        }
9753
9754        let mut refw = vec![0f32; rows * cols];
9755        dequant_q1t(&bytes, rows, cols, &mut refw);
9756        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
9757        // x exactly and matches the f32 reference (same trick as the q1 test).
9758        let x: Vec<f32> = (0..cols)
9759            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9760            .collect();
9761        let mut expect = vec![0f32; rows];
9762        for r in 0..rows {
9763            let mut a = 0.0f32;
9764            for j in 0..cols {
9765                a += refw[r * cols + j] * x[j];
9766            }
9767            expect[r] = a;
9768        }
9769        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
9770        let mut got = vec![0f32; rows];
9771        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
9772        for r in 0..rows {
9773            assert!(
9774                (got[r] - expect[r]).abs() < tol(expect[r]),
9775                "row {r}: {} vs {}",
9776                got[r],
9777                expect[r]
9778            );
9779        }
9780        // matmat (b=2, f32 decode path) must agree too.
9781        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
9782        let mut gm = vec![0f32; 2 * rows];
9783        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
9784        for r in 0..rows {
9785            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
9786            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
9787        }
9788        // Fused pair (q1t_matvec2) must equal two single matvecs
9789        // bit-for-bit: same unpack, same group order, same f32
9790        // accumulation per stream. Distinct x2 exercises both lanes.
9791        let xb: Vec<f32> = (0..cols)
9792            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
9793            .collect();
9794        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9795        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
9796        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
9797        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9798        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
9799        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
9800        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
9801    }
9802
9803    /// Pair == 2×matvec with an ODD group count (the kernel's tail
9804    /// group) and no overlay section.
9805    #[test]
9806    fn q1t_matvec2_odd_gpr_matches_singles() {
9807        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
9808        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
9809        let gpr = cols / GROUP_SIZE;
9810        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
9811        for r in 0..rows {
9812            for g in 0..gpr {
9813                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
9814                let mut c = [0u8; 7];
9815                for k in 0..GROUP_SIZE {
9816                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
9817                }
9818                bytes.extend_from_slice(&c);
9819            }
9820        }
9821        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
9822        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
9823        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9824        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9825        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
9826        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9827        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9828        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
9829        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
9830    }
9831
9832    // Speed A/B: fused pair (one unpack, two streams) vs two single
9833    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
9834    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
9835    #[test]
9836    #[ignore]
9837    fn q1t_matvec2_speed() {
9838        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
9839        use std::time::Instant;
9840        let (rows, cols) = (8192usize, 4096usize);
9841        let gpr = cols / GROUP_SIZE;
9842        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
9843        for r in 0..rows {
9844            for g in 0..gpr {
9845                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
9846                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
9847                let mut c = [0u8; 7];
9848                for k in 0..GROUP_SIZE {
9849                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
9850                }
9851                bytes.extend_from_slice(&c);
9852            }
9853        }
9854        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
9855        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
9856        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9857        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
9858        // Warm both paths once.
9859        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9860        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9861        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
9862        for _ in 0..8 {
9863            let t0 = Instant::now();
9864            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
9865            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
9866            let t1 = Instant::now();
9867            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
9868            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
9869            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
9870        }
9871        assert_eq!(p1, s1);
9872        assert_eq!(p2, s2);
9873        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
9874    }
9875
9876    // Speed A/B: the base-3-division decode (what the packing commit left in
9877    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
9878    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
9879    #[test]
9880    #[ignore]
9881    fn q1t_matvec_speed() {
9882        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
9883        use std::time::Instant;
9884        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
9885        let gpr = cols / GROUP_SIZE;
9886        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
9887        for r in 0..rows {
9888            for g in 0..gpr {
9889                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
9890                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
9891                let mut c = [0u8; 7];
9892                for k in 0..GROUP_SIZE {
9893                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
9894                }
9895                bytes.extend_from_slice(&c);
9896            }
9897        }
9898        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
9899        let mut row_ptr = vec![0u32; rows + 1];
9900        let mut idx = 0usize;
9901        while idx < n {
9902            row_ptr[idx / cols + 1] += 1;
9903            idx += stride;
9904        }
9905        for r in 0..rows {
9906            row_ptr[r + 1] += row_ptr[r];
9907        }
9908        for &p in &row_ptr {
9909            bytes.extend_from_slice(&p.to_le_bytes());
9910        }
9911        let mut idx = 0usize;
9912        while idx < n {
9913            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
9914            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
9915            idx += stride;
9916        }
9917        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
9918        // reference (the A/B is a timing check; values must still agree).
9919        let x: Vec<f32> = (0..cols)
9920            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9921            .collect();
9922        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
9923
9924        // "before": base-3 division decode into a buffer, then dot.
9925        let slow = |out: &mut [f32]| {
9926            let mut buf = vec![0f32; cols];
9927            for r in 0..rows {
9928                for g in 0..gpr {
9929                    let off = (r * gpr + g) * Q1T_TILE;
9930                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9931                    let codes = &bytes[off + 2..off + Q1T_TILE];
9932                    for k in 0..GROUP_SIZE {
9933                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
9934                            1 => s,
9935                            2 => -s,
9936                            _ => 0.0,
9937                        };
9938                    }
9939                }
9940                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
9941                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
9942            }
9943        };
9944        let iters = 5;
9945        let mut a = vec![0f32; rows];
9946        slow(&mut a); // warm
9947        let t = Instant::now();
9948        for _ in 0..iters {
9949            slow(&mut a);
9950        }
9951        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
9952
9953        let mut b = vec![0f32; rows];
9954        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
9955        let t = Instant::now();
9956        for _ in 0..iters {
9957            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
9958        }
9959        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
9960
9961        for r in 0..rows {
9962            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
9963        }
9964        println!(
9965            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
9966            slow_ms / fast_ms
9967        );
9968    }
9969}
cortiq_engine/qtensor.rs

cortiq_engine/
qtensor.rs
